[None][feat] Support the DMD2-distilled Cosmos3 4-step image-to-video checkpoint

docs/source/models/supported-models.md

@@ -173,6 +173,7 @@ For full documentation, see the [Visual Generation](./visual-generation.md) page
 | `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) |
 
 ### Feature Matrix
 

docs/source/models/visual-generation.md

@@ -42,6 +42,7 @@ TensorRT-LLM **VisualGen** provides a unified inference stack for diffusion mode
 | `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) |
 
 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.
 

examples/visual_gen/models/cosmos3/README.md

@@ -14,6 +14,7 @@ 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`).
 
 ## Guardrails
 
@@ -80,6 +81,14 @@ python cosmos3.py --model nvidia/Cosmos3-Super-Text2Image-4Step \
     --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
+
 # 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" \

examples/visual_gen/models/cosmos3/cosmos3.py

@@ -136,7 +136,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 +187,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

tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py

@@ -65,6 +65,7 @@
         "nvidia/Cosmos3-Nano",
         "nvidia/Cosmos3-Super",
         "nvidia/Cosmos3-Super-Image2Video",
+        "nvidia/Cosmos3-Super-Image2Video-4Step",
         "nvidia/Cosmos3-Super-Text2Image",
         "nvidia/Cosmos3-Super-Text2Image-4Step",
     ],
@@ -78,6 +79,7 @@ def __init__(self, pipeline_config):
         # Pre-load placeholder; load_standard_components derives the real
         # policy from the checkpoint's scheduler via from_scheduler().
         self.sampling = Cosmos3SamplingPolicy()
+        self.default_use_system_prompt = COSMOS3_EXTRA_SPECS["use_system_prompt"].default
         if getattr(
             primary_pretrained_config,
             "audio_gen",
@@ -108,6 +110,18 @@ 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)
+            )
+
         if self.audio_gen and PipelineComponent.SOUND_TOKENIZER not in skip_components:
             logger.info("Loading audio tokenizer...")
             self.audio_tokenizer = (
@@ -207,7 +221,11 @@ def default_generation_params(self):
 
     @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;
@@ -268,7 +286,8 @@ def resolved(value, field_name):
                 "use_resolution_template",
                 COSMOS3_EXTRA_SPECS["use_resolution_template"].default,
             ),
-            use_system_prompt=extra_params.get("use_system_prompt", False),
+            # None = unset; forward() resolves it to the checkpoint default.
+            use_system_prompt=extra_params.get("use_system_prompt"),
             use_guardrails=extra_params.get("use_guardrails", True),
             enable_audio=extra_params.get("enable_audio", False),
             output_type=output_type,
@@ -539,6 +558,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
     # =========================================================================
@@ -604,7 +647,7 @@ def forward(
         frame_rate: float = COSMOS3_720P_PARAMS["frame_rate"],
         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,
+        use_system_prompt: Optional[bool] = None,
         use_guardrails: bool = COSMOS3_EXTRA_SPECS["use_guardrails"].default,
         enable_audio: bool = COSMOS3_EXTRA_SPECS["enable_audio"].default,
         output_type: str = COSMOS3_EXTRA_SPECS["output_type"].default,
@@ -617,11 +660,18 @@ def forward(
         signature default from an explicit argument, so on distilled
         checkpoints (which fix steps/guidance and reject anything else) direct
         callers must pass checkpoint-valid sampling values.
+
+        ``use_system_prompt=None`` means "unset": it resolves to the
+        checkpoint-declared default, so warmup and other direct callers build
+        the same prompt as served requests.
         """
         pipeline_start = time.time()
         timer = CudaPhaseTimer()
         timer.mark_pre_start()
 
+        if use_system_prompt is None:
+            use_system_prompt = self.default_use_system_prompt
+
         use_guardrails = use_guardrails and not TRTLLM_DISABLE_COSMOS3_GUARDRAILS
 
         # Text-to-image mode: same checkpoint/forward path as T2V, but a single
@@ -634,13 +684,6 @@ def forward(
 
         self.sampling.validate_request(num_inference_steps, guidance_scale)
 
-        if image is not None and self.sampling.is_distilled:
-            raise ValueError(
-                "Image-conditioned generation is not supported on distilled Cosmos3 "
-                "checkpoints yet: the stochastic scheduler re-noises the conditioned "
-                "frame at every step, and this pipeline does not re-anchor it per step."
-            )
-
         guidance_interval = None
         if is_t2i:
             if image is not None:
@@ -660,6 +703,16 @@ def forward(
                 self.scheduler, self.sampling.checkpoint_flow_shift
             )
 
+        # 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]
         batch_size = len(prompt)
@@ -865,6 +918,7 @@ def forward_fn(
             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:
@@ -881,8 +935,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)
 

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."
+}