From db9e1edc93b98ed4fc85a3e4b126cdb0b8427d63 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Tue, 21 Jul 2026 13:57:16 -0700 Subject: [PATCH 1/5] [None][feat] Support the DMD2-distilled Cosmos3 4-step image-to-video checkpoint Register nvidia/Cosmos3-Super-Image2Video-4Step and add the one algorithmic piece distilled I2V needs: the stochastic FlowMatchEuler step re-noises every position each step, so the clean conditioning frame is re-anchored after every scheduler step via the denoise post_step_fn hook (matching the diffusers distilled loop, PR huggingface/diffusers#14177). Base UniPC sampling is unchanged: deterministic steps never move a zero-velocity frame. This lifts the temporary rejection of image-conditioned requests on distilled checkpoints, which was added to hold the line until per-step re-anchoring landed. Also raise on enable_audio=True when the checkpoint ships no audio tower (weight-presence guard, not workflow policy), make the final conditioning re-injection in-place instead of cloning the full latent tensor, and accept the I2V-4Step transformer config shape (sound_dim null, no action fields). Signed-off-by: Igor Shovkun --- .../models/cosmos3/pipeline_cosmos3.py | 47 +++- .../visual_gen/test_cosmos3_distilled.py | 257 +++++++++++++++++- .../visual_gen/test_cosmos3_transformer.py | 57 ++++ 3 files changed, 338 insertions(+), 23 deletions(-) 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 9987a380bedc..3dfe27fbfee2 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/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", ], @@ -539,6 +540,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 # ========================================================================= @@ -634,13 +659,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 +678,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 +893,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 +910,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/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py index 955eec6c29ab..6ce0aa9237ac 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py @@ -464,27 +464,34 @@ def test_forward_rejects_explicit_mismatch(self, bad_kwargs): with pytest.raises(ValueError, match="distilled"): pipeline.forward(prompt="x", seed=0, use_guardrails=False, **bad_kwargs) - def test_distilled_rejects_image_conditioning(self): - """Without per-step re-anchoring, the stochastic scheduler corrupts the - conditioned frame; the request must fail rather than degrade silently.""" - pipeline = _bare_pipeline(sampling=_distilled_policy()) - with pytest.raises(ValueError, match="re-anchor"): + @pytest.mark.parametrize( + "policy_factory, sampling_kwargs", + [ + (_base_policy, {}), + ( + _distilled_policy, + {"num_inference_steps": 4, "guidance_scale": DISTILLED_GUIDANCE_SCALE}, + ), + ], + ) + def test_image_conditioning_passes_validation(self, policy_factory, sampling_kwargs): + """No sampling policy rejects image conditioning by checkpoint kind. + + Distilled checkpoints keep the conditioned frame clean by re-anchoring + it after every scheduler step (see TestDistilledConditioningAnchor) + rather than by refusing the request. forward() therefore gets past + validation for both policies and only fails later on the bare double. + """ + pipeline = _bare_pipeline(sampling=policy_factory()) + with pytest.raises(AttributeError): pipeline.forward( prompt="x", seed=0, use_guardrails=False, image="frame.png", - num_inference_steps=4, - guidance_scale=1.0, + **sampling_kwargs, ) - def test_base_still_accepts_image_conditioning_path(self): - """Base checkpoints must not be caught by the distilled image rejection: - forward proceeds past validation (fails later on the bare test double).""" - pipeline = _bare_pipeline(sampling=_base_policy()) - with pytest.raises(AttributeError): - pipeline.forward(prompt="x", seed=0, use_guardrails=False, image="frame.png") - @pytest.mark.parametrize("bad_output_type", ["imgae", "png", "", "both"]) def test_invalid_output_type_raises(self, bad_output_type): pipeline = _bare_pipeline() @@ -590,6 +597,227 @@ def test_scheduler_step_kwargs_reach_every_step(self): 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 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: @@ -599,3 +827,4 @@ def test_model_index_class_name_dispatches(self, tmp_path): 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_transformer.py b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py index 3ebcb08290da..9293c371f459 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py @@ -505,3 +505,60 @@ def test_idempotent(self): 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.action_gen 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 From 0fc4a8448d8b140b5391cc03e00022121c066e9d Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Tue, 21 Jul 2026 13:57:36 -0700 Subject: [PATCH 2/5] [None][feat] Honor the checkpoint-declared Cosmos3 system-prompt default The distilled conversions declare default_use_system_prompt in model_index.json (the I2V-4Step checkpoint sets it to true, matching the diffusers distilled blocks); TRT-LLM previously hardcoded False. Read the declaration at load, reflect it in extra_param_specs so serve clients and default_params prefill see the truth, and use it as infer()'s fallback for an unset key. Checkpoints without the declaration keep the historical False. The example CLI flag becomes three-state (--use_system_prompt / --no-use_system_prompt / unset): omitting it no longer force-overwrites the checkpoint default with False. Signed-off-by: Igor Shovkun --- examples/visual_gen/models/cosmos3/cosmos3.py | 11 +++- .../models/cosmos3/pipeline_cosmos3.py | 21 ++++++- .../visual_gen/test_cosmos3_distilled.py | 55 +++++++++++++++++++ 3 files changed, 83 insertions(+), 4 deletions(-) diff --git a/examples/visual_gen/models/cosmos3/cosmos3.py b/examples/visual_gen/models/cosmos3/cosmos3.py index de9e9e5010ad..9c1f45be1abd 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3.py +++ b/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 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 3dfe27fbfee2..4ed264541c80 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -79,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", @@ -109,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 = ( @@ -208,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; @@ -269,7 +286,7 @@ 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), + 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, diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py index 6ce0aa9237ac..8ead5da2013a 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py @@ -83,6 +83,7 @@ def _bare_pipeline(**attrs) -> Cosmos3OmniMoTPipeline: audio_gen=False, action_gen=False, sampling=Cosmos3SamplingPolicy(), + default_use_system_prompt=False, ) defaults.update(attrs) for key, value in defaults.items(): @@ -773,6 +774,60 @@ def test_t2v_request_wires_no_anchor(self): 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.""" From 0deeeaebbccdafb01338041f8133d5cf56d0d92d Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Tue, 21 Jul 2026 13:57:54 -0700 Subject: [PATCH 3/5] [None][test] Add I2V-4Step smoke and diffusers-golden LPIPS gates Smoke: run the documented example invocation (deterministic PIL conditioning image, omni-default 720p x 189 shape) and assert a non-empty MP4. Quality gate: unlike the existing TRT-LLM self-goldens, the golden video is produced by the reference implementation (diffusers Cosmos3 distilled modular pipeline, huggingface/diffusers#14177, with its per-step SDE noise made generator-seeded), so the gate checks the denoising trajectory against the reference rather than regression against a past TRT-LLM run. Full provenance (diffusers commit, RNG patch, corrected modular index, generation parameters) is recorded in cosmos3_i2v_4step_lpips_golden_video.json. Threshold 0.10 = 0.0563 measured at golden creation plus headroom for the ~0.04 cross-host kernel drift documented in the harness; validated at 0.0588 on B200. Signed-off-by: Igor Shovkun --- .../cosmos3_i2v_4step_lpips_golden_video.json | 27 +++ .../visual_gen_lpips_golden_media.zip | 4 +- .../examples/visual_gen/test_visual_gen.py | 183 ++++++++++++++++++ .../test_lists/test-db/l0_b200.yml | 2 + 4 files changed, 214 insertions(+), 2 deletions(-) create mode 100644 tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_i2v_4step_lpips_golden_video.json 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 1b85d11a6f67..632d220fe72e 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:de95c0e23c15209a3f00f6955433917a156d057473fe6c5054450f9db38243ab -size 17636502 +oid sha256:2844dd5fcf63e55a98fa903137be192ddc7a912b76e9c32e1dd02524dedd3177 +size 17758973 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 2a729d93c132..af66db194c53 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen.py @@ -132,6 +132,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" @@ -2164,3 +2178,172 @@ def test_cosmos3_t2i_4step_example(_visual_gen_deps, llm_root, llm_venv): ) 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 _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 a0f5de033f09..bc2eaa6b06a4 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -267,6 +267,7 @@ l0_b200: - unittest/llmapi/test_llm_quant.py # 3.5 mins on B200 - unittest/disaggregated/test_openai_server_info.py - 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) - condition: ranges: system_gpu_count: @@ -342,6 +343,7 @@ l0_b200: - examples/visual_gen/test_visual_gen.py::test_qwen_image_layered_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) - 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) ---- From 0d2835772f99fd84e024611a3401cd65e7be225d Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Tue, 21 Jul 2026 13:57:55 -0700 Subject: [PATCH 4/5] [None][doc] Document Cosmos3-Super-Image2Video-4Step usage Model rows in the visual-generation and supported-models tables, and the README invocation: the omni default (720p x 189 frames) is the deployed shape so no dedicated config is needed; steps, guidance, and the system-prompt default come from the checkpoint. Signed-off-by: Igor Shovkun --- docs/source/models/supported-models.md | 1 + docs/source/models/visual-generation.md | 1 + examples/visual_gen/models/cosmos3/README.md | 9 +++++++++ 3 files changed, 11 insertions(+) diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index 195a80d4b4aa..ea06f99ca313 100644 --- a/docs/source/models/supported-models.md +++ b/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 diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index d7d59fdb9562..f8fb202f265e 100644 --- a/docs/source/models/visual-generation.md +++ b/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. diff --git a/examples/visual_gen/models/cosmos3/README.md b/examples/visual_gen/models/cosmos3/README.md index f5d740841dda..4d17cea608c8 100644 --- a/examples/visual_gen/models/cosmos3/README.md +++ b/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" \ From 290d24528fced152cf602cd9aa6c1e08eb629dd9 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 30 Jul 2026 22:13:35 -0700 Subject: [PATCH 5/5] [None][fix] Resolve the Cosmos3 system-prompt default in forward(), and run the distilled unit tests in CI forward() treated use_system_prompt as a plain bool defaulting to the static spec value, so only infer() picked up the checkpoint-declared default: warmup and other direct callers built a different prompt than served requests on a checkpoint that declares default_use_system_prompt=true. forward() now takes None as "unset" and resolves it, and infer() passes the request value straight through, so resolution lives in one place. tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py was never listed in any test-db block, and TRT-LLM CI only runs unit tests through explicit unittest/ bridge entries, so the whole file (sampling policy, conditioning anchor, system-prompt, audio guard) has never been exercised. Register it alongside the other cosmos3 unit tests. Signed-off-by: Igor Shovkun --- .../models/cosmos3/pipeline_cosmos3.py | 12 +- .../test_lists/test-db/l0_b200.yml | 1 + .../visual_gen/test_cosmos3_distilled.py | 109 ++++++++++++++---- 3 files changed, 95 insertions(+), 27 deletions(-) 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 4ed264541c80..f497b09f7c90 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -286,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", self.default_use_system_prompt), + # 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, @@ -646,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, @@ -659,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 diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index bc2eaa6b06a4..3efeff1ba746 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -229,6 +229,7 @@ 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_distilled.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 diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py index 8ead5da2013a..e6664dea569e 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py @@ -698,6 +698,32 @@ def test_without_anchor_the_conditioned_frame_drifts(self): assert torch.all(latent_input[:, :, 0:1] == self.CLEAN + step * 1.5) +def _forward_ready_pipeline(**attrs) -> Cosmos3OmniMoTPipeline: + """A pipeline stubbed just enough for forward() to run end to end.""" + defaults = dict( + 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), + ), + ) + defaults.update(attrs) + pipeline = _bare_pipeline(**defaults) + pipeline._tokenize_prompt = lambda *args, **kwargs: ( + torch.ones(1, 4, dtype=torch.long), + torch.ones(1, 4, dtype=torch.long), + ) + return pipeline + + class TestForwardConditioningWiring: """forward() must hand the denoise loop the anchor exactly when the checkpoint is distilled and the request carries image conditioning.""" @@ -705,26 +731,8 @@ class TestForwardConditioningWiring: 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), - ) + def _wiring_pipeline(self): + pipeline = _forward_ready_pipeline() pipeline._encode_conditioning_video = lambda *args, **kwargs: torch.full( (1, 4, self.T_LAT, self.H_LAT, self.W_LAT), self.CLEAN ) @@ -754,7 +762,7 @@ def _forward(self, pipeline, image): ) def test_i2v_request_wires_anchor_and_seeded_steps(self): - pipeline, captured = self._forward_ready_pipeline() + pipeline, captured = self._wiring_pipeline() self._forward(pipeline, image=torch.zeros(3, 32, 32)) post_step_fn = captured["post_step_fn"] @@ -769,7 +777,7 @@ def test_i2v_request_wires_anchor_and_seeded_steps(self): assert torch.all(captured["latents"][:, :, 0:1] == self.CLEAN) def test_t2v_request_wires_no_anchor(self): - pipeline, captured = self._forward_ready_pipeline() + pipeline, captured = self._wiring_pipeline() self._forward(pipeline, image=None) assert captured["post_step_fn"] is None @@ -816,17 +824,68 @@ def _captured_use_system_prompt(self, pipeline, extra_params): pipeline.infer(_fake_request("video", extra_params=extra_params)) return captured["use_system_prompt"] - def test_infer_unset_key_uses_checkpoint_default(self): + def test_infer_passes_unset_key_through_as_none(self): + """forward() owns the resolution, so infer() must forward "unset" + rather than pre-resolving it — otherwise the two entry points can + drift apart again.""" pipeline = _bare_pipeline(default_use_system_prompt=True) - assert self._captured_use_system_prompt(pipeline, {"output_type": "video"}) is True + assert self._captured_use_system_prompt(pipeline, {"output_type": "video"}) is None - def test_infer_explicit_false_preserved(self): + def test_infer_passes_explicit_false_through(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 + def _tokenized_use_system_prompt(self, pipeline, **forward_kwargs): + """Run forward() far enough to observe what it hands the tokenizer.""" + seen = {} + + class _Stop(Exception): + pass + + def fake_tokenize(prompt, max_sequence_length, use_system_prompt, system_prompt=None): + seen["value"] = use_system_prompt + raise _Stop + + pipeline._tokenize_prompt = fake_tokenize + with pytest.raises(_Stop): + pipeline.forward( + prompt="x", + seed=0, + use_guardrails=False, + num_inference_steps=4, + guidance_scale=DISTILLED_GUIDANCE_SCALE, + **forward_kwargs, + ) + return seen["value"] + + def test_forward_unset_resolves_to_checkpoint_default(self): + """Direct forward() callers (warmup included) must build the same + prompt as served requests.""" + pipeline = _forward_ready_pipeline(default_use_system_prompt=True) + assert self._tokenized_use_system_prompt(pipeline) is True + + def test_forward_explicit_false_overrides_checkpoint_default(self): + pipeline = _forward_ready_pipeline(default_use_system_prompt=True) + assert self._tokenized_use_system_prompt(pipeline, use_system_prompt=False) is False + + def test_forward_default_false_checkpoint_unchanged(self): + pipeline = _forward_ready_pipeline(default_use_system_prompt=False) + assert self._tokenized_use_system_prompt(pipeline) is False + + def test_warmup_leaves_system_prompt_unset(self): + """_run_warmup must not pin the historical False: leaving it unset is + what lets forward() resolve the checkpoint default.""" + pipeline = _bare_pipeline(sampling=_distilled_policy(), default_use_system_prompt=True) + captured = {} + pipeline.forward = lambda **kwargs: captured.update(kwargs) + + pipeline._run_warmup(height=720, width=1280, num_frames=9, steps=4) + + assert captured.get("use_system_prompt") is None + class TestAudioWeightPresenceGuard: """enable_audio=True must fail loudly when the checkpoint ships no audio