[None][feat] Support the DMD2-distilled Cosmos3 4-step image-to-video checkpoint - #16690
Conversation
|
/bot run --disable-fail-fast |
8c8e826 to
a5370b6
Compare
[None][feat] Support the DMD2-distilled Cosmos3 4-step image-to-video checkpoint|
/bot run --disable-fail-fast |
|
PR_Github #62299 [ run ] triggered by Bot. Commit: |
a5370b6 to
62de5be
Compare
|
/bot run --disable-fail-fast |
|
/bot run --disable-fail-fast |
|
PR_Github #62663 [ run ] triggered by Bot. Commit: |
|
PR_Github #62663 [ run ] completed with state |
6f048f6 to
612045c
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #62715 [ run ] triggered by Bot. Commit: |
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
tests/integration/defs/examples/visual_gen/test_visual_gen.py (1)
2234-2235: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueReuse the existing prompt constant instead of duplicating the literal.
COSMOS3_I2V_4STEP_LPIPS_PROMPT(lines 136-138) already holds this exact string; the smoke test duplicates it as a literal, so a future prompt update to the LPIPS constant would silently leave the smoke test using a stale prompt.♻️ Proposed fix
"--prompt", - "The orange sphere slowly rises while the camera pans right across the scene", + COSMOS3_I2V_4STEP_LPIPS_PROMPT, "--image_path",🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/integration/defs/examples/visual_gen/test_visual_gen.py` around lines 2234 - 2235, Replace the duplicated prompt literal in the smoke test’s “--prompt” argument with the existing COSMOS3_I2V_4STEP_LPIPS_PROMPT constant, preserving the current prompt value while ensuring future updates are shared.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py`:
- Around line 560-582: Update _conditioning_anchor_post_step to return
Callable[[torch.Tensor], torch.Tensor] | None, and import Callable from typing
alongside the existing typing imports.
In `@tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py`:
- Around line 478-494: Update test_image_conditioning_passes_validation to
assert the expected downstream AttributeError specifically references scheduler,
rather than accepting any AttributeError; preserve the existing pipeline.forward
invocation and ensure validation regressions cannot satisfy the test.
---
Nitpick comments:
In `@tests/integration/defs/examples/visual_gen/test_visual_gen.py`:
- Around line 2234-2235: Replace the duplicated prompt literal in the smoke
test’s “--prompt” argument with the existing COSMOS3_I2V_4STEP_LPIPS_PROMPT
constant, preserving the current prompt value while ensuring future updates are
shared.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 47021f13-ab4f-400a-9e18-4c7512b51316
⛔ Files ignored due to path filters (1)
tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zipis excluded by!**/*.zip
📒 Files selected for processing (10)
docs/source/models/supported-models.mddocs/source/models/visual-generation.mdexamples/visual_gen/models/cosmos3/README.mdexamples/visual_gen/models/cosmos3/cosmos3.pytensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.pytests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_i2v_4step_lpips_golden_video.jsontests/integration/defs/examples/visual_gen/test_visual_gen.pytests/integration/test_lists/test-db/l0_b200.ymltests/unittest/_torch/visual_gen/test_cosmos3_distilled.pytests/unittest/_torch/visual_gen/test_cosmos3_transformer.py
🚧 Files skipped from review as they are similar to previous changes (4)
- docs/source/models/supported-models.md
- examples/visual_gen/models/cosmos3/README.md
- tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py
- docs/source/models/visual-generation.md
|
PR_Github #62715 [ run ] completed with state
|
… 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 <ishovkun@nvidia.com>
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 <ishovkun@nvidia.com>
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 <ishovkun@nvidia.com>
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 <ishovkun@nvidia.com>
612045c to
0d28357
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #62800 [ run ] triggered by Bot. Commit: |
There was a problem hiding this comment.
♻️ Duplicate comments (1)
tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py (1)
468-495: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winStrengthen the
AttributeErrorassertion to confirm validation actually passed.
pytest.raises(AttributeError)without amatchwill also pass if validation regresses and throws anAttributeErrorfor an unrelated reason (e.g., a missing attribute earlier inforward()), silently hiding the exact regression this test is meant to catch. This mirrors an unresolved past review comment on this same test (previously at lines 478‑494) that requested matching onschedulerto prove execution actually reached the bare pipeline's downstream failure.🐛 Proposed fix
pipeline = _bare_pipeline(sampling=policy_factory()) - with pytest.raises(AttributeError): + with pytest.raises(AttributeError, match=r"\bscheduler\b"): pipeline.forward(🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py` around lines 468 - 495, Update test_image_conditioning_passes_validation to require the expected downstream AttributeError message by adding a match for “scheduler” to pytest.raises. Keep the existing parameterization and pipeline.forward invocation unchanged so the assertion proves validation passed before the bare pipeline failure.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Duplicate comments:
In `@tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py`:
- Around line 468-495: Update test_image_conditioning_passes_validation to
require the expected downstream AttributeError message by adding a match for
“scheduler” to pytest.raises. Keep the existing parameterization and
pipeline.forward invocation unchanged so the assertion proves validation passed
before the bare pipeline failure.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: b89a94bd-b04d-45ba-b018-f4db347bf529
⛔ Files ignored due to path filters (1)
tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zipis excluded by!**/*.zip
📒 Files selected for processing (10)
docs/source/models/supported-models.mddocs/source/models/visual-generation.mdexamples/visual_gen/models/cosmos3/README.mdexamples/visual_gen/models/cosmos3/cosmos3.pytensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.pytests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_i2v_4step_lpips_golden_video.jsontests/integration/defs/examples/visual_gen/test_visual_gen.pytests/integration/test_lists/test-db/l0_b200.ymltests/unittest/_torch/visual_gen/test_cosmos3_distilled.pytests/unittest/_torch/visual_gen/test_cosmos3_transformer.py
🚧 Files skipped from review as they are similar to previous changes (2)
- docs/source/models/visual-generation.md
- docs/source/models/supported-models.md
|
PR_Github #62800 [ run ] completed with state |
…nd 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 <ishovkun@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #62965 [ run ] triggered by Bot. Commit: |
|
PR_Github #62965 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63028 [ run ] triggered by Bot. Commit: |
|
PR_Github #63028 [ run ] completed with state |
mikeiovine
left a comment
There was a problem hiding this comment.
Rubber stamp for runtime reviewers, did not review as is visualgen-related
|
✅ LFS objects already in storage (1 file) — no sync needed. These LFS-tracked files are already present in this repository's LFS storage:
|
Brings in the DMD2-distilled Cosmos3 4-step I2V checkpoint (NVIDIA#16690), which lands in the same call paths as V2V. Non-obvious resolutions: - forward(): drop the guard rejecting image conditioning on distilled checkpoints. NVIDIA#16690 implements exactly that, so the guard would have rejected the feature it just merged. - denoise(): both sides added a post_step_fn argument to the same call and git took both, which does not compile. V2V pins the conditioning latents and distilled I2V re-anchors the conditioning frame; a request carries an image or a video, never both, so one selects the other. - use_system_prompt: upstream made it checkpoint-declared, this branch made it mode-dependent. Unset now resolves to "V2V uses it, otherwise the checkpoint decides". The extra-param spec keeps its None default because the executor materializes spec defaults into every request -- publishing the checkpoint's boolean there would destroy "unset" before forward() could see it, silently dropping the system prompt from served V2V. The checkpoint value stays discoverable as default_use_system_prompt. - T2I + audio: keep upstream's force-disable. This branch raised instead, left over from the action work already removed from this PR. - The LPIPS golden archive is unioned: ours carries the V2V reference and upstream's the I2V-4Step one, and taking either side whole would have deleted the other PR's baseline. Also scopes failure classification back to Cosmos3. classify_worker_error declared every ValueError from every pipeline a client error, changing the public failure type for Wan, LTX2, Qwen-Image and Flux -- 130 raise sites this PR never touched. It is now an opt-in BasePipeline hook, classify_request_failure, unclassified by default. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Dev Engineer Review
nvidia/Cosmos3-Super-Image2Video-4Step.default_use_system_promptfrom checkpoint metadata with three-state CLI handling.QA Engineer Review
Test-list changes
Modified
tests/integration/test_lists/test-db/l0_b200.yml:Test-code coverage
Added or updated:
The two integration tests are covered by
l0_b200.yml. Unit tests run through the unit-test suite and are not listed intest-db/.Verdict: needs follow-up
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: diffusersCosmos3DistilledModularPipeline(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_fndenoise 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. This lifts the temporary rejection of image-conditioned requests on distilled checkpoints that #16563 added to hold the line until this landed. The final pre-decode re-injection also drops its full-latentclone()in favor of an in-place one-frame write.Checkpoint-declared system-prompt default. Distilled conversions declare
default_use_system_promptinmodel_index.json(this checkpoint:true, matching diffusers' distilled blocks); TRT-LLM previously hardcodedFalse. The declaration is read at load, reflected inextra_param_specs(so serve clients anddefault_paramsprefill see the truth), and used asinfer()'s fallback for an unset key. Checkpoints without the declaration keep the historicalFalse. The example CLI flag becomes three-state (--use_system_prompt/--no-use_system_prompt/ unset = checkpoint default).Weight-presence guard.
enable_audio=Trueon 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 against 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
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 / explicitFalse); audio weight-presence guard (incl. T2I force-disable ordering); registry IDs.tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py): reduced I2V-4Step config shape constructs without audio/action towers (CPU-only).test_cosmos3_i2v_4step_example, post-merge B200): the documented example invocation with a deterministic conditioning image at the deployed 720p × 189-frame shape; asserts a non-empty MP4.test_cosmos3_i2v_4step_lpips_against_golden, B200): 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 ingolden/visual_gen_lpips/cosmos3_i2v_4step_lpips_golden_video.json; measured 0.0563 at golden creation, 0.0588 on validation.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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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.