[None][feat] Support the DMD2-distilled Cosmos3 4-step image-to-video checkpoint - #16690
Merged
Merged
Conversation
Contributor
Author
|
/bot run --disable-fail-fast |
1 task done
ishovkun
force-pushed
the
cosmos3_i2v_4step
branch
from
July 28, 2026 20:48
8c8e826 to
a5370b6
Compare
ishovkun
changed the title
[None][feat] Support the DMD2-distilled Cosmos3 4-step image-to-video checkpoint
Contributor
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #62299 [ run ] triggered by Bot. Commit: |
ishovkun
force-pushed
the
cosmos3_i2v_4step
branch
from
July 28, 2026 21:10
a5370b6 to
62de5be
Compare
Contributor
Author
|
/bot run --disable-fail-fast |
ishovkun
marked this pull request as ready for review
July 28, 2026 21:24
ishovkun
requested review from
BowenFu,
EmmaQiaoCh,
QiJune,
ZhanruiSunCh,
cascade812,
chang-l,
fredricz-20070104,
luyiyun1021,
o-stoner,
tongyuantongyu and
yufeiwu-nv
July 28, 2026 21:24
Collaborator
|
PR_Github #62663 [ run ] completed with state |
ishovkun
force-pushed
the
cosmos3_i2v_4step
branch
from
July 30, 2026 06:00
6f048f6 to
612045c
Compare
Contributor
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #62715 [ run ] triggered by Bot. Commit: |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
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
Collaborator
|
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>
ishovkun
force-pushed
the
cosmos3_i2v_4step
branch
from
July 30, 2026 16:04
612045c to
0d28357
Compare
Contributor
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #62800 [ run ] triggered by Bot. Commit: |
Contributor
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
Collaborator
|
PR_Github #62800 [ run ] completed with state |
chang-l
reviewed
Jul 30, 2026
…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>
Contributor
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #62965 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #62965 [ run ] completed with state
|
Contributor
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #63028 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #63028 [ run ] completed with state |
mikeiovine
approved these changes
Jul 31, 2026
mikeiovine
left a comment
mikeiovine
left a comment
Collaborator
There was a problem hiding this comment.
Rubber stamp for runtime reviewers, did not review as is visualgen-related
chang-l
approved these changes
Jul 31, 2026
|
✅ LFS objects already in storage (1 file) — no sync needed. These LFS-tracked files are already present in this repository's LFS storage:
|
ishovkun
added a commit
to ishovkun/TensorRT-LLM
that referenced
this pull request
Jul 31, 2026
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>
MaciejBalaNV
pushed a commit
to NVIDIA/cosmos
that referenced
this pull request
Sep 21, 2026
## Summary - add a TensorRT-LLM walkthrough for the two published DMD2-distilled Cosmos3-Super students, `nvidia/Cosmos3-Super-Text2Image-4Step` (text-to-image) and `nvidia/Cosmos3-Super-Image2Video-4Step` (image-to-video), each on a single GPU - document the distilled request contract: the checkpoint's scheduler config locks the schedule to four stochastic (SDE) steps and classifier-free guidance is baked into the weights, so requests omit `num_inference_steps` and `guidance_scale` entirely and the server fills both in; a conflicting value is rejected, not clamped - leave `use_system_prompt` unset in every request so the image-to-video student's checkpoint-declared `default_use_system_prompt: true` applies - add both `trtllm-serve` launch commands to the shared setup guide, and index the notebook from the root README - **separately**: reword two guardrail blocklist false positives in the shared audiovisual negative prompts, without which no image-to-video request can run at all (see below) This is the follow-up #310 deferred ("keep the four-step distilled T2I/I2V examples out of this PR pending separate output-quality follow-up"). The examples match the assets every other backend already uses for its distilled sections — `robot_draping` for text-to-image and `car_driving` for image-to-video — so the cookbook surface is consistent across Cosmos Framework, Diffusers, vLLM-Omni, SGLang, and now TensorRT-LLM. The shipped scope is text-to-image and image-to-video only. These students are task-specialized and no other backend advertises text-to-video, video-to-video, or synchronized audio for them either. ## The guardrail fix, and why it is in this PR The shared audiovisual negative prompts contained two substrings that the `cosmos_guardrail` 0.3.0 `Blocklist` censors. Any request carrying them fails with HTTP 500 and `Text guardrail blocked prompt` **before any denoising runs**: | Text | Censored word | Replacement | |---|---|---| | `color bleeding between elements` | `bleeding` | `color smearing between elements` | | `The scene feels lifeless and sterile` | `lifeless` | `The scene feels inert and sterile` | Neither describes anything the blocklist is meant to catch; both are ordinary image-artifact vocabulary. The replacement wording is not invented here. TensorRT-LLM already ships a corrected copy of this exact prompt as `examples/visual_gen/models/cosmos3/cosmos3_negative_prompt.json`, and that is the default every Cosmos3 request falls back to when it sends no negative prompt of its own. Upstream's copy and the cookbook's differ in **exactly four leaves, and all four are these two substrings** — so the cookbook asset is simply stale. Both files are now identical to the TensorRT-LLM default. Kept as separate commits so they can be split out if preferred. Text-to-image was unaffected because that path sends no negative prompt — which is why only the image-to-video example hit it. Three related issues found while diagnosing this, **not changed here**: 1. The two files are shared by the Diffusers, vLLM-Omni, SGLang and base TensorRT-LLM notebooks, so this also unblocks their image-to-video / text-to-video / video-to-video examples. 2. `assets/prompts/image2video/humanoid_robot.json` is blocked by a *different* rule (partial match, not the word blocklist). That is the prompt the merged `run_with_trt_llm.ipynb` uses for `i2v_nano` and `i2v_super`. Left alone deliberately: rewording a *positive* prompt changes generated content. 3. `transfer/assets/negative_prompt.json` carries the same two substrings. That file belongs to #310, which now rewords both, so it is covered there rather than duplicated here. ## Validation ### GPU inference Both notebook sections were executed **as written** with `jupyter nbconvert --to notebook --execute` against real `trtllm-serve` VisualGen servers — the shipped notebook's own cells, not a re-implementation. TensorRT-LLM at `505a02ce04`, whose `tensorrt_llm/_torch/visual_gen`, `tensorrt_llm/visual_gen` and `examples/visual_gen` trees are byte-identical to `main` (`2dbab44d4f`); the only delta is test files. Hardware: 1 x NVIDIA B300. **Text-to-image** — `Cosmos3-Super-Text2Image-4Step` with `configs/cosmos3-t2i-1gpu.yaml`: ``` Distilled Cosmos3 checkpoint: fixed 4-step schedule [1.0, 0.9375, 0.8333333333333334, 0.625], classifier-free guidance baked in. Running warmup for Cosmos3OmniMoTPipeline: 1 shape(s) [1024x1024x1], 4 steps Denoising: 4 steps, guidance=1.0 POST /v1/videos/generations HTTP/1.1" 200 OK ``` Artifact: 57,617 bytes, decoded with PyAV as H.264 1024x1024, 1 frame. **Image-to-video** — `Cosmos3-Super-Image2Video-4Step`, no config file: ``` Running warmup for Cosmos3OmniMoTPipeline: 1 shape(s) [720x1280x189], 4 steps Denoising: 4 steps, guidance=1.0 POST /v1/videos/generations HTTP/1.1" 200 OK ``` Artifact: 4,890,310 bytes, decoded as H.264 1280x720, 24 fps, 189 frames. The warmup shape confirms the claim that this checkpoint deploys at the default omni shape and needs no config file. **Negative control** (both servers): a request that sets `guidance_scale=6.0` is rejected rather than clamped — ``` HTTP 400: This is a distilled Cosmos3 checkpoint; classifier-free guidance is baked into the weights. guidance_scale must be 1.0 or left unset (got 6.0). ``` **Guardrail fix verified directly** against `cosmos_guardrail` 0.3.0's `Blocklist`: both reworded negative prompts and both positive prompts used by the notebook now report `safe=True`. The trigger words were located by diffing the censored output against the original, not guessed. Re-verified after adopting upstream's wording: `bleeding` and `lifeless` are present in NVIDIA's `blocklist`/`exact_match` word sets (383 and 1430 entries) while `smearing` and `inert` are not. **Visual review**: the text-to-image frame renders the prompt's robot-arm-draping-satin scene coherently. The image-to-video clip follows its `temporal_caption` — coastal dashcam view, rockfall beginning a few seconds in, boulders and dust building until the road is partially blocked — with no flicker or geometry snapping across the 189 frames. ### Local validation - parsed the new notebook as JSON - validated it against the nbformat schema with `nbformat==5.11.1` (the CI gate) - linted it with `ruff==0.16.4 check --select=E9,F63,F7,F82` (the CI gate) - verified unique cell IDs, empty outputs, null execution counts, and that every code cell parses - re-validated both edited JSON assets with `json.tool` - ran `git diff --check` ### Source contract audit The request contract is read from TensorRT-LLM `main` (`tensorrt_llm/_torch/visual_gen/models/cosmos3/sampling.py`): - `load_scheduler` selects `FlowMatchEulerDiscreteScheduler` for distilled checkpoints and `UniPCMultistepScheduler` for base ones; an unknown declaration is a load-time error - `Cosmos3SamplingPolicy.validate_request` rejects any `num_inference_steps` other than `len(t_list)` and any `guidance_scale` other than `1.0` - `generation_default_overrides` supplies both values when the request leaves them unset - `set_flow_shift` is a structural no-op for distilled checkpoints, so `flow_shift` does not apply - both checkpoints declare `fixed_step_sampler_config.t_list = [1.0, 0.9375, 0.8333333333333334, 0.625]` with `stochastic_sampling: true` ## Dependencies and review state - TensorRT-LLM support is already merged: NVIDIA/TensorRT-LLM#16563 (distilled text-to-image) and NVIDIA/TensorRT-LLM#16690 (distilled image-to-video). - No dependency on #310; the two PRs touch different notebooks. If #310 merges first there is no expected conflict. --------- Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
7 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.