diff --git a/.gitignore b/.gitignore index 47e39a1b2d71..1b512b0400af 100644 --- a/.gitignore +++ b/.gitignore @@ -122,3 +122,5 @@ tests/integration/defs/stress_test/artifacts/ .claude/agent-tests/perf-test-sync/report.html .claude/agent-tests/perf-test-sync/results.json .claude/settings.json + +.plans/ diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index 46f90a9df31c..3bd2d82f79a0 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -167,6 +167,9 @@ For full documentation, see the [Visual Generation](./visual-generation.md) page | `Qwen/Qwen-Image-2512` | Text-to-Image | | `nvidia/Cosmos3-Nano` | Text-to-Image, Text-to-Video, Image-to-Video | | `nvidia/Cosmos3-Super` | Text-to-Image, Text-to-Video, Image-to-Video | +| `nvidia/Cosmos3-Super-Text2Image-4Step` | Text-to-Image (DMD2-distilled, fixed 4-step schedule) | +| `nvidia/Cosmos3-Super-Image2Video-4Step` | Image-to-Video (DMD2-distilled, fixed 4-step schedule) | +| `nvidia/Cosmos3-Edge` | Text-to-Image, Text-to-Video, Image-to-Video (Nemotron-dense backbone, 480p-native) | ### Feature Matrix diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index d1e376b9a13e..71f4d354b255 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -39,6 +39,9 @@ TensorRT-LLM **VisualGen** provides a unified inference stack for diffusion mode | `Qwen/Qwen-Image-2512` | Text-to-Image | | `nvidia/Cosmos3-Nano` | Text-to-Image, Text-to-Video, Image-to-Video | | `nvidia/Cosmos3-Super` | Text-to-Image, Text-to-Video, Image-to-Video | +| `nvidia/Cosmos3-Super-Text2Image-4Step` | Text-to-Image (DMD2-distilled, fixed 4-step schedule) | +| `nvidia/Cosmos3-Super-Image2Video-4Step` | Image-to-Video (DMD2-distilled, fixed 4-step schedule) | +| `nvidia/Cosmos3-Edge` | Text-to-Image, Text-to-Video, Image-to-Video (Nemotron-dense backbone, 480p-native) | Models are auto-detected from the checkpoint directory. Diffusers-format models are detected via `model_index.json`; LTX-2 monolithic safetensors checkpoints are detected via embedded metadata. The `AutoPipeline` registry selects the appropriate pipeline class automatically. diff --git a/examples/visual_gen/configs/cosmos3-t2i-1gpu.yaml b/examples/visual_gen/configs/cosmos3-t2i-1gpu.yaml new file mode 100644 index 000000000000..6ea32adaebf7 --- /dev/null +++ b/examples/visual_gen/configs/cosmos3-t2i-1gpu.yaml @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-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. + +# 1-GPU Cosmos3 text-to-image deployment (base or distilled T2I checkpoints). +# Model: nvidia/Cosmos3-Super-Text2Image or nvidia/Cosmos3-Super-Text2Image-4Step +# Shared by offline examples (--visual_gen_args) and trtllm-serve. +# +# Warmup expresses the deployed workload: warms the 1024x1024 single-frame +# shape instead of the omni default (720p x 189-frame video). Requests should +# pass output_type="image". +attention_config: + backend: VANILLA +parallel_config: + cfg_size: 1 + ulysses_size: 1 +compilation_config: + resolutions: [[1024, 1024]] + num_frames: [1] diff --git a/examples/visual_gen/models/cosmos3/README.md b/examples/visual_gen/models/cosmos3/README.md index 69be21fe4880..f45bf7c14e6a 100644 --- a/examples/visual_gen/models/cosmos3/README.md +++ b/examples/visual_gen/models/cosmos3/README.md @@ -5,7 +5,7 @@ Cosmos3 supports four generation modes from a single checkpoint: - **T2V** — text-to-video (`prompts/t2v.json`). - **T2I** — text-to-image (`prompts/t2i.json`); emits a still frame (use `--output_type image` / a non-video `--output_path`). - **I2V / TI2V** — image-conditioned video (`prompts/i2v.json`). Condition on a reference frame via the prompt file's `vision_path` or `--image_path`. The image may be a local path, a `file://` / `http(s)://` URL, or a `data:` URI. -- **T2AV** — text-to-video with synchronized audio (`prompts/t2av.json` with `enable_audio: true`, or pass `--enable_audio`). Combine with a `vision_path` for image-conditioned audio-video (TI2AV). +- **T2AV** — text-to-video with synchronized audio (`prompts/t2av.json` with `enable_audio: true`, or pass `--enable_audio`). Combine with a `vision_path` for image-conditioned audio-video (TI2AV). Audio requires an audio-capable checkpoint (Nano / Super); Cosmos3-Edge has no audio tower. ## Checkpoints @@ -13,6 +13,9 @@ Pass the Hub ID or local path via `--model`: - [`nvidia/Cosmos3-Nano`](https://huggingface.co/nvidia/Cosmos3-Nano) - [`nvidia/Cosmos3-Super`](https://huggingface.co/nvidia/Cosmos3-Super) +- [`nvidia/Cosmos3-Super-Text2Image-4Step`](https://huggingface.co/nvidia/Cosmos3-Super-Text2Image-4Step) — DMD2-distilled text-to-image: fixed 4-step schedule with classifier-free guidance baked into the weights. Steps/guidance are read from the checkpoint; conflicting request values are rejected. Use with `configs/cosmos3-t2i-1gpu.yaml`. +- [`nvidia/Cosmos3-Super-Image2Video-4Step`](https://huggingface.co/nvidia/Cosmos3-Super-Image2Video-4Step) — DMD2-distilled image-to-video: same fixed 4-step, guidance-baked-in contract. The default omni video shape (720p × 189 frames) is the deployed shape, so no dedicated config is needed. This checkpoint declares `default_use_system_prompt: true` in its `model_index.json`, which the pipeline applies automatically (override with `--use_system_prompt` / `--no-use_system_prompt`). +- [`nvidia/Cosmos3-Edge`](https://huggingface.co/nvidia/Cosmos3-Edge) — 4B Nemotron-dense backbone supporting **T2I / T2V / I2V only**: no audio tower, and the checkpoint's action weights are not supported by this pipeline yet. 480p-native defaults (832×480 × 121 frames, 50 UniPC steps on the checkpoint-declared native flow schedule with shift 3.0, guidance 5.0; T2I defaults to 640×640), so no dedicated config is needed. The model card validates 256p/480p, 50–150 frames, and 12–30 FPS; requests outside that envelope run with an advisory log. ## Guardrails @@ -36,6 +39,7 @@ See `examples/visual_gen/configs/`: - `cosmos3-nano-1gpu.yaml` — 1 GPU - `cosmos3-super-4gpu.yaml` — 4 GPU, CFG + Ulysses + parallel VAE +- `cosmos3-t2i-1gpu.yaml` — 1 GPU, text-to-image deployments (base or distilled): warms the deployed 1024×1024 single-frame shape instead of the omni video shape. Example prompts live under `prompts/` (mirroring `cosmos3-internal/inputs/omni`). @@ -70,6 +74,28 @@ python cosmos3.py --model nvidia/Cosmos3-Nano \ --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml \ --output_path output.png +# T2I, distilled 4-step checkpoint (use the T2I config so warmup runs the +# image shape; steps/guidance come from the checkpoint automatically) +python cosmos3.py --model nvidia/Cosmos3-Super-Text2Image-4Step \ + --prompt_file prompts/t2i.json \ + --visual_gen_args ../../configs/cosmos3-t2i-1gpu.yaml \ + --output_type image \ + --output_path output.png + +# I2V, distilled 4-step checkpoint (steps/guidance and the system-prompt +# default come from the checkpoint automatically; defaults are the deployed +# 720p x 189-frame shape, so no config is required) +python cosmos3.py --model nvidia/Cosmos3-Super-Image2Video-4Step \ + --prompt "The camera slowly pans right across the scene" \ + --image_path https://example.com/frame.jpg \ + --output_path output.mp4 + +# Cosmos3-Edge image-to-video (480p-native defaults: 832x480 x 121 frames) +python cosmos3.py --model nvidia/Cosmos3-Edge \ + --prompt "The camera slowly pans right across the scene" \ + --image_path https://example.com/frame.jpg \ + --output_path output.mp4 + # Inline prompt (--prompt or a JSON file path) python cosmos3.py --model nvidia/Cosmos3-Nano \ --prompt "A cute puppy playing with a ball in a park" \ diff --git a/examples/visual_gen/models/cosmos3/cosmos3.py b/examples/visual_gen/models/cosmos3/cosmos3.py index de9e9e5010ad..182a97417ca6 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3.py +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -86,7 +86,8 @@ def main(): "--model", type=str, default="nvidia/Cosmos3-Nano", - help="Model path or HuggingFace Hub ID (nvidia/Cosmos3-Nano, nvidia/Cosmos3-Super)", + help="Model path or HuggingFace Hub ID " + "(nvidia/Cosmos3-Nano, nvidia/Cosmos3-Super, nvidia/Cosmos3-Edge)", ) parser.add_argument( "--visual_gen_args", @@ -136,7 +137,13 @@ def main(): help="Disable resolution metadata template (enabled by default, matching cosmos-framework CLI)", ) parser.add_argument( - "--use_system_prompt", action="store_true", help="Use system prompt in prompt" + "--use_system_prompt", + action=argparse.BooleanOptionalAction, + default=None, + help=( + "Prepend the Cosmos3 system prompt (--no-use_system_prompt to disable). " + "When omitted, the checkpoint's declared default applies." + ), ) parser.add_argument("--enable_audio", action="store_true", help="Enable audio generation") parser.add_argument( @@ -181,7 +188,8 @@ def main(): params.extra_params["use_duration_template"] = False if args.disable_resolution_template: params.extra_params["use_resolution_template"] = False - params.extra_params["use_system_prompt"] = args.use_system_prompt + if args.use_system_prompt is not None: + params.extra_params["use_system_prompt"] = args.use_system_prompt params.extra_params["enable_audio"] = enable_audio params.extra_params["use_guardrails"] = not args.disable_guardrails params.extra_params["output_type"] = output_type diff --git a/requirements.txt b/requirements.txt index dab577614f57..e7f300e9b274 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,7 +4,8 @@ accelerate>=1.7.0 build colored cuda-python>=13 -diffusers>=0.37.1 +# FlowMatchEuler respects a supplied generator starting in 0.39.0 (huggingface/diffusers#13678). +diffusers>=0.39.0 ftfy lark lazy_loader~=0.5 diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index f5747544946d..1de9e43bd6e8 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -35,10 +35,7 @@ "frame_rate": 24.0, } -# Text-to-image (``output_type="image"``) defaults. Applied by the pipeline when -# the corresponding request field still carries the merged video default, since -# the executor merges a single ``default_generation_params`` dict (the video -# params above) into the request before ``infer()`` runs. +# Text-to-image (``output_type="image"``) defaults; resolved in ``infer()``. COSMOS3_T2I_PARAMS = { "height": 1024, "width": 1024, @@ -48,6 +45,62 @@ "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. +COSMOS3_EDGE_VIDEO_PARAMS = { + "height": 480, + "width": 832, + "num_inference_steps": 50, + "guidance_scale": 5.0, + "max_sequence_length": 4096, + "num_frames": 121, + "frame_rate": 24.0, + "flow_shift": 3.0, +} + +COSMOS3_EDGE_T2I_PARAMS = { + "height": 640, + "width": 640, + "num_inference_steps": 50, + "guidance_scale": 4.0, + "flow_shift": 3.0, + "guidance_interval": None, +} + +# Model-card validated envelope for Edge; advisory only (the reference +# runtime accepts a wider range), surfaced as a log line per request. +COSMOS3_EDGE_ENVELOPE = { + "num_frames": (50, 150), + "frame_rate": (12.0, 30.0), + "max_sequence_length": 4096, + "resolutions": frozenset( + {(640, 640), (544, 736), (736, 544), (480, 832), (832, 480), (256, 256), (256, 320), + (320, 256), (192, 320), (320, 192)} + ), +} + +# (family, mode) → generation defaults. Family is the architecture recipe +# name resolved from the transformer config; mode is the request's output +# type — never inferred from the checkpoint name (a task-specialized +# checkpoint can still be asked to run any mode). +COSMOS3_GENERATION_DEFAULTS: Dict = { + ("qwen3", "video"): COSMOS3_720P_PARAMS, + ("qwen3", "image"): COSMOS3_T2I_PARAMS, + ("nemotron_dense", "video"): COSMOS3_EDGE_VIDEO_PARAMS, + ("nemotron_dense", "image"): COSMOS3_EDGE_T2I_PARAMS, +} + +# Families without an entry get no envelope advisory. +COSMOS3_ENVELOPES: Dict = { + "nemotron_dense": COSMOS3_EDGE_ENVELOPE, +} + + + + COSMOS3_EXTRA_SPECS: Dict[str, ExtraParamSchema] = { "use_duration_template": ExtraParamSchema( type="bool", 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 2dac241f2115..44f1809f050d 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -21,10 +21,10 @@ import PIL.Image import torch -from diffusers import AutoencoderKLWan, UniPCMultistepScheduler +from diffusers import AutoencoderKLWan from diffusers.utils.torch_utils import randn_tensor from diffusers.video_processor import VideoProcessor -from transformers import Qwen2Tokenizer +from transformers import AutoTokenizer from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline @@ -34,10 +34,11 @@ from tensorrt_llm.inputs.utils import load_image from tensorrt_llm.logger import logger -from .defaults import COSMOS3_720P_PARAMS, COSMOS3_EXTRA_SPECS, COSMOS3_T2I_PARAMS +from .defaults import COSMOS3_ENVELOPES, COSMOS3_EXTRA_SPECS, COSMOS3_GENERATION_DEFAULTS from .guardrails import check_video_safety, download_guardrail_checkpoint +from .sampling import Cosmos3SamplingPolicy, load_scheduler from .sound_tokenizer import LatentAutoEncoderV2 -from .transformer_cosmos3 import Cosmos3VFMTransformer +from .transformer_cosmos3 import NEMOTRON_DENSE_RECIPE, Cosmos3VFMTransformer, resolve_arch_recipe COSMOS3_DEFAULT_NEGATIVE_PROMPT = "" COSMOS3_DEFAULT_SYSTEM_PROMPT = ( @@ -53,13 +54,54 @@ TRTLLM_DISABLE_COSMOS3_GUARDRAILS = os.environ.get("TRTLLM_DISABLE_COSMOS3_GUARDRAILS", "0") == "1" +def _validate_sampling_recipe(family: str, use_native_flow_schedule: bool, sampling) -> None: + """Family, model_index schedule flag, and scheduler recipe must form a + known-supported combination — the pieces come from three different + checkpoint files and a mismatch samples the wrong trajectory silently. + """ + if family == NEMOTRON_DENSE_RECIPE.name: + if sampling.is_distilled: + raise ValueError( + "Distilled (fixed-sigma FlowMatchEuler) sampling is not supported for " + "the Edge (nemotron_dense) family; no such checkpoint exists." + ) + if not use_native_flow_schedule: + raise ValueError( + "Edge (nemotron_dense) checkpoints must declare " + "use_native_flow_schedule: true in model_index.json. Without it the " + "checkpoint's karras scheduler config would sample the wrong " + "trajectory; a missing flag means a broken or stale conversion." + ) + elif use_native_flow_schedule: + raise ValueError( + "use_native_flow_schedule is only supported for the Edge " + f"(nemotron_dense) family, but this checkpoint's family is {family!r}." + ) + + +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}." + ) + + @register_pipeline( "Cosmos3OmniMoTPipeline", hf_ids=[ "nvidia/Cosmos3-Nano", "nvidia/Cosmos3-Super", "nvidia/Cosmos3-Super-Image2Video", + "nvidia/Cosmos3-Super-Image2Video-4Step", "nvidia/Cosmos3-Super-Text2Image", + "nvidia/Cosmos3-Super-Text2Image-4Step", + "nvidia/Cosmos3-Edge", ], doc="Cosmos3 Omnimodal world models.", ) @@ -67,7 +109,14 @@ class Cosmos3OmniMoTPipeline(BasePipeline): def __init__(self, pipeline_config): primary_pretrained_config = pipeline_config.primary_pretrained_config self.audio_gen = False + # Checkpoint fact vs runtime capability: the checkpoint may ship + # action weights, but action generation is not implemented here. + self.has_action_weights = False self.action_gen = False + self.sampling = Cosmos3SamplingPolicy() + self.use_native_flow_schedule = False + self.default_use_system_prompt = COSMOS3_EXTRA_SPECS["use_system_prompt"].default + self.family = resolve_arch_recipe(primary_pretrained_config).name if getattr( primary_pretrained_config, "audio_gen", @@ -77,11 +126,75 @@ def __init__(self, pipeline_config): self.audio_gen = True if getattr(primary_pretrained_config, "action_gen", False): - logger.info("Initializing Cosmos3OmniMoTPipeline with action generation.") - self.action_gen = True + logger.info( + "Checkpoint declares action weights; action generation is not supported " + "by this pipeline (weights are skipped)." + ) + self.has_action_weights = True super().__init__(pipeline_config) + def _mode_params(self, mode: str) -> dict: + """Generation default table for this checkpoint family and request mode.""" + return COSMOS3_GENERATION_DEFAULTS[(self.family, mode)] + + def _resolve_generation_params(self, mode: str, **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(mode) + 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 + + def _log_envelope_advisory( + self, + *, + is_t2i: bool, + height: int, + width: int, + num_frames: int, + frame_rate: float, + max_sequence_length: int, + ) -> None: + """One advisory line for requests outside the model-card envelope. + + The envelope is documented support, not enforced validation: the + reference runtime accepts a wider range, so out-of-envelope requests + run — they just carry no quality claim. Families without a declared + envelope get no advisory. + """ + env = COSMOS3_ENVELOPES.get(self.family) + if env is None: + return + outside = [] + if (height, width) not in env["resolutions"]: + outside.append(f"{width}x{height} resolution") + if not is_t2i: + lo, hi = env["num_frames"] + if not lo <= num_frames <= hi: + outside.append(f"num_frames={num_frames} (validated: {lo}-{hi})") + lo, hi = env["frame_rate"] + if not lo <= frame_rate <= hi: + outside.append(f"frame_rate={frame_rate} (validated: {lo}-{hi})") + if max_sequence_length > env["max_sequence_length"]: + outside.append( + f"max_sequence_length={max_sequence_length} (validated: " + f"<= {env['max_sequence_length']})" + ) + if outside: + logger.warning( + "Request is outside the model-card validated envelope " + f"({'; '.join(outside)}); generation proceeds but quality may degrade." + ) + def _init_transformer(self) -> None: logger.info("Initializing Cosmos3VFMTransformer") model_config = self.pipeline_config.model_configs["transformer"] @@ -98,6 +211,21 @@ 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) + ) + if self.audio_gen and PipelineComponent.SOUND_TOKENIZER not in skip_components: logger.info("Loading audio tokenizer...") self.audio_tokenizer = ( @@ -112,7 +240,7 @@ def load_standard_components( if PipelineComponent.TOKENIZER not in skip_components: logger.info("Loading tokenizer...") - self.tokenizer = Qwen2Tokenizer.from_pretrained( + self.tokenizer = AutoTokenizer.from_pretrained( checkpoint_dir, subfolder="text_tokenizer", ) @@ -135,26 +263,33 @@ def load_standard_components( 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: logger.info("Loading scheduler...") - self.scheduler = UniPCMultistepScheduler.from_pretrained( - checkpoint_dir, - subfolder=PipelineComponent.SCHEDULER, - ) - # Snapshot the checkpoint scheduler config so the scheduler can be - # rebuilt at request time when a mode-specific ``flow_shift`` is - # needed (T2I uses shift=3.0; T2V/I2V keep the checkpoint default). - self._base_scheduler_config = self.scheduler.config - self._engine_init_flow_shift = float( - getattr(self.scheduler.config, "flow_shift", 1.0) or 1.0 + # The scheduler class comes from the checkpoint: UniPC for base + # checkpoints, FlowMatchEuler (fixed stochastic schedule) for + # 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._current_flow_shift = self._engine_init_flow_shift + _validate_sampling_recipe(self.family, self.use_native_flow_schedule, self.sampling) + # flow_shift is a per-mode table fact, not a request field, so the + # mode schedulers are prebuilt once; forward() only picks one. + base_scheduler = self.scheduler + self.mode_schedulers = {} + for mode in ("video", "image"): + shift = self._mode_params(mode).get("flow_shift") + if shift is None: + shift = self.sampling.checkpoint_flow_shift + self.mode_schedulers[mode] = self.sampling.with_flow_shift(base_scheduler, shift) + self.scheduler = self.mode_schedulers["video"] if self.audio_gen: - # Separate instance so video and audio scheduler states don't collide - # (UniPC mutates internal correction buffers on every .step() call). - self.audio_scheduler = UniPCMultistepScheduler.from_config(self.scheduler.config) + # Separate instance so video and audio scheduler states don't + # collide (schedulers mutate internal state on every .step()). + self.audio_scheduler = type(self.scheduler).from_config(self.scheduler.config) # Re-check the env var in case it was changed after initialization like in unit tests. guardrails_disabled = os.environ.get("TRTLLM_DISABLE_COSMOS3_GUARDRAILS", "0") == "1" @@ -185,40 +320,55 @@ def load_standard_components( self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) - def _set_flow_shift(self, target_shift: float) -> None: - """Rebuild the UniPC scheduler with ``flow_shift=target_shift`` if needed. - - T2I uses ``flow_shift=3.0`` while T2V/I2V use the checkpoint default. - ``self._current_flow_shift`` is tracked explicitly so a prior T2I rebuild - does not leak into a subsequent video request. - """ - if not hasattr(self, "_base_scheduler_config"): - return - target = float(target_shift) - if target == float(self._current_flow_shift): - return - self.scheduler = UniPCMultistepScheduler.from_config( - self._base_scheduler_config, flow_shift=target - ) - self._current_flow_shift = target - @property def default_warmup_resolutions(self): - return [(720, 1280)] + video = self._mode_params("video") + return [(video["height"], video["width"])] @property def default_warmup_num_frames(self): - return [189] + return [self._mode_params("video")["num_frames"]] + + @property + def default_warmup_steps(self): + # Distilled checkpoints only run their fixed schedule length. + return self.sampling.num_steps(super().default_warmup_steps) @property def default_generation_params(self): - return dict(COSMOS3_720P_PARAMS) + """Fields merged by the executor into every request. + + Mode-dependent values remain None until infer() selects the request + mode; key membership also declares these fields supported during + request validation. ``flow_shift`` is pipeline-internal, not a + request field. + """ + defaults = {k: v for k, v in self._mode_params("video").items() if k != "flow_shift"} + defaults.update( + { + "height": None, + "width": None, + "num_inference_steps": None, + "guidance_scale": None, + } + ) + return {**defaults, **self.sampling.generation_default_overrides()} @property def extra_param_specs(self): - return dict(COSMOS3_EXTRA_SPECS) + specs = dict(COSMOS3_EXTRA_SPECS) + specs["use_system_prompt"] = specs["use_system_prompt"].model_copy( + update={"default": self.default_use_system_prompt} + ) + return specs 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 + 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", @@ -227,51 +377,29 @@ def _run_warmup(self, height: int, width: int, num_frames: int, steps: int) -> N width=width, num_frames=num_frames, num_inference_steps=steps, - guidance_scale=COSMOS3_720P_PARAMS["guidance_scale"], + guidance_scale=guidance_scale, seed=42, - max_sequence_length=COSMOS3_720P_PARAMS["max_sequence_length"], + max_sequence_length=defaults["max_sequence_length"], use_guardrails=False, image=None, enable_audio=False, ) - @staticmethod - def _resolve_t2i_default(merged_value, video_default, t2i_default): - """Pick the T2I default when the field still carries the merged video default. - - The executor merges a single ``default_generation_params`` dict (the - video params) into the request before ``infer()``, so an unspecified - field arrives equal to its video default. For T2I we substitute the - T2I default in that case while honoring any explicit user override. - """ - return t2i_default if merged_value == video_default else merged_value - def infer(self, req): extra_params = req.params.extra_params or {} output_type = extra_params.get("output_type", "video") is_t2i = str(output_type).lower() == "image" - height = req.params.height - width = req.params.width - num_inference_steps = req.params.num_inference_steps - guidance_scale = req.params.guidance_scale - if is_t2i: - height = self._resolve_t2i_default( - height, COSMOS3_720P_PARAMS["height"], COSMOS3_T2I_PARAMS["height"] - ) - width = self._resolve_t2i_default( - width, COSMOS3_720P_PARAMS["width"], COSMOS3_T2I_PARAMS["width"] - ) - num_inference_steps = self._resolve_t2i_default( - num_inference_steps, - COSMOS3_720P_PARAMS["num_inference_steps"], - COSMOS3_T2I_PARAMS["num_inference_steps"], - ) - guidance_scale = self._resolve_t2i_default( - guidance_scale, - COSMOS3_720P_PARAMS["guidance_scale"], - COSMOS3_T2I_PARAMS["guidance_scale"], - ) + # None = unset; resolve by mode exactly once. Non-None values pass through. + mode_params = self._mode_params("image" if is_t2i else "video") + + def resolved(value, field_name): + return value if value is not None else mode_params[field_name] + + height = resolved(req.params.height, "height") + width = resolved(req.params.width, "width") + num_inference_steps = resolved(req.params.num_inference_steps, "num_inference_steps") + guidance_scale = resolved(req.params.guidance_scale, "guidance_scale") return self.forward( prompt=req.prompt, @@ -293,7 +421,7 @@ def infer(self, req): "use_resolution_template", COSMOS3_EXTRA_SPECS["use_resolution_template"].default, ), - use_system_prompt=extra_params.get("use_system_prompt", False), + use_system_prompt=extra_params.get("use_system_prompt", self.default_use_system_prompt), use_guardrails=extra_params.get("use_guardrails", True), enable_audio=extra_params.get("enable_audio", False), output_type=output_type, @@ -564,6 +692,30 @@ def _prepare_latents_i2v( velocity_mask = 1.0 - condition_mask return latents, velocity_mask, image_latent + def _conditioning_anchor_post_step(self, image_latent: Optional[torch.Tensor]): + """Per-step re-anchor of the conditioned frame for distilled sampling. + + The distilled FlowMatchEuler step is stochastic: it re-noises every + position, including the frame the velocity mask holds still, so the + conditioning frame the model reads as clean context degrades from step + 2 on. Writing the clean latent back after every scheduler step keeps + it clean (diffusers' distilled loop re-anchors the same way). + Deterministic UniPC steps never move a zero-velocity frame, so base + checkpoints need no per-step anchor and keep their exact behavior. + + Returns a ``post_step_fn`` for ``BasePipeline.denoise``, or ``None`` + when no anchoring is needed. + """ + if not self.sampling.is_distilled or image_latent is None: + return None + + def post_step_fn(latents: torch.Tensor) -> torch.Tensor: + # In-place: writes one latent frame, no full-tensor copies. + latents[:, :, 0:1] = image_latent + return latents + + return post_step_fn + # ========================================================================= # VAE decode # ========================================================================= @@ -620,13 +772,13 @@ def forward( seed: int, negative_prompt: Optional[str] = None, image: Optional[Union[PIL.Image.Image, torch.Tensor, str]] = None, - height: int = COSMOS3_720P_PARAMS["height"], - width: int = COSMOS3_720P_PARAMS["width"], - num_frames: int = COSMOS3_720P_PARAMS["num_frames"], - num_inference_steps: int = COSMOS3_720P_PARAMS["num_inference_steps"], - guidance_scale: float = COSMOS3_720P_PARAMS["guidance_scale"], - max_sequence_length: int = COSMOS3_720P_PARAMS["max_sequence_length"], - frame_rate: float = COSMOS3_720P_PARAMS["frame_rate"], + height: Optional[int] = None, + width: Optional[int] = None, + num_frames: Optional[int] = None, + num_inference_steps: Optional[int] = None, + guidance_scale: Optional[float] = None, + max_sequence_length: Optional[int] = None, + frame_rate: Optional[float] = None, use_duration_template: bool = COSMOS3_EXTRA_SPECS["use_duration_template"].default, use_resolution_template: bool = COSMOS3_EXTRA_SPECS["use_resolution_template"].default, use_system_prompt: bool = COSMOS3_EXTRA_SPECS["use_system_prompt"].default, @@ -634,6 +786,14 @@ def forward( enable_audio: bool = COSMOS3_EXTRA_SPECS["enable_audio"].default, output_type: str = COSMOS3_EXTRA_SPECS["output_type"].default, ): + """Run one generation. ``infer()`` is the resolved entry point. + + Production requests arrive through ``infer()`` with fully resolved + values; unset (None) numeric parameters resolve here from the same + per-variant mode tables, so direct internal callers get + checkpoint-appropriate values (including the fixed distilled + steps/guidance). + """ pipeline_start = time.time() timer = CudaPhaseTimer() timer.mark_pre_start() @@ -643,7 +803,41 @@ def forward( # Text-to-image mode: same checkpoint/forward path as T2V, but a single # latent frame, image-flavored prompt templates, flow_shift=3.0, a CFG # guidance interval, and an image (rather than video) output. - is_t2i = str(output_type).lower() == "image" + output_type = str(output_type).lower() + if output_type not in ("video", "image"): + raise ValueError(f"output_type must be 'video' or 'image', got {output_type!r}.") + is_t2i = output_type == "image" + + mode_params = self._mode_params(output_type) + resolved = self._resolve_generation_params( + output_type, + height=height, + width=width, + num_frames=num_frames, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + max_sequence_length=max_sequence_length, + frame_rate=frame_rate, + ) + height = resolved["height"] + width = resolved["width"] + num_frames = resolved["num_frames"] + num_inference_steps = resolved["num_inference_steps"] + guidance_scale = resolved["guidance_scale"] + max_sequence_length = resolved["max_sequence_length"] + frame_rate = resolved["frame_rate"] + + self.sampling.validate_request(num_inference_steps, guidance_scale) + + self._log_envelope_advisory( + is_t2i=is_t2i, + height=height, + width=width, + num_frames=num_frames, + frame_rate=frame_rate, + max_sequence_length=max_sequence_length, + ) + guidance_interval = None if is_t2i: if image is not None: @@ -652,12 +846,19 @@ def forward( ) num_frames = 1 enable_audio = False - guidance_interval = COSMOS3_T2I_PARAMS["guidance_interval"] - self._set_flow_shift(COSMOS3_T2I_PARAMS["flow_shift"]) - else: - # Restore the checkpoint flow_shift in case a prior T2I request - # rebuilt the scheduler with shift=3.0. - self._set_flow_shift(getattr(self, "_engine_init_flow_shift", 1.0)) + guidance_interval = mode_params["guidance_interval"] + + self.scheduler = self.mode_schedulers[output_type] + + # Weight-presence guard, not workflow policy: the request explicitly + # asks for audio, but the checkpoint ships no audio tower. Silently + # returning a silent video would hide the capability limit. + if enable_audio and not self.audio_gen: + raise ValueError( + "enable_audio=True, but this checkpoint has no audio tower " + "(transformer config declares sound_gen=false). Drop enable_audio " + "or use an audio-capable Cosmos3 checkpoint." + ) if isinstance(prompt, str): prompt = [prompt] @@ -781,7 +982,7 @@ def forward( video_shape = (T_latent, H_latent, W_latent) # 3. Set up scheduler - self.scheduler.set_timesteps(num_inference_steps, device=self.device) + self.sampling.set_timesteps(self.scheduler, num_inference_steps, device=self.device) # 3b. Audio noise init — latent length matches diffusers Cosmos3OmniPipeline.prepare_latents. do_audio = enable_audio and self.audio_gen and hasattr(self, "audio_tokenizer") @@ -798,7 +999,9 @@ def forward( dtype=latents.dtype, ) # Audio uses the same scheduler type/config as video. - self.audio_scheduler.set_timesteps(num_inference_steps, device=self.device) + self.sampling.set_timesteps( + self.audio_scheduler, num_inference_steps, device=self.device + ) # 4. Build forward_fn for the denoise loop def forward_fn( @@ -861,6 +1064,8 @@ def forward_fn( extra_cfg_tensors=extra_cfg_tensors, extra_streams=extra_streams, guidance_interval=guidance_interval, + scheduler_step_kwargs=self.sampling.scheduler_step_kwargs(generator), + post_step_fn=self._conditioning_anchor_post_step(image_latent), ) if extra_streams is not None: @@ -877,8 +1082,8 @@ def forward_fn( decode_start = time.time() if image_latent is not None: - latents = latents.clone() - latents[:, :, 0:1, :, :] = image_latent.to(device=latents.device, dtype=latents.dtype) + # In-place: the loop output is consumed only by the decode below. + latents[:, :, 0:1] = image_latent video = self.decode_latents(latents, self._decode_latents) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/sampling.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/sampling.py new file mode 100644 index 000000000000..a13a54e93839 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/sampling.py @@ -0,0 +1,253 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Per-checkpoint sampling policy for Cosmos3. + +Exactly two recipes are supported, read from the checkpoint's +``scheduler/scheduler_config.json``: + +* ``UniPCMultistepScheduler`` without fixed sigmas — base checkpoints: + request tables drive steps/guidance; per-mode flow shifts are table facts, + so the pipeline prebuilds one scheduler per mode at load time. +* ``FlowMatchEulerDiscreteScheduler`` with a nonempty + ``fixed_step_sampler_config.t_list`` — distilled checkpoints: the step + count is locked to the schedule, classifier-free guidance is baked into + the weights (scale 1.0), and the stochastic steps draw seeded noise. + +The pipeline owns the scheduler instances; :class:`Cosmos3SamplingPolicy` is +an immutable value object of config-derived facts whose methods take +schedulers as arguments. +""" + +import json +import os +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any, Optional + +import numpy as np +from diffusers import FlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler + +from tensorrt_llm.logger import logger + +# Distilled checkpoints bake classifier-free guidance into the weights; the +# only valid scale is 1.0 ("off": a single conditional forward per step). +DISTILLED_GUIDANCE_SCALE = 1.0 + + +def _config_get(config: Any, key: str, default: Any = None) -> Any: + """Fetch a key from a plain dict, a diffusers FrozenDict, or a config object.""" + if isinstance(config, Mapping): + return config.get(key, default) + 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. + + Base checkpoints declare ``UniPCMultistepScheduler`` (a missing declaration + also resolves to UniPC, preserving pre-declaration Cosmos3 behavior); + distilled ones declare ``FlowMatchEulerDiscreteScheduler``. An explicitly + unknown declaration is a load-time error — silently substituting UniPC + would sample the checkpoint with the wrong integrator. + """ + config_path = os.path.join(checkpoint_dir, subfolder, "scheduler_config.json") + class_name = "" + if os.path.exists(config_path): + with open(config_path) as f: + class_name = json.load(f).get("_class_name", "") + if class_name == "FlowMatchEulerDiscreteScheduler": + scheduler_cls = FlowMatchEulerDiscreteScheduler + elif class_name in ("", None, "UniPCMultistepScheduler"): + scheduler_cls = UniPCMultistepScheduler + else: + raise ValueError( + f"Unsupported Cosmos3 scheduler class {class_name!r}; supported: " + "UniPCMultistepScheduler (base), FlowMatchEulerDiscreteScheduler (distilled)." + ) + return scheduler_cls.from_pretrained(checkpoint_dir, subfolder=subfolder) + + +@dataclass(frozen=True) +class Cosmos3SamplingPolicy: + """Immutable sampling facts of a loaded Cosmos3 checkpoint. + + Methods take scheduler instances as arguments; the current flow shift is + read from the supplied scheduler's config rather than tracked here. + """ + + # Fixed distilled schedule (t_list); None for base checkpoints. + fixed_sigmas: "tuple[float, ...] | None" = None + # Checkpoint scheduler config, kept for flow-shift rebuilds (UniPC only). + unipc_base_config: Optional[Any] = None + # Checkpoint-declared (model_index): UniPC runs on explicit linear flow + # sigmas with a runtime shift instead of the config's karras grid. + native_flow_schedule: bool = False + + @classmethod + def from_scheduler( + cls, scheduler: Any, native_flow_schedule: bool = False + ) -> "Cosmos3SamplingPolicy": + """Derive the policy from a loaded scheduler's config. + + Valid recipes: UniPC without fixed sigmas (base) and FlowMatchEuler + with a nonempty ``fixed_step_sampler_config.t_list`` (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) + + if ( + _config_get(scheduler.config, "fixed_step_requires_explicit_sigmas", False) + and fixed_sigmas is None + ): + raise ValueError( + "Malformed distilled checkpoint: the scheduler config declares " + "fixed_step_requires_explicit_sigmas but carries no usable " + "fixed_step_sampler_config.t_list." + ) + + if is_unipc and fixed_sigmas is None: + return cls( + fixed_sigmas=None, + unipc_base_config=scheduler.config, + native_flow_schedule=native_flow_schedule, + ) + + if is_flow_match and fixed_sigmas is not None: + logger.info( + f"Distilled Cosmos3 checkpoint: fixed {len(fixed_sigmas)}-step schedule " + f"{list(fixed_sigmas)}, classifier-free guidance baked in." + ) + return cls(fixed_sigmas=fixed_sigmas, unipc_base_config=None) + + raise ValueError( + f"Unsupported Cosmos3 sampling recipe: {type(scheduler).__name__} with " + f"fixed sigmas {'present' if fixed_sigmas is not None else 'absent'}. " + "Supported: UniPCMultistepScheduler without fixed sigmas (base), " + "FlowMatchEulerDiscreteScheduler with fixed_step_sampler_config.t_list " + "(distilled)." + ) + + @property + 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 ``COSMOS3_720P_PARAMS`` 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 + + def validate_request( + self, num_inference_steps: Optional[int], guidance_scale: Optional[float] + ) -> None: + """Reject sampling parameters incompatible with a distilled checkpoint.""" + if not self.is_distilled: + return + distilled_steps = len(self.fixed_sigmas) + if num_inference_steps is not None and num_inference_steps != distilled_steps: + raise ValueError( + "This is a distilled Cosmos3 checkpoint; the step count is fixed by the " + f"scheduler's fixed_step_sampler_config.t_list ({distilled_steps} steps). " + f"num_inference_steps must be {distilled_steps} or left unset " + f"(got {num_inference_steps})." + ) + if guidance_scale is not None and float(guidance_scale) != DISTILLED_GUIDANCE_SCALE: + raise ValueError( + "This is a distilled Cosmos3 checkpoint; classifier-free guidance is baked " + f"into the weights. guidance_scale must be {DISTILLED_GUIDANCE_SCALE} or " + f"left unset (got {guidance_scale})." + ) + + def set_timesteps(self, scheduler: Any, num_inference_steps: int, device: Any) -> None: + """Program a scheduler for one generation: fixed sigmas or a step count.""" + if self.is_distilled: + scheduler.set_timesteps(sigmas=list(self.fixed_sigmas), device=device) + elif self.native_flow_schedule: + # The PyTorch-backend base grid: linear flow sigmas over + # (1 - 1/T, 0]. UniPC applies its flow_shift to provided sigmas; + # a numpy array is required (a list breaks diffusers 0.39). + num_train = int(_config_get(scheduler.config, "num_train_timesteps", 1000)) + sigmas = np.linspace(1.0 - 1.0 / num_train, 0.0, num_inference_steps + 1)[:-1] + scheduler.set_timesteps(num_inference_steps, device=device, sigmas=sigmas) + else: + scheduler.set_timesteps(num_inference_steps, device=device) + + def scheduler_step_kwargs(self, generator: Any) -> dict: + """Extra kwargs each ``scheduler.step()`` call requires. + + The distilled FlowMatchEuler scheduler is stochastic: every step draws + SDE noise, which must come from the request-seeded ``generator`` — + otherwise it comes from the process-global RNG, breaking seed + reproducibility and diverging the replicated latents across ranks + (each rank's global RNG state is independent). UniPC steps are + deterministic and accept no ``generator`` argument, so base + checkpoints pass nothing. + """ + if self.is_distilled: + return {"generator": generator} + return {} + + @property + def checkpoint_flow_shift(self) -> float: + """The flow shift the checkpoint shipped with (UniPC only; 1.0 otherwise).""" + if self.unipc_base_config is None: + return 1.0 + return float(_config_get(self.unipc_base_config, "flow_shift", 1.0) or 1.0) + + def with_flow_shift(self, scheduler: Any, target_shift: Optional[float]) -> Any: + """A scheduler configured with ``flow_shift=target_shift``. + + Reuses ``scheduler`` when its config already matches, otherwise builds + a new instance from the checkpoint config (diffusers scheduler configs + are frozen). Called once per mode at component-load time. Structural + no-op for distilled checkpoints (no UniPC base config) and for + ``target_shift=None``. + """ + if target_shift is None or self.unipc_base_config is None: + return scheduler + target_shift = float(target_shift) + current_shift = float(_config_get(scheduler.config, "flow_shift", 1.0) or 1.0) + karras_mismatch = self.native_flow_schedule and bool( + _config_get(scheduler.config, "use_karras_sigmas", False) + ) + if current_shift == target_shift and not karras_mismatch: + return scheduler + overrides = {"flow_shift": target_shift} + if self.native_flow_schedule: + overrides["use_karras_sigmas"] = False + return UniPCMultistepScheduler.from_config(self.unipc_base_config, **overrides) 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 b3df78a06ec1..58b946ad724c 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -15,7 +15,7 @@ import math from dataclasses import dataclass -from typing import Optional, Tuple +from typing import Optional, Tuple, TypeVar import torch import torch.nn as nn @@ -25,7 +25,9 @@ from tensorrt_llm._torch.attention_backend.interface import PredefinedAttentionMask from tensorrt_llm._torch.modules.embedding import Embedding from tensorrt_llm._torch.modules.gated_mlp import GatedMLP -from tensorrt_llm._torch.modules.linear import Linear +from tensorrt_llm._torch.modules.linear import Linear, WeightMode +from tensorrt_llm._torch.modules.mlp import MLP +from tensorrt_llm._torch.utils import relu2 from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig from tensorrt_llm._torch.visual_gen.models.modeling import BaseDiffusionModel from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode @@ -34,6 +36,160 @@ 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 + + +COSMOS3_EDGE_BACKBONE_TYPE = "cosmos3_edge_nemotron_dense" + + +def resolve_rope_axes_dim(pretrained_config) -> list: + """MRoPE axes: the explicit top-level ``rope_axes_dim`` wins (diffusers + precedence); the legacy ``rope_scaling["mrope_section"]`` is the fallback. + Both-declared-but-contradictory, wrong length, or a sum that does not + cover half the head dim are config errors. + """ + top_level = getattr(pretrained_config, "rope_axes_dim", None) + rope_scaling = getattr(pretrained_config, "rope_scaling", None) or {} + nested = rope_scaling.get("mrope_section") + + if top_level is not None and nested is not None and list(top_level) != list(nested): + raise ValueError( + f"Cosmos3 config declares contradictory MRoPE axes: rope_axes_dim=" + f"{list(top_level)} vs rope_scaling.mrope_section={list(nested)}." + ) + axes = top_level if top_level is not None else nested + if axes is None: + raise ValueError( + "Cosmos3 config declares neither rope_axes_dim nor rope_scaling.mrope_section." + ) + axes = list(axes) + half_head_dim = pretrained_config.head_dim // 2 + if len(axes) != 3 or sum(axes) != half_head_dim: + raise ValueError( + f"Cosmos3 MRoPE axes {axes} must have 3 entries summing to " + f"head_dim/2 = {half_head_dim}." + ) + return axes + + +@dataclass(frozen=True) +class Cosmos3ArchRecipe: + """Complete architecture combination selected by ``backbone_type``. + + Module construction derives from the recipe, never from individual config + flags, so a mixed config cannot half-apply. + """ + + name: str + gated_mlp: bool + und_qk_norm: bool + use_und_k_norm_for_gen: bool + nemotron_norms: bool + + +QWEN3_RECIPE = Cosmos3ArchRecipe( + name="qwen3", + gated_mlp=True, + und_qk_norm=True, + use_und_k_norm_for_gen=False, + nemotron_norms=False, +) + +NEMOTRON_DENSE_RECIPE = Cosmos3ArchRecipe( + name="nemotron_dense", + gated_mlp=False, + und_qk_norm=False, + use_und_k_norm_for_gen=True, + nemotron_norms=True, +) + + +def resolve_arch_recipe(pretrained_config) -> Cosmos3ArchRecipe: + """Select and validate the architecture recipe declared by the config.""" + backbone_type = getattr(pretrained_config, "backbone_type", None) + if backbone_type is None: + recipe = QWEN3_RECIPE + expected_flags = { + "hidden_act": (None, "silu"), + "qk_norm_for_text": (None, True), + "use_und_k_norm_for_gen": (None, False), + } + elif backbone_type == COSMOS3_EDGE_BACKBONE_TYPE: + recipe = NEMOTRON_DENSE_RECIPE + expected_flags = { + "hidden_act": ("relu2",), + "qk_norm_for_text": (False,), + "use_und_k_norm_for_gen": (True,), + "sound_gen": (False,), + "attention_bias": (False,), + "rms_norm_eps": (1e-5,), + } + else: + raise ValueError( + f"Unsupported Cosmos3 transformer backbone_type={backbone_type!r}; " + f"supported: absent (Qwen3 family) or {COSMOS3_EDGE_BACKBONE_TYPE!r}." + ) + + for key, allowed in expected_flags.items(): + actual = getattr(pretrained_config, key, None) + if actual not in allowed: + raise ValueError( + f"Cosmos3 config contradicts the {recipe.name!r} recipe: " + f"{key}={actual!r}, expected one of {allowed}." + ) + + # Latent-geometry invariants are validated for the Edge recipe only: every + # published Edge checkpoint declares 48/2, so a different value means a + # wrong or corrupt config. The qwen3 recipe is left unchecked because + # long-standing reduced-dimension test fixtures build it with tiny latent + # channels; real qwen3 checkpoints are protected by weight-shape checks. + if recipe is NEMOTRON_DENSE_RECIPE: + invariants = {"latent_channel": 48, "latent_patch_size": 2} + for key, expected in invariants.items(): + actual = getattr(pretrained_config, key, None) + if actual != expected: + raise ValueError( + f"Unsupported Cosmos3 Edge transformer config: {key}={actual!r}, " + f"expected {expected}." + ) + + # A declared patch_latent_dim must agree with the latent geometry it is + # derived from (the transformer recomputes it and would silently ignore an + # inconsistent declaration). + declared_patch_dim = getattr(pretrained_config, "patch_latent_dim", None) + latent_channel = getattr(pretrained_config, "latent_channel", None) + latent_patch_size = getattr(pretrained_config, "latent_patch_size", None) + if None not in (declared_patch_dim, latent_channel, latent_patch_size): + expected_patch_dim = (latent_patch_size**2) * latent_channel + if declared_patch_dim != expected_patch_dim: + raise ValueError( + f"Inconsistent Cosmos3 transformer config: patch_latent_dim=" + f"{declared_patch_dim}, but latent_patch_size**2 * latent_channel = " + f"{expected_patch_dim}." + ) + + return recipe + class Qwen3VLTextRMSNorm(nn.Module): def __init__( @@ -59,6 +215,19 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return output +class NemotronRMSNorm(Qwen3VLTextRMSNorm): + """RMSNorm with the weight multiply in float32 before downcast. + + ``F.rms_norm`` is bit-exact to this flavor for bf16 inputs, unlike the + parent's post-downcast bf16 weight multiply. + """ + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return F.rms_norm( + hidden_states, (hidden_states.shape[-1],), self.weight, self.variance_epsilon + ) + + @dataclass class TransformerOutput: """Velocity predictions from Cosmos3VFMTransformer.forward().""" @@ -232,6 +401,7 @@ def __init__( num_key_value_heads: int, head_dim: int, model_config: DiffusionModelConfig, + recipe: Cosmos3ArchRecipe, layer_idx: int = 0, module_name: Optional[str] = None, ): @@ -249,14 +419,28 @@ def __init__( module_name=module_name, enable_sequence_parallel=False, ) - self.norm_q = Qwen3VLTextRMSNorm(hidden_size=head_dim, dtype=torch.bfloat16) - self.norm_k = Qwen3VLTextRMSNorm(hidden_size=head_dim, dtype=torch.bfloat16) + # Attention Q/K norms run the fp32-weight-multiply flavor in both + # recipes (this path has always been F.rms_norm); only the layernorms + # differ per recipe. + eps = model_config.pretrained_config.rms_norm_eps + if recipe.und_qk_norm: + self.norm_q = NemotronRMSNorm(hidden_size=head_dim, eps=eps, dtype=torch.bfloat16) + self.norm_k = NemotronRMSNorm(hidden_size=head_dim, eps=eps, dtype=torch.bfloat16) + else: + self.norm_q = None + self.norm_k = None + if recipe.use_und_k_norm_for_gen: + self.k_norm_und_for_gen = NemotronRMSNorm( + hidden_size=head_dim, eps=eps, dtype=torch.bfloat16 + ) + else: + self.k_norm_und_for_gen = None def apply_qk_norm(self, q: torch.Tensor, k: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """Per-head RMSNorm on 4D tensors [B, S, H, D].""" - q = F.rms_norm(q, (q.shape[-1],), self.norm_q.weight, self.norm_q.variance_epsilon) - k = F.rms_norm(k, (k.shape[-1],), self.norm_k.weight, self.norm_k.variance_epsilon) - return q, k + if self.norm_q is None: + return q, k + return self.norm_q(q), self.norm_k(k) def forward_with_kv( self, @@ -273,9 +457,20 @@ def forward_with_kv( k = k.view(batch_size, seq_len, self.local_num_key_value_heads, self.head_dim) v = v.view(batch_size, seq_len, self.local_num_key_value_heads, self.head_dim) + # The gen tower consumes a separately normed view of the raw und keys; + # the und self-attention below must not see that norm. + k_for_gen = None + if self.k_norm_und_for_gen is not None: + k_for_gen = self.k_norm_und_for_gen(k) + q, k = self.apply_qk_norm(q, k) q, k = qwen3_apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin) + if k_for_gen is not None: + _, k_for_gen = qwen3_apply_rotary_pos_emb(q, k_for_gen, freqs_cos, freqs_sin) + else: + k_for_gen = k + out = self._attn_impl( q, k, @@ -284,7 +479,7 @@ def forward_with_kv( timestep=timestep, ) - return self.to_out[0](out), k, v + return self.to_out[0](out), k_for_gen, v def forward(self): raise NotImplementedError( @@ -310,6 +505,7 @@ def __init__( num_key_value_heads: int, head_dim: int, model_config: DiffusionModelConfig, + recipe: Cosmos3ArchRecipe, layer_idx: int = 0, module_name: Optional[str] = None, ): @@ -334,14 +530,15 @@ def __init__( ) model_config.attention.backend = original_backend - self.norm_q = Qwen3VLTextRMSNorm(hidden_size=head_dim, dtype=torch.bfloat16) - self.norm_k = Qwen3VLTextRMSNorm(hidden_size=head_dim, dtype=torch.bfloat16) + # Same flavor note as Cosmos3CausalAttention: attention Q/K norms are + # fp32-weight-multiply in both recipes. + eps = model_config.pretrained_config.rms_norm_eps + self.norm_q = NemotronRMSNorm(hidden_size=head_dim, eps=eps, dtype=torch.bfloat16) + self.norm_k = NemotronRMSNorm(hidden_size=head_dim, eps=eps, dtype=torch.bfloat16) def apply_qk_norm(self, q: torch.Tensor, k: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """Per-head RMSNorm on 4D tensors [B, S, H, D].""" - q = F.rms_norm(q, (q.shape[-1],), self.norm_q.weight, self.norm_q.variance_epsilon) - k = F.rms_norm(k, (k.shape[-1],), self.norm_k.weight, self.norm_k.variance_epsilon) - return q, k + return self.norm_q(q), self.norm_k(k) def forward( self, @@ -406,14 +603,46 @@ def forward( return self.to_out[0](out) +def _build_cosmos3_mlp( + model_config: DiffusionModelConfig, recipe: Cosmos3ArchRecipe, layer_idx: int +) -> nn.Module: + hidden_size = model_config.pretrained_config.hidden_size + intermediate_size = model_config.pretrained_config.intermediate_size + if recipe.gated_mlp: + return GatedMLP( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + bias=False, + dtype=torch.bfloat16, + config=model_config, + layer_idx=layer_idx, + reduce_output=model_config.mapping.tp_size > 1, + ) + return MLP( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + bias=False, + activation=relu2, + dtype=torch.bfloat16, + config=model_config, + layer_idx=layer_idx, + reduce_output=model_config.mapping.tp_size > 1, + ) + + +def _layer_norm_cls(recipe: Cosmos3ArchRecipe): + return NemotronRMSNorm if recipe.nemotron_norms else Qwen3VLTextRMSNorm + + class Cosmos3UndDecoderLayer(nn.Module): """Understanding pathway decoder layer: causal self-attention + MLP.""" - def __init__(self, model_config: DiffusionModelConfig, layer_idx: int): + def __init__( + self, model_config: DiffusionModelConfig, layer_idx: int, recipe: Cosmos3ArchRecipe + ): super().__init__() self.layer_idx = layer_idx hidden_size = model_config.pretrained_config.hidden_size - intermediate_size = model_config.pretrained_config.intermediate_size self.self_attn = Cosmos3CausalAttention( hidden_size=hidden_size, @@ -421,28 +650,22 @@ def __init__(self, model_config: DiffusionModelConfig, layer_idx: int): num_key_value_heads=model_config.pretrained_config.num_key_value_heads, head_dim=model_config.pretrained_config.head_dim, model_config=model_config, + recipe=recipe, layer_idx=layer_idx, module_name=f"layers.{layer_idx}.self_attn", ) - self.input_layernorm = Qwen3VLTextRMSNorm( + norm_cls = _layer_norm_cls(recipe) + self.input_layernorm = norm_cls( hidden_size=hidden_size, eps=model_config.pretrained_config.rms_norm_eps, dtype=torch.bfloat16, ) - self.post_attention_layernorm = Qwen3VLTextRMSNorm( + self.post_attention_layernorm = norm_cls( hidden_size=hidden_size, eps=model_config.pretrained_config.rms_norm_eps, dtype=torch.bfloat16, ) - self.mlp = GatedMLP( - hidden_size=hidden_size, - intermediate_size=intermediate_size, - bias=False, - dtype=torch.bfloat16, - config=model_config, - layer_idx=layer_idx, - reduce_output=model_config.mapping.tp_size > 1, - ) + self.mlp = _build_cosmos3_mlp(model_config, recipe, layer_idx) def forward( self, @@ -479,11 +702,12 @@ def forward( class Cosmos3GenDecoderLayer(nn.Module): """Generation pathway decoder layer: cross-attention (to UND K/V) + MLP.""" - def __init__(self, model_config: DiffusionModelConfig, layer_idx: int): + def __init__( + self, model_config: DiffusionModelConfig, layer_idx: int, recipe: Cosmos3ArchRecipe + ): super().__init__() self.layer_idx = layer_idx hidden_size = model_config.pretrained_config.hidden_size - intermediate_size = model_config.pretrained_config.intermediate_size self.cross_attention = Cosmos3CrossAttention( hidden_size=hidden_size, @@ -491,28 +715,22 @@ def __init__(self, model_config: DiffusionModelConfig, layer_idx: int): num_key_value_heads=model_config.pretrained_config.num_key_value_heads, head_dim=model_config.pretrained_config.head_dim, model_config=model_config, + recipe=recipe, layer_idx=layer_idx, module_name=f"layers.{layer_idx}.cross_attention", ) - self.input_layernorm = Qwen3VLTextRMSNorm( + norm_cls = _layer_norm_cls(recipe) + self.input_layernorm = norm_cls( hidden_size=hidden_size, eps=model_config.pretrained_config.rms_norm_eps, dtype=torch.bfloat16, ) - self.post_attention_layernorm = Qwen3VLTextRMSNorm( + self.post_attention_layernorm = norm_cls( hidden_size=hidden_size, eps=model_config.pretrained_config.rms_norm_eps, dtype=torch.bfloat16, ) - self.mlp = GatedMLP( - hidden_size=hidden_size, - intermediate_size=intermediate_size, - bias=False, - dtype=torch.bfloat16, - config=model_config, - layer_idx=layer_idx, - reduce_output=model_config.mapping.tp_size > 1, - ) + self.mlp = _build_cosmos3_mlp(model_config, recipe, layer_idx) def forward( self, @@ -593,11 +811,12 @@ def _compute_default_rope_parameters( class Qwen3VLTextRotaryEmbedding(nn.Module): def __init__(self, model_config: DiffusionModelConfig): super().__init__() - self.rope_type = model_config.pretrained_config.rope_scaling["rope_type"] + # Edge checkpoints omit rope_type from rope_scaling. + self.rope_type = model_config.pretrained_config.rope_scaling.get("rope_type", "default") self.max_seq_len_cached = model_config.pretrained_config.max_position_embeddings self.original_max_seq_len = model_config.pretrained_config.max_position_embeddings - self.mrope_section = model_config.pretrained_config.rope_scaling["mrope_section"] + self.mrope_section = resolve_rope_axes_dim(model_config.pretrained_config) inv_freq, self.attention_scaling = _compute_default_rope_parameters(model_config) self.register_buffer("inv_freq", inv_freq, persistent=False) @@ -654,7 +873,7 @@ class Cosmos3LanguageModel(nn.Module): computed once and reused across all sampling steps. """ - def __init__(self, model_config: DiffusionModelConfig): + def __init__(self, model_config: DiffusionModelConfig, recipe: Cosmos3ArchRecipe): super().__init__() hidden_size = model_config.pretrained_config.hidden_size num_hidden_layers = model_config.pretrained_config.num_hidden_layers @@ -667,7 +886,10 @@ def __init__(self, model_config: DiffusionModelConfig): ) self.rotary_emb = Qwen3VLTextRotaryEmbedding(model_config) self.layers = nn.ModuleList( - [Cosmos3UndDecoderLayer(model_config, layer_idx=i) for i in range(num_hidden_layers)] + [ + Cosmos3UndDecoderLayer(model_config, layer_idx=i, recipe=recipe) + for i in range(num_hidden_layers) + ] ) def forward( @@ -702,9 +924,14 @@ def forward( class Cosmos3VFMTransformer(BaseDiffusionModel): def __init__(self, model_config: DiffusionModelConfig): super().__init__(model_config) - pretrained_config = model_config.pretrained_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) + self.recipe = resolve_arch_recipe(pretrained_config) self.audio_gen = getattr(pretrained_config, "sound_gen", False) - self.action_gen = getattr(pretrained_config, "action_gen", False) + # Config fact only: the transformer never constructs action modules. + self.has_action_weights = getattr(pretrained_config, "action_gen", False) self.hidden_size = pretrained_config.hidden_size self.num_hidden_layers = pretrained_config.num_hidden_layers @@ -714,8 +941,9 @@ def __init__(self, model_config: DiffusionModelConfig): self.timestep_scale = pretrained_config.timestep_scale self.base_fps = pretrained_config.base_fps - # Comes from VAE. Updated after VAE is loaded. - self.temporal_compression_factor = 4 + # Config-declared (compat default 4); cross-checked against the VAE + # after component loading. + self.temporal_compression_factor = pretrained_config.temporal_compression_factor self.unified_3d_mrope_temporal_modality_margin = ( pretrained_config.unified_3d_mrope_temporal_modality_margin @@ -764,7 +992,7 @@ def __init__(self, model_config: DiffusionModelConfig): "Ring parallelism is not supported for Cosmos3 cross-attention." ) - self.language_model = Cosmos3LanguageModel(model_config) + self.language_model = Cosmos3LanguageModel(model_config, self.recipe) self.vae2llm = nn.Linear(self.patch_latent_dim, self.hidden_size) self.llm2vae = nn.Linear(self.hidden_size, self.patch_latent_dim) @@ -780,12 +1008,12 @@ def __init__(self, model_config: DiffusionModelConfig): self.gen_layers = nn.ModuleList( [ - Cosmos3GenDecoderLayer(model_config, layer_idx=i) + Cosmos3GenDecoderLayer(model_config, layer_idx=i, recipe=self.recipe) for i in range(self.num_hidden_layers) ] ) - self.norm_moe_gen = Qwen3VLTextRMSNorm( + self.norm_moe_gen = _layer_norm_cls(self.recipe)( hidden_size=self.hidden_size, eps=pretrained_config.rms_norm_eps, ) @@ -1161,19 +1389,22 @@ def load_weights(self, weights: dict) -> None: "action_modality_embed", "action_proj_", ) + skipped_keys = [] for key, value in weights.items(): k = key - if k.startswith(skip_prefixes): - continue - - # Normalize a leading "model." prefix up front so every remap below - # matches whether or not the checkpoint namespaces top-level tensors - # (e.g. "model.audio_proj_in.weight") under "model.". + # Normalize a leading "model." prefix up front so every skip and + # remap below matches whether or not the checkpoint namespaces + # top-level tensors (e.g. "model.audio_proj_in.weight") under + # "model.". if k.startswith("model."): k = k[len("model.") :] + if k.startswith(skip_prefixes): + skipped_keys.append(k) + continue + if k.startswith("proj_in."): remapped[k.replace("proj_in.", "vae2llm.", 1)] = value continue @@ -1200,11 +1431,17 @@ def load_weights(self, weights: dict) -> None: remapped[k] = value continue - # embed_tokens and norm → language_model.* - if k.startswith("embed_tokens.") or k.startswith("norm."): + # embed_tokens → language_model.* + if k.startswith("embed_tokens."): remapped[f"language_model.{k}"] = value continue + # Und final norm: normalizes the und hidden state consumed only by + # lm_head; the generation path uses per-layer und K/V exclusively. + if k.startswith("norm."): + skipped_keys.append(k) + continue + # norm_moe_gen stays at top level if k.startswith("norm_moe_gen."): remapped[k] = value @@ -1231,6 +1468,7 @@ def load_weights(self, weights: dict) -> None: "self_attn.to_out.": f"{und_lp}.self_attn.to_out.0.", "self_attn.norm_q.": f"{und_lp}.self_attn.norm_q.", "self_attn.norm_k.": f"{und_lp}.self_attn.norm_k.", + "self_attn.k_norm_und_for_gen.": f"{und_lp}.self_attn.k_norm_und_for_gen.", } # --- GEN attention → gen_layers.{i}.cross_attention.* --- @@ -1282,35 +1520,78 @@ def load_weights(self, weights: dict) -> None: } loader = DynamicLinearWeightLoader(self.model_config, params_map=params_map) + # Coverage is default-fail in both directions: every parameter of every + # constructed module must receive a checkpoint tensor, and every mapped + # checkpoint tensor must land on a parameter. The remap-stage skip list + # above is the only source of intentional omissions. + missing = [] + consumed = set() + for param_name, param in self._parameters.items(): - if param is not None and param_name in remapped: + if param is None: + continue + if param_name in remapped: param.data.copy_(remapped[param_name].to(param.dtype)) + consumed.add(param_name) + else: + missing.append(param_name) + + def _linear_source_prefixes(module: Linear, name: str) -> list: + weights_config = getattr(module, "weights_loading_config", None) + weight_mode = getattr(weights_config, "weight_mode", None) + if weight_mode in (WeightMode.FUSED_QKV_LINEAR, WeightMode.FUSED_GATE_UP_LINEAR): + parent = name.rsplit(".", 1)[0] + for suffix, sources in loader.params_map.items(): + if name == suffix or name.endswith("." + suffix): + return [f"{parent}.{source}." for source in sources] + return [f"{name}."] - loaded_linear = 0 - loaded_other = 0 - skipped_modules = [] for name, module in self.named_modules(): - if len(module._parameters) == 0: + if not name or len(module._parameters) == 0: continue if isinstance(module, Linear): weight_dicts = loader.get_linear_weights(module, name, remapped) - if weight_dicts: + # Fused sources yield one dict each; an empty dict means that + # source is absent, so a partially covered fusion also fails. + if weight_dicts and all(weight_dicts): loader.load_linear_weights(module, name, weight_dicts) - loaded_linear += 1 + prefixes = _linear_source_prefixes(module, name) + consumed.update(k for k in remapped if any(k.startswith(p) for p in prefixes)) else: - skipped_modules.append(f"{name}(Linear)") + missing.append(f"{name}(Linear)") else: module_weights = loader.filter_weights(name, remapped) - if module_weights: - loaded_other += 1 - else: - has_params = any(p is not None for p in module._parameters.values()) - if has_params and name: - skipped_modules.append(f"{name}({type(module).__name__})") for param_name, param in module._parameters.items(): - if param is not None and param_name in module_weights: + if param is None: + continue + if param_name in module_weights: param.data.copy_(module_weights[param_name].to(param.dtype)) + consumed.add(f"{name}.{param_name}") + else: + missing.append(f"{name}.{param_name}") + + if missing: + preview = ", ".join(missing[:10]) + suffix = " ..." if len(missing) > 10 else "" + raise ValueError( + f"Cosmos3 checkpoint is missing weights for {len(missing)} constructed " + f"parameter(s)/module(s): {preview}{suffix}" + ) + unconsumed = [k for k in remapped if k not in consumed] + if unconsumed: + preview = ", ".join(unconsumed[:10]) + suffix = " ..." if len(unconsumed) > 10 else "" + logger.warning( + f"{len(unconsumed)} mapped checkpoint tensor(s) matched no constructed " + f"parameter: {preview}{suffix}" + ) + if skipped_keys: + logger.info( + f"Skipped {len(skipped_keys)} intentionally unused checkpoint tensors " + f"(lm_head / und final norm / action heads): " + f"{', '.join(sorted(set(k.split('.')[0] for k in skipped_keys)))}" + ) def post_load_weights(self) -> None: """Post-load processing: dtype conversion and Linear finalization.""" diff --git a/tensorrt_llm/_torch/visual_gen/pipeline.py b/tensorrt_llm/_torch/visual_gen/pipeline.py index 13cbe77bdf12..08c8acae8b35 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline.py @@ -978,16 +978,22 @@ def _scheduler_step( timestep, scheduler, extra_stream_schedulers, + scheduler_step_kwargs=None, ): """Execute scheduler step for all streams.""" + step_kwargs = scheduler_step_kwargs or {} t_start = time.time() - latents = scheduler.step(noise_pred, timestep, latents, return_dict=False)[0] + latents = scheduler.step(noise_pred, timestep, latents, return_dict=False, **step_kwargs)[0] # Step schedulers for extra streams for name, noise_extra in extra_noise_preds.items(): if name in extra_stream_schedulers: extra_stream_latents[name] = extra_stream_schedulers[name].step( - noise_extra, timestep, extra_stream_latents[name], return_dict=False + noise_extra, + timestep, + extra_stream_latents[name], + return_dict=False, + **step_kwargs, )[0] t_sched = time.time() - t_start @@ -1010,6 +1016,7 @@ def denoise( boundary_timestep: Optional[float] = None, guidance_interval: Optional[Tuple[float, float]] = None, post_step_fn: Optional[Callable] = None, + scheduler_step_kwargs: Optional[Dict[str, Any]] = None, ): """Execute denoising loop with optional CFG parallel and TeaCache support. @@ -1045,6 +1052,8 @@ def denoise( post_step_fn: Optional callable applied to latents after each scheduler step. Signature: post_step_fn(latents) -> latents Use for constraints that must hold throughout denoising. + scheduler_step_kwargs: Extra keyword arguments forwarded to every + scheduler's ``step()`` call. Returns: Single latents if no extra_streams @@ -1171,6 +1180,7 @@ def denoise( t, scheduler, extra_stream_schedulers, + scheduler_step_kwargs=scheduler_step_kwargs, ) if post_step_fn is not None: diff --git a/tensorrt_llm/visual_gen/visual_gen.py b/tensorrt_llm/visual_gen/visual_gen.py index ce83f862326f..9648d26f87ce 100644 --- a/tensorrt_llm/visual_gen/visual_gen.py +++ b/tensorrt_llm/visual_gen/visual_gen.py @@ -283,11 +283,16 @@ def extra_param_specs(self) -> Dict[str, "ExtraParamSchema"]: @property def default_params(self) -> "VisualGenParams": - """Returns a ``VisualGenParams`` with all defaults resolved for the loaded pipeline. + """Returns a ``VisualGenParams`` with the loaded pipeline's defaults. Universal fields (height, width, etc.) are filled from the - pipeline's defaults. All declared ``extra_params`` keys are - included with their defaults (``None`` for params without one). + pipeline's defaults. Pipelines with mode-dependent defaults + (e.g. Cosmos3, where text-to-image and video requests use + different resolutions) leave such fields as ``None``; they are + resolved per request from the output mode, so ``None`` here + means "the mode's default", not "unset". All declared + ``extra_params`` keys are included with their defaults + (``None`` for params without one). Use this to inspect what the model will use, then modify and pass to ``generate()``:: diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_edge_i2v_lpips_golden_video.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_edge_i2v_lpips_golden_video.json new file mode 100644 index 000000000000..00f4e3ececa9 --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_edge_i2v_lpips_golden_video.json @@ -0,0 +1,27 @@ +{ + "source": "diffusers Cosmos3OmniPipeline on diffusers main (reference implementation, not a TRT-LLM self-golden)", + "diffusers_reference": "huggingface/diffusers#14181 'Cosmos3 edge support' + #14246 'Fix Cosmos3 Edge generator K normalization'", + "diffusers_version": "0.40.0.dev0", + "diffusers_commit": "2919c50968389232c527bdab1a3af69cef01ed07", + "scheduler_override": "UniPCMultistepScheduler.from_config(checkpoint config, use_karras_sigmas=False, flow_shift=3.0). With the checkpoint's use_native_flow_schedule=true this reproduces the cosmos-framework PyTorch backend schedule (fm_solvers_unipc @ 117c7d2) to fp32-ulp: timesteps bit-identical, full synthetic step() trajectories agree to <=1.6e-7 rel (see TestNativeFlowSchedule fixtures). Stock diffusers is NOT used as-is because its karras branch swallows the native flow sigmas.", + "prompt_text_matching": "The golden run passed pre-formatted cond AND uncond texts with add_duration_template=False and add_resolution_template=False. Both texts were produced by TRT-LLM's _format_prompt_with_metadata (keep-metadata negative-prompt semantics, matching cosmos-framework's CLI default rather than diffusers' inverse templates), so both stacks tokenize identical sequences in both CFG branches.", + "model": "Cosmos3-Edge", + "seed": 42, + "generator": "torch.Generator(device='cuda').manual_seed(42); initial latents match TRT-LLM's randn_tensor draw bit-for-bit (same shape/dtype/generator semantics)", + "use_system_prompt": false, + "torch_dtype": "bfloat16", + "lpips_net": "alex", + "video": "cosmos3_edge_i2v_lpips_golden_video.mp4", + "prompt": "The orange sphere slowly rises while the camera pans right across the scene", + "conditioning_image": "deterministic 832x480 image drawn by _write_cosmos3_edge_conditioning_image in test_visual_gen.py", + "height": 480, + "width": 832, + "num_frames": 29, + "num_inference_steps": 10, + "guidance_scale": 5.0, + "frame_rate": 24.0, + "lpips_threshold": 0.13, + "measured_lpips_at_creation": 0.0778, + "threshold_rationale": "0.0778 measured cross-stack at 10 steps (I2V accumulates cross-stack drift faster than T2V: 0.1105 at the deployed 50 steps), plus ~0.04 cross-host headroom. The failure signal is far away: a wrong-seed run against this golden measures LPIPS 0.858. The deployed 50-step I2V shape is exercised by test_cosmos3_edge_i2v_example.", + "notes": "Per-step masked-velocity parity vs diffusers is 0.8-1.5 percent rel (noisy frames); diffusers zeroes the conditioned frame's velocity while TRT-LLM masks it in the pipeline - equivalent for the scheduler." +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_edge_t2i_lpips_golden.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_edge_t2i_lpips_golden.json new file mode 100644 index 000000000000..0a5331b8bbc1 --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_edge_t2i_lpips_golden.json @@ -0,0 +1,24 @@ +{ + "source": "diffusers Cosmos3OmniPipeline on diffusers main (reference implementation, not a TRT-LLM self-golden)", + "diffusers_reference": "huggingface/diffusers#14181 'Cosmos3 edge support' + #14246 'Fix Cosmos3 Edge generator K normalization'", + "diffusers_version": "0.40.0.dev0", + "diffusers_commit": "2919c50968389232c527bdab1a3af69cef01ed07", + "scheduler_override": "UniPCMultistepScheduler.from_config(checkpoint config, use_karras_sigmas=False, flow_shift=3.0). With the checkpoint's use_native_flow_schedule=true this reproduces the cosmos-framework PyTorch backend schedule (fm_solvers_unipc @ 117c7d2) to fp32-ulp: timesteps bit-identical, full synthetic step() trajectories agree to <=1.6e-7 rel (see TestNativeFlowSchedule fixtures). Stock diffusers is NOT used as-is because its karras branch swallows the native flow sigmas.", + "prompt_text_matching": "The golden run passed pre-formatted cond AND uncond texts with add_duration_template=False and add_resolution_template=False. Both texts were produced by TRT-LLM's _format_prompt_with_metadata (keep-metadata negative-prompt semantics, matching cosmos-framework's CLI default rather than diffusers' inverse templates), so both stacks tokenize identical sequences in both CFG branches.", + "model": "Cosmos3-Edge", + "seed": 42, + "generator": "torch.Generator(device='cuda').manual_seed(42); initial latents match TRT-LLM's randn_tensor draw bit-for-bit (same shape/dtype/generator semantics)", + "use_system_prompt": false, + "torch_dtype": "bfloat16", + "lpips_net": "alex", + "image": "cosmos3_edge_t2i_lpips_golden.png", + "prompt": "A ceramic teapot pouring steaming tea into a cup, morning window light", + "height": 640, + "width": 640, + "num_frames": 1, + "num_inference_steps": 50, + "guidance_scale": 4.0, + "lpips_threshold": 0.05, + "measured_lpips_at_creation": 0.0056, + "threshold_rationale": "0.0056 measured cross-stack on B200; 0.05 matches the FLUX/QwenImage image-gate convention." +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_edge_t2v_lpips_golden_video.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_edge_t2v_lpips_golden_video.json new file mode 100644 index 000000000000..904374e075ef --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_edge_t2v_lpips_golden_video.json @@ -0,0 +1,25 @@ +{ + "source": "diffusers Cosmos3OmniPipeline on diffusers main (reference implementation, not a TRT-LLM self-golden)", + "diffusers_reference": "huggingface/diffusers#14181 'Cosmos3 edge support' + #14246 'Fix Cosmos3 Edge generator K normalization'", + "diffusers_version": "0.40.0.dev0", + "diffusers_commit": "2919c50968389232c527bdab1a3af69cef01ed07", + "scheduler_override": "UniPCMultistepScheduler.from_config(checkpoint config, use_karras_sigmas=False, flow_shift=3.0). With the checkpoint's use_native_flow_schedule=true this reproduces the cosmos-framework PyTorch backend schedule (fm_solvers_unipc @ 117c7d2) to fp32-ulp: timesteps bit-identical, full synthetic step() trajectories agree to <=1.6e-7 rel (see TestNativeFlowSchedule fixtures). Stock diffusers is NOT used as-is because its karras branch swallows the native flow sigmas.", + "prompt_text_matching": "The golden run passed pre-formatted cond AND uncond texts with add_duration_template=False and add_resolution_template=False. Both texts were produced by TRT-LLM's _format_prompt_with_metadata (keep-metadata negative-prompt semantics, matching cosmos-framework's CLI default rather than diffusers' inverse templates), so both stacks tokenize identical sequences in both CFG branches.", + "model": "Cosmos3-Edge", + "seed": 42, + "generator": "torch.Generator(device='cuda').manual_seed(42); initial latents match TRT-LLM's randn_tensor draw bit-for-bit (same shape/dtype/generator semantics)", + "use_system_prompt": false, + "torch_dtype": "bfloat16", + "lpips_net": "alex", + "video": "cosmos3_edge_t2v_lpips_golden_video.mp4", + "prompt": "A red ball rolls across a wooden floor, casting a soft shadow.", + "height": 480, + "width": 832, + "num_frames": 29, + "num_inference_steps": 50, + "guidance_scale": 5.0, + "frame_rate": 24.0, + "lpips_threshold": 0.1, + "measured_lpips_at_creation": 0.0447, + "threshold_rationale": "0.0447 measured cross-stack (TRT-LLM VANILLA attention vs diffusers main) on B200 with matched noise and matched CFG texts, plus headroom for the ~0.04 cross-host kernel drift documented in _preserve_lpips_candidate_on_failure." +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_i2v_4step_lpips_golden_video.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_i2v_4step_lpips_golden_video.json new file mode 100644 index 000000000000..1ebbc1f8f9c8 --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_i2v_4step_lpips_golden_video.json @@ -0,0 +1,27 @@ +{ + "video": "cosmos3_i2v_4step_lpips_golden_video.mp4", + "model": "Cosmos3-Super-Image2Video-4Step", + "source": "diffusers Cosmos3DistilledModularPipeline (reference implementation, not a TRT-LLM self-golden)", + "diffusers_reference": "huggingface/diffusers#14177 'Cosmos3 Distilled support' (merged 2026-07-17)", + "diffusers_version": "0.40.0.dev0", + "diffusers_commit": "6b5199f81dc0044fe0417ebe96bbad6c59c19874", + "reference_rng_patch": "Cosmos3DistilledVisionLoopSchedulerStep patched to declare InputParam('generator') and pass generator=block_state.generator into scheduler.step. The merged PR draws the per-step SDE noise from the global torch RNG, so unpatched same-seed runs are different samples (measured LPIPS 0.144 vs 0.056 patched); the patch makes the reference trajectory reproducible and comparable.", + "reference_pipeline_construction": "The checkpoint's own modular_model_index.json is stale (base Cosmos3OmniBlocks, refs/pr/1, no distilled_sigmas). The golden run loaded a corrected index selecting Cosmos3DistilledModularPipeline / Cosmos3DistilledBlocks with is_distilled=true, distilled_sigmas=[1.0, 0.9375, 0.8333333333333334, 0.625], and component specs pointing at the local checkpoint subfolders.", + "prompt": "The orange sphere slowly rises while the camera pans right across the scene", + "conditioning_image": "deterministic 1280x720 image drawn by _write_cosmos3_i2v_conditioning_image in test_visual_gen.py (gradient sky, orange circle, green rectangle, yellow triangle)", + "height": 720, + "width": 1280, + "num_frames": 29, + "num_inference_steps": 4, + "guidance_scale": 1.0, + "seed": 42, + "global_rng_seed": 42, + "frame_rate": 24.0, + "use_system_prompt": true, + "torch_dtype": "bfloat16", + "lpips_net": "alex", + "lpips_threshold": 0.1, + "measured_lpips_at_creation": 0.0563, + "threshold_rationale": "0.0563 measured cross-stack (TRT-LLM VANILLA attention vs diffusers) on B200 with matched noise trajectories, plus headroom for the ~0.04 cross-host kernel drift documented in _preserve_lpips_candidate_on_failure.", + "notes": "num_inference_steps and guidance_scale are fixed by the distilled checkpoint (scheduler fixed_step_sampler_config.t_list; CFG baked into the weights). use_system_prompt=true matches the checkpoint's default_use_system_prompt declaration in model_index.json." +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zip b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zip index b00403606a62..22aebc93c53f 100644 --- a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zip +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zip @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:6d38286a6ee60db331a89bb3563fcc0896309b6e21228d4867ab33990c5b47ed -size 14386487 +oid sha256:3c10ed17e353d9a8f02d744a0e475bd8b5c60614c72c296ce1cd3e13295cf216 +size 15364529 diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen.py b/tests/integration/defs/examples/visual_gen/test_visual_gen.py index 83191a0d4837..7c2cdec79ee4 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen.py @@ -122,6 +122,20 @@ COSMOS3_LPIPS_FRAME_RATE = 24.0 COSMOS3_LPIPS_THRESHOLD = 0.05 +COSMOS3_I2V_4STEP_MODEL_SUBPATH = "Cosmos3-Super-Image2Video-4Step" +COSMOS3_I2V_4STEP_LPIPS_PROMPT = ( + "The orange sphere slowly rises while the camera pans right across the scene" +) +COSMOS3_I2V_4STEP_LPIPS_NUM_FRAMES = 29 +# Fixed by the distilled checkpoint (scheduler t_list / CFG baked into weights). +COSMOS3_I2V_4STEP_LPIPS_NUM_INFERENCE_STEPS = 4 +COSMOS3_I2V_4STEP_LPIPS_GUIDANCE_SCALE = 1.0 +# Golden is diffusers-produced (cross-stack), not a TRT-LLM self-golden: +# 0.0563 measured at creation + headroom for ~0.04 cross-host kernel drift +# (see _preserve_lpips_candidate_on_failure). Provenance: +# golden/visual_gen_lpips/cosmos3_i2v_4step_lpips_golden_video.json. +COSMOS3_I2V_4STEP_LPIPS_THRESHOLD = 0.10 + # LTX-2 configuration LTX2_MODEL_CHECKPOINT_PATH = "LTX-2/ltx-2-19b-dev.safetensors" LTX2_TEXT_ENCODER_SUBPATH = "gemma-3-12b-it" @@ -1852,3 +1866,441 @@ def test_cosmos3_example(_visual_gen_deps, llm_root, llm_venv): env={"TRTLLM_DISABLE_COSMOS3_GUARDRAILS": "1"}, ) assert os.path.isfile(output_path), f"Example did not produce output at {output_path}" + + +def test_cosmos3_t2i_4step_example(_visual_gen_deps, llm_root, llm_venv): + """Run the distilled T2I checkpoint through the recommended invocation. + + Validates the documented deployment for ``Cosmos3-Super-Text2Image-4Step``: + the example script with ``configs/cosmos3-t2i-1gpu.yaml`` (T2I warmup + shapes) and ``--output_type image``. Steps/guidance come from the + checkpoint's fixed distilled schedule; the run must produce an image. + """ + model_path = _lpips_model_path("Cosmos3-Super-Text2Image-4Step") + _skip_if_missing(model_path, "Cosmos3-Super-Text2Image-4Step checkpoint", is_dir=True) + + out_dir = os.path.join( + llm_venv.get_working_directory(), "visual_gen_output", "cosmos3_t2i_4step_example" + ) + os.makedirs(out_dir, exist_ok=True) + output_path = os.path.join(out_dir, "cosmos3_t2i_4step_output.png") + if os.path.exists(output_path): + os.remove(output_path) + + script_path = os.path.join( + llm_root, "examples", "visual_gen", "models", "cosmos3", "cosmos3.py" + ) + config_path = os.path.join( + llm_root, "examples", "visual_gen", "configs", "cosmos3-t2i-1gpu.yaml" + ) + assert os.path.isfile(script_path), f"Example script not found: {script_path}" + assert os.path.isfile(config_path), f"Config not found: {config_path}" + + venv_check_call( + llm_venv, + [ + script_path, + "--model", + model_path, + "--visual_gen_args", + config_path, + "--prompt", + "A ceramic teapot pouring steaming tea into a cup, morning window light", + "--output_type", + "image", + "--output_path", + output_path, + ], + env={"TRTLLM_DISABLE_COSMOS3_GUARDRAILS": "1"}, + ) + assert os.path.isfile(output_path), f"Example did not produce output at {output_path}" + assert os.path.getsize(output_path) > 0, f"Example produced an empty image at {output_path}" + + +def _write_cosmos3_i2v_conditioning_image(path): + """Deterministic 1280x720 conditioning image for the I2V smoke test. + + Gradient sky plus simple shapes, so I2V has real structure to animate + without shipping an asset file. + """ + from PIL import Image, ImageDraw + + image = Image.new("RGB", (1280, 720)) + draw = ImageDraw.Draw(image) + for y in range(720): + draw.line([(0, y), (1280, y)], fill=(30, 60 + y // 8, 140)) + draw.ellipse([480, 200, 800, 520], fill=(230, 120, 40), outline=(255, 255, 255), width=6) + draw.rectangle([100, 500, 400, 680], fill=(40, 160, 90)) + draw.polygon([(1000, 600), (1120, 380), (1240, 600)], fill=(200, 200, 60)) + image.save(path) + + +def test_cosmos3_i2v_4step_example(_visual_gen_deps, llm_root, llm_venv): + """Run the distilled I2V checkpoint through the recommended invocation. + + Validates the documented deployment for ``Cosmos3-Super-Image2Video-4Step``: + the example script with a conditioning image and no config override (the + omni defaults — 720p x 189 frames — are the deployed shape). Steps, + guidance, and the system-prompt default come from the checkpoint; the run + must produce a video. + """ + model_path = _lpips_model_path("Cosmos3-Super-Image2Video-4Step") + _skip_if_missing(model_path, "Cosmos3-Super-Image2Video-4Step checkpoint", is_dir=True) + + out_dir = os.path.join( + llm_venv.get_working_directory(), "visual_gen_output", "cosmos3_i2v_4step_example" + ) + os.makedirs(out_dir, exist_ok=True) + image_path = os.path.join(out_dir, "conditioning.png") + _write_cosmos3_i2v_conditioning_image(image_path) + output_path = os.path.join(out_dir, "cosmos3_i2v_4step_output.mp4") + if os.path.exists(output_path): + os.remove(output_path) + + script_path = os.path.join( + llm_root, "examples", "visual_gen", "models", "cosmos3", "cosmos3.py" + ) + assert os.path.isfile(script_path), f"Example script not found: {script_path}" + + venv_check_call( + llm_venv, + [ + script_path, + "--model", + model_path, + "--prompt", + "The orange sphere slowly rises while the camera pans right across the scene", + "--image_path", + image_path, + "--output_path", + output_path, + ], + env={"TRTLLM_DISABLE_COSMOS3_GUARDRAILS": "1"}, + ) + assert os.path.isfile(output_path), f"Example did not produce output at {output_path}" + assert os.path.getsize(output_path) > 0, f"Example produced an empty video at {output_path}" + + +def _write_cosmos3_edge_conditioning_image(path): + """Deterministic 832x480 conditioning image (Edge's native 480p 16:9).""" + from PIL import Image, ImageDraw + + image = Image.new("RGB", (832, 480)) + draw = ImageDraw.Draw(image) + for y in range(480): + draw.line([(0, y), (832, y)], fill=(30, 60 + y // 6, 140)) + draw.ellipse([320, 130, 520, 330], fill=(230, 120, 40), outline=(255, 255, 255), width=4) + draw.rectangle([60, 330, 260, 450], fill=(40, 160, 90)) + image.save(path) + + +def test_cosmos3_edge_i2v_example(_visual_gen_deps, llm_root, llm_venv): + """Run the Edge checkpoint through the recommended invocation. + + Validates the documented deployment for ``Cosmos3-Edge``: the example + script with a conditioning image and no config override (the Edge + defaults — 480p x 121 frames, 50 UniPC steps on the native flow schedule, + guidance 5.0, shift 3.0 — are the deployed shape). The run must produce a + video. + """ + model_path = _lpips_model_path("Cosmos3-Edge") + _skip_if_missing(model_path, "Cosmos3-Edge checkpoint", is_dir=True) + + out_dir = os.path.join( + llm_venv.get_working_directory(), "visual_gen_output", "cosmos3_edge_i2v_example" + ) + os.makedirs(out_dir, exist_ok=True) + image_path = os.path.join(out_dir, "conditioning.png") + _write_cosmos3_edge_conditioning_image(image_path) + output_path = os.path.join(out_dir, "cosmos3_edge_i2v_output.mp4") + if os.path.exists(output_path): + os.remove(output_path) + + script_path = os.path.join( + llm_root, "examples", "visual_gen", "models", "cosmos3", "cosmos3.py" + ) + assert os.path.isfile(script_path), f"Example script not found: {script_path}" + + venv_check_call( + llm_venv, + [ + script_path, + "--model", + model_path, + "--prompt", + "The orange sphere slowly rises while the camera pans right across the scene", + "--image_path", + image_path, + "--output_path", + output_path, + ], + env={"TRTLLM_DISABLE_COSMOS3_GUARDRAILS": "1"}, + ) + assert os.path.isfile(output_path), f"Example did not produce output at {output_path}" + assert os.path.getsize(output_path) > 0, f"Example produced an empty video at {output_path}" + + +# Edge LPIPS gates compare against diffusers-main reference goldens with the +# scheduler patched to the cosmos-framework native flow schedule; full +# provenance in golden/visual_gen_lpips/cosmos3_edge_*.json. The I2V gate runs +# 10 steps (cross-stack drift accumulates per step; the deployed 50-step shape +# is covered by test_cosmos3_edge_i2v_example). +COSMOS3_EDGE_LPIPS_SEED = 42 +COSMOS3_EDGE_LPIPS_FRAME_RATE = 24.0 +COSMOS3_EDGE_LPIPS_NUM_FRAMES = 29 +COSMOS3_EDGE_T2V_LPIPS_PROMPT = "A red ball rolls across a wooden floor, casting a soft shadow." +COSMOS3_EDGE_T2V_LPIPS_STEPS = 50 +COSMOS3_EDGE_T2V_LPIPS_THRESHOLD = 0.1 +COSMOS3_EDGE_I2V_LPIPS_PROMPT = ( + "The orange sphere slowly rises while the camera pans right across the scene" +) +COSMOS3_EDGE_I2V_LPIPS_STEPS = 10 +COSMOS3_EDGE_I2V_LPIPS_THRESHOLD = 0.13 +COSMOS3_EDGE_T2I_LPIPS_PROMPT = ( + "A ceramic teapot pouring steaming tea into a cup, morning window light" +) +COSMOS3_EDGE_T2I_LPIPS_STEPS = 50 +COSMOS3_EDGE_T2I_LPIPS_THRESHOLD = 0.05 + + +def _run_cosmos3_edge_lpips_pipeline(**forward_kwargs): + """Run the Cosmos3-Edge pipeline and return the PipelineOutput. + + VANILLA attention, compile-off; guardrails disabled for the run. + """ + guardrails_env_key = "TRTLLM_DISABLE_COSMOS3_GUARDRAILS" + previous_guardrails_env = os.environ.get(guardrails_env_key) + os.environ[guardrails_env_key] = "1" + try: + from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader + from tensorrt_llm.visual_gen.args import ( + AttentionConfig, + CompilationConfig, + TorchCompileConfig, + VisualGenArgs, + ) + + model_path = _lpips_model_path("Cosmos3-Edge") + _skip_if_missing(model_path, "Cosmos3-Edge checkpoint", is_dir=True) + _disable_inductor_compile_worker_quiesce() + args = VisualGenArgs( + model=model_path, + compilation_config=CompilationConfig(skip_warmup=True), + torch_compile_config=TorchCompileConfig(enable=False), + attention_config=AttentionConfig(backend="VANILLA"), + ) + pipeline = PipelineLoader(args).load(skip_warmup=True) + try: + with torch.no_grad(): + result = pipeline.forward( + seed=COSMOS3_EDGE_LPIPS_SEED, + use_guardrails=False, + **forward_kwargs, + ) + if result is not None: + if result.video is not None: + result.video = result.video.detach().cpu() + if result.image is not None: + result.image = result.image.detach().cpu() + return result + finally: + del pipeline + _cleanup_cuda() + finally: + if previous_guardrails_env is None: + os.environ.pop(guardrails_env_key, None) + else: + os.environ[guardrails_env_key] = previous_guardrails_env + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_cosmos3_edge_t2v_lpips_against_golden(_visual_gen_deps, tmp_path): + generated_path = tmp_path / "cosmos3_edge_t2v_generated.mp4" + golden_path = _golden_media_path( + tmp_path, "cosmos3_edge_t2v_lpips_golden_video.mp4", "Cosmos3-Edge T2V LPIPS golden video" + ) + result = _run_cosmos3_edge_lpips_pipeline( + prompt=COSMOS3_EDGE_T2V_LPIPS_PROMPT, + height=480, + width=832, + num_frames=COSMOS3_EDGE_LPIPS_NUM_FRAMES, + num_inference_steps=COSMOS3_EDGE_T2V_LPIPS_STEPS, + guidance_scale=5.0, + frame_rate=COSMOS3_EDGE_LPIPS_FRAME_RATE, + ) + assert result is not None and result.video is not None, "Edge T2V produced no video" + _save_lpips_video_mp4(result.video, generated_path, frame_rate=COSMOS3_EDGE_LPIPS_FRAME_RATE) + score = _run_lpips_eval( + tmp_path, + "cosmos3_edge_t2v", + "video", + COSMOS3_EDGE_T2V_LPIPS_PROMPT, + golden_path, + generated_path, + ) + _assert_lpips_below_threshold(score, COSMOS3_EDGE_T2V_LPIPS_THRESHOLD) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_cosmos3_edge_i2v_lpips_against_golden(_visual_gen_deps, tmp_path): + generated_path = tmp_path / "cosmos3_edge_i2v_generated.mp4" + golden_path = _golden_media_path( + tmp_path, "cosmos3_edge_i2v_lpips_golden_video.mp4", "Cosmos3-Edge I2V LPIPS golden video" + ) + image_path = tmp_path / "cosmos3_edge_i2v_conditioning.png" + _write_cosmos3_edge_conditioning_image(str(image_path)) + result = _run_cosmos3_edge_lpips_pipeline( + prompt=COSMOS3_EDGE_I2V_LPIPS_PROMPT, + image=str(image_path), + height=480, + width=832, + num_frames=COSMOS3_EDGE_LPIPS_NUM_FRAMES, + num_inference_steps=COSMOS3_EDGE_I2V_LPIPS_STEPS, + guidance_scale=5.0, + frame_rate=COSMOS3_EDGE_LPIPS_FRAME_RATE, + ) + assert result is not None and result.video is not None, "Edge I2V produced no video" + _save_lpips_video_mp4(result.video, generated_path, frame_rate=COSMOS3_EDGE_LPIPS_FRAME_RATE) + score = _run_lpips_eval( + tmp_path, + "cosmos3_edge_i2v", + "video", + COSMOS3_EDGE_I2V_LPIPS_PROMPT, + golden_path, + generated_path, + ) + _assert_lpips_below_threshold(score, COSMOS3_EDGE_I2V_LPIPS_THRESHOLD) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_cosmos3_edge_t2i_lpips_against_golden(tmp_path): + from tensorrt_llm.media.encoding import save_image + + generated_path = tmp_path / "cosmos3_edge_t2i_generated.png" + golden_path = _golden_media_path( + tmp_path, "cosmos3_edge_t2i_lpips_golden.png", "Cosmos3-Edge T2I LPIPS golden image" + ) + result = _run_cosmos3_edge_lpips_pipeline( + prompt=COSMOS3_EDGE_T2I_LPIPS_PROMPT, + height=640, + width=640, + num_inference_steps=COSMOS3_EDGE_T2I_LPIPS_STEPS, + guidance_scale=4.0, + output_type="image", + ) + assert result is not None and result.image is not None, "Edge T2I produced no image" + save_image(result.image[0], str(generated_path)) + score = _run_lpips_eval( + tmp_path, + "cosmos3_edge_t2i", + "image", + COSMOS3_EDGE_T2I_LPIPS_PROMPT, + golden_path, + generated_path, + ) + _assert_lpips_below_threshold(score, COSMOS3_EDGE_T2I_LPIPS_THRESHOLD) + + +def _run_cosmos3_i2v_4step_lpips_pipeline(image_path): + """Run the distilled I2V pipeline on the deterministic conditioning image. + + VANILLA attention, compile-off. Returns the generated video tensor + ``(B, T, H, W, C)``, or ``None`` if generation produced no video. + """ + # Cosmos3 re-reads the guardrail flag in __init__; set it before the pipeline loads. + guardrails_env_key = "TRTLLM_DISABLE_COSMOS3_GUARDRAILS" + previous_guardrails_env = os.environ.get(guardrails_env_key) + os.environ[guardrails_env_key] = "1" + try: + from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader + from tensorrt_llm.visual_gen.args import ( + AttentionConfig, + CompilationConfig, + TorchCompileConfig, + VisualGenArgs, + ) + + model_path = _lpips_model_path(COSMOS3_I2V_4STEP_MODEL_SUBPATH) + _skip_if_missing(model_path, "Cosmos3-Super-Image2Video-4Step checkpoint", is_dir=True) + _disable_inductor_compile_worker_quiesce() + args = VisualGenArgs( + model=model_path, + compilation_config=CompilationConfig(skip_warmup=True), + torch_compile_config=TorchCompileConfig(enable=False), + attention_config=AttentionConfig(backend="VANILLA"), + ) + pipeline = PipelineLoader(args).load(skip_warmup=True) + try: + with torch.no_grad(): + result = pipeline.forward( + prompt=COSMOS3_I2V_4STEP_LPIPS_PROMPT, + seed=COSMOS3_LPIPS_SEED, + image=image_path, + height=COSMOS3_LPIPS_HEIGHT, + width=COSMOS3_LPIPS_WIDTH, + num_frames=COSMOS3_I2V_4STEP_LPIPS_NUM_FRAMES, + # Direct forward() calls must pass checkpoint-valid sampling + # values (the signature defaults are the base-checkpoint + # video table, which a distilled checkpoint rejects). + num_inference_steps=COSMOS3_I2V_4STEP_LPIPS_NUM_INFERENCE_STEPS, + guidance_scale=COSMOS3_I2V_4STEP_LPIPS_GUIDANCE_SCALE, + frame_rate=COSMOS3_LPIPS_FRAME_RATE, + # The checkpoint declares default_use_system_prompt=true and + # the golden was generated with it; forward()'s signature + # default is the historical False, so pass it explicitly. + use_system_prompt=True, + use_guardrails=False, + ) + if result is None or result.video is None: + return None + return result.video.detach().cpu() + finally: + del pipeline + _cleanup_cuda() + finally: + if previous_guardrails_env is None: + os.environ.pop(guardrails_env_key, None) + else: + os.environ[guardrails_env_key] = previous_guardrails_env + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_cosmos3_i2v_4step_lpips_against_golden(_visual_gen_deps, request, tmp_path): + """Quality gate for the distilled I2V checkpoint against a diffusers golden. + + Unlike the self-goldens of the other models, the golden video here was + produced by the reference implementation (diffusers modular pipeline, + PR #14177, with its per-step SDE noise made generator-seeded) — so this + gate checks the denoising trajectory against the reference, not just + regression against a past TRT-LLM run. Full provenance: + ``golden/visual_gen_lpips/cosmos3_i2v_4step_lpips_golden_video.json``. + """ + image_path = str(tmp_path / "cosmos3_i2v_4step_conditioning.png") + _write_cosmos3_i2v_conditioning_image(image_path) + generated_path = tmp_path / "cosmos3_i2v_4step_generated.mp4" + golden_path = _golden_media_path( + tmp_path, + "cosmos3_i2v_4step_lpips_golden_video.mp4", + "Cosmos3 I2V-4Step LPIPS golden video", + ) + + video = _run_cosmos3_i2v_4step_lpips_pipeline(image_path) + assert video is not None, "Cosmos3 I2V-4Step LPIPS run produced no video" + _save_lpips_video_mp4(video, generated_path, frame_rate=COSMOS3_LPIPS_FRAME_RATE) + + score = _run_lpips_eval( + tmp_path, + "cosmos3_i2v_4step", + "video", + COSMOS3_I2V_4STEP_LPIPS_PROMPT, + golden_path, + generated_path, + ) + _preserve_lpips_candidate_on_failure( + request, + score, + COSMOS3_I2V_4STEP_LPIPS_THRESHOLD, + generated_path, + "cosmos3_i2v_4step_lpips_golden_video.mp4", + ) + _assert_lpips_below_threshold(score, COSMOS3_I2V_4STEP_LPIPS_THRESHOLD) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 464c3e553f91..ef7a0bffd08e 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -259,12 +259,16 @@ l0_b200: - unittest/_torch/visual_gen/test_wan_transformer.py - unittest/_torch/visual_gen/test_cosmos3_transformer.py - unittest/_torch/visual_gen/test_cosmos3_pipeline.py + - unittest/_torch/visual_gen/test_cosmos3_edge.py - examples/visual_gen/test_visual_gen.py::test_wan_t2v_example - examples/visual_gen/test_visual_gen.py::test_flux1_example - examples/visual_gen/test_visual_gen.py::test_flux2_example - examples/visual_gen/test_visual_gen.py::test_ltx2_example - examples/visual_gen/test_visual_gen.py::test_wan_i2v_example - examples/visual_gen/test_visual_gen.py::test_cosmos3_example + - examples/visual_gen/test_visual_gen.py::test_cosmos3_t2i_4step_example TIMEOUT (30) + - examples/visual_gen/test_visual_gen.py::test_cosmos3_i2v_4step_example TIMEOUT (45) + - examples/visual_gen/test_visual_gen.py::test_cosmos3_edge_i2v_example TIMEOUT (30) - examples/visual_gen/test_visual_gen.py::test_qwen_image_example # - examples/visual_gen/test_visual_gen.py # ------------- Host perf module regression tests (6 representative scenarios) --------------- @@ -370,6 +374,10 @@ l0_b200: - examples/visual_gen/test_visual_gen.py::test_qwenimage_cuda_graph_lpips_against_golden TIMEOUT (10) - examples/visual_gen/test_visual_gen.py::test_cosmos3_nano_t2i_lpips_against_golden TIMEOUT (10) - examples/visual_gen/test_visual_gen.py::test_cosmos3_nano_t2v_lpips_against_golden TIMEOUT (15) + - examples/visual_gen/test_visual_gen.py::test_cosmos3_i2v_4step_lpips_against_golden TIMEOUT (20) + - examples/visual_gen/test_visual_gen.py::test_cosmos3_edge_t2v_lpips_against_golden TIMEOUT (20) + - examples/visual_gen/test_visual_gen.py::test_cosmos3_edge_i2v_lpips_against_golden TIMEOUT (15) + - examples/visual_gen/test_visual_gen.py::test_cosmos3_edge_t2i_lpips_against_golden TIMEOUT (15) - visual_gen/test_visual_gen_benchmark.py::test_offline_benchmark - visual_gen/test_visual_gen_benchmark.py::test_online_benchmark[openai-videos] # ---- moved to post-merge (MoE CI optimization) ---- diff --git a/tests/unittest/_torch/visual_gen/conftest.py b/tests/unittest/_torch/visual_gen/conftest.py new file mode 100644 index 000000000000..7ae74e66ff40 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/conftest.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared pytest configuration for the VisualGen unit tests.""" + +import os +from collections.abc import Iterator + +import pytest + +# Must be set before any test module imports tensorrt_llm (the MPI-vs-Ray +# orchestration choice is made at import time): these CPU/GPU unit tests never +# use MPI, and initializing a process-global MPI session risks +# fork-after-MPI_Init flakes under forked pytest workers. conftest is imported +# before the test modules, which makes this the single session-wide owner of +# the variable; ``setdefault`` keeps an explicit caller override intact. +os.environ.setdefault("TLLM_DISABLE_MPI", "1") + + +@pytest.fixture(scope="module") +def disable_cosmos3_guardrails() -> Iterator[None]: + """Disable Cosmos3 guardrails for the requesting module, leak-free. + + Patches both the environment variable (re-read by + ``load_standard_components`` on every call) and the pipeline module's + derived global (assigned by that same function), so teardown restores + both. Opt in per module via ``pytest.mark.usefixtures``. + """ + import tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 as pipe_mod + + patcher = pytest.MonkeyPatch() + patcher.setenv("TRTLLM_DISABLE_COSMOS3_GUARDRAILS", "1") + patcher.setattr(pipe_mod, "TRTLLM_DISABLE_COSMOS3_GUARDRAILS", True) + yield + patcher.undo() diff --git a/tests/unittest/_torch/visual_gen/cosmos3_edge_diffusers_parity.py b/tests/unittest/_torch/visual_gen/cosmos3_edge_diffusers_parity.py new file mode 100644 index 000000000000..3a6b1bf0ef7e --- /dev/null +++ b/tests/unittest/_torch/visual_gen/cosmos3_edge_diffusers_parity.py @@ -0,0 +1,131 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-step Edge transformer parity: TRT-LLM vs diffusers main. + +Not a pytest module — run as a subprocess by test_cosmos3_edge.py, because +diffusers main (>= 0.40, first with the Edge classes) cannot be imported into +a process that already imported the pinned diffusers. +Usage: DIFFUSERS_MAIN_PATH=/path/to/diffusers python cosmos3_edge_diffusers_parity.py + +Runs the diffusers Cosmos3OmniPipeline for 2 steps (guidance off, fixed +latents), captures each transformer call's (latent, timestep, velocity), +then replays the same inputs through the TRT-LLM Edge transformer and +compares velocities. +""" + +import os +import sys + +sys.path.insert(0, os.path.join(os.environ["DIFFUSERS_MAIN_PATH"], "src")) +os.environ.setdefault("TRTLLM_DISABLE_COSMOS3_GUARDRAILS", "1") + +CKPT = sys.argv[1] +DEV = "cuda" +H, W, FRAMES, STEPS = 192, 320, 9, 2 +PROMPT = "A red cube sits on a wooden table." + + +def main(): + import diffusers + import torch + + expected_root = os.path.realpath(os.environ["DIFFUSERS_MAIN_PATH"]) + assert os.path.realpath(diffusers.__file__).startswith(expected_root), diffusers.__file__ + print("diffusers:", diffusers.__version__) + from diffusers import Cosmos3OmniPipeline + + pipe = Cosmos3OmniPipeline.from_pretrained(CKPT, torch_dtype=torch.bfloat16).to(DEV) + + records = [] + orig_forward = pipe.transformer.forward + + def spy(**kwargs): + out = orig_forward(**kwargs) + velocity = out[0][0] if isinstance(out, tuple) else out.sample[0] + records.append( + { + "latent": kwargs["vision_tokens"][0].detach().clone(), + "t": kwargs["vision_timesteps"].detach().reshape(-1)[0].clone(), + "vel": velocity.detach().clone(), + } + ) + return out + + pipe.transformer.forward = spy + + latent_shape = (1, 48, (FRAMES - 1) // 4 + 1, H // 16, W // 16) + init_latents = torch.randn( + latent_shape, generator=torch.Generator().manual_seed(0), dtype=torch.float32 + ).to(DEV, torch.bfloat16) + + cond_ids, _ = pipe.tokenize_prompt( + PROMPT, negative_prompt=None, num_frames=FRAMES, height=H, width=W, fps=24.0 + ) + + pipe( + prompt=PROMPT, + num_frames=FRAMES, + height=H, + width=W, + fps=24.0, + num_inference_steps=STEPS, + guidance_scale=1.0, + generator=torch.Generator().manual_seed(0), + latents=init_latents.clone(), + output_type="latent", + ) + print(f"captured {len(records)} transformer calls") + + del pipe.transformer + del pipe + torch.cuda.empty_cache() + + from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineComponent, PipelineLoader + from tensorrt_llm.visual_gen.args import TorchCompileConfig, VisualGenArgs + + args = VisualGenArgs(model=CKPT, torch_compile_config=TorchCompileConfig(enable=False)) + trt_pipe = PipelineLoader(args).load( + skip_warmup=True, + skip_components=[ + PipelineComponent.VAE, + PipelineComponent.SCHEDULER, + PipelineComponent.TOKENIZER, + PipelineComponent.SOUND_TOKENIZER, + ], + ) + transformer = trt_pipe.transformer + + text_ids = torch.tensor([cond_ids], dtype=torch.long, device=DEV) + text_mask = torch.ones_like(text_ids) + latent_t, latent_h, latent_w = latent_shape[2:] + + transformer.reset_cache() + for step, rec in enumerate(records): + latent = rec["latent"].to(DEV, torch.bfloat16) + if latent.dim() == 4: + latent = latent.unsqueeze(0) + t = rec["t"].float().reshape(1).to(DEV) + with torch.inference_mode(): + out = transformer( + hidden_states=latent, + timestep=t / 1000.0, + raw_timestep=t, + text_ids=text_ids, + text_mask=text_mask, + video_shape=(latent_t, latent_h, latent_w), + fps=24.0, + ) + ours = out.video[0].float().cpu() + ref = rec["vel"].float().cpu() + if ref.dim() == 5: + ref = ref[0] + diff = (ours - ref).abs() + rel = diff.max() / ref.abs().max() + print( + f"step {step}: t={t.item():7.2f} max|d|={diff.max():.5f} " + f"mean|d|={diff.mean():.6f} rel={rel:.5f} |ref|max={ref.abs().max():.3f}" + ) + + +if __name__ == "__main__": + main() diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_cosmos3_transformer_parallel.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_cosmos3_transformer_parallel.py index e5b770f5dfa8..b7dc6691d06a 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_cosmos3_transformer_parallel.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_cosmos3_transformer_parallel.py @@ -78,7 +78,7 @@ num_attention_heads=8, num_key_value_heads=4, head_dim=64, - rope_scaling={"rope_type": "default", "mrope_section": [16, 12, 12]}, + rope_scaling={"rope_type": "default", "mrope_section": [12, 10, 10]}, rms_norm_eps=1e-6, vocab_size=1024, rope_theta=1_000_000.0, @@ -91,13 +91,14 @@ # attn2d needs an LSE-capable backend (FA4); FA4's CUTE kernels run with head_dim=128. # Same architecture as _COSMOS3_TEST_CONFIG otherwise (heads still divisible by -# Ulysses=2 for the attn2d+ulysses case). mrope_section is unchanged: the interleave -# slices clip to head_dim//2 (=64 here), so [16, 12, 12] stays valid. +# Ulysses=2 for the attn2d+ulysses case). mrope_section must sum to head_dim//2 +# (=64 here), so this config uses the real checkpoint sections. _COSMOS3_FA4_CONFIG = dict( _COSMOS3_TEST_CONFIG, head_dim=128, hidden_size=8 * 128, intermediate_size=1024, + rope_scaling={"rope_type": "default", "mrope_section": [24, 20, 20]}, ) # Video: [B, C, T, H, W]. patch_size=2 → seq_len = T * (H/2) * (W/2). diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py new file mode 100644 index 000000000000..ae9a43351def --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py @@ -0,0 +1,837 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for the Cosmos3 sampling policy (base vs distilled checkpoints) +and its pipeline wiring: scheduler loading, recipe validation, generation +defaults, mode resolution, and the guidance-1.0 denoise-loop contract.""" + +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +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.pipeline_cosmos3 import Cosmos3OmniMoTPipeline +from tensorrt_llm._torch.visual_gen.models.cosmos3.sampling import ( + DISTILLED_GUIDANCE_SCALE, + Cosmos3SamplingPolicy, + load_scheduler, +) +from tensorrt_llm._torch.visual_gen.models.cosmos3.transformer_cosmos3 import QWEN3_RECIPE +from tensorrt_llm._torch.visual_gen.pipeline_registry import PIPELINE_REGISTRY, AutoPipeline + +pytestmark = [pytest.mark.cosmos3, pytest.mark.usefixtures("disable_cosmos3_guardrails")] + +# The relevant subset of the 4-Step checkpoint's scheduler config +# (values verbatim; keys unrelated to distilled detection omitted). +DISTILLED_SIGMAS = (1.0, 0.9375, 0.8333333333333334, 0.625) +DISTILLED_SCHEDULER_CONFIG = { + "_class_name": "FlowMatchEulerDiscreteScheduler", + "num_train_timesteps": 1000, + "shift": 1.0, + "stochastic_sampling": True, + "use_karras_sigmas": False, + "fixed_step_requires_explicit_sigmas": True, + "fixed_step_sampler_config": { + "sample_type": "sde", + "t_list": list(DISTILLED_SIGMAS), + }, +} + +UNIPC_SCHEDULER_CONFIG = { + "_class_name": "UniPCMultistepScheduler", + "num_train_timesteps": 1000, + "flow_shift": 1.0, + "prediction_type": "flow_prediction", + "use_flow_sigmas": True, + "solver_order": 2, +} + +SKIP_NON_SCHEDULER = ["text_tokenizer", "tokenizer", "vae", "sound_tokenizer"] + + +def _write_scheduler_config(checkpoint_dir: Path, config: dict) -> None: + scheduler_dir = checkpoint_dir / "scheduler" + scheduler_dir.mkdir(parents=True, exist_ok=True) + with open(scheduler_dir / "scheduler_config.json", "w") as f: + json.dump(config, f) + + +def _distilled_policy() -> Cosmos3SamplingPolicy: + scheduler = FlowMatchEulerDiscreteScheduler.from_config(DISTILLED_SCHEDULER_CONFIG) + return Cosmos3SamplingPolicy.from_scheduler(scheduler) + + +def _base_scheduler() -> UniPCMultistepScheduler: + return UniPCMultistepScheduler.from_config(UNIPC_SCHEDULER_CONFIG) + + +def _base_policy() -> Cosmos3SamplingPolicy: + return Cosmos3SamplingPolicy.from_scheduler(_base_scheduler()) + + +def _bare_pipeline(**attrs) -> Cosmos3OmniMoTPipeline: + """A pipeline instance without heavyweight __init__; ``rank``/``dtype``/ + ``device`` are BasePipeline properties and must not be set here.""" + pipeline = object.__new__(Cosmos3OmniMoTPipeline) + defaults = dict( + audio_gen=False, + 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) + if "mode_schedulers" not in defaults: + scheduler = defaults.get("scheduler") + defaults["mode_schedulers"] = {"video": scheduler, "image": scheduler} + for key, value in defaults.items(): + setattr(pipeline, key, value) + return pipeline + + +def _fake_request(output_type: str = "video", **param_overrides) -> SimpleNamespace: + """A DiffusionRequest look-alike with executor-merged (None = unset) params.""" + params = SimpleNamespace( + height=None, + width=None, + num_inference_steps=None, + guidance_scale=None, + num_frames=COSMOS3_720P_PARAMS["num_frames"], + max_sequence_length=COSMOS3_720P_PARAMS["max_sequence_length"], + frame_rate=COSMOS3_720P_PARAMS["frame_rate"], + seed=0, + negative_prompt=None, + image=None, + extra_params={"output_type": output_type}, + ) + for key, value in param_overrides.items(): + setattr(params, key, value) + return SimpleNamespace(prompt="x", params=params) + + +class TestSchedulerLoading: + def test_flow_match_declared(self, tmp_path): + _write_scheduler_config(tmp_path, DISTILLED_SCHEDULER_CONFIG) + assert isinstance(load_scheduler(str(tmp_path)), FlowMatchEulerDiscreteScheduler) + + def test_unipc_declared(self, tmp_path): + _write_scheduler_config(tmp_path, UNIPC_SCHEDULER_CONFIG) + assert isinstance(load_scheduler(str(tmp_path)), UniPCMultistepScheduler) + + def test_missing_class_name_defaults_to_unipc(self, tmp_path): + config = {k: v for k, v in UNIPC_SCHEDULER_CONFIG.items() if k != "_class_name"} + _write_scheduler_config(tmp_path, config) + assert isinstance(load_scheduler(str(tmp_path)), UniPCMultistepScheduler) + + def test_unknown_class_name_raises(self, tmp_path): + """Silently substituting UniPC for an unknown declared scheduler would + sample the checkpoint with the wrong integrator.""" + _write_scheduler_config(tmp_path, {**UNIPC_SCHEDULER_CONFIG, "_class_name": "DDIM"}) + with pytest.raises(ValueError, match="DDIM"): + load_scheduler(str(tmp_path)) + + +class TestPolicyFacts: + def test_is_distilled(self): + assert _distilled_policy().is_distilled + 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 + + def test_checkpoint_flow_shift(self): + assert _base_policy().checkpoint_flow_shift == 1.0 + assert _distilled_policy().checkpoint_flow_shift == 1.0 # no UniPC config + + def test_scheduler_step_kwargs(self): + generator = torch.Generator().manual_seed(7) + assert _distilled_policy().scheduler_step_kwargs(generator) == {"generator": generator} + assert _base_policy().scheduler_step_kwargs(generator) == {} + + +class TestMalformedRecipeValidation: + """Only two recipes are valid; everything else must fail at load.""" + + @pytest.mark.parametrize("broken_fixed_step", [None, {}, {"t_list": []}]) + def test_required_sigmas_missing_raises(self, broken_fixed_step): + config = {k: v for k, v in DISTILLED_SCHEDULER_CONFIG.items()} + config.pop("fixed_step_sampler_config") + if broken_fixed_step is not None: + config["fixed_step_sampler_config"] = broken_fixed_step + scheduler = FlowMatchEulerDiscreteScheduler.from_config(config) + + with pytest.raises(ValueError, match="fixed_step_requires_explicit_sigmas"): + Cosmos3SamplingPolicy.from_scheduler(scheduler) + + def test_t_list_on_unipc_raises(self): + """UniPC cannot honor the distilled policy (no seeded step noise, + no baked-in guidance) even though its set_timesteps accepts sigmas.""" + config = { + **UNIPC_SCHEDULER_CONFIG, + "fixed_step_sampler_config": DISTILLED_SCHEDULER_CONFIG["fixed_step_sampler_config"], + } + scheduler = UniPCMultistepScheduler.from_config(config) + + with pytest.raises(ValueError, match="Unsupported Cosmos3 sampling recipe"): + Cosmos3SamplingPolicy.from_scheduler(scheduler) + + def test_unipc_with_declared_requirement_but_no_sigmas_raises(self): + config = {**UNIPC_SCHEDULER_CONFIG, "fixed_step_requires_explicit_sigmas": True} + scheduler = UniPCMultistepScheduler.from_config(config) + + with pytest.raises(ValueError, match="fixed_step_requires_explicit_sigmas"): + Cosmos3SamplingPolicy.from_scheduler(scheduler) + + def test_unipc_with_flag_and_sigmas_gets_unsupported_error(self): + config = { + **UNIPC_SCHEDULER_CONFIG, + "fixed_step_requires_explicit_sigmas": True, + "fixed_step_sampler_config": DISTILLED_SCHEDULER_CONFIG["fixed_step_sampler_config"], + } + scheduler = UniPCMultistepScheduler.from_config(config) + + with pytest.raises(ValueError, match="Unsupported Cosmos3 sampling recipe"): + Cosmos3SamplingPolicy.from_scheduler(scheduler) + + def test_flow_match_without_sigmas_raises(self): + config = {k: v for k, v in DISTILLED_SCHEDULER_CONFIG.items()} + config.pop("fixed_step_sampler_config") + config.pop("fixed_step_requires_explicit_sigmas") + scheduler = FlowMatchEulerDiscreteScheduler.from_config(config) + + with pytest.raises(ValueError, match="Unsupported Cosmos3 sampling recipe"): + Cosmos3SamplingPolicy.from_scheduler(scheduler) + + +class TestValidateRequest: + @pytest.mark.parametrize("steps", [None, 4]) + @pytest.mark.parametrize("guidance", [None, 1, 1.0]) + def test_valid_values_pass(self, steps, guidance): + _distilled_policy().validate_request(steps, guidance) + + @pytest.mark.parametrize("bad_steps", [1, 10, 35, 50, 100]) + def test_explicit_steps_mismatch_raises(self, bad_steps): + with pytest.raises(ValueError, match="distilled"): + _distilled_policy().validate_request(bad_steps, None) + + @pytest.mark.parametrize("bad_guidance", [0.5, 3.5, 6.0, 7.0]) + def test_explicit_guidance_mismatch_raises(self, bad_guidance): + with pytest.raises(ValueError, match="distilled"): + _distilled_policy().validate_request(None, bad_guidance) + + def test_base_policy_accepts_anything(self): + _base_policy().validate_request(17, 5.5) + _base_policy().validate_request(None, None) + + +class TestFlowShift: + def test_unipc_rebuilds_on_change(self): + policy = _base_policy() + scheduler = _base_scheduler() + + rebuilt = policy.with_flow_shift(scheduler, 3.0) + assert rebuilt is not scheduler + assert isinstance(rebuilt, UniPCMultistepScheduler) + assert float(rebuilt.config.flow_shift) == 3.0 + + def test_current_shift_read_from_scheduler_config(self): + """No separate shift-tracking state: a second call with the same target + on the rebuilt instance is a no-op; restoring rebuilds again.""" + policy = _base_policy() + rebuilt = policy.with_flow_shift(_base_scheduler(), 3.0) + assert policy.with_flow_shift(rebuilt, 3.0) is rebuilt + + restored = policy.with_flow_shift(rebuilt, 1.0) + assert restored is not rebuilt + assert float(restored.config.flow_shift) == 1.0 + + def test_distilled_is_structural_noop(self): + policy = _distilled_policy() + scheduler = FlowMatchEulerDiscreteScheduler.from_config(DISTILLED_SCHEDULER_CONFIG) + assert policy.with_flow_shift(scheduler, 3.0) is scheduler + + +class TestSetTimesteps: + def test_distilled_programs_fixed_sigmas(self): + policy = _distilled_policy() + scheduler = FlowMatchEulerDiscreteScheduler.from_config(DISTILLED_SCHEDULER_CONFIG) + policy.set_timesteps(scheduler, num_inference_steps=4, device="cpu") + expected = [s * 1000.0 for s in DISTILLED_SIGMAS] + assert torch.allclose(scheduler.timesteps.float(), torch.tensor(expected), atol=1e-3) + + def test_base_programs_step_count(self): + policy = _base_policy() + scheduler = _base_scheduler() + policy.set_timesteps(scheduler, num_inference_steps=7, device="cpu") + assert len(scheduler.timesteps) == 7 + + +class TestStochasticStepDeterminism: + """The seeded generator must fully determine the SDE noise trajectory.""" + + def _run_steps(self, seed): + policy = _distilled_policy() + scheduler = FlowMatchEulerDiscreteScheduler.from_config(DISTILLED_SCHEDULER_CONFIG) + policy.set_timesteps(scheduler, num_inference_steps=4, device="cpu") + generator = torch.Generator().manual_seed(seed) + kwargs = policy.scheduler_step_kwargs(generator) + + latents = torch.zeros(1, 4, 1, 2, 2) + velocity = torch.full_like(latents, 0.5) + for t in scheduler.timesteps: + latents = scheduler.step(velocity, t, latents, return_dict=False, **kwargs)[0] + return latents + + def test_same_seed_reproduces_sde_trajectory(self): + assert torch.equal(self._run_steps(seed=123), self._run_steps(seed=123)) + + def test_different_seeds_diverge(self): + assert not torch.equal(self._run_steps(seed=123), self._run_steps(seed=456)) + + +class TestGenerationDefaults: + def test_distilled_defaults_report_checkpoint_truth(self): + params = _bare_pipeline(sampling=_distilled_policy()).default_generation_params + assert params["num_inference_steps"] == 4 + assert params["guidance_scale"] == DISTILLED_GUIDANCE_SCALE + assert params["height"] is None # mode-dependent, resolved in infer() + assert params["num_frames"] == COSMOS3_720P_PARAMS["num_frames"] + + def test_base_defaults_leave_mode_dependent_fields_unset(self): + params = _bare_pipeline().default_generation_params + for field in ("height", "width", "num_inference_steps", "guidance_scale"): + assert params[field] is None + assert params["num_frames"] == COSMOS3_720P_PARAMS["num_frames"] + assert params["max_sequence_length"] == COSMOS3_720P_PARAMS["max_sequence_length"] + + +class TestInferModeResolution: + def _captured_forward_kwargs(self, pipeline, req): + captured = {} + pipeline.forward = lambda **kwargs: captured.update(kwargs) + pipeline.infer(req) + return captured + + def test_video_unset_resolves_to_video_table(self): + got = self._captured_forward_kwargs(_bare_pipeline(), _fake_request("video")) + assert got["height"] == COSMOS3_720P_PARAMS["height"] + assert got["width"] == COSMOS3_720P_PARAMS["width"] + assert got["num_inference_steps"] == COSMOS3_720P_PARAMS["num_inference_steps"] + assert got["guidance_scale"] == COSMOS3_720P_PARAMS["guidance_scale"] + + def test_t2i_unset_resolves_to_t2i_table(self): + got = self._captured_forward_kwargs(_bare_pipeline(), _fake_request("image")) + assert got["height"] == COSMOS3_T2I_PARAMS["height"] + assert got["width"] == COSMOS3_T2I_PARAMS["width"] + assert got["num_inference_steps"] == COSMOS3_T2I_PARAMS["num_inference_steps"] + assert got["guidance_scale"] == COSMOS3_T2I_PARAMS["guidance_scale"] + + def test_explicit_values_pass_through(self): + req = _fake_request("image", height=512, num_inference_steps=20) + got = self._captured_forward_kwargs(_bare_pipeline(), req) + assert got["height"] == 512 + assert got["num_inference_steps"] == 20 + assert got["width"] == COSMOS3_T2I_PARAMS["width"] + + def test_distilled_merged_defaults_pass_through(self): + req = _fake_request("image", num_inference_steps=4, guidance_scale=1.0) + got = self._captured_forward_kwargs(_bare_pipeline(sampling=_distilled_policy()), req) + assert got["num_inference_steps"] == 4 + assert got["guidance_scale"] == DISTILLED_GUIDANCE_SCALE + assert got["height"] == COSMOS3_T2I_PARAMS["height"] + + +class TestPipelineSchedulerLoading: + def test_distilled_checkpoint_loads_flow_match(self, tmp_path): + _write_scheduler_config(tmp_path, DISTILLED_SCHEDULER_CONFIG) + pipeline = _bare_pipeline() + + pipeline.load_standard_components( + str(tmp_path), torch.device("cpu"), skip_components=SKIP_NON_SCHEDULER + ) + + assert isinstance(pipeline.scheduler, FlowMatchEulerDiscreteScheduler) + assert pipeline.sampling.is_distilled + assert pipeline.sampling.fixed_sigmas == DISTILLED_SIGMAS + + def test_base_checkpoint_loads_unipc(self, tmp_path): + _write_scheduler_config(tmp_path, UNIPC_SCHEDULER_CONFIG) + pipeline = _bare_pipeline() + + pipeline.load_standard_components( + str(tmp_path), torch.device("cpu"), skip_components=SKIP_NON_SCHEDULER + ) + + assert isinstance(pipeline.scheduler, UniPCMultistepScheduler) + assert not pipeline.sampling.is_distilled + assert pipeline.sampling.checkpoint_flow_shift == 1.0 + + def test_audio_scheduler_is_separate_same_class_instance(self, tmp_path): + _write_scheduler_config(tmp_path, DISTILLED_SCHEDULER_CONFIG) + pipeline = _bare_pipeline(audio_gen=True) + + pipeline.load_standard_components( + str(tmp_path), torch.device("cpu"), skip_components=SKIP_NON_SCHEDULER + ) + + assert isinstance(pipeline.audio_scheduler, FlowMatchEulerDiscreteScheduler) + assert pipeline.audio_scheduler is not pipeline.scheduler + + +class TestWarmupAndForwardValidation: + def test_warmup_steps_follow_distilled_schedule(self): + assert _bare_pipeline(sampling=_distilled_policy()).default_warmup_steps == 4 + + def test_warmup_steps_base_default(self): + assert _bare_pipeline().default_warmup_steps == 2 # BasePipeline default + + @pytest.mark.parametrize( + "policy_factory, expected_guidance", + [(_distilled_policy, DISTILLED_GUIDANCE_SCALE), (_base_policy, 6.0)], + ) + def test_warmup_guidance_uses_pipeline_defaults(self, policy_factory, expected_guidance): + pipeline = _bare_pipeline(sampling=policy_factory()) + captured = {} + pipeline.forward = lambda **kwargs: captured.update(kwargs) + + pipeline._run_warmup(height=720, width=1280, num_frames=9, steps=4) + + assert captured["guidance_scale"] == expected_guidance + + @pytest.mark.parametrize( + "bad_kwargs", + [ + {"num_inference_steps": 10, "guidance_scale": 1.0}, + {"num_inference_steps": 4, "guidance_scale": 3.5}, + ], + ) + def test_forward_rejects_explicit_mismatch(self, bad_kwargs): + pipeline = _bare_pipeline(sampling=_distilled_policy()) + with pytest.raises(ValueError, match="distilled"): + pipeline.forward(prompt="x", seed=0, use_guardrails=False, **bad_kwargs) + + @pytest.mark.parametrize("bad_output_type", ["imgae", "png", "", "both"]) + def test_invalid_output_type_raises(self, bad_output_type): + pipeline = _bare_pipeline() + with pytest.raises(ValueError, match="output_type"): + pipeline.forward(prompt="x", seed=0, use_guardrails=False, output_type=bad_output_type) + + +class _AdditiveScheduler: + """step(v, t, x) = x + v, so final latents are the exact sum of all + predictions. The strict signature also pins that the loop passes no silent + extra step kwargs unless the caller supplies them.""" + + def __init__(self, timesteps): + self.timesteps = timesteps + + def step(self, model_output, timestep, sample, return_dict=False): + assert return_dict is False + return (sample + model_output,) + + +class _GeneratorRecordingScheduler(_AdditiveScheduler): + def __init__(self, timesteps): + super().__init__(timesteps) + self.generators = [] + + def step(self, model_output, timestep, sample, return_dict=False, generator=None): + self.generators.append(generator) + return super().step(model_output, timestep, sample, return_dict) + + +def _denoise_ready_pipeline() -> Cosmos3OmniMoTPipeline: + return _bare_pipeline( + pipeline_config=SimpleNamespace(visual_gen_mapping=None), + cache_accelerator=None, + _predenoise_pending=False, + _postdenoise_pending=False, + _is_warmup=False, + _profile_range=None, + ) + + +class TestDistilledDenoiseLoop: + POS_IDS = torch.arange(8).unsqueeze(0) + NEG_IDS = torch.arange(8).unsqueeze(0) + 100 + POS_MASK = torch.ones(1, 8, dtype=torch.long) + NEG_MASK = torch.zeros(1, 8, dtype=torch.long) + + def _run(self, scheduler=None, scheduler_step_kwargs=None): + pipeline = _denoise_ready_pipeline() + timesteps = torch.tensor([s * 1000.0 for s in DISTILLED_SIGMAS]) + scheduler = scheduler if scheduler is not None else _AdditiveScheduler(timesteps) + calls = [] + + def forward_fn(latent_input, extra_streams, step_index, timestep, embeds, extras): + calls.append( + { + "batch": latent_input.shape[0], + "timestep": float(timestep[0]), + "text_ids": extras["text_ids"], + } + ) + return torch.full_like(latent_input, 0.5) + + latents = torch.zeros(1, 4, 3, 2, 2) + result = pipeline.denoise( + latents=latents, + scheduler=scheduler, + prompt_embeds=self.POS_IDS, + neg_prompt_embeds=self.NEG_IDS, + guidance_scale=DISTILLED_GUIDANCE_SCALE, + forward_fn=forward_fn, + extra_cfg_tensors={ + "text_ids": (self.POS_IDS, self.NEG_IDS), + "text_mask": (self.POS_MASK, self.NEG_MASK), + }, + scheduler_step_kwargs=scheduler_step_kwargs, + ) + return result, calls + + def test_guidance_one_single_forward_per_step(self): + result, calls = self._run() + + assert len(calls) == 4, "one forward per distilled step, no CFG branch" + assert all(c["batch"] == 1 for c in calls), "no CFG batch duplication" + assert all(c["text_ids"] is self.POS_IDS for c in calls), "positive prompt only" + assert [c["timestep"] for c in calls] == pytest.approx( + [s * 1000.0 for s in DISTILLED_SIGMAS], abs=1e-3 + ) + assert torch.all(result == 2.0) # 4 additive steps of +0.5 from 0 + + def test_scheduler_step_kwargs_reach_every_step(self): + generator = torch.Generator().manual_seed(7) + timesteps = torch.tensor([s * 1000.0 for s in DISTILLED_SIGMAS]) + scheduler = _GeneratorRecordingScheduler(timesteps) + + result, calls = self._run( + scheduler=scheduler, + scheduler_step_kwargs=_distilled_policy().scheduler_step_kwargs(generator), + ) + + assert len(calls) == 4 + assert scheduler.generators == [generator] * 4 + assert torch.all(result == 2.0) + + +class _PerturbingScheduler(_AdditiveScheduler): + """step(v, t, x) = x + v + 1.0 — every position moves every step even + where the velocity is zero, emulating the distilled SDE step's + re-noising. A conditioned frame stays clean only if something re-anchors + it after each step.""" + + def step(self, model_output, timestep, sample, return_dict=False): + assert return_dict is False + return (sample + model_output + 1.0,) + + +class TestDistilledConditioningAnchor: + """Per-step re-anchoring of image-conditioned frames under SDE sampling.""" + + CLEAN = 7.0 # conditioned-frame latent value; drift is detected against it + + def _clean_frame(self): + return torch.full((1, 4, 1, 2, 2), self.CLEAN) + + def test_anchor_gating(self): + image_latent = self._clean_frame() + distilled = _bare_pipeline(sampling=_distilled_policy()) + assert callable(distilled._conditioning_anchor_post_step(image_latent)) + assert distilled._conditioning_anchor_post_step(None) is None + assert ( + _bare_pipeline(sampling=_base_policy())._conditioning_anchor_post_step(image_latent) + is None + ) + assert _bare_pipeline()._conditioning_anchor_post_step(image_latent) is None + + def test_anchor_writes_only_frame_zero_in_place(self): + pipeline = _bare_pipeline(sampling=_distilled_policy()) + post_step_fn = pipeline._conditioning_anchor_post_step(self._clean_frame()) + + latents = torch.arange(48, dtype=torch.float32).reshape(1, 4, 3, 2, 2) + untouched = latents[:, :, 1:].clone() + returned = post_step_fn(latents) + + assert returned is latents, "must write in place, not copy" + assert torch.all(latents[:, :, 0:1] == self.CLEAN) + assert torch.equal(latents[:, :, 1:], untouched) + + def _run_denoise(self, with_anchor: bool): + """Run the real BasePipeline.denoise loop with a perturbing scheduler, + recording what the transformer receives at every step.""" + pipeline = _denoise_ready_pipeline() + pipeline.sampling = _distilled_policy() + timesteps = torch.tensor([s * 1000.0 for s in DISTILLED_SIGMAS]) + scheduler = _PerturbingScheduler(timesteps) + + seen = [] + + def forward_fn(latent_input, extra_streams, step_index, timestep, embeds, extras): + seen.append(latent_input.clone()) + return torch.full_like(latent_input, 0.5) + + latents = torch.zeros(1, 4, 3, 2, 2) + latents[:, :, 0:1] = self.CLEAN # frame 0 pinned clean, rest noise-like + image_latent = self._clean_frame() + + post_step_fn = ( + pipeline._conditioning_anchor_post_step(image_latent) if with_anchor else None + ) + result = pipeline.denoise( + latents=latents, + scheduler=scheduler, + prompt_embeds=torch.arange(8).unsqueeze(0), + neg_prompt_embeds=torch.arange(8).unsqueeze(0) + 100, + guidance_scale=DISTILLED_GUIDANCE_SCALE, + forward_fn=forward_fn, + extra_cfg_tensors={}, + post_step_fn=post_step_fn, + ) + return result, seen + + def test_every_forward_sees_clean_conditioned_frame(self): + result, seen = self._run_denoise(with_anchor=True) + + assert len(seen) == 4 + for step, latent_input in enumerate(seen): + assert torch.all(latent_input[:, :, 0:1] == self.CLEAN), ( + f"transformer input at step {step} lost the clean conditioning frame" + ) + # The perturbing step really moved everything else: unconditioned + # frames accumulate (velocity 0.5 + drift 1.0) per completed step. + for step, latent_input in enumerate(seen): + assert torch.all(latent_input[:, :, 1:] == step * 1.5) + assert torch.all(result[:, :, 0:1] == self.CLEAN) + assert torch.all(result[:, :, 1:] == 4 * 1.5) + + def test_without_anchor_the_conditioned_frame_drifts(self): + """Control: the same loop without the anchor corrupts frame 0 from the + second forward on — the exact failure mode the anchor exists for.""" + _, seen = self._run_denoise(with_anchor=False) + + assert torch.all(seen[0][:, :, 0:1] == self.CLEAN) + for step, latent_input in enumerate(seen[1:], start=1): + assert torch.all(latent_input[:, :, 0:1] == self.CLEAN + step * 1.5) + + +class TestForwardConditioningWiring: + """forward() must hand the denoise loop the anchor exactly when the + checkpoint is distilled and the request carries image conditioning.""" + + T_LAT, H_LAT, W_LAT = 2, 2, 2 # from num_frames=5, 32x32, scale 4/16 + CLEAN = 7.0 + + def _forward_ready_pipeline(self): + pipeline = _bare_pipeline( + sampling=_distilled_policy(), + pipeline_config=SimpleNamespace(torch_dtype=torch.float32, visual_gen_mapping=None), + transformer=SimpleNamespace( + latent_channel_size=4, + reset_cache=lambda: None, + device=torch.device("cpu"), + ), + vae_scale_factor_temporal=4, + vae_scale_factor_spatial=16, + scheduler=SimpleNamespace( + set_timesteps=lambda *args, **kwargs: None, + config=SimpleNamespace(num_train_timesteps=1000), + ), + ) + pipeline._tokenize_prompt = lambda *args, **kwargs: ( + torch.ones(1, 4, dtype=torch.long), + torch.ones(1, 4, dtype=torch.long), + ) + pipeline._encode_conditioning_video = lambda *args, **kwargs: torch.full( + (1, 4, self.T_LAT, self.H_LAT, self.W_LAT), self.CLEAN + ) + pipeline.decode_latents = lambda latents, decode_fn: torch.zeros(1, 5, 32, 32, 3) + + captured = {} + + def denoise(**kwargs): + captured.update(kwargs) + return kwargs["latents"] + + pipeline.denoise = denoise + return pipeline, captured + + def _forward(self, pipeline, image): + return pipeline.forward( + prompt="x", + seed=0, + image=image, + height=32, + width=32, + num_frames=5, + num_inference_steps=4, + guidance_scale=DISTILLED_GUIDANCE_SCALE, + use_guardrails=False, + enable_audio=False, + ) + + def test_i2v_request_wires_anchor_and_seeded_steps(self): + pipeline, captured = self._forward_ready_pipeline() + self._forward(pipeline, image=torch.zeros(3, 32, 32)) + + post_step_fn = captured["post_step_fn"] + assert post_step_fn is not None + latents = torch.zeros(1, 4, self.T_LAT, self.H_LAT, self.W_LAT) + post_step_fn(latents) + assert torch.all(latents[:, :, 0:1] == self.CLEAN) + assert torch.all(latents[:, :, 1:] == 0.0) + + assert isinstance(captured["scheduler_step_kwargs"]["generator"], torch.Generator) + # Initial latents enter the loop with the clean frame already pinned. + assert torch.all(captured["latents"][:, :, 0:1] == self.CLEAN) + + def test_t2v_request_wires_no_anchor(self): + pipeline, captured = self._forward_ready_pipeline() + self._forward(pipeline, image=None) + assert captured["post_step_fn"] is None + + +class TestSystemPromptDefault: + """use_system_prompt defaults are checkpoint-declared via model_index.json.""" + + def _write_model_index(self, checkpoint_dir: Path, content: dict) -> None: + with open(checkpoint_dir / "model_index.json", "w") as f: + json.dump(content, f) + + def _loaded_pipeline(self, tmp_path) -> Cosmos3OmniMoTPipeline: + _write_scheduler_config(tmp_path, DISTILLED_SCHEDULER_CONFIG) + pipeline = _bare_pipeline() + pipeline.load_standard_components( + str(tmp_path), torch.device("cpu"), skip_components=SKIP_NON_SCHEDULER + ) + return pipeline + + 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.extra_param_specs["use_system_prompt"].default is True + # The shared spec table must stay untouched (model_copy, not mutation). + from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import COSMOS3_EXTRA_SPECS + + assert COSMOS3_EXTRA_SPECS["use_system_prompt"].default is False + + def test_missing_model_index_keeps_false(self, tmp_path): + pipeline = self._loaded_pipeline(tmp_path) + assert pipeline.default_use_system_prompt is False + assert pipeline.extra_param_specs["use_system_prompt"].default is False + + def test_model_index_without_field_keeps_false(self, tmp_path): + self._write_model_index(tmp_path, {"_class_name": "Cosmos3OmniPipeline"}) + pipeline = self._loaded_pipeline(tmp_path) + assert pipeline.default_use_system_prompt is False + + def _captured_use_system_prompt(self, pipeline, extra_params): + captured = {} + pipeline.forward = lambda **kwargs: captured.update(kwargs) + pipeline.infer(_fake_request("video", extra_params=extra_params)) + return captured["use_system_prompt"] + + def test_infer_unset_key_uses_checkpoint_default(self): + pipeline = _bare_pipeline(default_use_system_prompt=True) + assert self._captured_use_system_prompt(pipeline, {"output_type": "video"}) is True + + def test_infer_explicit_false_preserved(self): + pipeline = _bare_pipeline(default_use_system_prompt=True) + got = self._captured_use_system_prompt( + pipeline, {"output_type": "video", "use_system_prompt": False} + ) + assert got is False + + +class TestAudioWeightPresenceGuard: + """enable_audio=True must fail loudly when the checkpoint ships no audio + tower — a weight-presence guard, not a workflow restriction.""" + + def _pipeline(self, **attrs): + return _bare_pipeline(sampling=_distilled_policy(), scheduler=None, **attrs) + + def test_explicit_audio_on_audioless_checkpoint_raises(self): + with pytest.raises(ValueError, match="audio tower"): + self._pipeline().forward( + prompt="x", + seed=0, + use_guardrails=False, + enable_audio=True, + num_inference_steps=4, + guidance_scale=DISTILLED_GUIDANCE_SCALE, + ) + + def test_t2i_disables_audio_before_the_guard(self): + """T2I force-disables audio for every checkpoint (existing semantics); + the guard must not fire for it. The batch error proves forward got + past the guard.""" + with pytest.raises(ValueError, match="Batch generation"): + self._pipeline().forward( + prompt=["a", "b"], + seed=0, + use_guardrails=False, + enable_audio=True, + output_type="image", + num_inference_steps=4, + guidance_scale=DISTILLED_GUIDANCE_SCALE, + ) + + def test_audio_capable_checkpoint_passes_the_guard(self): + with pytest.raises(ValueError, match="Batch generation"): + self._pipeline(audio_gen=True).forward( + prompt=["a", "b"], + seed=0, + use_guardrails=False, + enable_audio=True, + num_inference_steps=4, + guidance_scale=DISTILLED_GUIDANCE_SCALE, + ) + + +class TestRegistryDispatch: + def test_model_index_class_name_dispatches(self, tmp_path): + with open(tmp_path / "model_index.json", "w") as f: + json.dump({"_class_name": "Cosmos3OmniPipeline"}, f) + assert AutoPipeline._detect_from_checkpoint(str(tmp_path)) == "Cosmos3OmniMoTPipeline" + + def test_hf_id_registered(self): + entry = PIPELINE_REGISTRY["Cosmos3OmniMoTPipeline"] + assert "nvidia/Cosmos3-Super-Text2Image-4Step" in entry.hf_ids + assert "nvidia/Cosmos3-Super-Image2Video-4Step" in entry.hf_ids diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py b/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py new file mode 100644 index 000000000000..131ad19dba9a --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py @@ -0,0 +1,1342 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Cosmos3-Edge (Nemotron-dense backbone) tests. + +Unit tests run on a reduced config mirroring the Edge checkpoint's exact key +set: recipe validation, Nemotron norm semantics, the generator-only und +K-norm, native flow schedule parity against cosmos-framework, strict weight +loading, and per-family defaults. Checkpoint-gated tests cover the real +checkpoint (tokenizer, recipe/scheduler wiring, load + forward). + +Override checkpoint: + DIFFUSION_MODEL_PATH_COSMOS3_EDGE=/path/to/Cosmos3-Edge \\ + pytest tests/unittest/_torch/visual_gen/test_cosmos3_edge.py -v +""" + +import gc +import os +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from diffusers import UniPCMultistepScheduler + +from tensorrt_llm._torch.modules.mlp import MLP +from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig +from tensorrt_llm._torch.visual_gen.models.cosmos3 import pipeline_cosmos3 as pipeline_module +from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( + COSMOS3_720P_PARAMS, + COSMOS3_EDGE_T2I_PARAMS, + COSMOS3_EDGE_VIDEO_PARAMS, + COSMOS3_GENERATION_DEFAULTS, +) +from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniMoTPipeline +from tensorrt_llm._torch.visual_gen.models.cosmos3.sampling import Cosmos3SamplingPolicy +from tensorrt_llm._torch.visual_gen.models.cosmos3.transformer_cosmos3 import ( + COSMOS3_EDGE_BACKBONE_TYPE, + NEMOTRON_DENSE_RECIPE, + QWEN3_RECIPE, + Cosmos3VFMTransformer, + NemotronRMSNorm, + Qwen3VLTextRMSNorm, + resolve_arch_recipe, +) +from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineComponent, PipelineLoader +from tensorrt_llm._torch.visual_gen.pipeline_registry import PIPELINE_REGISTRY +from tensorrt_llm.visual_gen.args import TorchCompileConfig, VisualGenArgs + +pytestmark = [pytest.mark.cosmos3, pytest.mark.usefixtures("disable_cosmos3_guardrails")] + +DEVICE = "cuda" + + +@pytest.fixture(autouse=True) +def _cleanup_gpu(): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + yield + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +def _require_edge_checkpoint() -> str: + """Resolve the Edge checkpoint lazily so unit tests collect and run on + machines without model storage; only checkpoint-gated tests skip.""" + path = os.environ.get("DIFFUSION_MODEL_PATH_COSMOS3_EDGE") + if not path: + root = Path(os.environ.get("LLM_MODELS_ROOT", "/home/scratch.trt_llm_data_ci/llm-models/")) + if not root.exists(): + root = Path("/scratch/trt_llm_data/llm-models/") + path = str(root / "Cosmos3-Edge") + if not os.path.isdir(path): + pytest.skip(f"Checkpoint not found: {path}") + return path + + +def _reduced_edge_config() -> SimpleNamespace: + # Key set mirrors the Edge checkpoint's transformer/config.json verbatim + # (including the missing rope_type in rope_scaling); only sizes shrink. + return SimpleNamespace( + action_dim=8, + action_gen=True, + attention_bias=False, + attention_dropout=0.0, + backbone_type=COSMOS3_EDGE_BACKBONE_TYPE, + base_fps=24, + enable_fps_modulation=True, + head_dim=8, + hidden_act="relu2", + hidden_size=32, + intermediate_size=64, + # Latent geometry stays at the real invariant values (validated for + # the Edge recipe); only backbone dimensions shrink. + latent_channel=48, + latent_patch_size=2, + num_attention_heads=4, + num_embodiment_domains=32, + num_hidden_layers=2, + num_key_value_heads=2, + patch_latent_dim=192, + qk_norm_for_text=False, + rms_norm_eps=1e-5, + rope_axes_dim=[2, 1, 1], + rope_scaling={"mrope_section": [2, 1, 1]}, + rope_theta=100000000, + sound_dim=None, + sound_gen=False, + temporal_compression_factor=4, + timestep_scale=0.001, + unified_3d_mrope_reset_spatial_ids=True, + unified_3d_mrope_temporal_modality_margin=15000, + use_und_k_norm_for_gen=True, + vocab_size=64, + ) + + +def _reduced_edge_model_config() -> DiffusionModelConfig: + model_config = DiffusionModelConfig(pretrained_config=_reduced_edge_config()) + model_config.attention.backend = "VANILLA" + return model_config + + +def _reduced_qwen3_config(**overrides) -> SimpleNamespace: + cfg = _reduced_edge_config() + cfg.backbone_type = None + cfg.hidden_act = "silu" + cfg.qk_norm_for_text = True + cfg.use_und_k_norm_for_gen = False + cfg.rms_norm_eps = 1e-6 + cfg.rope_theta = 5000000 + cfg.rope_scaling = {"mrope_section": [2, 1, 1], "rope_type": "default"} + for key, value in overrides.items(): + setattr(cfg, key, value) + return cfg + + +def _reduced_qwen3_model_config(**overrides) -> DiffusionModelConfig: + model_config = DiffusionModelConfig(pretrained_config=_reduced_qwen3_config(**overrides)) + model_config.attention.backend = "VANILLA" + return model_config + + +def _synthetic_state_dict(cfg: SimpleNamespace) -> dict: + """A complete synthetic checkpoint for either recipe (diffusers key layout).""" + h, d = cfg.hidden_size, cfg.head_dim + q_dim = cfg.num_attention_heads * d + kv_dim = cfg.num_key_value_heads * d + gated = cfg.hidden_act == "silu" + sd = { + "embed_tokens.weight": torch.randn(cfg.vocab_size, h), + "lm_head.weight": torch.randn(cfg.vocab_size, h), + "norm.weight": torch.ones(h), + "norm_moe_gen.weight": torch.ones(h), + "proj_in.weight": torch.randn(h, cfg.patch_latent_dim), + "proj_in.bias": torch.randn(h), + "proj_out.weight": torch.randn(cfg.patch_latent_dim, h), + "proj_out.bias": torch.randn(cfg.patch_latent_dim), + "time_embedder.linear_1.weight": torch.randn(h, 256), + "time_embedder.linear_1.bias": torch.randn(h), + "time_embedder.linear_2.weight": torch.randn(h, h), + "time_embedder.linear_2.bias": torch.randn(h), + } + if getattr(cfg, "sound_gen", False): + sd.update( + { + "audio_proj_in.weight": torch.randn(h, cfg.sound_dim), + "audio_proj_in.bias": torch.randn(h), + "audio_proj_out.weight": torch.randn(cfg.sound_dim, h), + "audio_proj_out.bias": torch.randn(cfg.sound_dim), + "audio_modality_embed": torch.randn(h), + } + ) + for i in range(cfg.num_hidden_layers): + p = f"layers.{i}" + sd.update( + { + f"{p}.self_attn.to_q.weight": torch.randn(q_dim, h), + f"{p}.self_attn.to_k.weight": torch.randn(kv_dim, h), + f"{p}.self_attn.to_v.weight": torch.randn(kv_dim, h), + f"{p}.self_attn.to_out.weight": torch.randn(h, q_dim), + f"{p}.self_attn.add_q_proj.weight": torch.randn(q_dim, h), + f"{p}.self_attn.add_k_proj.weight": torch.randn(kv_dim, h), + f"{p}.self_attn.add_v_proj.weight": torch.randn(kv_dim, h), + f"{p}.self_attn.to_add_out.weight": torch.randn(h, q_dim), + f"{p}.self_attn.norm_added_q.weight": torch.ones(d), + f"{p}.self_attn.norm_added_k.weight": torch.ones(d), + f"{p}.input_layernorm.weight": torch.ones(h), + f"{p}.input_layernorm_moe_gen.weight": torch.ones(h), + f"{p}.post_attention_layernorm.weight": torch.ones(h), + f"{p}.post_attention_layernorm_moe_gen.weight": torch.ones(h), + f"{p}.mlp.up_proj.weight": torch.randn(cfg.intermediate_size, h), + f"{p}.mlp.down_proj.weight": torch.randn(h, cfg.intermediate_size), + f"{p}.mlp_moe_gen.up_proj.weight": torch.randn(cfg.intermediate_size, h), + f"{p}.mlp_moe_gen.down_proj.weight": torch.randn(h, cfg.intermediate_size), + } + ) + if gated: + sd.update( + { + f"{p}.self_attn.norm_q.weight": torch.ones(d), + f"{p}.self_attn.norm_k.weight": torch.ones(d), + f"{p}.mlp.gate_proj.weight": torch.randn(cfg.intermediate_size, h), + f"{p}.mlp_moe_gen.gate_proj.weight": torch.randn(cfg.intermediate_size, h), + } + ) + else: + sd[f"{p}.self_attn.k_norm_und_for_gen.weight"] = torch.ones(d) + return sd + + +def _edge_state_dict(cfg: SimpleNamespace) -> dict: + return _synthetic_state_dict(cfg) + + +class TestArchRecipe: + def test_edge_config_resolves_nemotron_recipe(self): + recipe = resolve_arch_recipe(_reduced_edge_config()) + assert recipe is NEMOTRON_DENSE_RECIPE + + def test_qwen3_resolves_without_backbone_type(self): + cfg = SimpleNamespace(hidden_act="silu", qk_norm_for_text=True) + assert resolve_arch_recipe(cfg) is QWEN3_RECIPE + + def test_unknown_backbone_raises(self): + cfg = _reduced_edge_config() + cfg.backbone_type = "cosmos4_hybrid" + with pytest.raises(ValueError, match="cosmos4_hybrid"): + resolve_arch_recipe(cfg) + + @pytest.mark.parametrize( + "key,value", + [ + ("hidden_act", "silu"), + ("qk_norm_for_text", True), + ("use_und_k_norm_for_gen", False), + ("sound_gen", True), + ("sound_gen", None), + ("attention_bias", True), + ("rms_norm_eps", 1e-6), + ], + ) + def test_edge_flag_contradiction_raises(self, key, value): + cfg = _reduced_edge_config() + setattr(cfg, key, value) + with pytest.raises(ValueError, match=key): + resolve_arch_recipe(cfg) + + @pytest.mark.parametrize("make_config", [_reduced_edge_config, lambda: _reduced_qwen3_config()]) + def test_inconsistent_patch_latent_dim_raises(self, make_config): + cfg = make_config() + cfg.patch_latent_dim = cfg.patch_latent_dim + 1 + with pytest.raises(ValueError, match="patch_latent_dim"): + resolve_arch_recipe(cfg) + + +class TestRopeAxesResolution: + """The explicit top-level rope_axes_dim wins over the legacy + rope_scaling.mrope_section (diffusers precedence); disagreement and + head_dim mismatches are config errors.""" + + def _resolve(self, cfg): + from tensorrt_llm._torch.visual_gen.models.cosmos3.transformer_cosmos3 import ( + resolve_rope_axes_dim, + ) + + return resolve_rope_axes_dim(cfg) + + def test_top_level_only(self): + cfg = _reduced_edge_config() + cfg.rope_scaling = {} + assert self._resolve(cfg) == [2, 1, 1] + + def test_nested_only(self): + cfg = _reduced_edge_config() + del cfg.rope_axes_dim + assert self._resolve(cfg) == [2, 1, 1] + + def test_contradiction_raises(self): + cfg = _reduced_edge_config() + cfg.rope_axes_dim = [1, 2, 1] + with pytest.raises(ValueError, match="contradictory"): + self._resolve(cfg) + + @pytest.mark.parametrize("axes", [[2, 2, 1], [2, 2], [4]]) + def test_head_dim_mismatch_raises(self, axes): + cfg = _reduced_edge_config() + cfg.rope_axes_dim = axes + cfg.rope_scaling = {} + with pytest.raises(ValueError, match="head_dim"): + self._resolve(cfg) + + def test_neither_declared_raises(self): + cfg = _reduced_edge_config() + del cfg.rope_axes_dim + cfg.rope_scaling = {} + with pytest.raises(ValueError, match="neither"): + self._resolve(cfg) + + def test_transformer_builds_from_top_level_only(self): + cfg = _reduced_edge_config() + cfg.rope_scaling = {} + model_config = DiffusionModelConfig(pretrained_config=cfg) + model_config.attention.backend = "VANILLA" + model = Cosmos3VFMTransformer(model_config) + assert model.language_model.rotary_emb.mrope_section == [2, 1, 1] + + def test_qwen3_flag_contradiction_raises(self): + cfg = SimpleNamespace(hidden_act="relu2") + with pytest.raises(ValueError, match="hidden_act"): + resolve_arch_recipe(cfg) + + @pytest.mark.parametrize("key,value", [("latent_channel", 4), ("latent_patch_size", 4)]) + def test_edge_latent_geometry_invariants(self, key, value): + cfg = _reduced_edge_config() + setattr(cfg, key, value) + with pytest.raises(ValueError, match=key): + resolve_arch_recipe(cfg) + + +class TestNemotronNormSemantics: + """Pins the one intentional numerics fork: fp32 weight multiply, then + downcast — vs the Qwen flavor's bf16 multiply after downcast.""" + + def test_matches_fp32_weight_multiply_bit_exact(self): + torch.manual_seed(0) + x = (torch.randn(256, 8) * 3).bfloat16() + weight = (torch.randn(8) * 2 + 1.5).bfloat16() + + nemotron = NemotronRMSNorm(hidden_size=8, eps=1e-5) + qwen = Qwen3VLTextRMSNorm(hidden_size=8, eps=1e-5) + with torch.no_grad(): + nemotron.weight.copy_(weight) + qwen.weight.copy_(weight) + + xf = x.float() + normed = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + 1e-5) + reference = (weight.float() * normed).to(torch.bfloat16) + + assert torch.equal(nemotron(x), reference) + assert not torch.equal(qwen(x), reference) + + +class TestEdgeTransformerStructure: + @pytest.fixture(scope="class") + def model(self): + return Cosmos3VFMTransformer(_reduced_edge_model_config()) + + def test_recipe_and_flags(self, model): + assert model.recipe is NEMOTRON_DENSE_RECIPE + 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 + assert attn.norm_q is None + assert attn.norm_k is None + assert isinstance(attn.k_norm_und_for_gen, NemotronRMSNorm) + + def test_nemotron_norms_everywhere(self, model): + und = model.language_model.layers[0] + gen = model.gen_layers[0] + for norm in ( + und.input_layernorm, + und.post_attention_layernorm, + und.self_attn.k_norm_und_for_gen, + gen.input_layernorm, + gen.post_attention_layernorm, + gen.cross_attention.norm_q, + gen.cross_attention.norm_k, + model.norm_moe_gen, + ): + assert isinstance(norm, NemotronRMSNorm) + assert norm.variance_epsilon == 1e-5 + + def test_relu2_mlp_no_gate(self, model): + for layer in (model.language_model.layers[0], model.gen_layers[0]): + assert isinstance(layer.mlp, MLP) + assert not hasattr(layer.mlp, "gate_proj") + + def test_rope_defaults_without_rope_type(self, model): + assert model.language_model.rotary_emb.rope_type == "default" + + +class TestGeneratorOnlyKNorm: + """Mirror of diffusers' Edge regression: perturbing k_norm_und_for_gen + must not change the und causal attention output, only the gen-facing K.""" + + def test_k_norm_touches_only_gen_facing_keys(self): + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + torch.manual_seed(0) + model = Cosmos3VFMTransformer(_reduced_edge_model_config()).to(DEVICE).eval() + with torch.no_grad(): + for param in model.parameters(): + param.normal_(0, 0.02) + model.post_load_weights() + + attn = model.language_model.layers[0].self_attn + hidden = torch.randn(1, 4, 32, dtype=torch.bfloat16, device=DEVICE) + cos = torch.ones(1, 4, 1, 8, dtype=torch.bfloat16, device=DEVICE) + sin = torch.zeros(1, 4, 1, 8, dtype=torch.bfloat16, device=DEVICE) + + with torch.inference_mode(): + out_before, k_gen_before, v_before = attn.forward_with_kv(hidden, cos, sin) + # The rope above is identity (cos=1, sin=0), so the cached gen K + # must be exactly the Nemotron-normed raw K. + _, k_raw, _ = attn.get_qkv(hidden) + k_raw = k_raw.view(1, 4, attn.local_num_key_value_heads, attn.head_dim) + assert torch.equal(k_gen_before, attn.k_norm_und_for_gen(k_raw)) + + attn.k_norm_und_for_gen.weight.fill_(2.0) + out_after, k_gen_after, v_after = attn.forward_with_kv(hidden, cos, sin) + + assert torch.equal(out_before, out_after) + assert torch.equal(v_before, v_after) + assert not torch.equal(k_gen_before, k_gen_after) + + def test_nano_qk_norm_stays_byte_identical(self): + """Regression pin for the qwen3 recipe: apply_qk_norm must keep + today's exact ``F.rms_norm`` semantics (a future refactor routing it + through ``Qwen3VLTextRMSNorm.forward`` would change Nano numerics).""" + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + import torch.nn.functional as F + + torch.manual_seed(0) + model = Cosmos3VFMTransformer(_reduced_qwen3_model_config()).to(DEVICE).eval() + with torch.no_grad(): + for param in model.parameters(): + param.normal_(0, 0.02) + attn = model.language_model.layers[0].self_attn + + q = torch.randn(1, 4, 4, 8, dtype=torch.bfloat16, device=DEVICE) + k = torch.randn(1, 4, 2, 8, dtype=torch.bfloat16, device=DEVICE) + q_normed, k_normed = attn.apply_qk_norm(q, k) + assert torch.equal( + q_normed, F.rms_norm(q, (8,), attn.norm_q.weight, attn.norm_q.variance_epsilon) + ) + assert torch.equal( + k_normed, F.rms_norm(k, (8,), attn.norm_k.weight, attn.norm_k.variance_epsilon) + ) + + +# The relevant subset of the Edge checkpoint's scheduler config (values +# verbatim; identical to Nano's). +EDGE_UNIPC_CONFIG = { + "_class_name": "UniPCMultistepScheduler", + "num_train_timesteps": 1000, + "flow_shift": 1.0, + "prediction_type": "flow_prediction", + "use_flow_sigmas": True, + "use_karras_sigmas": True, + "sigma_max": 200.0, + "sigma_min": 0.147, + "solver_order": 2, + "solver_type": "bh2", + "final_sigmas_type": "zero", + "timestep_spacing": "linspace", + "lower_order_final": True, +} + +# Recorded from cosmos-framework 117c7d2 (`fm_solvers_unipc.py` +# FlowUniPCMultistepScheduler, num_train_timesteps=1000, shift=1.0, +# use_dynamic_shifting=False; set_timesteps(steps, shift=shift)). Each +# trajectory runs the synthetic velocity v = 0.05*x + 0.3*sin(t/1000) - 0.1 in +# float64 from x0 = linspace(-1, 1, 8).reshape(1, 2, 2, 2) through every +# step(). +COSMOS_FRAMEWORK_FIXTURES = { + (3.0, 10): { + "timesteps": [999, 963, 922, 874, 817, 749, 666, 562, 428, 249], + "sigmas": [ + 0.99966645, + 0.96394110, + 0.92272168, + 0.87463522, + 0.81780970, + 0.74962479, + 0.66629612, + 0.56214833, + 0.42826521, + 0.24979164, + 0.0, + ], + "final": [ + -0.995152214121, + -0.723385865647, + -0.451619517172, + -0.179853168698, + 0.091913179777, + 0.363679528251, + 0.635445876726, + 0.907212225200, + ], + }, + (10.0, 7): { + "timesteps": [999, 983, 961, 930, 882, 799, 624], + "sigmas": [ + 0.99989992, + 0.98349357, + 0.96140891, + 0.93008101, + 0.88217115, + 0.79977584, + 0.62472641, + 0.0, + ], + "final": [ + -1.040962681673, + -0.769320359241, + -0.497678036809, + -0.226035714377, + 0.045606608054, + 0.317248930486, + 0.588891252918, + 0.860533575350, + ], + }, + (5.0, 13): { + "timesteps": [999, 983, 964, 943, 918, 888, 853, 810, 757, 689, 599, 475, 293], + "sigmas": [ + 0.99979985, + 0.98339677, + 0.96469206, + 0.94316465, + 0.91812354, + 0.88863194, + 0.85338765, + 0.81052577, + 0.75727713, + 0.68934584, + 0.59968787, + 0.47589558, + 0.29389268, + 0.0, + ], + "final": [ + -0.998983254680, + -0.727228149256, + -0.455473043832, + -0.183717938407, + 0.088037167017, + 0.359792272441, + 0.631547377865, + 0.903302483290, + ], + }, +} + + +class TestNativeFlowSchedule: + def _native_policy_and_scheduler(self, shift: float): + scheduler = UniPCMultistepScheduler.from_config(EDGE_UNIPC_CONFIG) + policy = Cosmos3SamplingPolicy.from_scheduler(scheduler, native_flow_schedule=True) + return policy, policy.with_flow_shift(scheduler, shift) + + def test_with_flow_shift_disables_karras(self): + policy, scheduler = self._native_policy_and_scheduler(3.0) + assert float(scheduler.config.flow_shift) == 3.0 + assert scheduler.config.use_karras_sigmas is False + # Already matching → same instance. + assert policy.with_flow_shift(scheduler, 3.0) is scheduler + + def test_karras_rebuilt_even_at_checkpoint_shift(self): + """flow_shift 1.0 equals the checkpoint value, but the native flow + schedule still requires the karras grid off.""" + policy, scheduler = self._native_policy_and_scheduler(1.0) + assert scheduler.config.use_karras_sigmas is False + + def test_non_native_keeps_checkpoint_config(self): + scheduler = UniPCMultistepScheduler.from_config(EDGE_UNIPC_CONFIG) + policy = Cosmos3SamplingPolicy.from_scheduler(scheduler, native_flow_schedule=False) + assert policy.with_flow_shift(scheduler, 1.0) is scheduler + assert scheduler.config.use_karras_sigmas is True + + @pytest.mark.parametrize("shift,steps", sorted(COSMOS_FRAMEWORK_FIXTURES)) + def test_matches_cosmos_framework_reference(self, shift, steps): + fixture = COSMOS_FRAMEWORK_FIXTURES[(shift, steps)] + policy, scheduler = self._native_policy_and_scheduler(shift) + policy.set_timesteps(scheduler, num_inference_steps=steps, device="cpu") + + assert scheduler.timesteps.tolist() == fixture["timesteps"] + np.testing.assert_allclose(scheduler.sigmas.tolist(), fixture["sigmas"], atol=1e-6) + + x = torch.linspace(-1.0, 1.0, 8, dtype=torch.float64).reshape(1, 2, 2, 2) + for t in scheduler.timesteps: + v = 0.05 * x + 0.3 * torch.sin(t.double() / 1000.0) - 0.1 + x = scheduler.step(v, t, x, return_dict=False)[0] + np.testing.assert_allclose(x.flatten().tolist(), fixture["final"], atol=1e-9) + + +class TestStrictLoading: + def _model(self) -> Cosmos3VFMTransformer: + return Cosmos3VFMTransformer(_reduced_edge_model_config()) + + def test_full_checkpoint_loads(self): + cfg = _reduced_edge_config() + sd = _edge_state_dict(cfg) + sd["layers.0.self_attn.k_norm_und_for_gen.weight"] = torch.full((8,), 0.5) + model = self._model() + model.load_weights(sd) + loaded = model.language_model.layers[0].self_attn.k_norm_und_for_gen.weight + assert torch.equal(loaded.cpu().float(), torch.full((8,), 0.5)) + + @pytest.mark.parametrize( + "missing_key", + [ + "layers.0.self_attn.k_norm_und_for_gen.weight", + "layers.0.mlp.up_proj.weight", + "layers.1.mlp_moe_gen.down_proj.weight", + "layers.0.input_layernorm_moe_gen.weight", + "embed_tokens.weight", + ], + ) + def test_missing_weight_raises(self, missing_key): + sd = _edge_state_dict(_reduced_edge_config()) + del sd[missing_key] + with pytest.raises(ValueError, match="missing weights"): + self._model().load_weights(sd) + + def test_partial_fused_qkv_raises(self): + sd = _edge_state_dict(_reduced_edge_config()) + del sd["layers.0.self_attn.add_v_proj.weight"] + with pytest.raises(ValueError, match="missing weights"): + self._model().load_weights(sd) + + def test_intentional_skips_are_logged_with_names(self, monkeypatch): + """The skip log must name the skipped tensor families (dynamic + content, not just the logger's static category text).""" + import tensorrt_llm._torch.visual_gen.models.cosmos3.transformer_cosmos3 as tf_module + + infos = [] + monkeypatch.setattr(tf_module.logger, "info", infos.append) + cfg = _reduced_edge_config() + sd = _edge_state_dict(cfg) + sd.update( + { + "action_modality_embed": torch.randn(cfg.hidden_size), + "action_proj_in.fc.weight": torch.randn(4, 8), + "action_proj_in.bias.weight": torch.randn(4, 8), + "action_proj_out.fc.weight": torch.randn(4, 8), + "action_proj_out.bias.weight": torch.randn(4, 8), + } + ) + model = self._model() + model.load_weights(sd) + + param_names = {name for name, _ in model.named_parameters()} + assert not any("lm_head" in name for name in param_names) + assert "language_model.norm.weight" not in param_names + skip_logs = [m for m in infos if "intentionally unused" in m] + assert len(skip_logs) == 1 + # The dynamic name list follows the final colon; the static category + # text also mentions lm_head/norm, so assert the parsed set exactly. + skipped_families = {name.strip() for name in skip_logs[0].rsplit(": ", 1)[1].split(",")} + assert skipped_families == { + "action_modality_embed", + "action_proj_in", + "action_proj_out", + "lm_head", + "norm", + } + + def test_model_prefixed_skip_keys_are_intentional(self, monkeypatch): + """Checkpoints that namespace top-level tensors under "model." must + have their skip keys recognized (prefix normalization runs before the + skip check), not warned about as unknown.""" + import tensorrt_llm._torch.visual_gen.models.cosmos3.transformer_cosmos3 as tf_module + + warnings = [] + infos = [] + monkeypatch.setattr(tf_module.logger, "warning", warnings.append) + monkeypatch.setattr(tf_module.logger, "info", infos.append) + cfg = _reduced_edge_config() + sd = _edge_state_dict(cfg) + sd["model.lm_head.weight"] = torch.randn(cfg.vocab_size, cfg.hidden_size) + sd["model.action_modality_embed"] = torch.randn(cfg.hidden_size) + self._model().load_weights(sd) + + assert not any("unknown checkpoint key" in m for m in warnings) + skip_logs = [m for m in infos if "intentionally unused" in m] + assert len(skip_logs) == 1 + skipped_families = {name.strip() for name in skip_logs[0].rsplit(": ", 1)[1].split(",")} + assert {"lm_head", "action_modality_embed"} <= skipped_families + + def test_unconsumed_mapped_tensor_warns(self, monkeypatch): + """A checkpoint tensor that remaps to a module the recipe didn't + construct must warn, not vanish: Edge has no und norm_q, and a + sound_gen=false model constructs no audio projections.""" + import tensorrt_llm._torch.visual_gen.models.cosmos3.transformer_cosmos3 as tf_module + + warnings = [] + monkeypatch.setattr(tf_module.logger, "warning", warnings.append) + sd = _edge_state_dict(_reduced_edge_config()) + sd["layers.0.self_attn.norm_q.weight"] = torch.ones(8) + sd["audio_proj_in.weight"] = torch.randn(32, 4) + self._model().load_weights(sd) + + matched = [m for m in warnings if "matched no constructed parameter" in m] + assert len(matched) == 1 + assert "norm_q" in matched[0] + assert "audio2llm" in matched[0] + + def test_missing_root_parameter_raises(self): + """Nano-family audio checkpoints carry a root parameter + (audio_modality_embed); its absence must fail, not stay random.""" + cfg = _reduced_qwen3_config( + sound_gen=True, sound_dim=4, sound_latent_fps=25, temporal_compression_factor_sound=1 + ) + sd = _synthetic_state_dict(cfg) + model_config = DiffusionModelConfig(pretrained_config=cfg) + model_config.attention.backend = "VANILLA" + + Cosmos3VFMTransformer(model_config).load_weights(dict(sd)) + + del sd["audio_modality_embed"] + model_config = DiffusionModelConfig( + pretrained_config=_reduced_qwen3_config( + sound_gen=True, + sound_dim=4, + sound_latent_fps=25, + temporal_compression_factor_sound=1, + ) + ) + model_config.attention.backend = "VANILLA" + with pytest.raises(ValueError, match="audio_modality_embed"): + Cosmos3VFMTransformer(model_config).load_weights(sd) + + def test_qwen3_synthetic_checkpoint_loads(self): + """Nano-loading regression: the gated recipe (gate/up/down + und QK + norms, no k_norm_und_for_gen) loads cleanly through the same strict + coverage path.""" + cfg = _reduced_qwen3_config() + model_config = DiffusionModelConfig(pretrained_config=cfg) + model_config.attention.backend = "VANILLA" + model = Cosmos3VFMTransformer(model_config) + model.load_weights(_synthetic_state_dict(cfg)) + attn = model.language_model.layers[0].self_attn + assert attn.norm_q is not None + assert attn.k_norm_und_for_gen is None + + +def _bare_pipeline(family: str) -> Cosmos3OmniMoTPipeline: + pipeline = object.__new__(Cosmos3OmniMoTPipeline) + pipeline.family = family + pipeline.sampling = Cosmos3SamplingPolicy() + pipeline.audio_gen = False + pipeline.action_gen = False + pipeline.has_action_weights = family == NEMOTRON_DENSE_RECIPE.name + pipeline.use_native_flow_schedule = family == NEMOTRON_DENSE_RECIPE.name + return pipeline + + +class TestEdgeDefaults: + def test_generation_defaults_matrix(self): + edge_video = COSMOS3_GENERATION_DEFAULTS[(NEMOTRON_DENSE_RECIPE.name, "video")] + assert edge_video is COSMOS3_EDGE_VIDEO_PARAMS + assert (edge_video["height"], edge_video["width"]) == (480, 832) + 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 + + edge_t2i = COSMOS3_GENERATION_DEFAULTS[(NEMOTRON_DENSE_RECIPE.name, "image")] + assert edge_t2i is COSMOS3_EDGE_T2I_PARAMS + assert (edge_t2i["height"], edge_t2i["width"]) == (640, 640) + assert edge_t2i["guidance_scale"] == 4.0 + assert edge_t2i["guidance_interval"] is None + + assert COSMOS3_GENERATION_DEFAULTS[(QWEN3_RECIPE.name, "video")] is COSMOS3_720P_PARAMS + + def test_executor_defaults_are_edge_shaped(self): + params = _bare_pipeline(NEMOTRON_DENSE_RECIPE.name).default_generation_params + assert params["num_frames"] == 121 + assert params["max_sequence_length"] == 4096 + assert "flow_shift" not in params + for unresolved in ("height", "width", "num_inference_steps", "guidance_scale"): + assert params[unresolved] is None + + def test_warmup_shape_per_family(self): + edge = _bare_pipeline(NEMOTRON_DENSE_RECIPE.name) + assert edge.default_warmup_resolutions == [(480, 832)] + assert edge.default_warmup_num_frames == [121] + + qwen3 = _bare_pipeline(QWEN3_RECIPE.name) + assert qwen3.default_warmup_resolutions == [(720, 1280)] + assert qwen3.default_warmup_num_frames == [189] + + def test_hf_id_registered(self): + assert "nvidia/Cosmos3-Edge" in PIPELINE_REGISTRY["Cosmos3OmniMoTPipeline"].hf_ids + + 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, + } + # 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 + ) + 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 + + def test_warmup_runs_at_family_guidance(self, monkeypatch): + pipeline = _bare_pipeline(NEMOTRON_DENSE_RECIPE.name) + captured = {} + + def fake_forward(**kwargs): + captured.update(kwargs) + + monkeypatch.setattr(pipeline, "forward", fake_forward, raising=False) + pipeline._run_warmup(height=480, width=832, num_frames=121, steps=2) + 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 + together: the three facts come from different checkpoint files, and a + mismatch (e.g. a stale conversion missing use_native_flow_schedule) would + otherwise sample the wrong trajectory silently.""" + + BASE = Cosmos3SamplingPolicy() + DISTILLED = Cosmos3SamplingPolicy(fixed_sigmas=(1.0, 0.5)) + + @pytest.mark.parametrize( + "family,native,sampling,error", + [ + (QWEN3_RECIPE.name, False, BASE, None), + (QWEN3_RECIPE.name, False, DISTILLED, None), + (NEMOTRON_DENSE_RECIPE.name, True, BASE, None), + (NEMOTRON_DENSE_RECIPE.name, False, BASE, "use_native_flow_schedule"), + (QWEN3_RECIPE.name, True, BASE, "use_native_flow_schedule"), + (NEMOTRON_DENSE_RECIPE.name, True, DISTILLED, "istilled"), + (NEMOTRON_DENSE_RECIPE.name, False, DISTILLED, "istilled"), + ], + ) + def test_startup_matrix(self, family, native, sampling, error): + from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import ( + _validate_sampling_recipe, + ) + + if error is None: + _validate_sampling_recipe(family, native, sampling) + else: + with pytest.raises(ValueError, match=error): + _validate_sampling_recipe(family, native, sampling) + + def test_component_loading_rejects_edge_without_native_flag(self, tmp_path): + """End to end through load_standard_components: a UniPC Edge + checkpoint whose model_index omits the flag must fail at load.""" + import json + + scheduler_dir = tmp_path / "scheduler" + scheduler_dir.mkdir() + (scheduler_dir / "scheduler_config.json").write_text( + json.dumps( + { + "_class_name": "UniPCMultistepScheduler", + "num_train_timesteps": 1000, + "flow_shift": 1.0, + "prediction_type": "flow_prediction", + "use_flow_sigmas": True, + "use_karras_sigmas": True, + "solver_order": 2, + } + ) + ) + (tmp_path / "model_index.json").write_text( + json.dumps({"_class_name": "Cosmos3OmniPipeline"}) + ) + + 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 + return pipeline + + skip = ["text_tokenizer", "tokenizer", "vae", "sound_tokenizer"] + with pytest.raises(ValueError, match="use_native_flow_schedule"): + fresh_pipeline().load_standard_components(str(tmp_path), torch.device("cpu"), skip) + + (tmp_path / "model_index.json").write_text( + json.dumps({"_class_name": "Cosmos3OmniPipeline", "use_native_flow_schedule": True}) + ) + pipeline = fresh_pipeline() + pipeline.load_standard_components(str(tmp_path), torch.device("cpu"), skip) + assert pipeline.sampling.native_flow_schedule is True + assert pipeline.mode_schedulers["video"].config.use_karras_sigmas is False + assert float(pipeline.mode_schedulers["video"].config.flow_shift) == 3.0 + + +class TestEnvelopeAdvisory: + def _warnings(self, monkeypatch): + records = [] + monkeypatch.setattr(pipeline_module.logger, "warning", records.append) + return records + + def _advise(self, pipeline, **overrides): + kwargs = dict( + is_t2i=False, + height=480, + width=832, + num_frames=121, + frame_rate=24.0, + max_sequence_length=4096, + ) + kwargs.update(overrides) + pipeline._log_envelope_advisory(**kwargs) + + def test_in_envelope_is_silent(self, monkeypatch): + records = self._warnings(monkeypatch) + self._advise(_bare_pipeline(NEMOTRON_DENSE_RECIPE.name)) + assert records == [] + + def test_out_of_envelope_logs_once(self, monkeypatch): + records = self._warnings(monkeypatch) + self._advise(_bare_pipeline(NEMOTRON_DENSE_RECIPE.name), num_frames=25) + assert len(records) == 1 + assert "num_frames=25" in records[0] + + def test_family_without_envelope_never_logs(self, monkeypatch): + records = self._warnings(monkeypatch) + self._advise(_bare_pipeline(QWEN3_RECIPE.name), num_frames=25, height=13, width=17) + assert records == [] + + +class TestDiffusersParity: + """Per-step velocity parity against diffusers main (first release with the + Edge classes). Runs in a subprocess because diffusers main cannot be + imported next to the pinned diffusers; gated on DIFFUSERS_MAIN_PATH.""" + + def test_per_step_velocity_parity(self): + import re + import subprocess + import sys + + diffusers_main = os.environ.get("DIFFUSERS_MAIN_PATH") + if not diffusers_main: + pytest.skip("Set DIFFUSERS_MAIN_PATH to a diffusers checkout with Edge support") + checkpoint = _require_edge_checkpoint() + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + + script = Path(__file__).parent / "cosmos3_edge_diffusers_parity.py" + result = subprocess.run( + [sys.executable, str(script), checkpoint], + env={**os.environ, "DIFFUSERS_MAIN_PATH": diffusers_main}, + capture_output=True, + text=True, + timeout=900, + ) + assert result.returncode == 0, result.stderr[-2000:] + rels = [float(m) for m in re.findall(r"rel=([0-9.]+)", result.stdout)] + assert len(rels) == 2, result.stdout + # bf16 accumulation band across 28 layers with differing attention + # backends; observed 0.007 / 0.016 on B200. + assert all(rel < 0.05 for rel in rels), result.stdout + + +# Recorded from diffusers main 2919c5096 (`Cosmos3OmniPipeline.tokenize_prompt` +# on the real Edge checkpoint): prompt "A red cube sits on a wooden table.", +# num_frames=93, height=480, width=832, fps=24.0, system prompt off +# (checkpoint default), duration + resolution templates on. Cond branch only: +# the negative branch intentionally mirrors cosmos-framework's keep-metadata +# templates, not diffusers' inverse templates. +DIFFUSERS_COND_TOKEN_GOLDEN = [ + 1010, + 10, + 25708, + 1010, + 11, + 1010, + 10, + 3263, + 1010, + 1065, + 4804, + 50061, + 53048, + 1408, + 1261, + 32656, + 4234, + 1046, + 1531, + 7476, + 1395, + 1032, + 1051, + 1046, + 1057, + 12900, + 2730, + 1321, + 1395, + 1307, + 1032, + 1050, + 1052, + 1439, + 8148, + 1046, + 2409, + 7476, + 1395, + 1307, + 1032, + 1052, + 1056, + 1048, + 1120, + 1056, + 1051, + 1050, + 9617, + 1046, + 11, + 1010, + 10, + 1503, + 19464, + 1010, + 12, + 1010, + 11, + 20, +] + +# Uncond (CFG) branch pin: TRT-LLM deliberately mirrors cosmos-framework's +# keep-metadata negative-prompt semantics (same duration/resolution templates +# as the positive branch), which diverges from diffusers' inverse templates — +# so this is a self-golden recorded from this code path on the real Edge +# tokenizer (empty negative prompt, num_frames=93, height=480, width=832, +# fps=24.0, system prompt off). Not yet cross-checked against +# cosmos-framework's own tokenization. +UNCOND_TOKEN_GOLDEN = [ + 1010, + 10, + 25708, + 1010, + 11, + 1010, + 10, + 3263, + 1010, + 1784, + 7476, + 1395, + 1032, + 1051, + 1046, + 1057, + 12900, + 2730, + 1321, + 1395, + 1307, + 1032, + 1050, + 1052, + 1439, + 8148, + 1046, + 2409, + 7476, + 1395, + 1307, + 1032, + 1052, + 1056, + 1048, + 1120, + 1056, + 1051, + 1050, + 9617, + 1046, + 11, + 1010, + 10, + 1503, + 19464, + 1010, + 12, + 1010, + 11, + 20, +] + + +class TestEdgeCheckpoint: + """Gated on the real Cosmos3-Edge checkpoint.""" + + def test_tokenizer_specials_and_chat_template(self): + from transformers import AutoTokenizer + + checkpoint = _require_edge_checkpoint() + tokenizer = AutoTokenizer.from_pretrained(checkpoint, subfolder="text_tokenizer") + assert tokenizer.eos_token_id == 11 + assert tokenizer.pad_token_id == 11 + assert tokenizer.convert_tokens_to_ids("<|vision_start|>") == 20 + ids = tokenizer.apply_chat_template( + [{"role": "user", "content": "hello"}], + tokenize=True, + add_generation_prompt=True, + return_dict=False, + ) + assert isinstance(ids, list) and len(ids) > 0 + + def test_tokenization_matches_diffusers_golden(self): + from transformers import AutoTokenizer + + from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import ( + COSMOS3_DEFAULT_RESOLUTION_TEMPLATE, + COSMOS3_DURATION_TEMPLATE, + ) + + checkpoint = _require_edge_checkpoint() + + class _CpuPipeline(Cosmos3OmniMoTPipeline): + device = property(lambda self: torch.device("cpu")) + + pipeline = object.__new__(_CpuPipeline) + pipeline.tokenizer = AutoTokenizer.from_pretrained(checkpoint, subfolder="text_tokenizer") + text = pipeline._format_prompt_with_metadata( + "A red cube sits on a wooden table.", + height=480, + width=832, + num_frames=93, + frame_rate=24.0, + duration_template=COSMOS3_DURATION_TEMPLATE, + resolution_template=COSMOS3_DEFAULT_RESOLUTION_TEMPLATE, + ) + ids, mask = pipeline._tokenize_prompt(text, 128, use_system_prompt=False) + golden_length = len(DIFFUSERS_COND_TOKEN_GOLDEN) + assert ids[0, :golden_length].tolist() == DIFFUSERS_COND_TOKEN_GOLDEN + assert int(mask.sum()) == golden_length + + uncond_text = pipeline._format_prompt_with_metadata( + "", + height=480, + width=832, + num_frames=93, + frame_rate=24.0, + duration_template=COSMOS3_DURATION_TEMPLATE, + resolution_template=COSMOS3_DEFAULT_RESOLUTION_TEMPLATE, + ) + uncond_ids, uncond_mask = pipeline._tokenize_prompt( + uncond_text, 128, use_system_prompt=False + ) + uncond_length = len(UNCOND_TOKEN_GOLDEN) + assert uncond_ids[0, :uncond_length].tolist() == UNCOND_TOKEN_GOLDEN + assert int(uncond_mask.sum()) == uncond_length + + def test_model_index_detection(self): + from tensorrt_llm._torch.visual_gen.pipeline_registry import AutoPipeline + + checkpoint = _require_edge_checkpoint() + assert AutoPipeline._detect_from_checkpoint(checkpoint) == "Cosmos3OmniMoTPipeline" + + @pytest.fixture(scope="class") + def edge_pipeline(self): + checkpoint = _require_edge_checkpoint() + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + args = VisualGenArgs( + model=checkpoint, + torch_compile_config=TorchCompileConfig(enable=False), + ) + pipeline = PipelineLoader(args).load( + skip_warmup=True, + skip_components=[PipelineComponent.SOUND_TOKENIZER], + ) + yield pipeline + del pipeline + gc.collect() + torch.cuda.empty_cache() + + def test_recipe_and_scheduler_wiring(self, edge_pipeline): + assert edge_pipeline.family == NEMOTRON_DENSE_RECIPE.name + 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 + for mode in ("video", "image"): + scheduler = edge_pipeline.mode_schedulers[mode] + assert float(scheduler.config.flow_shift) == 3.0 + assert scheduler.config.use_karras_sigmas is False + + def test_load_weights_and_forward(self, edge_pipeline): + transformer = edge_pipeline.transformer + channels = transformer.latent_channel_size + latents = torch.randn(1, channels, 1, 16, 16, dtype=torch.bfloat16, device=DEVICE) + timestep = torch.tensor([999.0], device=DEVICE) + text_ids = torch.randint(0, 1000, (1, 32), device=DEVICE) + text_mask = torch.ones(1, 32, dtype=torch.long, device=DEVICE) + + transformer.reset_cache() + with torch.inference_mode(): + out = transformer( + hidden_states=latents, + timestep=timestep / 1000.0, + raw_timestep=timestep, + text_ids=text_ids, + text_mask=text_mask, + video_shape=(1, 16, 16), + fps=24.0, + ) + assert out.video.shape == latents.shape + assert torch.isfinite(out.video.float()).all() + + def test_direct_forward_resolves_edge_defaults(self, edge_pipeline, monkeypatch): + """A direct forward() call with unset numerics must reach the + generation path with Edge-table values (denoise and decode stubbed).""" + captured = {} + + def fake_denoise(**kwargs): + captured["latents"] = kwargs["latents"] + captured["guidance_scale"] = kwargs["guidance_scale"] + captured["scheduler"] = kwargs["scheduler"] + return kwargs["latents"] + + def fake_decode(latents, decode_fn, **kwargs): + captured["decoded"] = True + return torch.zeros(1, 2, 8, 8, 3, dtype=torch.uint8) + + monkeypatch.setattr(edge_pipeline, "denoise", fake_denoise, raising=False) + monkeypatch.setattr(edge_pipeline, "decode_latents", fake_decode, raising=False) + + edge_pipeline.forward(prompt="warmup-shaped request", seed=0, use_guardrails=False) + + 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.mode_schedulers["video"] + assert captured["scheduler"].num_inference_steps == 50 + assert captured["decoded"] is True + + def test_t2v_sanity_generation(self, edge_pipeline): + out = edge_pipeline.forward( + prompt="A red ball rolls across a wooden floor.", + seed=0, + height=192, + width=320, + num_frames=9, + num_inference_steps=2, + guidance_scale=5.0, + use_guardrails=False, + ) + video = out.video + assert video is not None + assert tuple(video.shape) == (1, 9, 192, 320, 3) + assert video.float().std() > 1.0, "generated video is (near-)constant" + + def test_i2v_sanity_generation(self, edge_pipeline): + from PIL import Image + + image = Image.new("RGB", (320, 192)) + for x in range(320): + for y in range(0, 192, 4): + image.putpixel((x, y), (x % 256, (2 * y) % 256, 120)) + out = edge_pipeline.forward( + prompt="The scene slowly brightens.", + seed=0, + image=image, + height=192, + width=320, + num_frames=9, + num_inference_steps=2, + guidance_scale=5.0, + use_guardrails=False, + ) + video = out.video + assert video is not None + assert tuple(video.shape) == (1, 9, 192, 320, 3) + assert video.float().std() > 1.0 + + def test_t2i_sanity_generation(self, edge_pipeline): + out = edge_pipeline.forward( + prompt="A ceramic teapot on a table.", + seed=0, + height=256, + width=256, + num_inference_steps=2, + guidance_scale=4.0, + use_guardrails=False, + output_type="image", + ) + image = out.image + assert image is not None + assert tuple(image.shape) == (1, 256, 256, 3) + assert image.float().std() > 1.0 diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index c0329b71a060..223777619eba 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -29,9 +29,6 @@ import os from pathlib import Path -os.environ["TLLM_DISABLE_MPI"] = "1" -os.environ["TRTLLM_DISABLE_COSMOS3_GUARDRAILS"] = "1" - import PIL.Image import pytest import torch @@ -46,13 +43,7 @@ from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader from tensorrt_llm.visual_gen.args import TorchCompileConfig, VisualGenArgs -pytestmark = pytest.mark.cosmos3 - - -@pytest.fixture(autouse=True, scope="module") -def _cleanup_mpi_env(): - yield - os.environ.pop("TLLM_DISABLE_MPI", None) +pytestmark = [pytest.mark.cosmos3, pytest.mark.usefixtures("disable_cosmos3_guardrails")] @pytest.fixture(autouse=True) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py index 76e6c1a01a8c..ef9f70a26275 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py @@ -20,26 +20,22 @@ import gc import os from pathlib import Path - -os.environ["TLLM_DISABLE_MPI"] = "1" -os.environ["TRTLLM_DISABLE_COSMOS3_GUARDRAILS"] = "1" +from types import SimpleNamespace import pytest import torch 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.transformer_cosmos3 import Cosmos3VFMTransformer +from tensorrt_llm._torch.visual_gen.models.cosmos3.transformer_cosmos3 import ( + PRETRAINED_CONFIG_COMPAT_DEFAULTS, + Cosmos3VFMTransformer, + 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 -pytestmark = pytest.mark.cosmos3 - - -@pytest.fixture(autouse=True, scope="module") -def _cleanup_mpi_env(): - yield - os.environ.pop("TLLM_DISABLE_MPI", None) +pytestmark = [pytest.mark.cosmos3, pytest.mark.usefixtures("disable_cosmos3_guardrails")] @pytest.fixture(autouse=True) @@ -468,3 +464,93 @@ def test_load_fp8_quantization(self, quant_algo: str): del pipeline gc.collect() torch.cuda.empty_cache() + + +# --- CPU-only coverage: checkpoint config schema compatibility --- + + +class TestConfigCompatDefaults: + """Newer diffusers conversions omit fields older ones carried explicitly.""" + + 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 + + def test_old_schema_untouched(self): + # Every field deliberately differs from its compat default, so an + # overwrite of any one of them fails its assertion. + config = SimpleNamespace( + position_embedding_type="rope_3d", + max_position_embeddings=12345, + temporal_compression_factor_sound=7, + ) + apply_pretrained_config_compat_defaults(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) + snapshot = vars(config).copy() + apply_pretrained_config_compat_defaults(config) + assert vars(config) == snapshot + + +class TestI2V4StepConfigShape: + """The Image2Video-4Step conversion drops the audio/action towers + (``sound_dim: null``, no ``action_*`` keys) and carries newer schema + fields (``qk_norm_for_text``, ``hidden_act``, nested ``rope_theta``). + The transformer must construct from that exact key set. CPU-only with + shrunk dimensions; the real 64B shape is covered by the checkpoint + integration test.""" + + def _reduced_i2v_config(self) -> SimpleNamespace: + # Key set mirrors the checkpoint's transformer/config.json verbatim; + # only the sizes are reduced (head_dim 8 -> mrope_section sums to 4). + return SimpleNamespace( + attention_bias=False, + attention_dropout=0.0, + base_fps=16, + enable_fps_modulation=True, + head_dim=8, + hidden_act="silu", + hidden_size=32, + intermediate_size=64, + latent_channel=4, + latent_patch_size=2, + num_attention_heads=4, + num_hidden_layers=2, + num_key_value_heads=2, + patch_latent_dim=16, + qk_norm_for_text=True, + rms_norm_eps=1e-6, + rope_axes_dim=[2, 1, 1], + rope_scaling={ + "mrope_interleaved": True, + "mrope_section": [2, 1, 1], + "rope_theta": 5000000, + "rope_type": "default", + }, + rope_theta=5000000, + sound_dim=None, + sound_gen=False, + sound_latent_fps=25, + timestep_scale=0.001, + unified_3d_mrope_reset_spatial_ids=True, + unified_3d_mrope_temporal_modality_margin=15000, + vocab_size=64, + ) + + def test_constructs_without_audio_or_action_towers(self): + model_config = DiffusionModelConfig(pretrained_config=self._reduced_i2v_config()) + model = Cosmos3VFMTransformer(model_config) + + assert model.audio_gen is False + assert model.has_action_weights is False + assert not hasattr(model, "audio2llm") + assert not hasattr(model, "audio_modality_embed") + assert model.base_fps == 16 + assert len(model.gen_layers) == 2