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

[ishovkun]

@coderabbitai summary

Description

Adds inference support for nvidia/Cosmos3-Super-Image2Video-4Step, the DMD2-distilled image-to-video counterpart of the Text2Image-4Step checkpoint enabled in #16563. Both checkpoints ship the identical distilled sampling recipe (FlowMatchEulerDiscreteScheduler, fixed 4-sigma schedule, classifier-free guidance baked into the weights), so this PR builds on the sampling policy introduced there and adds what distilled I2V specifically needs. Reference implementation: diffusers Cosmos3DistilledModularPipeline (huggingface/diffusers#14177).

Per-step conditioning re-anchor (the core change). The distilled FlowMatchEuler step is stochastic: it re-noises every latent position on each step, including the conditioning frame the velocity mask holds still — so from step 2 on the model attends to a corrupted "clean" frame. The clean image latent is now written back after every scheduler step via the existing post_step_fn denoise hook, gated on distilled sampling with image conditioning (matching the diffusers distilled loop). Base UniPC sampling is byte-identical: deterministic steps never move a zero-velocity frame. The final pre-decode re-injection also drops its full-latent clone() in favor of an in-place one-frame write.

Checkpoint-declared system-prompt default. Distilled conversions declare default_use_system_prompt in model_index.json (this checkpoint: true, matching diffusers' distilled blocks); TRT-LLM previously hardcoded False. The declaration is read at load, reflected in extra_param_specs (so serve clients and default_params prefill see the truth), and used 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 = checkpoint default).

Weight-presence guard. enable_audio=True on a checkpoint without an audio tower (sound_gen: false — this checkpoint ships no audio weights) now raises instead of silently returning a silent video. Workflow selection stays request-driven: no mode is rejected based on checkpoint identity.

Registration/config. HF-ID registration and acceptance of the I2V-4Step transformer config shape (sound_dim: null, no action fields, qk_norm_for_text: true, base_fps: 16).

Verified on 1×B200 (183 GB): the documented 720p × 189-frame invocation peaks at ~141 GB, 30.9 s generation (22.0 s denoise, 4 steps). Cross-stack parity vs the diffusers reference with matched noise trajectories: LPIPS 0.056.

Out of scope: video-conditioned (V2V) distilled generation — the V2V path is a separate PR; the re-anchor here is mask-generic, so distilled V2V becomes trivial once V2V merges. Note: the merged diffusers reference draws its per-step SDE noise from the global RNG; the LPIPS golden was generated with a one-line generator pass-through patch, recorded in the golden's provenance manifest.

Test Coverage

• Unit (tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py): re-anchor proven at the transformer boundary (a perturbing scheduler emulates SDE re-noising; every recorded forward input keeps the clean frame, plus a control test showing the drift without the anchor); forward wiring (anchor + seeded step kwargs exactly when distilled + image-conditioned); base/UniPC and distilled-T2V paths unchanged; system-prompt default resolution (declared / absent / explicit False); audio weight-presence guard (incl. T2I force-disable ordering); registry IDs.
• Unit (tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py): reduced I2V-4Step config shape constructs without audio/action towers (CPU-only).
• Integration smoke (test_cosmos3_i2v_4step_example, l0_b200.yml): the documented example invocation with a deterministic conditioning image at the deployed 720p × 189-frame shape; asserts a non-empty MP4.
• Integration quality gate (test_cosmos3_i2v_4step_lpips_against_golden, l0_b200.yml): LPIPS ≤ 0.10 against a diffusers-produced golden (not a TRT-LLM self-golden), so the gate checks the denoising trajectory against the reference implementation. Golden provenance (diffusers commit, RNG patch, pipeline construction, all generation parameters) is recorded in golden/visual_gen_lpips/cosmos3_i2v_4step_lpips_golden_video.json; measured 0.0563 at golden creation, 0.0588 on validation.
• Note for reviewers: both 4-step checkpoints must be staged in the CI llm-models/ storage, or these tests skip (currently also true for [None][feat] Support the DMD2-distilled Cosmos3-Super-Text2Image-4Step checkpoint #16563's test).

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[ishovkun]

/bot run --disable-fail-fast