[None][feat] Add Cosmos3-Edge (Nemotron-dense) support by ishovkun · Pull Request #16773 · NVIDIA/TensorRT-LLM

ishovkun · 2026-07-23T04:20:35Z

Summary by CodeRabbit

New Features
- Added support for Cosmos3-Super and Cosmos3-Edge visual generation models.
- Added text-to-image, image-to-video, and text-to-video examples, including single-GPU deployment guidance.
- Added checkpoint-aware sampling, generation defaults, and system-prompt behavior.
- Added support for distilled four-step generation workflows.
Documentation
- Expanded supported-model listings and Cosmos3 usage instructions.
Bug Fixes
- Improved compatibility with varied Cosmos3 checkpoint configurations and scheduler settings.
- Preserved checkpoint defaults when optional generation settings are omitted.
Tests
- Added comprehensive unit, parity, and LPIPS quality coverage for new Cosmos3 workflows.

Description

Adds inference support for nvidia/Cosmos3-Edge, the 4B Cosmos3 variant whose
reasoner and generator towers use the Nemotron-dense architecture instead of
Nano/Super's Qwen3. Supported tasks: T2I, T2V, I2V (480p-native).
Stacked on #16690 (distilled I2V-4Step): the sampling-policy and
checkpoint-declared-defaults plumbing introduced there is extended here.
Reference implementations: diffusers Cosmos3OmniTransformer/
Cosmos3OmniPipeline (huggingface/diffusers#14181 + #14246) for the
transformer, and cosmos-framework's PyTorch backend for the sampling schedule.

Architecture recipes (the core change). backbone_type in the transformer
config selects a complete, validated architecture recipe rather than
individual flags steering construction — a mixed config cannot half-apply and
fails at startup with the offending key named. The Nemotron-dense recipe:
non-gated squared-ReLU MLP (reusing _torch/modules/mlp.py::MLP with the
existing relu2 activation, as modeling_nemotron.py does — no Edge MLP
class), no reasoner Q/K norm, and a new per-layer k_norm_und_for_gen
applied to the reasoner's keys only where the generator consumes them: in the
cached-KV design, forward_with_kv attends with rope(k_raw) internally but
caches rope(k_norm_und_for_gen(k_raw)) for the generator layers. This
distinction is subtle enough that both public references shipped it wrong
initially (diffusers normed the reasoner pathway too, fixed in #14246;
vllm-omni read a transposed config key and skipped the norm entirely, fixed
in vllm-project/vllm-omni#5239) — it is pinned here by a dedicated
regression test. Edge norms use the Nemotron flavor (weight multiply in fp32
before downcast); F.rms_norm is bit-exact to that flavor, so
NemotronRMSNorm is a one-line fused wrapper, and the attention Q/K norms of
both recipes now route through norm modules instead of a hardcoded
functional — byte-identical for Nano by construction (pinned by test).

Native flow schedule. Edge's model_index.json declares
use_native_flow_schedule: true; the intended schedule (per the model card
and cosmos-framework's PyTorch backend) is linear flow sigmas with a runtime
shift of 3.0. Stock diffusers does not reproduce it — its UniPC karras branch
discards the provided sigmas (the checkpoint config ships
use_karras_sigmas: true), which is why the schedule is validated against
cosmos-framework directly. The implementation configures diffusers' own
Apache-2.0 UniPCMultistepScheduler (karras off, flow_shift from the
per-mode table, explicit linspace sigmas): timesteps are bit-identical to
cosmos-framework's FlowUniPCMultistepScheduler and full multi-step step()
trajectories agree to ≤1.6e-7 relative across three (shift, steps) fixtures
recorded in the tests — no OpenMDW code is copied. Per-mode schedulers are
prebuilt once at load; forward() just picks by output type.

Strict weight loading (both recipes benefit). load_weights previously
collected skipped modules and dropped the list on return — a missing tensor
silently kept its init values. Coverage is now default-fail in both
directions at parameter granularity: any parameter of a constructed module
(root parameters included) without a checkpoint tensor raises with names, a
partially covered fused QKV raises, and mapped tensors that land on no
constructed parameter warn with names. The explicit skip list (lm_head, the
reasoner's final norm — text-output path VisualGen never computes — and the
action heads) is logged, never silent.

Pipeline. The tokenizer loads via AutoTokenizer (Edge ships a Nemotron
PreTrainedTokenizerFast; eos = pad = 11, <|vision_start|> = 20).
Generation defaults become a (family, mode) table — family resolved from
the same recipe function, mode from the request's output type, never from the
checkpoint name — with Edge's model-card values (video 832×480 × 121 frames,
50 steps, guidance 5.0, shift 3.0; T2I 640×640, guidance 4.0). forward()'s
numeric parameters are now None-defaulted and resolve through the same
tables, so direct callers get checkpoint-appropriate values (this also lets
the sampling policy fill fixed distilled steps/guidance for direct callers).
Requests outside the model-card envelope (256p/480p, 50–150 frames, 12–30
fps) log one advisory line and proceed — the reference runtime accepts a
wider range, so the envelope is documented support, not enforced validation.
The checkpoint's action weights are skipped loudly and the init log says
action generation is not supported; enable_safety_checker: false in the
model_index is deliberately not honored (guardrail policy unchanged).

Verified on 1×B200: per-step transformer velocity parity vs diffusers main is
0.7–1.6% relative (bf16, both CFG branches, masked I2V path included); E2E
LPIPS vs reference goldens: T2V 0.045, I2V 0.078, T2I 0.006.

Out of scope: action generation (lands on the in-flight cosmos3 action work),
audio (the checkpoint has no audio tower — requesting it raises), and
V2V/transfer (not an Edge capability; generator inputs are text, image, and
action trajectories only).

Test Coverage

Unit (tests/unittest/_torch/visual_gen/test_cosmos3_edge.py, 61 tests, no
checkpoint required for the unit tier): recipe selection/validation with
rejection cases (unknown backbone_type, contradicting flags, latent-geometry
invariants, inconsistent patch_latent_dim); Nemotron RMSNorm bit-exactness
vs the Qwen flavor; the generator-only K-norm regression (perturbing
k_norm_und_for_gen leaves reasoner attention bit-identical) plus a
numeric identity check; Nano Q/K byte-identical pin; native-flow schedule
parity against three recorded cosmos-framework fixtures (bit-equal
timesteps, full step() trajectories to 1e-9); strict-loading coverage in
both directions (missing generic/architecture-specific/root tensors raise,
partial fused raises, skip list and unconsumed-tensor warnings asserted by
content); per-family defaults/warmup/advisory resolution and the
executor-defaults shape.
Unit, checkpoint-gated: tokenizer specials + chat template; cond token ids
vs a diffusers-main golden and the uncond CFG branch vs a recorded
self-golden (keep-metadata semantics documented); model-index detection;
recipe/scheduler wiring (both mode schedulers at shift 3.0, karras off);
full 549-tensor load + forward; direct forward(None) resolution reaching
the generation path with Edge values; T2V/I2V/T2I sanity generations with
shape and non-black checks.
Parity (env-gated on DIFFUSERS_MAIN_PATH, subprocess because the pinned
diffusers predates the Edge classes): per-step velocity comparison vs
diffusers main on the real checkpoint, threshold 5% (observed 0.7–1.6%).
Integration smoke (test_cosmos3_edge_i2v_example, l0_b200.yml): the
documented example invocation at the deployed 480p × 121-frame, 50-step
shape; asserts a non-empty MP4.
Integration quality gates (test_cosmos3_edge_{t2v,i2v,t2i}_lpips_against_golden,
l0_b200.yml): LPIPS vs goldens produced by the reference implementation
(diffusers main with the scheduler patched to the native flow schedule —
equivalent to cosmos-framework at fp32-ulp), with matched initial noise and
matched cond/uncond prompt texts, so the gates check the denoising
trajectory against the reference rather than regression against a past
TRT-LLM run. Thresholds 0.10 / 0.13 / 0.05 (measured 0.0447 / 0.0778 /
0.0056 at creation; a wrong-seed I2V run measures 0.858). The I2V gate runs
10 steps — cross-stack drift accumulates per step and the deployed 50-step
shape is covered by the example test. Full provenance in
golden/visual_gen_lpips/cosmos3_edge_*.json.
Pre-existing cosmos3 suite (115 tests) passes unchanged; Nano behavior is
pinned byte-identical through the norm-routing and loader refactors.
Note for reviewers: Cosmos3-Edge (8.7 GB) must be staged in the CI
llm-models/ storage, or the checkpoint-gated tests skip.

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.

Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

…chastic sampling Distilled Cosmos3 checkpoints sample with FlowMatchEulerDiscreteScheduler's stochastic (SDE) step. diffusers respects a caller-supplied torch.Generator in that step starting in 0.39.0 (huggingface/diffusers#13678); earlier versions silently break seed reproducibility. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

…conversions Newer Cosmos3OmniTransformer checkpoint configs omit position_embedding_type, max_position_embeddings, and temporal_compression_factor_sound. Fill these schema gaps with their historical values at model construction (idempotent) instead of failing to load. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

…checkpoint - Load the scheduler class the checkpoint declares: UniPC for base checkpoints, FlowMatchEulerDiscreteScheduler for distilled ones; an explicitly unknown declaration is a load-time error. - Introduce Cosmos3SamplingPolicy, an immutable value object holding the checkpoint's sampling facts. Distilled checkpoints run their fixed 4-sigma stochastic schedule with classifier-free guidance baked into the weights (one forward per step); requests that conflict with the distilled recipe are rejected, and malformed recipes fail at load. - Generation defaults report the checkpoint's true steps/guidance; mode-dependent fields stay unset until infer() resolves the request mode exactly once. - Thread scheduler_step_kwargs through the shared denoise loop so the seeded generator reaches every stochastic scheduler step. - Register nvidia/Cosmos3-Super-Text2Image-4Step. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

Add a 1-GPU text-to-image example config that warms up the deployed image shape, README coverage with the exact invocation, and the model row in the visual-generation docs. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

Unit coverage for scheduler loading, recipe validation (only the two known recipes load), request validation, flow-shift handling, fixed-sigma timesteps, SDE seed determinism, generation defaults, infer() mode resolution, and the guidance-1.0 denoise-loop contract. A shared conftest owns TLLM_DISABLE_MPI for the VisualGen unit tests and provides a leak-free guardrail-disable fixture. Add a B200 integration test that runs the documented example invocation against the real checkpoint. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

… 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. 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>

Introduces the `nemotron_dense` architecture recipe alongside the existing `qwen3` recipe for Cosmos3, enabling the Edge checkpoint variant: - Add `resolve_arch_recipe()` to select and validate arch from `backbone_type` config field; drives all module construction - Add `NemotronRMSNorm` (float32 weight multiply), optional QK norms, and `k_norm_und_for_gen` path in `Cosmos3CausalAttention` - Add `relu2`-activated `MLP` path in the transformer block alongside the existing `GatedMLP` path - Add Edge-specific generation defaults and model-card envelope advisory in `defaults.py`; replace flat `COSMOS3_PIPELINE_DEFAULTS` with `COSMOS3_GENERATION_DEFAULTS[(family, mode)]` table - Add `native_flow_schedule` flag: builds explicit linear flow sigmas for UniPC instead of the Karras grid - Rename `set_flow_shift` → `with_flow_shift` to reflect that it returns a new scheduler instance - Switch tokenizer loading from `Qwen2Tokenizer` to `AutoTokenizer` for tokenizer-agnostic checkpoint loading - Register `nvidia/Cosmos3-Edge` in the pipeline model list Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

Recipe validation now covers the full Edge combination: declared-required sound_gen/attention_bias/rms_norm_eps, literal latent geometry (48-channel, patch 2), and patch_latent_dim consistency for both recipes. Weight loading is default-fail in both directions: root parameters join coverage (a missing audio_modality_embed raises), and mapped checkpoint tensors that land on no constructed parameter warn with names instead of vanishing. Attention Q/K norms route through NemotronRMSNorm modules in both recipes - byte-identical to the previous hardcoded F.rms_norm path - and the generation-parameter resolution and VAE temporal-factor check are extracted into testable helpers. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

Unit tier (no checkpoint required): recipe validation including rejection cases, Nemotron RMSNorm bit-exactness vs the Qwen flavor, the generator-only und K-norm regression, Nano byte-identical Q/K pinning, native flow schedule parity against recorded cosmos-framework 117c7d2 fixtures (three shift/step pairs, full trajectories through step()), strict-loading coverage in both directions, and per-family defaults/warmup/advisory resolution. Checkpoint-gated: tokenizer specials and chat template, cond token ids vs a diffusers-main golden plus an uncond self-golden (cosmos-framework keep-metadata semantics), model-index detection, recipe/scheduler wiring, load + forward, direct-forward default resolution, and T2V/I2V/T2I sanity generations with non-black checks. A DIFFUSERS_MAIN_PATH-gated subprocess compares per-step velocities against diffusers main (observed 0.7-1.6 percent relative in bf16, threshold 5 percent). The example invocation joins the B200 test list alongside the unit file. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

Model rows in the visual-generation and supported-models tables, and the example README: Edge supports T2I/T2V/I2V only (no audio tower; the checkpoint's action weights are not supported yet), with 480p-native defaults (832x480 x 121 frames, 50 UniPC steps on the checkpoint-declared native flow schedule with shift 3.0, guidance 5.0; T2I 640x640) and an advisory log for requests outside the model-card envelope. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

T2V, I2V, and T2I golden videos/image produced by the reference implementation (diffusers main, Cosmos3OmniPipeline) with the scheduler patched to the cosmos-framework native flow schedule (karras off, flow_shift 3.0) - verified equivalent to fm_solvers_unipc at fp32-ulp. Initial noise, prompts, and both CFG-branch texts (including the keep-metadata negative prompt) are matched between stacks, so the gates check the denoising trajectory against the reference rather than regression against a past TRT-LLM run. Measured cross-stack on B200: T2V 0.0447 (threshold 0.10, 50 steps), I2V 0.0778 (threshold 0.13, 10 steps - drift accumulates per step and the deployed 50-step shape is covered by the example test; a wrong-seed run measures 0.858), T2I 0.0056 (threshold 0.05). Full provenance in the golden JSONs. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

coderabbitai · 2026-07-23T04:34:05Z

📝 Walkthrough

Walkthrough

Cosmos3 support now includes Edge architecture recipes, distilled and native-flow sampling policies, checkpoint-aware pipeline defaults, scheduler kwargs propagation, new deployment examples, and expanded unit, parity, integration, and LPIPS coverage.

Changes

Cosmos3 visual generation

Layer / File(s)	Summary
Architecture recipes and checkpoint loading `tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py`	Adds Qwen3 and Nemotron-dense architecture recipes, recipe-specific normalization and MLP behavior, temporal-compression compatibility, generation-only key normalization, and strict checkpoint coverage validation.
Checkpoint-aware sampling and pipeline execution `tensorrt_llm/_torch/visual_gen/models/cosmos3/{defaults.py,sampling.py,pipeline_cosmos3.py}`	Adds checkpoint scheduler detection, distilled fixed-sigma policies, native flow scheduling, mode-specific defaults, envelope advisories, system-prompt metadata handling, conditioning anchors, audio validation, and mode-specific execution.
Shared denoising API and example enablement `tensorrt_llm/_torch/visual_gen/pipeline.py`, `examples/visual_gen/...`, `docs/source/models/...`, `requirements.txt`, `.gitignore`	Forwards scheduler-specific kwargs through denoising, updates CLI system-prompt handling, adds Cosmos3 examples and documentation, raises the Diffusers minimum version, and ignores `.plans/`.
Unit and parity validation `tests/unittest/_torch/visual_gen/*`	Adds coverage for sampling, architecture recipes, strict loading, scheduler parity, tokenizer behavior, pipeline defaults, conditioning, and test-environment setup.
Integration and LPIPS coverage `tests/integration/defs/examples/visual_gen/*`, `tests/integration/test_lists/test-db/l0_b200.yml`	Adds Cosmos3 Edge and distilled checkpoint example tests, LPIPS golden metadata, quality gates, and B200 test-list entries.

Estimated code review effort: 5 (Critical) | ~120 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Checkpoint
  participant Cosmos3Pipeline
  participant SamplingPolicy
  participant Denoiser
  participant Scheduler
  Checkpoint->>Cosmos3Pipeline: Load model and scheduler metadata
  Cosmos3Pipeline->>SamplingPolicy: Resolve checkpoint sampling policy
  Cosmos3Pipeline->>SamplingPolicy: Set timesteps and step kwargs
  Cosmos3Pipeline->>Denoiser: Run denoising with policy data
  Denoiser->>Scheduler: Execute scheduler steps

Possibly related PRs

NVIDIA/TensorRT-LLM#16683 — Refactors the VisualGen integration test module touched by the new Cosmos3 LPIPS coverage.

Suggested labels: api-compatible

Suggested reviewers: mikeiovine

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Out of Scope Changes check	⚠️ Warning	The `.gitignore` update adding `.plans/` is unrelated to Cosmos3-Edge support and the linked bugfix.	Move the `.plans/` ignore change to a separate housekeeping PR, or explain why it is required for this feature work.
Docstring Coverage	⚠️ Warning	Docstring coverage is 25.97% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is concise and accurately summarizes the main Cosmos3-Edge support change.
Description check	✅ Passed	The PR description follows the template and includes description, test coverage, and checklist sections.
Linked Issues check	✅ Passed	The Edge K-norm config uses `use_und_k_norm_for_gen`, and tests cover the corrected generator-only behavior required by `#5239`.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

Warning

There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure.

🔧 Checkov (3.3.8)

examples/visual_gen/configs/cosmos3-t2i-1gpu.yaml

Traceback (most recent call last):
File "/usr/local/bin/checkov", line 2, in
from checkov.main import Checkov
ModuleNotFoundError: No module named 'checkov'

requirements.txt

Traceback (most recent call last):
File "/usr/local/bin/checkov", line 2, in
from checkov.main import Checkov
ModuleNotFoundError: No module named 'checkov'

tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_edge_i2v_lpips_golden_video.json

Traceback (most recent call last):
File "/usr/local/bin/checkov", line 2, in
from checkov.main import Checkov
ModuleNotFoundError: No module named 'checkov'

4 others

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

coderabbitai

🧹 Nitpick comments (2)

tests/integration/defs/examples/visual_gen/test_visual_gen.py (1)
2116-2141: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Edge LPIPS gates lack on-failure candidate preservation. The three Edge gates assert the threshold but never call _preserve_lpips_candidate_on_failure, so a cross-host failure leaves no archived media to refresh the golden — yet the Edge goldens' threshold_rationale explicitly banks on that helper for the ~0.04 cross-host headroom. The distilled I2V gate (Lines 2299-2305) already wires this up; mirror it here (each gate takes request and passes its generated path + golden basename).

tests/integration/defs/examples/visual_gen/test_visual_gen.py#L2116-L2141: add request to the signature and call _preserve_lpips_candidate_on_failure before _assert_lpips_below_threshold with generated_path / cosmos3_edge_t2v_lpips_golden_video.mp4.

tests/integration/defs/examples/visual_gen/test_visual_gen.py#L2144-L2172: same, with cosmos3_edge_i2v_lpips_golden_video.mp4.

tests/integration/defs/examples/visual_gen/test_visual_gen.py#L2175-L2201: same, with cosmos3_edge_t2i_lpips_golden.png (also add the _visual_gen_deps/request fixtures already noted separately).
🤖 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
2116 - 2141, Update the three Edge LPIPS gates in
tests/integration/defs/examples/visual_gen/test_visual_gen.py at lines
2116-2141, 2144-2172, and 2175-2201: add the request fixture to each test
signature, call _preserve_lpips_candidate_on_failure with the generated path and
the corresponding golden basename before _assert_lpips_below_threshold, and add
the _visual_gen_deps fixture to the T2I gate as needed. Use
cosmos3_edge_t2v_lpips_golden_video.mp4,
cosmos3_edge_i2v_lpips_golden_video.mp4, and cosmos3_edge_t2i_lpips_golden.png
respectively.
.gitignore (1)
125-126: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

.plans/ looks unrelated to this PR's scope.

This ignore entry isn't tied to Cosmos3 Edge/distilled support and reads like a local tooling artifact. Consider dropping it here and adding it in a dedicated change (or your personal global gitignore) to keep the PR focused.

As per coding guidelines: "Keep each pull request focused on one concern and avoid scope creep; split unrelated changes into separate pull requests."
🤖 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 @.gitignore around lines 125 - 126, Remove the unrelated .plans/ entry from
.gitignore to keep this change focused on Cosmos3 Edge/distilled support; do not
replace it with another ignore rule.
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.

Nitpick comments:
In @.gitignore:
- Around line 125-126: Remove the unrelated .plans/ entry from .gitignore to
keep this change focused on Cosmos3 Edge/distilled support; do not replace it
with another ignore rule.

In `@tests/integration/defs/examples/visual_gen/test_visual_gen.py`:
- Around line 2116-2141: Update the three Edge LPIPS gates in
tests/integration/defs/examples/visual_gen/test_visual_gen.py at lines
2116-2141, 2144-2172, and 2175-2201: add the request fixture to each test
signature, call _preserve_lpips_candidate_on_failure with the generated path and
the corresponding golden basename before _assert_lpips_below_threshold, and add
the _visual_gen_deps fixture to the T2I gate as needed. Use
cosmos3_edge_t2v_lpips_golden_video.mp4,
cosmos3_edge_i2v_lpips_golden_video.mp4, and cosmos3_edge_t2i_lpips_golden.png
respectively.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 82c2a32f-30dc-4081-8a65-f056cb82ae4c

📥 Commits

Reviewing files that changed from the base of the PR and between 0a401b4 and 0165d4f.

⛔ Files ignored due to path filters (1)

tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zip is excluded by !**/*.zip

📒 Files selected for processing (24)

.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
examples/visual_gen/models/cosmos3/cosmos3.py
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/golden/visual_gen_lpips/cosmos3_edge_i2v_lpips_golden_video.json
tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_edge_t2i_lpips_golden.json
tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_edge_t2v_lpips_golden_video.json
tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_i2v_4step_lpips_golden_video.json
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/cosmos3_edge_diffusers_parity.py
tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py
tests/unittest/_torch/visual_gen/test_cosmos3_edge.py
tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py
tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py

…n at load The architecture family, the model_index use_native_flow_schedule flag, and the scheduler recipe come from three different checkpoint files and were read independently, so cross-component mismatches loaded silently. The worst case reproduces the upstream diffusers schedule bug: an Edge conversion whose model_index omits the flag would sample the karras trajectory instead of the native flow schedule with no indication. load_standard_components now validates the triple: Edge requires the declared native flow schedule and rejects distilled sampling (no such checkpoint exists); the flag is rejected outside the Edge family. Covered by a parametrized startup matrix plus an end-to-end load_standard_components case for the missing-flag conversion. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

MRoPE axes were read only from the legacy rope_scaling.mrope_section; diffusers gives the explicit top-level rope_axes_dim precedence. The official checkpoints declare both identically, which hid the gap: a top-level-only config failed and contradictory declarations silently used the nested value. The resolver now prefers the top-level field, falls back to the nested one, and rejects contradictions, wrong axis counts, and sums that do not cover half the head dimension. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

…contract The docstring promised all defaults resolved, but pipelines with mode-dependent defaults (Cosmos3: text-to-image and video requests use different resolutions/steps/guidance) deliberately leave those fields None until the output mode is known per request. Document that None means the mode's default rather than unset; runtime behavior is unchanged. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

The reduced parallel-test configs declared mrope_section [16, 12, 12] against head_dim 64 and 128, relying on the interleave slices silently clipping to head_dim//2; the new rope-axes validation rejects sections that do not sum to half the head dim. The base config now uses [12, 10, 10] (sum 32) and the FA4 config the real checkpoint sections [24, 20, 20] (sum 64), and the stale clipping comment is gone. Both configs verified to construct through the resolver. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

The pipeline and transformer set action_gen=True when the checkpoint declares action weights, despite action generation being unsupported (the weights are skipped at load). Split the flag: has_action_weights records the checkpoint fact, while action_gen stays False until action generation is actually implemented. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

- Normalize the "model." checkpoint prefix before the intentional-skip check so namespaced lm_head/action tensors are classified as skips rather than unknown keys, matching the comment's claim; record the normalized key so the skip log reports canonical family names. Add a regression test for model.-prefixed skip keys. - List Cosmos3-Edge in the example CLI --model help. - Drop review-conversation wording from a test comment. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

ishovkun added 20 commits July 17, 2026 15:40

Add .plans/ to .gitignore

50f95f1

Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

Add Cosmos3-Super-Text2Image-4Step to supported models

7dd763f

Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

[None][doc] Document Cosmos3-Super-Text2Image-4Step usage

6f33ce8

Add a 1-GPU text-to-image example config that warms up the deployed image shape, README coverage with the exact invocation, and the model row in the visual-generation docs. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

Add output file size check to Cosmos3 T2I test

c942161

Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

Add type annotations to VisualGen test helpers and fixtures

860a470

Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

Fix test_old_schema_untouched to verify all non-default fields

ea72ebb

Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

Add TypeVar annotation to preserve config type in compat defaults

e3ec19c

Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

ishovkun requested review from a team as code owners July 23, 2026 04:20

ishovkun requested review from BowenFu, QiJune, karljang and o-stoner July 23, 2026 04:20

ishovkun requested review from arysef, fredricz-20070104, mzweilz, niukuo and ruodil July 23, 2026 04:20

github-actions Bot assigned ishovkun Jul 23, 2026

github-actions Bot added the VisualGen label Jul 23, 2026

coderabbitai Bot reviewed Jul 23, 2026

View reviewed changes

ishovkun added 3 commits July 22, 2026 21:58

ishovkun requested a review from a team as a code owner July 23, 2026 05:17

ishovkun requested a review from zhaoyangwang-nvidia July 23, 2026 05:17

ishovkun added 3 commits July 22, 2026 22:19

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[None][feat] Add Cosmos3-Edge (Nemotron-dense) support#16773

[None][feat] Add Cosmos3-Edge (Nemotron-dense) support#16773
ishovkun wants to merge 26 commits into
NVIDIA:mainfrom
ishovkun:cosmos3_edge

ishovkun commented Jul 23, 2026 •

edited by coderabbitai Bot

Loading

Uh oh!

coderabbitai Bot commented Jul 23, 2026

Walkthrough

Changes

Sequence Diagram(s)

❌ Failed checks (2 warnings)

Uh oh!

coderabbitai Bot left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

ishovkun commented Jul 23, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

coderabbitai Bot commented Jul 23, 2026

Walkthrough

Changes

Sequence Diagram(s)

❌ Failed checks (2 warnings)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

ishovkun commented Jul 23, 2026 •

edited by coderabbitai Bot

Loading