From f5b718cc10854ef6bc4d2d7966035f10d9677946 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Fri, 30 Oct 2026 14:06:44 -0700 Subject: [PATCH 1/8] [None][feat] Cosmos3: resolve checkpoint, serving and request settings through one precedence table Signed-off-by: Igor Shovkun --- .../visual_gen/models/cosmos3/action.py | 7 +- .../models/cosmos3/config_resolver.py | 409 ++++++++++++ .../visual_gen/models/cosmos3/defaults.py | 240 ++----- .../models/cosmos3/pipeline_cosmos3.py | 177 +++-- .../models/cosmos3/request_resolver.py | 487 ++++++++++++++ .../models/cosmos3/transformer_cosmos3.py | 18 +- .../integration/test_lists/test-db/l0_cpu.yml | 2 + .../visual_gen/cosmos3_request_golden.py | 154 +++++ .../_torch/visual_gen/test_cosmos3_action.py | 4 +- .../test_cosmos3_config_resolver.py | 218 ++++++ .../visual_gen/test_cosmos3_distilled.py | 39 +- .../_torch/visual_gen/test_cosmos3_edge.py | 9 +- .../visual_gen/test_cosmos3_pipeline.py | 8 +- .../test_cosmos3_request_resolver.py | 107 +++ .../_torch/visual_gen/test_data/README.md | 12 + .../test_data/cosmos3_checkpoint_configs.json | 226 +++++++ .../test_data/cosmos3_request_defaults.json | 631 ++++++++++++++++++ 17 files changed, 2443 insertions(+), 305 deletions(-) create mode 100644 tensorrt_llm/_torch/visual_gen/models/cosmos3/config_resolver.py create mode 100644 tensorrt_llm/_torch/visual_gen/models/cosmos3/request_resolver.py create mode 100644 tests/unittest/_torch/visual_gen/cosmos3_request_golden.py create mode 100644 tests/unittest/_torch/visual_gen/test_cosmos3_config_resolver.py create mode 100644 tests/unittest/_torch/visual_gen/test_cosmos3_request_resolver.py create mode 100644 tests/unittest/_torch/visual_gen/test_data/cosmos3_checkpoint_configs.json create mode 100644 tests/unittest/_torch/visual_gen/test_data/cosmos3_request_defaults.json diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py index 2bb2284f3b55..fe13630c1ade 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py @@ -95,8 +95,11 @@ def resolve_raw_action_dim( return widths.pop() if len(widths) == 1 else None -# Canonical ``W,H`` labels; every action canvas is one of these bucket shapes. -ACTION_ASPECT_RATIO_LABELS = ("1,1", "4,3", "3,4", "16,9", "9,16") +# The reference's ``W,H`` aspect buckets: every canvas is one of these shapes, and a +# request's bucket string is carried into the caption verbatim, so a resolved frame +# size maps back to the nearest bucket rather than to its exact reduced ratio +# (832x480 reduces to "15,26", which the model never saw). +COSMOS3_ASPECT_RATIO_BUCKETS = ("1,1", "4,3", "3,4", "16,9", "9,16") VIDEO_RES_SIZE_INFO: dict[str, dict[str, tuple[int, int]]] = { diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/config_resolver.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/config_resolver.py new file mode 100644 index 000000000000..f2da4db2f230 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/config_resolver.py @@ -0,0 +1,409 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""One place that turns a Cosmos3 checkpoint and the serving config into settings. + +A Cosmos3 export carries the same facts in several files, written by different +exporters that are not kept in sync: the training config (``config.json``), the +diffusers pipeline index (``model_index.json``), the export metadata +(``checkpoint.json``) and the component configs (``transformer/``, ``scheduler/``, +``vae/``). :data:`RULES` names, for every setting, the sources in priority order +(the first one present wins) and what a disagreeing lower source means. +:class:`Cosmos3ConfigResolver` applies the table and returns a frozen +:class:`Cosmos3Config` whose every value remembers where it came from. + +A setting no source provides and no default covers is an error. Every +:class:`Default` in the table is a value agreed on for checkpoints that omit the key. +""" + +from __future__ import annotations + +import json +import math +import os +from collections.abc import Mapping +from dataclasses import dataclass +from enum import Enum +from typing import Any, Optional + +from tensorrt_llm.logger import logger + +CONFIG = "config.json" +CHECKPOINT = "checkpoint.json" +MODEL_INDEX = "model_index.json" +TRANSFORMER = "transformer/config.json" +SCHEDULER = "scheduler/scheduler_config.json" +VAE = "vae/config.json" + +_MISSING = object() + + +class Conflict(Enum): + """What a lower-priority source holding a different value means.""" + + ERROR = "error" + WARN = "warn" + IGNORE = "ignore" + + +@dataclass(frozen=True) +class Key: + """A dotted key path in one checkpoint file.""" + + file: str + path: str + + @property + def label(self) -> str: + return f"{self.file}:{self.path}" + + +@dataclass(frozen=True) +class Serving: + """An attribute path in the serving config (``VisualGenArgs`` as the pipeline sees it).""" + + path: str + + @property + def label(self) -> str: + return f"serving:{self.path}" + + +@dataclass(frozen=True) +class Default: + """The value used when no source declares the setting.""" + + value: Any + why: str + + @property + def label(self) -> str: + return f"default ({self.why})" + + +@dataclass(frozen=True) +class Rule: + sources: tuple + conflict: Conflict = Conflict.WARN + required: bool = True + """``False``: an unresolved setting is ``None`` instead of an error.""" + nullable: bool = False + """``True``: a source holding ``null`` declares the setting (e.g. "no window"); + otherwise ``null`` counts as absent.""" + + +@dataclass(frozen=True) +class Resolved: + value: Any + source: str + disagreeing: tuple[str, ...] = () + """Lower-priority sources that held a different value.""" + + +# Settings every Cosmos3 checkpoint resolves. Priority: serving config, then the +# training config (config.json), then the component configs, then the export +# metadata and the diffusers index. +RULES: dict[str, Rule] = { + "architecture.position_embedding_type": Rule( + ( + Key(TRANSFORMER, "position_embedding_type"), + Default("unified_3d_mrope", "exports that predate the key"), + ) + ), + "architecture.max_position_embeddings": Rule( + ( + Key(TRANSFORMER, "max_position_embeddings"), + Default(262144, "exports that predate the key"), + ) + ), + "vae.temporal_compression": Rule( + ( + Key(VAE, "scale_factor_temporal"), + Key(TRANSFORMER, "temporal_compression_factor"), + Default(4, "exports that predate the key"), + ), + conflict=Conflict.ERROR, + ), + "vae.spatial_compression": Rule((Key(VAE, "scale_factor_spatial"),), conflict=Conflict.ERROR), + "audio.temporal_compression": Rule( + ( + Key(TRANSFORMER, "temporal_compression_factor_sound"), + Default(1, "exports that predate the key"), + ) + ), + "prompt.use_system_prompt": Rule( + ( + Key(CONFIG, "model.config.vlm_config.use_system_prompt"), + Key(MODEL_INDEX, "default_use_system_prompt"), + ) + ), + "sampling.native_flow_schedule": Rule( + ( + Key(MODEL_INDEX, "use_native_flow_schedule"), + Default(False, "only Edge exports declare the native grid"), + ) + ), + "sampling.fixed_sigmas": Rule( + (Key(SCHEDULER, "fixed_step_sampler_config.t_list"),), required=False + ), + "action.policy": Rule((Key(CHECKPOINT, "policy"),), required=False), +} + + +@dataclass(frozen=True) +class Cosmos3ArchitectureSettings: + position_embedding_type: str + max_position_embeddings: int + + +@dataclass(frozen=True) +class Cosmos3VAESettings: + temporal_compression: int + """Pixel frames per latent frame after the first, which is encoded alone.""" + spatial_compression: int + + +@dataclass(frozen=True) +class Cosmos3AudioSettings: + temporal_compression: int + + +@dataclass(frozen=True) +class Cosmos3SamplingSettings: + fixed_sigmas: Optional[tuple[float, ...]] + """The distilled schedule (``t_list``); ``None`` on base checkpoints.""" + native_flow_schedule: bool + + +@dataclass(frozen=True) +class Cosmos3PromptSettings: + use_system_prompt: bool + + +@dataclass(frozen=True) +class Cosmos3ActionSettings: + policy: Optional[Mapping[str, Any]] + """The ``checkpoint.json`` policy block, unnormalized.""" + + +@dataclass(frozen=True) +class Cosmos3Config: + architecture: Cosmos3ArchitectureSettings + vae: Cosmos3VAESettings + audio: Cosmos3AudioSettings + sampling: Cosmos3SamplingSettings + prompt: Cosmos3PromptSettings + action: Cosmos3ActionSettings + resolved: Mapping[str, Resolved] + + def source_of(self, name: str) -> str: + return self.resolved[name].source + + def describe(self) -> str: + lines = ["Cosmos3 config:"] + for name, r in sorted(self.resolved.items()): + value = r.value + if isinstance(value, Mapping): + value = f"<{len(value)} entries>" + line = f" {name} = {value!r} [{r.source}]" + if r.disagreeing: + line += f" (disagrees: {', '.join(r.disagreeing)})" + lines.append(line) + return "\n".join(lines) + + +def _step(obj: Any, key: str) -> Any: + if isinstance(obj, Mapping): + return obj.get(key, _MISSING) + if obj is None or isinstance(obj, (str, int, float, bool, list, tuple)): + return _MISSING + return getattr(obj, key, _MISSING) + + +def _walk(doc: Any, path: str) -> Any: + node = doc + for part in path.split("."): + node = _step(node, part) + if node is _MISSING: + return _MISSING + return node + + +def same_value(a: Any, b: Any) -> bool: + """Whether two settings agree: bools by identity, numbers within 1e-8, lists and + mappings elementwise.""" + if isinstance(a, bool) or isinstance(b, bool): + return a is b or a == b + if isinstance(a, (int, float)) and isinstance(b, (int, float)): + return math.isclose(a, b, rel_tol=0.0, abs_tol=1e-8) + if isinstance(a, (list, tuple)) and isinstance(b, (list, tuple)): + return len(a) == len(b) and all(same_value(x, y) for x, y in zip(a, b)) + if isinstance(a, Mapping) and isinstance(b, Mapping): + return a.keys() == b.keys() and all(same_value(a[k], b[k]) for k in a) + return a == b + + +class Cosmos3ConfigResolver: + """Applies :data:`RULES` to a checkpoint directory and the serving config. + + ``checkpoint_dir`` may be ``None``: only ``pretrained_config`` (standing in for + ``transformer/config.json``) and ``serving`` are consulted then. + """ + + def __init__( + self, + checkpoint_dir: Optional[str], + pretrained_config: Any = None, + serving: Any = None, + ) -> None: + self._dir = checkpoint_dir + self._pretrained = pretrained_config + self._serving = serving + self._docs: dict[str, Any] = {} + + def _doc(self, file: str) -> Any: + if file == TRANSFORMER and self._pretrained is not None: + return self._pretrained + if file not in self._docs: + doc = None + if self._dir is not None: + path = os.path.join(self._dir, file) + if os.path.exists(path): + with open(path) as f: + doc = json.load(f) + self._docs[file] = doc + return self._docs[file] + + def _lookup(self, source: Any) -> Any: + if isinstance(source, Key): + doc = self._doc(source.file) + return _MISSING if doc is None else _walk(doc, source.path) + if isinstance(source, Serving): + return _MISSING if self._serving is None else _walk(self._serving, source.path) + if isinstance(source, Default): + return source.value + raise TypeError(f"unknown source {source!r}") + + def resolve_one(self, name: str, rule: Rule) -> Optional[Resolved]: + winner: Optional[Resolved] = None + disagreeing = [] + for source in rule.sources: + if winner is not None and isinstance(source, Default): + break + value = self._lookup(source) + if value is _MISSING or (value is None and not rule.nullable): + continue + if value is None and isinstance(source, Serving): + continue # an unset serving field defers to the checkpoint + if winner is None: + winner = Resolved(value, source.label) + if isinstance(source, Serving): + break # an override is meant to differ from the checkpoint + elif not same_value(value, winner.value): + disagreeing.append(f"{source.label}={value!r}") + if winner is None: + if rule.required: + labels = ", ".join(s.label for s in rule.sources) + raise ValueError(f"Cosmos3 config: cannot resolve {name}; none of {labels} is set.") + return None + if disagreeing: + message = ( + f"Cosmos3 config: {name} = {winner.value!r} from {winner.source}, but " + + "; ".join(disagreeing) + ) + if rule.conflict is Conflict.ERROR: + raise ValueError(message) + if rule.conflict is Conflict.WARN: + logger.warning(message) + return Resolved(winner.value, winner.source, tuple(disagreeing)) + + def resolve(self) -> Cosmos3Config: + resolved: dict[str, Resolved] = {} + for name, rule in RULES.items(): + r = self.resolve_one(name, rule) + if r is not None: + resolved[name] = r + + def get(name: str) -> Any: + r = resolved.get(name) + return None if r is None else r.value + + sigmas = get("sampling.fixed_sigmas") + if sigmas is not None and ( + any(not 0.0 < float(s) <= 1.0 for s in sigmas) + or [float(s) for s in sigmas] != sorted((float(s) for s in sigmas), reverse=True) + ): + raise ValueError( + f"Cosmos3 config: fixed_step_sampler_config.t_list must decrease within (0, 1]: " + f"{list(sigmas)}." + ) + return Cosmos3Config( + architecture=Cosmos3ArchitectureSettings( + position_embedding_type=str(get("architecture.position_embedding_type")), + max_position_embeddings=int(get("architecture.max_position_embeddings")), + ), + vae=Cosmos3VAESettings( + temporal_compression=int(get("vae.temporal_compression")), + spatial_compression=int(get("vae.spatial_compression")), + ), + audio=Cosmos3AudioSettings(temporal_compression=int(get("audio.temporal_compression"))), + sampling=Cosmos3SamplingSettings( + fixed_sigmas=None if not sigmas else tuple(float(s) for s in sigmas), + native_flow_schedule=bool(get("sampling.native_flow_schedule")), + ), + prompt=Cosmos3PromptSettings(use_system_prompt=bool(get("prompt.use_system_prompt"))), + action=Cosmos3ActionSettings(policy=get("action.policy")), + resolved=resolved, + ) + + +ARCHITECTURE_ATTRIBUTES = { + "position_embedding_type": "architecture.position_embedding_type", + "max_position_embeddings": "architecture.max_position_embeddings", + "temporal_compression_factor": "vae.temporal_compression", + "temporal_compression_factor_sound": "audio.temporal_compression", +} +"""``pretrained_config`` attributes the transformer reads, and the setting each holds.""" + + +def architecture_defaults() -> dict[str, Any]: + """The table's default for each attribute in :data:`ARCHITECTURE_ATTRIBUTES`.""" + out = {} + for attr, name in ARCHITECTURE_ATTRIBUTES.items(): + out[attr] = next(s.value for s in RULES[name].sources if isinstance(s, Default)) + return out + + +def fill_pretrained_config(pretrained_config: Any) -> Any: + """Set the transformer-construction attributes of ``pretrained_config`` that the + export omits, from :data:`RULES`. Idempotent; returns the config. + + Runs before the pipeline loads its components, so only ``pretrained_config`` + itself is consulted. + """ + resolver = Cosmos3ConfigResolver(None, pretrained_config) + for attr, name in ARCHITECTURE_ATTRIBUTES.items(): + if getattr(pretrained_config, attr, None) is not None: + continue + rule = RULES[name] + transformer_only = Rule( + tuple(s for s in rule.sources if not isinstance(s, Key) or s.file == TRANSFORMER), + rule.conflict, + rule.required, + rule.nullable, + ) + r = resolver.resolve_one(name, transformer_only) + if r is not None: + setattr(pretrained_config, attr, r.value) + return pretrained_config diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index 18ec9b117671..bc27c35ce241 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -23,8 +23,8 @@ from tensorrt_llm._torch.visual_gen.models.cosmos3.action import ( COSMOS3_ACTION_RESOLUTIONS, EMBODIMENT_TO_DOMAIN_ID, + VIDEO_RES_SIZE_INFO, normalize_action_resolution, - resolve_raw_action_dim, ) from tensorrt_llm._torch.visual_gen.pipeline import ExtraParamSchema from tensorrt_llm.inputs.media_io import sniff_media_kind @@ -36,37 +36,6 @@ # Cosmos3 output resolution buckets keyed by target level, then aspect ratio; # each value is (width, height). A source frame maps onto the bucket whose # aspect ratio is closest (see ``find_closest_target_size`` in ``transfer.py``). -VIDEO_RES_SIZE_INFO = { - "256": { - "1,1": (256, 256), - "4,3": (320, 256), - "3,4": (256, 320), - "16,9": (320, 192), - "9,16": (192, 320), - }, - "480": { - "1,1": (640, 640), - "4,3": (736, 544), - "3,4": (544, 736), - "16,9": (832, 480), - "9,16": (480, 832), - }, - "704": { - "1,1": (960, 960), - "4,3": (1088, 832), - "3,4": (832, 1088), - "16,9": (1280, 704), - "9,16": (704, 1280), - }, - "720": { - "1,1": (960, 960), - "4,3": (1104, 832), - "3,4": (832, 1104), - "16,9": (1280, 720), - "9,16": (720, 1280), - }, -} - # The default video resolution is the 720p 16:9 bucket, ``(width, height)``. _DEFAULT_VIDEO_W, _DEFAULT_VIDEO_H = VIDEO_RES_SIZE_INFO["720"]["16,9"] COSMOS3_720P_PARAMS = { @@ -557,166 +526,71 @@ def resolve_domain_action_config( action_fps: float | None = None, use_state: bool | None = None, checkpoint_policy_defaults: Mapping[str, Any] | None = None, + embodiments: Mapping[str, Any] | None = None, + input_action_dim: int | None = None, ) -> dict[str, Any]: - """Merge user action params with domain presets and generic fallbacks. - - A recognized ``domain_name`` (or a uniquely mapped ``domain_id``) fills - whichever of ``action_chunk_size``, ``action_resolution`` and ``frame_rate`` - the caller left unset; an explicit value wins but is reported in - ``warnings`` when it differs from the preset. ``num_frames`` is derived as - ``action_chunk_size + 1`` and is never a preset field. + """The action fields a request runs with, by the request rules + (``request_resolver.ACTION_RULES``): the caller's value, else the checkpoint policy + block when its domain is the request's, else the checkpoint's embodiment table + (``embodiments``) or the code's width table, else the domain preset, else the + generic action defaults. ``num_frames`` is always ``action_chunk_size + 1``. A + caller's value that differs from the recommendation is reported in ``warnings``. """ - checkpoint_policy_defaults = dict(checkpoint_policy_defaults or {}) - checkpoint_domain = checkpoint_policy_defaults.get("domain_name") - warnings: list[str] = [] - - if checkpoint_domain is not None: - checkpoint_domain = str(checkpoint_domain).strip().lower() - explicit_domain_name = ( - str(domain_name).strip().lower() - if domain_name is not None and str(domain_name).strip() - else None - ) - checkpoint_domain_id = EMBODIMENT_TO_DOMAIN_ID.get(checkpoint_domain) - domain_id_matches = domain_id is not None and checkpoint_domain_id == int(domain_id) - if explicit_domain_name is None and (domain_id is None or domain_id_matches): - domain_name = checkpoint_domain - elif explicit_domain_name != checkpoint_domain: - warnings.append( - f"Cosmos3 domain_name={domain_name!r}, domain_id={domain_id!r} overrides " - f"checkpoint policy domain_name={checkpoint_domain!r}; checkpoint domain " - "defaults will not be applied." - ) - - policy_domain_matches = checkpoint_domain is not None and ( - (domain_name is not None and str(domain_name).strip().lower() == checkpoint_domain) - or ( - domain_name is None - and domain_id is not None - and EMBODIMENT_TO_DOMAIN_ID.get(checkpoint_domain) == int(domain_id) - ) + from .request_resolver import RequestContext, embodiment_entry, policy_domain, resolve_action + + policy_defaults = dict(checkpoint_policy_defaults or {}) + domain_name, policy_active, warnings = policy_domain(policy_defaults, domain_name, domain_id) + ctx = RequestContext( + family="", + action_mode="action", + policy_defaults=policy_defaults, + policy_active=policy_active, + domain_name=domain_name, + domain_id=None if domain_id is None else int(domain_id), + embodiment=embodiment_entry(embodiments or {}, input_action_dim, domain_name, domain_id), + checkpoint_has_embodiments=bool(embodiments), ) - policy_defaults = checkpoint_policy_defaults if policy_domain_matches else {} - - preset_key = canonical_domain_preset_key(domain_name, domain_id) - preset = COSMOS3_DOMAIN_PRESETS.get(preset_key) if preset_key else None - - domain_requested = (domain_name is not None and str(domain_name).strip() != "") or ( - domain_id is not None and str(domain_id).strip() not in {"", "0"} + resolved = resolve_action( + ctx, + raw_action_dim=raw_action_dim, + action_chunk_size=action_chunk_size, + action_resolution=action_resolution, + frame_rate=frame_rate, + action_fps=action_fps, + use_state=use_state, ) - if domain_requested and preset is None: - warnings.append( - "Cosmos3 action domain preset was not found for " - f"domain_name={domain_name!r}, domain_id={domain_id!r}; " - "using generic action defaults for omitted fields." - ) - - def _resolve_field( - field: str, - current: Any, - *, - fallback: Any = None, - ) -> Any: - recommended = policy_defaults.get(field) - if recommended is None: - recommended = preset.get(field) if preset else None - if current is not None: - if recommended is not None and current != recommended: - warnings.append( - f"Cosmos3 {field}={current} differs from recommended " - f"{recommended} for domain {preset_key!r}." - ) - return current - if recommended is not None: - return recommended - return fallback - - # The action width is canonical per embodiment, so it comes from the - # embodiment table rather than the (alias-shared) sampling preset. - canonical_raw_action_dim = resolve_raw_action_dim(domain_name=domain_name, domain_id=domain_id) - recommended_raw_action_dim = policy_defaults.get("raw_action_dim", canonical_raw_action_dim) - if raw_action_dim is not None: - if recommended_raw_action_dim is not None and int(raw_action_dim) != int( - recommended_raw_action_dim - ): - warnings.append( - f"Cosmos3 raw_action_dim={raw_action_dim} differs from the recommended width " - f"{recommended_raw_action_dim} for domain_name={domain_name!r}." - ) - resolved_raw_action_dim = raw_action_dim - else: - resolved_raw_action_dim = recommended_raw_action_dim - if domain_requested and recommended_raw_action_dim is None: - warnings.append( - "Cosmos3 has no canonical action width for " - f"domain_name={domain_name!r}, domain_id={domain_id!r}; " - "pass raw_action_dim explicitly for policy/inverse_dynamics." - ) - resolved_chunk = _resolve_field( - "action_chunk_size", - action_chunk_size, - fallback=COSMOS3_ACTION_PARAMS["action_chunk_size"], - ) - resolved_resolution = normalize_action_resolution( - _resolve_field( - "action_resolution", - action_resolution, - fallback=480, - ) - ) - resolved_frame_rate = _resolve_field( - "frame_rate", - frame_rate, - fallback=COSMOS3_ACTION_PARAMS["frame_rate"], - ) - # Always derived: an action clip is the chunk plus its initial frame. Both - # references fix this, and diffusers rejects a caller-supplied num_frames - # for action runs outright, so a preset must not pin it independently of - # an overridden action_chunk_size. - resolved_num_frames = int(resolved_chunk) + 1 - resolved_action_fps = ( - float(action_fps) - if action_fps is not None - else float(policy_defaults.get("action_fps", resolved_frame_rate)) - ) - recommended_use_state = bool(policy_defaults.get("use_state", False)) - resolved_use_state = recommended_use_state if use_state is None else bool(use_state) - if ( - use_state is not None - and "use_state" in policy_defaults - and resolved_use_state != recommended_use_state - ): - warnings.append( - f"Cosmos3 use_state={resolved_use_state} differs from checkpoint policy " - f"default {recommended_use_state}." - ) - if resolved_raw_action_dim is not None and int(resolved_raw_action_dim) <= 0: - raise ValueError(f"Cosmos3 raw_action_dim must be positive, got {resolved_raw_action_dim}.") - if int(resolved_chunk) <= 0: - raise ValueError(f"Cosmos3 action_chunk_size must be positive, got {resolved_chunk}.") - if float(resolved_frame_rate) <= 0.0: - raise ValueError(f"Cosmos3 frame_rate must be positive, got {resolved_frame_rate}.") - if resolved_action_fps <= 0.0: - raise ValueError(f"Cosmos3 action_fps must be positive, got {resolved_action_fps}.") - if int(resolved_num_frames) <= 0: - raise ValueError(f"Cosmos3 num_frames must be positive, got {resolved_num_frames}.") - + width = resolved["raw_action_dim"] return { "domain_name": domain_name, - "checkpoint_policy_domain_matches": policy_domain_matches, - "raw_action_dim": resolved_raw_action_dim, - "action_chunk_size": int(resolved_chunk), - "action_resolution": resolved_resolution, - "frame_rate": float(resolved_frame_rate), - "action_fps": resolved_action_fps, - "use_state": resolved_use_state, - "num_frames": int(resolved_num_frames), - "preset_key": preset_key, - "warnings": warnings, + "checkpoint_policy_domain_matches": policy_active, + "raw_action_dim": None if width is None else int(width), + "action_chunk_size": int(resolved["action_chunk_size"]), + "action_resolution": resolved["action_resolution"], + "frame_rate": float(resolved["frame_rate"]), + "action_fps": float(resolved["action_fps"]), + "use_state": bool(resolved["use_state"]), + "num_frames": int(resolved["num_frames"]), + "preset_key": ctx.preset_key, + "warnings": warnings + resolved.warnings, + "sources": {name: r.source for name, r in resolved.values.items()}, } -COSMOS3_V2V_DEFAULT_FLOW_SHIFT = 10.0 +# Video-to-video and transfer sample on the shifted flow grid with this shift; the +# other modes carry ``flow_shift`` in their tables above or take the checkpoint's. +COSMOS3_V2V_FLOW_SHIFT = 10.0 + +# Prompt format the checkpoints were trained with. The "give" typo is the training +# text; keep it. Metadata sentences are appended to the caption the same way. +COSMOS3_DEFAULT_SYSTEM_PROMPT = ( + "You are a helpful assistant who will generate videos from a give prompt." +) +COSMOS3_T2I_SYSTEM_PROMPT = ( + "You are a helpful assistant who will generate images from a give prompt." +) +COSMOS3_DURATION_TEMPLATE = "The video is {duration:.1f} seconds long and is of {fps:.0f} FPS." +COSMOS3_DEFAULT_RESOLUTION_TEMPLATE = "This video is of {height}x{width} resolution." +COSMOS3_IMAGE_RESOLUTION_TEMPLATE = "This image is of {height}x{width} resolution." COSMOS3_EXTRA_SPECS: Dict[str, ExtraParamSchema] = { "use_duration_template": ExtraParamSchema( diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index 1bf29e825979..bd5862904c31 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -55,6 +55,7 @@ def tqdm(iterable, **kwargs): from .action import ( ACTION_MODE_INVERSE_DYNAMICS, ACTION_MODE_POLICY, + COSMOS3_ASPECT_RATIO_BUCKETS, action_reference_frame_step, action_reference_size, action_start_frame_offset, @@ -68,12 +69,18 @@ def tqdm(iterable, **kwargs): resolve_action_size, resolve_domain_id, ) +from .config_resolver import Cosmos3ConfigResolver from .defaults import ( COSMOS3_720P_PARAMS, + COSMOS3_DEFAULT_RESOLUTION_TEMPLATE, + COSMOS3_DEFAULT_SYSTEM_PROMPT, + COSMOS3_DURATION_TEMPLATE, COSMOS3_ENVELOPES, COSMOS3_EXTRA_SPECS, COSMOS3_GENERATION_DEFAULTS, - COSMOS3_V2V_DEFAULT_FLOW_SHIFT, + COSMOS3_IMAGE_RESOLUTION_TEMPLATE, + COSMOS3_T2I_SYSTEM_PROMPT, + COSMOS3_V2V_FLOW_SHIFT, _normalize_condition_video_keep, _normalize_condition_video_latent_indexes, resolve_checkpoint_policy_defaults, @@ -81,6 +88,7 @@ def tqdm(iterable, **kwargs): ) from .guardrails import check_video_safety, download_guardrail_checkpoint from .negative_prompt import COSMOS3_VIDEO_NEGATIVE_PROMPT +from .request_resolver import RequestContext, resolve_generation from .sampling import DISTILLED_GUIDANCE_SCALE, Cosmos3SamplingPolicy, load_scheduler from .sound_tokenizer import LatentAutoEncoderV2 from .transfer import ( @@ -130,14 +138,15 @@ def _resolve_use_system_prompt( is_v2v: bool, is_transfer: bool, ) -> bool: - """Resolve the request override against workflow and checkpoint defaults.""" - if requested is not None: - return bool(requested) - # Transfer opts out for reference parity; vllm-omni's transfer path also - # defaults use_system_prompt to False. - if is_transfer: - return False - return is_v2v or checkpoint_default + """The request rule: the caller's value, else transfer never (reference parity), + else video-to-video always, else the checkpoint's declaration.""" + ctx = RequestContext( + family="", + checkpoint_use_system_prompt=checkpoint_default, + is_v2v=is_v2v, + is_transfer=is_transfer, + ) + return bool(resolve_generation(ctx, use_system_prompt=requested)["use_system_prompt"]) def _validate_request_compatibility( @@ -223,22 +232,11 @@ def _validate_request_compatibility( ) -# NOTE: Intentional typo in "give" instead of "given" to match training setup. -COSMOS3_DEFAULT_SYSTEM_PROMPT = ( - "You are a helpful assistant who will generate videos from a give prompt." -) -COSMOS3_T2I_SYSTEM_PROMPT = ( - "You are a helpful assistant who will generate images from a give prompt." -) -COSMOS3_V2V_FLOW_SHIFT = 10.0 # Fraction of a transfer hint that may be mirror-padding before it is worth a # warning. A few frames of tail ping-pong is normal when clips differ slightly; # beyond this the control is mostly invented and the client likely sent clips of # different videos. CONTROL_LENGTH_MISMATCH_RATIO = 0.10 -COSMOS3_DURATION_TEMPLATE = "The video is {duration:.1f} seconds long and is of {fps:.0f} FPS." -COSMOS3_DEFAULT_RESOLUTION_TEMPLATE = "This video is of {height}x{width} resolution." -COSMOS3_IMAGE_RESOLUTION_TEMPLATE = "This image is of {height}x{width} resolution." TRTLLM_DISABLE_COSMOS3_GUARDRAILS = os.environ.get("TRTLLM_DISABLE_COSMOS3_GUARDRAILS", "0") == "1" @@ -258,13 +256,6 @@ def _validate_request_compatibility( (COSMOS3_GENERATOR_OFFLOAD_COMPONENT,), ) -# ``W,H`` bucket names the reference builds requests from. A request there starts -# as a (resolution, bucket) pair and the bucket string is carried into the prompt -# verbatim; we only ever see the resolved height/width, so map back to the nearest -# bucket. Emitting the exact reduced ratio instead would put a string the model -# never saw in training into the caption (832x480 reduces to "15,26", not "16,9"). -COSMOS3_ASPECT_RATIO_BUCKETS = ("1,1", "4,3", "3,4", "16,9", "9,16") - def _aspect_ratio_bucket(height: int, width: int) -> str: """Nearest reference aspect-ratio bucket for a resolved frame size.""" @@ -491,6 +482,32 @@ def _mode_params( } return params + def _request_context( + self, + output_type: str = "video", + *, + action_mode: Optional[str] = None, + checkpoint_policy_domain_matches: bool = True, + **workflow, + ) -> RequestContext: + normalized_action_mode = normalize_action_mode(action_mode) + policy_defaults = getattr(self, "checkpoint_policy_defaults", {}) or {} + fixed_sigmas = getattr(self.sampling, "fixed_sigmas", None) + return RequestContext( + family=self.family, + mode=str(output_type).lower(), + action_mode=normalized_action_mode, + distilled_steps=None if not fixed_sigmas else len(fixed_sigmas), + policy_defaults=policy_defaults, + policy_active=( + normalized_action_mode == ACTION_MODE_POLICY + and checkpoint_policy_domain_matches + and bool(policy_defaults) + ), + checkpoint_use_system_prompt=getattr(self, "default_use_system_prompt", False), + **workflow, + ) + def _resolve_generation_params( self, output_type: str, @@ -499,24 +516,15 @@ def _resolve_generation_params( checkpoint_policy_domain_matches: bool = True, **values, ) -> dict: - """Fill None values: sampling-policy overrides win, then the mode - table, then the video table (for fields the image table omits).""" - mode_params = self._mode_params( + """Fill ``None`` values by the request rules: the distilled schedule, then the + policy block, then the mode table, then the video table.""" + ctx = self._request_context( output_type, action_mode=action_mode, checkpoint_policy_domain_matches=checkpoint_policy_domain_matches, ) - video_params = self._mode_params("video") - sampling_overrides = self.sampling.generation_default_overrides() - resolved = {} - for key, value in values.items(): - if value is None: - if key in sampling_overrides: - value = sampling_overrides[key] - else: - value = mode_params.get(key, video_params.get(key)) - resolved[key] = value - return resolved + resolved = resolve_generation(ctx, **values) + return {name: resolved[name] for name in values} def _log_envelope_advisory( self, @@ -643,32 +651,16 @@ def load_standard_components( ) -> None: skip_components = skip_components or [] - # Prompting defaults are checkpoint-declared: distilled conversions - # carry ``default_use_system_prompt`` in model_index.json (diffusers' - # distilled blocks default it to True); older checkpoints omit it and - # keep the historical False. - model_index_path = os.path.join(checkpoint_dir, "model_index.json") - if os.path.exists(model_index_path): - with open(model_index_path) as f: - model_index = json.load(f) - self.default_use_system_prompt = bool( - model_index.get("default_use_system_prompt", self.default_use_system_prompt) - ) - self.use_native_flow_schedule = bool( - model_index.get("use_native_flow_schedule", self.use_native_flow_schedule) - ) - - checkpoint_metadata_path = os.path.join(checkpoint_dir, "checkpoint.json") - if os.path.exists(checkpoint_metadata_path): - with open(checkpoint_metadata_path) as f: - checkpoint_metadata = json.load(f) - self.checkpoint_policy_defaults = resolve_checkpoint_policy_defaults( - checkpoint_metadata.get("policy") - ) - if self.checkpoint_policy_defaults: - logger.info( - f"Loaded Cosmos3 checkpoint policy defaults: {self.checkpoint_policy_defaults}" - ) + # The transformer config is read from the checkpoint directory, not from the + # loaded pretrained config, whose omitted keys the transformer has already filled. + self.cosmos3_config = Cosmos3ConfigResolver( + checkpoint_dir, serving=getattr(self, "pipeline_config", None) + ).resolve() + logger.info(self.cosmos3_config.describe()) + cfg = self.cosmos3_config + self.default_use_system_prompt = cfg.prompt.use_system_prompt + self.use_native_flow_schedule = cfg.sampling.native_flow_schedule + self.checkpoint_policy_defaults = resolve_checkpoint_policy_defaults(cfg.action.policy) if self.audio_gen and PipelineComponent.SOUND_TOKENIZER not in skip_components: logger.info("Loading audio tokenizer...") @@ -689,9 +681,8 @@ def load_standard_components( subfolder="text_tokenizer", ) - # Cosmos3 canonical defaults — overwritten if VAE is loaded - self.vae_scale_factor_temporal = 4 - self.vae_scale_factor_spatial = 16 + self.vae_scale_factor_temporal = cfg.vae.temporal_compression + self.vae_scale_factor_spatial = cfg.vae.spatial_compression if PipelineComponent.VAE not in skip_components: logger.info("Loading VAE...") @@ -706,12 +697,6 @@ def load_standard_components( dtype=torch.bfloat16, ) - self.vae_scale_factor_temporal = getattr( - self.vae.config, "scale_factor_temporal", self.vae_scale_factor_temporal - ) - self.vae_scale_factor_spatial = getattr( - self.vae.config, "scale_factor_spatial", self.vae_scale_factor_spatial - ) _validate_temporal_compression(self.transformer, self.vae_scale_factor_temporal) self.transformer.temporal_compression_factor = self.vae_scale_factor_temporal @@ -918,7 +903,7 @@ def _cacheable_flow_shifts(self) -> frozenset: cached = getattr(self, "_cacheable_flow_shifts_cache", None) if cached is not None: return cached - shifts = {COSMOS3_V2V_FLOW_SHIFT, COSMOS3_V2V_DEFAULT_FLOW_SHIFT} + shifts = {COSMOS3_V2V_FLOW_SHIFT} checkpoint_shift = getattr(self.sampling, "checkpoint_flow_shift", None) if checkpoint_shift is not None: shifts.add(float(checkpoint_shift)) @@ -1380,26 +1365,25 @@ def post_step_fn(latents: torch.Tensor, extra_stream_latents): # VAE decode # ========================================================================= - def _decode_latents_raw(self, latents): + def _scale_latents_for_decode(self, latents): + """Latents as the VAE decoder expects them: the checkpoint's per-channel + statistics undone, or its scaling factor.""" latents = latents.to(self.vae.dtype) - if hasattr(self.vae.config, "latents_mean") and hasattr(self.vae.config, "latents_std"): - if not hasattr(self, "_latents_mean"): - self._latents_mean = ( - torch.tensor(self.vae.config.latents_mean) - .view(1, -1, 1, 1, 1) - .to(self.device, self.vae.dtype) - ) - self._latents_std = ( - torch.tensor(self.vae.config.latents_std) - .view(1, -1, 1, 1, 1) - .to(self.device, self.vae.dtype) + stats = getattr(self, "_latents_stats", None) + if stats is None or stats[0].device != latents.device: + stats = tuple( + torch.tensor(values).view(1, -1, 1, 1, 1).to(latents.device, self.vae.dtype) + for values in (self.vae.config.latents_std, self.vae.config.latents_mean) ) - latents = (latents * self._latents_std) + self._latents_mean - else: - scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0) - latents = latents / scaling_factor + self._latents_stats = stats + std, mean = stats + return (latents * std) + mean + scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0) + return latents / scaling_factor + def _decode_latents_raw(self, latents): + latents = self._scale_latents_for_decode(latents) with self.offloader.context_if_requested(PipelineComponent.VAE.value): return self.vae.decode(latents, return_dict=False)[0] @@ -1744,6 +1728,11 @@ def _resolve_request( checkpoint_policy_defaults=checkpoint_policy_defaults, ) checkpoint_policy_domain_matches = action_cfg["checkpoint_policy_domain_matches"] + if self.rank == 0: + logger.debug( + "Cosmos3 action fields: " + + ", ".join(f"{k} [{v}]" for k, v in action_cfg["sources"].items()) + ) mode_params = self._mode_params( output_type, @@ -3012,7 +3001,7 @@ def _forward_transfer( flow_shift_target = float( transfer_config.flow_shift if transfer_config.flow_shift is not None - else COSMOS3_V2V_DEFAULT_FLOW_SHIFT + else COSMOS3_V2V_FLOW_SHIFT ) max_sequence_length = max_sequence_length or COSMOS3_720P_PARAMS["max_sequence_length"] self._guidance_scale = guidance_scale diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/request_resolver.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/request_resolver.py new file mode 100644 index 000000000000..90093a1489a6 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/request_resolver.py @@ -0,0 +1,487 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Per-request resolution of Cosmos3 generation and action settings. + +The engine-level :mod:`config_resolver` settles what the checkpoint and the serving +config say. This module settles what one request runs with: for every field, +:data:`GENERATION_RULES` and :data:`ACTION_RULES` name the sources in priority order, +the caller's value first, then the checkpoint's facts (the distilled schedule, the +policy block, the embodiment table), then the per-mode and per-domain tables, then a +default. The first source holding a value wins; a lower source that disagrees with a +caller's value is reported as a warning (or an error, where the checkpoint is the +authority), and every resolved value remembers its source. +""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field +from typing import Any, Optional + +from .action import EMBODIMENT_TO_DOMAIN_ID, normalize_action_resolution, resolve_raw_action_dim +from .config_resolver import Conflict, Default, Resolved, same_value +from .defaults import ( + COSMOS3_ACTION_PARAMS, + COSMOS3_DOMAIN_PRESETS, + COSMOS3_GENERATION_DEFAULTS, + canonical_domain_preset_key, +) +from .sampling import DISTILLED_GUIDANCE_SCALE + +_MISSING = object() + + +@dataclass(frozen=True) +class RequestContext: + """What a request resolves against: the checkpoint's facts and the request's mode.""" + + family: str + mode: str = "video" + """``"video"`` or ``"image"``; ignored for the table when ``action_mode`` is set.""" + action_mode: Optional[str] = None + distilled_steps: Optional[int] = None + """Length of the fixed schedule on a distilled checkpoint.""" + policy_defaults: Mapping[str, Any] = field(default_factory=dict) + """The normalized ``checkpoint.json`` policy block.""" + policy_active: bool = False + """The policy block applies: policy mode, and the request's domain is the block's.""" + domain_name: Optional[str] = None + domain_id: Optional[int] = None + embodiment: Mapping[str, Any] = field(default_factory=dict) + """The checkpoint's entry for the request's embodiment: ``domain_id``, + ``action_dim``. Empty when the checkpoint declares no embodiment table.""" + checkpoint_has_embodiments: bool = False + checkpoint_use_system_prompt: bool = False + is_v2v: bool = False + is_transfer: bool = False + + @property + def table_mode(self) -> str: + return "action" if self.action_mode is not None else self.mode + + @property + def preset_key(self) -> Optional[str]: + return canonical_domain_preset_key(self.domain_name, self.domain_id) + + +# ----------------------------------------------------------------------------- sources + + +class Given: + """The caller's value for the field being resolved.""" + + label = "request" + + def lookup(self, ctx: RequestContext, field_name: str, given: Mapping[str, Any], values) -> Any: + value = given.get(field_name) + return _MISSING if value is None else value + + +@dataclass(frozen=True) +class Distilled: + """A fact of the distilled schedule: its step count, or the guidance it bakes in.""" + + what: str + + @property + def label(self) -> str: + return f"checkpoint:fixed_step_sampler_config ({self.what})" + + def lookup(self, ctx, field_name, given, values): + if ctx.distilled_steps is None: + return _MISSING + return ctx.distilled_steps if self.what == "steps" else DISTILLED_GUIDANCE_SCALE + + +@dataclass(frozen=True) +class Policy: + key: str + + @property + def label(self) -> str: + return f"checkpoint.json:policy.{self.key}" + + def lookup(self, ctx, field_name, given, values): + if not ctx.policy_active: + return _MISSING + value = ctx.policy_defaults.get(self.key) + return _MISSING if value is None else value + + +@dataclass(frozen=True) +class Embodiment: + """The checkpoint's embodiment table.""" + + key: str + + @property + def label(self) -> str: + return f"checkpoint:conditioning.embodiments.{self.key}" + + def lookup(self, ctx, field_name, given, values): + value = ctx.embodiment.get(self.key) + return _MISSING if value is None else value + + +class KnownWidth: + """The code's table of action widths per embodiment; silent when the checkpoint + declares its own embodiments.""" + + label = "embodiment table (defaults)" + + def lookup(self, ctx, field_name, given, values): + if ctx.checkpoint_has_embodiments: + return _MISSING + width = resolve_raw_action_dim(domain_name=ctx.domain_name, domain_id=ctx.domain_id) + return _MISSING if width is None else width + + +@dataclass(frozen=True) +class Preset: + key: str + + @property + def label(self) -> str: + return f"domain preset.{self.key}" + + def lookup(self, ctx, field_name, given, values): + preset = COSMOS3_DOMAIN_PRESETS.get(ctx.preset_key) if ctx.preset_key else None + value = preset.get(self.key) if preset else None + return _MISSING if value is None else value + + +@dataclass(frozen=True) +class Table: + """The (family, mode) generation table; ``mode=None`` is the request's own mode.""" + + key: str + mode: Optional[str] = None + + @property + def label(self) -> str: + return f"{self.mode or 'mode'} table.{self.key}" + + def lookup(self, ctx, field_name, given, values): + table = COSMOS3_GENERATION_DEFAULTS.get((ctx.family, self.mode or ctx.table_mode)) + value = table.get(self.key) if table else None + return _MISSING if value is None else value + + +@dataclass(frozen=True) +class Workflow: + """A value a workflow fixes: transfer never prepends the system prompt, V2V does.""" + + label: str + pick: Callable[[RequestContext], Any] + + def lookup(self, ctx, field_name, given, values): + value = self.pick(ctx) + return _MISSING if value is None else value + + +@dataclass(frozen=True) +class Checkpoint: + label: str + pick: Callable[[RequestContext], Any] + + def lookup(self, ctx, field_name, given, values): + value = self.pick(ctx) + return _MISSING if value is None else value + + +@dataclass(frozen=True) +class Derived: + """A value computed from fields resolved before it.""" + + label: str + compute: Callable[[Mapping[str, Any]], Any] + + def lookup(self, ctx, field_name, given, values): + value = self.compute(values) + return _MISSING if value is None else value + + +@dataclass(frozen=True) +class RequestRule: + sources: tuple + conflict: Conflict = Conflict.IGNORE + required: bool = True + + +# ------------------------------------------------------------------------------- rules + +GENERATION_RULES: dict[str, RequestRule] = { + "num_inference_steps": RequestRule( + ( + Given(), + Distilled("steps"), + Policy("num_inference_steps"), + Table("num_inference_steps"), + Table("num_inference_steps", "video"), + ) + ), + "guidance_scale": RequestRule( + ( + Given(), + Distilled("guidance"), + Policy("guidance_scale"), + Table("guidance_scale"), + Table("guidance_scale", "video"), + ) + ), + "guidance_interval": RequestRule( + ( + Given(), + Policy("guidance_interval"), + Table("guidance_interval"), + Table("guidance_interval", "video"), + ), + required=False, + ), + "max_sequence_length": RequestRule( + (Given(), Table("max_sequence_length"), Table("max_sequence_length", "video")) + ), + "height": RequestRule((Given(), Table("height"), Table("height", "video"))), + "width": RequestRule((Given(), Table("width"), Table("width", "video"))), + "num_frames": RequestRule((Given(), Table("num_frames"), Table("num_frames", "video"))), + "frame_rate": RequestRule((Given(), Table("frame_rate"), Table("frame_rate", "video"))), + "use_system_prompt": RequestRule( + ( + Given(), + Workflow("transfer (never)", lambda ctx: False if ctx.is_transfer else None), + Workflow("video-to-video (always)", lambda ctx: True if ctx.is_v2v else None), + Checkpoint( + "checkpoint:use_system_prompt", lambda ctx: ctx.checkpoint_use_system_prompt + ), + ) + ), +} + +ACTION_RULES: dict[str, RequestRule] = { + "raw_action_dim": RequestRule( + (Given(), Policy("raw_action_dim"), Embodiment("action_dim"), KnownWidth()), + conflict=Conflict.WARN, + required=False, + ), + "action_chunk_size": RequestRule( + ( + Given(), + Policy("action_chunk_size"), + Preset("action_chunk_size"), + Default(COSMOS3_ACTION_PARAMS["action_chunk_size"], "generic action horizon"), + ), + conflict=Conflict.WARN, + ), + "action_resolution": RequestRule( + ( + Given(), + Policy("action_resolution"), + Preset("action_resolution"), + Default(480, "the bucket the action models were trained at"), + ), + conflict=Conflict.WARN, + ), + "frame_rate": RequestRule( + ( + Given(), + Policy("frame_rate"), + Preset("frame_rate"), + Default(COSMOS3_ACTION_PARAMS["frame_rate"], "generic action frame rate"), + ), + conflict=Conflict.WARN, + ), + "action_fps": RequestRule( + (Given(), Policy("action_fps"), Derived("frame_rate", lambda v: v["frame_rate"])) + ), + "use_state": RequestRule( + (Given(), Policy("use_state"), Default(False, "policies without a state row")), + conflict=Conflict.WARN, + ), + "num_frames": RequestRule( + (Derived("action_chunk_size + 1", lambda v: int(v["action_chunk_size"]) + 1),) + ), +} + + +@dataclass +class ResolvedRequest: + values: dict[str, Resolved] = field(default_factory=dict) + warnings: list[str] = field(default_factory=list) + + def __getitem__(self, name: str) -> Any: + return self.values[name].value + + def get(self, name: str, default: Any = None) -> Any: + r = self.values.get(name) + return default if r is None else r.value + + def source_of(self, name: str) -> str: + return self.values[name].source + + +def _resolve( + rules: Mapping[str, RequestRule], + ctx: RequestContext, + given: Mapping[str, Any], + domain_label: str, +) -> ResolvedRequest: + out = ResolvedRequest() + plain: dict[str, Any] = {} + for name, rule in rules.items(): + winner: Optional[Resolved] = None + disagreeing: list[str] = [] + for source in rule.sources: + if winner is not None and isinstance(source, Default): + break + if isinstance(source, Default): + value = source.value + else: + value = source.lookup(ctx, name, given, plain) + if value is _MISSING: + continue + if winner is None: + winner = Resolved(value, source.label) + continue + # Only the first recommendation below the caller's value is compared. + if not same_value(value, winner.value): + disagreeing.append(f"{value}") + break + if winner is None: + if rule.required: + raise ValueError(f"Cosmos3 request: cannot resolve {name} for {domain_label}.") + out.values[name] = Resolved(None, "unresolved") + plain[name] = None + continue + if disagreeing and winner.source == Given.label: + message = ( + f"Cosmos3 {name}={winner.value} differs from recommended " + f"{' / '.join(disagreeing)} for domain {domain_label}." + ) + if rule.conflict is Conflict.ERROR: + raise ValueError(message) + if rule.conflict is Conflict.WARN: + out.warnings.append(message) + out.values[name] = Resolved(winner.value, winner.source, tuple(disagreeing)) + plain[name] = winner.value + return out + + +def resolve_generation(ctx: RequestContext, **given: Any) -> ResolvedRequest: + """Resolve the generation fields in ``given`` (``None`` = unset).""" + rules = {name: GENERATION_RULES[name] for name in given} + return _resolve(rules, ctx, given, repr(ctx.table_mode)) + + +def policy_domain( + policy_defaults: Mapping[str, Any], + domain_name: Optional[str], + domain_id: Optional[int], +) -> tuple[Optional[str], bool, list[str]]: + """Reconcile the request's domain with the checkpoint policy block's: returns the + domain name to use, whether the block applies, and a warning when the request + names a different domain (the block's defaults are then not applied).""" + checkpoint_domain = policy_defaults.get("domain_name") + warnings: list[str] = [] + if checkpoint_domain is None: + return domain_name, False, warnings + checkpoint_domain = str(checkpoint_domain).strip().lower() + explicit = ( + str(domain_name).strip().lower() if domain_name and str(domain_name).strip() else None + ) + checkpoint_domain_id = EMBODIMENT_TO_DOMAIN_ID.get(checkpoint_domain) + id_matches = domain_id is not None and checkpoint_domain_id == int(domain_id) + if explicit is None and (domain_id is None or id_matches): + return checkpoint_domain, True, warnings + if explicit == checkpoint_domain: + return domain_name, True, warnings + warnings.append( + f"Cosmos3 domain_name={domain_name!r}, domain_id={domain_id!r} overrides " + f"checkpoint policy domain_name={checkpoint_domain!r}; checkpoint domain " + "defaults will not be applied." + ) + return domain_name, False, warnings + + +def resolve_action(ctx: RequestContext, **given: Any) -> ResolvedRequest: + """Resolve the action fields; ``ctx.domain_name`` / ``domain_id`` already reconciled + with the policy block (:func:`policy_domain`).""" + domain_label = ( + repr(ctx.preset_key) + if ctx.preset_key + else (f"domain_name={ctx.domain_name!r}, domain_id={ctx.domain_id!r}") + ) + resolved = _resolve(ACTION_RULES, ctx, given, domain_label) + domain_requested = bool(ctx.domain_name and str(ctx.domain_name).strip()) or ( + ctx.domain_id is not None and str(ctx.domain_id).strip() not in {"", "0"} + ) + if domain_requested and ctx.preset_key is None: + resolved.warnings.insert( + 0, + "Cosmos3 action domain preset was not found for " + f"domain_name={ctx.domain_name!r}, domain_id={ctx.domain_id!r}; " + "using generic action defaults for omitted fields.", + ) + if domain_requested and resolved["raw_action_dim"] is None: + resolved.warnings.append( + "Cosmos3 has no canonical action width for " + f"domain_name={ctx.domain_name!r}, domain_id={ctx.domain_id!r}; " + "pass raw_action_dim explicitly for policy/inverse_dynamics." + ) + resolved.values["action_resolution"] = Resolved( + normalize_action_resolution(resolved["action_resolution"]), + resolved.source_of("action_resolution"), + ) + if resolved["raw_action_dim"] is not None and int(resolved["raw_action_dim"]) <= 0: + raise ValueError( + f"Cosmos3 raw_action_dim must be positive, got {resolved['raw_action_dim']}." + ) + if int(resolved["action_chunk_size"]) <= 0: + raise ValueError( + f"Cosmos3 action_chunk_size must be positive, got {resolved['action_chunk_size']}." + ) + if float(resolved["frame_rate"]) <= 0.0: + raise ValueError(f"Cosmos3 frame_rate must be positive, got {resolved['frame_rate']}.") + if float(resolved["action_fps"]) <= 0.0: + raise ValueError(f"Cosmos3 action_fps must be positive, got {resolved['action_fps']}.") + return resolved + + +def embodiment_entry( + embodiments: Mapping[str, Any], input_action_dim: Optional[int], domain_name, domain_id +): + """The checkpoint's embodiment facts for a request: ``{"domain_id", "action_dim"}``, + found by name or by id; empty when the checkpoint lists no such embodiment.""" + name = str(domain_name).strip().lower() if domain_name and str(domain_name).strip() else None + spec = None + if name is not None: + spec = embodiments.get(name) + elif domain_id is not None: + spec = next( + ( + s + for s in embodiments.values() + if s.get("domain_id") is not None and int(s["domain_id"]) == int(domain_id) + ), + None, + ) + if not spec: + return {} + width = spec.get("input_action_dim") + if width is None: + width = spec.get("raw_action_dim", input_action_dim) + return { + "domain_id": None if spec.get("domain_id") is None else int(spec["domain_id"]), + "action_dim": None if width is None else int(width), + } diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py index fa4fa56881af..ceab1a72eb9f 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -43,14 +43,11 @@ from tensorrt_llm.logger import logger from tensorrt_llm.models.modeling_utils import QuantConfig -# Some Cosmos3OmniTransformer checkpoint configs omit these fields; the values -# match what other conversions carry explicitly. -PRETRAINED_CONFIG_COMPAT_DEFAULTS = { - "position_embedding_type": "unified_3d_mrope", - "max_position_embeddings": 262144, - "temporal_compression_factor": 4, - "temporal_compression_factor_sound": 1, -} +from .config_resolver import architecture_defaults, fill_pretrained_config + +# Some Cosmos3OmniTransformer checkpoint configs omit these fields; the values come +# from the precedence table in config_resolver, the one place that declares them. +PRETRAINED_CONFIG_COMPAT_DEFAULTS = architecture_defaults() _PretrainedConfigT = TypeVar("_PretrainedConfigT") @@ -60,10 +57,7 @@ def apply_pretrained_config_compat_defaults( pretrained_config: _PretrainedConfigT, ) -> _PretrainedConfigT: """Fill missing schema fields in place (idempotent); returns the config.""" - for key, value in PRETRAINED_CONFIG_COMPAT_DEFAULTS.items(): - if getattr(pretrained_config, key, None) is None: - setattr(pretrained_config, key, value) - return pretrained_config + return fill_pretrained_config(pretrained_config) def uses_static_fp8(model_config: DiffusionModelConfig) -> bool: diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 8a7f47c7be2c..fe568354ebf9 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -161,6 +161,8 @@ l0_cpu: - unittest/_torch/visual_gen/test_attention_integration.py - unittest/_torch/visual_gen/test_attention_importable.py - unittest/_torch/visual_gen/test_cache_dit.py + - unittest/_torch/visual_gen/test_cosmos3_config_resolver.py + - unittest/_torch/visual_gen/test_cosmos3_request_resolver.py - unittest/_torch/visual_gen/test_executor_lifecycle.py - unittest/_torch/visual_gen/test_flux2_load_weights.py - unittest/_torch/visual_gen/test_flux_infer.py diff --git a/tests/unittest/_torch/visual_gen/cosmos3_request_golden.py b/tests/unittest/_torch/visual_gen/cosmos3_request_golden.py new file mode 100644 index 000000000000..6029ff9c7faa --- /dev/null +++ b/tests/unittest/_torch/visual_gen/cosmos3_request_golden.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The request-level defaults every Cosmos3 family, mode and domain resolves to when +the caller sets nothing. + +``record()`` evaluates the grid with the pipeline's current resolution helpers; +``python cosmos3_request_golden.py > test_data/cosmos3_request_defaults.json`` stores +it. ``test_cosmos3_request_resolver.py`` checks the request resolver against the +stored values, so a refactor of the defaults is a visible diff of that file. +""" + +import json +import sys + +from diffusers import FlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler + +from tensorrt_llm._torch.visual_gen.models.cosmos3.action import EMBODIMENT_TO_DOMAIN_ID +from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( + COSMOS3_DOMAIN_PRESETS, + resolve_checkpoint_policy_defaults, + resolve_domain_action_config, +) +from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import ( + Cosmos3OmniMoTPipeline, + _resolve_use_system_prompt, +) +from tensorrt_llm._torch.visual_gen.models.cosmos3.sampling import Cosmos3SamplingPolicy + +FAMILIES = ("qwen3", "nemotron_dense") +MODES = ("video", "image") +GENERATION_FIELDS = ( + "height", + "width", + "num_frames", + "frame_rate", + "num_inference_steps", + "guidance_scale", + "guidance_interval", + "max_sequence_length", +) +DROID_POLICY = {"domain_name": "droid_lerobot", "action_chunk_size": 32, "conditioning_fps": 15.0} + +UNIPC = { + "_class_name": "UniPCMultistepScheduler", + "num_train_timesteps": 1000, + "flow_shift": 1.0, + "prediction_type": "flow_prediction", + "use_flow_sigmas": True, + "use_karras_sigmas": True, + "solver_order": 2, +} +DISTILLED = { + "_class_name": "FlowMatchEulerDiscreteScheduler", + "num_train_timesteps": 1000, + "shift": 1.0, + "stochastic_sampling": True, + "fixed_step_sampler_config": {"sample_type": "sde", "t_list": [1.0, 0.9375, 0.8333, 0.625]}, +} + + +def _policy(distilled: bool) -> Cosmos3SamplingPolicy: + if distilled: + return Cosmos3SamplingPolicy.from_scheduler( + FlowMatchEulerDiscreteScheduler.from_config(DISTILLED) + ) + return Cosmos3SamplingPolicy.from_scheduler(UniPCMultistepScheduler.from_config(UNIPC)) + + +def _pipeline(family: str, distilled: bool, policy_block=None) -> Cosmos3OmniMoTPipeline: + pipe = object.__new__(Cosmos3OmniMoTPipeline) + pipe.family = family + pipe.sampling = _policy(distilled) + pipe.checkpoint_policy_defaults = resolve_checkpoint_policy_defaults(policy_block) + return pipe + + +def _jsonable(value): + if isinstance(value, tuple): + return list(value) + return value + + +def record() -> dict: + out = {"generation": {}, "action": {}, "system_prompt": {}} + for family in FAMILIES: + for distilled in (False, True): + if family == "nemotron_dense" and distilled: + continue # no such checkpoint; the pipeline refuses the combination + pipe = _pipeline(family, distilled) + for mode in MODES: + values = pipe._resolve_generation_params( + mode, **{f: None for f in GENERATION_FIELDS} + ) + key = f"{family}/{mode}/{'distilled' if distilled else 'base'}" + out["generation"][key] = {k: _jsonable(v) for k, v in values.items()} + # Action sampling fields with and without the DROID policy block. + for label, block in (("", None), ("/droid-policy", DROID_POLICY)): + pipe = _pipeline(family, distilled, block) + for action_mode in ("forward_dynamics", "policy"): + values = pipe._resolve_generation_params( + "video", + action_mode=action_mode, + num_inference_steps=None, + guidance_scale=None, + guidance_interval=None, + max_sequence_length=None, + ) + key = f"{family}/action:{action_mode}/{'distilled' if distilled else 'base'}{label}" + out["generation"][key] = {k: _jsonable(v) for k, v in values.items()} + domains = sorted(set(EMBODIMENT_TO_DOMAIN_ID) | set(COSMOS3_DOMAIN_PRESETS)) + [None] + for policy_label, block in (("", None), ("/droid-policy", DROID_POLICY)): + policy_defaults = resolve_checkpoint_policy_defaults(block) + for domain in domains: + try: + cfg = resolve_domain_action_config( + domain_name=domain, checkpoint_policy_defaults=policy_defaults + ) + cfg = {k: _jsonable(v) for k, v in cfg.items() if k not in ("warnings", "sources")} + cfg["num_warnings"] = len( + resolve_domain_action_config( + domain_name=domain, checkpoint_policy_defaults=policy_defaults + )["warnings"] + ) + except ValueError as exc: + cfg = {"error": str(exc)} + out["action"][f"{domain}{policy_label}"] = cfg + for checkpoint_default in (False, True): + for is_v2v in (False, True): + for is_transfer in (False, True): + key = f"checkpoint={checkpoint_default}/v2v={is_v2v}/transfer={is_transfer}" + out["system_prompt"][key] = _resolve_use_system_prompt( + None, + checkpoint_default=checkpoint_default, + is_v2v=is_v2v, + is_transfer=is_transfer, + ) + return out + + +if __name__ == "__main__": + json.dump(record(), sys.stdout, indent=1, sort_keys=True) + sys.stdout.write("\n") diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py index b6f1c026341b..3c53eaa1413a 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py @@ -13,7 +13,7 @@ import torch from tensorrt_llm._torch.visual_gen.models.cosmos3.action import ( - ACTION_ASPECT_RATIO_LABELS, + COSMOS3_ASPECT_RATIO_BUCKETS, EMBODIMENT_TO_DOMAIN_ID, EMBODIMENT_TO_RAW_ACTION_DIM, VIDEO_RES_SIZE_INFO, @@ -75,7 +75,7 @@ def test_unknown_resolution_raises(self): def test_all_buckets_have_aspect_entries(self, action_resolution): # find_closest_target_size picks from whatever entries exist, so a bucket # that lost "9,16" would silently land portrait sources on another canvas. - assert set(VIDEO_RES_SIZE_INFO[action_resolution]) == set(ACTION_ASPECT_RATIO_LABELS) + assert set(VIDEO_RES_SIZE_INFO[action_resolution]) == set(COSMOS3_ASPECT_RATIO_BUCKETS) class TestResolveActionSize: diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_config_resolver.py b/tests/unittest/_torch/visual_gen/test_cosmos3_config_resolver.py new file mode 100644 index 000000000000..ea453e5bd50a --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_config_resolver.py @@ -0,0 +1,218 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The Cosmos3 precedence table, on synthetic files and on the published checkpoints.""" + +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.visual_gen.models.cosmos3.config_resolver import ( + CHECKPOINT, + CONFIG, + MODEL_INDEX, + SCHEDULER, + TRANSFORMER, + VAE, + Conflict, + Cosmos3ConfigResolver, + Default, + Key, + Rule, + Serving, + architecture_defaults, + fill_pretrained_config, +) + +pytestmark = pytest.mark.cpu_only + +FIXTURE = Path(__file__).parent / "test_data" / "cosmos3_checkpoint_configs.json" +T_LIST = [1.0, 0.9375, 0.8333333333333334, 0.625] + + +def _write(checkpoint_dir: Path, files: dict) -> Path: + for name, doc in files.items(): + path = checkpoint_dir / name + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(doc)) + return checkpoint_dir + + +def _resolver(tmp_path: Path, files: dict, serving=None) -> Cosmos3ConfigResolver: + _write(tmp_path, files) + return Cosmos3ConfigResolver(str(tmp_path), files.get(TRANSFORMER, {}), serving) + + +class TestRule: + def test_first_present_source_wins(self, tmp_path): + r = _resolver(tmp_path, {CONFIG: {"a": 1}, MODEL_INDEX: {"a": 1}}) + rule = Rule((Key(CONFIG, "a"), Key(MODEL_INDEX, "a"), Default(7, "x"))) + assert r.resolve_one("a", rule).value == 1 + assert r.resolve_one("a", rule).source == "config.json:a" + + def test_falls_through_missing_and_null_sources(self, tmp_path): + r = _resolver(tmp_path, {CONFIG: {"a": None}, MODEL_INDEX: {"a": 2}}) + rule = Rule((Key(CONFIG, "a"), Key(MODEL_INDEX, "a"), Key(CHECKPOINT, "a"))) + assert r.resolve_one("a", rule).value == 2 + + def test_nullable_keeps_an_explicit_null(self, tmp_path): + r = _resolver(tmp_path, {CONFIG: {"a": None}, MODEL_INDEX: {"a": 2}}) + rule = Rule( + (Key(CONFIG, "a"), Key(MODEL_INDEX, "a")), conflict=Conflict.IGNORE, nullable=True + ) + assert r.resolve_one("a", rule).value is None + + def test_default_only_when_nothing_is_declared(self, tmp_path): + r = _resolver(tmp_path, {}) + rule = Rule((Key(CONFIG, "a"), Default(7, "x"))) + resolved = r.resolve_one("a", rule) + assert (resolved.value, resolved.source) == (7, "default (x)") + + def test_unresolvable_required_setting_is_an_error(self, tmp_path): + r = _resolver(tmp_path, {}) + with pytest.raises(ValueError, match="cannot resolve a"): + r.resolve_one("a", Rule((Key(CONFIG, "a"),))) + assert r.resolve_one("a", Rule((Key(CONFIG, "a"),), required=False)) is None + + def test_conflict_policies(self, tmp_path): + r = _resolver(tmp_path, {CONFIG: {"a": 1}, MODEL_INDEX: {"a": 2}}) + sources = (Key(CONFIG, "a"), Key(MODEL_INDEX, "a")) + with pytest.raises(ValueError, match="model_index.json:a=2"): + r.resolve_one("a", Rule(sources, conflict=Conflict.ERROR)) + warned = r.resolve_one("a", Rule(sources, conflict=Conflict.WARN)) + assert warned.value == 1 and warned.disagreeing == ("model_index.json:a=2",) + + def test_float_lists_compare_within_tolerance(self, tmp_path): + r = _resolver( + tmp_path, {CONFIG: {"a": [0.5, 1 / 3]}, MODEL_INDEX: {"a": [0.5, 0.3333333333]}} + ) + rule = Rule((Key(CONFIG, "a"), Key(MODEL_INDEX, "a")), conflict=Conflict.ERROR) + assert r.resolve_one("a", rule).disagreeing == () + + def test_serving_config_is_a_source(self, tmp_path): + serving = SimpleNamespace(sampling=SimpleNamespace(flow_shift=0.5)) + r = _resolver(tmp_path, {}, serving) + rule = Rule((Serving("sampling.flow_shift"),)) + assert r.resolve_one("f", rule).value == 0.5 + + +def _files(): + """A distilled checkpoint whose training config and diffusers index disagree on + the system prompt.""" + return { + TRANSFORMER: {"temporal_compression_factor": 4}, + SCHEDULER: {"fixed_step_sampler_config": {"t_list": T_LIST}}, + CONFIG: {"model": {"config": {"vlm_config": {"use_system_prompt": False}}}}, + MODEL_INDEX: {"default_use_system_prompt": True}, + VAE: {"scale_factor_temporal": 4, "scale_factor_spatial": 16}, + } + + +class TestResolve: + def test_training_config_beats_the_diffusers_index(self, tmp_path): + cfg = _resolver(tmp_path, _files()).resolve() + assert cfg.prompt.use_system_prompt is False + assert cfg.source_of("prompt.use_system_prompt") == ( + "config.json:model.config.vlm_config.use_system_prompt" + ) + assert cfg.resolved["prompt.use_system_prompt"].disagreeing == ( + "model_index.json:default_use_system_prompt=True", + ) + + def test_sections(self, tmp_path): + cfg = _resolver(tmp_path, _files()).resolve() + assert (cfg.vae.temporal_compression, cfg.vae.spatial_compression) == (4, 16) + assert cfg.sampling.fixed_sigmas == tuple(T_LIST) + assert cfg.sampling.native_flow_schedule is False + assert cfg.source_of("sampling.native_flow_schedule").startswith("default (") + assert cfg.action.policy is None + + def test_vae_compression_copies_must_agree(self, tmp_path): + files = _files() + files[TRANSFORMER]["temporal_compression_factor"] = 8 + with pytest.raises(ValueError, match="vae.temporal_compression"): + _resolver(tmp_path, files).resolve() + + def test_schedule_must_decrease_within_unit_interval(self, tmp_path): + files = _files() + files[SCHEDULER]["fixed_step_sampler_config"]["t_list"] = [0.5, 1.0] + with pytest.raises(ValueError, match="t_list must decrease"): + _resolver(tmp_path, files).resolve() + + def test_base_checkpoint_has_no_schedule(self, tmp_path): + files = {k: v for k, v in _files().items() if k != SCHEDULER} + cfg = _resolver(tmp_path, files).resolve() + assert cfg.sampling.fixed_sigmas is None + + def test_system_prompt_must_be_declared(self, tmp_path): + files = _files() + del files[CONFIG]["model"]["config"]["vlm_config"]["use_system_prompt"] + del files[MODEL_INDEX]["default_use_system_prompt"] + with pytest.raises(ValueError, match="cannot resolve prompt.use_system_prompt"): + _resolver(tmp_path, files).resolve() + + def test_spatial_compression_must_be_declared(self, tmp_path): + files = _files() + del files[VAE]["scale_factor_spatial"] + with pytest.raises(ValueError, match="cannot resolve vae.spatial_compression"): + _resolver(tmp_path, files).resolve() + + +def test_fill_pretrained_config_sets_only_missing_attributes(): + declared = SimpleNamespace(temporal_compression_factor=8) + fill_pretrained_config(declared) + assert declared.temporal_compression_factor == 8 + assert declared.position_embedding_type == architecture_defaults()["position_embedding_type"] + assert declared.max_position_embeddings == 262144 + assert declared.temporal_compression_factor_sound == 1 + + +# Disagreements between a checkpoint's own files that the table expects and resolves. +EXPECTED_DISAGREEMENTS = { + "nvidia/Cosmos3-Nano-Sim-Bimanual": {"prompt.use_system_prompt"}, +} + +EXPECTED = { + # repo: (use_system_prompt, native_flow_schedule, distilled) + "nvidia/Cosmos3-Nano": (False, False, False), + "nvidia/Cosmos3-Super": (False, False, False), + "nvidia/Cosmos3-Edge": (False, True, False), + "nvidia/Cosmos3-Super-Image2Video": (False, False, False), + "nvidia/Cosmos3-Super-Image2Video-4Step": (False, False, True), + "nvidia/Cosmos3-Super-Text2Image": (False, False, False), + "nvidia/Cosmos3-Super-Text2Image-4Step": (False, False, True), + "nvidia/Cosmos3-Nano-Policy-DROID": (False, False, False), + "nvidia/Cosmos3-Edge-Policy-DROID": (False, True, False), + "nvidia/Cosmos3-Nano-Sim-Bimanual": (False, False, True), +} + + +@pytest.mark.parametrize("repo", sorted(EXPECTED)) +def test_published_checkpoint_resolves(tmp_path, repo): + files = json.loads(FIXTURE.read_text())[repo] + cfg = _resolver(tmp_path, files).resolve() + system_prompt, native, distilled = EXPECTED[repo] + assert cfg.prompt.use_system_prompt is system_prompt + assert cfg.sampling.native_flow_schedule is native + assert (cfg.sampling.fixed_sigmas is not None) is distilled + assert (cfg.vae.temporal_compression, cfg.vae.spatial_compression) == (4, 16) + assert cfg.architecture.position_embedding_type == "unified_3d_mrope" + disagreeing = {name for name, r in cfg.resolved.items() if r.disagreeing} + assert disagreeing == EXPECTED_DISAGREEMENTS.get(repo, set()) + + +def test_fixture_covers_every_expected_checkpoint(): + assert set(json.loads(FIXTURE.read_text())) == set(EXPECTED) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py index 85804719ea67..2e4c381cb7b0 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py @@ -69,6 +69,17 @@ def _write_scheduler_config(checkpoint_dir: Path, config: dict) -> None: scheduler_dir.mkdir(parents=True, exist_ok=True) with open(scheduler_dir / "scheduler_config.json", "w") as f: json.dump(config, f) + vae_dir = checkpoint_dir / "vae" + vae_dir.mkdir(parents=True, exist_ok=True) + with open(vae_dir / "config.json", "w") as f: + json.dump({"scale_factor_temporal": 4, "scale_factor_spatial": 16}, f) + + +def _write_trained_config(checkpoint_dir: Path, use_system_prompt) -> None: + """The training config's prompt flag; ``None`` writes a config without it.""" + vlm_config = {} if use_system_prompt is None else {"use_system_prompt": use_system_prompt} + with open(checkpoint_dir / "config.json", "w") as f: + json.dump({"model": {"config": {"vlm_config": vlm_config}}}, f) def _distilled_policy() -> Cosmos3SamplingPolicy: @@ -599,6 +610,7 @@ def test_full_chain_mode_switch_resolves_to_t2i(self): class TestPipelineSchedulerLoading: def test_distilled_checkpoint_loads_flow_match(self, tmp_path): _write_scheduler_config(tmp_path, DISTILLED_SCHEDULER_CONFIG) + _write_trained_config(tmp_path, False) pipeline = _bare_pipeline() pipeline.load_standard_components( @@ -611,6 +623,7 @@ def test_distilled_checkpoint_loads_flow_match(self, tmp_path): def test_base_checkpoint_loads_unipc(self, tmp_path): _write_scheduler_config(tmp_path, UNIPC_SCHEDULER_CONFIG) + _write_trained_config(tmp_path, False) pipeline = _bare_pipeline() pipeline.load_standard_components( @@ -623,6 +636,7 @@ def test_base_checkpoint_loads_unipc(self, tmp_path): def test_audio_scheduler_is_separate_same_class_instance(self, tmp_path): _write_scheduler_config(tmp_path, DISTILLED_SCHEDULER_CONFIG) + _write_trained_config(tmp_path, False) pipeline = _bare_pipeline(audio_gen=True) pipeline.load_standard_components( @@ -1042,7 +1056,7 @@ def resolve_with_interval( class TestSystemPromptDefault: - """use_system_prompt defaults are checkpoint-declared via model_index.json.""" + """use_system_prompt default: config.json's vlm_config, else model_index.json.""" def _write_model_index(self, checkpoint_dir: Path, content: dict) -> None: with open(checkpoint_dir / "model_index.json", "w") as f: @@ -1071,15 +1085,26 @@ def test_checkpoint_declared_true(self, tmp_path): assert COSMOS3_EXTRA_SPECS["use_system_prompt"].default is None - def test_missing_model_index_keeps_false(self, tmp_path): + def test_undeclared_system_prompt_is_a_load_error(self, tmp_path): + with pytest.raises(ValueError, match="cannot resolve prompt.use_system_prompt"): + self._loaded_pipeline(tmp_path) + self._write_model_index(tmp_path, {"_class_name": "Cosmos3OmniPipeline"}) + _write_trained_config(tmp_path, None) + with pytest.raises(ValueError, match="cannot resolve prompt.use_system_prompt"): + self._loaded_pipeline(tmp_path) + + @pytest.mark.parametrize("trained", [False, True]) + def test_trained_config_wins_over_model_index(self, tmp_path, trained): + self._write_model_index(tmp_path, {"default_use_system_prompt": not trained}) + _write_trained_config(tmp_path, trained) pipeline = self._loaded_pipeline(tmp_path) - assert pipeline.default_use_system_prompt is False - assert pipeline.extra_param_specs["use_system_prompt"].default is None + assert pipeline.default_use_system_prompt is trained - def test_model_index_without_field_keeps_false(self, tmp_path): - self._write_model_index(tmp_path, {"_class_name": "Cosmos3OmniPipeline"}) + def test_trained_config_without_key_keeps_model_index(self, tmp_path): + self._write_model_index(tmp_path, {"default_use_system_prompt": True}) + _write_trained_config(tmp_path, None) pipeline = self._loaded_pipeline(tmp_path) - assert pipeline.default_use_system_prompt is False + assert pipeline.default_use_system_prompt is True def _captured_use_system_prompt(self, pipeline, extra_params): captured = {} diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py b/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py index 6957f1646fc3..901628f8fedf 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py @@ -1129,6 +1129,13 @@ def test_component_loading_rejects_edge_without_native_flag(self, tmp_path): (tmp_path / "model_index.json").write_text( json.dumps({"_class_name": "Cosmos3OmniPipeline"}) ) + (tmp_path / "config.json").write_text( + json.dumps({"model": {"config": {"vlm_config": {"use_system_prompt": False}}}}) + ) + (tmp_path / "vae").mkdir() + (tmp_path / "vae" / "config.json").write_text( + json.dumps({"scale_factor_temporal": 4, "scale_factor_spatial": 16}) + ) def fresh_pipeline(): pipeline = _bare_pipeline(NEMOTRON_DENSE_RECIPE.name) @@ -1802,7 +1809,7 @@ def test_tokenizer_specials_and_chat_template(self): def test_tokenization_matches_diffusers_golden(self): from transformers import AutoTokenizer - from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import ( + from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( COSMOS3_DEFAULT_RESOLUTION_TEMPLATE, COSMOS3_DURATION_TEMPLATE, ) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index 4797c1b61357..2cece22b55b0 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -41,15 +41,15 @@ COSMOS3_ACTION_PARAMS, COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP, COSMOS3_DEFAULT_CONDITION_VIDEO_LATENT_INDEXES, + COSMOS3_DEFAULT_RESOLUTION_TEMPLATE, + COSMOS3_DEFAULT_SYSTEM_PROMPT, + COSMOS3_DURATION_TEMPLATE, COSMOS3_EXTRA_SPECS, + COSMOS3_IMAGE_RESOLUTION_TEMPLATE, COSMOS3_T2I_PARAMS, _normalize_condition_video_keep, ) from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import ( - COSMOS3_DEFAULT_RESOLUTION_TEMPLATE, - COSMOS3_DEFAULT_SYSTEM_PROMPT, - COSMOS3_DURATION_TEMPLATE, - COSMOS3_IMAGE_RESOLUTION_TEMPLATE, Cosmos3OmniMoTPipeline, _condition_pixel_frame_count, _load_reference_image, diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_request_resolver.py b/tests/unittest/_torch/visual_gen/test_cosmos3_request_resolver.py new file mode 100644 index 000000000000..ca6779c1b9d8 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_request_resolver.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The request rules: the recorded defaults of every family, mode and domain, and the +source each value comes from.""" + +import json +from pathlib import Path + +import pytest +from cosmos3_request_golden import record + +from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import resolve_domain_action_config +from tensorrt_llm._torch.visual_gen.models.cosmos3.request_resolver import ( + RequestContext, + resolve_action, + resolve_generation, +) + +pytestmark = pytest.mark.cpu_only + +FIXTURE = Path(__file__).parent / "test_data" / "cosmos3_request_defaults.json" + + +def test_defaults_match_the_recorded_grid(): + """Every (family, mode, domain) combination resolves to the values recorded before + the rules existed; a change here is a visible change of a default.""" + recorded = json.loads(FIXTURE.read_text()) + current = json.loads(json.dumps(record())) + for section in ("generation", "action", "system_prompt"): + assert current[section] == recorded[section], section + + +class TestSources: + def test_generation_values_name_their_source(self): + ctx = RequestContext(family="qwen3", mode="image") + r = resolve_generation(ctx, num_inference_steps=None, num_frames=None, height=20) + assert r.source_of("num_inference_steps") == "mode table.num_inference_steps" + assert r["num_inference_steps"] == 50 + assert r.source_of("num_frames") == "video table.num_frames" # the image table has none + assert (r["height"], r.source_of("height")) == (20, "request") + + def test_distilled_schedule_beats_the_tables(self): + ctx = RequestContext(family="qwen3", distilled_steps=4) + r = resolve_generation(ctx, num_inference_steps=None, guidance_scale=None) + assert (r["num_inference_steps"], r["guidance_scale"]) == (4, 1.0) + assert r.source_of("guidance_scale").startswith("checkpoint:fixed_step_sampler_config") + + def test_policy_block_applies_only_to_its_domain(self): + policy = {"domain_name": "droid_lerobot", "action_chunk_size": 32, "num_inference_steps": 4} + matched = resolve_domain_action_config(checkpoint_policy_defaults=policy) + assert matched["domain_name"] == "droid_lerobot" + assert matched["action_chunk_size"] == 32 + assert matched["sources"]["action_chunk_size"] == "checkpoint.json:policy.action_chunk_size" + other = resolve_domain_action_config(domain_name="av", checkpoint_policy_defaults=policy) + assert other["action_chunk_size"] == 60 + assert other["sources"]["action_chunk_size"] == "domain preset.action_chunk_size" + assert any("overrides checkpoint policy domain_name" in w for w in other["warnings"]) + + def test_caller_value_wins_and_the_recommendation_is_reported(self): + cfg = resolve_domain_action_config(domain_name="av", action_chunk_size=16) + assert cfg["action_chunk_size"] == 16 + assert cfg["sources"]["action_chunk_size"] == "request" + assert cfg["warnings"] == [ + "Cosmos3 action_chunk_size=16 differs from recommended 60 for domain 'av'." + ] + + def test_checkpoint_embodiments_replace_the_code_table(self): + embodiments = {"agibotworld": {"domain_id": 15, "input_action_dim": 59}} + cfg = resolve_domain_action_config(domain_name="agibotworld", embodiments=embodiments) + assert cfg["raw_action_dim"] == 59 + assert cfg["sources"]["raw_action_dim"] == "checkpoint:conditioning.embodiments.action_dim" + assert cfg["warnings"] == [] + # The same request against a checkpoint without a table: the code's width. + plain = resolve_domain_action_config(domain_name="agibotworld") + assert plain["raw_action_dim"] == 29 + assert plain["sources"]["raw_action_dim"] == "embodiment table (defaults)" + + def test_unresolvable_field_is_an_error(self): + ctx = RequestContext(family="no-such-family") + with pytest.raises(ValueError, match="cannot resolve num_frames"): + resolve_generation(ctx, num_frames=None) + + def test_action_fps_follows_the_frame_rate(self): + ctx = RequestContext(family="qwen3", action_mode="forward_dynamics", domain_name="av") + r = resolve_action( + ctx, + raw_action_dim=None, + action_chunk_size=None, + action_resolution=None, + frame_rate=None, + action_fps=None, + use_state=None, + ) + assert (r["frame_rate"], r["action_fps"], r["num_frames"]) == (10.0, 10.0, 61) + assert r.source_of("action_fps") == "frame_rate" diff --git a/tests/unittest/_torch/visual_gen/test_data/README.md b/tests/unittest/_torch/visual_gen/test_data/README.md index 369e5c0c5189..6d3232e56189 100644 --- a/tests/unittest/_torch/visual_gen/test_data/README.md +++ b/tests/unittest/_torch/visual_gen/test_data/README.md @@ -57,3 +57,15 @@ ffmpeg -y -framerate 24 -i frame_%02d.png \ -x264-params b_adapt=0:scenecut=0 \ cosmos3_v2v_ref_9f_bframes.avi ``` + +## `cosmos3_checkpoint_configs.json` + +The settings `config_resolver.RULES` reads, extracted from the +published Cosmos3 checkpoints on Hugging Face (`main`, Fri 23 Oct 2026): for each +repo, file and key path a rule names, the value at that path, nested back into its +file. Nothing else from the checkpoints is kept. `test_cosmos3_config_resolver.py` +resolves every entry, so a new export whose files disagree in a way the precedence +table does not expect fails the test. + +Regenerate by fetching `https://huggingface.co//raw/main/` for every +`Key` in the table and keeping only those paths. diff --git a/tests/unittest/_torch/visual_gen/test_data/cosmos3_checkpoint_configs.json b/tests/unittest/_torch/visual_gen/test_data/cosmos3_checkpoint_configs.json new file mode 100644 index 000000000000..ca68ec07d0aa --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_data/cosmos3_checkpoint_configs.json @@ -0,0 +1,226 @@ +{ + "nvidia/Cosmos3-Edge": { + "model_index.json": { + "default_use_system_prompt": false, + "use_native_flow_schedule": true + }, + "transformer/config.json": { + "temporal_compression_factor": 4 + }, + "vae/config.json": { + "scale_factor_temporal": 4, + "scale_factor_spatial": 16 + } + }, + "nvidia/Cosmos3-Edge-Policy-DROID": { + "checkpoint.json": { + "policy": { + "action_chunk_size": 32, + "conditioning_fps": 15.0, + "domain_name": "droid_lerobot" + } + }, + "config.json": { + "model": { + "config": { + "vlm_config": { + "use_system_prompt": false + } + } + } + }, + "model_index.json": { + "default_use_system_prompt": false, + "use_native_flow_schedule": true + }, + "transformer/config.json": { + "temporal_compression_factor": 4 + }, + "vae/config.json": { + "scale_factor_temporal": 4, + "scale_factor_spatial": 16 + } + }, + "nvidia/Cosmos3-Nano": { + "config.json": { + "model": { + "config": { + "vlm_config": { + "use_system_prompt": false + } + } + } + }, + "transformer/config.json": { + "position_embedding_type": "unified_3d_mrope", + "max_position_embeddings": 262144, + "temporal_compression_factor_sound": 1 + }, + "vae/config.json": { + "scale_factor_temporal": 4, + "scale_factor_spatial": 16 + } + }, + "nvidia/Cosmos3-Nano-Policy-DROID": { + "config.json": { + "model": { + "config": { + "vlm_config": { + "use_system_prompt": false + } + } + } + }, + "transformer/config.json": { + "position_embedding_type": "unified_3d_mrope", + "max_position_embeddings": 262144, + "temporal_compression_factor_sound": 1 + }, + "vae/config.json": { + "scale_factor_temporal": 4, + "scale_factor_spatial": 16 + } + }, + "nvidia/Cosmos3-Nano-Sim-Bimanual": { + "config.json": { + "model": { + "config": { + "vlm_config": { + "use_system_prompt": false + } + } + } + }, + "model_index.json": { + "default_use_system_prompt": true, + "use_native_flow_schedule": false + }, + "scheduler/scheduler_config.json": { + "fixed_step_sampler_config": { + "t_list": [ + 1.0, + 0.9375, + 0.8333333333333334, + 0.625 + ] + } + }, + "transformer/config.json": { + "temporal_compression_factor": 4 + }, + "vae/config.json": { + "scale_factor_temporal": 4, + "scale_factor_spatial": 16 + } + }, + "nvidia/Cosmos3-Super": { + "config.json": { + "model": { + "config": { + "vlm_config": { + "use_system_prompt": false + } + } + } + }, + "transformer/config.json": { + "position_embedding_type": "unified_3d_mrope", + "max_position_embeddings": 262144, + "temporal_compression_factor_sound": 1 + }, + "vae/config.json": { + "scale_factor_temporal": 4, + "scale_factor_spatial": 16 + } + }, + "nvidia/Cosmos3-Super-Image2Video": { + "config.json": { + "model": { + "config": { + "vlm_config": { + "use_system_prompt": false + } + } + } + }, + "transformer/config.json": { + "position_embedding_type": "unified_3d_mrope", + "max_position_embeddings": 262144, + "temporal_compression_factor_sound": 1 + }, + "vae/config.json": { + "scale_factor_temporal": 4, + "scale_factor_spatial": 16 + } + }, + "nvidia/Cosmos3-Super-Image2Video-4Step": { + "config.json": { + "model": { + "config": { + "vlm_config": { + "use_system_prompt": false + } + } + } + }, + "scheduler/scheduler_config.json": { + "fixed_step_sampler_config": { + "t_list": [ + 1.0, + 0.9375, + 0.8333333333333334, + 0.625 + ] + } + }, + "vae/config.json": { + "scale_factor_temporal": 4, + "scale_factor_spatial": 16 + } + }, + "nvidia/Cosmos3-Super-Text2Image": { + "config.json": { + "model": { + "config": { + "vlm_config": { + "use_system_prompt": false + } + } + } + }, + "transformer/config.json": { + "position_embedding_type": "unified_3d_mrope", + "max_position_embeddings": 262144, + "temporal_compression_factor_sound": 1 + }, + "vae/config.json": { + "scale_factor_temporal": 4, + "scale_factor_spatial": 16 + } + }, + "nvidia/Cosmos3-Super-Text2Image-4Step": { + "config.json": { + "model": { + "config": { + "vlm_config": { + "use_system_prompt": false + } + } + } + }, + "scheduler/scheduler_config.json": { + "fixed_step_sampler_config": { + "t_list": [ + 1.0, + 0.9375, + 0.8333333333333334, + 0.625 + ] + } + }, + "vae/config.json": { + "scale_factor_temporal": 4, + "scale_factor_spatial": 16 + } + } +} diff --git a/tests/unittest/_torch/visual_gen/test_data/cosmos3_request_defaults.json b/tests/unittest/_torch/visual_gen/test_data/cosmos3_request_defaults.json new file mode 100644 index 000000000000..d38e8627f0cc --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_data/cosmos3_request_defaults.json @@ -0,0 +1,631 @@ +{ + "action": { + "None": { + "action_chunk_size": 16, + "action_fps": 24.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": false, + "domain_name": null, + "frame_rate": 24.0, + "num_frames": 17, + "num_warnings": 0, + "preset_key": null, + "raw_action_dim": null, + "use_state": false + }, + "None/droid-policy": { + "action_chunk_size": 32, + "action_fps": 15.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": true, + "domain_name": "droid_lerobot", + "frame_rate": 15.0, + "num_frames": 33, + "num_warnings": 0, + "preset_key": "droid_lerobot", + "raw_action_dim": 8, + "use_state": true + }, + "agibot_gear_gripper": { + "action_chunk_size": 16, + "action_fps": 10.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": false, + "domain_name": "agibot_gear_gripper", + "frame_rate": 10.0, + "num_frames": 17, + "num_warnings": 0, + "preset_key": "agibotworld", + "raw_action_dim": 29, + "use_state": false + }, + "agibot_gear_gripper/droid-policy": { + "action_chunk_size": 16, + "action_fps": 10.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": false, + "domain_name": "agibot_gear_gripper", + "frame_rate": 10.0, + "num_frames": 17, + "num_warnings": 1, + "preset_key": "agibotworld", + "raw_action_dim": 29, + "use_state": false + }, + "agibot_gear_gripper_ext": { + "action_chunk_size": 16, + "action_fps": 10.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": false, + "domain_name": "agibot_gear_gripper_ext", + "frame_rate": 10.0, + "num_frames": 17, + "num_warnings": 0, + "preset_key": "agibotworld", + "raw_action_dim": 29, + "use_state": false + }, + "agibot_gear_gripper_ext/droid-policy": { + "action_chunk_size": 16, + "action_fps": 10.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": false, + "domain_name": "agibot_gear_gripper_ext", + "frame_rate": 10.0, + "num_frames": 17, + "num_warnings": 1, + "preset_key": "agibotworld", + "raw_action_dim": 29, + "use_state": false + }, + "agibotworld": { + "action_chunk_size": 16, + "action_fps": 10.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": false, + "domain_name": "agibotworld", + "frame_rate": 10.0, + "num_frames": 17, + "num_warnings": 0, + "preset_key": "agibotworld", + "raw_action_dim": 29, + "use_state": false + }, + "agibotworld/droid-policy": { + "action_chunk_size": 16, + "action_fps": 10.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": false, + "domain_name": "agibotworld", + "frame_rate": 10.0, + "num_frames": 17, + "num_warnings": 1, + "preset_key": "agibotworld", + "raw_action_dim": 29, + "use_state": false + }, + "av": { + "action_chunk_size": 60, + "action_fps": 10.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": false, + "domain_name": "av", + "frame_rate": 10.0, + "num_frames": 61, + "num_warnings": 0, + "preset_key": "av", + "raw_action_dim": 9, + "use_state": false + }, + "av/droid-policy": { + "action_chunk_size": 60, + "action_fps": 10.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": false, + "domain_name": "av", + "frame_rate": 10.0, + "num_frames": 61, + "num_warnings": 1, + "preset_key": "av", + "raw_action_dim": 9, + "use_state": false + }, + "bridge_orig_lerobot": { + "action_chunk_size": 16, + "action_fps": 5.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": false, + "domain_name": "bridge_orig_lerobot", + "frame_rate": 5.0, + "num_frames": 17, + "num_warnings": 0, + "preset_key": "bridge_orig_lerobot", + "raw_action_dim": 10, + "use_state": false + }, + "bridge_orig_lerobot/droid-policy": { + "action_chunk_size": 16, + "action_fps": 5.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": false, + "domain_name": "bridge_orig_lerobot", + "frame_rate": 5.0, + "num_frames": 17, + "num_warnings": 1, + "preset_key": "bridge_orig_lerobot", + "raw_action_dim": 10, + "use_state": false + }, + "camera_pose": { + "action_chunk_size": 60, + "action_fps": 30.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": false, + "domain_name": "camera_pose", + "frame_rate": 30.0, + "num_frames": 61, + "num_warnings": 0, + "preset_key": "camera_pose", + "raw_action_dim": 9, + "use_state": false + }, + "camera_pose/droid-policy": { + "action_chunk_size": 60, + "action_fps": 30.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": false, + "domain_name": "camera_pose", + "frame_rate": 30.0, + "num_frames": 61, + "num_warnings": 1, + "preset_key": "camera_pose", + "raw_action_dim": 9, + "use_state": false + }, + "droid_lerobot": { + "action_chunk_size": 16, + "action_fps": 15.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": false, + "domain_name": "droid_lerobot", + "frame_rate": 15.0, + "num_frames": 17, + "num_warnings": 0, + "preset_key": "droid_lerobot", + "raw_action_dim": 10, + "use_state": false + }, + "droid_lerobot/droid-policy": { + "action_chunk_size": 32, + "action_fps": 15.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": true, + "domain_name": "droid_lerobot", + "frame_rate": 15.0, + "num_frames": 33, + "num_warnings": 0, + "preset_key": "droid_lerobot", + "raw_action_dim": 8, + "use_state": true + }, + "fractal": { + "action_chunk_size": 16, + "action_fps": 5.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": false, + "domain_name": "fractal", + "frame_rate": 5.0, + "num_frames": 17, + "num_warnings": 0, + "preset_key": "fractal", + "raw_action_dim": 10, + "use_state": false + }, + "fractal/droid-policy": { + "action_chunk_size": 16, + "action_fps": 5.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": false, + "domain_name": "fractal", + "frame_rate": 5.0, + "num_frames": 17, + "num_warnings": 1, + "preset_key": "fractal", + "raw_action_dim": 10, + "use_state": false + }, + "galbot": { + "action_chunk_size": 16, + "action_fps": 10.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": false, + "domain_name": "galbot", + "frame_rate": 10.0, + "num_frames": 17, + "num_warnings": 0, + "preset_key": "agibotworld", + "raw_action_dim": 30, + "use_state": false + }, + "galbot/droid-policy": { + "action_chunk_size": 16, + "action_fps": 10.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": false, + "domain_name": "galbot", + "frame_rate": 10.0, + "num_frames": 17, + "num_warnings": 1, + "preset_key": "agibotworld", + "raw_action_dim": 30, + "use_state": false + }, + "hand_pose": { + "action_chunk_size": 16, + "action_fps": 24.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": false, + "domain_name": "hand_pose", + "frame_rate": 24.0, + "num_frames": 17, + "num_warnings": 0, + "preset_key": "hand_pose", + "raw_action_dim": 57, + "use_state": false + }, + "hand_pose/droid-policy": { + "action_chunk_size": 16, + "action_fps": 24.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": false, + "domain_name": "hand_pose", + "frame_rate": 24.0, + "num_frames": 17, + "num_warnings": 1, + "preset_key": "hand_pose", + "raw_action_dim": 57, + "use_state": false + }, + "libero": { + "action_chunk_size": 16, + "action_fps": 10.0, + "action_resolution": 256, + "checkpoint_policy_domain_matches": false, + "domain_name": "libero", + "frame_rate": 10.0, + "num_frames": 17, + "num_warnings": 1, + "preset_key": "libero", + "raw_action_dim": null, + "use_state": false + }, + "libero/droid-policy": { + "action_chunk_size": 16, + "action_fps": 10.0, + "action_resolution": 256, + "checkpoint_policy_domain_matches": false, + "domain_name": "libero", + "frame_rate": 10.0, + "num_frames": 17, + "num_warnings": 2, + "preset_key": "libero", + "raw_action_dim": null, + "use_state": false + }, + "no_action": { + "action_chunk_size": 16, + "action_fps": 24.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": false, + "domain_name": "no_action", + "frame_rate": 24.0, + "num_frames": 17, + "num_warnings": 2, + "preset_key": null, + "raw_action_dim": null, + "use_state": false + }, + "no_action/droid-policy": { + "action_chunk_size": 16, + "action_fps": 24.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": false, + "domain_name": "no_action", + "frame_rate": 24.0, + "num_frames": 17, + "num_warnings": 3, + "preset_key": null, + "raw_action_dim": null, + "use_state": false + }, + "pusht": { + "action_chunk_size": 16, + "action_fps": 10.0, + "action_resolution": 256, + "checkpoint_policy_domain_matches": false, + "domain_name": "pusht", + "frame_rate": 10.0, + "num_frames": 17, + "num_warnings": 0, + "preset_key": "pusht", + "raw_action_dim": 2, + "use_state": false + }, + "pusht/droid-policy": { + "action_chunk_size": 16, + "action_fps": 10.0, + "action_resolution": 256, + "checkpoint_policy_domain_matches": false, + "domain_name": "pusht", + "frame_rate": 10.0, + "num_frames": 17, + "num_warnings": 1, + "preset_key": "pusht", + "raw_action_dim": 2, + "use_state": false + }, + "robomind-franka": { + "action_chunk_size": 16, + "action_fps": 15.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": false, + "domain_name": "robomind-franka", + "frame_rate": 15.0, + "num_frames": 17, + "num_warnings": 0, + "preset_key": "droid_lerobot", + "raw_action_dim": 10, + "use_state": false + }, + "robomind-franka-dual": { + "action_chunk_size": 16, + "action_fps": 15.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": false, + "domain_name": "robomind-franka-dual", + "frame_rate": 15.0, + "num_frames": 17, + "num_warnings": 0, + "preset_key": "droid_lerobot", + "raw_action_dim": 20, + "use_state": false + }, + "robomind-franka-dual/droid-policy": { + "action_chunk_size": 16, + "action_fps": 15.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": false, + "domain_name": "robomind-franka-dual", + "frame_rate": 15.0, + "num_frames": 17, + "num_warnings": 1, + "preset_key": "droid_lerobot", + "raw_action_dim": 20, + "use_state": false + }, + "robomind-franka/droid-policy": { + "action_chunk_size": 16, + "action_fps": 15.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": false, + "domain_name": "robomind-franka", + "frame_rate": 15.0, + "num_frames": 17, + "num_warnings": 1, + "preset_key": "droid_lerobot", + "raw_action_dim": 10, + "use_state": false + }, + "robomind-ur": { + "action_chunk_size": 16, + "action_fps": 15.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": false, + "domain_name": "robomind-ur", + "frame_rate": 15.0, + "num_frames": 17, + "num_warnings": 0, + "preset_key": "droid_lerobot", + "raw_action_dim": 10, + "use_state": false + }, + "robomind-ur/droid-policy": { + "action_chunk_size": 16, + "action_fps": 15.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": false, + "domain_name": "robomind-ur", + "frame_rate": 15.0, + "num_frames": 17, + "num_warnings": 1, + "preset_key": "droid_lerobot", + "raw_action_dim": 10, + "use_state": false + }, + "umi": { + "action_chunk_size": 16, + "action_fps": 10.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": false, + "domain_name": "umi", + "frame_rate": 10.0, + "num_frames": 17, + "num_warnings": 0, + "preset_key": "umi", + "raw_action_dim": 10, + "use_state": false + }, + "umi/droid-policy": { + "action_chunk_size": 16, + "action_fps": 10.0, + "action_resolution": 480, + "checkpoint_policy_domain_matches": false, + "domain_name": "umi", + "frame_rate": 10.0, + "num_frames": 17, + "num_warnings": 1, + "preset_key": "umi", + "raw_action_dim": 10, + "use_state": false + } + }, + "generation": { + "nemotron_dense/action:forward_dynamics/base": { + "guidance_interval": null, + "guidance_scale": 1.0, + "max_sequence_length": 4096, + "num_inference_steps": 30 + }, + "nemotron_dense/action:forward_dynamics/base/droid-policy": { + "guidance_interval": null, + "guidance_scale": 1.0, + "max_sequence_length": 4096, + "num_inference_steps": 30 + }, + "nemotron_dense/action:policy/base": { + "guidance_interval": null, + "guidance_scale": 1.0, + "max_sequence_length": 4096, + "num_inference_steps": 30 + }, + "nemotron_dense/action:policy/base/droid-policy": { + "guidance_interval": [ + 960.0, + 1001.0 + ], + "guidance_scale": 3.0, + "max_sequence_length": 4096, + "num_inference_steps": 4 + }, + "nemotron_dense/image/base": { + "frame_rate": 24.0, + "guidance_interval": null, + "guidance_scale": 4.0, + "height": 640, + "max_sequence_length": 4096, + "num_frames": 121, + "num_inference_steps": 50, + "width": 640 + }, + "nemotron_dense/video/base": { + "frame_rate": 24.0, + "guidance_interval": null, + "guidance_scale": 5.0, + "height": 480, + "max_sequence_length": 4096, + "num_frames": 121, + "num_inference_steps": 50, + "width": 832 + }, + "qwen3/action:forward_dynamics/base": { + "guidance_interval": null, + "guidance_scale": 1.0, + "max_sequence_length": 4096, + "num_inference_steps": 30 + }, + "qwen3/action:forward_dynamics/base/droid-policy": { + "guidance_interval": null, + "guidance_scale": 1.0, + "max_sequence_length": 4096, + "num_inference_steps": 30 + }, + "qwen3/action:forward_dynamics/distilled": { + "guidance_interval": null, + "guidance_scale": 1.0, + "max_sequence_length": 4096, + "num_inference_steps": 4 + }, + "qwen3/action:forward_dynamics/distilled/droid-policy": { + "guidance_interval": null, + "guidance_scale": 1.0, + "max_sequence_length": 4096, + "num_inference_steps": 4 + }, + "qwen3/action:policy/base": { + "guidance_interval": null, + "guidance_scale": 1.0, + "max_sequence_length": 4096, + "num_inference_steps": 30 + }, + "qwen3/action:policy/base/droid-policy": { + "guidance_interval": [ + 960.0, + 1001.0 + ], + "guidance_scale": 3.0, + "max_sequence_length": 4096, + "num_inference_steps": 4 + }, + "qwen3/action:policy/distilled": { + "guidance_interval": null, + "guidance_scale": 1.0, + "max_sequence_length": 4096, + "num_inference_steps": 4 + }, + "qwen3/action:policy/distilled/droid-policy": { + "guidance_interval": [ + 960.0, + 1001.0 + ], + "guidance_scale": 1.0, + "max_sequence_length": 4096, + "num_inference_steps": 4 + }, + "qwen3/image/base": { + "frame_rate": 24.0, + "guidance_interval": [ + 400.0, + 1000.0 + ], + "guidance_scale": 7.0, + "height": 1024, + "max_sequence_length": 4096, + "num_frames": 189, + "num_inference_steps": 50, + "width": 1024 + }, + "qwen3/image/distilled": { + "frame_rate": 24.0, + "guidance_interval": [ + 400.0, + 1000.0 + ], + "guidance_scale": 1.0, + "height": 1024, + "max_sequence_length": 4096, + "num_frames": 189, + "num_inference_steps": 4, + "width": 1024 + }, + "qwen3/video/base": { + "frame_rate": 24.0, + "guidance_interval": null, + "guidance_scale": 6.0, + "height": 720, + "max_sequence_length": 4096, + "num_frames": 189, + "num_inference_steps": 35, + "width": 1280 + }, + "qwen3/video/distilled": { + "frame_rate": 24.0, + "guidance_interval": null, + "guidance_scale": 1.0, + "height": 720, + "max_sequence_length": 4096, + "num_frames": 189, + "num_inference_steps": 4, + "width": 1280 + } + }, + "system_prompt": { + "checkpoint=False/v2v=False/transfer=False": false, + "checkpoint=False/v2v=False/transfer=True": false, + "checkpoint=False/v2v=True/transfer=False": true, + "checkpoint=False/v2v=True/transfer=True": false, + "checkpoint=True/v2v=False/transfer=False": true, + "checkpoint=True/v2v=False/transfer=True": false, + "checkpoint=True/v2v=True/transfer=False": true, + "checkpoint=True/v2v=True/transfer=True": false + } +} From 60c88326b128b17b65e6af763652d33c41df923c Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Fri, 30 Oct 2026 23:12:32 -0700 Subject: [PATCH 2/8] [None][refactor] Cosmos3: define the video geometry tables once, in defaults.py Signed-off-by: Igor Shovkun --- .../visual_gen/models/cosmos3/action.py | 80 +--------------- .../visual_gen/models/cosmos3/defaults.py | 92 +++++++++++++++++-- .../models/cosmos3/pipeline_cosmos3.py | 4 +- .../models/cosmos3/request_resolver.py | 3 +- .../visual_gen/models/cosmos3/transfer.py | 27 ------ .../_torch/visual_gen/test_cosmos3_action.py | 10 +- .../test_cosmos3_example_prompts.py | 4 +- .../visual_gen/test_cosmos3_transfer.py | 4 +- 8 files changed, 97 insertions(+), 127 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py index fe13630c1ade..7c3ce4172b07 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py @@ -13,6 +13,8 @@ import torch from diffusers.utils.torch_utils import randn_tensor +from .defaults import find_closest_target_size + ACTION_MODE_POLICY = "policy" ACTION_MODE_FORWARD_DYNAMICS = "forward_dynamics" ACTION_MODE_INVERSE_DYNAMICS = "inverse_dynamics" @@ -95,65 +97,6 @@ def resolve_raw_action_dim( return widths.pop() if len(widths) == 1 else None -# The reference's ``W,H`` aspect buckets: every canvas is one of these shapes, and a -# request's bucket string is carried into the caption verbatim, so a resolved frame -# size maps back to the nearest bucket rather than to its exact reduced ratio -# (832x480 reduces to "15,26", which the model never saw). -COSMOS3_ASPECT_RATIO_BUCKETS = ("1,1", "4,3", "3,4", "16,9", "9,16") - - -VIDEO_RES_SIZE_INFO: dict[str, dict[str, tuple[int, int]]] = { - "256": { - "1,1": (256, 256), - "4,3": (320, 256), - "3,4": (256, 320), - "16,9": (320, 192), - "9,16": (192, 320), - }, - "480": { - "1,1": (640, 640), - "4,3": (736, 544), - "3,4": (544, 736), - "16,9": (832, 480), - "9,16": (480, 832), - }, - "704": { - "1,1": (960, 960), - "4,3": (1088, 832), - "3,4": (832, 1088), - "16,9": (1280, 704), - "9,16": (704, 1280), - }, - "720": { - "1,1": (960, 960), - "4,3": (1104, 832), - "3,4": (832, 1104), - "16,9": (1280, 720), - "9,16": (720, 1280), - }, -} - - -COSMOS3_ACTION_RESOLUTIONS = tuple(int(key) for key in sorted(VIDEO_RES_SIZE_INFO, key=int)) - - -def normalize_action_resolution(resolution: Any) -> int: - if resolution is None: - raise ValueError("Cosmos3 action_resolution is required for action generation.") - try: - bucket = int(resolution) - except (TypeError, ValueError) as exc: - raise ValueError( - f"Cosmos3 action_resolution must be an int bucket, got {resolution!r}." - ) from exc - if bucket not in COSMOS3_ACTION_RESOLUTIONS: - raise ValueError( - f"Unknown Cosmos3 action_resolution={bucket}; " - f"expected one of {COSMOS3_ACTION_RESOLUTIONS}." - ) - return bucket - - def normalize_action_mode(mode: Any) -> str | None: if mode is None: return None @@ -324,25 +267,6 @@ def load_action_state(action: Any = None) -> torch.Tensor: return tensor -def find_closest_target_size(h: int, w: int, resolution: str | int) -> tuple[int, int]: - key = str(resolution) - if key not in VIDEO_RES_SIZE_INFO: - raise ValueError( - f"Unknown Cosmos3 action resolution={resolution!r}; " - f"expected one of {sorted(VIDEO_RES_SIZE_INFO)}." - ) - input_ratio = h / w - best_size = None - best_diff = float("inf") - for cand_w, cand_h in VIDEO_RES_SIZE_INFO[key].values(): - diff = abs(input_ratio - cand_h / cand_w) - if diff < best_diff: - best_diff = diff - best_size = (cand_w, cand_h) - assert best_size is not None - return best_size - - ACTION_IMAGE_EXTENSIONS = frozenset({".png", ".jpg", ".jpeg", ".webp", ".bmp"}) ACTION_VIDEO_EXTENSIONS = frozenset({".mp4", ".avi"}) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index bc27c35ce241..884c5dd7664e 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -20,12 +20,6 @@ from collections.abc import Mapping from typing import Any, Dict, Iterable, TypedDict -from tensorrt_llm._torch.visual_gen.models.cosmos3.action import ( - COSMOS3_ACTION_RESOLUTIONS, - EMBODIMENT_TO_DOMAIN_ID, - VIDEO_RES_SIZE_INFO, - normalize_action_resolution, -) from tensorrt_llm._torch.visual_gen.pipeline import ExtraParamSchema from tensorrt_llm.inputs.media_io import sniff_media_kind @@ -33,9 +27,87 @@ # Constant tables # --------------------------------------------------------------------------- -# Cosmos3 output resolution buckets keyed by target level, then aspect ratio; -# each value is (width, height). A source frame maps onto the bucket whose -# aspect ratio is closest (see ``find_closest_target_size`` in ``transfer.py``). +# The reference's ``W,H`` aspect buckets: every canvas is one of these shapes, and a +# request's bucket string is carried into the caption verbatim, so a resolved frame +# size maps back to the nearest bucket rather than to its exact reduced ratio +# (832x480 reduces to "15,26", which the model never saw). +COSMOS3_ASPECT_RATIO_BUCKETS = ("1,1", "4,3", "3,4", "16,9", "9,16") + + +# Output sizes per resolution level and aspect bucket, as (width, height). +VIDEO_RES_SIZE_INFO: dict[str, dict[str, tuple[int, int]]] = { + "256": { + "1,1": (256, 256), + "4,3": (320, 256), + "3,4": (256, 320), + "16,9": (320, 192), + "9,16": (192, 320), + }, + "480": { + "1,1": (640, 640), + "4,3": (736, 544), + "3,4": (544, 736), + "16,9": (832, 480), + "9,16": (480, 832), + }, + "704": { + "1,1": (960, 960), + "4,3": (1088, 832), + "3,4": (832, 1088), + "16,9": (1280, 704), + "9,16": (704, 1280), + }, + "720": { + "1,1": (960, 960), + "4,3": (1104, 832), + "3,4": (832, 1104), + "16,9": (1280, 720), + "9,16": (720, 1280), + }, +} + + +COSMOS3_ACTION_RESOLUTIONS = tuple(int(key) for key in sorted(VIDEO_RES_SIZE_INFO, key=int)) + + +def normalize_action_resolution(resolution: Any) -> int: + if resolution is None: + raise ValueError("Cosmos3 action_resolution is required for action generation.") + try: + bucket = int(resolution) + except (TypeError, ValueError) as exc: + raise ValueError( + f"Cosmos3 action_resolution must be an int bucket, got {resolution!r}." + ) from exc + if bucket not in COSMOS3_ACTION_RESOLUTIONS: + raise ValueError( + f"Unknown Cosmos3 action_resolution={bucket}; " + f"expected one of {COSMOS3_ACTION_RESOLUTIONS}." + ) + return bucket + + +def find_closest_target_size(h: int, w: int, resolution: str | int) -> tuple[int, int]: + """The ``resolution`` bucket whose aspect ratio is closest to ``h/w``, as + ``(width, height)``, so a source frame lands on a supported size undistorted.""" + key = str(resolution) + if key not in VIDEO_RES_SIZE_INFO: + raise ValueError( + f"Unknown Cosmos3 resolution={resolution!r}; " + f"expected one of {sorted(VIDEO_RES_SIZE_INFO)}." + ) + input_ratio = h / w + best_size = None + best_diff = float("inf") + for cand_w, cand_h in VIDEO_RES_SIZE_INFO[key].values(): + diff = abs(input_ratio - cand_h / cand_w) + if diff < best_diff: + best_diff = diff + best_size = (cand_w, cand_h) + assert best_size is not None + return best_size + + # The default video resolution is the 720p 16:9 bucket, ``(width, height)``. _DEFAULT_VIDEO_W, _DEFAULT_VIDEO_H = VIDEO_RES_SIZE_INFO["720"]["16,9"] COSMOS3_720P_PARAMS = { @@ -492,6 +564,8 @@ def canonical_domain_preset_key( if resolved_id == 0: return None + from .action import EMBODIMENT_TO_DOMAIN_ID + candidates: list[str] = [] for name, mapped_id in EMBODIMENT_TO_DOMAIN_ID.items(): if mapped_id != resolved_id: diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index bd5862904c31..d7c39b39d2f0 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -55,7 +55,6 @@ def tqdm(iterable, **kwargs): from .action import ( ACTION_MODE_INVERSE_DYNAMICS, ACTION_MODE_POLICY, - COSMOS3_ASPECT_RATIO_BUCKETS, action_reference_frame_step, action_reference_size, action_start_frame_offset, @@ -72,6 +71,7 @@ def tqdm(iterable, **kwargs): from .config_resolver import Cosmos3ConfigResolver from .defaults import ( COSMOS3_720P_PARAMS, + COSMOS3_ASPECT_RATIO_BUCKETS, COSMOS3_DEFAULT_RESOLUTION_TEMPLATE, COSMOS3_DEFAULT_SYSTEM_PROMPT, COSMOS3_DURATION_TEMPLATE, @@ -83,6 +83,7 @@ def tqdm(iterable, **kwargs): COSMOS3_V2V_FLOW_SHIFT, _normalize_condition_video_keep, _normalize_condition_video_latent_indexes, + find_closest_target_size, resolve_checkpoint_policy_defaults, resolve_domain_action_config, ) @@ -95,7 +96,6 @@ def tqdm(iterable, **kwargs): TRANSFER_DEFAULTS, Cosmos3TransferConfig, decode_media_to_uint8_cthw, - find_closest_target_size, load_or_compute_control_frames, pad_temporal_frames, resolve_transfer_config, diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/request_resolver.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/request_resolver.py index 90093a1489a6..960d1aa7da53 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/request_resolver.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/request_resolver.py @@ -30,13 +30,14 @@ from dataclasses import dataclass, field from typing import Any, Optional -from .action import EMBODIMENT_TO_DOMAIN_ID, normalize_action_resolution, resolve_raw_action_dim +from .action import EMBODIMENT_TO_DOMAIN_ID, resolve_raw_action_dim from .config_resolver import Conflict, Default, Resolved, same_value from .defaults import ( COSMOS3_ACTION_PARAMS, COSMOS3_DOMAIN_PRESETS, COSMOS3_GENERATION_DEFAULTS, canonical_domain_preset_key, + normalize_action_resolution, ) from .sampling import DISTILLED_GUIDANCE_SCALE diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transfer.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transfer.py index bf3b232ff19f..ec497c06b5fb 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transfer.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transfer.py @@ -22,33 +22,6 @@ from tensorrt_llm._utils import nvtx_range from tensorrt_llm.media.decoding import decode_video_reference_window -from .defaults import VIDEO_RES_SIZE_INFO - - -def find_closest_target_size(h: int, w: int, resolution: str | int) -> tuple[int, int]: - """Pick the ``resolution`` bucket whose aspect ratio best matches ``h/w``. - - Returns ``(target_w, target_h)`` so a source frame maps onto a supported - output size without distorting its aspect ratio. - """ - key = str(resolution) - if key not in VIDEO_RES_SIZE_INFO: - raise ValueError( - f"Unknown Cosmos3 transfer resolution={resolution!r}; " - f"expected one of {sorted(VIDEO_RES_SIZE_INFO)}." - ) - input_ratio = h / w - best_size = None - best_diff = float("inf") - for cand_w, cand_h in VIDEO_RES_SIZE_INFO[key].values(): - diff = abs(input_ratio - cand_h / cand_w) - if diff < best_diff: - best_diff = diff - best_size = (cand_w, cand_h) - assert best_size is not None - return best_size - - # --------------------------------------------------------------------------- # Transfer sampling + per-control-hint defaults, ported from vllm-omni's # Cosmos3 transfer implementation; the guidance and edge/blur presets are diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py index 3c53eaa1413a..de8d4fdbc0ae 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py @@ -13,15 +13,11 @@ import torch from tensorrt_llm._torch.visual_gen.models.cosmos3.action import ( - COSMOS3_ASPECT_RATIO_BUCKETS, EMBODIMENT_TO_DOMAIN_ID, EMBODIMENT_TO_RAW_ACTION_DIM, - VIDEO_RES_SIZE_INFO, action_reference_frame_step, action_reference_size, crop_action_latent, - find_closest_target_size, - normalize_action_resolution, prepare_action_latents, resize_and_pad_action_image, resolve_action_content_size, @@ -30,11 +26,15 @@ resolve_raw_action_dim, ) from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( + COSMOS3_ASPECT_RATIO_BUCKETS, COSMOS3_DOMAIN_PRESET_ALIASES, COSMOS3_DOMAIN_PRESETS, COSMOS3_EXTRA_SPECS, COSMOS3_POLICY_SAMPLING_PARAMS, + VIDEO_RES_SIZE_INFO, + find_closest_target_size, get_domain_preset, + normalize_action_resolution, resolve_checkpoint_policy_defaults, resolve_domain_action_config, ) @@ -68,7 +68,7 @@ def test_accepts_string_and_int_resolution_keys(self): ) def test_unknown_resolution_raises(self): - with pytest.raises(ValueError, match="Unknown Cosmos3 action resolution"): + with pytest.raises(ValueError, match="Unknown Cosmos3 resolution"): find_closest_target_size(480, 832, 1080) @pytest.mark.parametrize("action_resolution", sorted(VIDEO_RES_SIZE_INFO)) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_example_prompts.py b/tests/unittest/_torch/visual_gen/test_cosmos3_example_prompts.py index 96c680dc7ff9..4aa1f47c00dd 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_example_prompts.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_example_prompts.py @@ -15,11 +15,9 @@ import pytest -from tensorrt_llm._torch.visual_gen.models.cosmos3.action import ( +from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( VIDEO_RES_SIZE_INFO, find_closest_target_size, -) -from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( resolve_checkpoint_policy_defaults, resolve_domain_action_config, ) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py b/tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py index a8fce467b784..ef3014b4ed86 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py @@ -31,6 +31,7 @@ COSMOS3_EXTRA_SPECS, COSMOS3_GENERATION_DEFAULTS, VIDEO_RES_SIZE_INFO, + find_closest_target_size, ) from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniMoTPipeline from tensorrt_llm._torch.visual_gen.models.cosmos3.sampling import DISTILLED_GUIDANCE_SCALE @@ -40,7 +41,6 @@ BILATERAL_SIGMA_SPACE, _scaled_bilateral_params, decode_media_to_uint8_cthw, - find_closest_target_size, load_or_compute_control_frames, pad_temporal_frames, resolve_transfer_config, @@ -1310,5 +1310,5 @@ def test_resolution_accepts_int_or_str(self): ) def test_unknown_resolution_raises(self): - with pytest.raises(ValueError, match="Unknown Cosmos3 transfer resolution"): + with pytest.raises(ValueError, match="Unknown Cosmos3 resolution"): find_closest_target_size(720, 1280, 1080) From f8d773927a5049f0742b60e12e3b492375a2caf2 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Sat, 31 Oct 2026 00:24:16 -0700 Subject: [PATCH 3/8] [None][refactor] Cosmos3: one request-resolution path, drop the dead resolver sources and duplicate checks Signed-off-by: Igor Shovkun --- .../visual_gen/models/cosmos3/action.py | 2 +- .../models/cosmos3/config_resolver.py | 68 ++----- .../visual_gen/models/cosmos3/defaults.py | 156 +--------------- .../visual_gen/models/cosmos3/geometry.py | 101 +++++++++++ .../models/cosmos3/pipeline_cosmos3.py | 95 +++------- .../models/cosmos3/request_resolver.py | 170 ++++++++---------- .../visual_gen/models/cosmos3/sampling.py | 26 +-- .../models/cosmos3/transformer_cosmos3.py | 24 +-- .../visual_gen/cosmos3_request_golden.py | 23 +-- .../_torch/visual_gen/test_cosmos3_action.py | 10 +- .../test_cosmos3_config_resolver.py | 21 +-- .../visual_gen/test_cosmos3_distilled.py | 11 +- .../_torch/visual_gen/test_cosmos3_edge.py | 37 +--- .../test_cosmos3_example_prompts.py | 6 +- .../visual_gen/test_cosmos3_pipeline.py | 18 -- .../test_cosmos3_request_resolver.py | 29 +-- .../visual_gen/test_cosmos3_transfer.py | 13 +- .../visual_gen/test_cosmos3_transformer.py | 19 +- 18 files changed, 305 insertions(+), 524 deletions(-) create mode 100644 tensorrt_llm/_torch/visual_gen/models/cosmos3/geometry.py diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py index 7c3ce4172b07..7172c01ac67e 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py @@ -13,7 +13,7 @@ import torch from diffusers.utils.torch_utils import randn_tensor -from .defaults import find_closest_target_size +from .geometry import find_closest_target_size ACTION_MODE_POLICY = "policy" ACTION_MODE_FORWARD_DYNAMICS = "forward_dynamics" diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/config_resolver.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/config_resolver.py index f2da4db2f230..b5756e584e35 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/config_resolver.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/config_resolver.py @@ -12,7 +12,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""One place that turns a Cosmos3 checkpoint and the serving config into settings. +"""One place that turns a Cosmos3 checkpoint's files into settings. A Cosmos3 export carries the same facts in several files, written by different exporters that are not kept in sync: the training config (``config.json``), the @@ -33,7 +33,7 @@ import math import os from collections.abc import Mapping -from dataclasses import dataclass +from dataclasses import dataclass, replace from enum import Enum from typing import Any, Optional @@ -69,17 +69,6 @@ def label(self) -> str: return f"{self.file}:{self.path}" -@dataclass(frozen=True) -class Serving: - """An attribute path in the serving config (``VisualGenArgs`` as the pipeline sees it).""" - - path: str - - @property - def label(self) -> str: - return f"serving:{self.path}" - - @dataclass(frozen=True) class Default: """The value used when no source declares the setting.""" @@ -98,9 +87,6 @@ class Rule: conflict: Conflict = Conflict.WARN required: bool = True """``False``: an unresolved setting is ``None`` instead of an error.""" - nullable: bool = False - """``True``: a source holding ``null`` declares the setting (e.g. "no window"); - otherwise ``null`` counts as absent.""" @dataclass(frozen=True) @@ -111,9 +97,9 @@ class Resolved: """Lower-priority sources that held a different value.""" -# Settings every Cosmos3 checkpoint resolves. Priority: serving config, then the -# training config (config.json), then the component configs, then the export -# metadata and the diffusers index. +# Settings every Cosmos3 checkpoint resolves. Priority: the training config +# (config.json), then the component configs, then the export metadata and the +# diffusers index. RULES: dict[str, Rule] = { "architecture.position_embedding_type": Rule( ( @@ -228,7 +214,7 @@ def _step(obj: Any, key: str) -> Any: return obj.get(key, _MISSING) if obj is None or isinstance(obj, (str, int, float, bool, list, tuple)): return _MISSING - return getattr(obj, key, _MISSING) + return getattr(obj, key, _MISSING) # a pretrained config object standing in for a file def _walk(doc: Any, path: str) -> Any: @@ -255,21 +241,15 @@ def same_value(a: Any, b: Any) -> bool: class Cosmos3ConfigResolver: - """Applies :data:`RULES` to a checkpoint directory and the serving config. + """Applies :data:`RULES` to a checkpoint directory. - ``checkpoint_dir`` may be ``None``: only ``pretrained_config`` (standing in for - ``transformer/config.json``) and ``serving`` are consulted then. + ``pretrained_config`` stands in for ``transformer/config.json`` when given; + ``checkpoint_dir`` may then be ``None``. """ - def __init__( - self, - checkpoint_dir: Optional[str], - pretrained_config: Any = None, - serving: Any = None, - ) -> None: + def __init__(self, checkpoint_dir: Optional[str], pretrained_config: Any = None) -> None: self._dir = checkpoint_dir self._pretrained = pretrained_config - self._serving = serving self._docs: dict[str, Any] = {} def _doc(self, file: str) -> Any: @@ -289,8 +269,6 @@ def _lookup(self, source: Any) -> Any: if isinstance(source, Key): doc = self._doc(source.file) return _MISSING if doc is None else _walk(doc, source.path) - if isinstance(source, Serving): - return _MISSING if self._serving is None else _walk(self._serving, source.path) if isinstance(source, Default): return source.value raise TypeError(f"unknown source {source!r}") @@ -302,14 +280,10 @@ def resolve_one(self, name: str, rule: Rule) -> Optional[Resolved]: if winner is not None and isinstance(source, Default): break value = self._lookup(source) - if value is _MISSING or (value is None and not rule.nullable): - continue - if value is None and isinstance(source, Serving): - continue # an unset serving field defers to the checkpoint + if value is _MISSING or value is None: + continue # a null in a file counts as absent if winner is None: winner = Resolved(value, source.label) - if isinstance(source, Serving): - break # an override is meant to differ from the checkpoint elif not same_value(value, winner.value): disagreeing.append(f"{source.label}={value!r}") if winner is None: @@ -377,14 +351,6 @@ def get(name: str) -> Any: """``pretrained_config`` attributes the transformer reads, and the setting each holds.""" -def architecture_defaults() -> dict[str, Any]: - """The table's default for each attribute in :data:`ARCHITECTURE_ATTRIBUTES`.""" - out = {} - for attr, name in ARCHITECTURE_ATTRIBUTES.items(): - out[attr] = next(s.value for s in RULES[name].sources if isinstance(s, Default)) - return out - - def fill_pretrained_config(pretrained_config: Any) -> Any: """Set the transformer-construction attributes of ``pretrained_config`` that the export omits, from :data:`RULES`. Idempotent; returns the config. @@ -397,11 +363,11 @@ def fill_pretrained_config(pretrained_config: Any) -> Any: if getattr(pretrained_config, attr, None) is not None: continue rule = RULES[name] - transformer_only = Rule( - tuple(s for s in rule.sources if not isinstance(s, Key) or s.file == TRANSFORMER), - rule.conflict, - rule.required, - rule.nullable, + transformer_only = replace( + rule, + sources=tuple( + s for s in rule.sources if not isinstance(s, Key) or s.file == TRANSFORMER + ), ) r = resolver.resolve_one(name, transformer_only) if r is not None: diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index 884c5dd7664e..4761a02ca29c 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -12,14 +12,20 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Per-model default generation parameters for Cosmos3 pipelines. - -Shared by the Cosmos3 OmniMoT text-to-video and image-to-video generation paths. +"""The Cosmos3 tables: generation defaults per family and mode, the action and +policy presets per embodiment domain, the prompt templates, and the schema of the +request's extra parameters. The request resolver reads them in precedence order. """ from collections.abc import Mapping from typing import Any, Dict, Iterable, TypedDict +from tensorrt_llm._torch.visual_gen.models.cosmos3.action import EMBODIMENT_TO_DOMAIN_ID +from tensorrt_llm._torch.visual_gen.models.cosmos3.geometry import ( + COSMOS3_ACTION_RESOLUTIONS, + VIDEO_RES_SIZE_INFO, + normalize_action_resolution, +) from tensorrt_llm._torch.visual_gen.pipeline import ExtraParamSchema from tensorrt_llm.inputs.media_io import sniff_media_kind @@ -27,87 +33,6 @@ # Constant tables # --------------------------------------------------------------------------- -# The reference's ``W,H`` aspect buckets: every canvas is one of these shapes, and a -# request's bucket string is carried into the caption verbatim, so a resolved frame -# size maps back to the nearest bucket rather than to its exact reduced ratio -# (832x480 reduces to "15,26", which the model never saw). -COSMOS3_ASPECT_RATIO_BUCKETS = ("1,1", "4,3", "3,4", "16,9", "9,16") - - -# Output sizes per resolution level and aspect bucket, as (width, height). -VIDEO_RES_SIZE_INFO: dict[str, dict[str, tuple[int, int]]] = { - "256": { - "1,1": (256, 256), - "4,3": (320, 256), - "3,4": (256, 320), - "16,9": (320, 192), - "9,16": (192, 320), - }, - "480": { - "1,1": (640, 640), - "4,3": (736, 544), - "3,4": (544, 736), - "16,9": (832, 480), - "9,16": (480, 832), - }, - "704": { - "1,1": (960, 960), - "4,3": (1088, 832), - "3,4": (832, 1088), - "16,9": (1280, 704), - "9,16": (704, 1280), - }, - "720": { - "1,1": (960, 960), - "4,3": (1104, 832), - "3,4": (832, 1104), - "16,9": (1280, 720), - "9,16": (720, 1280), - }, -} - - -COSMOS3_ACTION_RESOLUTIONS = tuple(int(key) for key in sorted(VIDEO_RES_SIZE_INFO, key=int)) - - -def normalize_action_resolution(resolution: Any) -> int: - if resolution is None: - raise ValueError("Cosmos3 action_resolution is required for action generation.") - try: - bucket = int(resolution) - except (TypeError, ValueError) as exc: - raise ValueError( - f"Cosmos3 action_resolution must be an int bucket, got {resolution!r}." - ) from exc - if bucket not in COSMOS3_ACTION_RESOLUTIONS: - raise ValueError( - f"Unknown Cosmos3 action_resolution={bucket}; " - f"expected one of {COSMOS3_ACTION_RESOLUTIONS}." - ) - return bucket - - -def find_closest_target_size(h: int, w: int, resolution: str | int) -> tuple[int, int]: - """The ``resolution`` bucket whose aspect ratio is closest to ``h/w``, as - ``(width, height)``, so a source frame lands on a supported size undistorted.""" - key = str(resolution) - if key not in VIDEO_RES_SIZE_INFO: - raise ValueError( - f"Unknown Cosmos3 resolution={resolution!r}; " - f"expected one of {sorted(VIDEO_RES_SIZE_INFO)}." - ) - input_ratio = h / w - best_size = None - best_diff = float("inf") - for cand_w, cand_h in VIDEO_RES_SIZE_INFO[key].values(): - diff = abs(input_ratio - cand_h / cand_w) - if diff < best_diff: - best_diff = diff - best_size = (cand_w, cand_h) - assert best_size is not None - return best_size - - # The default video resolution is the 720p 16:9 bucket, ``(width, height)``. _DEFAULT_VIDEO_W, _DEFAULT_VIDEO_H = VIDEO_RES_SIZE_INFO["720"]["16,9"] COSMOS3_720P_PARAMS = { @@ -564,8 +489,6 @@ def canonical_domain_preset_key( if resolved_id == 0: return None - from .action import EMBODIMENT_TO_DOMAIN_ID - candidates: list[str] = [] for name, mapped_id in EMBODIMENT_TO_DOMAIN_ID.items(): if mapped_id != resolved_id: @@ -589,67 +512,6 @@ def get_domain_preset( return COSMOS3_DOMAIN_PRESETS[key] -def resolve_domain_action_config( - *, - domain_name: str | None = None, - domain_id: str | int | None = None, - raw_action_dim: int | None = None, - action_chunk_size: int | None = None, - action_resolution: int | None = None, - frame_rate: float | None = None, - action_fps: float | None = None, - use_state: bool | None = None, - checkpoint_policy_defaults: Mapping[str, Any] | None = None, - embodiments: Mapping[str, Any] | None = None, - input_action_dim: int | None = None, -) -> dict[str, Any]: - """The action fields a request runs with, by the request rules - (``request_resolver.ACTION_RULES``): the caller's value, else the checkpoint policy - block when its domain is the request's, else the checkpoint's embodiment table - (``embodiments``) or the code's width table, else the domain preset, else the - generic action defaults. ``num_frames`` is always ``action_chunk_size + 1``. A - caller's value that differs from the recommendation is reported in ``warnings``. - """ - from .request_resolver import RequestContext, embodiment_entry, policy_domain, resolve_action - - policy_defaults = dict(checkpoint_policy_defaults or {}) - domain_name, policy_active, warnings = policy_domain(policy_defaults, domain_name, domain_id) - ctx = RequestContext( - family="", - action_mode="action", - policy_defaults=policy_defaults, - policy_active=policy_active, - domain_name=domain_name, - domain_id=None if domain_id is None else int(domain_id), - embodiment=embodiment_entry(embodiments or {}, input_action_dim, domain_name, domain_id), - checkpoint_has_embodiments=bool(embodiments), - ) - resolved = resolve_action( - ctx, - raw_action_dim=raw_action_dim, - action_chunk_size=action_chunk_size, - action_resolution=action_resolution, - frame_rate=frame_rate, - action_fps=action_fps, - use_state=use_state, - ) - width = resolved["raw_action_dim"] - return { - "domain_name": domain_name, - "checkpoint_policy_domain_matches": policy_active, - "raw_action_dim": None if width is None else int(width), - "action_chunk_size": int(resolved["action_chunk_size"]), - "action_resolution": resolved["action_resolution"], - "frame_rate": float(resolved["frame_rate"]), - "action_fps": float(resolved["action_fps"]), - "use_state": bool(resolved["use_state"]), - "num_frames": int(resolved["num_frames"]), - "preset_key": ctx.preset_key, - "warnings": warnings + resolved.warnings, - "sources": {name: r.source for name, r in resolved.values.items()}, - } - - # Video-to-video and transfer sample on the shifted flow grid with this shift; the # other modes carry ``flow_shift`` in their tables above or take the checkpoint's. COSMOS3_V2V_FLOW_SHIFT = 10.0 diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/geometry.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/geometry.py new file mode 100644 index 000000000000..33ee45c1a444 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/geometry.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The Cosmos3 output geometry: the aspect buckets, the output size per resolution +level and bucket, and the lookups every mode shares. A leaf module: nothing here +depends on the rest of the package.""" + +from __future__ import annotations + +from typing import Any + +# The reference's ``W,H`` aspect buckets: every canvas is one of these shapes, and a +# request's bucket string is carried into the caption verbatim, so a resolved frame +# size maps back to the nearest bucket rather than to its exact reduced ratio +# (832x480 reduces to "15,26", which the model never saw). +COSMOS3_ASPECT_RATIO_BUCKETS = ("1,1", "4,3", "3,4", "16,9", "9,16") + + +# Output sizes per resolution level and aspect bucket, as (width, height). +VIDEO_RES_SIZE_INFO: dict[str, dict[str, tuple[int, int]]] = { + "256": { + "1,1": (256, 256), + "4,3": (320, 256), + "3,4": (256, 320), + "16,9": (320, 192), + "9,16": (192, 320), + }, + "480": { + "1,1": (640, 640), + "4,3": (736, 544), + "3,4": (544, 736), + "16,9": (832, 480), + "9,16": (480, 832), + }, + "704": { + "1,1": (960, 960), + "4,3": (1088, 832), + "3,4": (832, 1088), + "16,9": (1280, 704), + "9,16": (704, 1280), + }, + "720": { + "1,1": (960, 960), + "4,3": (1104, 832), + "3,4": (832, 1104), + "16,9": (1280, 720), + "9,16": (720, 1280), + }, +} + + +COSMOS3_ACTION_RESOLUTIONS = tuple(int(key) for key in sorted(VIDEO_RES_SIZE_INFO, key=int)) + + +def normalize_action_resolution(resolution: Any) -> int: + if resolution is None: + raise ValueError("Cosmos3 action_resolution is required for action generation.") + try: + bucket = int(resolution) + except (TypeError, ValueError) as exc: + raise ValueError( + f"Cosmos3 action_resolution must be an int bucket, got {resolution!r}." + ) from exc + if bucket not in COSMOS3_ACTION_RESOLUTIONS: + raise ValueError( + f"Unknown Cosmos3 action_resolution={bucket}; " + f"expected one of {COSMOS3_ACTION_RESOLUTIONS}." + ) + return bucket + + +def find_closest_target_size(h: int, w: int, resolution: str | int) -> tuple[int, int]: + """The ``resolution`` bucket whose aspect ratio is closest to ``h/w``, as + ``(width, height)``, so a source frame lands on a supported size undistorted.""" + key = str(resolution) + if key not in VIDEO_RES_SIZE_INFO: + raise ValueError( + f"Unknown Cosmos3 resolution={resolution!r}; " + f"expected one of {sorted(VIDEO_RES_SIZE_INFO)}." + ) + input_ratio = h / w + best_size = None + best_diff = float("inf") + for cand_w, cand_h in VIDEO_RES_SIZE_INFO[key].values(): + diff = abs(input_ratio - cand_h / cand_w) + if diff < best_diff: + best_diff = diff + best_size = (cand_w, cand_h) + assert best_size is not None + return best_size diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index d7c39b39d2f0..b42eac2af35f 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -71,7 +71,6 @@ def tqdm(iterable, **kwargs): from .config_resolver import Cosmos3ConfigResolver from .defaults import ( COSMOS3_720P_PARAMS, - COSMOS3_ASPECT_RATIO_BUCKETS, COSMOS3_DEFAULT_RESOLUTION_TEMPLATE, COSMOS3_DEFAULT_SYSTEM_PROMPT, COSMOS3_DURATION_TEMPLATE, @@ -83,13 +82,12 @@ def tqdm(iterable, **kwargs): COSMOS3_V2V_FLOW_SHIFT, _normalize_condition_video_keep, _normalize_condition_video_latent_indexes, - find_closest_target_size, resolve_checkpoint_policy_defaults, - resolve_domain_action_config, ) +from .geometry import COSMOS3_ASPECT_RATIO_BUCKETS, find_closest_target_size from .guardrails import check_video_safety, download_guardrail_checkpoint from .negative_prompt import COSMOS3_VIDEO_NEGATIVE_PROMPT -from .request_resolver import RequestContext, resolve_generation +from .request_resolver import RequestContext, resolve_domain_action_config, resolve_generation from .sampling import DISTILLED_GUIDANCE_SCALE, Cosmos3SamplingPolicy, load_scheduler from .sound_tokenizer import LatentAutoEncoderV2 from .transfer import ( @@ -312,19 +310,6 @@ def _assert_anchor_matches(image_latent: torch.Tensor, latents: torch.Tensor) -> ) -def _validate_temporal_compression(transformer, vae_scale_factor_temporal: int) -> None: - """A config-declared temporal compression factor must match the VAE.""" - if ( - getattr(transformer, "temporal_compression_factor_declared", False) - and transformer.temporal_compression_factor != vae_scale_factor_temporal - ): - raise ValueError( - f"Transformer config declares temporal_compression_factor=" - f"{transformer.temporal_compression_factor}, but the VAE reports " - f"scale_factor_temporal={vae_scale_factor_temporal}." - ) - - def _condition_pixel_frame_count( condition_video_latent_indexes: Iterable[int], temporal_compression: int, @@ -454,33 +439,12 @@ def __init__(self, pipeline_config): super().__init__(pipeline_config) - def _mode_params( - self, - output_type: str, - *, - action_mode: Optional[str] = None, - checkpoint_policy_domain_matches: bool = True, - ) -> dict: - """Generation defaults selected from the request's output and action mode.""" - normalized_action_mode = normalize_action_mode(action_mode) - request_mode = "action" if normalized_action_mode is not None else output_type - params = COSMOS3_GENERATION_DEFAULTS[(self.family, request_mode)] - checkpoint_policy_defaults = getattr(self, "checkpoint_policy_defaults", {}) - if ( - normalized_action_mode == ACTION_MODE_POLICY - and checkpoint_policy_domain_matches - and checkpoint_policy_defaults - ): - sampling_keys = {"num_inference_steps", "guidance_scale", "guidance_interval"} - return { - **params, - **{ - key: value - for key, value in checkpoint_policy_defaults.items() - if key in sampling_keys - }, - } - return params + def _mode_table(self, output_type: str, action_mode: Optional[str] = None) -> dict: + """The family's generation table for the request's output and action mode. The + request rules read it through :func:`resolve_generation`; the pipeline reads + only ``flow_shift`` from it directly.""" + request_mode = "action" if normalize_action_mode(action_mode) is not None else output_type + return COSMOS3_GENERATION_DEFAULTS[(self.family, request_mode)] def _request_context( self, @@ -488,7 +452,6 @@ def _request_context( *, action_mode: Optional[str] = None, checkpoint_policy_domain_matches: bool = True, - **workflow, ) -> RequestContext: normalized_action_mode = normalize_action_mode(action_mode) policy_defaults = getattr(self, "checkpoint_policy_defaults", {}) or {} @@ -505,7 +468,6 @@ def _request_context( and bool(policy_defaults) ), checkpoint_use_system_prompt=getattr(self, "default_use_system_prompt", False), - **workflow, ) def _resolve_generation_params( @@ -653,9 +615,7 @@ def load_standard_components( # The transformer config is read from the checkpoint directory, not from the # loaded pretrained config, whose omitted keys the transformer has already filled. - self.cosmos3_config = Cosmos3ConfigResolver( - checkpoint_dir, serving=getattr(self, "pipeline_config", None) - ).resolve() + self.cosmos3_config = Cosmos3ConfigResolver(checkpoint_dir).resolve() logger.info(self.cosmos3_config.describe()) cfg = self.cosmos3_config self.default_use_system_prompt = cfg.prompt.use_system_prompt @@ -697,7 +657,6 @@ def load_standard_components( dtype=torch.bfloat16, ) - _validate_temporal_compression(self.transformer, self.vae_scale_factor_temporal) self.transformer.temporal_compression_factor = self.vae_scale_factor_temporal if PipelineComponent.SCHEDULER not in skip_components: @@ -707,7 +666,9 @@ def load_standard_components( # distilled ones. The policy holds the derived immutable facts. self.scheduler = load_scheduler(checkpoint_dir) self.sampling = Cosmos3SamplingPolicy.from_scheduler( - self.scheduler, native_flow_schedule=self.use_native_flow_schedule + self.scheduler, + native_flow_schedule=self.use_native_flow_schedule, + fixed_sigmas=cfg.sampling.fixed_sigmas, ) _validate_sampling_recipe(self.family, self.use_native_flow_schedule, self.sampling) # Each stream's variants derive from that stream's own instance, so @@ -756,12 +717,12 @@ def load_standard_components( @property def default_warmup_resolutions(self): - video = self._mode_params("video") + video = self.default_generation_params return [(video["height"], video["width"])] @property def default_warmup_num_frames(self): - return [self._mode_params("video")["num_frames"]] + return [self.default_generation_params["num_frames"]] @property def default_warmup_steps(self): @@ -778,15 +739,15 @@ def default_warmup_steps(self): def default_generation_params(self): """Fields merged by the executor into every request. - These are the video-mode values — what an unmodified request runs. - A request that selects another mode re-resolves them in ``infer()``, - which tells a merged default from a caller-supplied value via - ``model_fields_set``. Key membership also declares these fields - supported during request validation. ``flow_shift`` is + These are the video-mode values by the request rules — what an unmodified + request runs, the distilled schedule included. A request that selects another + mode re-resolves them in ``infer()``, which tells a merged default from a + caller-supplied value via ``model_fields_set``. Key membership also declares + these fields supported during request validation. ``flow_shift`` is pipeline-internal, not a request field. """ - defaults = {k: v for k, v in self._mode_params("video").items() if k != "flow_shift"} - return {**defaults, **self.sampling.generation_default_overrides()} + fields = [k for k in self._mode_table("video") if k != "flow_shift"] + return self._resolve_generation_params("video", **{f: None for f in fields}) def classify_request_failure(self, exc: BaseException) -> Optional[str]: """Cosmos3 rejects unusable request content with ``ValueError`` and @@ -822,12 +783,8 @@ def ref_slot_specs(self) -> dict[str, RefSlotSpec]: } def _run_warmup(self, height: int, width: int, num_frames: int, steps: int) -> None: - # Checkpoint-aware guidance: distilled defaults carry a concrete 1.0; - # base defaults leave it None ("by mode") — warmup runs the video mode. - defaults = self.default_generation_params + defaults = self.default_generation_params # the video mode, distilled schedule included guidance_scale = defaults["guidance_scale"] - if guidance_scale is None: - guidance_scale = self._mode_params("video")["guidance_scale"] with torch.no_grad(): self.forward( prompt="warmup", @@ -1734,11 +1691,7 @@ def _resolve_request( + ", ".join(f"{k} [{v}]" for k, v in action_cfg["sources"].items()) ) - mode_params = self._mode_params( - output_type, - action_mode=normalized_action_mode, - checkpoint_policy_domain_matches=checkpoint_policy_domain_matches, - ) + mode_table = self._mode_table(output_type, normalized_action_mode) values_to_resolve = { "num_inference_steps": num_inference_steps, "guidance_scale": guidance_scale, @@ -1818,7 +1771,7 @@ def _resolve_request( use_state = action_cfg["use_state"] enable_audio = False - mode_shift = mode_params.get("flow_shift") + mode_shift = mode_table.get("flow_shift") if mode_shift is None: mode_shift = self.sampling.checkpoint_flow_shift if is_v2v: diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/request_resolver.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/request_resolver.py index 960d1aa7da53..960fd0f9410d 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/request_resolver.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/request_resolver.py @@ -14,14 +14,13 @@ # limitations under the License. """Per-request resolution of Cosmos3 generation and action settings. -The engine-level :mod:`config_resolver` settles what the checkpoint and the serving -config say. This module settles what one request runs with: for every field, -:data:`GENERATION_RULES` and :data:`ACTION_RULES` name the sources in priority order, -the caller's value first, then the checkpoint's facts (the distilled schedule, the -policy block, the embodiment table), then the per-mode and per-domain tables, then a -default. The first source holding a value wins; a lower source that disagrees with a -caller's value is reported as a warning (or an error, where the checkpoint is the -authority), and every resolved value remembers its source. +The engine-level :mod:`config_resolver` settles what the checkpoint says. This module +settles what one request runs with: for every field, :data:`GENERATION_RULES` and +:data:`ACTION_RULES` name the sources in priority order, the caller's value first, +then the checkpoint's facts (the distilled schedule, the policy block), then the +per-mode and per-domain tables, then a default. The first source holding a value +wins; where a rule says so, a recommendation that disagrees with the caller's value +is reported as a warning, and every resolved value remembers its source. """ from __future__ import annotations @@ -37,8 +36,8 @@ COSMOS3_DOMAIN_PRESETS, COSMOS3_GENERATION_DEFAULTS, canonical_domain_preset_key, - normalize_action_resolution, ) +from .geometry import normalize_action_resolution from .sampling import DISTILLED_GUIDANCE_SCALE _MISSING = object() @@ -60,10 +59,6 @@ class RequestContext: """The policy block applies: policy mode, and the request's domain is the block's.""" domain_name: Optional[str] = None domain_id: Optional[int] = None - embodiment: Mapping[str, Any] = field(default_factory=dict) - """The checkpoint's entry for the request's embodiment: ``domain_id``, - ``action_dim``. Empty when the checkpoint declares no embodiment table.""" - checkpoint_has_embodiments: bool = False checkpoint_use_system_prompt: bool = False is_v2v: bool = False is_transfer: bool = False @@ -121,30 +116,12 @@ def lookup(self, ctx, field_name, given, values): return _MISSING if value is None else value -@dataclass(frozen=True) -class Embodiment: - """The checkpoint's embodiment table.""" - - key: str - - @property - def label(self) -> str: - return f"checkpoint:conditioning.embodiments.{self.key}" - - def lookup(self, ctx, field_name, given, values): - value = ctx.embodiment.get(self.key) - return _MISSING if value is None else value - - class KnownWidth: - """The code's table of action widths per embodiment; silent when the checkpoint - declares its own embodiments.""" + """The code's table of action widths per embodiment.""" label = "embodiment table (defaults)" def lookup(self, ctx, field_name, given, values): - if ctx.checkpoint_has_embodiments: - return _MISSING width = resolve_raw_action_dim(domain_name=ctx.domain_name, domain_id=ctx.domain_id) return _MISSING if width is None else width @@ -176,13 +153,15 @@ def label(self) -> str: def lookup(self, ctx, field_name, given, values): table = COSMOS3_GENERATION_DEFAULTS.get((ctx.family, self.mode or ctx.table_mode)) - value = table.get(self.key) if table else None - return _MISSING if value is None else value + if not table or self.key not in table: + return _MISSING + return table[self.key] # an explicit ``None`` in a table is a value: the field is off @dataclass(frozen=True) -class Workflow: - """A value a workflow fixes: transfer never prepends the system prompt, V2V does.""" +class Context: + """A value the request context fixes: a workflow rule (transfer never prepends the + system prompt, video-to-video always does) or a checkpoint fact.""" label: str pick: Callable[[RequestContext], Any] @@ -192,16 +171,6 @@ def lookup(self, ctx, field_name, given, values): return _MISSING if value is None else value -@dataclass(frozen=True) -class Checkpoint: - label: str - pick: Callable[[RequestContext], Any] - - def lookup(self, ctx, field_name, given, values): - value = self.pick(ctx) - return _MISSING if value is None else value - - @dataclass(frozen=True) class Derived: """A value computed from fields resolved before it.""" @@ -218,6 +187,8 @@ def lookup(self, ctx, field_name, given, values): class RequestRule: sources: tuple conflict: Conflict = Conflict.IGNORE + """``WARN``: a caller's value that differs from the first recommendation below it + is reported in the result's warnings.""" required: bool = True @@ -261,18 +232,16 @@ class RequestRule: "use_system_prompt": RequestRule( ( Given(), - Workflow("transfer (never)", lambda ctx: False if ctx.is_transfer else None), - Workflow("video-to-video (always)", lambda ctx: True if ctx.is_v2v else None), - Checkpoint( - "checkpoint:use_system_prompt", lambda ctx: ctx.checkpoint_use_system_prompt - ), + Context("transfer (never)", lambda ctx: False if ctx.is_transfer else None), + Context("video-to-video (always)", lambda ctx: True if ctx.is_v2v else None), + Context("checkpoint:use_system_prompt", lambda ctx: ctx.checkpoint_use_system_prompt), ) ), } ACTION_RULES: dict[str, RequestRule] = { "raw_action_dim": RequestRule( - (Given(), Policy("raw_action_dim"), Embodiment("action_dim"), KnownWidth()), + (Given(), Policy("raw_action_dim"), KnownWidth()), conflict=Conflict.WARN, required=False, ), @@ -324,10 +293,6 @@ class ResolvedRequest: def __getitem__(self, name: str) -> Any: return self.values[name].value - def get(self, name: str, default: Any = None) -> Any: - r = self.values.get(name) - return default if r is None else r.value - def source_of(self, name: str) -> str: return self.values[name].source @@ -342,7 +307,7 @@ def _resolve( plain: dict[str, Any] = {} for name, rule in rules.items(): winner: Optional[Resolved] = None - disagreeing: list[str] = [] + disagreeing: list[tuple[str, Any]] = [] for source in rule.sources: if winner is not None and isinstance(source, Default): break @@ -357,7 +322,7 @@ def _resolve( continue # Only the first recommendation below the caller's value is compared. if not same_value(value, winner.value): - disagreeing.append(f"{value}") + disagreeing.append((source.label, value)) break if winner is None: if rule.required: @@ -365,16 +330,14 @@ def _resolve( out.values[name] = Resolved(None, "unresolved") plain[name] = None continue - if disagreeing and winner.source == Given.label: - message = ( + if disagreeing and winner.source == Given.label and rule.conflict is Conflict.WARN: + out.warnings.append( f"Cosmos3 {name}={winner.value} differs from recommended " - f"{' / '.join(disagreeing)} for domain {domain_label}." + f"{' / '.join(str(v) for _, v in disagreeing)} for domain {domain_label}." ) - if rule.conflict is Conflict.ERROR: - raise ValueError(message) - if rule.conflict is Conflict.WARN: - out.warnings.append(message) - out.values[name] = Resolved(winner.value, winner.source, tuple(disagreeing)) + out.values[name] = Resolved( + winner.value, winner.source, tuple(f"{label}={v!r}" for label, v in disagreeing) + ) plain[name] = winner.value return out @@ -459,30 +422,55 @@ def resolve_action(ctx: RequestContext, **given: Any) -> ResolvedRequest: return resolved -def embodiment_entry( - embodiments: Mapping[str, Any], input_action_dim: Optional[int], domain_name, domain_id -): - """The checkpoint's embodiment facts for a request: ``{"domain_id", "action_dim"}``, - found by name or by id; empty when the checkpoint lists no such embodiment.""" - name = str(domain_name).strip().lower() if domain_name and str(domain_name).strip() else None - spec = None - if name is not None: - spec = embodiments.get(name) - elif domain_id is not None: - spec = next( - ( - s - for s in embodiments.values() - if s.get("domain_id") is not None and int(s["domain_id"]) == int(domain_id) - ), - None, - ) - if not spec: - return {} - width = spec.get("input_action_dim") - if width is None: - width = spec.get("raw_action_dim", input_action_dim) +def resolve_domain_action_config( + *, + domain_name: str | None = None, + domain_id: str | int | None = None, + raw_action_dim: int | None = None, + action_chunk_size: int | None = None, + action_resolution: int | None = None, + frame_rate: float | None = None, + action_fps: float | None = None, + use_state: bool | None = None, + checkpoint_policy_defaults: Mapping[str, Any] | None = None, +) -> dict[str, Any]: + """The action fields a request runs with, by :data:`ACTION_RULES`: the caller's + value, else the checkpoint policy block when its domain is the request's, else the + code's width table, else the domain preset, else the generic action defaults. + ``num_frames`` is always ``action_chunk_size + 1``. A caller's value that differs + from the recommendation is reported in ``warnings``. + """ + policy_defaults = dict(checkpoint_policy_defaults or {}) + domain_name, policy_active, warnings = policy_domain(policy_defaults, domain_name, domain_id) + ctx = RequestContext( + family="", + action_mode="action", + policy_defaults=policy_defaults, + policy_active=policy_active, + domain_name=domain_name, + domain_id=None if domain_id is None else int(domain_id), + ) + resolved = resolve_action( + ctx, + raw_action_dim=raw_action_dim, + action_chunk_size=action_chunk_size, + action_resolution=action_resolution, + frame_rate=frame_rate, + action_fps=action_fps, + use_state=use_state, + ) + width = resolved["raw_action_dim"] return { - "domain_id": None if spec.get("domain_id") is None else int(spec["domain_id"]), - "action_dim": None if width is None else int(width), + "domain_name": domain_name, + "checkpoint_policy_domain_matches": policy_active, + "raw_action_dim": None if width is None else int(width), + "action_chunk_size": int(resolved["action_chunk_size"]), + "action_resolution": resolved["action_resolution"], + "frame_rate": float(resolved["frame_rate"]), + "action_fps": float(resolved["action_fps"]), + "use_state": bool(resolved["use_state"]), + "num_frames": int(resolved["num_frames"]), + "preset_key": ctx.preset_key, + "warnings": warnings + resolved.warnings, + "sources": {name: r.source for name, r in resolved.values.items()}, } diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/sampling.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/sampling.py index cfdbc3531356..4534e57aa45b 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/sampling.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/sampling.py @@ -113,15 +113,21 @@ class Cosmos3SamplingPolicy: @classmethod def from_scheduler( - cls, scheduler: Any, native_flow_schedule: bool = False + cls, + scheduler: Any, + native_flow_schedule: bool = False, + fixed_sigmas: "tuple[float, ...] | None" = None, ) -> "Cosmos3SamplingPolicy": """Derive the policy from a loaded scheduler's config. Valid recipes: UniPC without fixed sigmas (base) and stochastic FlowMatchEuler with a nonempty ``fixed_step_sampler_config.t_list`` - (distilled); anything else fails here, at load time. + (distilled); anything else fails here, at load time. ``fixed_sigmas`` is the + schedule as the checkpoint resolver read it; when omitted it is read from the + scheduler's own config. """ - fixed_sigmas = _resolve_distilled_sigmas(scheduler.config) + if fixed_sigmas is None: + fixed_sigmas = _resolve_distilled_sigmas(scheduler.config) is_unipc = isinstance(scheduler, UniPCMultistepScheduler) is_flow_match = isinstance(scheduler, FlowMatchEulerDiscreteScheduler) @@ -175,20 +181,6 @@ def from_scheduler( def is_distilled(self) -> bool: return self.fixed_sigmas is not None - def generation_default_overrides(self) -> dict: - """Checkpoint-mandated overrides of the table generation defaults. - - Merged over the checkpoint family's video table by the pipeline's - ``default_generation_params``, so executor-merged requests arrive - carrying the checkpoint's true defaults. - """ - if not self.is_distilled: - return {} - return { - "num_inference_steps": len(self.fixed_sigmas), - "guidance_scale": DISTILLED_GUIDANCE_SCALE, - } - def num_steps(self, default: int) -> int: """The only step count this policy can run: fixed for distilled, else ``default``.""" return len(self.fixed_sigmas) if self.is_distilled else default diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py index ceab1a72eb9f..c1f9e0e66502 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -17,7 +17,7 @@ from contextlib import nullcontext from dataclasses import dataclass from types import SimpleNamespace -from typing import Any, Callable, ContextManager, Optional, Tuple, TypeVar +from typing import Any, Callable, ContextManager, Optional, Tuple import torch import torch.nn as nn @@ -43,21 +43,7 @@ from tensorrt_llm.logger import logger from tensorrt_llm.models.modeling_utils import QuantConfig -from .config_resolver import architecture_defaults, fill_pretrained_config - -# Some Cosmos3OmniTransformer checkpoint configs omit these fields; the values come -# from the precedence table in config_resolver, the one place that declares them. -PRETRAINED_CONFIG_COMPAT_DEFAULTS = architecture_defaults() - - -_PretrainedConfigT = TypeVar("_PretrainedConfigT") - - -def apply_pretrained_config_compat_defaults( - pretrained_config: _PretrainedConfigT, -) -> _PretrainedConfigT: - """Fill missing schema fields in place (idempotent); returns the config.""" - return fill_pretrained_config(pretrained_config) +from .config_resolver import fill_pretrained_config def uses_static_fp8(model_config: DiffusionModelConfig) -> bool: @@ -1205,10 +1191,8 @@ def forward( class Cosmos3VFMTransformer(BaseDiffusionModel): def __init__(self, model_config: DiffusionModelConfig): super().__init__(model_config) - self.temporal_compression_factor_declared = ( - getattr(model_config.pretrained_config, "temporal_compression_factor", None) is not None - ) - pretrained_config = apply_pretrained_config_compat_defaults(model_config.pretrained_config) + # Some exports omit a few schema fields; the checkpoint resolver's table fills them. + pretrained_config = fill_pretrained_config(model_config.pretrained_config) self.recipe = resolve_arch_recipe(pretrained_config) # Installed in post_load_weights when the checkpoint is static FP8. self.step_precision_controller: Optional[StepPrecisionController] = None diff --git a/tests/unittest/_torch/visual_gen/cosmos3_request_golden.py b/tests/unittest/_torch/visual_gen/cosmos3_request_golden.py index 6029ff9c7faa..f0e189304000 100644 --- a/tests/unittest/_torch/visual_gen/cosmos3_request_golden.py +++ b/tests/unittest/_torch/visual_gen/cosmos3_request_golden.py @@ -15,10 +15,11 @@ """The request-level defaults every Cosmos3 family, mode and domain resolves to when the caller sets nothing. -``record()`` evaluates the grid with the pipeline's current resolution helpers; +``record()`` evaluates the grid through the request resolver; ``python cosmos3_request_golden.py > test_data/cosmos3_request_defaults.json`` stores -it. ``test_cosmos3_request_resolver.py`` checks the request resolver against the -stored values, so a refactor of the defaults is a visible diff of that file. +it. ``test_cosmos3_request_resolver.py`` compares the live grid with the stored one, +so a change of a table, a rule or a precedence is a visible diff of that file that +has to be regenerated on purpose. """ import json @@ -30,12 +31,14 @@ from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( COSMOS3_DOMAIN_PRESETS, resolve_checkpoint_policy_defaults, - resolve_domain_action_config, ) from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import ( Cosmos3OmniMoTPipeline, _resolve_use_system_prompt, ) +from tensorrt_llm._torch.visual_gen.models.cosmos3.request_resolver import ( + resolve_domain_action_config, +) from tensorrt_llm._torch.visual_gen.models.cosmos3.sampling import Cosmos3SamplingPolicy FAMILIES = ("qwen3", "nemotron_dense") @@ -124,15 +127,13 @@ def record() -> dict: policy_defaults = resolve_checkpoint_policy_defaults(block) for domain in domains: try: - cfg = resolve_domain_action_config( + resolved = resolve_domain_action_config( domain_name=domain, checkpoint_policy_defaults=policy_defaults ) - cfg = {k: _jsonable(v) for k, v in cfg.items() if k not in ("warnings", "sources")} - cfg["num_warnings"] = len( - resolve_domain_action_config( - domain_name=domain, checkpoint_policy_defaults=policy_defaults - )["warnings"] - ) + cfg = { + k: _jsonable(v) for k, v in resolved.items() if k not in ("warnings", "sources") + } + cfg["num_warnings"] = len(resolved["warnings"]) except ValueError as exc: cfg = {"error": str(exc)} out["action"][f"{domain}{policy_label}"] = cfg diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py index de8d4fdbc0ae..82fe1b83cebf 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py @@ -26,16 +26,20 @@ resolve_raw_action_dim, ) from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( - COSMOS3_ASPECT_RATIO_BUCKETS, COSMOS3_DOMAIN_PRESET_ALIASES, COSMOS3_DOMAIN_PRESETS, COSMOS3_EXTRA_SPECS, COSMOS3_POLICY_SAMPLING_PARAMS, + get_domain_preset, + resolve_checkpoint_policy_defaults, +) +from tensorrt_llm._torch.visual_gen.models.cosmos3.geometry import ( + COSMOS3_ASPECT_RATIO_BUCKETS, VIDEO_RES_SIZE_INFO, find_closest_target_size, - get_domain_preset, normalize_action_resolution, - resolve_checkpoint_policy_defaults, +) +from tensorrt_llm._torch.visual_gen.models.cosmos3.request_resolver import ( resolve_domain_action_config, ) from tensorrt_llm._torch.visual_gen.models.cosmos3.transformer_cosmos3 import ( diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_config_resolver.py b/tests/unittest/_torch/visual_gen/test_cosmos3_config_resolver.py index ea453e5bd50a..b1e8e6393670 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_config_resolver.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_config_resolver.py @@ -32,8 +32,6 @@ Default, Key, Rule, - Serving, - architecture_defaults, fill_pretrained_config, ) @@ -51,9 +49,9 @@ def _write(checkpoint_dir: Path, files: dict) -> Path: return checkpoint_dir -def _resolver(tmp_path: Path, files: dict, serving=None) -> Cosmos3ConfigResolver: +def _resolver(tmp_path: Path, files: dict) -> Cosmos3ConfigResolver: _write(tmp_path, files) - return Cosmos3ConfigResolver(str(tmp_path), files.get(TRANSFORMER, {}), serving) + return Cosmos3ConfigResolver(str(tmp_path), files.get(TRANSFORMER, {})) class TestRule: @@ -68,13 +66,6 @@ def test_falls_through_missing_and_null_sources(self, tmp_path): rule = Rule((Key(CONFIG, "a"), Key(MODEL_INDEX, "a"), Key(CHECKPOINT, "a"))) assert r.resolve_one("a", rule).value == 2 - def test_nullable_keeps_an_explicit_null(self, tmp_path): - r = _resolver(tmp_path, {CONFIG: {"a": None}, MODEL_INDEX: {"a": 2}}) - rule = Rule( - (Key(CONFIG, "a"), Key(MODEL_INDEX, "a")), conflict=Conflict.IGNORE, nullable=True - ) - assert r.resolve_one("a", rule).value is None - def test_default_only_when_nothing_is_declared(self, tmp_path): r = _resolver(tmp_path, {}) rule = Rule((Key(CONFIG, "a"), Default(7, "x"))) @@ -102,12 +93,6 @@ def test_float_lists_compare_within_tolerance(self, tmp_path): rule = Rule((Key(CONFIG, "a"), Key(MODEL_INDEX, "a")), conflict=Conflict.ERROR) assert r.resolve_one("a", rule).disagreeing == () - def test_serving_config_is_a_source(self, tmp_path): - serving = SimpleNamespace(sampling=SimpleNamespace(flow_shift=0.5)) - r = _resolver(tmp_path, {}, serving) - rule = Rule((Serving("sampling.flow_shift"),)) - assert r.resolve_one("f", rule).value == 0.5 - def _files(): """A distilled checkpoint whose training config and diffusers index disagree on @@ -175,7 +160,7 @@ def test_fill_pretrained_config_sets_only_missing_attributes(): declared = SimpleNamespace(temporal_compression_factor=8) fill_pretrained_config(declared) assert declared.temporal_compression_factor == 8 - assert declared.position_embedding_type == architecture_defaults()["position_embedding_type"] + assert declared.position_embedding_type == "unified_3d_mrope" assert declared.max_position_embeddings == 262144 assert declared.temporal_compression_factor_sound == 1 diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py index 2e4c381cb7b0..d26a67a3126d 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py @@ -193,13 +193,6 @@ def test_sigma_values_coerced_to_floats(self): scheduler = FlowMatchEulerDiscreteScheduler.from_config(config) assert Cosmos3SamplingPolicy.from_scheduler(scheduler).fixed_sigmas == (1.0, 0.5) - def test_generation_default_overrides(self): - assert _distilled_policy().generation_default_overrides() == { - "num_inference_steps": 4, - "guidance_scale": DISTILLED_GUIDANCE_SCALE, - } - assert _base_policy().generation_default_overrides() == {} - def test_num_steps(self): assert _distilled_policy().num_steps(2) == 4 assert _base_policy().num_steps(2) == 2 @@ -1409,8 +1402,8 @@ def test_unset_resolves_to_checkpoint_not_base_table(self, monkeypatch): assert f"guidance_scale={COSMOS3_720P_PARAMS['guidance_scale']:.2f}" not in dims def test_base_checkpoint_still_uses_the_mode_table(self, monkeypatch): - """The checkpoint-first step must be a no-op for a base checkpoint: - ``generation_default_overrides()`` is empty, so the mode tables win.""" + """The checkpoint-first step must be a no-op for a base checkpoint: there is no + fixed schedule, so the mode tables win.""" import tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 as mod pipeline = _bare_pipeline() diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py b/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py index 901628f8fedf..8f7b9259aed2 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py @@ -381,7 +381,6 @@ def test_recipe_and_flags(self, model): assert model.audio_gen is False assert model.has_action_weights is True assert model.temporal_compression_factor == 4 - assert model.temporal_compression_factor_declared is True def test_und_attention_norms(self, model): attn = model.language_model.layers[0].self_attn @@ -855,11 +854,6 @@ def test_policy_manifest_overrides_policy_sampling_only(self): } ) - policy = pipeline._mode_params("video", action_mode=ACTION_MODE_POLICY) - - assert policy["num_inference_steps"] == 4 - assert policy["guidance_scale"] == 3.0 - assert policy["guidance_interval"] == COSMOS3_POLICY_SAMPLING_PARAMS["guidance_interval"] resolved_policy = pipeline._resolve_generation_params( "video", action_mode=ACTION_MODE_POLICY, @@ -872,14 +866,8 @@ def test_policy_manifest_overrides_policy_sampling_only(self): "guidance_scale": 3.0, "guidance_interval": COSMOS3_POLICY_SAMPLING_PARAMS["guidance_interval"], } - assert ( - pipeline._mode_params("video", action_mode=ACTION_MODE_FORWARD_DYNAMICS) - is COSMOS3_ACTION_PARAMS - ) - assert ( - pipeline._mode_params("video", action_mode=ACTION_MODE_INVERSE_DYNAMICS) - is COSMOS3_ACTION_PARAMS - ) + assert pipeline._mode_table("video", ACTION_MODE_FORWARD_DYNAMICS) is COSMOS3_ACTION_PARAMS + assert pipeline._mode_table("video", ACTION_MODE_INVERSE_DYNAMICS) is COSMOS3_ACTION_PARAMS for action_mode in (ACTION_MODE_FORWARD_DYNAMICS, ACTION_MODE_INVERSE_DYNAMICS): resolved = pipeline._resolve_generation_params( "video", @@ -893,7 +881,7 @@ def test_policy_manifest_overrides_policy_sampling_only(self): "guidance_scale": 1.0, "guidance_interval": None, } - assert pipeline._mode_params("video") is COSMOS3_EDGE_VIDEO_PARAMS + assert pipeline._mode_table("video") is COSMOS3_EDGE_VIDEO_PARAMS @pytest.mark.parametrize( "action_mode", @@ -1056,23 +1044,6 @@ def fake_forward(**kwargs): assert captured["guidance_scale"] == 5.0 assert captured["max_sequence_length"] == 4096 - def test_declared_temporal_factor_must_match_vae(self): - from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import ( - _validate_temporal_compression, - ) - - declared = SimpleNamespace( - temporal_compression_factor=4, temporal_compression_factor_declared=True - ) - _validate_temporal_compression(declared, 4) - with pytest.raises(ValueError, match="temporal_compression_factor"): - _validate_temporal_compression(declared, 8) - # Nano-style configs don't declare it; the VAE value simply wins. - undeclared = SimpleNamespace( - temporal_compression_factor=4, temporal_compression_factor_declared=False - ) - _validate_temporal_compression(undeclared, 8) - class TestSamplingRecipeMatrix: """Family + model_index schedule flag + scheduler recipe must be validated @@ -1884,7 +1855,7 @@ def test_recipe_and_scheduler_wiring(self, edge_pipeline): # the two modes resolve to the same (shift, karras) pair and share one # cached scheduler rather than building two identical ones. schedulers = { - mode: edge_pipeline._scheduler_for(edge_pipeline._mode_params(mode)["flow_shift"]) + mode: edge_pipeline._scheduler_for(edge_pipeline._mode_table(mode)["flow_shift"]) for mode in ("video", "image") } for scheduler in schedulers.values(): diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_example_prompts.py b/tests/unittest/_torch/visual_gen/test_cosmos3_example_prompts.py index 4aa1f47c00dd..6f4951b87155 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_example_prompts.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_example_prompts.py @@ -16,9 +16,13 @@ import pytest from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( + resolve_checkpoint_policy_defaults, +) +from tensorrt_llm._torch.visual_gen.models.cosmos3.geometry import ( VIDEO_RES_SIZE_INFO, find_closest_target_size, - resolve_checkpoint_policy_defaults, +) +from tensorrt_llm._torch.visual_gen.models.cosmos3.request_resolver import ( resolve_domain_action_config, ) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index 2cece22b55b0..ba653f97b4f4 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -928,9 +928,6 @@ class FakeSampling: def validate_request(self, num_inference_steps, guidance_scale): return None - def generation_default_overrides(self): - return {} - def set_flow_shift(self, scheduler, target, *, use_karras_sigmas=None): calls.append((target, use_karras_sigmas)) return scheduler @@ -990,9 +987,6 @@ class FakeSampling: def validate_request(self, num_inference_steps, guidance_scale): return None - def generation_default_overrides(self): - return {} - def set_flow_shift(self, scheduler, target, *, use_karras_sigmas=None): rebuilt.append((scheduler.name, target, use_karras_sigmas)) return scheduler @@ -1054,9 +1048,6 @@ class FakeSampling: def validate_request(self, num_inference_steps, guidance_scale): return None - def generation_default_overrides(self): - return {} - def set_flow_shift(self, scheduler, target, *, use_karras_sigmas=None): return scheduler @@ -1112,9 +1103,6 @@ class FakeSampling: def validate_request(self, num_inference_steps, guidance_scale): return None - def generation_default_overrides(self): - return {} - def set_flow_shift(self, scheduler, target, *, use_karras_sigmas=None): calls.append((scheduler.name, target, use_karras_sigmas)) return scheduler @@ -1219,9 +1207,6 @@ class FakeSampling: def validate_request(self, num_inference_steps, guidance_scale): return None - def generation_default_overrides(self): - return {} - pipeline.sampling = FakeSampling() pipeline._forward_transfer = lambda **kwargs: None # A precomputed control and no video: an existing guard already rejects @@ -1274,9 +1259,6 @@ class FakeSampling: def validate_request(self, num_inference_steps, guidance_scale): return None - def generation_default_overrides(self): - return {} - pipeline.sampling = FakeSampling() captured = {} diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_request_resolver.py b/tests/unittest/_torch/visual_gen/test_cosmos3_request_resolver.py index ca6779c1b9d8..25a461be9202 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_request_resolver.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_request_resolver.py @@ -21,10 +21,10 @@ import pytest from cosmos3_request_golden import record -from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import resolve_domain_action_config from tensorrt_llm._torch.visual_gen.models.cosmos3.request_resolver import ( RequestContext, resolve_action, + resolve_domain_action_config, resolve_generation, ) @@ -34,8 +34,9 @@ def test_defaults_match_the_recorded_grid(): - """Every (family, mode, domain) combination resolves to the values recorded before - the rules existed; a change here is a visible change of a default.""" + """Every (family, mode, domain) combination resolves to the recorded values, so a + change of any default, or of the precedence behind it, is a visible diff of the + fixture rather than a silent change of what requests run with.""" recorded = json.loads(FIXTURE.read_text()) current = json.loads(json.dumps(record())) for section in ("generation", "action", "system_prompt"): @@ -76,16 +77,18 @@ def test_caller_value_wins_and_the_recommendation_is_reported(self): "Cosmos3 action_chunk_size=16 differs from recommended 60 for domain 'av'." ] - def test_checkpoint_embodiments_replace_the_code_table(self): - embodiments = {"agibotworld": {"domain_id": 15, "input_action_dim": 59}} - cfg = resolve_domain_action_config(domain_name="agibotworld", embodiments=embodiments) - assert cfg["raw_action_dim"] == 59 - assert cfg["sources"]["raw_action_dim"] == "checkpoint:conditioning.embodiments.action_dim" - assert cfg["warnings"] == [] - # The same request against a checkpoint without a table: the code's width. - plain = resolve_domain_action_config(domain_name="agibotworld") - assert plain["raw_action_dim"] == 29 - assert plain["sources"]["raw_action_dim"] == "embodiment table (defaults)" + def test_explicit_none_in_a_mode_table_is_a_value(self): + """Edge text-to-image switches the guidance window off with an explicit ``None``; + that must not fall through to the video table.""" + ctx = RequestContext(family="nemotron_dense", mode="image") + r = resolve_generation(ctx, guidance_interval=None) + assert r["guidance_interval"] is None + assert r.source_of("guidance_interval") == "mode table.guidance_interval" + + def test_disagreeing_recommendation_names_its_source(self): + cfg = resolve_domain_action_config(domain_name="av", action_chunk_size=16) + assert cfg["action_chunk_size"] == 16 + assert cfg["sources"]["action_chunk_size"] == "request" def test_unresolvable_field_is_an_error(self): ctx = RequestContext(family="no-such-family") diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py b/tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py index ef3014b4ed86..50f256cec23a 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py @@ -30,11 +30,12 @@ COSMOS3_720P_PARAMS, COSMOS3_EXTRA_SPECS, COSMOS3_GENERATION_DEFAULTS, +) +from tensorrt_llm._torch.visual_gen.models.cosmos3.geometry import ( VIDEO_RES_SIZE_INFO, find_closest_target_size, ) from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniMoTPipeline -from tensorrt_llm._torch.visual_gen.models.cosmos3.sampling import DISTILLED_GUIDANCE_SCALE from tensorrt_llm._torch.visual_gen.models.cosmos3.transfer import ( BILATERAL_D, BILATERAL_SIGMA_COLOR, @@ -230,16 +231,6 @@ def scheduler_step_kwargs(self, generator): self.step_kwargs_calls += 1 return {"generator": generator} if self.is_distilled else {} - def generation_default_overrides(self): - # Mirrors the real policy rather than returning {}, so a distilled stub - # still overrides the table the way the checkpoint would. - if not self.is_distilled: - return {} - return { - "num_inference_steps": len(self.fixed_sigmas), - "guidance_scale": DISTILLED_GUIDANCE_SCALE, - } - def num_steps(self, default): return len(self.fixed_sigmas) if self.is_distilled else default diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py index 7eed96138f04..53e41a5e58bd 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py @@ -25,11 +25,10 @@ from tensorrt_llm._torch.modules.linear import Linear from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig, DiffusionPipelineConfig +from tensorrt_llm._torch.visual_gen.models.cosmos3.config_resolver import fill_pretrained_config from tensorrt_llm._torch.visual_gen.models.cosmos3.transformer_cosmos3 import ( - PRETRAINED_CONFIG_COMPAT_DEFAULTS, Cosmos3VFMTransformer, Qwen3VLTextRotaryEmbedding, - apply_pretrained_config_compat_defaults, ) from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineComponent, PipelineLoader from tensorrt_llm.visual_gen.args import TorchCompileConfig, VisualGenArgs @@ -773,9 +772,11 @@ class TestConfigCompatDefaults: def test_new_schema_gets_defaults(self): config = SimpleNamespace(hidden_size=64, rope_axes_dim=[4, 2, 2]) - apply_pretrained_config_compat_defaults(config) - for key, value in PRETRAINED_CONFIG_COMPAT_DEFAULTS.items(): - assert getattr(config, key) == value + fill_pretrained_config(config) + assert config.position_embedding_type == "unified_3d_mrope" + assert config.max_position_embeddings == 262144 + assert config.temporal_compression_factor == 4 + assert config.temporal_compression_factor_sound == 1 def test_old_schema_untouched(self): # Every field deliberately differs from its compat default, so an @@ -785,16 +786,16 @@ def test_old_schema_untouched(self): max_position_embeddings=12345, temporal_compression_factor_sound=7, ) - apply_pretrained_config_compat_defaults(config) + fill_pretrained_config(config) assert config.position_embedding_type == "rope_3d" assert config.max_position_embeddings == 12345 assert config.temporal_compression_factor_sound == 7 def test_idempotent(self): config = SimpleNamespace(hidden_size=64) - apply_pretrained_config_compat_defaults(config) + fill_pretrained_config(config) snapshot = vars(config).copy() - apply_pretrained_config_compat_defaults(config) + fill_pretrained_config(config) assert vars(config) == snapshot @pytest.mark.parametrize("rope_scaling", [_OMITTED, None, {}], ids=["omitted", "none", "empty"]) @@ -818,7 +819,7 @@ def test_rope_type_tolerates_missing_rope_scaling(self, rope_scaling: object) -> if rope_scaling is not _OMITTED: config.rope_scaling = rope_scaling assert hasattr(config, "rope_scaling") is (rope_scaling is not _OMITTED) - apply_pretrained_config_compat_defaults(config) + fill_pretrained_config(config) # Construct for real: the point is that __init__ reaches the resolver # instead of raising AttributeError on the missing block. embedding = tf.Qwen3VLTextRotaryEmbedding(SimpleNamespace(pretrained_config=config)) From e27621ebcbe6ea53bc5aec83b627edc1914a1101 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Sun, 1 Nov 2026 12:20:52 -0800 Subject: [PATCH 4/8] [None][refactor] Cosmos3: keep the embodiment and geometry tables in defaults.py Signed-off-by: Igor Shovkun --- .../visual_gen/models/cosmos3/action.py | 74 +------- .../visual_gen/models/cosmos3/defaults.py | 162 +++++++++++++++++- .../visual_gen/models/cosmos3/geometry.py | 101 ----------- .../models/cosmos3/pipeline_cosmos3.py | 3 +- .../models/cosmos3/request_resolver.py | 5 +- .../visual_gen/cosmos3_request_golden.py | 2 +- .../_torch/visual_gen/test_cosmos3_action.py | 14 +- .../test_cosmos3_example_prompts.py | 4 +- .../visual_gen/test_cosmos3_transfer.py | 2 - 9 files changed, 169 insertions(+), 198 deletions(-) delete mode 100644 tensorrt_llm/_torch/visual_gen/models/cosmos3/geometry.py diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py index 7172c01ac67e..adb0fa50e8c0 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py @@ -13,7 +13,7 @@ import torch from diffusers.utils.torch_utils import randn_tensor -from .geometry import find_closest_target_size +from .defaults import EMBODIMENT_TO_DOMAIN_ID, find_closest_target_size ACTION_MODE_POLICY = "policy" ACTION_MODE_FORWARD_DYNAMICS = "forward_dynamics" @@ -24,78 +24,6 @@ ACTION_MODE_INVERSE_DYNAMICS, } -EMBODIMENT_TO_DOMAIN_ID: dict[str, int] = { - "no_action": 0, - "av": 1, - "camera_pose": 2, - "hand_pose": 3, - "pusht": 4, - "libero": 5, - "umi": 6, - "bridge_orig_lerobot": 7, - "droid_lerobot": 8, - "robomind-franka": 8, - "galbot": 9, - "robomind-franka-dual": 12, - "robomind-ur": 13, - "agibotworld": 15, - "agibot_gear_gripper": 15, - "agibot_gear_gripper_ext": 15, - "fractal": 20, -} - -# Canonical unpadded action width per embodiment. Widths compose the Cosmos3 -# unified action representation from shared geometric blocks: a 9-D pose (3-D -# translation + 6-D rotation), a 1-D grasp state, and a 15-D fingertip state. -# One arm is 9 + 1 = 10; a dual-arm setup is 20; the AgiBot humanoid is -# 9 + 2 x (9 + 1) = 29; two-hand egocentric motion is 9 + 2 x (9 + 15) = 57. -# -# This is a property of the embodiment, not a tunable, so it is keyed by the -# real domain name rather than by the sampling presets in ``defaults.py`` (where -# several of these names share one preset). ``libero`` is absent on purpose: -# its width depends on the dataset's rotation space (7/10/13), so callers must -# pass ``raw_action_dim`` explicitly. -EMBODIMENT_TO_RAW_ACTION_DIM: dict[str, int] = { - "av": 9, - "camera_pose": 9, - "hand_pose": 57, - "pusht": 2, - "umi": 10, - "bridge_orig_lerobot": 10, - "droid_lerobot": 10, - "robomind-franka": 10, - "robomind-franka-dual": 20, - "robomind-ur": 10, - "galbot": 30, - "agibotworld": 29, - "agibot_gear_gripper": 29, - "agibot_gear_gripper_ext": 29, - "fractal": 10, -} - - -def resolve_raw_action_dim( - domain_name: Any = None, - domain_id: Any = None, -) -> int | None: - """Look up the canonical action width, or None when it cannot be determined. - - Resolves by name first. A bare ``domain_id`` is only usable when every - embodiment sharing that id agrees on a width (true for all current ids). - """ - if domain_name is not None and str(domain_name).strip(): - return EMBODIMENT_TO_RAW_ACTION_DIM.get(str(domain_name).strip().lower()) - - if domain_id is None: - return None - - widths = { - EMBODIMENT_TO_RAW_ACTION_DIM[name] - for name, mapped_id in EMBODIMENT_TO_DOMAIN_ID.items() - if mapped_id == int(domain_id) and name in EMBODIMENT_TO_RAW_ACTION_DIM - } - return widths.pop() if len(widths) == 1 else None - def normalize_action_mode(mode: Any) -> str | None: if mode is None: diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index 4761a02ca29c..29ff55588973 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -20,12 +20,6 @@ from collections.abc import Mapping from typing import Any, Dict, Iterable, TypedDict -from tensorrt_llm._torch.visual_gen.models.cosmos3.action import EMBODIMENT_TO_DOMAIN_ID -from tensorrt_llm._torch.visual_gen.models.cosmos3.geometry import ( - COSMOS3_ACTION_RESOLUTIONS, - VIDEO_RES_SIZE_INFO, - normalize_action_resolution, -) from tensorrt_llm._torch.visual_gen.pipeline import ExtraParamSchema from tensorrt_llm.inputs.media_io import sniff_media_kind @@ -33,6 +27,160 @@ # Constant tables # --------------------------------------------------------------------------- +# The reference's ``W,H`` aspect buckets: every canvas is one of these shapes, and a +# request's bucket string is carried into the caption verbatim, so a resolved frame +# size maps back to the nearest bucket rather than to its exact reduced ratio +# (832x480 reduces to "15,26", which the model never saw). +COSMOS3_ASPECT_RATIO_BUCKETS = ("1,1", "4,3", "3,4", "16,9", "9,16") + + +# Output sizes per resolution level and aspect bucket, as (width, height). +VIDEO_RES_SIZE_INFO: dict[str, dict[str, tuple[int, int]]] = { + "256": { + "1,1": (256, 256), + "4,3": (320, 256), + "3,4": (256, 320), + "16,9": (320, 192), + "9,16": (192, 320), + }, + "480": { + "1,1": (640, 640), + "4,3": (736, 544), + "3,4": (544, 736), + "16,9": (832, 480), + "9,16": (480, 832), + }, + "704": { + "1,1": (960, 960), + "4,3": (1088, 832), + "3,4": (832, 1088), + "16,9": (1280, 704), + "9,16": (704, 1280), + }, + "720": { + "1,1": (960, 960), + "4,3": (1104, 832), + "3,4": (832, 1104), + "16,9": (1280, 720), + "9,16": (720, 1280), + }, +} + + +COSMOS3_ACTION_RESOLUTIONS = tuple(int(key) for key in sorted(VIDEO_RES_SIZE_INFO, key=int)) + + +def normalize_action_resolution(resolution: Any) -> int: + if resolution is None: + raise ValueError("Cosmos3 action_resolution is required for action generation.") + try: + bucket = int(resolution) + except (TypeError, ValueError) as exc: + raise ValueError( + f"Cosmos3 action_resolution must be an int bucket, got {resolution!r}." + ) from exc + if bucket not in COSMOS3_ACTION_RESOLUTIONS: + raise ValueError( + f"Unknown Cosmos3 action_resolution={bucket}; " + f"expected one of {COSMOS3_ACTION_RESOLUTIONS}." + ) + return bucket + + +def find_closest_target_size(h: int, w: int, resolution: str | int) -> tuple[int, int]: + """The ``resolution`` bucket whose aspect ratio is closest to ``h/w``, as + ``(width, height)``, so a source frame lands on a supported size undistorted.""" + key = str(resolution) + if key not in VIDEO_RES_SIZE_INFO: + raise ValueError( + f"Unknown Cosmos3 resolution={resolution!r}; " + f"expected one of {sorted(VIDEO_RES_SIZE_INFO)}." + ) + input_ratio = h / w + best_size = None + best_diff = float("inf") + for cand_w, cand_h in VIDEO_RES_SIZE_INFO[key].values(): + diff = abs(input_ratio - cand_h / cand_w) + if diff < best_diff: + best_diff = diff + best_size = (cand_w, cand_h) + assert best_size is not None + return best_size + + +EMBODIMENT_TO_DOMAIN_ID: dict[str, int] = { + "no_action": 0, + "av": 1, + "camera_pose": 2, + "hand_pose": 3, + "pusht": 4, + "libero": 5, + "umi": 6, + "bridge_orig_lerobot": 7, + "droid_lerobot": 8, + "robomind-franka": 8, + "galbot": 9, + "robomind-franka-dual": 12, + "robomind-ur": 13, + "agibotworld": 15, + "agibot_gear_gripper": 15, + "agibot_gear_gripper_ext": 15, + "fractal": 20, +} + +# Canonical unpadded action width per embodiment. Widths compose the Cosmos3 +# unified action representation from shared geometric blocks: a 9-D pose (3-D +# translation + 6-D rotation), a 1-D grasp state, and a 15-D fingertip state. +# One arm is 9 + 1 = 10; a dual-arm setup is 20; the AgiBot humanoid is +# 9 + 2 x (9 + 1) = 29; two-hand egocentric motion is 9 + 2 x (9 + 15) = 57. +# +# This is a property of the embodiment, not a tunable, so it is keyed by the +# real domain name rather than by the sampling presets below (where several of +# these names share one preset). ``libero`` is absent on purpose: +# its width depends on the dataset's rotation space (7/10/13), so callers must +# pass ``raw_action_dim`` explicitly. +EMBODIMENT_TO_RAW_ACTION_DIM: dict[str, int] = { + "av": 9, + "camera_pose": 9, + "hand_pose": 57, + "pusht": 2, + "umi": 10, + "bridge_orig_lerobot": 10, + "droid_lerobot": 10, + "robomind-franka": 10, + "robomind-franka-dual": 20, + "robomind-ur": 10, + "galbot": 30, + "agibotworld": 29, + "agibot_gear_gripper": 29, + "agibot_gear_gripper_ext": 29, + "fractal": 10, +} + + +def resolve_raw_action_dim( + domain_name: Any = None, + domain_id: Any = None, +) -> int | None: + """Look up the canonical action width, or None when it cannot be determined. + + Resolves by name first. A bare ``domain_id`` is only usable when every + embodiment sharing that id agrees on a width (true for all current ids). + """ + if domain_name is not None and str(domain_name).strip(): + return EMBODIMENT_TO_RAW_ACTION_DIM.get(str(domain_name).strip().lower()) + + if domain_id is None: + return None + + widths = { + EMBODIMENT_TO_RAW_ACTION_DIM[name] + for name, mapped_id in EMBODIMENT_TO_DOMAIN_ID.items() + if mapped_id == int(domain_id) and name in EMBODIMENT_TO_RAW_ACTION_DIM + } + return widths.pop() if len(widths) == 1 else None + + # The default video resolution is the 720p 16:9 bucket, ``(width, height)``. _DEFAULT_VIDEO_W, _DEFAULT_VIDEO_H = VIDEO_RES_SIZE_INFO["720"]["16,9"] COSMOS3_720P_PARAMS = { @@ -385,7 +533,7 @@ class Cosmos3DomainPreset(TypedDict, total=False): """Recommended action sampling settings for a trained embodiment. Sampling settings only — the embodiment's action width lives in - ``action.EMBODIMENT_TO_RAW_ACTION_DIM``, keyed by the unaliased domain name. + ``EMBODIMENT_TO_RAW_ACTION_DIM``, keyed by the unaliased domain name. """ action_chunk_size: int diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/geometry.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/geometry.py deleted file mode 100644 index 33ee45c1a444..000000000000 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/geometry.py +++ /dev/null @@ -1,101 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""The Cosmos3 output geometry: the aspect buckets, the output size per resolution -level and bucket, and the lookups every mode shares. A leaf module: nothing here -depends on the rest of the package.""" - -from __future__ import annotations - -from typing import Any - -# The reference's ``W,H`` aspect buckets: every canvas is one of these shapes, and a -# request's bucket string is carried into the caption verbatim, so a resolved frame -# size maps back to the nearest bucket rather than to its exact reduced ratio -# (832x480 reduces to "15,26", which the model never saw). -COSMOS3_ASPECT_RATIO_BUCKETS = ("1,1", "4,3", "3,4", "16,9", "9,16") - - -# Output sizes per resolution level and aspect bucket, as (width, height). -VIDEO_RES_SIZE_INFO: dict[str, dict[str, tuple[int, int]]] = { - "256": { - "1,1": (256, 256), - "4,3": (320, 256), - "3,4": (256, 320), - "16,9": (320, 192), - "9,16": (192, 320), - }, - "480": { - "1,1": (640, 640), - "4,3": (736, 544), - "3,4": (544, 736), - "16,9": (832, 480), - "9,16": (480, 832), - }, - "704": { - "1,1": (960, 960), - "4,3": (1088, 832), - "3,4": (832, 1088), - "16,9": (1280, 704), - "9,16": (704, 1280), - }, - "720": { - "1,1": (960, 960), - "4,3": (1104, 832), - "3,4": (832, 1104), - "16,9": (1280, 720), - "9,16": (720, 1280), - }, -} - - -COSMOS3_ACTION_RESOLUTIONS = tuple(int(key) for key in sorted(VIDEO_RES_SIZE_INFO, key=int)) - - -def normalize_action_resolution(resolution: Any) -> int: - if resolution is None: - raise ValueError("Cosmos3 action_resolution is required for action generation.") - try: - bucket = int(resolution) - except (TypeError, ValueError) as exc: - raise ValueError( - f"Cosmos3 action_resolution must be an int bucket, got {resolution!r}." - ) from exc - if bucket not in COSMOS3_ACTION_RESOLUTIONS: - raise ValueError( - f"Unknown Cosmos3 action_resolution={bucket}; " - f"expected one of {COSMOS3_ACTION_RESOLUTIONS}." - ) - return bucket - - -def find_closest_target_size(h: int, w: int, resolution: str | int) -> tuple[int, int]: - """The ``resolution`` bucket whose aspect ratio is closest to ``h/w``, as - ``(width, height)``, so a source frame lands on a supported size undistorted.""" - key = str(resolution) - if key not in VIDEO_RES_SIZE_INFO: - raise ValueError( - f"Unknown Cosmos3 resolution={resolution!r}; " - f"expected one of {sorted(VIDEO_RES_SIZE_INFO)}." - ) - input_ratio = h / w - best_size = None - best_diff = float("inf") - for cand_w, cand_h in VIDEO_RES_SIZE_INFO[key].values(): - diff = abs(input_ratio - cand_h / cand_w) - if diff < best_diff: - best_diff = diff - best_size = (cand_w, cand_h) - assert best_size is not None - return best_size diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index b42eac2af35f..c62b48dd16d6 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -71,6 +71,7 @@ def tqdm(iterable, **kwargs): from .config_resolver import Cosmos3ConfigResolver from .defaults import ( COSMOS3_720P_PARAMS, + COSMOS3_ASPECT_RATIO_BUCKETS, COSMOS3_DEFAULT_RESOLUTION_TEMPLATE, COSMOS3_DEFAULT_SYSTEM_PROMPT, COSMOS3_DURATION_TEMPLATE, @@ -82,9 +83,9 @@ def tqdm(iterable, **kwargs): COSMOS3_V2V_FLOW_SHIFT, _normalize_condition_video_keep, _normalize_condition_video_latent_indexes, + find_closest_target_size, resolve_checkpoint_policy_defaults, ) -from .geometry import COSMOS3_ASPECT_RATIO_BUCKETS, find_closest_target_size from .guardrails import check_video_safety, download_guardrail_checkpoint from .negative_prompt import COSMOS3_VIDEO_NEGATIVE_PROMPT from .request_resolver import RequestContext, resolve_domain_action_config, resolve_generation diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/request_resolver.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/request_resolver.py index 960fd0f9410d..f274963144b8 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/request_resolver.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/request_resolver.py @@ -29,15 +29,16 @@ from dataclasses import dataclass, field from typing import Any, Optional -from .action import EMBODIMENT_TO_DOMAIN_ID, resolve_raw_action_dim from .config_resolver import Conflict, Default, Resolved, same_value from .defaults import ( COSMOS3_ACTION_PARAMS, COSMOS3_DOMAIN_PRESETS, COSMOS3_GENERATION_DEFAULTS, + EMBODIMENT_TO_DOMAIN_ID, canonical_domain_preset_key, + normalize_action_resolution, + resolve_raw_action_dim, ) -from .geometry import normalize_action_resolution from .sampling import DISTILLED_GUIDANCE_SCALE _MISSING = object() diff --git a/tests/unittest/_torch/visual_gen/cosmos3_request_golden.py b/tests/unittest/_torch/visual_gen/cosmos3_request_golden.py index f0e189304000..367f2a625991 100644 --- a/tests/unittest/_torch/visual_gen/cosmos3_request_golden.py +++ b/tests/unittest/_torch/visual_gen/cosmos3_request_golden.py @@ -27,9 +27,9 @@ from diffusers import FlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler -from tensorrt_llm._torch.visual_gen.models.cosmos3.action import EMBODIMENT_TO_DOMAIN_ID from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( COSMOS3_DOMAIN_PRESETS, + EMBODIMENT_TO_DOMAIN_ID, resolve_checkpoint_policy_defaults, ) from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import ( diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py index 82fe1b83cebf..141e9f7ce126 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py @@ -13,8 +13,6 @@ import torch from tensorrt_llm._torch.visual_gen.models.cosmos3.action import ( - EMBODIMENT_TO_DOMAIN_ID, - EMBODIMENT_TO_RAW_ACTION_DIM, action_reference_frame_step, action_reference_size, crop_action_latent, @@ -23,21 +21,21 @@ resolve_action_content_size, resolve_action_size, resolve_domain_id, - resolve_raw_action_dim, ) from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( + COSMOS3_ASPECT_RATIO_BUCKETS, COSMOS3_DOMAIN_PRESET_ALIASES, COSMOS3_DOMAIN_PRESETS, COSMOS3_EXTRA_SPECS, COSMOS3_POLICY_SAMPLING_PARAMS, - get_domain_preset, - resolve_checkpoint_policy_defaults, -) -from tensorrt_llm._torch.visual_gen.models.cosmos3.geometry import ( - COSMOS3_ASPECT_RATIO_BUCKETS, + EMBODIMENT_TO_DOMAIN_ID, + EMBODIMENT_TO_RAW_ACTION_DIM, VIDEO_RES_SIZE_INFO, find_closest_target_size, + get_domain_preset, normalize_action_resolution, + resolve_checkpoint_policy_defaults, + resolve_raw_action_dim, ) from tensorrt_llm._torch.visual_gen.models.cosmos3.request_resolver import ( resolve_domain_action_config, diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_example_prompts.py b/tests/unittest/_torch/visual_gen/test_cosmos3_example_prompts.py index 6f4951b87155..e8011bbd7bca 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_example_prompts.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_example_prompts.py @@ -16,11 +16,9 @@ import pytest from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( - resolve_checkpoint_policy_defaults, -) -from tensorrt_llm._torch.visual_gen.models.cosmos3.geometry import ( VIDEO_RES_SIZE_INFO, find_closest_target_size, + resolve_checkpoint_policy_defaults, ) from tensorrt_llm._torch.visual_gen.models.cosmos3.request_resolver import ( resolve_domain_action_config, diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py b/tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py index 50f256cec23a..17cbf76148b3 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py @@ -30,8 +30,6 @@ COSMOS3_720P_PARAMS, COSMOS3_EXTRA_SPECS, COSMOS3_GENERATION_DEFAULTS, -) -from tensorrt_llm._torch.visual_gen.models.cosmos3.geometry import ( VIDEO_RES_SIZE_INFO, find_closest_target_size, ) From 5c0fad9053d12bea11fbec1c644fd2587cd8787e Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 2 Nov 2026 09:44:24 -0800 Subject: [PATCH 5/8] [None][test] Cosmos3: drop the request-defaults golden fixture Signed-off-by: Igor Shovkun --- .../visual_gen/cosmos3_request_golden.py | 155 ----- .../test_cosmos3_request_resolver.py | 20 +- .../test_data/cosmos3_request_defaults.json | 631 ------------------ 3 files changed, 2 insertions(+), 804 deletions(-) delete mode 100644 tests/unittest/_torch/visual_gen/cosmos3_request_golden.py delete mode 100644 tests/unittest/_torch/visual_gen/test_data/cosmos3_request_defaults.json diff --git a/tests/unittest/_torch/visual_gen/cosmos3_request_golden.py b/tests/unittest/_torch/visual_gen/cosmos3_request_golden.py deleted file mode 100644 index 367f2a625991..000000000000 --- a/tests/unittest/_torch/visual_gen/cosmos3_request_golden.py +++ /dev/null @@ -1,155 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""The request-level defaults every Cosmos3 family, mode and domain resolves to when -the caller sets nothing. - -``record()`` evaluates the grid through the request resolver; -``python cosmos3_request_golden.py > test_data/cosmos3_request_defaults.json`` stores -it. ``test_cosmos3_request_resolver.py`` compares the live grid with the stored one, -so a change of a table, a rule or a precedence is a visible diff of that file that -has to be regenerated on purpose. -""" - -import json -import sys - -from diffusers import FlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler - -from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( - COSMOS3_DOMAIN_PRESETS, - EMBODIMENT_TO_DOMAIN_ID, - resolve_checkpoint_policy_defaults, -) -from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import ( - Cosmos3OmniMoTPipeline, - _resolve_use_system_prompt, -) -from tensorrt_llm._torch.visual_gen.models.cosmos3.request_resolver import ( - resolve_domain_action_config, -) -from tensorrt_llm._torch.visual_gen.models.cosmos3.sampling import Cosmos3SamplingPolicy - -FAMILIES = ("qwen3", "nemotron_dense") -MODES = ("video", "image") -GENERATION_FIELDS = ( - "height", - "width", - "num_frames", - "frame_rate", - "num_inference_steps", - "guidance_scale", - "guidance_interval", - "max_sequence_length", -) -DROID_POLICY = {"domain_name": "droid_lerobot", "action_chunk_size": 32, "conditioning_fps": 15.0} - -UNIPC = { - "_class_name": "UniPCMultistepScheduler", - "num_train_timesteps": 1000, - "flow_shift": 1.0, - "prediction_type": "flow_prediction", - "use_flow_sigmas": True, - "use_karras_sigmas": True, - "solver_order": 2, -} -DISTILLED = { - "_class_name": "FlowMatchEulerDiscreteScheduler", - "num_train_timesteps": 1000, - "shift": 1.0, - "stochastic_sampling": True, - "fixed_step_sampler_config": {"sample_type": "sde", "t_list": [1.0, 0.9375, 0.8333, 0.625]}, -} - - -def _policy(distilled: bool) -> Cosmos3SamplingPolicy: - if distilled: - return Cosmos3SamplingPolicy.from_scheduler( - FlowMatchEulerDiscreteScheduler.from_config(DISTILLED) - ) - return Cosmos3SamplingPolicy.from_scheduler(UniPCMultistepScheduler.from_config(UNIPC)) - - -def _pipeline(family: str, distilled: bool, policy_block=None) -> Cosmos3OmniMoTPipeline: - pipe = object.__new__(Cosmos3OmniMoTPipeline) - pipe.family = family - pipe.sampling = _policy(distilled) - pipe.checkpoint_policy_defaults = resolve_checkpoint_policy_defaults(policy_block) - return pipe - - -def _jsonable(value): - if isinstance(value, tuple): - return list(value) - return value - - -def record() -> dict: - out = {"generation": {}, "action": {}, "system_prompt": {}} - for family in FAMILIES: - for distilled in (False, True): - if family == "nemotron_dense" and distilled: - continue # no such checkpoint; the pipeline refuses the combination - pipe = _pipeline(family, distilled) - for mode in MODES: - values = pipe._resolve_generation_params( - mode, **{f: None for f in GENERATION_FIELDS} - ) - key = f"{family}/{mode}/{'distilled' if distilled else 'base'}" - out["generation"][key] = {k: _jsonable(v) for k, v in values.items()} - # Action sampling fields with and without the DROID policy block. - for label, block in (("", None), ("/droid-policy", DROID_POLICY)): - pipe = _pipeline(family, distilled, block) - for action_mode in ("forward_dynamics", "policy"): - values = pipe._resolve_generation_params( - "video", - action_mode=action_mode, - num_inference_steps=None, - guidance_scale=None, - guidance_interval=None, - max_sequence_length=None, - ) - key = f"{family}/action:{action_mode}/{'distilled' if distilled else 'base'}{label}" - out["generation"][key] = {k: _jsonable(v) for k, v in values.items()} - domains = sorted(set(EMBODIMENT_TO_DOMAIN_ID) | set(COSMOS3_DOMAIN_PRESETS)) + [None] - for policy_label, block in (("", None), ("/droid-policy", DROID_POLICY)): - policy_defaults = resolve_checkpoint_policy_defaults(block) - for domain in domains: - try: - resolved = resolve_domain_action_config( - domain_name=domain, checkpoint_policy_defaults=policy_defaults - ) - cfg = { - k: _jsonable(v) for k, v in resolved.items() if k not in ("warnings", "sources") - } - cfg["num_warnings"] = len(resolved["warnings"]) - except ValueError as exc: - cfg = {"error": str(exc)} - out["action"][f"{domain}{policy_label}"] = cfg - for checkpoint_default in (False, True): - for is_v2v in (False, True): - for is_transfer in (False, True): - key = f"checkpoint={checkpoint_default}/v2v={is_v2v}/transfer={is_transfer}" - out["system_prompt"][key] = _resolve_use_system_prompt( - None, - checkpoint_default=checkpoint_default, - is_v2v=is_v2v, - is_transfer=is_transfer, - ) - return out - - -if __name__ == "__main__": - json.dump(record(), sys.stdout, indent=1, sort_keys=True) - sys.stdout.write("\n") diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_request_resolver.py b/tests/unittest/_torch/visual_gen/test_cosmos3_request_resolver.py index 25a461be9202..612dc0458586 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_request_resolver.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_request_resolver.py @@ -12,14 +12,10 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""The request rules: the recorded defaults of every family, mode and domain, and the -source each value comes from.""" - -import json -from pathlib import Path +"""The request rules: precedence between the caller, the checkpoint's facts and the +tables, and the source each value reports.""" import pytest -from cosmos3_request_golden import record from tensorrt_llm._torch.visual_gen.models.cosmos3.request_resolver import ( RequestContext, @@ -30,18 +26,6 @@ pytestmark = pytest.mark.cpu_only -FIXTURE = Path(__file__).parent / "test_data" / "cosmos3_request_defaults.json" - - -def test_defaults_match_the_recorded_grid(): - """Every (family, mode, domain) combination resolves to the recorded values, so a - change of any default, or of the precedence behind it, is a visible diff of the - fixture rather than a silent change of what requests run with.""" - recorded = json.loads(FIXTURE.read_text()) - current = json.loads(json.dumps(record())) - for section in ("generation", "action", "system_prompt"): - assert current[section] == recorded[section], section - class TestSources: def test_generation_values_name_their_source(self): diff --git a/tests/unittest/_torch/visual_gen/test_data/cosmos3_request_defaults.json b/tests/unittest/_torch/visual_gen/test_data/cosmos3_request_defaults.json deleted file mode 100644 index d38e8627f0cc..000000000000 --- a/tests/unittest/_torch/visual_gen/test_data/cosmos3_request_defaults.json +++ /dev/null @@ -1,631 +0,0 @@ -{ - "action": { - "None": { - "action_chunk_size": 16, - "action_fps": 24.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": false, - "domain_name": null, - "frame_rate": 24.0, - "num_frames": 17, - "num_warnings": 0, - "preset_key": null, - "raw_action_dim": null, - "use_state": false - }, - "None/droid-policy": { - "action_chunk_size": 32, - "action_fps": 15.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": true, - "domain_name": "droid_lerobot", - "frame_rate": 15.0, - "num_frames": 33, - "num_warnings": 0, - "preset_key": "droid_lerobot", - "raw_action_dim": 8, - "use_state": true - }, - "agibot_gear_gripper": { - "action_chunk_size": 16, - "action_fps": 10.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": false, - "domain_name": "agibot_gear_gripper", - "frame_rate": 10.0, - "num_frames": 17, - "num_warnings": 0, - "preset_key": "agibotworld", - "raw_action_dim": 29, - "use_state": false - }, - "agibot_gear_gripper/droid-policy": { - "action_chunk_size": 16, - "action_fps": 10.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": false, - "domain_name": "agibot_gear_gripper", - "frame_rate": 10.0, - "num_frames": 17, - "num_warnings": 1, - "preset_key": "agibotworld", - "raw_action_dim": 29, - "use_state": false - }, - "agibot_gear_gripper_ext": { - "action_chunk_size": 16, - "action_fps": 10.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": false, - "domain_name": "agibot_gear_gripper_ext", - "frame_rate": 10.0, - "num_frames": 17, - "num_warnings": 0, - "preset_key": "agibotworld", - "raw_action_dim": 29, - "use_state": false - }, - "agibot_gear_gripper_ext/droid-policy": { - "action_chunk_size": 16, - "action_fps": 10.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": false, - "domain_name": "agibot_gear_gripper_ext", - "frame_rate": 10.0, - "num_frames": 17, - "num_warnings": 1, - "preset_key": "agibotworld", - "raw_action_dim": 29, - "use_state": false - }, - "agibotworld": { - "action_chunk_size": 16, - "action_fps": 10.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": false, - "domain_name": "agibotworld", - "frame_rate": 10.0, - "num_frames": 17, - "num_warnings": 0, - "preset_key": "agibotworld", - "raw_action_dim": 29, - "use_state": false - }, - "agibotworld/droid-policy": { - "action_chunk_size": 16, - "action_fps": 10.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": false, - "domain_name": "agibotworld", - "frame_rate": 10.0, - "num_frames": 17, - "num_warnings": 1, - "preset_key": "agibotworld", - "raw_action_dim": 29, - "use_state": false - }, - "av": { - "action_chunk_size": 60, - "action_fps": 10.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": false, - "domain_name": "av", - "frame_rate": 10.0, - "num_frames": 61, - "num_warnings": 0, - "preset_key": "av", - "raw_action_dim": 9, - "use_state": false - }, - "av/droid-policy": { - "action_chunk_size": 60, - "action_fps": 10.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": false, - "domain_name": "av", - "frame_rate": 10.0, - "num_frames": 61, - "num_warnings": 1, - "preset_key": "av", - "raw_action_dim": 9, - "use_state": false - }, - "bridge_orig_lerobot": { - "action_chunk_size": 16, - "action_fps": 5.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": false, - "domain_name": "bridge_orig_lerobot", - "frame_rate": 5.0, - "num_frames": 17, - "num_warnings": 0, - "preset_key": "bridge_orig_lerobot", - "raw_action_dim": 10, - "use_state": false - }, - "bridge_orig_lerobot/droid-policy": { - "action_chunk_size": 16, - "action_fps": 5.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": false, - "domain_name": "bridge_orig_lerobot", - "frame_rate": 5.0, - "num_frames": 17, - "num_warnings": 1, - "preset_key": "bridge_orig_lerobot", - "raw_action_dim": 10, - "use_state": false - }, - "camera_pose": { - "action_chunk_size": 60, - "action_fps": 30.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": false, - "domain_name": "camera_pose", - "frame_rate": 30.0, - "num_frames": 61, - "num_warnings": 0, - "preset_key": "camera_pose", - "raw_action_dim": 9, - "use_state": false - }, - "camera_pose/droid-policy": { - "action_chunk_size": 60, - "action_fps": 30.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": false, - "domain_name": "camera_pose", - "frame_rate": 30.0, - "num_frames": 61, - "num_warnings": 1, - "preset_key": "camera_pose", - "raw_action_dim": 9, - "use_state": false - }, - "droid_lerobot": { - "action_chunk_size": 16, - "action_fps": 15.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": false, - "domain_name": "droid_lerobot", - "frame_rate": 15.0, - "num_frames": 17, - "num_warnings": 0, - "preset_key": "droid_lerobot", - "raw_action_dim": 10, - "use_state": false - }, - "droid_lerobot/droid-policy": { - "action_chunk_size": 32, - "action_fps": 15.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": true, - "domain_name": "droid_lerobot", - "frame_rate": 15.0, - "num_frames": 33, - "num_warnings": 0, - "preset_key": "droid_lerobot", - "raw_action_dim": 8, - "use_state": true - }, - "fractal": { - "action_chunk_size": 16, - "action_fps": 5.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": false, - "domain_name": "fractal", - "frame_rate": 5.0, - "num_frames": 17, - "num_warnings": 0, - "preset_key": "fractal", - "raw_action_dim": 10, - "use_state": false - }, - "fractal/droid-policy": { - "action_chunk_size": 16, - "action_fps": 5.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": false, - "domain_name": "fractal", - "frame_rate": 5.0, - "num_frames": 17, - "num_warnings": 1, - "preset_key": "fractal", - "raw_action_dim": 10, - "use_state": false - }, - "galbot": { - "action_chunk_size": 16, - "action_fps": 10.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": false, - "domain_name": "galbot", - "frame_rate": 10.0, - "num_frames": 17, - "num_warnings": 0, - "preset_key": "agibotworld", - "raw_action_dim": 30, - "use_state": false - }, - "galbot/droid-policy": { - "action_chunk_size": 16, - "action_fps": 10.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": false, - "domain_name": "galbot", - "frame_rate": 10.0, - "num_frames": 17, - "num_warnings": 1, - "preset_key": "agibotworld", - "raw_action_dim": 30, - "use_state": false - }, - "hand_pose": { - "action_chunk_size": 16, - "action_fps": 24.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": false, - "domain_name": "hand_pose", - "frame_rate": 24.0, - "num_frames": 17, - "num_warnings": 0, - "preset_key": "hand_pose", - "raw_action_dim": 57, - "use_state": false - }, - "hand_pose/droid-policy": { - "action_chunk_size": 16, - "action_fps": 24.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": false, - "domain_name": "hand_pose", - "frame_rate": 24.0, - "num_frames": 17, - "num_warnings": 1, - "preset_key": "hand_pose", - "raw_action_dim": 57, - "use_state": false - }, - "libero": { - "action_chunk_size": 16, - "action_fps": 10.0, - "action_resolution": 256, - "checkpoint_policy_domain_matches": false, - "domain_name": "libero", - "frame_rate": 10.0, - "num_frames": 17, - "num_warnings": 1, - "preset_key": "libero", - "raw_action_dim": null, - "use_state": false - }, - "libero/droid-policy": { - "action_chunk_size": 16, - "action_fps": 10.0, - "action_resolution": 256, - "checkpoint_policy_domain_matches": false, - "domain_name": "libero", - "frame_rate": 10.0, - "num_frames": 17, - "num_warnings": 2, - "preset_key": "libero", - "raw_action_dim": null, - "use_state": false - }, - "no_action": { - "action_chunk_size": 16, - "action_fps": 24.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": false, - "domain_name": "no_action", - "frame_rate": 24.0, - "num_frames": 17, - "num_warnings": 2, - "preset_key": null, - "raw_action_dim": null, - "use_state": false - }, - "no_action/droid-policy": { - "action_chunk_size": 16, - "action_fps": 24.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": false, - "domain_name": "no_action", - "frame_rate": 24.0, - "num_frames": 17, - "num_warnings": 3, - "preset_key": null, - "raw_action_dim": null, - "use_state": false - }, - "pusht": { - "action_chunk_size": 16, - "action_fps": 10.0, - "action_resolution": 256, - "checkpoint_policy_domain_matches": false, - "domain_name": "pusht", - "frame_rate": 10.0, - "num_frames": 17, - "num_warnings": 0, - "preset_key": "pusht", - "raw_action_dim": 2, - "use_state": false - }, - "pusht/droid-policy": { - "action_chunk_size": 16, - "action_fps": 10.0, - "action_resolution": 256, - "checkpoint_policy_domain_matches": false, - "domain_name": "pusht", - "frame_rate": 10.0, - "num_frames": 17, - "num_warnings": 1, - "preset_key": "pusht", - "raw_action_dim": 2, - "use_state": false - }, - "robomind-franka": { - "action_chunk_size": 16, - "action_fps": 15.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": false, - "domain_name": "robomind-franka", - "frame_rate": 15.0, - "num_frames": 17, - "num_warnings": 0, - "preset_key": "droid_lerobot", - "raw_action_dim": 10, - "use_state": false - }, - "robomind-franka-dual": { - "action_chunk_size": 16, - "action_fps": 15.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": false, - "domain_name": "robomind-franka-dual", - "frame_rate": 15.0, - "num_frames": 17, - "num_warnings": 0, - "preset_key": "droid_lerobot", - "raw_action_dim": 20, - "use_state": false - }, - "robomind-franka-dual/droid-policy": { - "action_chunk_size": 16, - "action_fps": 15.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": false, - "domain_name": "robomind-franka-dual", - "frame_rate": 15.0, - "num_frames": 17, - "num_warnings": 1, - "preset_key": "droid_lerobot", - "raw_action_dim": 20, - "use_state": false - }, - "robomind-franka/droid-policy": { - "action_chunk_size": 16, - "action_fps": 15.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": false, - "domain_name": "robomind-franka", - "frame_rate": 15.0, - "num_frames": 17, - "num_warnings": 1, - "preset_key": "droid_lerobot", - "raw_action_dim": 10, - "use_state": false - }, - "robomind-ur": { - "action_chunk_size": 16, - "action_fps": 15.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": false, - "domain_name": "robomind-ur", - "frame_rate": 15.0, - "num_frames": 17, - "num_warnings": 0, - "preset_key": "droid_lerobot", - "raw_action_dim": 10, - "use_state": false - }, - "robomind-ur/droid-policy": { - "action_chunk_size": 16, - "action_fps": 15.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": false, - "domain_name": "robomind-ur", - "frame_rate": 15.0, - "num_frames": 17, - "num_warnings": 1, - "preset_key": "droid_lerobot", - "raw_action_dim": 10, - "use_state": false - }, - "umi": { - "action_chunk_size": 16, - "action_fps": 10.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": false, - "domain_name": "umi", - "frame_rate": 10.0, - "num_frames": 17, - "num_warnings": 0, - "preset_key": "umi", - "raw_action_dim": 10, - "use_state": false - }, - "umi/droid-policy": { - "action_chunk_size": 16, - "action_fps": 10.0, - "action_resolution": 480, - "checkpoint_policy_domain_matches": false, - "domain_name": "umi", - "frame_rate": 10.0, - "num_frames": 17, - "num_warnings": 1, - "preset_key": "umi", - "raw_action_dim": 10, - "use_state": false - } - }, - "generation": { - "nemotron_dense/action:forward_dynamics/base": { - "guidance_interval": null, - "guidance_scale": 1.0, - "max_sequence_length": 4096, - "num_inference_steps": 30 - }, - "nemotron_dense/action:forward_dynamics/base/droid-policy": { - "guidance_interval": null, - "guidance_scale": 1.0, - "max_sequence_length": 4096, - "num_inference_steps": 30 - }, - "nemotron_dense/action:policy/base": { - "guidance_interval": null, - "guidance_scale": 1.0, - "max_sequence_length": 4096, - "num_inference_steps": 30 - }, - "nemotron_dense/action:policy/base/droid-policy": { - "guidance_interval": [ - 960.0, - 1001.0 - ], - "guidance_scale": 3.0, - "max_sequence_length": 4096, - "num_inference_steps": 4 - }, - "nemotron_dense/image/base": { - "frame_rate": 24.0, - "guidance_interval": null, - "guidance_scale": 4.0, - "height": 640, - "max_sequence_length": 4096, - "num_frames": 121, - "num_inference_steps": 50, - "width": 640 - }, - "nemotron_dense/video/base": { - "frame_rate": 24.0, - "guidance_interval": null, - "guidance_scale": 5.0, - "height": 480, - "max_sequence_length": 4096, - "num_frames": 121, - "num_inference_steps": 50, - "width": 832 - }, - "qwen3/action:forward_dynamics/base": { - "guidance_interval": null, - "guidance_scale": 1.0, - "max_sequence_length": 4096, - "num_inference_steps": 30 - }, - "qwen3/action:forward_dynamics/base/droid-policy": { - "guidance_interval": null, - "guidance_scale": 1.0, - "max_sequence_length": 4096, - "num_inference_steps": 30 - }, - "qwen3/action:forward_dynamics/distilled": { - "guidance_interval": null, - "guidance_scale": 1.0, - "max_sequence_length": 4096, - "num_inference_steps": 4 - }, - "qwen3/action:forward_dynamics/distilled/droid-policy": { - "guidance_interval": null, - "guidance_scale": 1.0, - "max_sequence_length": 4096, - "num_inference_steps": 4 - }, - "qwen3/action:policy/base": { - "guidance_interval": null, - "guidance_scale": 1.0, - "max_sequence_length": 4096, - "num_inference_steps": 30 - }, - "qwen3/action:policy/base/droid-policy": { - "guidance_interval": [ - 960.0, - 1001.0 - ], - "guidance_scale": 3.0, - "max_sequence_length": 4096, - "num_inference_steps": 4 - }, - "qwen3/action:policy/distilled": { - "guidance_interval": null, - "guidance_scale": 1.0, - "max_sequence_length": 4096, - "num_inference_steps": 4 - }, - "qwen3/action:policy/distilled/droid-policy": { - "guidance_interval": [ - 960.0, - 1001.0 - ], - "guidance_scale": 1.0, - "max_sequence_length": 4096, - "num_inference_steps": 4 - }, - "qwen3/image/base": { - "frame_rate": 24.0, - "guidance_interval": [ - 400.0, - 1000.0 - ], - "guidance_scale": 7.0, - "height": 1024, - "max_sequence_length": 4096, - "num_frames": 189, - "num_inference_steps": 50, - "width": 1024 - }, - "qwen3/image/distilled": { - "frame_rate": 24.0, - "guidance_interval": [ - 400.0, - 1000.0 - ], - "guidance_scale": 1.0, - "height": 1024, - "max_sequence_length": 4096, - "num_frames": 189, - "num_inference_steps": 4, - "width": 1024 - }, - "qwen3/video/base": { - "frame_rate": 24.0, - "guidance_interval": null, - "guidance_scale": 6.0, - "height": 720, - "max_sequence_length": 4096, - "num_frames": 189, - "num_inference_steps": 35, - "width": 1280 - }, - "qwen3/video/distilled": { - "frame_rate": 24.0, - "guidance_interval": null, - "guidance_scale": 1.0, - "height": 720, - "max_sequence_length": 4096, - "num_frames": 189, - "num_inference_steps": 4, - "width": 1280 - } - }, - "system_prompt": { - "checkpoint=False/v2v=False/transfer=False": false, - "checkpoint=False/v2v=False/transfer=True": false, - "checkpoint=False/v2v=True/transfer=False": true, - "checkpoint=False/v2v=True/transfer=True": false, - "checkpoint=True/v2v=False/transfer=False": true, - "checkpoint=True/v2v=False/transfer=True": false, - "checkpoint=True/v2v=True/transfer=False": true, - "checkpoint=True/v2v=True/transfer=True": false - } -} From c99415189a298916626551c894370789a060fb6b Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 2 Nov 2026 09:58:16 -0800 Subject: [PATCH 6/8] [None][refactor] Cosmos3: one parameter resolver for checkpoint and request settings Signed-off-by: Igor Shovkun --- .../models/cosmos3/config_resolver.py | 375 -------- .../models/cosmos3/parameter_resolver.py | 852 ++++++++++++++++++ .../models/cosmos3/pipeline_cosmos3.py | 10 +- .../models/cosmos3/request_resolver.py | 477 ---------- .../models/cosmos3/transformer_cosmos3.py | 2 +- .../integration/test_lists/test-db/l0_cpu.yml | 3 +- .../_torch/visual_gen/test_cosmos3_action.py | 2 +- .../test_cosmos3_config_resolver.py | 203 ----- .../test_cosmos3_example_prompts.py | 2 +- .../test_cosmos3_parameter_resolver.py | 301 +++++++ .../test_cosmos3_request_resolver.py | 94 -- .../visual_gen/test_cosmos3_transformer.py | 2 +- .../_torch/visual_gen/test_data/README.md | 4 +- 13 files changed, 1167 insertions(+), 1160 deletions(-) delete mode 100644 tensorrt_llm/_torch/visual_gen/models/cosmos3/config_resolver.py create mode 100644 tensorrt_llm/_torch/visual_gen/models/cosmos3/parameter_resolver.py delete mode 100644 tensorrt_llm/_torch/visual_gen/models/cosmos3/request_resolver.py delete mode 100644 tests/unittest/_torch/visual_gen/test_cosmos3_config_resolver.py create mode 100644 tests/unittest/_torch/visual_gen/test_cosmos3_parameter_resolver.py delete mode 100644 tests/unittest/_torch/visual_gen/test_cosmos3_request_resolver.py diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/config_resolver.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/config_resolver.py deleted file mode 100644 index b5756e584e35..000000000000 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/config_resolver.py +++ /dev/null @@ -1,375 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""One place that turns a Cosmos3 checkpoint's files into settings. - -A Cosmos3 export carries the same facts in several files, written by different -exporters that are not kept in sync: the training config (``config.json``), the -diffusers pipeline index (``model_index.json``), the export metadata -(``checkpoint.json``) and the component configs (``transformer/``, ``scheduler/``, -``vae/``). :data:`RULES` names, for every setting, the sources in priority order -(the first one present wins) and what a disagreeing lower source means. -:class:`Cosmos3ConfigResolver` applies the table and returns a frozen -:class:`Cosmos3Config` whose every value remembers where it came from. - -A setting no source provides and no default covers is an error. Every -:class:`Default` in the table is a value agreed on for checkpoints that omit the key. -""" - -from __future__ import annotations - -import json -import math -import os -from collections.abc import Mapping -from dataclasses import dataclass, replace -from enum import Enum -from typing import Any, Optional - -from tensorrt_llm.logger import logger - -CONFIG = "config.json" -CHECKPOINT = "checkpoint.json" -MODEL_INDEX = "model_index.json" -TRANSFORMER = "transformer/config.json" -SCHEDULER = "scheduler/scheduler_config.json" -VAE = "vae/config.json" - -_MISSING = object() - - -class Conflict(Enum): - """What a lower-priority source holding a different value means.""" - - ERROR = "error" - WARN = "warn" - IGNORE = "ignore" - - -@dataclass(frozen=True) -class Key: - """A dotted key path in one checkpoint file.""" - - file: str - path: str - - @property - def label(self) -> str: - return f"{self.file}:{self.path}" - - -@dataclass(frozen=True) -class Default: - """The value used when no source declares the setting.""" - - value: Any - why: str - - @property - def label(self) -> str: - return f"default ({self.why})" - - -@dataclass(frozen=True) -class Rule: - sources: tuple - conflict: Conflict = Conflict.WARN - required: bool = True - """``False``: an unresolved setting is ``None`` instead of an error.""" - - -@dataclass(frozen=True) -class Resolved: - value: Any - source: str - disagreeing: tuple[str, ...] = () - """Lower-priority sources that held a different value.""" - - -# Settings every Cosmos3 checkpoint resolves. Priority: the training config -# (config.json), then the component configs, then the export metadata and the -# diffusers index. -RULES: dict[str, Rule] = { - "architecture.position_embedding_type": Rule( - ( - Key(TRANSFORMER, "position_embedding_type"), - Default("unified_3d_mrope", "exports that predate the key"), - ) - ), - "architecture.max_position_embeddings": Rule( - ( - Key(TRANSFORMER, "max_position_embeddings"), - Default(262144, "exports that predate the key"), - ) - ), - "vae.temporal_compression": Rule( - ( - Key(VAE, "scale_factor_temporal"), - Key(TRANSFORMER, "temporal_compression_factor"), - Default(4, "exports that predate the key"), - ), - conflict=Conflict.ERROR, - ), - "vae.spatial_compression": Rule((Key(VAE, "scale_factor_spatial"),), conflict=Conflict.ERROR), - "audio.temporal_compression": Rule( - ( - Key(TRANSFORMER, "temporal_compression_factor_sound"), - Default(1, "exports that predate the key"), - ) - ), - "prompt.use_system_prompt": Rule( - ( - Key(CONFIG, "model.config.vlm_config.use_system_prompt"), - Key(MODEL_INDEX, "default_use_system_prompt"), - ) - ), - "sampling.native_flow_schedule": Rule( - ( - Key(MODEL_INDEX, "use_native_flow_schedule"), - Default(False, "only Edge exports declare the native grid"), - ) - ), - "sampling.fixed_sigmas": Rule( - (Key(SCHEDULER, "fixed_step_sampler_config.t_list"),), required=False - ), - "action.policy": Rule((Key(CHECKPOINT, "policy"),), required=False), -} - - -@dataclass(frozen=True) -class Cosmos3ArchitectureSettings: - position_embedding_type: str - max_position_embeddings: int - - -@dataclass(frozen=True) -class Cosmos3VAESettings: - temporal_compression: int - """Pixel frames per latent frame after the first, which is encoded alone.""" - spatial_compression: int - - -@dataclass(frozen=True) -class Cosmos3AudioSettings: - temporal_compression: int - - -@dataclass(frozen=True) -class Cosmos3SamplingSettings: - fixed_sigmas: Optional[tuple[float, ...]] - """The distilled schedule (``t_list``); ``None`` on base checkpoints.""" - native_flow_schedule: bool - - -@dataclass(frozen=True) -class Cosmos3PromptSettings: - use_system_prompt: bool - - -@dataclass(frozen=True) -class Cosmos3ActionSettings: - policy: Optional[Mapping[str, Any]] - """The ``checkpoint.json`` policy block, unnormalized.""" - - -@dataclass(frozen=True) -class Cosmos3Config: - architecture: Cosmos3ArchitectureSettings - vae: Cosmos3VAESettings - audio: Cosmos3AudioSettings - sampling: Cosmos3SamplingSettings - prompt: Cosmos3PromptSettings - action: Cosmos3ActionSettings - resolved: Mapping[str, Resolved] - - def source_of(self, name: str) -> str: - return self.resolved[name].source - - def describe(self) -> str: - lines = ["Cosmos3 config:"] - for name, r in sorted(self.resolved.items()): - value = r.value - if isinstance(value, Mapping): - value = f"<{len(value)} entries>" - line = f" {name} = {value!r} [{r.source}]" - if r.disagreeing: - line += f" (disagrees: {', '.join(r.disagreeing)})" - lines.append(line) - return "\n".join(lines) - - -def _step(obj: Any, key: str) -> Any: - if isinstance(obj, Mapping): - return obj.get(key, _MISSING) - if obj is None or isinstance(obj, (str, int, float, bool, list, tuple)): - return _MISSING - return getattr(obj, key, _MISSING) # a pretrained config object standing in for a file - - -def _walk(doc: Any, path: str) -> Any: - node = doc - for part in path.split("."): - node = _step(node, part) - if node is _MISSING: - return _MISSING - return node - - -def same_value(a: Any, b: Any) -> bool: - """Whether two settings agree: bools by identity, numbers within 1e-8, lists and - mappings elementwise.""" - if isinstance(a, bool) or isinstance(b, bool): - return a is b or a == b - if isinstance(a, (int, float)) and isinstance(b, (int, float)): - return math.isclose(a, b, rel_tol=0.0, abs_tol=1e-8) - if isinstance(a, (list, tuple)) and isinstance(b, (list, tuple)): - return len(a) == len(b) and all(same_value(x, y) for x, y in zip(a, b)) - if isinstance(a, Mapping) and isinstance(b, Mapping): - return a.keys() == b.keys() and all(same_value(a[k], b[k]) for k in a) - return a == b - - -class Cosmos3ConfigResolver: - """Applies :data:`RULES` to a checkpoint directory. - - ``pretrained_config`` stands in for ``transformer/config.json`` when given; - ``checkpoint_dir`` may then be ``None``. - """ - - def __init__(self, checkpoint_dir: Optional[str], pretrained_config: Any = None) -> None: - self._dir = checkpoint_dir - self._pretrained = pretrained_config - self._docs: dict[str, Any] = {} - - def _doc(self, file: str) -> Any: - if file == TRANSFORMER and self._pretrained is not None: - return self._pretrained - if file not in self._docs: - doc = None - if self._dir is not None: - path = os.path.join(self._dir, file) - if os.path.exists(path): - with open(path) as f: - doc = json.load(f) - self._docs[file] = doc - return self._docs[file] - - def _lookup(self, source: Any) -> Any: - if isinstance(source, Key): - doc = self._doc(source.file) - return _MISSING if doc is None else _walk(doc, source.path) - if isinstance(source, Default): - return source.value - raise TypeError(f"unknown source {source!r}") - - def resolve_one(self, name: str, rule: Rule) -> Optional[Resolved]: - winner: Optional[Resolved] = None - disagreeing = [] - for source in rule.sources: - if winner is not None and isinstance(source, Default): - break - value = self._lookup(source) - if value is _MISSING or value is None: - continue # a null in a file counts as absent - if winner is None: - winner = Resolved(value, source.label) - elif not same_value(value, winner.value): - disagreeing.append(f"{source.label}={value!r}") - if winner is None: - if rule.required: - labels = ", ".join(s.label for s in rule.sources) - raise ValueError(f"Cosmos3 config: cannot resolve {name}; none of {labels} is set.") - return None - if disagreeing: - message = ( - f"Cosmos3 config: {name} = {winner.value!r} from {winner.source}, but " - + "; ".join(disagreeing) - ) - if rule.conflict is Conflict.ERROR: - raise ValueError(message) - if rule.conflict is Conflict.WARN: - logger.warning(message) - return Resolved(winner.value, winner.source, tuple(disagreeing)) - - def resolve(self) -> Cosmos3Config: - resolved: dict[str, Resolved] = {} - for name, rule in RULES.items(): - r = self.resolve_one(name, rule) - if r is not None: - resolved[name] = r - - def get(name: str) -> Any: - r = resolved.get(name) - return None if r is None else r.value - - sigmas = get("sampling.fixed_sigmas") - if sigmas is not None and ( - any(not 0.0 < float(s) <= 1.0 for s in sigmas) - or [float(s) for s in sigmas] != sorted((float(s) for s in sigmas), reverse=True) - ): - raise ValueError( - f"Cosmos3 config: fixed_step_sampler_config.t_list must decrease within (0, 1]: " - f"{list(sigmas)}." - ) - return Cosmos3Config( - architecture=Cosmos3ArchitectureSettings( - position_embedding_type=str(get("architecture.position_embedding_type")), - max_position_embeddings=int(get("architecture.max_position_embeddings")), - ), - vae=Cosmos3VAESettings( - temporal_compression=int(get("vae.temporal_compression")), - spatial_compression=int(get("vae.spatial_compression")), - ), - audio=Cosmos3AudioSettings(temporal_compression=int(get("audio.temporal_compression"))), - sampling=Cosmos3SamplingSettings( - fixed_sigmas=None if not sigmas else tuple(float(s) for s in sigmas), - native_flow_schedule=bool(get("sampling.native_flow_schedule")), - ), - prompt=Cosmos3PromptSettings(use_system_prompt=bool(get("prompt.use_system_prompt"))), - action=Cosmos3ActionSettings(policy=get("action.policy")), - resolved=resolved, - ) - - -ARCHITECTURE_ATTRIBUTES = { - "position_embedding_type": "architecture.position_embedding_type", - "max_position_embeddings": "architecture.max_position_embeddings", - "temporal_compression_factor": "vae.temporal_compression", - "temporal_compression_factor_sound": "audio.temporal_compression", -} -"""``pretrained_config`` attributes the transformer reads, and the setting each holds.""" - - -def fill_pretrained_config(pretrained_config: Any) -> Any: - """Set the transformer-construction attributes of ``pretrained_config`` that the - export omits, from :data:`RULES`. Idempotent; returns the config. - - Runs before the pipeline loads its components, so only ``pretrained_config`` - itself is consulted. - """ - resolver = Cosmos3ConfigResolver(None, pretrained_config) - for attr, name in ARCHITECTURE_ATTRIBUTES.items(): - if getattr(pretrained_config, attr, None) is not None: - continue - rule = RULES[name] - transformer_only = replace( - rule, - sources=tuple( - s for s in rule.sources if not isinstance(s, Key) or s.file == TRANSFORMER - ), - ) - r = resolver.resolve_one(name, transformer_only) - if r is not None: - setattr(pretrained_config, attr, r.value) - return pretrained_config diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/parameter_resolver.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/parameter_resolver.py new file mode 100644 index 000000000000..dd36c097837d --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/parameter_resolver.py @@ -0,0 +1,852 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""One precedence table for every Cosmos3 setting: request over checkpoint over default. + +A rule names the sources of one setting in priority order; the first source holding a +value wins, the next one holding a value is compared with it, and every resolved +value remembers where it came from. The same engine serves two moments: + +* at load, :data:`CHECKPOINT_RULES` turn the checkpoint's files (``config.json``, + ``model_index.json``, ``checkpoint.json``, the component configs) into a + :class:`Cosmos3Config`: a Cosmos3 export carries the same facts in several files + written by exporters that are not kept in sync, so each rule says which file wins + and what a disagreement means; +* per request, :data:`GENERATION_RULES` and :data:`ACTION_RULES` turn the caller's + values, the checkpoint's facts (the distilled schedule, the policy block) and the + per-family and per-domain tables into what the request runs with. + +In every rule the caller's value comes first, then the checkpoint, then the code's +tables, then an explicit :class:`Default` with a reason. A setting no source provides +and no default covers is an error, not a guess. +""" + +from __future__ import annotations + +import json +import math +import os +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field, replace +from enum import Enum +from typing import Any, Optional + +from tensorrt_llm.logger import logger + +from .defaults import ( + COSMOS3_ACTION_PARAMS, + COSMOS3_DOMAIN_PRESETS, + COSMOS3_GENERATION_DEFAULTS, + EMBODIMENT_TO_DOMAIN_ID, + canonical_domain_preset_key, + normalize_action_resolution, + resolve_raw_action_dim, +) +from .sampling import DISTILLED_GUIDANCE_SCALE + +_MISSING = object() + + +# ------------------------------------------------------------------------------ engine + + +class Conflict(Enum): + """What a lower-priority source holding a different value means.""" + + ERROR = "error" + WARN = "warn" + IGNORE = "ignore" + + +@dataclass(frozen=True) +class Default: + """The value used when no source declares the setting; always the last source.""" + + value: Any + why: str + + @property + def label(self) -> str: + return f"default ({self.why})" + + def lookup(self, ctx, name, given, values): + return self.value + + +@dataclass(frozen=True) +class Rule: + sources: tuple + conflict: Conflict = Conflict.IGNORE + """``WARN`` / ``ERROR``: a source below the winner that holds a different value is + reported / rejected.""" + required: bool = True + """``False``: an unresolved setting is ``None`` instead of an error.""" + + +@dataclass(frozen=True) +class Resolved: + value: Any + source: str + disagreeing: tuple[str, ...] = () + """``label=value`` of the lower-priority source that held a different value.""" + + +def same_value(a: Any, b: Any) -> bool: + """Whether two settings agree: bools by identity, numbers within 1e-8, lists and + mappings elementwise.""" + if isinstance(a, bool) or isinstance(b, bool): + return a is b or a == b + if isinstance(a, (int, float)) and isinstance(b, (int, float)): + return math.isclose(a, b, rel_tol=0.0, abs_tol=1e-8) + if isinstance(a, (list, tuple)) and isinstance(b, (list, tuple)): + return len(a) == len(b) and all(same_value(x, y) for x, y in zip(a, b)) + if isinstance(a, Mapping) and isinstance(b, Mapping): + return a.keys() == b.keys() and all(same_value(a[k], b[k]) for k in a) + return a == b + + +def _pick( + name: str, rule: Rule, ctx: Any, given: Mapping[str, Any], values: Mapping[str, Any] +) -> tuple[Optional[Resolved], list[tuple[str, Any]]]: + """Walk ``rule.sources`` in order. The first source holding a value wins; the next + one holding a value is compared with it. A :class:`Default` is consulted only when + nothing above it holds a value.""" + winner: Optional[Resolved] = None + disagreeing: list[tuple[str, Any]] = [] + for source in rule.sources: + if winner is not None and isinstance(source, Default): + break + value = source.lookup(ctx, name, given, values) + if value is _MISSING: + continue + if winner is None: + winner = Resolved(value, source.label) + continue + if not same_value(value, winner.value): + disagreeing.append((source.label, value)) + break + return winner, disagreeing + + +def _record(winner: Resolved, disagreeing: list[tuple[str, Any]]) -> Resolved: + return Resolved( + winner.value, winner.source, tuple(f"{label}={v!r}" for label, v in disagreeing) + ) + + +# -------------------------------------------------------------------------- checkpoint + +CONFIG = "config.json" +CHECKPOINT = "checkpoint.json" +MODEL_INDEX = "model_index.json" +TRANSFORMER = "transformer/config.json" +SCHEDULER = "scheduler/scheduler_config.json" +VAE = "vae/config.json" + + +class CheckpointFiles: + """The checkpoint's files, read on demand. ``pretrained_config`` stands in for + ``transformer/config.json`` when given; ``checkpoint_dir`` may then be ``None``.""" + + def __init__(self, checkpoint_dir: Optional[str], pretrained_config: Any = None) -> None: + self._dir = checkpoint_dir + self._pretrained = pretrained_config + self._docs: dict[str, Any] = {} + + def doc(self, file: str) -> Any: + if file == TRANSFORMER and self._pretrained is not None: + return self._pretrained + if file not in self._docs: + doc = None + if self._dir is not None: + path = os.path.join(self._dir, file) + if os.path.exists(path): + with open(path) as f: + doc = json.load(f) + self._docs[file] = doc + return self._docs[file] + + +def _step(obj: Any, key: str) -> Any: + if isinstance(obj, Mapping): + return obj.get(key, _MISSING) + if obj is None or isinstance(obj, (str, int, float, bool, list, tuple)): + return _MISSING + return getattr(obj, key, _MISSING) # a pretrained config object standing in for a file + + +def _walk(doc: Any, path: str) -> Any: + node = doc + for part in path.split("."): + node = _step(node, part) + if node is _MISSING: + return _MISSING + return node + + +@dataclass(frozen=True) +class Key: + """A dotted key path in one checkpoint file; a ``null`` there counts as absent.""" + + file: str + path: str + + @property + def label(self) -> str: + return f"{self.file}:{self.path}" + + def lookup(self, ctx: CheckpointFiles, name, given, values): + doc = ctx.doc(self.file) + value = _MISSING if doc is None else _walk(doc, self.path) + return _MISSING if value is None else value + + +# Priority among the files: the training config (config.json), then the component +# configs, then the export metadata and the diffusers index. +CHECKPOINT_RULES: dict[str, Rule] = { + "architecture.position_embedding_type": Rule( + ( + Key(TRANSFORMER, "position_embedding_type"), + Default("unified_3d_mrope", "exports that predate the key"), + ) + ), + "architecture.max_position_embeddings": Rule( + ( + Key(TRANSFORMER, "max_position_embeddings"), + Default(262144, "exports that predate the key"), + ) + ), + "vae.temporal_compression": Rule( + ( + Key(VAE, "scale_factor_temporal"), + Key(TRANSFORMER, "temporal_compression_factor"), + Default(4, "exports that predate the key"), + ), + conflict=Conflict.ERROR, + ), + "vae.spatial_compression": Rule((Key(VAE, "scale_factor_spatial"),), conflict=Conflict.ERROR), + "audio.temporal_compression": Rule( + ( + Key(TRANSFORMER, "temporal_compression_factor_sound"), + Default(1, "exports that predate the key"), + ) + ), + "prompt.use_system_prompt": Rule( + ( + Key(CONFIG, "model.config.vlm_config.use_system_prompt"), + Key(MODEL_INDEX, "default_use_system_prompt"), + ), + conflict=Conflict.WARN, + ), + "sampling.native_flow_schedule": Rule( + ( + Key(MODEL_INDEX, "use_native_flow_schedule"), + Default(False, "only Edge exports declare the native grid"), + ) + ), + "sampling.fixed_sigmas": Rule( + (Key(SCHEDULER, "fixed_step_sampler_config.t_list"),), required=False + ), + "action.policy": Rule((Key(CHECKPOINT, "policy"),), required=False), +} + + +@dataclass(frozen=True) +class Cosmos3ArchitectureSettings: + position_embedding_type: str + max_position_embeddings: int + + +@dataclass(frozen=True) +class Cosmos3VAESettings: + temporal_compression: int + """Pixel frames per latent frame after the first, which is encoded alone.""" + spatial_compression: int + + +@dataclass(frozen=True) +class Cosmos3AudioSettings: + temporal_compression: int + + +@dataclass(frozen=True) +class Cosmos3SamplingSettings: + fixed_sigmas: Optional[tuple[float, ...]] + """The distilled schedule (``t_list``); ``None`` on base checkpoints.""" + native_flow_schedule: bool + + +@dataclass(frozen=True) +class Cosmos3PromptSettings: + use_system_prompt: bool + + +@dataclass(frozen=True) +class Cosmos3ActionSettings: + policy: Optional[Mapping[str, Any]] + """The ``checkpoint.json`` policy block, unnormalized.""" + + +@dataclass(frozen=True) +class Cosmos3Config: + architecture: Cosmos3ArchitectureSettings + vae: Cosmos3VAESettings + audio: Cosmos3AudioSettings + sampling: Cosmos3SamplingSettings + prompt: Cosmos3PromptSettings + action: Cosmos3ActionSettings + resolved: Mapping[str, Resolved] + + def source_of(self, name: str) -> str: + return self.resolved[name].source + + def describe(self) -> str: + lines = ["Cosmos3 config:"] + for name, r in sorted(self.resolved.items()): + value = r.value + if isinstance(value, Mapping): + value = f"<{len(value)} entries>" + line = f" {name} = {value!r} [{r.source}]" + if r.disagreeing: + line += f" (disagrees: {', '.join(r.disagreeing)})" + lines.append(line) + return "\n".join(lines) + + +def resolve_setting(name: str, rule: Rule, files: CheckpointFiles) -> Optional[Resolved]: + """One checkpoint setting. A required setting nothing declares is an error naming + the sources; a disagreement is rejected or logged as the rule says.""" + winner, disagreeing = _pick(name, rule, files, {}, {}) + if winner is None: + if rule.required: + labels = ", ".join(s.label for s in rule.sources) + raise ValueError(f"Cosmos3 config: cannot resolve {name}; none of {labels} is set.") + return None + if disagreeing and rule.conflict is not Conflict.IGNORE: + message = ( + f"Cosmos3 config: {name} = {winner.value!r} from {winner.source}, but " + + "; ".join(f"{label}={v!r}" for label, v in disagreeing) + ) + if rule.conflict is Conflict.ERROR: + raise ValueError(message) + logger.warning(message) + return _record(winner, disagreeing) + + +def resolve_checkpoint( + checkpoint_dir: Optional[str], pretrained_config: Any = None +) -> Cosmos3Config: + """Apply :data:`CHECKPOINT_RULES` to the checkpoint's files.""" + files = CheckpointFiles(checkpoint_dir, pretrained_config) + resolved: dict[str, Resolved] = {} + for name, rule in CHECKPOINT_RULES.items(): + r = resolve_setting(name, rule, files) + if r is not None: + resolved[name] = r + + def get(name: str) -> Any: + r = resolved.get(name) + return None if r is None else r.value + + sigmas = get("sampling.fixed_sigmas") + if sigmas is not None and ( + any(not 0.0 < float(s) <= 1.0 for s in sigmas) + or [float(s) for s in sigmas] != sorted((float(s) for s in sigmas), reverse=True) + ): + raise ValueError( + f"Cosmos3 config: fixed_step_sampler_config.t_list must decrease within (0, 1]: " + f"{list(sigmas)}." + ) + return Cosmos3Config( + architecture=Cosmos3ArchitectureSettings( + position_embedding_type=str(get("architecture.position_embedding_type")), + max_position_embeddings=int(get("architecture.max_position_embeddings")), + ), + vae=Cosmos3VAESettings( + temporal_compression=int(get("vae.temporal_compression")), + spatial_compression=int(get("vae.spatial_compression")), + ), + audio=Cosmos3AudioSettings(temporal_compression=int(get("audio.temporal_compression"))), + sampling=Cosmos3SamplingSettings( + fixed_sigmas=None if not sigmas else tuple(float(s) for s in sigmas), + native_flow_schedule=bool(get("sampling.native_flow_schedule")), + ), + prompt=Cosmos3PromptSettings(use_system_prompt=bool(get("prompt.use_system_prompt"))), + action=Cosmos3ActionSettings(policy=get("action.policy")), + resolved=resolved, + ) + + +ARCHITECTURE_ATTRIBUTES = { + "position_embedding_type": "architecture.position_embedding_type", + "max_position_embeddings": "architecture.max_position_embeddings", + "temporal_compression_factor": "vae.temporal_compression", + "temporal_compression_factor_sound": "audio.temporal_compression", +} +"""``pretrained_config`` attributes the transformer reads, and the setting each holds.""" + + +def fill_pretrained_config(pretrained_config: Any) -> Any: + """Set the transformer-construction attributes of ``pretrained_config`` that the + export omits, from :data:`CHECKPOINT_RULES`. Idempotent; returns the config. + + Runs before the pipeline loads its components, so only ``pretrained_config`` + itself is consulted. + """ + files = CheckpointFiles(None, pretrained_config) + for attr, name in ARCHITECTURE_ATTRIBUTES.items(): + if getattr(pretrained_config, attr, None) is not None: + continue + rule = CHECKPOINT_RULES[name] + transformer_only = replace( + rule, + sources=tuple( + s for s in rule.sources if not isinstance(s, Key) or s.file == TRANSFORMER + ), + ) + r = resolve_setting(name, transformer_only, files) + if r is not None: + setattr(pretrained_config, attr, r.value) + return pretrained_config + + +# ----------------------------------------------------------------------------- request + + +@dataclass(frozen=True) +class RequestContext: + """What a request resolves against: the checkpoint's facts and the request's mode.""" + + family: str + mode: str = "video" + """``"video"`` or ``"image"``; ignored for the table when ``action_mode`` is set.""" + action_mode: Optional[str] = None + distilled_steps: Optional[int] = None + """Length of the fixed schedule on a distilled checkpoint.""" + policy_defaults: Mapping[str, Any] = field(default_factory=dict) + """The normalized ``checkpoint.json`` policy block.""" + policy_active: bool = False + """The policy block applies: policy mode, and the request's domain is the block's.""" + domain_name: Optional[str] = None + domain_id: Optional[int] = None + checkpoint_use_system_prompt: bool = False + is_v2v: bool = False + is_transfer: bool = False + + @property + def table_mode(self) -> str: + return "action" if self.action_mode is not None else self.mode + + @property + def preset_key(self) -> Optional[str]: + return canonical_domain_preset_key(self.domain_name, self.domain_id) + + +class Given: + """The caller's value for the field being resolved; always the first source.""" + + label = "request" + + def lookup(self, ctx: RequestContext, name: str, given: Mapping[str, Any], values) -> Any: + value = given.get(name) + return _MISSING if value is None else value + + +@dataclass(frozen=True) +class Distilled: + """A fact of the distilled schedule: its step count, or the guidance it bakes in.""" + + what: str + + @property + def label(self) -> str: + return f"checkpoint:fixed_step_sampler_config ({self.what})" + + def lookup(self, ctx, name, given, values): + if ctx.distilled_steps is None: + return _MISSING + return ctx.distilled_steps if self.what == "steps" else DISTILLED_GUIDANCE_SCALE + + +@dataclass(frozen=True) +class Policy: + """A field of the checkpoint's policy block, when the block applies to the request.""" + + key: str + + @property + def label(self) -> str: + return f"checkpoint.json:policy.{self.key}" + + def lookup(self, ctx, name, given, values): + if not ctx.policy_active: + return _MISSING + value = ctx.policy_defaults.get(self.key) + return _MISSING if value is None else value + + +@dataclass(frozen=True) +class Context: + """A value the request context fixes: a workflow rule (transfer never prepends the + system prompt, video-to-video always does) or a checkpoint fact.""" + + label: str + pick: Callable[[RequestContext], Any] + + def lookup(self, ctx, name, given, values): + value = self.pick(ctx) + return _MISSING if value is None else value + + +class KnownWidth: + """The code's table of action widths per embodiment.""" + + label = "embodiment table (defaults)" + + def lookup(self, ctx, name, given, values): + width = resolve_raw_action_dim(domain_name=ctx.domain_name, domain_id=ctx.domain_id) + return _MISSING if width is None else width + + +@dataclass(frozen=True) +class Preset: + """A field of the request's domain preset.""" + + key: str + + @property + def label(self) -> str: + return f"domain preset.{self.key}" + + def lookup(self, ctx, name, given, values): + preset = COSMOS3_DOMAIN_PRESETS.get(ctx.preset_key) if ctx.preset_key else None + value = preset.get(self.key) if preset else None + return _MISSING if value is None else value + + +@dataclass(frozen=True) +class Table: + """The (family, mode) generation table; ``mode=None`` is the request's own mode.""" + + key: str + mode: Optional[str] = None + + @property + def label(self) -> str: + return f"{self.mode or 'mode'} table.{self.key}" + + def lookup(self, ctx, name, given, values): + table = COSMOS3_GENERATION_DEFAULTS.get((ctx.family, self.mode or ctx.table_mode)) + if not table or self.key not in table: + return _MISSING + return table[self.key] # an explicit ``None`` in a table is a value: the field is off + + +@dataclass(frozen=True) +class Derived: + """A value computed from fields resolved before it.""" + + label: str + compute: Callable[[Mapping[str, Any]], Any] + + def lookup(self, ctx, name, given, values): + value = self.compute(values) + return _MISSING if value is None else value + + +# Precedence ranks: request, then the checkpoint's facts, then the code's tables, then +# a default. ``check_precedence`` holds every rule to this order. +PRECEDENCE: tuple[tuple[type, ...], ...] = ( + (Given,), + (Key, Distilled, Policy, Context), + (KnownWidth, Preset, Table, Derived), + (Default,), +) + + +def check_precedence(rules: Mapping[str, Rule]) -> None: + """Raise unless every rule orders its sources request, checkpoint, tables, default.""" + rank = {cls: i for i, group in enumerate(PRECEDENCE) for cls in group} + for name, rule in rules.items(): + ranks = [rank[type(s)] for s in rule.sources] + if ranks != sorted(ranks): + raise ValueError(f"rule {name!r} breaks the request > checkpoint > default order") + + +GENERATION_RULES: dict[str, Rule] = { + "num_inference_steps": Rule( + ( + Given(), + Distilled("steps"), + Policy("num_inference_steps"), + Table("num_inference_steps"), + Table("num_inference_steps", "video"), + ) + ), + "guidance_scale": Rule( + ( + Given(), + Distilled("guidance"), + Policy("guidance_scale"), + Table("guidance_scale"), + Table("guidance_scale", "video"), + ) + ), + "guidance_interval": Rule( + ( + Given(), + Policy("guidance_interval"), + Table("guidance_interval"), + Table("guidance_interval", "video"), + ), + required=False, + ), + "max_sequence_length": Rule( + (Given(), Table("max_sequence_length"), Table("max_sequence_length", "video")) + ), + "height": Rule((Given(), Table("height"), Table("height", "video"))), + "width": Rule((Given(), Table("width"), Table("width", "video"))), + "num_frames": Rule((Given(), Table("num_frames"), Table("num_frames", "video"))), + "frame_rate": Rule((Given(), Table("frame_rate"), Table("frame_rate", "video"))), + "use_system_prompt": Rule( + ( + Given(), + Context("transfer (never)", lambda ctx: False if ctx.is_transfer else None), + Context("video-to-video (always)", lambda ctx: True if ctx.is_v2v else None), + Context("checkpoint:use_system_prompt", lambda ctx: ctx.checkpoint_use_system_prompt), + ) + ), +} + +ACTION_RULES: dict[str, Rule] = { + "raw_action_dim": Rule( + (Given(), Policy("raw_action_dim"), KnownWidth()), + conflict=Conflict.WARN, + required=False, + ), + "action_chunk_size": Rule( + ( + Given(), + Policy("action_chunk_size"), + Preset("action_chunk_size"), + Default(COSMOS3_ACTION_PARAMS["action_chunk_size"], "generic action horizon"), + ), + conflict=Conflict.WARN, + ), + "action_resolution": Rule( + ( + Given(), + Policy("action_resolution"), + Preset("action_resolution"), + Default(480, "the bucket the action models were trained at"), + ), + conflict=Conflict.WARN, + ), + "frame_rate": Rule( + ( + Given(), + Policy("frame_rate"), + Preset("frame_rate"), + Default(COSMOS3_ACTION_PARAMS["frame_rate"], "generic action frame rate"), + ), + conflict=Conflict.WARN, + ), + "action_fps": Rule( + (Given(), Policy("action_fps"), Derived("frame_rate", lambda v: v["frame_rate"])) + ), + "use_state": Rule( + (Given(), Policy("use_state"), Default(False, "policies without a state row")), + conflict=Conflict.WARN, + ), + "num_frames": Rule( + (Derived("action_chunk_size + 1", lambda v: int(v["action_chunk_size"]) + 1),) + ), +} + +for _rules in (CHECKPOINT_RULES, GENERATION_RULES, ACTION_RULES): + check_precedence(_rules) + + +@dataclass +class ResolvedRequest: + values: dict[str, Resolved] = field(default_factory=dict) + warnings: list[str] = field(default_factory=list) + + def __getitem__(self, name: str) -> Any: + return self.values[name].value + + def source_of(self, name: str) -> str: + return self.values[name].source + + +def _resolve_request( + rules: Mapping[str, Rule], + ctx: RequestContext, + given: Mapping[str, Any], + domain_label: str, +) -> ResolvedRequest: + """Request fields: a caller's value that differs from the recommendation below it is + reported in ``warnings`` where the rule says ``WARN``.""" + out = ResolvedRequest() + plain: dict[str, Any] = {} + for name, rule in rules.items(): + winner, disagreeing = _pick(name, rule, ctx, given, plain) + if winner is None: + if rule.required: + raise ValueError(f"Cosmos3 request: cannot resolve {name} for {domain_label}.") + out.values[name] = Resolved(None, "unresolved") + plain[name] = None + continue + if disagreeing and winner.source == Given.label and rule.conflict is Conflict.WARN: + out.warnings.append( + f"Cosmos3 {name}={winner.value} differs from recommended " + f"{' / '.join(str(v) for _, v in disagreeing)} for domain {domain_label}." + ) + out.values[name] = _record(winner, disagreeing) + plain[name] = winner.value + return out + + +def resolve_generation(ctx: RequestContext, **given: Any) -> ResolvedRequest: + """Resolve the generation fields in ``given`` (``None`` = unset).""" + rules = {name: GENERATION_RULES[name] for name in given} + return _resolve_request(rules, ctx, given, repr(ctx.table_mode)) + + +def policy_domain( + policy_defaults: Mapping[str, Any], + domain_name: Optional[str], + domain_id: Optional[int], +) -> tuple[Optional[str], bool, list[str]]: + """Reconcile the request's domain with the checkpoint policy block's: returns the + domain name to use, whether the block applies, and a warning when the request + names a different domain (the block's defaults are then not applied).""" + checkpoint_domain = policy_defaults.get("domain_name") + warnings: list[str] = [] + if checkpoint_domain is None: + return domain_name, False, warnings + checkpoint_domain = str(checkpoint_domain).strip().lower() + explicit = ( + str(domain_name).strip().lower() if domain_name and str(domain_name).strip() else None + ) + checkpoint_domain_id = EMBODIMENT_TO_DOMAIN_ID.get(checkpoint_domain) + id_matches = domain_id is not None and checkpoint_domain_id == int(domain_id) + if explicit is None and (domain_id is None or id_matches): + return checkpoint_domain, True, warnings + if explicit == checkpoint_domain: + return domain_name, True, warnings + warnings.append( + f"Cosmos3 domain_name={domain_name!r}, domain_id={domain_id!r} overrides " + f"checkpoint policy domain_name={checkpoint_domain!r}; checkpoint domain " + "defaults will not be applied." + ) + return domain_name, False, warnings + + +def resolve_action(ctx: RequestContext, **given: Any) -> ResolvedRequest: + """Resolve the action fields; ``ctx.domain_name`` / ``domain_id`` already reconciled + with the policy block (:func:`policy_domain`).""" + domain_label = ( + repr(ctx.preset_key) + if ctx.preset_key + else (f"domain_name={ctx.domain_name!r}, domain_id={ctx.domain_id!r}") + ) + resolved = _resolve_request(ACTION_RULES, ctx, given, domain_label) + domain_requested = bool(ctx.domain_name and str(ctx.domain_name).strip()) or ( + ctx.domain_id is not None and str(ctx.domain_id).strip() not in {"", "0"} + ) + if domain_requested and ctx.preset_key is None: + resolved.warnings.insert( + 0, + "Cosmos3 action domain preset was not found for " + f"domain_name={ctx.domain_name!r}, domain_id={ctx.domain_id!r}; " + "using generic action defaults for omitted fields.", + ) + if domain_requested and resolved["raw_action_dim"] is None: + resolved.warnings.append( + "Cosmos3 has no canonical action width for " + f"domain_name={ctx.domain_name!r}, domain_id={ctx.domain_id!r}; " + "pass raw_action_dim explicitly for policy/inverse_dynamics." + ) + resolved.values["action_resolution"] = Resolved( + normalize_action_resolution(resolved["action_resolution"]), + resolved.source_of("action_resolution"), + ) + if resolved["raw_action_dim"] is not None and int(resolved["raw_action_dim"]) <= 0: + raise ValueError( + f"Cosmos3 raw_action_dim must be positive, got {resolved['raw_action_dim']}." + ) + if int(resolved["action_chunk_size"]) <= 0: + raise ValueError( + f"Cosmos3 action_chunk_size must be positive, got {resolved['action_chunk_size']}." + ) + if float(resolved["frame_rate"]) <= 0.0: + raise ValueError(f"Cosmos3 frame_rate must be positive, got {resolved['frame_rate']}.") + if float(resolved["action_fps"]) <= 0.0: + raise ValueError(f"Cosmos3 action_fps must be positive, got {resolved['action_fps']}.") + return resolved + + +def resolve_domain_action_config( + *, + domain_name: str | None = None, + domain_id: str | int | None = None, + raw_action_dim: int | None = None, + action_chunk_size: int | None = None, + action_resolution: int | None = None, + frame_rate: float | None = None, + action_fps: float | None = None, + use_state: bool | None = None, + checkpoint_policy_defaults: Mapping[str, Any] | None = None, +) -> dict[str, Any]: + """The action fields a request runs with, by :data:`ACTION_RULES`: the caller's + value, else the checkpoint policy block when its domain is the request's, else the + code's width table, else the domain preset, else the generic action defaults. + ``num_frames`` is always ``action_chunk_size + 1``. A caller's value that differs + from the recommendation is reported in ``warnings``. + """ + policy_defaults = dict(checkpoint_policy_defaults or {}) + domain_name, policy_active, warnings = policy_domain(policy_defaults, domain_name, domain_id) + ctx = RequestContext( + family="", + action_mode="action", + policy_defaults=policy_defaults, + policy_active=policy_active, + domain_name=domain_name, + domain_id=None if domain_id is None else int(domain_id), + ) + resolved = resolve_action( + ctx, + raw_action_dim=raw_action_dim, + action_chunk_size=action_chunk_size, + action_resolution=action_resolution, + frame_rate=frame_rate, + action_fps=action_fps, + use_state=use_state, + ) + width = resolved["raw_action_dim"] + return { + "domain_name": domain_name, + "checkpoint_policy_domain_matches": policy_active, + "raw_action_dim": None if width is None else int(width), + "action_chunk_size": int(resolved["action_chunk_size"]), + "action_resolution": resolved["action_resolution"], + "frame_rate": float(resolved["frame_rate"]), + "action_fps": float(resolved["action_fps"]), + "use_state": bool(resolved["use_state"]), + "num_frames": int(resolved["num_frames"]), + "preset_key": ctx.preset_key, + "warnings": warnings + resolved.warnings, + "sources": {name: r.source for name, r in resolved.values.items()}, + } diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index c62b48dd16d6..02834b441954 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -68,7 +68,6 @@ def tqdm(iterable, **kwargs): resolve_action_size, resolve_domain_id, ) -from .config_resolver import Cosmos3ConfigResolver from .defaults import ( COSMOS3_720P_PARAMS, COSMOS3_ASPECT_RATIO_BUCKETS, @@ -88,7 +87,12 @@ def tqdm(iterable, **kwargs): ) from .guardrails import check_video_safety, download_guardrail_checkpoint from .negative_prompt import COSMOS3_VIDEO_NEGATIVE_PROMPT -from .request_resolver import RequestContext, resolve_domain_action_config, resolve_generation +from .parameter_resolver import ( + RequestContext, + resolve_checkpoint, + resolve_domain_action_config, + resolve_generation, +) from .sampling import DISTILLED_GUIDANCE_SCALE, Cosmos3SamplingPolicy, load_scheduler from .sound_tokenizer import LatentAutoEncoderV2 from .transfer import ( @@ -616,7 +620,7 @@ def load_standard_components( # The transformer config is read from the checkpoint directory, not from the # loaded pretrained config, whose omitted keys the transformer has already filled. - self.cosmos3_config = Cosmos3ConfigResolver(checkpoint_dir).resolve() + self.cosmos3_config = resolve_checkpoint(checkpoint_dir) logger.info(self.cosmos3_config.describe()) cfg = self.cosmos3_config self.default_use_system_prompt = cfg.prompt.use_system_prompt diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/request_resolver.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/request_resolver.py deleted file mode 100644 index f274963144b8..000000000000 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/request_resolver.py +++ /dev/null @@ -1,477 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Per-request resolution of Cosmos3 generation and action settings. - -The engine-level :mod:`config_resolver` settles what the checkpoint says. This module -settles what one request runs with: for every field, :data:`GENERATION_RULES` and -:data:`ACTION_RULES` name the sources in priority order, the caller's value first, -then the checkpoint's facts (the distilled schedule, the policy block), then the -per-mode and per-domain tables, then a default. The first source holding a value -wins; where a rule says so, a recommendation that disagrees with the caller's value -is reported as a warning, and every resolved value remembers its source. -""" - -from __future__ import annotations - -from collections.abc import Callable, Mapping -from dataclasses import dataclass, field -from typing import Any, Optional - -from .config_resolver import Conflict, Default, Resolved, same_value -from .defaults import ( - COSMOS3_ACTION_PARAMS, - COSMOS3_DOMAIN_PRESETS, - COSMOS3_GENERATION_DEFAULTS, - EMBODIMENT_TO_DOMAIN_ID, - canonical_domain_preset_key, - normalize_action_resolution, - resolve_raw_action_dim, -) -from .sampling import DISTILLED_GUIDANCE_SCALE - -_MISSING = object() - - -@dataclass(frozen=True) -class RequestContext: - """What a request resolves against: the checkpoint's facts and the request's mode.""" - - family: str - mode: str = "video" - """``"video"`` or ``"image"``; ignored for the table when ``action_mode`` is set.""" - action_mode: Optional[str] = None - distilled_steps: Optional[int] = None - """Length of the fixed schedule on a distilled checkpoint.""" - policy_defaults: Mapping[str, Any] = field(default_factory=dict) - """The normalized ``checkpoint.json`` policy block.""" - policy_active: bool = False - """The policy block applies: policy mode, and the request's domain is the block's.""" - domain_name: Optional[str] = None - domain_id: Optional[int] = None - checkpoint_use_system_prompt: bool = False - is_v2v: bool = False - is_transfer: bool = False - - @property - def table_mode(self) -> str: - return "action" if self.action_mode is not None else self.mode - - @property - def preset_key(self) -> Optional[str]: - return canonical_domain_preset_key(self.domain_name, self.domain_id) - - -# ----------------------------------------------------------------------------- sources - - -class Given: - """The caller's value for the field being resolved.""" - - label = "request" - - def lookup(self, ctx: RequestContext, field_name: str, given: Mapping[str, Any], values) -> Any: - value = given.get(field_name) - return _MISSING if value is None else value - - -@dataclass(frozen=True) -class Distilled: - """A fact of the distilled schedule: its step count, or the guidance it bakes in.""" - - what: str - - @property - def label(self) -> str: - return f"checkpoint:fixed_step_sampler_config ({self.what})" - - def lookup(self, ctx, field_name, given, values): - if ctx.distilled_steps is None: - return _MISSING - return ctx.distilled_steps if self.what == "steps" else DISTILLED_GUIDANCE_SCALE - - -@dataclass(frozen=True) -class Policy: - key: str - - @property - def label(self) -> str: - return f"checkpoint.json:policy.{self.key}" - - def lookup(self, ctx, field_name, given, values): - if not ctx.policy_active: - return _MISSING - value = ctx.policy_defaults.get(self.key) - return _MISSING if value is None else value - - -class KnownWidth: - """The code's table of action widths per embodiment.""" - - label = "embodiment table (defaults)" - - def lookup(self, ctx, field_name, given, values): - width = resolve_raw_action_dim(domain_name=ctx.domain_name, domain_id=ctx.domain_id) - return _MISSING if width is None else width - - -@dataclass(frozen=True) -class Preset: - key: str - - @property - def label(self) -> str: - return f"domain preset.{self.key}" - - def lookup(self, ctx, field_name, given, values): - preset = COSMOS3_DOMAIN_PRESETS.get(ctx.preset_key) if ctx.preset_key else None - value = preset.get(self.key) if preset else None - return _MISSING if value is None else value - - -@dataclass(frozen=True) -class Table: - """The (family, mode) generation table; ``mode=None`` is the request's own mode.""" - - key: str - mode: Optional[str] = None - - @property - def label(self) -> str: - return f"{self.mode or 'mode'} table.{self.key}" - - def lookup(self, ctx, field_name, given, values): - table = COSMOS3_GENERATION_DEFAULTS.get((ctx.family, self.mode or ctx.table_mode)) - if not table or self.key not in table: - return _MISSING - return table[self.key] # an explicit ``None`` in a table is a value: the field is off - - -@dataclass(frozen=True) -class Context: - """A value the request context fixes: a workflow rule (transfer never prepends the - system prompt, video-to-video always does) or a checkpoint fact.""" - - label: str - pick: Callable[[RequestContext], Any] - - def lookup(self, ctx, field_name, given, values): - value = self.pick(ctx) - return _MISSING if value is None else value - - -@dataclass(frozen=True) -class Derived: - """A value computed from fields resolved before it.""" - - label: str - compute: Callable[[Mapping[str, Any]], Any] - - def lookup(self, ctx, field_name, given, values): - value = self.compute(values) - return _MISSING if value is None else value - - -@dataclass(frozen=True) -class RequestRule: - sources: tuple - conflict: Conflict = Conflict.IGNORE - """``WARN``: a caller's value that differs from the first recommendation below it - is reported in the result's warnings.""" - required: bool = True - - -# ------------------------------------------------------------------------------- rules - -GENERATION_RULES: dict[str, RequestRule] = { - "num_inference_steps": RequestRule( - ( - Given(), - Distilled("steps"), - Policy("num_inference_steps"), - Table("num_inference_steps"), - Table("num_inference_steps", "video"), - ) - ), - "guidance_scale": RequestRule( - ( - Given(), - Distilled("guidance"), - Policy("guidance_scale"), - Table("guidance_scale"), - Table("guidance_scale", "video"), - ) - ), - "guidance_interval": RequestRule( - ( - Given(), - Policy("guidance_interval"), - Table("guidance_interval"), - Table("guidance_interval", "video"), - ), - required=False, - ), - "max_sequence_length": RequestRule( - (Given(), Table("max_sequence_length"), Table("max_sequence_length", "video")) - ), - "height": RequestRule((Given(), Table("height"), Table("height", "video"))), - "width": RequestRule((Given(), Table("width"), Table("width", "video"))), - "num_frames": RequestRule((Given(), Table("num_frames"), Table("num_frames", "video"))), - "frame_rate": RequestRule((Given(), Table("frame_rate"), Table("frame_rate", "video"))), - "use_system_prompt": RequestRule( - ( - Given(), - Context("transfer (never)", lambda ctx: False if ctx.is_transfer else None), - Context("video-to-video (always)", lambda ctx: True if ctx.is_v2v else None), - Context("checkpoint:use_system_prompt", lambda ctx: ctx.checkpoint_use_system_prompt), - ) - ), -} - -ACTION_RULES: dict[str, RequestRule] = { - "raw_action_dim": RequestRule( - (Given(), Policy("raw_action_dim"), KnownWidth()), - conflict=Conflict.WARN, - required=False, - ), - "action_chunk_size": RequestRule( - ( - Given(), - Policy("action_chunk_size"), - Preset("action_chunk_size"), - Default(COSMOS3_ACTION_PARAMS["action_chunk_size"], "generic action horizon"), - ), - conflict=Conflict.WARN, - ), - "action_resolution": RequestRule( - ( - Given(), - Policy("action_resolution"), - Preset("action_resolution"), - Default(480, "the bucket the action models were trained at"), - ), - conflict=Conflict.WARN, - ), - "frame_rate": RequestRule( - ( - Given(), - Policy("frame_rate"), - Preset("frame_rate"), - Default(COSMOS3_ACTION_PARAMS["frame_rate"], "generic action frame rate"), - ), - conflict=Conflict.WARN, - ), - "action_fps": RequestRule( - (Given(), Policy("action_fps"), Derived("frame_rate", lambda v: v["frame_rate"])) - ), - "use_state": RequestRule( - (Given(), Policy("use_state"), Default(False, "policies without a state row")), - conflict=Conflict.WARN, - ), - "num_frames": RequestRule( - (Derived("action_chunk_size + 1", lambda v: int(v["action_chunk_size"]) + 1),) - ), -} - - -@dataclass -class ResolvedRequest: - values: dict[str, Resolved] = field(default_factory=dict) - warnings: list[str] = field(default_factory=list) - - def __getitem__(self, name: str) -> Any: - return self.values[name].value - - def source_of(self, name: str) -> str: - return self.values[name].source - - -def _resolve( - rules: Mapping[str, RequestRule], - ctx: RequestContext, - given: Mapping[str, Any], - domain_label: str, -) -> ResolvedRequest: - out = ResolvedRequest() - plain: dict[str, Any] = {} - for name, rule in rules.items(): - winner: Optional[Resolved] = None - disagreeing: list[tuple[str, Any]] = [] - for source in rule.sources: - if winner is not None and isinstance(source, Default): - break - if isinstance(source, Default): - value = source.value - else: - value = source.lookup(ctx, name, given, plain) - if value is _MISSING: - continue - if winner is None: - winner = Resolved(value, source.label) - continue - # Only the first recommendation below the caller's value is compared. - if not same_value(value, winner.value): - disagreeing.append((source.label, value)) - break - if winner is None: - if rule.required: - raise ValueError(f"Cosmos3 request: cannot resolve {name} for {domain_label}.") - out.values[name] = Resolved(None, "unresolved") - plain[name] = None - continue - if disagreeing and winner.source == Given.label and rule.conflict is Conflict.WARN: - out.warnings.append( - f"Cosmos3 {name}={winner.value} differs from recommended " - f"{' / '.join(str(v) for _, v in disagreeing)} for domain {domain_label}." - ) - out.values[name] = Resolved( - winner.value, winner.source, tuple(f"{label}={v!r}" for label, v in disagreeing) - ) - plain[name] = winner.value - return out - - -def resolve_generation(ctx: RequestContext, **given: Any) -> ResolvedRequest: - """Resolve the generation fields in ``given`` (``None`` = unset).""" - rules = {name: GENERATION_RULES[name] for name in given} - return _resolve(rules, ctx, given, repr(ctx.table_mode)) - - -def policy_domain( - policy_defaults: Mapping[str, Any], - domain_name: Optional[str], - domain_id: Optional[int], -) -> tuple[Optional[str], bool, list[str]]: - """Reconcile the request's domain with the checkpoint policy block's: returns the - domain name to use, whether the block applies, and a warning when the request - names a different domain (the block's defaults are then not applied).""" - checkpoint_domain = policy_defaults.get("domain_name") - warnings: list[str] = [] - if checkpoint_domain is None: - return domain_name, False, warnings - checkpoint_domain = str(checkpoint_domain).strip().lower() - explicit = ( - str(domain_name).strip().lower() if domain_name and str(domain_name).strip() else None - ) - checkpoint_domain_id = EMBODIMENT_TO_DOMAIN_ID.get(checkpoint_domain) - id_matches = domain_id is not None and checkpoint_domain_id == int(domain_id) - if explicit is None and (domain_id is None or id_matches): - return checkpoint_domain, True, warnings - if explicit == checkpoint_domain: - return domain_name, True, warnings - warnings.append( - f"Cosmos3 domain_name={domain_name!r}, domain_id={domain_id!r} overrides " - f"checkpoint policy domain_name={checkpoint_domain!r}; checkpoint domain " - "defaults will not be applied." - ) - return domain_name, False, warnings - - -def resolve_action(ctx: RequestContext, **given: Any) -> ResolvedRequest: - """Resolve the action fields; ``ctx.domain_name`` / ``domain_id`` already reconciled - with the policy block (:func:`policy_domain`).""" - domain_label = ( - repr(ctx.preset_key) - if ctx.preset_key - else (f"domain_name={ctx.domain_name!r}, domain_id={ctx.domain_id!r}") - ) - resolved = _resolve(ACTION_RULES, ctx, given, domain_label) - domain_requested = bool(ctx.domain_name and str(ctx.domain_name).strip()) or ( - ctx.domain_id is not None and str(ctx.domain_id).strip() not in {"", "0"} - ) - if domain_requested and ctx.preset_key is None: - resolved.warnings.insert( - 0, - "Cosmos3 action domain preset was not found for " - f"domain_name={ctx.domain_name!r}, domain_id={ctx.domain_id!r}; " - "using generic action defaults for omitted fields.", - ) - if domain_requested and resolved["raw_action_dim"] is None: - resolved.warnings.append( - "Cosmos3 has no canonical action width for " - f"domain_name={ctx.domain_name!r}, domain_id={ctx.domain_id!r}; " - "pass raw_action_dim explicitly for policy/inverse_dynamics." - ) - resolved.values["action_resolution"] = Resolved( - normalize_action_resolution(resolved["action_resolution"]), - resolved.source_of("action_resolution"), - ) - if resolved["raw_action_dim"] is not None and int(resolved["raw_action_dim"]) <= 0: - raise ValueError( - f"Cosmos3 raw_action_dim must be positive, got {resolved['raw_action_dim']}." - ) - if int(resolved["action_chunk_size"]) <= 0: - raise ValueError( - f"Cosmos3 action_chunk_size must be positive, got {resolved['action_chunk_size']}." - ) - if float(resolved["frame_rate"]) <= 0.0: - raise ValueError(f"Cosmos3 frame_rate must be positive, got {resolved['frame_rate']}.") - if float(resolved["action_fps"]) <= 0.0: - raise ValueError(f"Cosmos3 action_fps must be positive, got {resolved['action_fps']}.") - return resolved - - -def resolve_domain_action_config( - *, - domain_name: str | None = None, - domain_id: str | int | None = None, - raw_action_dim: int | None = None, - action_chunk_size: int | None = None, - action_resolution: int | None = None, - frame_rate: float | None = None, - action_fps: float | None = None, - use_state: bool | None = None, - checkpoint_policy_defaults: Mapping[str, Any] | None = None, -) -> dict[str, Any]: - """The action fields a request runs with, by :data:`ACTION_RULES`: the caller's - value, else the checkpoint policy block when its domain is the request's, else the - code's width table, else the domain preset, else the generic action defaults. - ``num_frames`` is always ``action_chunk_size + 1``. A caller's value that differs - from the recommendation is reported in ``warnings``. - """ - policy_defaults = dict(checkpoint_policy_defaults or {}) - domain_name, policy_active, warnings = policy_domain(policy_defaults, domain_name, domain_id) - ctx = RequestContext( - family="", - action_mode="action", - policy_defaults=policy_defaults, - policy_active=policy_active, - domain_name=domain_name, - domain_id=None if domain_id is None else int(domain_id), - ) - resolved = resolve_action( - ctx, - raw_action_dim=raw_action_dim, - action_chunk_size=action_chunk_size, - action_resolution=action_resolution, - frame_rate=frame_rate, - action_fps=action_fps, - use_state=use_state, - ) - width = resolved["raw_action_dim"] - return { - "domain_name": domain_name, - "checkpoint_policy_domain_matches": policy_active, - "raw_action_dim": None if width is None else int(width), - "action_chunk_size": int(resolved["action_chunk_size"]), - "action_resolution": resolved["action_resolution"], - "frame_rate": float(resolved["frame_rate"]), - "action_fps": float(resolved["action_fps"]), - "use_state": bool(resolved["use_state"]), - "num_frames": int(resolved["num_frames"]), - "preset_key": ctx.preset_key, - "warnings": warnings + resolved.warnings, - "sources": {name: r.source for name, r in resolved.values.items()}, - } diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py index c1f9e0e66502..fb51c248693b 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -43,7 +43,7 @@ from tensorrt_llm.logger import logger from tensorrt_llm.models.modeling_utils import QuantConfig -from .config_resolver import fill_pretrained_config +from .parameter_resolver import fill_pretrained_config def uses_static_fp8(model_config: DiffusionModelConfig) -> bool: diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 7882fb418edc..c1d0e1ef37f0 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -173,8 +173,7 @@ l0_cpu: - unittest/_torch/visual_gen/test_attention_integration.py - unittest/_torch/visual_gen/test_attention_importable.py - unittest/_torch/visual_gen/test_cache_dit.py - - unittest/_torch/visual_gen/test_cosmos3_config_resolver.py - - unittest/_torch/visual_gen/test_cosmos3_request_resolver.py + - unittest/_torch/visual_gen/test_cosmos3_parameter_resolver.py - unittest/_torch/visual_gen/test_executor_lifecycle.py - unittest/_torch/visual_gen/test_flux2_load_weights.py - unittest/_torch/visual_gen/test_flux_infer.py diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py index 141e9f7ce126..f8a8ec19cfa0 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py @@ -37,7 +37,7 @@ resolve_checkpoint_policy_defaults, resolve_raw_action_dim, ) -from tensorrt_llm._torch.visual_gen.models.cosmos3.request_resolver import ( +from tensorrt_llm._torch.visual_gen.models.cosmos3.parameter_resolver import ( resolve_domain_action_config, ) from tensorrt_llm._torch.visual_gen.models.cosmos3.transformer_cosmos3 import ( diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_config_resolver.py b/tests/unittest/_torch/visual_gen/test_cosmos3_config_resolver.py deleted file mode 100644 index b1e8e6393670..000000000000 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_config_resolver.py +++ /dev/null @@ -1,203 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""The Cosmos3 precedence table, on synthetic files and on the published checkpoints.""" - -import json -from pathlib import Path -from types import SimpleNamespace - -import pytest - -from tensorrt_llm._torch.visual_gen.models.cosmos3.config_resolver import ( - CHECKPOINT, - CONFIG, - MODEL_INDEX, - SCHEDULER, - TRANSFORMER, - VAE, - Conflict, - Cosmos3ConfigResolver, - Default, - Key, - Rule, - fill_pretrained_config, -) - -pytestmark = pytest.mark.cpu_only - -FIXTURE = Path(__file__).parent / "test_data" / "cosmos3_checkpoint_configs.json" -T_LIST = [1.0, 0.9375, 0.8333333333333334, 0.625] - - -def _write(checkpoint_dir: Path, files: dict) -> Path: - for name, doc in files.items(): - path = checkpoint_dir / name - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(json.dumps(doc)) - return checkpoint_dir - - -def _resolver(tmp_path: Path, files: dict) -> Cosmos3ConfigResolver: - _write(tmp_path, files) - return Cosmos3ConfigResolver(str(tmp_path), files.get(TRANSFORMER, {})) - - -class TestRule: - def test_first_present_source_wins(self, tmp_path): - r = _resolver(tmp_path, {CONFIG: {"a": 1}, MODEL_INDEX: {"a": 1}}) - rule = Rule((Key(CONFIG, "a"), Key(MODEL_INDEX, "a"), Default(7, "x"))) - assert r.resolve_one("a", rule).value == 1 - assert r.resolve_one("a", rule).source == "config.json:a" - - def test_falls_through_missing_and_null_sources(self, tmp_path): - r = _resolver(tmp_path, {CONFIG: {"a": None}, MODEL_INDEX: {"a": 2}}) - rule = Rule((Key(CONFIG, "a"), Key(MODEL_INDEX, "a"), Key(CHECKPOINT, "a"))) - assert r.resolve_one("a", rule).value == 2 - - def test_default_only_when_nothing_is_declared(self, tmp_path): - r = _resolver(tmp_path, {}) - rule = Rule((Key(CONFIG, "a"), Default(7, "x"))) - resolved = r.resolve_one("a", rule) - assert (resolved.value, resolved.source) == (7, "default (x)") - - def test_unresolvable_required_setting_is_an_error(self, tmp_path): - r = _resolver(tmp_path, {}) - with pytest.raises(ValueError, match="cannot resolve a"): - r.resolve_one("a", Rule((Key(CONFIG, "a"),))) - assert r.resolve_one("a", Rule((Key(CONFIG, "a"),), required=False)) is None - - def test_conflict_policies(self, tmp_path): - r = _resolver(tmp_path, {CONFIG: {"a": 1}, MODEL_INDEX: {"a": 2}}) - sources = (Key(CONFIG, "a"), Key(MODEL_INDEX, "a")) - with pytest.raises(ValueError, match="model_index.json:a=2"): - r.resolve_one("a", Rule(sources, conflict=Conflict.ERROR)) - warned = r.resolve_one("a", Rule(sources, conflict=Conflict.WARN)) - assert warned.value == 1 and warned.disagreeing == ("model_index.json:a=2",) - - def test_float_lists_compare_within_tolerance(self, tmp_path): - r = _resolver( - tmp_path, {CONFIG: {"a": [0.5, 1 / 3]}, MODEL_INDEX: {"a": [0.5, 0.3333333333]}} - ) - rule = Rule((Key(CONFIG, "a"), Key(MODEL_INDEX, "a")), conflict=Conflict.ERROR) - assert r.resolve_one("a", rule).disagreeing == () - - -def _files(): - """A distilled checkpoint whose training config and diffusers index disagree on - the system prompt.""" - return { - TRANSFORMER: {"temporal_compression_factor": 4}, - SCHEDULER: {"fixed_step_sampler_config": {"t_list": T_LIST}}, - CONFIG: {"model": {"config": {"vlm_config": {"use_system_prompt": False}}}}, - MODEL_INDEX: {"default_use_system_prompt": True}, - VAE: {"scale_factor_temporal": 4, "scale_factor_spatial": 16}, - } - - -class TestResolve: - def test_training_config_beats_the_diffusers_index(self, tmp_path): - cfg = _resolver(tmp_path, _files()).resolve() - assert cfg.prompt.use_system_prompt is False - assert cfg.source_of("prompt.use_system_prompt") == ( - "config.json:model.config.vlm_config.use_system_prompt" - ) - assert cfg.resolved["prompt.use_system_prompt"].disagreeing == ( - "model_index.json:default_use_system_prompt=True", - ) - - def test_sections(self, tmp_path): - cfg = _resolver(tmp_path, _files()).resolve() - assert (cfg.vae.temporal_compression, cfg.vae.spatial_compression) == (4, 16) - assert cfg.sampling.fixed_sigmas == tuple(T_LIST) - assert cfg.sampling.native_flow_schedule is False - assert cfg.source_of("sampling.native_flow_schedule").startswith("default (") - assert cfg.action.policy is None - - def test_vae_compression_copies_must_agree(self, tmp_path): - files = _files() - files[TRANSFORMER]["temporal_compression_factor"] = 8 - with pytest.raises(ValueError, match="vae.temporal_compression"): - _resolver(tmp_path, files).resolve() - - def test_schedule_must_decrease_within_unit_interval(self, tmp_path): - files = _files() - files[SCHEDULER]["fixed_step_sampler_config"]["t_list"] = [0.5, 1.0] - with pytest.raises(ValueError, match="t_list must decrease"): - _resolver(tmp_path, files).resolve() - - def test_base_checkpoint_has_no_schedule(self, tmp_path): - files = {k: v for k, v in _files().items() if k != SCHEDULER} - cfg = _resolver(tmp_path, files).resolve() - assert cfg.sampling.fixed_sigmas is None - - def test_system_prompt_must_be_declared(self, tmp_path): - files = _files() - del files[CONFIG]["model"]["config"]["vlm_config"]["use_system_prompt"] - del files[MODEL_INDEX]["default_use_system_prompt"] - with pytest.raises(ValueError, match="cannot resolve prompt.use_system_prompt"): - _resolver(tmp_path, files).resolve() - - def test_spatial_compression_must_be_declared(self, tmp_path): - files = _files() - del files[VAE]["scale_factor_spatial"] - with pytest.raises(ValueError, match="cannot resolve vae.spatial_compression"): - _resolver(tmp_path, files).resolve() - - -def test_fill_pretrained_config_sets_only_missing_attributes(): - declared = SimpleNamespace(temporal_compression_factor=8) - fill_pretrained_config(declared) - assert declared.temporal_compression_factor == 8 - assert declared.position_embedding_type == "unified_3d_mrope" - assert declared.max_position_embeddings == 262144 - assert declared.temporal_compression_factor_sound == 1 - - -# Disagreements between a checkpoint's own files that the table expects and resolves. -EXPECTED_DISAGREEMENTS = { - "nvidia/Cosmos3-Nano-Sim-Bimanual": {"prompt.use_system_prompt"}, -} - -EXPECTED = { - # repo: (use_system_prompt, native_flow_schedule, distilled) - "nvidia/Cosmos3-Nano": (False, False, False), - "nvidia/Cosmos3-Super": (False, False, False), - "nvidia/Cosmos3-Edge": (False, True, False), - "nvidia/Cosmos3-Super-Image2Video": (False, False, False), - "nvidia/Cosmos3-Super-Image2Video-4Step": (False, False, True), - "nvidia/Cosmos3-Super-Text2Image": (False, False, False), - "nvidia/Cosmos3-Super-Text2Image-4Step": (False, False, True), - "nvidia/Cosmos3-Nano-Policy-DROID": (False, False, False), - "nvidia/Cosmos3-Edge-Policy-DROID": (False, True, False), - "nvidia/Cosmos3-Nano-Sim-Bimanual": (False, False, True), -} - - -@pytest.mark.parametrize("repo", sorted(EXPECTED)) -def test_published_checkpoint_resolves(tmp_path, repo): - files = json.loads(FIXTURE.read_text())[repo] - cfg = _resolver(tmp_path, files).resolve() - system_prompt, native, distilled = EXPECTED[repo] - assert cfg.prompt.use_system_prompt is system_prompt - assert cfg.sampling.native_flow_schedule is native - assert (cfg.sampling.fixed_sigmas is not None) is distilled - assert (cfg.vae.temporal_compression, cfg.vae.spatial_compression) == (4, 16) - assert cfg.architecture.position_embedding_type == "unified_3d_mrope" - disagreeing = {name for name, r in cfg.resolved.items() if r.disagreeing} - assert disagreeing == EXPECTED_DISAGREEMENTS.get(repo, set()) - - -def test_fixture_covers_every_expected_checkpoint(): - assert set(json.loads(FIXTURE.read_text())) == set(EXPECTED) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_example_prompts.py b/tests/unittest/_torch/visual_gen/test_cosmos3_example_prompts.py index e8011bbd7bca..0b44a2bb7868 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_example_prompts.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_example_prompts.py @@ -20,7 +20,7 @@ find_closest_target_size, resolve_checkpoint_policy_defaults, ) -from tensorrt_llm._torch.visual_gen.models.cosmos3.request_resolver import ( +from tensorrt_llm._torch.visual_gen.models.cosmos3.parameter_resolver import ( resolve_domain_action_config, ) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_parameter_resolver.py b/tests/unittest/_torch/visual_gen/test_cosmos3_parameter_resolver.py new file mode 100644 index 000000000000..cc34cddc53ae --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_parameter_resolver.py @@ -0,0 +1,301 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The Cosmos3 precedence table: the engine on synthetic sources, the checkpoint rules +on synthetic files and on the published checkpoints, and the request rules.""" + +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.visual_gen.models.cosmos3.parameter_resolver import ( + ACTION_RULES, + CHECKPOINT, + CHECKPOINT_RULES, + CONFIG, + GENERATION_RULES, + MODEL_INDEX, + SCHEDULER, + TRANSFORMER, + VAE, + CheckpointFiles, + Conflict, + Default, + Given, + Key, + RequestContext, + Rule, + Table, + check_precedence, + fill_pretrained_config, + resolve_action, + resolve_checkpoint, + resolve_domain_action_config, + resolve_generation, + resolve_setting, +) + +pytestmark = pytest.mark.cpu_only + +FIXTURE = Path(__file__).parent / "test_data" / "cosmos3_checkpoint_configs.json" +T_LIST = [1.0, 0.9375, 0.8333333333333334, 0.625] + + +def _write(checkpoint_dir: Path, files: dict) -> Path: + for name, doc in files.items(): + path = checkpoint_dir / name + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(doc)) + return checkpoint_dir + + +def _files_on_disk(tmp_path: Path, files: dict) -> CheckpointFiles: + _write(tmp_path, files) + return CheckpointFiles(str(tmp_path), files.get(TRANSFORMER, {})) + + +def _resolve(tmp_path: Path, files: dict): + _write(tmp_path, files) + return resolve_checkpoint(str(tmp_path), files.get(TRANSFORMER, {})) + + +class TestEngine: + def test_first_present_source_wins(self, tmp_path): + files = _files_on_disk(tmp_path, {CONFIG: {"a": 1}, MODEL_INDEX: {"a": 1}}) + rule = Rule((Key(CONFIG, "a"), Key(MODEL_INDEX, "a"), Default(7, "x"))) + assert resolve_setting("a", rule, files).value == 1 + assert resolve_setting("a", rule, files).source == "config.json:a" + + def test_falls_through_missing_and_null_sources(self, tmp_path): + files = _files_on_disk(tmp_path, {CONFIG: {"a": None}, MODEL_INDEX: {"a": 2}}) + rule = Rule((Key(CONFIG, "a"), Key(MODEL_INDEX, "a"), Key(CHECKPOINT, "a"))) + assert resolve_setting("a", rule, files).value == 2 + + def test_default_only_when_nothing_is_declared(self, tmp_path): + files = _files_on_disk(tmp_path, {}) + rule = Rule((Key(CONFIG, "a"), Default(7, "x"))) + resolved = resolve_setting("a", rule, files) + assert (resolved.value, resolved.source) == (7, "default (x)") + + def test_unresolvable_required_setting_is_an_error(self, tmp_path): + files = _files_on_disk(tmp_path, {}) + with pytest.raises(ValueError, match="cannot resolve a"): + resolve_setting("a", Rule((Key(CONFIG, "a"),)), files) + assert resolve_setting("a", Rule((Key(CONFIG, "a"),), required=False), files) is None + + def test_conflict_policies(self, tmp_path): + files = _files_on_disk(tmp_path, {CONFIG: {"a": 1}, MODEL_INDEX: {"a": 2}}) + sources = (Key(CONFIG, "a"), Key(MODEL_INDEX, "a")) + with pytest.raises(ValueError, match="model_index.json:a=2"): + resolve_setting("a", Rule(sources, conflict=Conflict.ERROR), files) + warned = resolve_setting("a", Rule(sources, conflict=Conflict.WARN), files) + assert warned.value == 1 and warned.disagreeing == ("model_index.json:a=2",) + ignored = resolve_setting("a", Rule(sources), files) + assert ignored.value == 1 and ignored.disagreeing == ("model_index.json:a=2",) + + def test_float_lists_compare_within_tolerance(self, tmp_path): + files = _files_on_disk( + tmp_path, {CONFIG: {"a": [0.5, 1 / 3]}, MODEL_INDEX: {"a": [0.5, 0.3333333333]}} + ) + rule = Rule((Key(CONFIG, "a"), Key(MODEL_INDEX, "a")), conflict=Conflict.ERROR) + assert resolve_setting("a", rule, files).disagreeing == () + + def test_every_rule_keeps_request_over_checkpoint_over_default(self): + for rules in (CHECKPOINT_RULES, GENERATION_RULES, ACTION_RULES): + check_precedence(rules) + with pytest.raises(ValueError, match="breaks the request > checkpoint > default order"): + check_precedence({"x": Rule((Default(1, "first"), Given()))}) + with pytest.raises(ValueError, match="breaks the request > checkpoint > default order"): + check_precedence({"x": Rule((Table("x"), Given()))}) + + +def _files(): + """A distilled checkpoint whose training config and diffusers index disagree on + the system prompt.""" + return { + TRANSFORMER: {"temporal_compression_factor": 4}, + SCHEDULER: {"fixed_step_sampler_config": {"t_list": T_LIST}}, + CONFIG: {"model": {"config": {"vlm_config": {"use_system_prompt": False}}}}, + MODEL_INDEX: {"default_use_system_prompt": True}, + VAE: {"scale_factor_temporal": 4, "scale_factor_spatial": 16}, + } + + +class TestCheckpoint: + def test_training_config_beats_the_diffusers_index(self, tmp_path): + cfg = _resolve(tmp_path, _files()) + assert cfg.prompt.use_system_prompt is False + assert cfg.source_of("prompt.use_system_prompt") == ( + "config.json:model.config.vlm_config.use_system_prompt" + ) + assert cfg.resolved["prompt.use_system_prompt"].disagreeing == ( + "model_index.json:default_use_system_prompt=True", + ) + + def test_sections(self, tmp_path): + cfg = _resolve(tmp_path, _files()) + assert (cfg.vae.temporal_compression, cfg.vae.spatial_compression) == (4, 16) + assert cfg.sampling.fixed_sigmas == tuple(T_LIST) + assert cfg.sampling.native_flow_schedule is False + assert cfg.source_of("sampling.native_flow_schedule").startswith("default (") + assert cfg.action.policy is None + + def test_vae_compression_copies_must_agree(self, tmp_path): + files = _files() + files[TRANSFORMER]["temporal_compression_factor"] = 8 + with pytest.raises(ValueError, match="vae.temporal_compression"): + _resolve(tmp_path, files) + + def test_schedule_must_decrease_within_unit_interval(self, tmp_path): + files = _files() + files[SCHEDULER]["fixed_step_sampler_config"]["t_list"] = [0.5, 1.0] + with pytest.raises(ValueError, match="t_list must decrease"): + _resolve(tmp_path, files) + + def test_base_checkpoint_has_no_schedule(self, tmp_path): + files = {k: v for k, v in _files().items() if k != SCHEDULER} + assert _resolve(tmp_path, files).sampling.fixed_sigmas is None + + def test_system_prompt_must_be_declared(self, tmp_path): + files = _files() + del files[CONFIG]["model"]["config"]["vlm_config"]["use_system_prompt"] + del files[MODEL_INDEX]["default_use_system_prompt"] + with pytest.raises(ValueError, match="cannot resolve prompt.use_system_prompt"): + _resolve(tmp_path, files) + + def test_spatial_compression_must_be_declared(self, tmp_path): + files = _files() + del files[VAE]["scale_factor_spatial"] + with pytest.raises(ValueError, match="cannot resolve vae.spatial_compression"): + _resolve(tmp_path, files) + + def test_fill_pretrained_config_sets_only_missing_attributes(self): + declared = SimpleNamespace(temporal_compression_factor=8) + fill_pretrained_config(declared) + assert declared.temporal_compression_factor == 8 + assert declared.position_embedding_type == "unified_3d_mrope" + assert declared.max_position_embeddings == 262144 + assert declared.temporal_compression_factor_sound == 1 + + +# Disagreements between a checkpoint's own files that the table expects and resolves. +EXPECTED_DISAGREEMENTS = { + "nvidia/Cosmos3-Nano-Sim-Bimanual": {"prompt.use_system_prompt"}, +} + +EXPECTED = { + # repo: (use_system_prompt, native_flow_schedule, distilled) + "nvidia/Cosmos3-Nano": (False, False, False), + "nvidia/Cosmos3-Super": (False, False, False), + "nvidia/Cosmos3-Edge": (False, True, False), + "nvidia/Cosmos3-Super-Image2Video": (False, False, False), + "nvidia/Cosmos3-Super-Image2Video-4Step": (False, False, True), + "nvidia/Cosmos3-Super-Text2Image": (False, False, False), + "nvidia/Cosmos3-Super-Text2Image-4Step": (False, False, True), + "nvidia/Cosmos3-Nano-Policy-DROID": (False, False, False), + "nvidia/Cosmos3-Edge-Policy-DROID": (False, True, False), + "nvidia/Cosmos3-Nano-Sim-Bimanual": (False, False, True), +} + + +@pytest.mark.parametrize("repo", sorted(EXPECTED)) +def test_published_checkpoint_resolves(tmp_path, repo): + files = json.loads(FIXTURE.read_text())[repo] + cfg = _resolve(tmp_path, files) + system_prompt, native, distilled = EXPECTED[repo] + assert cfg.prompt.use_system_prompt is system_prompt + assert cfg.sampling.native_flow_schedule is native + assert (cfg.sampling.fixed_sigmas is not None) is distilled + assert (cfg.vae.temporal_compression, cfg.vae.spatial_compression) == (4, 16) + assert cfg.architecture.position_embedding_type == "unified_3d_mrope" + disagreeing = {name for name, r in cfg.resolved.items() if r.disagreeing} + assert disagreeing == EXPECTED_DISAGREEMENTS.get(repo, set()) + + +def test_fixture_covers_every_expected_checkpoint(): + assert set(json.loads(FIXTURE.read_text())) == set(EXPECTED) + + +class TestRequest: + def test_generation_values_name_their_source(self): + ctx = RequestContext(family="qwen3", mode="image") + r = resolve_generation(ctx, num_inference_steps=None, num_frames=None, height=20) + assert r.source_of("num_inference_steps") == "mode table.num_inference_steps" + assert r["num_inference_steps"] == 50 + assert r.source_of("num_frames") == "video table.num_frames" # the image table has none + assert (r["height"], r.source_of("height")) == (20, "request") + + def test_distilled_schedule_beats_the_tables(self): + ctx = RequestContext(family="qwen3", distilled_steps=4) + r = resolve_generation(ctx, num_inference_steps=None, guidance_scale=None) + assert (r["num_inference_steps"], r["guidance_scale"]) == (4, 1.0) + assert r.source_of("guidance_scale").startswith("checkpoint:fixed_step_sampler_config") + + def test_request_beats_the_distilled_schedule(self): + """The caller's value wins over the checkpoint, and the checkpoint's value is + recorded as the disagreeing source.""" + ctx = RequestContext(family="qwen3", distilled_steps=4) + r = resolve_generation(ctx, num_inference_steps=8) + assert (r["num_inference_steps"], r.source_of("num_inference_steps")) == (8, "request") + assert r.values["num_inference_steps"].disagreeing == ( + "checkpoint:fixed_step_sampler_config (steps)=4", + ) + + def test_policy_block_applies_only_to_its_domain(self): + policy = {"domain_name": "droid_lerobot", "action_chunk_size": 32, "num_inference_steps": 4} + matched = resolve_domain_action_config(checkpoint_policy_defaults=policy) + assert matched["domain_name"] == "droid_lerobot" + assert matched["action_chunk_size"] == 32 + assert matched["sources"]["action_chunk_size"] == "checkpoint.json:policy.action_chunk_size" + other = resolve_domain_action_config(domain_name="av", checkpoint_policy_defaults=policy) + assert other["action_chunk_size"] == 60 + assert other["sources"]["action_chunk_size"] == "domain preset.action_chunk_size" + assert any("overrides checkpoint policy domain_name" in w for w in other["warnings"]) + + def test_caller_value_wins_and_the_recommendation_is_reported(self): + cfg = resolve_domain_action_config(domain_name="av", action_chunk_size=16) + assert cfg["action_chunk_size"] == 16 + assert cfg["sources"]["action_chunk_size"] == "request" + assert cfg["warnings"] == [ + "Cosmos3 action_chunk_size=16 differs from recommended 60 for domain 'av'." + ] + + def test_explicit_none_in_a_mode_table_is_a_value(self): + """Edge text-to-image switches the guidance window off with an explicit ``None``; + that must not fall through to the video table.""" + ctx = RequestContext(family="nemotron_dense", mode="image") + r = resolve_generation(ctx, guidance_interval=None) + assert r["guidance_interval"] is None + assert r.source_of("guidance_interval") == "mode table.guidance_interval" + + def test_unresolvable_field_is_an_error(self): + ctx = RequestContext(family="no-such-family") + with pytest.raises(ValueError, match="cannot resolve num_frames"): + resolve_generation(ctx, num_frames=None) + + def test_action_fps_follows_the_frame_rate(self): + ctx = RequestContext(family="qwen3", action_mode="forward_dynamics", domain_name="av") + r = resolve_action( + ctx, + raw_action_dim=None, + action_chunk_size=None, + action_resolution=None, + frame_rate=None, + action_fps=None, + use_state=None, + ) + assert (r["frame_rate"], r["action_fps"], r["num_frames"]) == (10.0, 10.0, 61) + assert r.source_of("action_fps") == "frame_rate" diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_request_resolver.py b/tests/unittest/_torch/visual_gen/test_cosmos3_request_resolver.py deleted file mode 100644 index 612dc0458586..000000000000 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_request_resolver.py +++ /dev/null @@ -1,94 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""The request rules: precedence between the caller, the checkpoint's facts and the -tables, and the source each value reports.""" - -import pytest - -from tensorrt_llm._torch.visual_gen.models.cosmos3.request_resolver import ( - RequestContext, - resolve_action, - resolve_domain_action_config, - resolve_generation, -) - -pytestmark = pytest.mark.cpu_only - - -class TestSources: - def test_generation_values_name_their_source(self): - ctx = RequestContext(family="qwen3", mode="image") - r = resolve_generation(ctx, num_inference_steps=None, num_frames=None, height=20) - assert r.source_of("num_inference_steps") == "mode table.num_inference_steps" - assert r["num_inference_steps"] == 50 - assert r.source_of("num_frames") == "video table.num_frames" # the image table has none - assert (r["height"], r.source_of("height")) == (20, "request") - - def test_distilled_schedule_beats_the_tables(self): - ctx = RequestContext(family="qwen3", distilled_steps=4) - r = resolve_generation(ctx, num_inference_steps=None, guidance_scale=None) - assert (r["num_inference_steps"], r["guidance_scale"]) == (4, 1.0) - assert r.source_of("guidance_scale").startswith("checkpoint:fixed_step_sampler_config") - - def test_policy_block_applies_only_to_its_domain(self): - policy = {"domain_name": "droid_lerobot", "action_chunk_size": 32, "num_inference_steps": 4} - matched = resolve_domain_action_config(checkpoint_policy_defaults=policy) - assert matched["domain_name"] == "droid_lerobot" - assert matched["action_chunk_size"] == 32 - assert matched["sources"]["action_chunk_size"] == "checkpoint.json:policy.action_chunk_size" - other = resolve_domain_action_config(domain_name="av", checkpoint_policy_defaults=policy) - assert other["action_chunk_size"] == 60 - assert other["sources"]["action_chunk_size"] == "domain preset.action_chunk_size" - assert any("overrides checkpoint policy domain_name" in w for w in other["warnings"]) - - def test_caller_value_wins_and_the_recommendation_is_reported(self): - cfg = resolve_domain_action_config(domain_name="av", action_chunk_size=16) - assert cfg["action_chunk_size"] == 16 - assert cfg["sources"]["action_chunk_size"] == "request" - assert cfg["warnings"] == [ - "Cosmos3 action_chunk_size=16 differs from recommended 60 for domain 'av'." - ] - - def test_explicit_none_in_a_mode_table_is_a_value(self): - """Edge text-to-image switches the guidance window off with an explicit ``None``; - that must not fall through to the video table.""" - ctx = RequestContext(family="nemotron_dense", mode="image") - r = resolve_generation(ctx, guidance_interval=None) - assert r["guidance_interval"] is None - assert r.source_of("guidance_interval") == "mode table.guidance_interval" - - def test_disagreeing_recommendation_names_its_source(self): - cfg = resolve_domain_action_config(domain_name="av", action_chunk_size=16) - assert cfg["action_chunk_size"] == 16 - assert cfg["sources"]["action_chunk_size"] == "request" - - def test_unresolvable_field_is_an_error(self): - ctx = RequestContext(family="no-such-family") - with pytest.raises(ValueError, match="cannot resolve num_frames"): - resolve_generation(ctx, num_frames=None) - - def test_action_fps_follows_the_frame_rate(self): - ctx = RequestContext(family="qwen3", action_mode="forward_dynamics", domain_name="av") - r = resolve_action( - ctx, - raw_action_dim=None, - action_chunk_size=None, - action_resolution=None, - frame_rate=None, - action_fps=None, - use_state=None, - ) - assert (r["frame_rate"], r["action_fps"], r["num_frames"]) == (10.0, 10.0, 61) - assert r.source_of("action_fps") == "frame_rate" diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py index 53e41a5e58bd..09327dfc6d3f 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py @@ -25,7 +25,7 @@ from tensorrt_llm._torch.modules.linear import Linear from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig, DiffusionPipelineConfig -from tensorrt_llm._torch.visual_gen.models.cosmos3.config_resolver import fill_pretrained_config +from tensorrt_llm._torch.visual_gen.models.cosmos3.parameter_resolver import fill_pretrained_config from tensorrt_llm._torch.visual_gen.models.cosmos3.transformer_cosmos3 import ( Cosmos3VFMTransformer, Qwen3VLTextRotaryEmbedding, diff --git a/tests/unittest/_torch/visual_gen/test_data/README.md b/tests/unittest/_torch/visual_gen/test_data/README.md index 6d3232e56189..cfe1ad65314c 100644 --- a/tests/unittest/_torch/visual_gen/test_data/README.md +++ b/tests/unittest/_torch/visual_gen/test_data/README.md @@ -60,10 +60,10 @@ ffmpeg -y -framerate 24 -i frame_%02d.png \ ## `cosmos3_checkpoint_configs.json` -The settings `config_resolver.RULES` reads, extracted from the +The settings `parameter_resolver.CHECKPOINT_RULES` reads, extracted from the published Cosmos3 checkpoints on Hugging Face (`main`, Fri 23 Oct 2026): for each repo, file and key path a rule names, the value at that path, nested back into its -file. Nothing else from the checkpoints is kept. `test_cosmos3_config_resolver.py` +file. Nothing else from the checkpoints is kept. `test_cosmos3_parameter_resolver.py` resolves every entry, so a new export whose files disagree in a way the precedence table does not expect fails the test. From b69c3d547a45304c0319f4ed30252efdb80a0d56 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 2 Nov 2026 13:34:23 -0800 Subject: [PATCH 7/8] [None][refactor] Cosmos3: one resolver with three sources, request over checkpoint over default Every setting now goes through select(request, checkpoint, default); the rule table, the per-source classes and the precedence checker are gone. The resolver holds one Cosmos3Config read from the checkpoint files; the sampling policy takes the distilled schedule from it instead of reading the scheduler config itself. flow_shift follows the same rule: the request's value, else the checkpoint scheduler's. The table entries (3.0 for text-to-image and Edge video) and the 10.0 video-to-video constant that used to outrank the checkpoint are removed; the LPIPS golden tests pass the shift their goldens were cut at. vae.spatial_compression is required only when the VAE component is loaded. The transfer/V2V system-prompt rules are tagged "workflow", not "request". Tests build checkpoint settings by hand through cosmos3_fixtures.py. Signed-off-by: Igor Shovkun --- .../visual_gen/models/cosmos3/defaults.py | 26 +- .../models/cosmos3/parameter_resolver.py | 1056 +++++++---------- .../models/cosmos3/pipeline_cosmos3.py | 347 ++---- .../visual_gen/models/cosmos3/sampling.py | 22 +- .../visual_gen/test_visual_gen_cosmos3.py | 13 +- .../_torch/visual_gen/cosmos3_fixtures.py | 73 ++ .../_torch/visual_gen/test_cosmos3_action.py | 31 +- .../visual_gen/test_cosmos3_distilled.py | 88 +- .../_torch/visual_gen/test_cosmos3_edge.py | 178 +-- .../test_cosmos3_example_prompts.py | 39 +- .../_torch/visual_gen/test_cosmos3_fp8.py | 2 +- .../test_cosmos3_parameter_resolver.py | 295 +++-- .../visual_gen/test_cosmos3_pipeline.py | 20 +- .../visual_gen/test_cosmos3_transfer.py | 2 + .../_torch/visual_gen/test_data/README.md | 10 +- 15 files changed, 1009 insertions(+), 1193 deletions(-) create mode 100644 tests/unittest/_torch/visual_gen/cosmos3_fixtures.py diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index 29ff55588973..23b4ce61baf8 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -14,7 +14,8 @@ # limitations under the License. """The Cosmos3 tables: generation defaults per family and mode, the action and policy presets per embodiment domain, the prompt templates, and the schema of the -request's extra parameters. The request resolver reads them in precedence order. +request's extra parameters. The parameter resolver reads them as the defaults, after +the request and the checkpoint. """ from collections.abc import Mapping @@ -351,21 +352,20 @@ def _validate_non_negative_frames(value: Any) -> None: raise ValueError(f"must be a non-negative frame count, got {value}.") -# Text-to-image (``output_type="image"``) defaults; resolved in ``infer()``. +# Text-to-image (``output_type="image"``) defaults. COSMOS3_T2I_PARAMS = { "height": 1024, "width": 1024, "num_inference_steps": 50, "guidance_scale": 7.0, - "flow_shift": 3.0, "guidance_interval": (400.0, 1000.0), } # Edge (Nemotron-dense backbone) is 480p-native. Video values follow the -# model-card I2V command (T2V mirrors it — the model card documents I2V only); -# ``flow_shift`` rides the checkpoint-declared native flow schedule. T2I values -# are the cosmos-framework t2i mode defaults at Edge's native resolution -# (480p at 1:1 aspect), with full-range CFG. +# model-card I2V command (T2V mirrors it — the model card documents I2V only). +# T2I values are the cosmos-framework t2i mode defaults at Edge's native +# resolution (480p at 1:1 aspect), with full-range CFG. No table carries a +# ``flow_shift``: a request that omits it runs at the checkpoint scheduler's. COSMOS3_EDGE_VIDEO_PARAMS = { "height": 480, "width": 832, @@ -374,7 +374,6 @@ def _validate_non_negative_frames(value: Any) -> None: "max_sequence_length": 4096, "num_frames": 121, "frame_rate": 24.0, - "flow_shift": 3.0, } COSMOS3_EDGE_T2I_PARAMS = { @@ -382,7 +381,6 @@ def _validate_non_negative_frames(value: Any) -> None: "width": 640, "num_inference_steps": 50, "guidance_scale": 4.0, - "flow_shift": 3.0, "guidance_interval": None, } @@ -420,9 +418,9 @@ def _validate_non_negative_frames(value: Any) -> None: } # Action's table carries the sampling recipe only; the canvas, clip length and -# frame rate resolve from the embodiment preset in the pipeline, and anything -# missing here falls back to the family's video table inside -# _resolve_generation_params. Same recipe for both families. +# frame rate come from the embodiment preset, and anything missing here falls +# back to the family's video table in the parameter resolver. Same recipe for +# both families. # Families without an entry get no envelope advisory. @@ -660,10 +658,6 @@ def get_domain_preset( return COSMOS3_DOMAIN_PRESETS[key] -# Video-to-video and transfer sample on the shifted flow grid with this shift; the -# other modes carry ``flow_shift`` in their tables above or take the checkpoint's. -COSMOS3_V2V_FLOW_SHIFT = 10.0 - # Prompt format the checkpoints were trained with. The "give" typo is the training # text; keep it. Metadata sentences are appended to the caption the same way. COSMOS3_DEFAULT_SYSTEM_PROMPT = ( diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/parameter_resolver.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/parameter_resolver.py index dd36c097837d..be4a2a718fb1 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/parameter_resolver.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/parameter_resolver.py @@ -12,24 +12,23 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""One precedence table for every Cosmos3 setting: request over checkpoint over default. +"""Every Cosmos3 setting comes from one of three places, tried in this order: the +request, the checkpoint, a default. :func:`select` is the only thing that orders +them, and every resolved value remembers which one it came from. -A rule names the sources of one setting in priority order; the first source holding a -value wins, the next one holding a value is compared with it, and every resolved -value remembers where it came from. The same engine serves two moments: +:class:`ParameterResolver` owns one checkpoint's settings and serves two moments: -* at load, :data:`CHECKPOINT_RULES` turn the checkpoint's files (``config.json``, +* at load, :func:`resolve_checkpoint` reads the checkpoint's files (``config.json``, ``model_index.json``, ``checkpoint.json``, the component configs) into a - :class:`Cosmos3Config`: a Cosmos3 export carries the same facts in several files - written by exporters that are not kept in sync, so each rule says which file wins - and what a disagreement means; -* per request, :data:`GENERATION_RULES` and :data:`ACTION_RULES` turn the caller's - values, the checkpoint's facts (the distilled schedule, the policy block) and the - per-family and per-domain tables into what the request runs with. - -In every rule the caller's value comes first, then the checkpoint, then the code's -tables, then an explicit :class:`Default` with a reason. A setting no source provides -and no default covers is an error, not a guess. + :class:`Cosmos3Config`. A Cosmos3 export carries the same fact in several files + written by exporters that are not kept in sync, so each setting says which file + wins and what a disagreement means; +* per request, ``request()`` turns a :class:`RequestParameters` (the caller's values, + ``None`` = unset) into the same object with every field the mode uses filled: the + caller's value, else the checkpoint (the distilled schedule, the policy block), else + the per-family tables and the embodiment presets. + +A setting nothing provides is an error, not a guess. """ from __future__ import annotations @@ -37,69 +36,57 @@ import json import math import os -from collections.abc import Callable, Mapping -from dataclasses import dataclass, field, replace +from collections.abc import Mapping +from dataclasses import dataclass, replace from enum import Enum -from typing import Any, Optional +from typing import Any, NamedTuple, Optional from tensorrt_llm.logger import logger +from .action import ACTION_MODE_POLICY from .defaults import ( - COSMOS3_ACTION_PARAMS, COSMOS3_DOMAIN_PRESETS, COSMOS3_GENERATION_DEFAULTS, EMBODIMENT_TO_DOMAIN_ID, canonical_domain_preset_key, normalize_action_resolution, + resolve_checkpoint_policy_defaults, resolve_raw_action_dim, ) from .sampling import DISTILLED_GUIDANCE_SCALE -_MISSING = object() - - -# ------------------------------------------------------------------------------ engine - - -class Conflict(Enum): - """What a lower-priority source holding a different value means.""" - - ERROR = "error" - WARN = "warn" - IGNORE = "ignore" +REQUEST = "request" +CHECKPOINT = "checkpoint" +DEFAULT = "default" +WORKFLOW = "workflow" +"""Origin of a field the kind of request fixes (transfer, video-to-video) before the +checkpoint is consulted; only ``use_system_prompt`` has such rules.""" +UNSET = "unset" +"""Origin of a field that none of the three provides and that may stay ``None``.""" @dataclass(frozen=True) -class Default: - """The value used when no source declares the setting; always the last source.""" - +class Resolved: value: Any - why: str + origin: str + """``REQUEST``, ``CHECKPOINT``, ``DEFAULT`` or ``WORKFLOW``.""" + where: str = "" + """For checkpoint settings, the file and key that held the value.""" + disagreeing: tuple[str, ...] = () + """``file:key=value`` of the checkpoint files that held a different value.""" @property - def label(self) -> str: - return f"default ({self.why})" - - def lookup(self, ctx, name, given, values): - return self.value + def source(self) -> str: + return f"{self.origin} ({self.where})" if self.where else self.origin -@dataclass(frozen=True) -class Rule: - sources: tuple - conflict: Conflict = Conflict.IGNORE - """``WARN`` / ``ERROR``: a source below the winner that holds a different value is - reported / rejected.""" - required: bool = True - """``False``: an unresolved setting is ``None`` instead of an error.""" - - -@dataclass(frozen=True) -class Resolved: - value: Any - source: str - disagreeing: tuple[str, ...] = () - """``label=value`` of the lower-priority source that held a different value.""" +def select(request: Any = None, checkpoint: Any = None, default: Any = None) -> Optional[Resolved]: + """The first of request, checkpoint, default that is set (not ``None``), tagged + with its origin; ``None`` when nothing is set.""" + for origin, value in ((REQUEST, request), (CHECKPOINT, checkpoint), (DEFAULT, default)): + if value is not None: + return Resolved(value, origin) + return None def same_value(a: Any, b: Any) -> bool: @@ -116,43 +103,42 @@ def same_value(a: Any, b: Any) -> bool: return a == b -def _pick( - name: str, rule: Rule, ctx: Any, given: Mapping[str, Any], values: Mapping[str, Any] -) -> tuple[Optional[Resolved], list[tuple[str, Any]]]: - """Walk ``rule.sources`` in order. The first source holding a value wins; the next - one holding a value is compared with it. A :class:`Default` is consulted only when - nothing above it holds a value.""" - winner: Optional[Resolved] = None - disagreeing: list[tuple[str, Any]] = [] - for source in rule.sources: - if winner is not None and isinstance(source, Default): - break - value = source.lookup(ctx, name, given, values) - if value is _MISSING: - continue - if winner is None: - winner = Resolved(value, source.label) - continue - if not same_value(value, winner.value): - disagreeing.append((source.label, value)) - break - return winner, disagreeing +# -------------------------------------------------------------------------- checkpoint +CONFIG_FILE = "config.json" +CHECKPOINT_FILE = "checkpoint.json" +MODEL_INDEX_FILE = "model_index.json" +TRANSFORMER_FILE = "transformer/config.json" +SCHEDULER_FILE = "scheduler/scheduler_config.json" +VAE_FILE = "vae/config.json" -def _record(winner: Resolved, disagreeing: list[tuple[str, Any]]) -> Resolved: - return Resolved( - winner.value, winner.source, tuple(f"{label}={v!r}" for label, v in disagreeing) - ) +class Conflict(Enum): + """What it means when a lower-priority checkpoint file holds a different value.""" + + ERROR = "error" + WARN = "warn" + IGNORE = "ignore" -# -------------------------------------------------------------------------- checkpoint -CONFIG = "config.json" -CHECKPOINT = "checkpoint.json" -MODEL_INDEX = "model_index.json" -TRANSFORMER = "transformer/config.json" -SCHEDULER = "scheduler/scheduler_config.json" -VAE = "vae/config.json" +_MISSING = object() + + +def _step(obj: Any, key: str) -> Any: + if isinstance(obj, Mapping): + return obj.get(key, _MISSING) + if obj is None or isinstance(obj, (str, int, float, bool, list, tuple)): + return _MISSING + return getattr(obj, key, _MISSING) # a pretrained config object standing in for a file + + +def _walk(doc: Any, path: str) -> Any: + node = doc + for part in path.split("."): + node = _step(node, part) + if node is _MISSING: + return _MISSING + return node class CheckpointFiles: @@ -165,7 +151,7 @@ def __init__(self, checkpoint_dir: Optional[str], pretrained_config: Any = None) self._docs: dict[str, Any] = {} def doc(self, file: str) -> Any: - if file == TRANSFORMER and self._pretrained is not None: + if file == TRANSFORMER_FILE and self._pretrained is not None: return self._pretrained if file not in self._docs: doc = None @@ -177,89 +163,43 @@ def doc(self, file: str) -> Any: self._docs[file] = doc return self._docs[file] - -def _step(obj: Any, key: str) -> Any: - if isinstance(obj, Mapping): - return obj.get(key, _MISSING) - if obj is None or isinstance(obj, (str, int, float, bool, list, tuple)): - return _MISSING - return getattr(obj, key, _MISSING) # a pretrained config object standing in for a file - - -def _walk(doc: Any, path: str) -> Any: - node = doc - for part in path.split("."): - node = _step(node, part) - if node is _MISSING: - return _MISSING - return node - - -@dataclass(frozen=True) -class Key: - """A dotted key path in one checkpoint file; a ``null`` there counts as absent.""" - - file: str - path: str - - @property - def label(self) -> str: - return f"{self.file}:{self.path}" - - def lookup(self, ctx: CheckpointFiles, name, given, values): - doc = ctx.doc(self.file) - value = _MISSING if doc is None else _walk(doc, self.path) - return _MISSING if value is None else value - - -# Priority among the files: the training config (config.json), then the component -# configs, then the export metadata and the diffusers index. -CHECKPOINT_RULES: dict[str, Rule] = { - "architecture.position_embedding_type": Rule( - ( - Key(TRANSFORMER, "position_embedding_type"), - Default("unified_3d_mrope", "exports that predate the key"), - ) - ), - "architecture.max_position_embeddings": Rule( - ( - Key(TRANSFORMER, "max_position_embeddings"), - Default(262144, "exports that predate the key"), - ) - ), - "vae.temporal_compression": Rule( - ( - Key(VAE, "scale_factor_temporal"), - Key(TRANSFORMER, "temporal_compression_factor"), - Default(4, "exports that predate the key"), - ), - conflict=Conflict.ERROR, - ), - "vae.spatial_compression": Rule((Key(VAE, "scale_factor_spatial"),), conflict=Conflict.ERROR), - "audio.temporal_compression": Rule( - ( - Key(TRANSFORMER, "temporal_compression_factor_sound"), - Default(1, "exports that predate the key"), - ) - ), - "prompt.use_system_prompt": Rule( - ( - Key(CONFIG, "model.config.vlm_config.use_system_prompt"), - Key(MODEL_INDEX, "default_use_system_prompt"), - ), - conflict=Conflict.WARN, - ), - "sampling.native_flow_schedule": Rule( - ( - Key(MODEL_INDEX, "use_native_flow_schedule"), - Default(False, "only Edge exports declare the native grid"), - ) - ), - "sampling.fixed_sigmas": Rule( - (Key(SCHEDULER, "fixed_step_sampler_config.t_list"),), required=False - ), - "action.policy": Rule((Key(CHECKPOINT, "policy"),), required=False), -} + def value(self, key: str) -> Any: + """The value at ``"file:dotted.path"``; a ``null`` there counts as absent.""" + file, path = key.split(":", 1) + doc = self.doc(file) + value = _MISSING if doc is None else _walk(doc, path) + return None if value is _MISSING else value + + def read( + self, name: str, *keys: str, conflict: Conflict = Conflict.IGNORE + ) -> Optional[Resolved]: + """The setting ``name`` from the first of ``keys`` that holds it, tagged + ``CHECKPOINT`` with the key it came from; ``None`` when no file holds it. The + next file holding it is compared with the winner and a disagreement is + rejected, logged or only recorded, as ``conflict`` says.""" + winner: Optional[Resolved] = None + disagreeing: list[str] = [] + for key in keys: + value = self.value(key) + if value is None: + continue + if winner is None: + winner = Resolved(value, CHECKPOINT, key) + continue + if not same_value(value, winner.value): + disagreeing.append(f"{key}={value!r}") + break + if winner is None: + return None + if disagreeing and conflict is not Conflict.IGNORE: + message = ( + f"Cosmos3 config: {name} = {winner.value!r} from {winner.where}, but " + + "; ".join(disagreeing) + ) + if conflict is Conflict.ERROR: + raise ValueError(message) + logger.warning(message) + return replace(winner, disagreeing=tuple(disagreeing)) @dataclass(frozen=True) @@ -272,7 +212,9 @@ class Cosmos3ArchitectureSettings: class Cosmos3VAESettings: temporal_compression: int """Pixel frames per latent frame after the first, which is encoded alone.""" - spatial_compression: int + spatial_compression: Optional[int] + """``None`` when the export ships no ``vae/config.json``; a pipeline that loads the + VAE needs it, one that skips the VAE does not.""" @dataclass(frozen=True) @@ -285,6 +227,9 @@ class Cosmos3SamplingSettings: fixed_sigmas: Optional[tuple[float, ...]] """The distilled schedule (``t_list``); ``None`` on base checkpoints.""" native_flow_schedule: bool + flow_shift: Optional[float] + """The scheduler's flow shift; ``None`` on distilled checkpoints, whose fixed + schedule has no shift.""" @dataclass(frozen=True) @@ -307,6 +252,7 @@ class Cosmos3Config: prompt: Cosmos3PromptSettings action: Cosmos3ActionSettings resolved: Mapping[str, Resolved] + """Every setting above by dotted name (``vae.temporal_compression``), with origin.""" def source_of(self, name: str) -> str: return self.resolved[name].source @@ -324,40 +270,80 @@ def describe(self) -> str: return "\n".join(lines) -def resolve_setting(name: str, rule: Rule, files: CheckpointFiles) -> Optional[Resolved]: - """One checkpoint setting. A required setting nothing declares is an error naming - the sources; a disagreement is rejected or logged as the rule says.""" - winner, disagreeing = _pick(name, rule, files, {}, {}) - if winner is None: - if rule.required: - labels = ", ".join(s.label for s in rule.sources) - raise ValueError(f"Cosmos3 config: cannot resolve {name}; none of {labels} is set.") - return None - if disagreeing and rule.conflict is not Conflict.IGNORE: - message = ( - f"Cosmos3 config: {name} = {winner.value!r} from {winner.source}, but " - + "; ".join(f"{label}={v!r}" for label, v in disagreeing) - ) - if rule.conflict is Conflict.ERROR: - raise ValueError(message) - logger.warning(message) - return _record(winner, disagreeing) +ARCHITECTURE_KEYS: dict[str, tuple[str, Any]] = { + "position_embedding_type": (f"{TRANSFORMER_FILE}:position_embedding_type", "unified_3d_mrope"), + "max_position_embeddings": (f"{TRANSFORMER_FILE}:max_position_embeddings", 262144), + "temporal_compression_factor": (f"{TRANSFORMER_FILE}:temporal_compression_factor", 4), + "temporal_compression_factor_sound": ( + f"{TRANSFORMER_FILE}:temporal_compression_factor_sound", + 1, + ), +} +"""The ``pretrained_config`` attributes the transformer is built from: the key in +``transformer/config.json`` and the value for exports that predate it.""" + +REQUIRED_SETTINGS = ("prompt.use_system_prompt",) +"""Settings some checkpoint file must declare; nothing stands in for them.""" def resolve_checkpoint( checkpoint_dir: Optional[str], pretrained_config: Any = None ) -> Cosmos3Config: - """Apply :data:`CHECKPOINT_RULES` to the checkpoint's files.""" + """Read the checkpoint's files into a :class:`Cosmos3Config`. Raises on a required + setting no file declares, or on a disagreement between files that is an error.""" files = CheckpointFiles(checkpoint_dir, pretrained_config) resolved: dict[str, Resolved] = {} - for name, rule in CHECKPOINT_RULES.items(): - r = resolve_setting(name, rule, files) + + def setting( + name: str, *keys: str, conflict: Conflict = Conflict.IGNORE, default: Any = None + ) -> None: + r = files.read(name, *keys, conflict=conflict) or select(default=default) if r is not None: resolved[name] = r + def transformer_setting( + name: str, attr: str, *before: str, conflict: Conflict = Conflict.IGNORE + ) -> None: + """A setting the transformer is built from: the files in ``before`` first, then + the transformer config, then the compatibility default.""" + key, default = ARCHITECTURE_KEYS[attr] + setting(name, *before, key, conflict=conflict, default=default) + + transformer_setting("architecture.position_embedding_type", "position_embedding_type") + transformer_setting("architecture.max_position_embeddings", "max_position_embeddings") + transformer_setting( + "vae.temporal_compression", + "temporal_compression_factor", + f"{VAE_FILE}:scale_factor_temporal", + conflict=Conflict.ERROR, + ) + setting("vae.spatial_compression", f"{VAE_FILE}:scale_factor_spatial") + transformer_setting("audio.temporal_compression", "temporal_compression_factor_sound") + # The training config over the diffusers index; published exports do disagree. + setting( + "prompt.use_system_prompt", + f"{CONFIG_FILE}:model.config.vlm_config.use_system_prompt", + f"{MODEL_INDEX_FILE}:default_use_system_prompt", + conflict=Conflict.WARN, + ) + setting( + "sampling.native_flow_schedule", + f"{MODEL_INDEX_FILE}:use_native_flow_schedule", + default=False, + ) + setting("sampling.fixed_sigmas", f"{SCHEDULER_FILE}:fixed_step_sampler_config.t_list") + setting("sampling.flow_shift", f"{SCHEDULER_FILE}:flow_shift") + setting("action.policy", f"{CHECKPOINT_FILE}:policy") + + for name in REQUIRED_SETTINGS: + if name not in resolved: + raise ValueError( + f"Cosmos3 config: cannot resolve {name}; no checkpoint file declares it." + ) + def get(name: str) -> Any: - r = resolved.get(name) - return None if r is None else r.value + item = resolved.get(name) + return None if item is None else item.value sigmas = get("sampling.fixed_sigmas") if sigmas is not None and ( @@ -368,6 +354,8 @@ def get(name: str) -> Any: f"Cosmos3 config: fixed_step_sampler_config.t_list must decrease within (0, 1]: " f"{list(sigmas)}." ) + spatial = get("vae.spatial_compression") + flow_shift = get("sampling.flow_shift") return Cosmos3Config( architecture=Cosmos3ArchitectureSettings( position_embedding_type=str(get("architecture.position_embedding_type")), @@ -375,12 +363,13 @@ def get(name: str) -> Any: ), vae=Cosmos3VAESettings( temporal_compression=int(get("vae.temporal_compression")), - spatial_compression=int(get("vae.spatial_compression")), + spatial_compression=None if spatial is None else int(spatial), ), audio=Cosmos3AudioSettings(temporal_compression=int(get("audio.temporal_compression"))), sampling=Cosmos3SamplingSettings( fixed_sigmas=None if not sigmas else tuple(float(s) for s in sigmas), native_flow_schedule=bool(get("sampling.native_flow_schedule")), + flow_shift=None if flow_shift is None else float(flow_shift), ), prompt=Cosmos3PromptSettings(use_system_prompt=bool(get("prompt.use_system_prompt"))), action=Cosmos3ActionSettings(policy=get("action.policy")), @@ -388,36 +377,18 @@ def get(name: str) -> Any: ) -ARCHITECTURE_ATTRIBUTES = { - "position_embedding_type": "architecture.position_embedding_type", - "max_position_embeddings": "architecture.max_position_embeddings", - "temporal_compression_factor": "vae.temporal_compression", - "temporal_compression_factor_sound": "audio.temporal_compression", -} -"""``pretrained_config`` attributes the transformer reads, and the setting each holds.""" - - def fill_pretrained_config(pretrained_config: Any) -> Any: """Set the transformer-construction attributes of ``pretrained_config`` that the - export omits, from :data:`CHECKPOINT_RULES`. Idempotent; returns the config. + export omits (:data:`ARCHITECTURE_KEYS`). Idempotent; returns the config. Runs before the pipeline loads its components, so only ``pretrained_config`` - itself is consulted. - """ + itself is consulted.""" files = CheckpointFiles(None, pretrained_config) - for attr, name in ARCHITECTURE_ATTRIBUTES.items(): - if getattr(pretrained_config, attr, None) is not None: - continue - rule = CHECKPOINT_RULES[name] - transformer_only = replace( - rule, - sources=tuple( - s for s in rule.sources if not isinstance(s, Key) or s.file == TRANSFORMER - ), - ) - r = resolve_setting(name, transformer_only, files) - if r is not None: - setattr(pretrained_config, attr, r.value) + for attr, (key, default) in ARCHITECTURE_KEYS.items(): + if getattr(pretrained_config, attr, None) is None: + resolved = select(checkpoint=files.value(key), default=default) + assert resolved is not None # every entry has a default + setattr(pretrained_config, attr, resolved.value) return pretrained_config @@ -425,303 +396,80 @@ def fill_pretrained_config(pretrained_config: Any) -> Any: @dataclass(frozen=True) -class RequestContext: - """What a request resolves against: the checkpoint's facts and the request's mode.""" +class RequestParameters: + """A request's settable fields plus the facts that pick the tables. Built by the + caller with the values the request set (``None`` = unset); + :meth:`ParameterResolver.request` returns a copy with every field this mode uses + filled. Fields another mode owns (the action fields outside action mode, the + canvas inside it) stay ``None``.""" - family: str mode: str = "video" - """``"video"`` or ``"image"``; ignored for the table when ``action_mode`` is set.""" + """``"video"`` or ``"image"``; the action table is used instead when ``action_mode`` is set.""" action_mode: Optional[str] = None - distilled_steps: Optional[int] = None - """Length of the fixed schedule on a distilled checkpoint.""" - policy_defaults: Mapping[str, Any] = field(default_factory=dict) - """The normalized ``checkpoint.json`` policy block.""" - policy_active: bool = False - """The policy block applies: policy mode, and the request's domain is the block's.""" - domain_name: Optional[str] = None - domain_id: Optional[int] = None - checkpoint_use_system_prompt: bool = False is_v2v: bool = False is_transfer: bool = False + domain_name: Optional[str] = None + """The embodiment. A policy request that names none comes back with the checkpoint + policy block's domain; one that names another keeps its own and the block is not applied.""" + domain_id: Optional[int] = None - @property - def table_mode(self) -> str: - return "action" if self.action_mode is not None else self.mode - - @property - def preset_key(self) -> Optional[str]: - return canonical_domain_preset_key(self.domain_name, self.domain_id) - - -class Given: - """The caller's value for the field being resolved; always the first source.""" - - label = "request" - - def lookup(self, ctx: RequestContext, name: str, given: Mapping[str, Any], values) -> Any: - value = given.get(name) - return _MISSING if value is None else value - - -@dataclass(frozen=True) -class Distilled: - """A fact of the distilled schedule: its step count, or the guidance it bakes in.""" - - what: str - - @property - def label(self) -> str: - return f"checkpoint:fixed_step_sampler_config ({self.what})" - - def lookup(self, ctx, name, given, values): - if ctx.distilled_steps is None: - return _MISSING - return ctx.distilled_steps if self.what == "steps" else DISTILLED_GUIDANCE_SCALE - - -@dataclass(frozen=True) -class Policy: - """A field of the checkpoint's policy block, when the block applies to the request.""" - - key: str - - @property - def label(self) -> str: - return f"checkpoint.json:policy.{self.key}" - - def lookup(self, ctx, name, given, values): - if not ctx.policy_active: - return _MISSING - value = ctx.policy_defaults.get(self.key) - return _MISSING if value is None else value - - -@dataclass(frozen=True) -class Context: - """A value the request context fixes: a workflow rule (transfer never prepends the - system prompt, video-to-video always does) or a checkpoint fact.""" - - label: str - pick: Callable[[RequestContext], Any] - - def lookup(self, ctx, name, given, values): - value = self.pick(ctx) - return _MISSING if value is None else value - - -class KnownWidth: - """The code's table of action widths per embodiment.""" - - label = "embodiment table (defaults)" - - def lookup(self, ctx, name, given, values): - width = resolve_raw_action_dim(domain_name=ctx.domain_name, domain_id=ctx.domain_id) - return _MISSING if width is None else width - - -@dataclass(frozen=True) -class Preset: - """A field of the request's domain preset.""" - - key: str - - @property - def label(self) -> str: - return f"domain preset.{self.key}" - - def lookup(self, ctx, name, given, values): - preset = COSMOS3_DOMAIN_PRESETS.get(ctx.preset_key) if ctx.preset_key else None - value = preset.get(self.key) if preset else None - return _MISSING if value is None else value - - -@dataclass(frozen=True) -class Table: - """The (family, mode) generation table; ``mode=None`` is the request's own mode.""" - - key: str - mode: Optional[str] = None - - @property - def label(self) -> str: - return f"{self.mode or 'mode'} table.{self.key}" - - def lookup(self, ctx, name, given, values): - table = COSMOS3_GENERATION_DEFAULTS.get((ctx.family, self.mode or ctx.table_mode)) - if not table or self.key not in table: - return _MISSING - return table[self.key] # an explicit ``None`` in a table is a value: the field is off - - -@dataclass(frozen=True) -class Derived: - """A value computed from fields resolved before it.""" - - label: str - compute: Callable[[Mapping[str, Any]], Any] + height: Optional[int] = None + width: Optional[int] = None + num_frames: Optional[int] = None + frame_rate: Optional[float] = None + num_inference_steps: Optional[int] = None + guidance_scale: Optional[float] = None + guidance_interval: Optional[tuple[float, float]] = None + flow_shift: Optional[float] = None + max_sequence_length: Optional[int] = None + use_system_prompt: Optional[bool] = None + raw_action_dim: Optional[int] = None + action_chunk_size: Optional[int] = None + action_resolution: Optional[int] = None + action_fps: Optional[float] = None + use_state: Optional[bool] = None - def lookup(self, ctx, name, given, values): - value = self.compute(values) - return _MISSING if value is None else value + policy_active: bool = False + """Set by the resolver: the checkpoint's policy block applied to this request.""" + preset_key: Optional[str] = None + """Set by the resolver: the domain preset the action fields came from, if any.""" -# Precedence ranks: request, then the checkpoint's facts, then the code's tables, then -# a default. ``check_precedence`` holds every rule to this order. -PRECEDENCE: tuple[tuple[type, ...], ...] = ( - (Given,), - (Key, Distilled, Policy, Context), - (KnownWidth, Preset, Table, Derived), - (Default,), +ACTION_FIELDS = ( + "raw_action_dim", + "action_chunk_size", + "action_resolution", + "action_fps", + "use_state", ) +"""Fields only an action request resolves; ``None`` otherwise.""" +CANVAS_FIELDS = ("height", "width") +"""Fields only a video or image request resolves; an action request's canvas is the +resolution bucket (``action_resolution``), so they stay ``None`` there.""" -def check_precedence(rules: Mapping[str, Rule]) -> None: - """Raise unless every rule orders its sources request, checkpoint, tables, default.""" - rank = {cls: i for i, group in enumerate(PRECEDENCE) for cls in group} - for name, rule in rules.items(): - ranks = [rank[type(s)] for s in rule.sources] - if ranks != sorted(ranks): - raise ValueError(f"rule {name!r} breaks the request > checkpoint > default order") +OPTIONAL_FIELDS = ("guidance_interval", "flow_shift", "raw_action_dim") +"""Fields that may legitimately end up unset: a mode without a guidance window, a +distilled checkpoint without a flow shift, an embodiment without a canonical action width.""" +ACTION_RESOLUTION_DEFAULT = 480 +"""The bucket the action models were trained at.""" -GENERATION_RULES: dict[str, Rule] = { - "num_inference_steps": Rule( - ( - Given(), - Distilled("steps"), - Policy("num_inference_steps"), - Table("num_inference_steps"), - Table("num_inference_steps", "video"), - ) - ), - "guidance_scale": Rule( - ( - Given(), - Distilled("guidance"), - Policy("guidance_scale"), - Table("guidance_scale"), - Table("guidance_scale", "video"), - ) - ), - "guidance_interval": Rule( - ( - Given(), - Policy("guidance_interval"), - Table("guidance_interval"), - Table("guidance_interval", "video"), - ), - required=False, - ), - "max_sequence_length": Rule( - (Given(), Table("max_sequence_length"), Table("max_sequence_length", "video")) - ), - "height": Rule((Given(), Table("height"), Table("height", "video"))), - "width": Rule((Given(), Table("width"), Table("width", "video"))), - "num_frames": Rule((Given(), Table("num_frames"), Table("num_frames", "video"))), - "frame_rate": Rule((Given(), Table("frame_rate"), Table("frame_rate", "video"))), - "use_system_prompt": Rule( - ( - Given(), - Context("transfer (never)", lambda ctx: False if ctx.is_transfer else None), - Context("video-to-video (always)", lambda ctx: True if ctx.is_v2v else None), - Context("checkpoint:use_system_prompt", lambda ctx: ctx.checkpoint_use_system_prompt), - ) - ), -} -ACTION_RULES: dict[str, Rule] = { - "raw_action_dim": Rule( - (Given(), Policy("raw_action_dim"), KnownWidth()), - conflict=Conflict.WARN, - required=False, - ), - "action_chunk_size": Rule( - ( - Given(), - Policy("action_chunk_size"), - Preset("action_chunk_size"), - Default(COSMOS3_ACTION_PARAMS["action_chunk_size"], "generic action horizon"), - ), - conflict=Conflict.WARN, - ), - "action_resolution": Rule( - ( - Given(), - Policy("action_resolution"), - Preset("action_resolution"), - Default(480, "the bucket the action models were trained at"), - ), - conflict=Conflict.WARN, - ), - "frame_rate": Rule( - ( - Given(), - Policy("frame_rate"), - Preset("frame_rate"), - Default(COSMOS3_ACTION_PARAMS["frame_rate"], "generic action frame rate"), - ), - conflict=Conflict.WARN, - ), - "action_fps": Rule( - (Given(), Policy("action_fps"), Derived("frame_rate", lambda v: v["frame_rate"])) - ), - "use_state": Rule( - (Given(), Policy("use_state"), Default(False, "policies without a state row")), - conflict=Conflict.WARN, - ), - "num_frames": Rule( - (Derived("action_chunk_size + 1", lambda v: int(v["action_chunk_size"]) + 1),) - ), -} - -for _rules in (CHECKPOINT_RULES, GENERATION_RULES, ACTION_RULES): - check_precedence(_rules) - - -@dataclass -class ResolvedRequest: - values: dict[str, Resolved] = field(default_factory=dict) - warnings: list[str] = field(default_factory=list) - - def __getitem__(self, name: str) -> Any: - return self.values[name].value - - def source_of(self, name: str) -> str: - return self.values[name].source - - -def _resolve_request( - rules: Mapping[str, Rule], - ctx: RequestContext, - given: Mapping[str, Any], - domain_label: str, -) -> ResolvedRequest: - """Request fields: a caller's value that differs from the recommendation below it is - reported in ``warnings`` where the rule says ``WARN``.""" - out = ResolvedRequest() - plain: dict[str, Any] = {} - for name, rule in rules.items(): - winner, disagreeing = _pick(name, rule, ctx, given, plain) - if winner is None: - if rule.required: - raise ValueError(f"Cosmos3 request: cannot resolve {name} for {domain_label}.") - out.values[name] = Resolved(None, "unresolved") - plain[name] = None - continue - if disagreeing and winner.source == Given.label and rule.conflict is Conflict.WARN: - out.warnings.append( - f"Cosmos3 {name}={winner.value} differs from recommended " - f"{' / '.join(str(v) for _, v in disagreeing)} for domain {domain_label}." - ) - out.values[name] = _record(winner, disagreeing) - plain[name] = winner.value - return out +class Candidates(NamedTuple): + """What one field may take, by origin. ``recommended`` and ``fallback`` are both + defaults; a caller's value that differs from ``recommended`` (the domain preset, + the embodiment width) is worth a warning, one that differs from ``fallback`` (the + generation tables) is not.""" + request: Any = None + checkpoint: Any = None + recommended: Any = None + fallback: Any = None -def resolve_generation(ctx: RequestContext, **given: Any) -> ResolvedRequest: - """Resolve the generation fields in ``given`` (``None`` = unset).""" - rules = {name: GENERATION_RULES[name] for name in given} - return _resolve_request(rules, ctx, given, repr(ctx.table_mode)) + def pick(self) -> Optional[Resolved]: + default = self.fallback if self.recommended is None else self.recommended + return select(self.request, self.checkpoint, default) def policy_domain( @@ -754,99 +502,215 @@ def policy_domain( return domain_name, False, warnings -def resolve_action(ctx: RequestContext, **given: Any) -> ResolvedRequest: - """Resolve the action fields; ``ctx.domain_name`` / ``domain_id`` already reconciled - with the policy block (:func:`policy_domain`).""" - domain_label = ( - repr(ctx.preset_key) - if ctx.preset_key - else (f"domain_name={ctx.domain_name!r}, domain_id={ctx.domain_id!r}") - ) - resolved = _resolve_request(ACTION_RULES, ctx, given, domain_label) - domain_requested = bool(ctx.domain_name and str(ctx.domain_name).strip()) or ( - ctx.domain_id is not None and str(ctx.domain_id).strip() not in {"", "0"} - ) - if domain_requested and ctx.preset_key is None: - resolved.warnings.insert( - 0, - "Cosmos3 action domain preset was not found for " - f"domain_name={ctx.domain_name!r}, domain_id={ctx.domain_id!r}; " - "using generic action defaults for omitted fields.", - ) - if domain_requested and resolved["raw_action_dim"] is None: - resolved.warnings.append( - "Cosmos3 has no canonical action width for " - f"domain_name={ctx.domain_name!r}, domain_id={ctx.domain_id!r}; " - "pass raw_action_dim explicitly for policy/inverse_dynamics." - ) - resolved.values["action_resolution"] = Resolved( - normalize_action_resolution(resolved["action_resolution"]), - resolved.source_of("action_resolution"), - ) - if resolved["raw_action_dim"] is not None and int(resolved["raw_action_dim"]) <= 0: - raise ValueError( - f"Cosmos3 raw_action_dim must be positive, got {resolved['raw_action_dim']}." +class ParameterResolver: + """The settings of one Cosmos3 checkpoint, and the per-request resolution against + them. Build it with :meth:`from_checkpoint`, or from a :class:`Cosmos3Config`.""" + + def __init__(self, family: str, config: Cosmos3Config, *, log: bool = True) -> None: + self.family = family + self._config = config + self._log = log + sigmas = config.sampling.fixed_sigmas + self.distilled_steps: Optional[int] = None if sigmas is None else len(sigmas) + self.policy_defaults: Mapping[str, Any] = resolve_checkpoint_policy_defaults( + config.action.policy ) - if int(resolved["action_chunk_size"]) <= 0: - raise ValueError( - f"Cosmos3 action_chunk_size must be positive, got {resolved['action_chunk_size']}." + """The normalized ``checkpoint.json`` policy block; empty for non-policy checkpoints.""" + + @classmethod + def from_checkpoint( + cls, checkpoint_dir: str, family: str, pretrained_config: Any = None, *, log: bool = True + ) -> "ParameterResolver": + """Read the checkpoint's files. Raises on a required setting nothing declares, + or on a disagreement between files that is an error.""" + return cls(family, resolve_checkpoint(checkpoint_dir, pretrained_config), log=log) + + def config(self) -> Cosmos3Config: + """The checkpoint-level settings, defaults filled where the files are silent.""" + return self._config + + def describe(self) -> str: + return self._config.describe() + + def request(self, params: RequestParameters) -> RequestParameters: + """``params`` with every field this mode uses filled: the caller's value where + set, else the checkpoint, else the tables and defaults. ``domain_name`` and + ``policy_active`` come back reconciled with the policy block. Recommendations + the caller's values differ from are logged as warnings.""" + resolved, warnings, origins = self._resolve(params) + if self._log: + for warning in warnings: + logger.warning(warning) + logger.debug( + "Cosmos3 request fields: " + ", ".join(f"{k} [{v}]" for k, v in origins.items()) + ) + return resolved + + def _resolve(self, p: RequestParameters) -> tuple[RequestParameters, list[str], dict[str, str]]: + in_action = p.action_mode is not None + policy_active = False + warnings: list[str] = [] + if in_action and p.action_mode == ACTION_MODE_POLICY and self.policy_defaults: + domain_name, policy_active, warnings = policy_domain( + self.policy_defaults, p.domain_name, p.domain_id + ) + p = replace(p, domain_name=domain_name) + policy: Mapping[str, Any] = self.policy_defaults if policy_active else {} + + mode = "action" if in_action else p.mode + table = COSMOS3_GENERATION_DEFAULTS.get((self.family, mode), {}) + video_table = COSMOS3_GENERATION_DEFAULTS.get((self.family, "video"), {}) + + def tables(key: str) -> Any: + # An explicit ``None`` in the mode's table switches the field off for that + # mode; only a key the table does not mention falls through to the video table. + return table[key] if key in table else video_table.get(key) + + preset_key = canonical_domain_preset_key(p.domain_name, p.domain_id) if in_action else None + preset: Mapping[str, Any] = COSMOS3_DOMAIN_PRESETS.get(preset_key, {}) if preset_key else {} + label = ( + repr(preset_key) + if preset_key + else ( + f"domain_name={p.domain_name!r}, domain_id={p.domain_id!r}" + if in_action + else repr(mode) + ) ) - if float(resolved["frame_rate"]) <= 0.0: - raise ValueError(f"Cosmos3 frame_rate must be positive, got {resolved['frame_rate']}.") - if float(resolved["action_fps"]) <= 0.0: - raise ValueError(f"Cosmos3 action_fps must be positive, got {resolved['action_fps']}.") - return resolved - - -def resolve_domain_action_config( - *, - domain_name: str | None = None, - domain_id: str | int | None = None, - raw_action_dim: int | None = None, - action_chunk_size: int | None = None, - action_resolution: int | None = None, - frame_rate: float | None = None, - action_fps: float | None = None, - use_state: bool | None = None, - checkpoint_policy_defaults: Mapping[str, Any] | None = None, -) -> dict[str, Any]: - """The action fields a request runs with, by :data:`ACTION_RULES`: the caller's - value, else the checkpoint policy block when its domain is the request's, else the - code's width table, else the domain preset, else the generic action defaults. - ``num_frames`` is always ``action_chunk_size + 1``. A caller's value that differs - from the recommendation is reported in ``warnings``. - """ - policy_defaults = dict(checkpoint_policy_defaults or {}) - domain_name, policy_active, warnings = policy_domain(policy_defaults, domain_name, domain_id) - ctx = RequestContext( - family="", - action_mode="action", - policy_defaults=policy_defaults, - policy_active=policy_active, - domain_name=domain_name, - domain_id=None if domain_id is None else int(domain_id), - ) - resolved = resolve_action( - ctx, - raw_action_dim=raw_action_dim, - action_chunk_size=action_chunk_size, - action_resolution=action_resolution, - frame_rate=frame_rate, - action_fps=action_fps, - use_state=use_state, - ) - width = resolved["raw_action_dim"] - return { - "domain_name": domain_name, - "checkpoint_policy_domain_matches": policy_active, - "raw_action_dim": None if width is None else int(width), - "action_chunk_size": int(resolved["action_chunk_size"]), - "action_resolution": resolved["action_resolution"], - "frame_rate": float(resolved["frame_rate"]), - "action_fps": float(resolved["action_fps"]), - "use_state": bool(resolved["use_state"]), - "num_frames": int(resolved["num_frames"]), - "preset_key": ctx.preset_key, - "warnings": warnings + resolved.warnings, - "sources": {name: r.source for name, r in resolved.values.items()}, - } + + distilled_guidance = None if self.distilled_steps is None else DISTILLED_GUIDANCE_SCALE + candidates: dict[str, Candidates] = { + "num_inference_steps": Candidates( + p.num_inference_steps, + self.distilled_steps + if self.distilled_steps is not None + else policy.get("num_inference_steps"), + fallback=tables("num_inference_steps"), + ), + "guidance_scale": Candidates( + p.guidance_scale, + distilled_guidance + if distilled_guidance is not None + else policy.get("guidance_scale"), + fallback=tables("guidance_scale"), + ), + "guidance_interval": Candidates( + p.guidance_interval, + policy.get("guidance_interval"), + fallback=tables("guidance_interval"), + ), + "flow_shift": Candidates(p.flow_shift, self._config.sampling.flow_shift), + "max_sequence_length": Candidates( + p.max_sequence_length, fallback=tables("max_sequence_length") + ), + } + if in_action: + known_width = resolve_raw_action_dim(domain_name=p.domain_name, domain_id=p.domain_id) + candidates["raw_action_dim"] = Candidates( + p.raw_action_dim, policy.get("raw_action_dim"), known_width + ) + candidates["action_chunk_size"] = Candidates( + p.action_chunk_size, + policy.get("action_chunk_size"), + preset.get("action_chunk_size"), + tables("action_chunk_size"), + ) + candidates["action_resolution"] = Candidates( + p.action_resolution, + policy.get("action_resolution"), + preset.get("action_resolution"), + ACTION_RESOLUTION_DEFAULT, + ) + candidates["frame_rate"] = Candidates( + p.frame_rate, + policy.get("frame_rate"), + preset.get("frame_rate"), + tables("frame_rate"), + ) + frame_rate = candidates["frame_rate"].pick() + candidates["action_fps"] = Candidates( + p.action_fps, + policy.get("action_fps"), + fallback=None if frame_rate is None else frame_rate.value, + ) + candidates["use_state"] = Candidates( + p.use_state, policy.get("use_state"), fallback=False + ) + chunk = candidates["action_chunk_size"].pick() + candidates["num_frames"] = Candidates( + p.num_frames, fallback=None if chunk is None else int(chunk.value) + 1 + ) + else: + candidates["height"] = Candidates(p.height, fallback=tables("height")) + candidates["width"] = Candidates(p.width, fallback=tables("width")) + candidates["frame_rate"] = Candidates(p.frame_rate, fallback=tables("frame_rate")) + candidates["num_frames"] = Candidates(p.num_frames, fallback=tables("num_frames")) + + picked: dict[str, Optional[Resolved]] = {name: c.pick() for name, c in candidates.items()} + # Transfer never prepends the system prompt and video-to-video always does. Those + # rules come from the kind of request, so they sit between the request and the + # checkpoint. + workflow_prompt = False if p.is_transfer else (True if p.is_v2v else None) + if p.use_system_prompt is None and workflow_prompt is not None: + picked["use_system_prompt"] = Resolved( + workflow_prompt, WORKFLOW, "transfer" if p.is_transfer else "video-to-video" + ) + else: + picked["use_system_prompt"] = select( + p.use_system_prompt, self._config.prompt.use_system_prompt + ) + unused = CANVAS_FIELDS if in_action else ACTION_FIELDS + picked.update({name: None for name in unused}) + + values: dict[str, Any] = {} + origins: dict[str, str] = {} + for name, item in picked.items(): + if item is None: + if name in OPTIONAL_FIELDS or name in unused: + values[name] = None + origins[name] = UNSET + continue + raise ValueError(f"Cosmos3 request: cannot resolve {name} for {label}.") + values[name] = item.value + origins[name] = item.origin + + if in_action: + # In action mode the policy block and the presets are recommendations: a + # caller's value that differs from them is worth a warning. + for name, c in candidates.items(): + suggestion = c.checkpoint if c.checkpoint is not None else c.recommended + if ( + suggestion is not None + and origins[name] == REQUEST + and not same_value(values[name], suggestion) + ): + warnings.append( + f"Cosmos3 {name}={values[name]} differs from recommended " + f"{suggestion} for domain {label}." + ) + domain_requested = bool(p.domain_name and str(p.domain_name).strip()) or ( + p.domain_id is not None and str(p.domain_id).strip() not in {"", "0"} + ) + if domain_requested and preset_key is None: + warnings.insert( + 0, + "Cosmos3 action domain preset was not found for " + f"domain_name={p.domain_name!r}, domain_id={p.domain_id!r}; " + "using generic action defaults for omitted fields.", + ) + if domain_requested and values["raw_action_dim"] is None: + warnings.append( + "Cosmos3 has no canonical action width for " + f"domain_name={p.domain_name!r}, domain_id={p.domain_id!r}; " + "pass raw_action_dim explicitly for policy/inverse_dynamics." + ) + values["action_resolution"] = normalize_action_resolution(values["action_resolution"]) + if values["raw_action_dim"] is not None and int(values["raw_action_dim"]) <= 0: + raise ValueError( + f"Cosmos3 raw_action_dim must be positive, got {values['raw_action_dim']}." + ) + for name in ("action_chunk_size", "frame_rate", "action_fps"): + if float(values[name]) <= 0.0: + raise ValueError(f"Cosmos3 {name} must be positive, got {values[name]}.") + resolved = replace(p, policy_active=policy_active, preset_key=preset_key, **values) + return resolved, warnings, origins diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index 02834b441954..c9b5fbdba8c0 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -79,20 +79,13 @@ def tqdm(iterable, **kwargs): COSMOS3_GENERATION_DEFAULTS, COSMOS3_IMAGE_RESOLUTION_TEMPLATE, COSMOS3_T2I_SYSTEM_PROMPT, - COSMOS3_V2V_FLOW_SHIFT, _normalize_condition_video_keep, _normalize_condition_video_latent_indexes, find_closest_target_size, - resolve_checkpoint_policy_defaults, ) from .guardrails import check_video_safety, download_guardrail_checkpoint from .negative_prompt import COSMOS3_VIDEO_NEGATIVE_PROMPT -from .parameter_resolver import ( - RequestContext, - resolve_checkpoint, - resolve_domain_action_config, - resolve_generation, -) +from .parameter_resolver import ParameterResolver, RequestParameters from .sampling import DISTILLED_GUIDANCE_SCALE, Cosmos3SamplingPolicy, load_scheduler from .sound_tokenizer import LatentAutoEncoderV2 from .transfer import ( @@ -134,24 +127,6 @@ def default_negative_prompt(output_type: str) -> str: ) -def _resolve_use_system_prompt( - requested: Optional[bool], - *, - checkpoint_default: bool, - is_v2v: bool, - is_transfer: bool, -) -> bool: - """The request rule: the caller's value, else transfer never (reference parity), - else video-to-video always, else the checkpoint's declaration.""" - ctx = RequestContext( - family="", - checkpoint_use_system_prompt=checkpoint_default, - is_v2v=is_v2v, - is_transfer=is_transfer, - ) - return bool(resolve_generation(ctx, use_system_prompt=requested)["use_system_prompt"]) - - def _validate_request_compatibility( *, output_type: str, @@ -417,16 +392,12 @@ def __init__(self, pipeline_config): # only when the transformer config declares and loads its action heads. self.has_action_weights = False self.action_gen = False - self.checkpoint_policy_defaults: dict[str, Any] = {} - # Pre-load placeholder; load_standard_components derives the real - # policy from the checkpoint's scheduler via from_scheduler(). + # Pre-load placeholder; load_standard_components derives the real policy + # from the checkpoint's scheduler. self.sampling = Cosmos3SamplingPolicy() - # Independent of the extra-param spec, whose default stays None so an - # omitted value survives the executor's default merge and reaches - # forward() as "unset". - self.default_use_system_prompt = False self.use_native_flow_schedule = False self.family = resolve_arch_recipe(primary_pretrained_config).name + self._resolver: Optional[ParameterResolver] = None # Schedulers are identified by their resolved (flow_shift, karras) pair self._scheduler_cache: dict = {} if getattr( @@ -444,54 +415,17 @@ def __init__(self, pipeline_config): super().__init__(pipeline_config) - def _mode_table(self, output_type: str, action_mode: Optional[str] = None) -> dict: - """The family's generation table for the request's output and action mode. The - request rules read it through :func:`resolve_generation`; the pipeline reads - only ``flow_shift`` from it directly.""" - request_mode = "action" if normalize_action_mode(action_mode) is not None else output_type - return COSMOS3_GENERATION_DEFAULTS[(self.family, request_mode)] + @property + def resolver(self) -> ParameterResolver: + """The checkpoint's settings; available once ``load_standard_components`` has + read the checkpoint.""" + if self._resolver is None: + raise RuntimeError("Cosmos3 pipeline: the checkpoint has not been loaded yet.") + return self._resolver - def _request_context( - self, - output_type: str = "video", - *, - action_mode: Optional[str] = None, - checkpoint_policy_domain_matches: bool = True, - ) -> RequestContext: - normalized_action_mode = normalize_action_mode(action_mode) - policy_defaults = getattr(self, "checkpoint_policy_defaults", {}) or {} - fixed_sigmas = getattr(self.sampling, "fixed_sigmas", None) - return RequestContext( - family=self.family, - mode=str(output_type).lower(), - action_mode=normalized_action_mode, - distilled_steps=None if not fixed_sigmas else len(fixed_sigmas), - policy_defaults=policy_defaults, - policy_active=( - normalized_action_mode == ACTION_MODE_POLICY - and checkpoint_policy_domain_matches - and bool(policy_defaults) - ), - checkpoint_use_system_prompt=getattr(self, "default_use_system_prompt", False), - ) - - def _resolve_generation_params( - self, - output_type: str, - *, - action_mode: Optional[str] = None, - checkpoint_policy_domain_matches: bool = True, - **values, - ) -> dict: - """Fill ``None`` values by the request rules: the distilled schedule, then the - policy block, then the mode table, then the video table.""" - ctx = self._request_context( - output_type, - action_mode=action_mode, - checkpoint_policy_domain_matches=checkpoint_policy_domain_matches, - ) - resolved = resolve_generation(ctx, **values) - return {name: resolved[name] for name in values} + @resolver.setter + def resolver(self, resolver: ParameterResolver) -> None: + self._resolver = resolver def _log_envelope_advisory( self, @@ -620,12 +554,13 @@ def load_standard_components( # The transformer config is read from the checkpoint directory, not from the # loaded pretrained config, whose omitted keys the transformer has already filled. - self.cosmos3_config = resolve_checkpoint(checkpoint_dir) - logger.info(self.cosmos3_config.describe()) - cfg = self.cosmos3_config - self.default_use_system_prompt = cfg.prompt.use_system_prompt + self.resolver = ParameterResolver.from_checkpoint( + checkpoint_dir, self.family, log=self.rank == 0 + ) + cfg = self.resolver.config() + if self.rank == 0: + logger.info(cfg.describe()) self.use_native_flow_schedule = cfg.sampling.native_flow_schedule - self.checkpoint_policy_defaults = resolve_checkpoint_policy_defaults(cfg.action.policy) if self.audio_gen and PipelineComponent.SOUND_TOKENIZER not in skip_components: logger.info("Loading audio tokenizer...") @@ -650,6 +585,11 @@ def load_standard_components( self.vae_scale_factor_spatial = cfg.vae.spatial_compression if PipelineComponent.VAE not in skip_components: + if self.vae_scale_factor_spatial is None: + raise ValueError( + f"Cosmos3 config: {checkpoint_dir} ships a VAE but vae/config.json declares " + "no scale_factor_spatial." + ) logger.info("Loading VAE...") vae_device = ( torch.device("cpu") @@ -718,7 +658,8 @@ def load_standard_components( self.safety_checker = CosmosSafetyChecker() self.safety_checker.to(device) - self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) + if self.vae_scale_factor_spatial is not None: + self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) @property def default_warmup_resolutions(self): @@ -748,11 +689,10 @@ def default_generation_params(self): request runs, the distilled schedule included. A request that selects another mode re-resolves them in ``infer()``, which tells a merged default from a caller-supplied value via ``model_fields_set``. Key membership also declares - these fields supported during request validation. ``flow_shift`` is - pipeline-internal, not a request field. + these fields supported during request validation. """ - fields = [k for k in self._mode_table("video") if k != "flow_shift"] - return self._resolve_generation_params("video", **{f: None for f in fields}) + params = self.resolver.request(RequestParameters(mode="video")) + return {k: getattr(params, k) for k in COSMOS3_GENERATION_DEFAULTS[(self.family, "video")]} def classify_request_failure(self, exc: BaseException) -> Optional[str]: """Cosmos3 rejects unusable request content with ``ValueError`` and @@ -765,10 +705,9 @@ def extra_param_specs(self): # ``use_system_prompt`` keeps its None default here on purpose: the # executor materializes these into every request, so publishing the # checkpoint's boolean would destroy "unset" before forward() can - # resolve it by mode. The checkpoint value is exposed separately as - # ``default_use_system_prompt``. + # resolve it by mode. The checkpoint value lives in the resolver. specs = dict(COSMOS3_EXTRA_SPECS) - if getattr(self, "checkpoint_policy_defaults", {}): + if self.resolver.policy_defaults: specs["action_mode"] = specs["action_mode"].model_copy( update={"default": ACTION_MODE_POLICY} ) @@ -849,31 +788,17 @@ def _scheduler_for( return cache[key] def _cacheable_flow_shifts(self) -> frozenset: - """Flow shifts reachable without a caller override, for this checkpoint. - - Derived rather than fixed so a new family or mode table is picked up - automatically: the per-mode generation tables, the checkpoint's own - shift, V2V's stronger shift, and the per-hint transfer presets. Entries - are still created lazily, so a mode that is never served never builds - one. - - Transfer's shifts are included even though every hint currently declares - the same value as V2V: tuning one off that value would otherwise drop it - out of the cache silently and rebuild a scheduler on every request for - that hint. - """ + """Flow shifts reachable without a caller override, for this checkpoint: + the checkpoint's own shift and the per-hint transfer presets. Entries are + still created lazily, so a configuration that is never served never builds + one.""" cached = getattr(self, "_cacheable_flow_shifts_cache", None) if cached is not None: return cached - shifts = {COSMOS3_V2V_FLOW_SHIFT} + shifts: set[float] = set() checkpoint_shift = getattr(self.sampling, "checkpoint_flow_shift", None) if checkpoint_shift is not None: shifts.add(float(checkpoint_shift)) - for mode in ("video", "image"): - table = COSMOS3_GENERATION_DEFAULTS.get((self.family, mode)) or {} - mode_shift = table.get("flow_shift") - if mode_shift is not None: - shifts.add(float(mode_shift)) for hint_defaults in TRANSFER_DEFAULTS.values(): hint_shift = hint_defaults.get("flow_shift") if hint_shift is not None: @@ -931,7 +856,7 @@ def as_given(field_name): refs_v = req.params.video_reference video = refs_v[0].content if refs_v else None action_mode = extra_params.get("action_mode") - if action_mode is None and getattr(self, "checkpoint_policy_defaults", {}): + if action_mode is None and self.resolver.policy_defaults: # A checkpoint with a policy manifest has one supported action # workflow, so callers need not restate it on every request. action_mode = ACTION_MODE_POLICY @@ -950,7 +875,7 @@ def as_given(field_name): frame_rate = as_given("frame_rate") else: # Source-derived sizes go in as non-None, so they win over the mode - # table: _resolve_generation_params only fills what is still None. + # table: the resolver only fills what is still None. height, width = as_given("height"), as_given("width") frame_rate = req.params.frame_rate wants_source_size = "height" not in specified and "width" not in specified @@ -988,17 +913,19 @@ def as_given(field_name): f"Cosmos3 following the source: " f"{width}x{height} (WxH) @ {frame_rate} fps" ) - resolved = self._resolve_generation_params( - "image" if is_t2i else "video", - height=height, - width=width, - num_inference_steps=as_given("num_inference_steps"), - guidance_scale=as_given("guidance_scale"), + resolved = self.resolver.request( + RequestParameters( + mode="image" if is_t2i else "video", + height=height, + width=width, + num_inference_steps=as_given("num_inference_steps"), + guidance_scale=as_given("guidance_scale"), + ) ) - height = resolved["height"] - width = resolved["width"] - num_inference_steps = resolved["num_inference_steps"] - guidance_scale = resolved["guidance_scale"] + height = resolved.height + width = resolved.width + num_inference_steps = resolved.num_inference_steps + guidance_scale = resolved.guidance_scale refs_i = req.params.image_reference @@ -1327,25 +1254,26 @@ def post_step_fn(latents: torch.Tensor, extra_stream_latents): # VAE decode # ========================================================================= - def _scale_latents_for_decode(self, latents): - """Latents as the VAE decoder expects them: the checkpoint's per-channel - statistics undone, or its scaling factor.""" + def _decode_latents_raw(self, latents): latents = latents.to(self.vae.dtype) + if hasattr(self.vae.config, "latents_mean") and hasattr(self.vae.config, "latents_std"): - stats = getattr(self, "_latents_stats", None) - if stats is None or stats[0].device != latents.device: - stats = tuple( - torch.tensor(values).view(1, -1, 1, 1, 1).to(latents.device, self.vae.dtype) - for values in (self.vae.config.latents_std, self.vae.config.latents_mean) + if not hasattr(self, "_latents_mean"): + self._latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, -1, 1, 1, 1) + .to(self.device, self.vae.dtype) + ) + self._latents_std = ( + torch.tensor(self.vae.config.latents_std) + .view(1, -1, 1, 1, 1) + .to(self.device, self.vae.dtype) ) - self._latents_stats = stats - std, mean = stats - return (latents * std) + mean - scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0) - return latents / scaling_factor + latents = (latents * self._latents_std) + self._latents_mean + else: + scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0) + latents = latents / scaling_factor - def _decode_latents_raw(self, latents): - latents = self._scale_latents_for_decode(latents) with self.offloader.context_if_requested(PipelineComponent.VAE.value): return self.vae.decode(latents, return_dict=False)[0] @@ -1670,61 +1598,42 @@ def _resolve_request( transfer_config=transfer_config, ) - action_cfg: dict[str, Any] = {} - checkpoint_policy_domain_matches = True - if do_action: - checkpoint_policy_defaults = ( - getattr(self, "checkpoint_policy_defaults", {}) - if normalized_action_mode == ACTION_MODE_POLICY - else {} - ) - action_cfg = resolve_domain_action_config( + is_v2v = video is not None and not is_t2i and not do_action + params = self.resolver.request( + RequestParameters( + mode=output_type, + action_mode=normalized_action_mode, + is_v2v=is_v2v, + is_transfer=transfer_config is not None, domain_name=domain_name, domain_id=domain_id, + height=None if do_action else height, + width=None if do_action else width, + num_frames=None if do_action else num_frames, + frame_rate=frame_rate, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + flow_shift=flow_shift, + max_sequence_length=max_sequence_length, + use_system_prompt=use_system_prompt, raw_action_dim=raw_action_dim, action_chunk_size=action_chunk_size, action_resolution=action_resolution, - frame_rate=frame_rate, action_fps=action_fps, use_state=use_state, - checkpoint_policy_defaults=checkpoint_policy_defaults, - ) - checkpoint_policy_domain_matches = action_cfg["checkpoint_policy_domain_matches"] - if self.rank == 0: - logger.debug( - "Cosmos3 action fields: " - + ", ".join(f"{k} [{v}]" for k, v in action_cfg["sources"].items()) - ) - - mode_table = self._mode_table(output_type, normalized_action_mode) - values_to_resolve = { - "num_inference_steps": num_inference_steps, - "guidance_scale": guidance_scale, - "guidance_interval": None, - "max_sequence_length": max_sequence_length, - } - if not do_action: - values_to_resolve.update( - height=height, - width=width, - num_frames=num_frames, - frame_rate=frame_rate, ) - resolved = self._resolve_generation_params( - output_type, - action_mode=normalized_action_mode, - checkpoint_policy_domain_matches=checkpoint_policy_domain_matches, - **values_to_resolve, ) - num_inference_steps = resolved["num_inference_steps"] - guidance_scale = resolved["guidance_scale"] - guidance_interval = resolved["guidance_interval"] - max_sequence_length = resolved["max_sequence_length"] + num_inference_steps = params.num_inference_steps + guidance_scale = params.guidance_scale + guidance_interval = params.guidance_interval + flow_shift = params.flow_shift + max_sequence_length = params.max_sequence_length + use_system_prompt = params.use_system_prompt if not do_action: - height = resolved["height"] - width = resolved["width"] - num_frames = resolved["num_frames"] - frame_rate = resolved["frame_rate"] + height = params.height + width = params.width + num_frames = params.num_frames + frame_rate = params.frame_rate self.sampling.validate_request(num_inference_steps, guidance_scale) if not do_action: @@ -1737,54 +1646,37 @@ def _resolve_request( max_sequence_length=max_sequence_length, ) - is_v2v = video is not None and not is_t2i and not do_action - use_system_prompt = _resolve_use_system_prompt( - use_system_prompt, - checkpoint_default=getattr(self, "default_use_system_prompt", False), - is_v2v=is_v2v, - is_transfer=transfer_config is not None, - ) - resolved_action_fps: Optional[float] = None if is_t2i: num_frames = 1 enable_audio = False if do_action: - if self.rank == 0: - for warning in action_cfg["warnings"]: - logger.warning(warning) - if action_cfg["preset_key"] is not None: - logger.info( - f"Cosmos3 action domain preset {action_cfg['preset_key']!r}: " - f"raw_action_dim={action_cfg['raw_action_dim']}, " - f"action_chunk_size={action_cfg['action_chunk_size']}, " - f"action_resolution={action_cfg['action_resolution']}, " - f"frame_rate={action_cfg['frame_rate']:.1f}, " - f"action_fps={action_cfg['action_fps']:.1f}, " - f"use_state={action_cfg['use_state']}, " - f"num_frames={action_cfg['num_frames']}" - ) - - domain_name = action_cfg["domain_name"] - raw_action_dim = action_cfg["raw_action_dim"] - action_chunk_size = action_cfg["action_chunk_size"] - action_resolution = action_cfg["action_resolution"] - num_frames = action_cfg["num_frames"] - frame_rate = action_cfg["frame_rate"] - resolved_action_fps = action_cfg["action_fps"] - use_state = action_cfg["use_state"] + if self.rank == 0 and params.preset_key is not None: + logger.info( + f"Cosmos3 action domain preset {params.preset_key!r}: " + f"raw_action_dim={params.raw_action_dim}, " + f"action_chunk_size={params.action_chunk_size}, " + f"action_resolution={params.action_resolution}, " + f"frame_rate={params.frame_rate:.1f}, " + f"action_fps={params.action_fps:.1f}, " + f"use_state={params.use_state}, " + f"num_frames={params.num_frames}" + ) + domain_name = params.domain_name + raw_action_dim = params.raw_action_dim + action_chunk_size = params.action_chunk_size + action_resolution = params.action_resolution + num_frames = params.num_frames + frame_rate = params.frame_rate + resolved_action_fps = params.action_fps + use_state = params.use_state enable_audio = False - mode_shift = mode_table.get("flow_shift") - if mode_shift is None: - mode_shift = self.sampling.checkpoint_flow_shift - if is_v2v: - target_shift = COSMOS3_V2V_FLOW_SHIFT if flow_shift is None else flow_shift - target_karras = False - else: - target_shift = mode_shift if flow_shift is None else flow_shift - target_karras = None + # Video-to-video samples on the uniform sigma grid; every mode runs at the + # resolved flow shift (the request's, else the checkpoint scheduler's). + target_shift = flow_shift + target_karras = False if is_v2v else None if transfer_config is None: self.scheduler = self._scheduler_for(target_shift, target_karras) if getattr(self, "audio_scheduler", None) is not None: @@ -2956,11 +2848,8 @@ def _forward_transfer( f"{DISTILLED_GUIDANCE_SCALE}; per-hint guidance presets do not apply." ) guidance_scale = DISTILLED_GUIDANCE_SCALE - flow_shift_target = float( - transfer_config.flow_shift - if transfer_config.flow_shift is not None - else COSMOS3_V2V_FLOW_SHIFT - ) + # The hint preset's shift, or the request's; ``None`` keeps the checkpoint's. + flow_shift_target = transfer_config.flow_shift max_sequence_length = max_sequence_length or COSMOS3_720P_PARAMS["max_sequence_length"] self._guidance_scale = guidance_scale self._num_timesteps = num_inference_steps diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/sampling.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/sampling.py index 4534e57aa45b..e45a02bb09f2 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/sampling.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/sampling.py @@ -54,15 +54,6 @@ def _config_get(config: Any, key: str, default: Any = None) -> Any: return getattr(config, key, default) -def _resolve_distilled_sigmas(scheduler_config: Any) -> "tuple[float, ...] | None": - """``fixed_step_sampler_config.t_list`` as floats, or None (base checkpoints).""" - fixed_step_cfg = _config_get(scheduler_config, "fixed_step_sampler_config") - t_list = _config_get(fixed_step_cfg, "t_list") if fixed_step_cfg else None - if not t_list: - return None - return tuple(float(sigma) for sigma in t_list) - - def load_scheduler(checkpoint_dir: str, subfolder: str = "scheduler") -> Any: """Instantiate the scheduler class the checkpoint declares. @@ -118,16 +109,13 @@ def from_scheduler( native_flow_schedule: bool = False, fixed_sigmas: "tuple[float, ...] | None" = None, ) -> "Cosmos3SamplingPolicy": - """Derive the policy from a loaded scheduler's config. + """Derive the policy from a loaded scheduler and the checkpoint's schedule. - Valid recipes: UniPC without fixed sigmas (base) and stochastic - FlowMatchEuler with a nonempty ``fixed_step_sampler_config.t_list`` - (distilled); anything else fails here, at load time. ``fixed_sigmas`` is the - schedule as the checkpoint resolver read it; when omitted it is read from the - scheduler's own config. + ``fixed_sigmas`` is the distilled schedule as the checkpoint resolver read it + (``fixed_step_sampler_config.t_list``), ``None`` on base checkpoints. Valid + recipes: UniPC without fixed sigmas (base) and stochastic FlowMatchEuler with + fixed sigmas (distilled); anything else fails here, at load time. """ - if fixed_sigmas is None: - fixed_sigmas = _resolve_distilled_sigmas(scheduler.config) is_unipc = isinstance(scheduler, UniPCMultistepScheduler) is_flow_match = isinstance(scheduler, FlowMatchEulerDiscreteScheduler) diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_cosmos3.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_cosmos3.py index 3831a7181a12..9df6cddbd241 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_cosmos3.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_cosmos3.py @@ -163,6 +163,7 @@ def _run_cosmos3_lpips_pipeline( width=COSMOS3_LPIPS_WIDTH, num_inference_steps=COSMOS3_LPIPS_NUM_INFERENCE_STEPS, output_type="video", + flow_shift=None, ): """Run a Cosmos3 pipeline (default setting, VANILLA attn, compile-off). @@ -224,6 +225,7 @@ def _run_cosmos3_lpips_pipeline( video=video, image=image, output_type=output_type, + flow_shift=flow_shift, ) if result is None: return None @@ -269,8 +271,10 @@ def _generate_cosmos3_v2v_lpips_frame(output_path): """Generate the Cosmos3-Nano video-to-video LPIPS sample (free frame only).""" from tensorrt_llm.media.encoding import save_image + # The golden was cut at flow shift 10.0, the value the cosmos-framework cookbook + # sends for video-to-video; an omitted value runs at the checkpoint scheduler's. video = _run_cosmos3_lpips_pipeline( - COSMOS3_LPIPS_V2V_NUM_FRAMES, video=_cosmos3_v2v_lpips_reference_bytes() + COSMOS3_LPIPS_V2V_NUM_FRAMES, video=_cosmos3_v2v_lpips_reference_bytes(), flow_shift=10.0 ) assert video is not None, "Cosmos3-Nano V2V LPIPS run produced no video" # video is (B, T, H, W, C); take the free frame -> (H, W, C) for save_image. @@ -281,7 +285,9 @@ def _generate_cosmos3_lpips_image(output_path): """Generate the Cosmos3-Nano text-to-image LPIPS sample (single frame).""" from tensorrt_llm.media.encoding import save_image - video = _run_cosmos3_lpips_pipeline(COSMOS3_LPIPS_T2I_NUM_FRAMES) + # The golden was cut at flow shift 3.0, the value the cosmos-framework cookbook + # sends for text-to-image; an omitted value runs at the checkpoint scheduler's. + video = _run_cosmos3_lpips_pipeline(COSMOS3_LPIPS_T2I_NUM_FRAMES, flow_shift=3.0) assert video is not None, "Cosmos3-Nano T2I LPIPS run produced no frame" # video is (B, T, H, W, C); take the single frame -> (H, W, C) for save_image. save_image(video[0, 0], output_path) @@ -967,6 +973,7 @@ def test_cosmos3_edge_t2v_lpips_against_golden(_visual_gen_deps, request, tmp_pa width=832, num_frames=COSMOS3_EDGE_LPIPS_NUM_FRAMES, num_inference_steps=COSMOS3_EDGE_T2V_LPIPS_STEPS, + flow_shift=3.0, # the shift the golden was cut at; omitted, the checkpoint's applies guidance_scale=5.0, frame_rate=COSMOS3_EDGE_LPIPS_FRAME_RATE, ) @@ -1008,6 +1015,7 @@ def test_cosmos3_edge_i2v_lpips_against_golden(_visual_gen_deps, request, tmp_pa width=832, num_frames=COSMOS3_EDGE_LPIPS_NUM_FRAMES, num_inference_steps=COSMOS3_EDGE_I2V_LPIPS_STEPS, + flow_shift=3.0, # the shift the golden was cut at; omitted, the checkpoint's applies guidance_scale=5.0, frame_rate=COSMOS3_EDGE_LPIPS_FRAME_RATE, ) @@ -1044,6 +1052,7 @@ def test_cosmos3_edge_t2i_lpips_against_golden(request, tmp_path): height=640, width=640, num_inference_steps=COSMOS3_EDGE_T2I_LPIPS_STEPS, + flow_shift=3.0, # the shift the golden was cut at; omitted, the checkpoint's applies guidance_scale=4.0, output_type="image", ) diff --git a/tests/unittest/_torch/visual_gen/cosmos3_fixtures.py b/tests/unittest/_torch/visual_gen/cosmos3_fixtures.py new file mode 100644 index 000000000000..715031739781 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/cosmos3_fixtures.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Cosmos3 checkpoint settings built by hand, for tests that run a pipeline or a +resolver without a checkpoint on disk.""" + +from collections.abc import Mapping +from typing import Any, Optional + +from tensorrt_llm._torch.visual_gen.models.cosmos3.parameter_resolver import ( + DEFAULT, + Cosmos3ActionSettings, + Cosmos3ArchitectureSettings, + Cosmos3AudioSettings, + Cosmos3Config, + Cosmos3PromptSettings, + Cosmos3SamplingSettings, + Cosmos3VAESettings, + ParameterResolver, + Resolved, +) + + +def bare_config( + *, + fixed_sigmas: Optional[tuple[float, ...]] = None, + policy: Optional[Mapping[str, Any]] = None, + use_system_prompt: bool = False, + native_flow_schedule: bool = False, + flow_shift: Optional[float] = None, +) -> Cosmos3Config: + """A base-checkpoint :class:`Cosmos3Config` with the given facts overridden.""" + settings = { + "architecture.position_embedding_type": "unified_3d_mrope", + "architecture.max_position_embeddings": 262144, + "vae.temporal_compression": 4, + "vae.spatial_compression": 16, + "audio.temporal_compression": 1, + "sampling.fixed_sigmas": fixed_sigmas, + "sampling.native_flow_schedule": native_flow_schedule, + "sampling.flow_shift": flow_shift, + "prompt.use_system_prompt": use_system_prompt, + "action.policy": policy, + } + return Cosmos3Config( + architecture=Cosmos3ArchitectureSettings("unified_3d_mrope", 262144), + vae=Cosmos3VAESettings(temporal_compression=4, spatial_compression=16), + audio=Cosmos3AudioSettings(temporal_compression=1), + sampling=Cosmos3SamplingSettings( + fixed_sigmas=fixed_sigmas, + native_flow_schedule=native_flow_schedule, + flow_shift=flow_shift, + ), + prompt=Cosmos3PromptSettings(use_system_prompt=use_system_prompt), + action=Cosmos3ActionSettings(policy=policy), + resolved={k: Resolved(v, DEFAULT) for k, v in settings.items() if v is not None}, + ) + + +def bare_resolver(family: str, **facts: Any) -> ParameterResolver: + """A silent :class:`ParameterResolver` over :func:`bare_config`'s settings.""" + return ParameterResolver(family, bare_config(**facts), log=False) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py index f8a8ec19cfa0..ab7ef834e472 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py @@ -7,12 +7,16 @@ pytest tests/unittest/_torch/visual_gen/test_cosmos3_action.py -v """ +import dataclasses + import numpy as np import PIL.Image import pytest import torch +from cosmos3_fixtures import bare_resolver from tensorrt_llm._torch.visual_gen.models.cosmos3.action import ( + ACTION_MODE_POLICY, action_reference_frame_step, action_reference_size, crop_action_latent, @@ -37,9 +41,7 @@ resolve_checkpoint_policy_defaults, resolve_raw_action_dim, ) -from tensorrt_llm._torch.visual_gen.models.cosmos3.parameter_resolver import ( - resolve_domain_action_config, -) +from tensorrt_llm._torch.visual_gen.models.cosmos3.parameter_resolver import RequestParameters from tensorrt_llm._torch.visual_gen.models.cosmos3.transformer_cosmos3 import ( compute_mrope_position_ids_action, compute_mrope_position_ids_vision, @@ -48,6 +50,24 @@ pytestmark = pytest.mark.cosmos3 +def resolve_domain_action_config(checkpoint_policy=None, **given) -> dict: + """The action fields a request runs with, as a dict with the warnings and origins: + against the raw ``checkpoint.json`` policy block when one is given, else the domain + preset.""" + resolver = bare_resolver("qwen3", policy=checkpoint_policy) + params, warnings, origins = resolver._resolve( + RequestParameters( + action_mode=ACTION_MODE_POLICY if checkpoint_policy else "action", **given + ) + ) + return { + **dataclasses.asdict(params), + "checkpoint_policy_domain_matches": params.policy_active, + "warnings": warnings, + "origins": origins, + } + + class TestFindClosestTargetSize: @pytest.mark.parametrize( "input_h,input_w,action_resolution,expected", @@ -241,7 +261,7 @@ def test_edge_policy_metadata_resolves_reference_recipe(self): "action_fps": 15.0, } - cfg = resolve_domain_action_config(checkpoint_policy_defaults=policy) + cfg = resolve_domain_action_config(checkpoint_policy=self.METADATA) assert cfg["domain_name"] == "droid_lerobot" assert cfg["raw_action_dim"] == 8 assert cfg["action_chunk_size"] == 32 @@ -253,10 +273,9 @@ def test_edge_policy_metadata_resolves_reference_recipe(self): assert cfg["warnings"] == [] def test_explicit_non_policy_domain_does_not_inherit_droid_width(self): - policy = resolve_checkpoint_policy_defaults(self.METADATA) cfg = resolve_domain_action_config( domain_name="bridge_orig_lerobot", - checkpoint_policy_defaults=policy, + checkpoint_policy=self.METADATA, ) assert cfg["raw_action_dim"] == 10 assert cfg["action_chunk_size"] == 16 diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py index d26a67a3126d..28b62f1296c8 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py @@ -5,6 +5,7 @@ and its pipeline wiring: scheduler loading, recipe validation, generation defaults, mode resolution, and the guidance-1.0 denoise-loop contract.""" +import dataclasses import json import pickle from pathlib import Path @@ -13,15 +14,16 @@ import pytest import torch +from cosmos3_fixtures import bare_resolver from diffusers import FlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( COSMOS3_720P_PARAMS, COSMOS3_T2I_PARAMS, ) +from tensorrt_llm._torch.visual_gen.models.cosmos3.parameter_resolver import RequestParameters from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import ( Cosmos3OmniMoTPipeline, - _resolve_use_system_prompt, _validate_request_compatibility, ) from tensorrt_llm._torch.visual_gen.models.cosmos3.sampling import ( @@ -84,7 +86,7 @@ def _write_trained_config(checkpoint_dir: Path, use_system_prompt) -> None: def _distilled_policy() -> Cosmos3SamplingPolicy: scheduler = FlowMatchEulerDiscreteScheduler.from_config(DISTILLED_SCHEDULER_CONFIG) - return Cosmos3SamplingPolicy.from_scheduler(scheduler) + return Cosmos3SamplingPolicy.from_scheduler(scheduler, fixed_sigmas=DISTILLED_SIGMAS) def _base_scheduler() -> UniPCMultistepScheduler: @@ -96,7 +98,9 @@ def _base_policy() -> Cosmos3SamplingPolicy: def _bare_pipeline(**attrs) -> Cosmos3OmniMoTPipeline: - """A pipeline instance without heavyweight ``__init__``.""" + """A pipeline instance without heavyweight ``__init__``. ``default_use_system_prompt`` + and ``checkpoint_policy`` (the raw ``checkpoint.json`` block) go into the resolver, as + the checkpoint would put them; the resolver sees the distilled schedule of ``sampling``.""" pipeline = object.__new__(Cosmos3OmniMoTPipeline) defaults = dict( _device=torch.device("cpu"), @@ -104,13 +108,24 @@ def _bare_pipeline(**attrs) -> Cosmos3OmniMoTPipeline: action_gen=False, has_action_weights=False, sampling=Cosmos3SamplingPolicy(), - default_use_system_prompt=False, family=QWEN3_RECIPE.name, use_native_flow_schedule=False, ) defaults.update(attrs) defaults.setdefault("_scheduler_cache", {}) defaults.setdefault("_base_scheduler", defaults.get("scheduler")) + use_system_prompt = defaults.pop("default_use_system_prompt", False) + policy = defaults.pop("checkpoint_policy", None) + sigmas = getattr(defaults["sampling"], "fixed_sigmas", None) + defaults.setdefault( + "resolver", + bare_resolver( + defaults["family"], + fixed_sigmas=sigmas, + policy=policy, + use_system_prompt=use_system_prompt, + ), + ) for key, value in defaults.items(): setattr(pipeline, key, value) return pipeline @@ -176,23 +191,6 @@ def test_is_distilled(self): assert not _base_policy().is_distilled assert not Cosmos3SamplingPolicy().is_distilled - def test_diffusers_retains_unexpected_config_keys(self): - """Canary: diffusers must keep the unexpected fixed_step_sampler_config - key in scheduler.config; if an upgrade drops it, distilled detection - silently breaks.""" - scheduler = FlowMatchEulerDiscreteScheduler.from_config(DISTILLED_SCHEDULER_CONFIG) - policy = Cosmos3SamplingPolicy.from_scheduler(scheduler) - assert policy.fixed_sigmas == DISTILLED_SIGMAS - assert scheduler.config.stochastic_sampling is True - - def test_sigma_values_coerced_to_floats(self): - config = { - **DISTILLED_SCHEDULER_CONFIG, - "fixed_step_sampler_config": {"t_list": [1, "0.5"]}, - } - scheduler = FlowMatchEulerDiscreteScheduler.from_config(config) - assert Cosmos3SamplingPolicy.from_scheduler(scheduler).fixed_sigmas == (1.0, 0.5) - def test_num_steps(self): assert _distilled_policy().num_steps(2) == 4 assert _base_policy().num_steps(2) == 2 @@ -231,7 +229,7 @@ def test_t_list_on_unipc_raises(self): scheduler = UniPCMultistepScheduler.from_config(config) with pytest.raises(ValueError, match="Unsupported Cosmos3 sampling recipe"): - Cosmos3SamplingPolicy.from_scheduler(scheduler) + Cosmos3SamplingPolicy.from_scheduler(scheduler, fixed_sigmas=DISTILLED_SIGMAS) def test_unipc_with_declared_requirement_but_no_sigmas_raises(self): config = {**UNIPC_SCHEDULER_CONFIG, "fixed_step_requires_explicit_sigmas": True} @@ -249,7 +247,7 @@ def test_unipc_with_flag_and_sigmas_gets_unsupported_error(self): scheduler = UniPCMultistepScheduler.from_config(config) with pytest.raises(ValueError, match="Unsupported Cosmos3 sampling recipe"): - Cosmos3SamplingPolicy.from_scheduler(scheduler) + Cosmos3SamplingPolicy.from_scheduler(scheduler, fixed_sigmas=DISTILLED_SIGMAS) def test_flow_match_without_sigmas_raises(self): config = {k: v for k, v in DISTILLED_SCHEDULER_CONFIG.items()} @@ -272,7 +270,7 @@ def test_non_stochastic_fixed_schedule_raises(self, stochastic): scheduler = FlowMatchEulerDiscreteScheduler.from_config(config) with pytest.raises(ValueError, match="stochastic_sampling"): - Cosmos3SamplingPolicy.from_scheduler(scheduler) + Cosmos3SamplingPolicy.from_scheduler(scheduler, fixed_sigmas=DISTILLED_SIGMAS) def test_declared_non_sde_sample_type_raises(self): config = { @@ -285,7 +283,7 @@ def test_declared_non_sde_sample_type_raises(self): scheduler = FlowMatchEulerDiscreteScheduler.from_config(config) with pytest.raises(ValueError, match="sample_type"): - Cosmos3SamplingPolicy.from_scheduler(scheduler) + Cosmos3SamplingPolicy.from_scheduler(scheduler, fixed_sigmas=DISTILLED_SIGMAS) class TestValidateRequest: @@ -1029,19 +1027,16 @@ def test_resolved_guidance_interval_reaches_denoise( self, monkeypatch: pytest.MonkeyPatch ) -> None: pipeline, captured = self._wiring_pipeline() - original_resolver = pipeline._resolve_generation_params + original_request = pipeline.resolver.request expected_interval = (123.0, 456.0) - def resolve_with_interval( - mode: str, **values: int | float | None - ) -> dict[str, int | float | tuple[float, float] | None]: - assert values["guidance_interval"] is None - return { - **original_resolver(mode, **values), - "guidance_interval": expected_interval, - } + def request_with_interval(params: RequestParameters) -> RequestParameters: + assert params.guidance_interval is None + return dataclasses.replace( + original_request(params), guidance_interval=expected_interval + ) - monkeypatch.setattr(pipeline, "_resolve_generation_params", resolve_with_interval) + monkeypatch.setattr(pipeline.resolver, "request", request_with_interval) self._forward(pipeline, image=None) @@ -1067,11 +1062,11 @@ def test_checkpoint_declared_true(self, tmp_path): self._write_model_index(tmp_path, {"default_use_system_prompt": True}) pipeline = self._loaded_pipeline(tmp_path) - assert pipeline.default_use_system_prompt is True + assert pipeline.resolver.config().prompt.use_system_prompt is True # The published spec default stays None ("model decides"): the executor # materializes spec defaults into every request, so publishing the # checkpoint boolean here would destroy "unset" before forward() could - # resolve it by mode. The checkpoint value lives on the pipeline. + # resolve it by mode. The checkpoint value lives in the resolver. assert pipeline.extra_param_specs["use_system_prompt"].default is None # The shared spec table must stay untouched. from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import COSMOS3_EXTRA_SPECS @@ -1091,13 +1086,13 @@ def test_trained_config_wins_over_model_index(self, tmp_path, trained): self._write_model_index(tmp_path, {"default_use_system_prompt": not trained}) _write_trained_config(tmp_path, trained) pipeline = self._loaded_pipeline(tmp_path) - assert pipeline.default_use_system_prompt is trained + assert pipeline.resolver.config().prompt.use_system_prompt is trained def test_trained_config_without_key_keeps_model_index(self, tmp_path): self._write_model_index(tmp_path, {"default_use_system_prompt": True}) _write_trained_config(tmp_path, None) pipeline = self._loaded_pipeline(tmp_path) - assert pipeline.default_use_system_prompt is True + assert pipeline.resolver.config().prompt.use_system_prompt is True def _captured_use_system_prompt(self, pipeline, extra_params): captured = {} @@ -1183,15 +1178,11 @@ class TestSystemPromptResolution: def test_resolution_precedence( self, requested, checkpoint_default, is_v2v, is_transfer, expected ): - assert ( - _resolve_use_system_prompt( - requested, - checkpoint_default=checkpoint_default, - is_v2v=is_v2v, - is_transfer=is_transfer, - ) - is expected + resolver = bare_resolver(QWEN3_RECIPE.name, use_system_prompt=checkpoint_default) + resolved = resolver.request( + RequestParameters(is_v2v=is_v2v, is_transfer=is_transfer, use_system_prompt=requested) ) + assert resolved.use_system_prompt is expected class TestRequestCompatibility: @@ -1453,8 +1444,7 @@ def _merged_extra_params(self, extra_params): from tensorrt_llm._torch.visual_gen.executor import DiffusionExecutor from tensorrt_llm.visual_gen.params import VisualGenParams - pipeline = _bare_pipeline() - pipeline.default_use_system_prompt = False # base checkpoint + pipeline = _bare_pipeline() # base checkpoint: no system prompt pipeline.derive_output_size_from_reference = False executor = object.__new__(DiffusionExecutor) executor.pipeline = pipeline diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py b/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py index 8f7b9259aed2..8dcedcd9364d 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py @@ -19,6 +19,7 @@ import PIL.Image import pytest import torch +from cosmos3_fixtures import bare_resolver from diffusers import UniPCMultistepScheduler from utils.llm_data import get_checkpoint @@ -39,7 +40,10 @@ COSMOS3_GENERATION_DEFAULTS, COSMOS3_POLICY_SAMPLING_PARAMS, COSMOS3_T2I_PARAMS, - resolve_checkpoint_policy_defaults, +) +from tensorrt_llm._torch.visual_gen.models.cosmos3.parameter_resolver import ( + ParameterResolver, + RequestParameters, ) from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniMoTPipeline from tensorrt_llm._torch.visual_gen.models.cosmos3.sampling import Cosmos3SamplingPolicy @@ -770,9 +774,17 @@ def _bare_pipeline(family: str) -> Cosmos3OmniMoTPipeline: pipeline.action_gen = False pipeline.has_action_weights = family == NEMOTRON_DENSE_RECIPE.name pipeline.use_native_flow_schedule = family == NEMOTRON_DENSE_RECIPE.name + pipeline.resolver = bare_resolver(family) return pipeline +DROID_POLICY = {"action_chunk_size": 32, "conditioning_fps": 15.0, "domain_name": "droid_lerobot"} + + +def _policy_resolver(family: str) -> ParameterResolver: + return bare_resolver(family, policy=DROID_POLICY) + + class TestEdgePolicyReferenceInput: def test_encoded_image_bytes_reach_action_preprocessing_as_rgb(self): """L1: worker-side PNG decoding preserves the Policy input pixels.""" @@ -809,7 +821,7 @@ def test_generation_defaults_matrix(self): assert edge_video["num_frames"] == 121 assert edge_video["num_inference_steps"] == 50 assert edge_video["guidance_scale"] == 5.0 - assert edge_video["flow_shift"] == 3.0 + assert "flow_shift" not in edge_video # the checkpoint scheduler's shift applies edge_t2i = COSMOS3_GENERATION_DEFAULTS[(NEMOTRON_DENSE_RECIPE.name, "image")] assert edge_t2i is COSMOS3_EDGE_T2I_PARAMS @@ -846,42 +858,23 @@ def test_hf_id_registered(self): def test_policy_manifest_overrides_policy_sampling_only(self): pipeline = _bare_pipeline(NEMOTRON_DENSE_RECIPE.name) - pipeline.checkpoint_policy_defaults = resolve_checkpoint_policy_defaults( - { - "action_chunk_size": 32, - "conditioning_fps": 15.0, - "domain_name": "droid_lerobot", - } + pipeline.resolver = _policy_resolver(NEMOTRON_DENSE_RECIPE.name) + + policy = pipeline.resolver.request(RequestParameters(action_mode=ACTION_MODE_POLICY)) + assert (policy.num_inference_steps, policy.guidance_scale) == (4, 3.0) + assert policy.guidance_interval == COSMOS3_POLICY_SAMPLING_PARAMS["guidance_interval"] + assert ( + COSMOS3_GENERATION_DEFAULTS[(NEMOTRON_DENSE_RECIPE.name, "action")] + is COSMOS3_ACTION_PARAMS ) - - resolved_policy = pipeline._resolve_generation_params( - "video", - action_mode=ACTION_MODE_POLICY, - num_inference_steps=None, - guidance_scale=None, - guidance_interval=None, - ) - assert resolved_policy == { - "num_inference_steps": 4, - "guidance_scale": 3.0, - "guidance_interval": COSMOS3_POLICY_SAMPLING_PARAMS["guidance_interval"], - } - assert pipeline._mode_table("video", ACTION_MODE_FORWARD_DYNAMICS) is COSMOS3_ACTION_PARAMS - assert pipeline._mode_table("video", ACTION_MODE_INVERSE_DYNAMICS) is COSMOS3_ACTION_PARAMS for action_mode in (ACTION_MODE_FORWARD_DYNAMICS, ACTION_MODE_INVERSE_DYNAMICS): - resolved = pipeline._resolve_generation_params( - "video", - action_mode=action_mode, - num_inference_steps=None, - guidance_scale=None, - guidance_interval=None, - ) - assert resolved == { - "num_inference_steps": 30, - "guidance_scale": 1.0, - "guidance_interval": None, - } - assert pipeline._mode_table("video") is COSMOS3_EDGE_VIDEO_PARAMS + resolved = pipeline.resolver.request(RequestParameters(action_mode=action_mode)) + assert (resolved.num_inference_steps, resolved.guidance_scale) == (30, 1.0) + assert resolved.guidance_interval is None + assert ( + COSMOS3_GENERATION_DEFAULTS[(NEMOTRON_DENSE_RECIPE.name, "video")] + is COSMOS3_EDGE_VIDEO_PARAMS + ) @pytest.mark.parametrize( "action_mode", @@ -928,13 +921,7 @@ def record_prompt(text, max_sequence_length, use_system_prompt, system_prompt=No def test_policy_manifest_sampling_does_not_cross_domain_override(self, monkeypatch): pipeline = _bare_pipeline(NEMOTRON_DENSE_RECIPE.name) pipeline.action_gen = True - pipeline.checkpoint_policy_defaults = resolve_checkpoint_policy_defaults( - { - "action_chunk_size": 32, - "conditioning_fps": 15.0, - "domain_name": "droid_lerobot", - } - ) + pipeline.resolver = _policy_resolver(NEMOTRON_DENSE_RECIPE.name) monkeypatch.setattr(pipeline, "_scheduler_for", lambda *args, **kwargs: object()) monkeypatch.setattr(pipeline_module, "action_reference_size", lambda **kwargs: (480, 640)) @@ -975,62 +962,42 @@ def test_guidance_interval_resolves_for_visual_and_audio_modes(self) -> None: # Audio is an extra stream on a video request and therefore shares the # video request's resolved CFG interval. - video_and_audio = pipeline._resolve_generation_params("video", guidance_interval=None) - image = pipeline._resolve_generation_params("image", guidance_interval=None) + video_and_audio = pipeline.resolver.request(RequestParameters(mode="video")) + image = pipeline.resolver.request(RequestParameters(mode="image")) - assert video_and_audio["guidance_interval"] is None - assert image["guidance_interval"] == COSMOS3_T2I_PARAMS["guidance_interval"] + assert video_and_audio.guidance_interval is None + assert image.guidance_interval == COSMOS3_T2I_PARAMS["guidance_interval"] def test_policy_manifest_declares_policy_as_the_action_mode_default(self): pipeline = _bare_pipeline(NEMOTRON_DENSE_RECIPE.name) - pipeline.checkpoint_policy_defaults = resolve_checkpoint_policy_defaults( - { - "action_chunk_size": 32, - "conditioning_fps": 15.0, - "domain_name": "droid_lerobot", - } - ) + pipeline.resolver = _policy_resolver(NEMOTRON_DENSE_RECIPE.name) assert pipeline.extra_param_specs["action_mode"].default == ACTION_MODE_POLICY assert COSMOS3_EXTRA_SPECS["action_mode"].default is None def test_none_params_resolve_from_edge_tables(self): pipeline = _bare_pipeline(NEMOTRON_DENSE_RECIPE.name) - resolved = pipeline._resolve_generation_params( - "video", - height=None, - width=None, - num_frames=None, - num_inference_steps=None, - guidance_scale=None, - max_sequence_length=None, - frame_rate=None, - ) - assert resolved == { - "height": 480, - "width": 832, - "num_frames": 121, - "num_inference_steps": 50, - "guidance_scale": 5.0, - "max_sequence_length": 4096, - "frame_rate": 24.0, - } + resolved = pipeline.resolver.request(RequestParameters(mode="video")) + assert (resolved.height, resolved.width, resolved.num_frames) == (480, 832, 121) + assert (resolved.num_inference_steps, resolved.guidance_scale) == (50, 5.0) + assert (resolved.max_sequence_length, resolved.frame_rate) == (4096, 24.0) # Image mode falls back to the video table for fields the image table # omits, and explicit values always win. - image = pipeline._resolve_generation_params( - "image", height=None, max_sequence_length=None, num_inference_steps=8 + image = pipeline.resolver.request(RequestParameters(mode="image", num_inference_steps=8)) + assert (image.height, image.max_sequence_length, image.num_inference_steps) == ( + 640, + 4096, + 8, ) - assert image == {"height": 640, "max_sequence_length": 4096, "num_inference_steps": 8} def test_sampling_overrides_beat_tables(self): pipeline = _bare_pipeline(NEMOTRON_DENSE_RECIPE.name) pipeline.sampling = Cosmos3SamplingPolicy(fixed_sigmas=(1.0, 0.5)) - resolved = pipeline._resolve_generation_params( - "video", num_inference_steps=None, guidance_scale=None, height=None - ) - assert resolved["num_inference_steps"] == 2 - assert resolved["guidance_scale"] == 1.0 - assert resolved["height"] == 480 + pipeline.resolver = bare_resolver(NEMOTRON_DENSE_RECIPE.name, fixed_sigmas=(1.0, 0.5)) + resolved = pipeline.resolver.request(RequestParameters(mode="video")) + assert resolved.num_inference_steps == 2 + assert resolved.guidance_scale == 1.0 + assert resolved.height == 480 def test_warmup_runs_at_family_guidance(self, monkeypatch): pipeline = _bare_pipeline(NEMOTRON_DENSE_RECIPE.name) @@ -1110,7 +1077,6 @@ def test_component_loading_rejects_edge_without_native_flag(self, tmp_path): def fresh_pipeline(): pipeline = _bare_pipeline(NEMOTRON_DENSE_RECIPE.name) - pipeline.default_use_system_prompt = False # Mirror __init__: the flag starts False and only the checkpoint's # model_index may turn it on. pipeline.use_native_flow_schedule = False @@ -1155,34 +1121,13 @@ def _pipeline( pipeline._cacheable_flow_shifts_cache = None return pipeline - def test_every_cacheable_shift_source_contributes(self, monkeypatch) -> None: - """Each source must be read independently. - - Distinct values throughout, so dropping any one source -- or hard-coding - today's Edge values, where checkpoint and both mode tables all say 3.0 -- - fails instead of passing by coincidence. - """ - monkeypatch.setitem( - pipeline_module.COSMOS3_GENERATION_DEFAULTS, - ("nemotron_dense", "video"), - {"flow_shift": 4.5}, - ) - monkeypatch.setitem( - pipeline_module.COSMOS3_GENERATION_DEFAULTS, - ("nemotron_dense", "image"), - {"flow_shift": 6.25}, - ) + def test_every_cacheable_shift_source_contributes(self) -> None: + """The checkpoint's shift and every transfer hint preset; nothing else.""" + hint_shifts = { + float(defaults["flow_shift"]) for defaults in pipeline_module.TRANSFER_DEFAULTS.values() + } pipeline = self._pipeline(checkpoint_shift=2.75) - assert pipeline._cacheable_flow_shifts() == frozenset( - {2.75, 4.5, 6.25, pipeline_module.COSMOS3_V2V_FLOW_SHIFT} - ) - - def test_cacheable_set_for_the_shipped_edge_tables(self) -> None: - pipeline = self._pipeline() - # Edge declares 3.0 in both mode tables; V2V contributes its stronger shift. - assert pipeline._cacheable_flow_shifts() == frozenset( - {3.0, pipeline_module.COSMOS3_V2V_FLOW_SHIFT} - ) + assert pipeline._cacheable_flow_shifts() == frozenset({2.75} | hint_shifts) def test_arbitrary_shifts_never_enter_the_cache(self) -> None: pipeline = self._pipeline() @@ -1851,15 +1796,14 @@ def test_recipe_and_scheduler_wiring(self, edge_pipeline): assert edge_pipeline.transformer.recipe is NEMOTRON_DENSE_RECIPE assert edge_pipeline.use_native_flow_schedule is True assert edge_pipeline.sampling.native_flow_schedule is True - # Both Edge tables declare shift 3.0 on the native flow schedule, so - # the two modes resolve to the same (shift, karras) pair and share one + # Both modes run at the checkpoint scheduler's shift on the native flow + # schedule, so they resolve to the same (shift, karras) pair and share one # cached scheduler rather than building two identical ones. - schedulers = { - mode: edge_pipeline._scheduler_for(edge_pipeline._mode_table(mode)["flow_shift"]) - for mode in ("video", "image") - } + shift = edge_pipeline.resolver.config().sampling.flow_shift + assert shift == edge_pipeline.sampling.checkpoint_flow_shift + schedulers = {mode: edge_pipeline._scheduler_for(shift) for mode in ("video", "image")} for scheduler in schedulers.values(): - assert float(scheduler.config.flow_shift) == 3.0 + assert float(scheduler.config.flow_shift) == shift assert scheduler.config.use_karras_sigmas is False assert schedulers["video"] is schedulers["image"] diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_example_prompts.py b/tests/unittest/_torch/visual_gen/test_cosmos3_example_prompts.py index 0b44a2bb7868..3eda63929a8d 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_example_prompts.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_example_prompts.py @@ -14,15 +14,14 @@ from pathlib import Path import pytest +from cosmos3_fixtures import bare_resolver +from tensorrt_llm._torch.visual_gen.models.cosmos3.action import ACTION_MODE_POLICY from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( VIDEO_RES_SIZE_INFO, find_closest_target_size, - resolve_checkpoint_policy_defaults, -) -from tensorrt_llm._torch.visual_gen.models.cosmos3.parameter_resolver import ( - resolve_domain_action_config, ) +from tensorrt_llm._torch.visual_gen.models.cosmos3.parameter_resolver import RequestParameters _PROJECT_ROOT = Path(__file__).resolve().parents[4] _EXAMPLE_DIR = _PROJECT_ROOT / "examples" / "visual_gen" / "models" / "cosmos3" @@ -275,7 +274,9 @@ def test_bundled_prompt_files_load(self, name: str) -> None: def test_action_prompt_matches_example_domain(self, name: str, domain_name: str) -> None: data = cosmos3.load_prompt_file(f"prompts/{name}.json") prompt = json.loads(data["prompt"]) - config = resolve_domain_action_config(domain_name=domain_name) + config = bare_resolver("qwen3").request( + RequestParameters(action_mode="action", domain_name=domain_name) + ) assert list(prompt) == [ "cinematography", @@ -288,14 +289,14 @@ def test_action_prompt_matches_example_domain(self, name: str, domain_name: str) assert prompt["cinematography"] == { "framing": "This video is captured from a first-person perspective looking at the scene." } - duration_seconds = config["num_frames"] / config["frame_rate"] + duration_seconds = config.num_frames / config.frame_rate minutes, seconds = divmod(round(duration_seconds), 60) assert prompt["duration"] == f"{int(duration_seconds)}s" assert prompt["actions"][0]["time"] == f"0:00-{minutes}:{seconds:02d}" - assert prompt["fps"] == config["frame_rate"] - width, height = find_closest_target_size(720, 1280, config["action_resolution"]) + assert prompt["fps"] == config.frame_rate + width, height = find_closest_target_size(720, 1280, config.action_resolution) assert prompt["resolution"] == {"H": height, "W": width} - assert VIDEO_RES_SIZE_INFO[str(config["action_resolution"])][prompt["aspect_ratio"]] == ( + assert VIDEO_RES_SIZE_INFO[str(config.action_resolution)][prompt["aspect_ratio"]] == ( width, height, ) @@ -303,25 +304,19 @@ def test_action_prompt_matches_example_domain(self, name: str, domain_name: str) def test_edge_policy_prompt_matches_checkpoint_defaults(self): data = cosmos3.load_prompt_file("prompts/action_edge_policy_droid.json") prompt = json.loads(data["prompt"]) - policy_defaults = resolve_checkpoint_policy_defaults( - { - "action_chunk_size": 32, - "conditioning_fps": 15.0, - "domain_name": "droid_lerobot", - } + policy = {"action_chunk_size": 32, "conditioning_fps": 15.0, "domain_name": "droid_lerobot"} + config = bare_resolver("nemotron_dense", policy=policy).request( + RequestParameters(action_mode=ACTION_MODE_POLICY) ) - config = resolve_domain_action_config(checkpoint_policy_defaults=policy_defaults) - duration_seconds = config["num_frames"] / config["frame_rate"] + duration_seconds = config.num_frames / config.frame_rate minutes, seconds = divmod(round(duration_seconds), 60) assert prompt["duration"] == f"{int(duration_seconds)}s" assert prompt["actions"][0]["time"] == f"0:00-{minutes}:{seconds:02d}" - assert prompt["fps"] == config["frame_rate"] - width, height = find_closest_target_size(544, 736, config["action_resolution"]) + assert prompt["fps"] == config.frame_rate + width, height = find_closest_target_size(544, 736, config.action_resolution) assert prompt["resolution"] == {"H": height, "W": width} - width, height = VIDEO_RES_SIZE_INFO[str(config["action_resolution"])][ - prompt["aspect_ratio"] - ] + width, height = VIDEO_RES_SIZE_INFO[str(config.action_resolution)][prompt["aspect_ratio"]] assert prompt["resolution"] == {"H": height, "W": width} def test_default_prompt_file(self): diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_fp8.py b/tests/unittest/_torch/visual_gen/test_cosmos3_fp8.py index 64c6c0a37ba6..f4b439e7548b 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_fp8.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_fp8.py @@ -736,7 +736,7 @@ def _run_distilled_loop(model, device) -> torch.Tensor: from tensorrt_llm._torch.visual_gen.models.cosmos3.sampling import Cosmos3SamplingPolicy scheduler = FlowMatchEulerDiscreteScheduler.from_config(_DISTILLED_SCHEDULER_CONFIG) - policy = Cosmos3SamplingPolicy.from_scheduler(scheduler) + policy = Cosmos3SamplingPolicy.from_scheduler(scheduler, fixed_sigmas=_DISTILLED_SIGMAS) policy.set_timesteps(scheduler, num_inference_steps=len(_DISTILLED_SIGMAS), device=device) step_kwargs = policy.scheduler_step_kwargs(torch.Generator().manual_seed(_SEED_SDE)) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_parameter_resolver.py b/tests/unittest/_torch/visual_gen/test_cosmos3_parameter_resolver.py index cc34cddc53ae..6462b9ce94b4 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_parameter_resolver.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_parameter_resolver.py @@ -12,40 +12,36 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""The Cosmos3 precedence table: the engine on synthetic sources, the checkpoint rules -on synthetic files and on the published checkpoints, and the request rules.""" +"""The Cosmos3 parameter resolver: ``select`` and the checkpoint-file reader on +synthetic inputs, the checkpoint settings on synthetic files and on the published +checkpoints, and the per-request settings.""" import json from pathlib import Path from types import SimpleNamespace import pytest +from cosmos3_fixtures import bare_resolver from tensorrt_llm._torch.visual_gen.models.cosmos3.parameter_resolver import ( - ACTION_RULES, CHECKPOINT, - CHECKPOINT_RULES, - CONFIG, - GENERATION_RULES, - MODEL_INDEX, - SCHEDULER, - TRANSFORMER, - VAE, + CHECKPOINT_FILE, + CONFIG_FILE, + DEFAULT, + MODEL_INDEX_FILE, + REQUEST, + SCHEDULER_FILE, + TRANSFORMER_FILE, + UNSET, + VAE_FILE, + WORKFLOW, CheckpointFiles, Conflict, - Default, - Given, - Key, - RequestContext, - Rule, - Table, - check_precedence, + ParameterResolver, + RequestParameters, fill_pretrained_config, - resolve_action, resolve_checkpoint, - resolve_domain_action_config, - resolve_generation, - resolve_setting, + select, ) pytestmark = pytest.mark.cpu_only @@ -64,73 +60,66 @@ def _write(checkpoint_dir: Path, files: dict) -> Path: def _files_on_disk(tmp_path: Path, files: dict) -> CheckpointFiles: _write(tmp_path, files) - return CheckpointFiles(str(tmp_path), files.get(TRANSFORMER, {})) + return CheckpointFiles(str(tmp_path), files.get(TRANSFORMER_FILE, {})) def _resolve(tmp_path: Path, files: dict): _write(tmp_path, files) - return resolve_checkpoint(str(tmp_path), files.get(TRANSFORMER, {})) + return resolve_checkpoint(str(tmp_path), files.get(TRANSFORMER_FILE, {})) -class TestEngine: - def test_first_present_source_wins(self, tmp_path): - files = _files_on_disk(tmp_path, {CONFIG: {"a": 1}, MODEL_INDEX: {"a": 1}}) - rule = Rule((Key(CONFIG, "a"), Key(MODEL_INDEX, "a"), Default(7, "x"))) - assert resolve_setting("a", rule, files).value == 1 - assert resolve_setting("a", rule, files).source == "config.json:a" +class TestSelect: + def test_request_over_checkpoint_over_default(self): + assert select(1, 2, 3).origin == REQUEST + assert select(None, 2, 3).origin == CHECKPOINT + assert select(None, None, 3).origin == DEFAULT + assert select(None, None, None) is None - def test_falls_through_missing_and_null_sources(self, tmp_path): - files = _files_on_disk(tmp_path, {CONFIG: {"a": None}, MODEL_INDEX: {"a": 2}}) - rule = Rule((Key(CONFIG, "a"), Key(MODEL_INDEX, "a"), Key(CHECKPOINT, "a"))) - assert resolve_setting("a", rule, files).value == 2 + def test_false_and_zero_are_values(self): + assert select(False, True).value is False + assert select(0, 5).value == 0 - def test_default_only_when_nothing_is_declared(self, tmp_path): - files = _files_on_disk(tmp_path, {}) - rule = Rule((Key(CONFIG, "a"), Default(7, "x"))) - resolved = resolve_setting("a", rule, files) - assert (resolved.value, resolved.source) == (7, "default (x)") - def test_unresolvable_required_setting_is_an_error(self, tmp_path): - files = _files_on_disk(tmp_path, {}) - with pytest.raises(ValueError, match="cannot resolve a"): - resolve_setting("a", Rule((Key(CONFIG, "a"),)), files) - assert resolve_setting("a", Rule((Key(CONFIG, "a"),), required=False), files) is None +class TestCheckpointFiles: + def test_first_file_holding_the_key_wins(self, tmp_path): + files = _files_on_disk(tmp_path, {CONFIG_FILE: {"a": 1}, MODEL_INDEX_FILE: {"a": 1}}) + r = files.read("a", f"{CONFIG_FILE}:a", f"{MODEL_INDEX_FILE}:a") + assert (r.value, r.where, r.disagreeing) == (1, "config.json:a", ()) + + def test_missing_and_null_fall_through(self, tmp_path): + files = _files_on_disk(tmp_path, {CONFIG_FILE: {"a": None}, MODEL_INDEX_FILE: {"a": 2}}) + r = files.read("a", f"{CONFIG_FILE}:a", f"{MODEL_INDEX_FILE}:a", f"{CHECKPOINT_FILE}:a") + assert r.value == 2 + assert files.read("b", f"{CONFIG_FILE}:b") is None def test_conflict_policies(self, tmp_path): - files = _files_on_disk(tmp_path, {CONFIG: {"a": 1}, MODEL_INDEX: {"a": 2}}) - sources = (Key(CONFIG, "a"), Key(MODEL_INDEX, "a")) + files = _files_on_disk(tmp_path, {CONFIG_FILE: {"a": 1}, MODEL_INDEX_FILE: {"a": 2}}) + keys = (f"{CONFIG_FILE}:a", f"{MODEL_INDEX_FILE}:a") with pytest.raises(ValueError, match="model_index.json:a=2"): - resolve_setting("a", Rule(sources, conflict=Conflict.ERROR), files) - warned = resolve_setting("a", Rule(sources, conflict=Conflict.WARN), files) + files.read("a", *keys, conflict=Conflict.ERROR) + warned = files.read("a", *keys, conflict=Conflict.WARN) assert warned.value == 1 and warned.disagreeing == ("model_index.json:a=2",) - ignored = resolve_setting("a", Rule(sources), files) + ignored = files.read("a", *keys) assert ignored.value == 1 and ignored.disagreeing == ("model_index.json:a=2",) def test_float_lists_compare_within_tolerance(self, tmp_path): files = _files_on_disk( - tmp_path, {CONFIG: {"a": [0.5, 1 / 3]}, MODEL_INDEX: {"a": [0.5, 0.3333333333]}} + tmp_path, + {CONFIG_FILE: {"a": [0.5, 1 / 3]}, MODEL_INDEX_FILE: {"a": [0.5, 0.3333333333]}}, ) - rule = Rule((Key(CONFIG, "a"), Key(MODEL_INDEX, "a")), conflict=Conflict.ERROR) - assert resolve_setting("a", rule, files).disagreeing == () - - def test_every_rule_keeps_request_over_checkpoint_over_default(self): - for rules in (CHECKPOINT_RULES, GENERATION_RULES, ACTION_RULES): - check_precedence(rules) - with pytest.raises(ValueError, match="breaks the request > checkpoint > default order"): - check_precedence({"x": Rule((Default(1, "first"), Given()))}) - with pytest.raises(ValueError, match="breaks the request > checkpoint > default order"): - check_precedence({"x": Rule((Table("x"), Given()))}) + r = files.read("a", f"{CONFIG_FILE}:a", f"{MODEL_INDEX_FILE}:a", conflict=Conflict.ERROR) + assert r.disagreeing == () def _files(): """A distilled checkpoint whose training config and diffusers index disagree on the system prompt.""" return { - TRANSFORMER: {"temporal_compression_factor": 4}, - SCHEDULER: {"fixed_step_sampler_config": {"t_list": T_LIST}}, - CONFIG: {"model": {"config": {"vlm_config": {"use_system_prompt": False}}}}, - MODEL_INDEX: {"default_use_system_prompt": True}, - VAE: {"scale_factor_temporal": 4, "scale_factor_spatial": 16}, + TRANSFORMER_FILE: {"temporal_compression_factor": 4}, + SCHEDULER_FILE: {"fixed_step_sampler_config": {"t_list": T_LIST}}, + CONFIG_FILE: {"model": {"config": {"vlm_config": {"use_system_prompt": False}}}}, + MODEL_INDEX_FILE: {"default_use_system_prompt": True}, + VAE_FILE: {"scale_factor_temporal": 4, "scale_factor_spatial": 16}, } @@ -139,7 +128,7 @@ def test_training_config_beats_the_diffusers_index(self, tmp_path): cfg = _resolve(tmp_path, _files()) assert cfg.prompt.use_system_prompt is False assert cfg.source_of("prompt.use_system_prompt") == ( - "config.json:model.config.vlm_config.use_system_prompt" + "checkpoint (config.json:model.config.vlm_config.use_system_prompt)" ) assert cfg.resolved["prompt.use_system_prompt"].disagreeing == ( "model_index.json:default_use_system_prompt=True", @@ -150,37 +139,41 @@ def test_sections(self, tmp_path): assert (cfg.vae.temporal_compression, cfg.vae.spatial_compression) == (4, 16) assert cfg.sampling.fixed_sigmas == tuple(T_LIST) assert cfg.sampling.native_flow_schedule is False - assert cfg.source_of("sampling.native_flow_schedule").startswith("default (") + assert cfg.source_of("sampling.native_flow_schedule") == DEFAULT assert cfg.action.policy is None def test_vae_compression_copies_must_agree(self, tmp_path): files = _files() - files[TRANSFORMER]["temporal_compression_factor"] = 8 + files[TRANSFORMER_FILE]["temporal_compression_factor"] = 8 with pytest.raises(ValueError, match="vae.temporal_compression"): _resolve(tmp_path, files) def test_schedule_must_decrease_within_unit_interval(self, tmp_path): files = _files() - files[SCHEDULER]["fixed_step_sampler_config"]["t_list"] = [0.5, 1.0] + files[SCHEDULER_FILE]["fixed_step_sampler_config"]["t_list"] = [0.5, 1.0] with pytest.raises(ValueError, match="t_list must decrease"): _resolve(tmp_path, files) def test_base_checkpoint_has_no_schedule(self, tmp_path): - files = {k: v for k, v in _files().items() if k != SCHEDULER} + files = {k: v for k, v in _files().items() if k != SCHEDULER_FILE} assert _resolve(tmp_path, files).sampling.fixed_sigmas is None def test_system_prompt_must_be_declared(self, tmp_path): files = _files() - del files[CONFIG]["model"]["config"]["vlm_config"]["use_system_prompt"] - del files[MODEL_INDEX]["default_use_system_prompt"] + del files[CONFIG_FILE]["model"]["config"]["vlm_config"]["use_system_prompt"] + del files[MODEL_INDEX_FILE]["default_use_system_prompt"] with pytest.raises(ValueError, match="cannot resolve prompt.use_system_prompt"): _resolve(tmp_path, files) - def test_spatial_compression_must_be_declared(self, tmp_path): + def test_vae_spatial_compression_is_optional(self, tmp_path): files = _files() - del files[VAE]["scale_factor_spatial"] - with pytest.raises(ValueError, match="cannot resolve vae.spatial_compression"): - _resolve(tmp_path, files) + del files[VAE_FILE]["scale_factor_spatial"] + assert _resolve(tmp_path, files).vae.spatial_compression is None + + def test_sigmas_are_coerced_to_floats(self, tmp_path): + files = _files() + files[SCHEDULER_FILE]["fixed_step_sampler_config"]["t_list"] = [1, "0.5"] + assert _resolve(tmp_path, files).sampling.fixed_sigmas == (1.0, 0.5) def test_fill_pretrained_config_sets_only_missing_attributes(self): declared = SimpleNamespace(temporal_compression_factor=8) @@ -191,7 +184,7 @@ def test_fill_pretrained_config_sets_only_missing_attributes(self): assert declared.temporal_compression_factor_sound == 1 -# Disagreements between a checkpoint's own files that the table expects and resolves. +# Disagreements between a checkpoint's own files that the resolver expects and resolves. EXPECTED_DISAGREEMENTS = { "nvidia/Cosmos3-Nano-Sim-Bimanual": {"prompt.use_system_prompt"}, } @@ -229,73 +222,125 @@ def test_fixture_covers_every_expected_checkpoint(): assert set(json.loads(FIXTURE.read_text())) == set(EXPECTED) +def _resolver(**facts) -> ParameterResolver: + return bare_resolver(facts.pop("family", "qwen3"), **facts) + + class TestRequest: - def test_generation_values_name_their_source(self): - ctx = RequestContext(family="qwen3", mode="image") - r = resolve_generation(ctx, num_inference_steps=None, num_frames=None, height=20) - assert r.source_of("num_inference_steps") == "mode table.num_inference_steps" - assert r["num_inference_steps"] == 50 - assert r.source_of("num_frames") == "video table.num_frames" # the image table has none - assert (r["height"], r.source_of("height")) == (20, "request") + def test_generation_values_come_from_the_tables(self): + r = _resolver().request(RequestParameters(mode="image", height=20)) + assert r.num_inference_steps == 50 # the image table + assert r.num_frames == 189 # the image table has none: the video table + assert r.height == 20 # the caller's def test_distilled_schedule_beats_the_tables(self): - ctx = RequestContext(family="qwen3", distilled_steps=4) - r = resolve_generation(ctx, num_inference_steps=None, guidance_scale=None) - assert (r["num_inference_steps"], r["guidance_scale"]) == (4, 1.0) - assert r.source_of("guidance_scale").startswith("checkpoint:fixed_step_sampler_config") + r = _resolver(fixed_sigmas=(1.0, 0.75, 0.5, 0.25)).request(RequestParameters()) + assert (r.num_inference_steps, r.guidance_scale) == (4, 1.0) def test_request_beats_the_distilled_schedule(self): - """The caller's value wins over the checkpoint, and the checkpoint's value is - recorded as the disagreeing source.""" - ctx = RequestContext(family="qwen3", distilled_steps=4) - r = resolve_generation(ctx, num_inference_steps=8) - assert (r["num_inference_steps"], r.source_of("num_inference_steps")) == (8, "request") - assert r.values["num_inference_steps"].disagreeing == ( - "checkpoint:fixed_step_sampler_config (steps)=4", + r = _resolver(fixed_sigmas=(1.0, 0.75, 0.5, 0.25)).request( + RequestParameters(num_inference_steps=8) + ) + assert r.num_inference_steps == 8 + + def test_origins(self): + _, _, origins = _resolver(fixed_sigmas=(1.0, 0.75, 0.5, 0.25))._resolve( + RequestParameters(height=20) + ) + assert (origins["height"], origins["num_inference_steps"], origins["width"]) == ( + REQUEST, + CHECKPOINT, + DEFAULT, + ) + + def test_action_fields_stay_none_outside_action_mode(self): + r = _resolver().request(RequestParameters(mode="video")) + assert (r.raw_action_dim, r.action_chunk_size, r.action_resolution) == (None, None, None) + assert (r.action_fps, r.use_state, r.policy_active, r.preset_key) == ( + None, + None, + False, + None, ) def test_policy_block_applies_only_to_its_domain(self): policy = {"domain_name": "droid_lerobot", "action_chunk_size": 32, "num_inference_steps": 4} - matched = resolve_domain_action_config(checkpoint_policy_defaults=policy) - assert matched["domain_name"] == "droid_lerobot" - assert matched["action_chunk_size"] == 32 - assert matched["sources"]["action_chunk_size"] == "checkpoint.json:policy.action_chunk_size" - other = resolve_domain_action_config(domain_name="av", checkpoint_policy_defaults=policy) - assert other["action_chunk_size"] == 60 - assert other["sources"]["action_chunk_size"] == "domain preset.action_chunk_size" - assert any("overrides checkpoint policy domain_name" in w for w in other["warnings"]) + resolver = _resolver(policy=policy) + matched, _, origins = resolver._resolve(RequestParameters(action_mode="policy")) + assert matched.domain_name == "droid_lerobot" + assert matched.action_chunk_size == 32 + assert origins["action_chunk_size"] == CHECKPOINT + other, warnings, origins = resolver._resolve( + RequestParameters(action_mode="policy", domain_name="av") + ) + assert other.action_chunk_size == 60 + assert origins["action_chunk_size"] == DEFAULT + assert any("overrides checkpoint policy domain_name" in w for w in warnings) + + def test_policy_block_needs_policy_mode(self): + policy = {"domain_name": "droid_lerobot", "action_chunk_size": 32, "num_inference_steps": 4} + resolver = _resolver(family="nemotron_dense", policy=policy) + in_policy = resolver.request(RequestParameters(action_mode="policy")) + assert ( + in_policy.policy_active, + in_policy.action_chunk_size, + in_policy.num_inference_steps, + ) == (True, 32, 4) + dynamics = resolver.request(RequestParameters(action_mode="forward_dynamics")) + assert dynamics.policy_active is False + assert dynamics.num_inference_steps == 30 # the action table, not the block def test_caller_value_wins_and_the_recommendation_is_reported(self): - cfg = resolve_domain_action_config(domain_name="av", action_chunk_size=16) - assert cfg["action_chunk_size"] == 16 - assert cfg["sources"]["action_chunk_size"] == "request" - assert cfg["warnings"] == [ + params, warnings, origins = _resolver()._resolve( + RequestParameters(action_mode="action", domain_name="av", action_chunk_size=16) + ) + assert params.action_chunk_size == 16 + assert origins["action_chunk_size"] == REQUEST + assert warnings == [ "Cosmos3 action_chunk_size=16 differs from recommended 60 for domain 'av'." ] + def test_table_values_are_not_recommendations(self): + """A caller's height that differs from the video table is not a warning: only + the policy block, the domain preset and the width table recommend.""" + _, warnings, _ = _resolver()._resolve(RequestParameters(height=20)) + assert warnings == [] + def test_explicit_none_in_a_mode_table_is_a_value(self): """Edge text-to-image switches the guidance window off with an explicit ``None``; that must not fall through to the video table.""" - ctx = RequestContext(family="nemotron_dense", mode="image") - r = resolve_generation(ctx, guidance_interval=None) - assert r["guidance_interval"] is None - assert r.source_of("guidance_interval") == "mode table.guidance_interval" + resolver = _resolver(family="nemotron_dense") + _, _, origins = resolver._resolve(RequestParameters(mode="image")) + assert resolver.request(RequestParameters(mode="image")).guidance_interval is None + assert origins["guidance_interval"] == UNSET + + def test_flow_shift_is_the_checkpoint_scheduler_unless_requested(self): + resolver = _resolver(flow_shift=5.0) + assert resolver.request(RequestParameters()).flow_shift == 5.0 + assert resolver.request(RequestParameters(mode="image")).flow_shift == 5.0 + assert resolver.request(RequestParameters(flow_shift=3.0)).flow_shift == 3.0 + assert _resolver().request(RequestParameters()).flow_shift is None # distilled + + def test_system_prompt_rule(self): + resolver = _resolver(use_system_prompt=True) + assert resolver.request(RequestParameters()).use_system_prompt is True + assert resolver.request(RequestParameters(is_transfer=True)).use_system_prompt is False + assert _resolver().request(RequestParameters(is_v2v=True)).use_system_prompt is True + assert ( + resolver.request(RequestParameters(use_system_prompt=False)).use_system_prompt is False + ) def test_unresolvable_field_is_an_error(self): - ctx = RequestContext(family="no-such-family") - with pytest.raises(ValueError, match="cannot resolve num_frames"): - resolve_generation(ctx, num_frames=None) + with pytest.raises(ValueError, match="cannot resolve num_inference_steps"): + _resolver(family="no-such-family").request(RequestParameters()) def test_action_fps_follows_the_frame_rate(self): - ctx = RequestContext(family="qwen3", action_mode="forward_dynamics", domain_name="av") - r = resolve_action( - ctx, - raw_action_dim=None, - action_chunk_size=None, - action_resolution=None, - frame_rate=None, - action_fps=None, - use_state=None, - ) - assert (r["frame_rate"], r["action_fps"], r["num_frames"]) == (10.0, 10.0, 61) - assert r.source_of("action_fps") == "frame_rate" + r = _resolver().request(RequestParameters(action_mode="forward_dynamics", domain_name="av")) + assert (r.frame_rate, r.action_fps, r.num_frames) == (10.0, 10.0, 61) + assert r.preset_key == "av" + + def test_workflow_rules_are_tagged(self): + _, _, origins = _resolver()._resolve(RequestParameters(is_transfer=True)) + assert origins["use_system_prompt"] == WORKFLOW + _, _, origins = _resolver()._resolve(RequestParameters()) + assert origins["use_system_prompt"] == CHECKPOINT diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index ba653f97b4f4..9af56d146b5e 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -33,6 +33,7 @@ import PIL.Image import pytest import torch +from cosmos3_fixtures import bare_resolver from utils.llm_data import get_checkpoint import tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 as pipe_mod @@ -908,6 +909,7 @@ def test_v2v_flow_shift_override_request_path(self): # Bypassing __init__ means the generation-default family has to be set # here; forward() reads its per-mode table through it. pipeline.family = QWEN3_RECIPE.name + pipeline.resolver = bare_resolver(QWEN3_RECIPE.name) calls = [] token_calls = [] @@ -975,6 +977,7 @@ def test_v2v_rebuilds_the_audio_scheduler_too(self): # Bypassing __init__ means the generation-default family has to be set # here; forward() reads its per-mode table through it. pipeline.family = QWEN3_RECIPE.name + pipeline.resolver = bare_resolver(QWEN3_RECIPE.name) rebuilt = [] class StopAfterTokenize(Exception): @@ -1034,7 +1037,7 @@ def test_action_video_is_not_classified_as_v2v(self): ) pipeline.audio_gen = False pipeline.action_gen = True - pipeline.default_use_system_prompt = False + pipeline.resolver = bare_resolver(QWEN3_RECIPE.name) pipeline.family = QWEN3_RECIPE.name token_calls = [] @@ -1089,7 +1092,7 @@ def test_action_restores_the_checkpoint_flow_shift(self): ) pipeline.audio_gen = False pipeline.action_gen = True - pipeline.default_use_system_prompt = False + pipeline.resolver = bare_resolver(QWEN3_RECIPE.name) pipeline.family = QWEN3_RECIPE.name calls = [] @@ -1142,6 +1145,7 @@ def test_scheduler_for_serves_every_stream(self): from that stream's own base.""" pipeline = Cosmos3OmniMoTPipeline.__new__(Cosmos3OmniMoTPipeline) pipeline.family = QWEN3_RECIPE.name + pipeline.resolver = bare_resolver(QWEN3_RECIPE.name) rebuilt = [] class FakeSampling: @@ -1198,6 +1202,7 @@ def test_transfer_rejects_an_image_reference(self): pipeline.action_gen = False # __new__ skips __init__, where the real pipeline resolves this. pipeline.family = QWEN3_RECIPE.name + pipeline.resolver = bare_resolver(QWEN3_RECIPE.name) pipeline.audio_gen = False class FakeSampling: @@ -1250,6 +1255,7 @@ def test_transfer_use_system_prompt_defaults_off(self): pipeline.action_gen = False # __new__ skips __init__, where the real pipeline resolves this. pipeline.family = QWEN3_RECIPE.name + pipeline.resolver = bare_resolver(QWEN3_RECIPE.name) pipeline.audio_gen = False class FakeSampling: @@ -1308,11 +1314,8 @@ def test_t2i_smoke(self, cosmos3_pipeline): ) assert result.video is None _assert_valid_image(result.image, height=T2I_HEIGHT, width=T2I_WIDTH) - _assert_scheduler_config( - cosmos3_pipeline, - flow_shift=COSMOS3_T2I_PARAMS["flow_shift"], - use_karras_sigmas=_base_use_karras_sigmas(cosmos3_pipeline), - ) + # Image mode runs at the checkpoint scheduler's shift like every other mode. + _assert_default_video_scheduler_config(cosmos3_pipeline) @pytest.mark.integration @@ -1633,9 +1636,10 @@ def _blocked_pipeline(): pipeline.audio_gen = False pipeline.audio_scheduler = None pipeline.family = QWEN3_RECIPE.name + pipeline.resolver = bare_resolver(QWEN3_RECIPE.name) # All-None policy is the documented pre-load placeholder. pipeline.sampling = Cosmos3SamplingPolicy() - pipeline.default_use_system_prompt = False + pipeline.resolver = bare_resolver(QWEN3_RECIPE.name) # ``rank`` is a read-only property that resolves to 0 with no # distributed init. The guardrail model is not under test, so a # checker that reports "unsafe" for any input drives the path diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py b/tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py index 17cbf76148b3..99b3def3fe39 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py @@ -23,6 +23,7 @@ import pytest import torch import torch.nn as nn +from cosmos3_fixtures import bare_resolver from tensorrt_llm._torch.visual_gen.models.cosmos3 import pipeline_cosmos3 as pipeline_module from tensorrt_llm._torch.visual_gen.models.cosmos3 import transfer as transfer_module @@ -249,6 +250,7 @@ def _make_pipeline(sampling=None): # __new__ skips __init__, which is where the real pipeline resolves this # from the transformer config; the mode-defaults tables are keyed on it. pipeline.family = QWEN3_RECIPE.name + pipeline.resolver = bare_resolver(QWEN3_RECIPE.name) pipeline.scheduler = StubScheduler() pipeline.sampling = sampling or StubSamplingPolicy() pipeline.safety_checker = None diff --git a/tests/unittest/_torch/visual_gen/test_data/README.md b/tests/unittest/_torch/visual_gen/test_data/README.md index cfe1ad65314c..5d8ec43bd82b 100644 --- a/tests/unittest/_torch/visual_gen/test_data/README.md +++ b/tests/unittest/_torch/visual_gen/test_data/README.md @@ -60,12 +60,12 @@ ffmpeg -y -framerate 24 -i frame_%02d.png \ ## `cosmos3_checkpoint_configs.json` -The settings `parameter_resolver.CHECKPOINT_RULES` reads, extracted from the +The settings `parameter_resolver.resolve_checkpoint` reads, extracted from the published Cosmos3 checkpoints on Hugging Face (`main`, Fri 23 Oct 2026): for each -repo, file and key path a rule names, the value at that path, nested back into its +repo, file and key path it names, the value at that path, nested back into its file. Nothing else from the checkpoints is kept. `test_cosmos3_parameter_resolver.py` -resolves every entry, so a new export whose files disagree in a way the precedence -table does not expect fails the test. +resolves every entry, so a new export whose files disagree in a way the resolver +does not expect fails the test. Regenerate by fetching `https://huggingface.co//raw/main/` for every -`Key` in the table and keeping only those paths. +`file:key` that `resolve_checkpoint` names and keeping only those paths. From 9832c6907ebce0a4679bd1de752c878412640a84 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 2 Nov 2026 14:13:02 -0800 Subject: [PATCH 8/8] [None][test] Cosmos3: assert the checkpoint flow shift where the tables used to override it Signed-off-by: Igor Shovkun --- .../unittest/_torch/visual_gen/test_cosmos3_edge.py | 4 +++- .../_torch/visual_gen/test_cosmos3_pipeline.py | 13 +++++++------ 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py b/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py index 8dcedcd9364d..140ac56fb786 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py @@ -1852,7 +1852,9 @@ def fake_decode(latents, decode_fn, **kwargs): assert captured["guidance_scale"] == 5.0 # Edge video defaults: 121 frames -> 31 latent frames, 480x832 -> 30x52. assert tuple(captured["latents"].shape) == (1, 48, 31, 30, 52) - assert captured["scheduler"] is edge_pipeline._scheduler_for(3.0) + assert captured["scheduler"] is edge_pipeline._scheduler_for( + edge_pipeline.sampling.checkpoint_flow_shift + ) assert captured["scheduler"].num_inference_steps == 50 assert captured["decoded"] is True diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index 9af56d146b5e..1d132bd0099e 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -867,7 +867,7 @@ def test_v2v_smoke(self, cosmos3_pipeline): assert result.frame_rate == FRAME_RATE _assert_scheduler_config( cosmos3_pipeline, - flow_shift=10.0, + flow_shift=cosmos3_pipeline.sampling.checkpoint_flow_shift, use_karras_sigmas=False, ) @@ -897,7 +897,7 @@ def test_v2v_keep_last_smoke(self, cosmos3_pipeline): ) _assert_scheduler_config( cosmos3_pipeline, - flow_shift=10.0, + flow_shift=cosmos3_pipeline.sampling.checkpoint_flow_shift, use_karras_sigmas=False, ) @@ -977,7 +977,7 @@ def test_v2v_rebuilds_the_audio_scheduler_too(self): # Bypassing __init__ means the generation-default family has to be set # here; forward() reads its per-mode table through it. pipeline.family = QWEN3_RECIPE.name - pipeline.resolver = bare_resolver(QWEN3_RECIPE.name) + pipeline.resolver = bare_resolver(QWEN3_RECIPE.name, flow_shift=1.0) rebuilt = [] class StopAfterTokenize(Exception): @@ -1023,7 +1023,8 @@ def fake_tokenize_prompt(text, max_sequence_length, use_system_prompt, system_pr enable_audio=True, ) - assert rebuilt == [("video", 10.0, False), ("audio", 10.0, False)] + # Video-to-video keeps the checkpoint shift and forces the uniform sigma grid. + assert rebuilt == [("video", 1.0, False), ("audio", 1.0, False)] def test_action_video_is_not_classified_as_v2v(self): """An action reference arrives as the same `video` bytes V2V uses, but @@ -1092,7 +1093,7 @@ def test_action_restores_the_checkpoint_flow_shift(self): ) pipeline.audio_gen = False pipeline.action_gen = True - pipeline.resolver = bare_resolver(QWEN3_RECIPE.name) + pipeline.resolver = bare_resolver(QWEN3_RECIPE.name, flow_shift=7.0) pipeline.family = QWEN3_RECIPE.name calls = [] @@ -1343,7 +1344,7 @@ def test_v2v_audio_smoke(self, cosmos3_pipeline): _assert_valid_audio(result.audio, result.audio_sample_rate) _assert_scheduler_config( cosmos3_pipeline, - flow_shift=10.0, + flow_shift=cosmos3_pipeline.sampling.checkpoint_flow_shift, use_karras_sigmas=False, )