diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py index 2bb2284f3b55..adb0fa50e8c0 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 EMBODIMENT_TO_DOMAIN_ID, find_closest_target_size + ACTION_MODE_POLICY = "policy" ACTION_MODE_FORWARD_DYNAMICS = "forward_dynamics" ACTION_MODE_INVERSE_DYNAMICS = "inverse_dynamics" @@ -22,134 +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 - - -# 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") - - -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: @@ -321,25 +195,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 18ec9b117671..23b4ce61baf8 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -12,20 +12,15 @@ # 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 parameter resolver reads them as the defaults, after +the request and the checkpoint. """ 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, - 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 @@ -33,10 +28,15 @@ # 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``). -VIDEO_RES_SIZE_INFO = { +# 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), @@ -67,6 +67,121 @@ }, } + +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 = { @@ -237,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, @@ -260,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 = { @@ -268,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, } @@ -306,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. @@ -419,7 +531,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 @@ -546,177 +658,17 @@ 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, -) -> 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. - """ - 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) - ) - ) - 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"} - ) - 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}.") - - 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, - } - - -COSMOS3_V2V_DEFAULT_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/parameter_resolver.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/parameter_resolver.py new file mode 100644 index 000000000000..be4a2a718fb1 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/parameter_resolver.py @@ -0,0 +1,716 @@ +# 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. +"""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. + +:class:`ParameterResolver` owns one checkpoint's settings and serves two moments: + +* 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 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 + +import json +import math +import os +from collections.abc import Mapping +from dataclasses import dataclass, replace +from enum import Enum +from typing import Any, NamedTuple, Optional + +from tensorrt_llm.logger import logger + +from .action import ACTION_MODE_POLICY +from .defaults import ( + 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 + +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 Resolved: + value: Any + 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 source(self) -> str: + return f"{self.origin} ({self.where})" if self.where else self.origin + + +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: + """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 + + +# -------------------------------------------------------------------------- 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" + + +class Conflict(Enum): + """What it means when a lower-priority checkpoint file holds a different value.""" + + ERROR = "error" + WARN = "warn" + IGNORE = "ignore" + + +_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: + """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_FILE 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 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) +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: 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) +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 + flow_shift: Optional[float] + """The scheduler's flow shift; ``None`` on distilled checkpoints, whose fixed + schedule has no shift.""" + + +@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] + """Every setting above by dotted name (``vae.temporal_compression``), with origin.""" + + 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) + + +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: + """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] = {} + + 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: + item = resolved.get(name) + return None if item is None else item.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)}." + ) + spatial = get("vae.spatial_compression") + flow_shift = get("sampling.flow_shift") + 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=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")), + resolved=resolved, + ) + + +def fill_pretrained_config(pretrained_config: Any) -> Any: + """Set the transformer-construction attributes of ``pretrained_config`` that the + export omits (:data:`ARCHITECTURE_KEYS`). 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, (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 + + +# ----------------------------------------------------------------------------- request + + +@dataclass(frozen=True) +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``.""" + + mode: str = "video" + """``"video"`` or ``"image"``; the action table is used instead when ``action_mode`` is set.""" + action_mode: Optional[str] = None + 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 + + 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 + + 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.""" + + +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.""" + +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.""" + + +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 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( + 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 + + +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 + ) + """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) + ) + ) + + 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 1bf29e825979..c9b5fbdba8c0 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -70,24 +70,28 @@ def tqdm(iterable, **kwargs): ) from .defaults import ( COSMOS3_720P_PARAMS, + COSMOS3_ASPECT_RATIO_BUCKETS, + 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, _normalize_condition_video_keep, _normalize_condition_video_latent_indexes, - resolve_checkpoint_policy_defaults, - resolve_domain_action_config, + find_closest_target_size, ) from .guardrails import check_video_safety, download_guardrail_checkpoint from .negative_prompt import COSMOS3_VIDEO_NEGATIVE_PROMPT +from .parameter_resolver import ParameterResolver, RequestParameters from .sampling import DISTILLED_GUIDANCE_SCALE, Cosmos3SamplingPolicy, load_scheduler from .sound_tokenizer import LatentAutoEncoderV2 from .transfer import ( TRANSFER_DEFAULTS, Cosmos3TransferConfig, decode_media_to_uint8_cthw, - find_closest_target_size, load_or_compute_control_frames, pad_temporal_frames, resolve_transfer_config, @@ -123,23 +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: - """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 - - def _validate_request_compatibility( *, output_type: str, @@ -223,22 +210,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 +234,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.""" @@ -321,19 +290,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, @@ -436,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( @@ -463,60 +415,17 @@ 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 + @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 _resolve_generation_params( - self, - output_type: str, - *, - action_mode: Optional[str] = None, - 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( - 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 + @resolver.setter + def resolver(self, resolver: ParameterResolver) -> None: + self._resolver = resolver def _log_envelope_advisory( self, @@ -643,32 +552,15 @@ 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.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 if self.audio_gen and PipelineComponent.SOUND_TOKENIZER not in skip_components: logger.info("Loading audio tokenizer...") @@ -689,11 +581,15 @@ 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: + 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") @@ -706,13 +602,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 if PipelineComponent.SCHEDULER not in skip_components: @@ -722,7 +611,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 @@ -767,16 +658,17 @@ 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): - 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): @@ -793,15 +685,14 @@ 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 - pipeline-internal, not a request field. + 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. """ - defaults = {k: v for k, v in self._mode_params("video").items() if k != "flow_shift"} - return {**defaults, **self.sampling.generation_default_overrides()} + 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 @@ -814,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} ) @@ -837,12 +727,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", @@ -902,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, COSMOS3_V2V_DEFAULT_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: @@ -984,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 @@ -1003,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 @@ -1041,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 @@ -1724,60 +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"] - - mode_params = self._mode_params( - output_type, - action_mode=normalized_action_mode, - checkpoint_policy_domain_matches=checkpoint_policy_domain_matches, - ) - 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: @@ -1790,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_params.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: @@ -3009,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_DEFAULT_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 cfdbc3531356..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. @@ -113,15 +104,18 @@ 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. + """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 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. """ - fixed_sigmas = _resolve_distilled_sigmas(scheduler.config) is_unipc = isinstance(scheduler, UniPCMultistepScheduler) is_flow_match = isinstance(scheduler, FlowMatchEulerDiscreteScheduler) @@ -175,20 +169,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/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/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py index fa4fa56881af..fb51c248693b 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,27 +43,7 @@ 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, -} - - -_PretrainedConfigT = TypeVar("_PretrainedConfigT") - - -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 +from .parameter_resolver import fill_pretrained_config def uses_static_fp8(model_config: DiffusionModelConfig) -> bool: @@ -1211,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/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/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 246e4c0a25e6..c1d0e1ef37f0 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -173,6 +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_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/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 b6f1c026341b..ab7ef834e472 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py @@ -7,37 +7,41 @@ 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_ASPECT_RATIO_LABELS, - EMBODIMENT_TO_DOMAIN_ID, - EMBODIMENT_TO_RAW_ACTION_DIM, - VIDEO_RES_SIZE_INFO, + ACTION_MODE_POLICY, 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, 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, + 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_domain_action_config, + resolve_raw_action_dim, ) +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, @@ -46,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", @@ -68,14 +90,14 @@ 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)) 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: @@ -239,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 @@ -251,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 85804719ea67..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 ( @@ -69,11 +71,22 @@ 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: 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: @@ -85,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"), @@ -93,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 @@ -165,30 +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_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 @@ -227,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} @@ -245,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()} @@ -268,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 = { @@ -281,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: @@ -599,6 +601,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 +614,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 +627,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( @@ -1022,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) @@ -1042,7 +1044,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: @@ -1060,26 +1062,37 @@ 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 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.resolver.config().prompt.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.resolver.config().prompt.use_system_prompt is True def _captured_use_system_prompt(self, pipeline, extra_params): captured = {} @@ -1165,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: @@ -1384,8 +1393,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() @@ -1435,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 6957f1646fc3..140ac56fb786 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 @@ -381,7 +385,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 @@ -771,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.""" @@ -810,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 @@ -847,53 +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", - } - ) - - policy = pipeline._mode_params("video", action_mode=ACTION_MODE_POLICY) + pipeline.resolver = _policy_resolver(NEMOTRON_DENSE_RECIPE.name) - 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, - 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"], - } + 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 ( - pipeline._mode_params("video", action_mode=ACTION_MODE_FORWARD_DYNAMICS) + COSMOS3_GENERATION_DEFAULTS[(NEMOTRON_DENSE_RECIPE.name, "action")] is COSMOS3_ACTION_PARAMS ) + for action_mode in (ACTION_MODE_FORWARD_DYNAMICS, ACTION_MODE_INVERSE_DYNAMICS): + 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 ( - pipeline._mode_params("video", action_mode=ACTION_MODE_INVERSE_DYNAMICS) - is COSMOS3_ACTION_PARAMS + COSMOS3_GENERATION_DEFAULTS[(NEMOTRON_DENSE_RECIPE.name, "video")] + is COSMOS3_EDGE_VIDEO_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_params("video") is COSMOS3_EDGE_VIDEO_PARAMS @pytest.mark.parametrize( "action_mode", @@ -940,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)) @@ -987,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) @@ -1056,23 +1011,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 @@ -1129,10 +1067,16 @@ 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) - 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 @@ -1177,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() @@ -1802,7 +1725,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, ) @@ -1873,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_params(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"] @@ -1930,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_example_prompts.py b/tests/unittest/_torch/visual_gen/test_cosmos3_example_prompts.py index 96c680dc7ff9..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 ( +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, ) -from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( - resolve_checkpoint_policy_defaults, - 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 new file mode 100644 index 000000000000..6462b9ce94b4 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_parameter_resolver.py @@ -0,0 +1,346 @@ +# 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 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 ( + CHECKPOINT, + CHECKPOINT_FILE, + CONFIG_FILE, + DEFAULT, + MODEL_INDEX_FILE, + REQUEST, + SCHEDULER_FILE, + TRANSFORMER_FILE, + UNSET, + VAE_FILE, + WORKFLOW, + CheckpointFiles, + Conflict, + ParameterResolver, + RequestParameters, + fill_pretrained_config, + resolve_checkpoint, + select, +) + +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_FILE, {})) + + +def _resolve(tmp_path: Path, files: dict): + _write(tmp_path, files) + return resolve_checkpoint(str(tmp_path), files.get(TRANSFORMER_FILE, {})) + + +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_false_and_zero_are_values(self): + assert select(False, True).value is False + assert select(0, 5).value == 0 + + +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_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"): + 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 = 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_FILE: {"a": [0.5, 1 / 3]}, MODEL_INDEX_FILE: {"a": [0.5, 0.3333333333]}}, + ) + 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_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}, + } + + +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") == ( + "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", + ) + + 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") == DEFAULT + assert cfg.action.policy is None + + def test_vae_compression_copies_must_agree(self, tmp_path): + files = _files() + 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_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_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_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_vae_spatial_compression_is_optional(self, tmp_path): + files = _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) + 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 resolver 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) + + +def _resolver(**facts) -> ParameterResolver: + return bare_resolver(facts.pop("family", "qwen3"), **facts) + + +class TestRequest: + 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): + 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): + 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} + 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): + 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.""" + 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): + 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): + 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 4797c1b61357..1d132bd0099e 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 @@ -41,15 +42,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, @@ -866,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, ) @@ -896,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, ) @@ -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 = [] @@ -928,9 +930,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 @@ -978,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, flow_shift=1.0) rebuilt = [] class StopAfterTokenize(Exception): @@ -990,9 +990,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 @@ -1026,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 @@ -1040,7 +1038,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 = [] @@ -1054,9 +1052,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 @@ -1098,7 +1093,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, flow_shift=7.0) pipeline.family = QWEN3_RECIPE.name calls = [] @@ -1112,9 +1107,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 @@ -1154,6 +1146,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: @@ -1210,6 +1203,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: @@ -1219,9 +1213,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 @@ -1265,6 +1256,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: @@ -1274,9 +1266,6 @@ class FakeSampling: def validate_request(self, num_inference_steps, guidance_scale): return None - def generation_default_overrides(self): - return {} - pipeline.sampling = FakeSampling() captured = {} @@ -1326,11 +1315,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 @@ -1358,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, ) @@ -1651,9 +1637,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 a8fce467b784..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 @@ -31,16 +32,15 @@ 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 from tensorrt_llm._torch.visual_gen.models.cosmos3.transfer import ( BILATERAL_D, BILATERAL_SIGMA_COLOR, 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, @@ -230,16 +230,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 @@ -260,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 @@ -1310,5 +1301,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) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py index 7eed96138f04..09327dfc6d3f 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.parameter_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)) diff --git a/tests/unittest/_torch/visual_gen/test_data/README.md b/tests/unittest/_torch/visual_gen/test_data/README.md index 369e5c0c5189..5d8ec43bd82b 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 `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 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 resolver +does not expect fails the test. + +Regenerate by fetching `https://huggingface.co//raw/main/` for every +`file:key` that `resolve_checkpoint` names 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 + } + } +}