[None][feat] Support the DMD2-distilled Cosmos3-Super-Text2Image-4Step checkpoint

[ishovkun]

Summary by CodeRabbit

• New Features

  • Added support for NVIDIA Cosmos3 Super Text-to-Image, including a distilled four-step generation option.
  • Added a one-GPU 1024×1024 image-generation configuration and usage example.
  • Improved checkpoint-specific sampling, validation, warmup behavior, and image output handling.
  • Added compatibility defaults for newer Cosmos3 checkpoint configurations.

• Documentation

  • Documented the newly supported model, deployment configuration, usage constraints, and example command.

• Bug Fixes

  • Improved scheduler behavior and seeded generation consistency across supported Cosmos3 workflows.

• Tests

  • Added unit and end-to-end coverage for distilled generation, validation, configuration compatibility, and image output.

Description

Adds support for nvidia/Cosmos3-Super-Text2Image-4Step, a DMD2-distilled text-to-image Cosmos3 checkpoint. Unlike the base checkpoints, it samples with FlowMatchEulerDiscreteScheduler on a fixed 4-sigma stochastic (SDE) schedule declared in the checkpoint's scheduler config, with classifier-free guidance baked into the weights (one forward per step). It also ships a newer diffusers conversion of the transformer config that omits a few schema fields older conversions carried.

What changed:

• Scheduler loading (sampling.py): the pipeline instantiates the scheduler class the checkpoint declares — UniPC for base checkpoints (a missing declaration also resolves to UniPC, preserving existing behavior), FlowMatchEuler for distilled ones. An explicitly unknown declaration is a load-time error rather than a silent UniPC substitution.
• Cosmos3SamplingPolicy (sampling.py): an immutable value object holding the checkpoint's sampling facts (fixed sigmas, distilled detection, UniPC base config for flow-shift rebuilds). Only two recipes are valid — UniPC without fixed sigmas (base) and FlowMatchEuler with fixed sigmas plus stochastic_sampling=true (distilled); malformed combinations, including non-stochastic or non-SDE declarations, fail at load. Requests that conflict with a distilled checkpoint's fixed steps/guidance are rejected with a clear error, as are image-conditioned requests (correct distilled conditioning needs per-step re-anchoring, which lands with I2V-4Step support).
• Honest generation defaults: default_generation_params reports the checkpoint's true steps/guidance (4 / 1.0 for distilled). Mode-dependent fields (height, width, num_inference_steps, guidance_scale) stay None until infer() resolves the request mode (video vs. image) exactly once; explicit request values pass through unchanged.
• Seed determinism (pipeline.py): the shared denoise loop now threads scheduler_step_kwargs into every scheduler.step() call so the request-seeded torch.Generator drives the stochastic step's noise. Requires diffusers>=0.39.0 — earlier versions ignore a caller-supplied generator in the stochastic branch (Fix ignored generator in FlowMatchEulerDiscreteScheduler huggingface/diffusers#13678). Same-seed runs produce bit-identical images.
• Transformer config schema compat (transformer_cosmos3.py): newer conversions omit position_embedding_type, max_position_embeddings, and temporal_compression_factor_sound; these are filled with their historical values at model construction (idempotent).
• Deployment config + docs: examples/visual_gen/configs/cosmos3-t2i-1gpu.yaml warms up the deployed 1024×1024 single-frame shape (warmup follows the workflow, not the checkpoint name); README and model docs updated with the exact invocation.

Behavior for base checkpoints (Cosmos3-Nano, Cosmos3-Super) is unchanged: same UniPC scheduler, same defaults, no new step kwargs.

Test Coverage

• tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py (new, 69 tests): scheduler loading from real config files, recipe-matrix validation (malformed combinations fail at load), distilled request validation, flow-shift rebuild/restore semantics, fixed-sigma timestep programming, SDE seed determinism (same seed reproduces, different seeds diverge), generation defaults, infer() mode resolution, guidance-1.0 denoise-loop contract (single forward per step, step kwargs reach every scheduler.step), registry dispatch. Includes a canary pinning that diffusers retains unknown scheduler-config keys, which distilled detection depends on.
• tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py: config schema compat-default tests.
• tests/integration/defs/examples/visual_gen/test_visual_gen.py::test_cosmos3_t2i_4step_example (in l0_b200.yml, B200 post-merge): runs the documented example invocation against the real checkpoint and asserts an image is produced.
• Manually verified on B200: e2e image generation via the documented invocation (fixed 4-step schedule detected from the checkpoint, warmup on the 1024×1024/1-frame shape, denoise ~1.4 s), and bit-identical outputs across same-seed runs. Full cosmos3 unit suite (105 tests across the three modules) passes in a single process in both module orders.

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.

[ishovkun]

Could a maintainer please add the api-compatible label? No public API signature changes in this PR — the observable change is that VisualGen.default_params now reports None for mode-dependent fields (resolved per request) and the checkpoint's true steps/guidance for distilled checkpoints.

[coderabbitai]

Walkthrough

Adds checkpoint-aware Cosmos3 sampling for distilled fixed-step text-to-image generation, integrates the new checkpoint and deployment configuration, adds transformer compatibility defaults, updates scheduler denoising support, and expands unit and integration coverage.

Changes

Cosmos3 distilled text-to-image

Layer / File(s)	Summary
Checkpoint-aware sampling policy tensorrt_llm/_torch/visual_gen/models/cosmos3/sampling.py	Loads UniPC or FlowMatch schedulers, derives distilled fixed sigmas, validates fixed steps and guidance, programs timesteps, and supplies scheduler step arguments.
Pipeline sampling and request integration tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py, tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py, tensorrt_llm/_torch/visual_gen/pipeline.py	Integrates policy-driven defaults, mode resolution, scheduler setup, flow-shift handling, output validation, warmup, and denoising kwargs.
Transformer configuration compatibility tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py	Populates missing pretrained configuration fields before transformer initialization.
Distilled checkpoint deployment guidance examples/visual_gen/configs/cosmos3-t2i-1gpu.yaml, examples/visual_gen/models/cosmos3/README.md, docs/source/models/*.md, requirements.txt	Documents the distilled checkpoint, adds a one-GPU 1024×1024 image configuration and command, updates supported-model tables, and raises the diffusers minimum version.
Sampling and deployment validation tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py, tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py, tests/unittest/_torch/visual_gen/conftest.py, tests/integration/defs/examples/visual_gen/test_visual_gen.py, tests/integration/test_lists/test-db/l0_b200.yml	Adds coverage for scheduler policies, request validation, denoising, compatibility defaults, registry dispatch, and end-to-end image generation.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant User
  participant Cosmos3OmniMoTPipeline
  participant Cosmos3SamplingPolicy
  participant BasePipeline
  participant FlowMatchEulerDiscreteScheduler
  User->>Cosmos3OmniMoTPipeline: request distilled text-to-image
  Cosmos3OmniMoTPipeline->>Cosmos3SamplingPolicy: validate fixed steps and guidance
  Cosmos3OmniMoTPipeline->>Cosmos3SamplingPolicy: configure timesteps and step kwargs
  Cosmos3OmniMoTPipeline->>BasePipeline: start denoising
  BasePipeline->>FlowMatchEulerDiscreteScheduler: step with generator
  FlowMatchEulerDiscreteScheduler-->>Cosmos3OmniMoTPipeline: updated latents
  Cosmos3OmniMoTPipeline-->>User: output image

Loading

Suggested reviewers: dc3671

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 10.77% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title matches the main change and follows the repository’s [None][feat] format.
Description check	✅ Passed	The description includes the required Description, Test Coverage, and PR Checklist sections and is sufficiently detailed.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 3

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py (1)

189-247: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

Default distilled requests and warmup still select video mode.

For Cosmos3-Super-Text2Image-4Step, omitted output_type resolves to "video", while warmup remains 720×1280×189 and also invokes the video path. Default distilled requests should select "image", and warmup should use the T2I resolution with one frame; explicitly reject video mode if this checkpoint does not support it.

🤖 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 `@tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py` around
lines 189 - 247, Update Cosmos3-Super-Text2Image-4Step defaults so omitted
output_type resolves to "image" rather than "video". Adjust
default_warmup_resolutions and default_warmup_num_frames, and _run_warmup, to
use the T2I resolution with one frame. In infer, explicitly reject video mode
for this checkpoint while preserving supported image generation behavior.

🧹 Nitpick comments (1)

tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py (1)

46-51: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add parameter and return annotations to the compatibility helper.

Use the concrete pretrained-config type accepted by DiffusionModelConfig.

As per coding guidelines, “Annotate every function.”

🤖 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 `@tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py` around
lines 46 - 51, Annotate apply_pretrained_config_compat_defaults with the
concrete pretrained-config type accepted by DiffusionModelConfig and its return
type, reflecting that it mutates and returns the same configuration object.
Preserve the existing idempotent default-filling behavior.

Source: Coding guidelines

🤖 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 `@tests/integration/defs/examples/visual_gen/test_visual_gen.py`:
- Around line 1868-1900: Update the test around the output_path assertion in the
Cosmos3 visual generation case to remove any existing PNG before invoking
venv_check_call, then verify the newly generated artifact exists and is
non-empty. Keep the existing output path and invocation unchanged, and state
that coverage is sufficient only when the fresh file check passes.

In `@tests/unittest/_torch/visual_gen/conftest.py`:
- Around line 19-33: Annotate every new callable with complete, precise types
and avoid Any: in tests/unittest/_torch/visual_gen/conftest.py lines 19-33,
update disable_cosmos3_guardrails with a concrete iterator/generator return
type; in tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py lines
57-552, annotate all helper parameters and returns and add -> None to test
methods; in tests/integration/defs/examples/visual_gen/test_visual_gen.py lines
1857-1900, annotate fixture parameters and add -> None; and in
tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py lines 475-495, add
-> None to each new test method.

In `@tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py`:
- Around line 481-488: Extend test_old_schema_untouched in
apply_pretrained_config_compat_defaults coverage so every explicitly provided
field uses a non-default sentinel and has its own preservation assertion,
including position_embedding_type and temporal_compression_factor_sound
alongside max_position_embeddings. Confirm TensorRT-LLM coverage is not relying
on only one sentinel, and add equivalent assertions there if its compatibility
tests cover this helper.

---

Outside diff comments:
In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py`:
- Around line 189-247: Update Cosmos3-Super-Text2Image-4Step defaults so omitted
output_type resolves to "image" rather than "video". Adjust
default_warmup_resolutions and default_warmup_num_frames, and _run_warmup, to
use the T2I resolution with one frame. In infer, explicitly reject video mode
for this checkpoint while preserving supported image generation behavior.

---

Nitpick comments:
In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py`:
- Around line 46-51: Annotate apply_pretrained_config_compat_defaults with the
concrete pretrained-config type accepted by DiffusionModelConfig and its return
type, reflecting that it mutates and returns the same configuration object.
Preserve the existing idempotent default-filling behavior.

🪄 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: 938dea2c-547e-404e-92bc-4778696ac546

📥 Commits

Reviewing files that changed from the base of the PR and between f1434b7 and 02d90b1.

📒 Files selected for processing (17)

• .gitignore
• docs/source/models/supported-models.md
• docs/source/models/visual-generation.md
• examples/visual_gen/configs/cosmos3-t2i-1gpu.yaml
• examples/visual_gen/models/cosmos3/README.md
• requirements.txt
• tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py
• tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py
• tensorrt_llm/_torch/visual_gen/models/cosmos3/sampling.py
• tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py
• tensorrt_llm/_torch/visual_gen/pipeline.py
• tests/integration/defs/examples/visual_gen/test_visual_gen.py
• tests/integration/test_lists/test-db/l0_b200.yml
• tests/unittest/_torch/visual_gen/conftest.py
• tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py
• tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py
• tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py

[ishovkun]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #60087 [ run ] triggered by Bot. Commit: 02d90b1 Link to invocation

[ishovkun]

Responses to the two CodeRabbit items without inline threads:

Nitpick — annotate apply_pretrained_config_compat_defaults with the concrete pretrained-config type (transformer_cosmos3.py): addressed in e3ec19c. There is no concrete class to name — DiffusionModelConfig.pretrained_config is declared Optional[Any] and the loader builds a SimpleNamespace from the checkpoint JSON — so the helper is annotated with an identity TypeVar (_PretrainedConfigT -> _PretrainedConfigT), which states precisely that it mutates and returns the same configuration object without over-narrowing a deliberately structural field.

Outside-diff — make omitted output_type resolve to image, default warmup to the T2I shape, and reject video mode for this checkpoint (pipeline_cosmos3.py): not applying, by design. The checkpoint declares its sampling recipe (distilled FlowMatchEuler + fixed sigmas), not its modality — "distilled ⇒ image-only" does not hold (a distilled image-to-video checkpoint exists), so implementing this would require gating on the checkpoint name, which this codebase deliberately avoids. Request validation rejects what the checkpoint actually declares (conflicting steps/guidance). Warmup shape follows the deployed workflow rather than the checkpoint: configs/cosmos3-t2i-1gpu.yaml sets compilation_config to 1024×1024 / 1 frame, and the documented invocation, the README, and the B200 integration test all use it together with an explicit --output_type image.

[tensorrt-cicd]

PR_Github #60087 [ run ] completed with state SUCCESS. Commit: 02d90b1
/LLM/main/L0_MergeRequest_PR pipeline #48474 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[BowenFu]

LGTM — additive DMD2-distilled Cosmos3 support; the shared pipeline.py change adds an optional scheduler_step_kwargs (default None→{}, so existing models are unchanged) and there's no public tensorrt_llm/visual_gen API change. Only cross-cutting item: the diffusers floor bump 0.37.1→0.39.0 raises the minimum for all diffusers-based VisualGen models — CI-covered, worth a heads-up.

[ishovkun]

/bot run --disable-fail-fast

[ishovkun]

Rebased onto current main (cfebf191f3) — the branch was 86 commits behind, and the two x86 failures (unittest/_torch/executor, unittest/_torch/speculative/test_eagle3.py) were in areas main has since fixed, e.g. #16805 (draft token accounting) and #16571 (exclude ADP/cuda-graph dummy requests from speculative-decode acceptance stats). Neither area is touched by this PR. The SBSA multi-GPU failure was infra (agent offline → 240 min Slurm walltime).

Also worth noting from the last run: test_flux_attention.py passed, confirming the conftest fix in the previous commit.

/bot run --disable-fail-fast

[ishovkun]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #62229 [ run ] triggered by Bot. Commit: a9a7dac Link to invocation

[tensorrt-cicd]

PR_Github #62229 [ run ] completed with state FAILURE. Commit: a9a7dac
/LLM/main/L0_MergeRequest_PR pipeline #50397 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[ishovkun]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #62576 [ run ] triggered by Bot. Commit: a9a7dac Link to invocation

[ishovkun]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #62694 [ run ] triggered by Bot. Commit: a9a7dac Link to invocation

[tensorrt-cicd]

PR_Github #62576 [ run ] completed with state ABORTED. Commit: a9a7dac

Link to invocation

[tensorrt-cicd]

PR_Github #62694 [ run ] completed with state SUCCESS. Commit: a9a7dac
/LLM/main/L0_MergeRequest_PR pipeline #50832 completed with status: 'SUCCESS'

CI Report

Link to invocation