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

.gitignore

@@ -122,3 +122,5 @@ tests/integration/defs/stress_test/artifacts/
 .claude/agent-tests/perf-test-sync/report.html
 .claude/agent-tests/perf-test-sync/results.json
 .claude/settings.json
+
+.plans/

docs/source/models/supported-models.md

@@ -167,6 +167,9 @@ For full documentation, see the [Visual Generation](./visual-generation.md) page
 | `Qwen/Qwen-Image-2512` | Text-to-Image |
 | `nvidia/Cosmos3-Nano` | Text-to-Image, Text-to-Video, Image-to-Video |
 | `nvidia/Cosmos3-Super` | Text-to-Image, Text-to-Video, Image-to-Video |
+| `nvidia/Cosmos3-Super-Text2Image-4Step` | Text-to-Image (DMD2-distilled, fixed 4-step schedule) |
+| `nvidia/Cosmos3-Super-Image2Video-4Step` | Image-to-Video (DMD2-distilled, fixed 4-step schedule) |
+| `nvidia/Cosmos3-Edge` | Text-to-Image, Text-to-Video, Image-to-Video (Nemotron-dense backbone, 480p-native) |
 
 ### Feature Matrix
 

docs/source/models/visual-generation.md

@@ -39,6 +39,9 @@ TensorRT-LLM **VisualGen** provides a unified inference stack for diffusion mode
 | `Qwen/Qwen-Image-2512` | Text-to-Image |
 | `nvidia/Cosmos3-Nano` | Text-to-Image, Text-to-Video, Image-to-Video |
 | `nvidia/Cosmos3-Super` | Text-to-Image, Text-to-Video, Image-to-Video |
+| `nvidia/Cosmos3-Super-Text2Image-4Step` | Text-to-Image (DMD2-distilled, fixed 4-step schedule) |
+| `nvidia/Cosmos3-Super-Image2Video-4Step` | Image-to-Video (DMD2-distilled, fixed 4-step schedule) |
+| `nvidia/Cosmos3-Edge` | Text-to-Image, Text-to-Video, Image-to-Video (Nemotron-dense backbone, 480p-native) |
 
 Models are auto-detected from the checkpoint directory. Diffusers-format models are detected via `model_index.json`; LTX-2 monolithic safetensors checkpoints are detected via embedded metadata. The `AutoPipeline` registry selects the appropriate pipeline class automatically.
 

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

@@ -0,0 +1,30 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# 1-GPU Cosmos3 text-to-image deployment (base or distilled T2I checkpoints).
+# Model: nvidia/Cosmos3-Super-Text2Image or nvidia/Cosmos3-Super-Text2Image-4Step
+# Shared by offline examples (--visual_gen_args) and trtllm-serve.
+#
+# Warmup expresses the deployed workload: warms the 1024x1024 single-frame
+# shape instead of the omni default (720p x 189-frame video). Requests should
+# pass output_type="image".
+attention_config:
+  backend: VANILLA
+parallel_config:
+  cfg_size: 1
+  ulysses_size: 1
+compilation_config:
+  resolutions: [[1024, 1024]]
+  num_frames: [1]

examples/visual_gen/models/cosmos3/README.md

@@ -5,14 +5,17 @@ Cosmos3 supports four generation modes from a single checkpoint:
 - **T2V** — text-to-video (`prompts/t2v.json`).
 - **T2I** — text-to-image (`prompts/t2i.json`); emits a still frame (use `--output_type image` / a non-video `--output_path`).
 - **I2V / TI2V** — image-conditioned video (`prompts/i2v.json`). Condition on a reference frame via the prompt file's `vision_path` or `--image_path`. The image may be a local path, a `file://` / `http(s)://` URL, or a `data:` URI.
-- **T2AV** — text-to-video with synchronized audio (`prompts/t2av.json` with `enable_audio: true`, or pass `--enable_audio`). Combine with a `vision_path` for image-conditioned audio-video (TI2AV).
+- **T2AV** — text-to-video with synchronized audio (`prompts/t2av.json` with `enable_audio: true`, or pass `--enable_audio`). Combine with a `vision_path` for image-conditioned audio-video (TI2AV). Audio requires an audio-capable checkpoint (Nano / Super); Cosmos3-Edge has no audio tower.
 
 ## Checkpoints
 
 Pass the Hub ID or local path via `--model`:
 
 - [`nvidia/Cosmos3-Nano`](https://huggingface.co/nvidia/Cosmos3-Nano)
 - [`nvidia/Cosmos3-Super`](https://huggingface.co/nvidia/Cosmos3-Super)
+- [`nvidia/Cosmos3-Super-Text2Image-4Step`](https://huggingface.co/nvidia/Cosmos3-Super-Text2Image-4Step) — DMD2-distilled text-to-image: fixed 4-step schedule with classifier-free guidance baked into the weights. Steps/guidance are read from the checkpoint; conflicting request values are rejected. Use with `configs/cosmos3-t2i-1gpu.yaml`.
+- [`nvidia/Cosmos3-Super-Image2Video-4Step`](https://huggingface.co/nvidia/Cosmos3-Super-Image2Video-4Step) — DMD2-distilled image-to-video: same fixed 4-step, guidance-baked-in contract. The default omni video shape (720p × 189 frames) is the deployed shape, so no dedicated config is needed. This checkpoint declares `default_use_system_prompt: true` in its `model_index.json`, which the pipeline applies automatically (override with `--use_system_prompt` / `--no-use_system_prompt`).
+- [`nvidia/Cosmos3-Edge`](https://huggingface.co/nvidia/Cosmos3-Edge) — 4B Nemotron-dense backbone supporting **T2I / T2V / I2V only**: no audio tower, and the checkpoint's action weights are not supported by this pipeline yet. 480p-native defaults (832×480 × 121 frames, 50 UniPC steps on the checkpoint-declared native flow schedule with shift 3.0, guidance 5.0; T2I defaults to 640×640), so no dedicated config is needed. The model card validates 256p/480p, 50–150 frames, and 12–30 FPS; requests outside that envelope run with an advisory log.
 
 ## Guardrails
 
@@ -36,6 +39,7 @@ See `examples/visual_gen/configs/`:
 
 - `cosmos3-nano-1gpu.yaml` — 1 GPU
 - `cosmos3-super-4gpu.yaml` — 4 GPU, CFG + Ulysses + parallel VAE
+- `cosmos3-t2i-1gpu.yaml` — 1 GPU, text-to-image deployments (base or distilled): warms the deployed 1024×1024 single-frame shape instead of the omni video shape.
 
 Example prompts live under `prompts/` (mirroring `cosmos3-internal/inputs/omni`).
 
@@ -70,6 +74,28 @@ python cosmos3.py --model nvidia/Cosmos3-Nano \
     --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml \
     --output_path output.png
 
+# T2I, distilled 4-step checkpoint (use the T2I config so warmup runs the
+# image shape; steps/guidance come from the checkpoint automatically)
+python cosmos3.py --model nvidia/Cosmos3-Super-Text2Image-4Step \
+    --prompt_file prompts/t2i.json \
+    --visual_gen_args ../../configs/cosmos3-t2i-1gpu.yaml \
+    --output_type image \
+    --output_path output.png
+
+# I2V, distilled 4-step checkpoint (steps/guidance and the system-prompt
+# default come from the checkpoint automatically; defaults are the deployed
+# 720p x 189-frame shape, so no config is required)
+python cosmos3.py --model nvidia/Cosmos3-Super-Image2Video-4Step \
+    --prompt "The camera slowly pans right across the scene" \
+    --image_path https://example.com/frame.jpg \
+    --output_path output.mp4
+
+# Cosmos3-Edge image-to-video (480p-native defaults: 832x480 x 121 frames)
+python cosmos3.py --model nvidia/Cosmos3-Edge \
+    --prompt "The camera slowly pans right across the scene" \
+    --image_path https://example.com/frame.jpg \
+    --output_path output.mp4
+
 # Inline prompt (--prompt or a JSON file path)
 python cosmos3.py --model nvidia/Cosmos3-Nano \
     --prompt "A cute puppy playing with a ball in a park" \

examples/visual_gen/models/cosmos3/cosmos3.py

@@ -86,7 +86,8 @@ def main():
         "--model",
         type=str,
         default="nvidia/Cosmos3-Nano",
-        help="Model path or HuggingFace Hub ID (nvidia/Cosmos3-Nano, nvidia/Cosmos3-Super)",
+        help="Model path or HuggingFace Hub ID "
+        "(nvidia/Cosmos3-Nano, nvidia/Cosmos3-Super, nvidia/Cosmos3-Edge)",
     )
     parser.add_argument(
         "--visual_gen_args",
@@ -136,7 +137,13 @@ def main():
         help="Disable resolution metadata template (enabled by default, matching cosmos-framework CLI)",
     )
     parser.add_argument(
-        "--use_system_prompt", action="store_true", help="Use system prompt in prompt"
+        "--use_system_prompt",
+        action=argparse.BooleanOptionalAction,
+        default=None,
+        help=(
+            "Prepend the Cosmos3 system prompt (--no-use_system_prompt to disable). "
+            "When omitted, the checkpoint's declared default applies."
+        ),
     )
     parser.add_argument("--enable_audio", action="store_true", help="Enable audio generation")
     parser.add_argument(
@@ -181,7 +188,8 @@ def main():
         params.extra_params["use_duration_template"] = False
     if args.disable_resolution_template:
         params.extra_params["use_resolution_template"] = False
-    params.extra_params["use_system_prompt"] = args.use_system_prompt
+    if args.use_system_prompt is not None:
+        params.extra_params["use_system_prompt"] = args.use_system_prompt
     params.extra_params["enable_audio"] = enable_audio
     params.extra_params["use_guardrails"] = not args.disable_guardrails
     params.extra_params["output_type"] = output_type

requirements.txt

@@ -4,7 +4,8 @@ accelerate>=1.7.0
 build
 colored
 cuda-python>=13
-diffusers>=0.37.1
+# FlowMatchEuler respects a supplied generator starting in 0.39.0 (huggingface/diffusers#13678).
+diffusers>=0.39.0
 ftfy
 lark
 lazy_loader~=0.5

tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py

@@ -35,10 +35,7 @@
     "frame_rate": 24.0,
 }
 
-# Text-to-image (``output_type="image"``) defaults. Applied by the pipeline when
-# the corresponding request field still carries the merged video default, since
-# the executor merges a single ``default_generation_params`` dict (the video
-# params above) into the request before ``infer()`` runs.
+# Text-to-image (``output_type="image"``) defaults; resolved in ``infer()``.
 COSMOS3_T2I_PARAMS = {
     "height": 1024,
     "width": 1024,
@@ -48,6 +45,62 @@
     "guidance_interval": (400.0, 1000.0),
 }
 
+# Edge (Nemotron-dense backbone) is 480p-native. Video values follow the
+# model-card I2V command (T2V mirrors it — the model card documents I2V only);
+# ``flow_shift`` rides the checkpoint-declared native flow schedule. T2I values
+# are the cosmos-framework t2i mode defaults at Edge's native resolution
+# (480p at 1:1 aspect), with full-range CFG.
+COSMOS3_EDGE_VIDEO_PARAMS = {
+    "height": 480,
+    "width": 832,
+    "num_inference_steps": 50,
+    "guidance_scale": 5.0,
+    "max_sequence_length": 4096,
+    "num_frames": 121,
+    "frame_rate": 24.0,
+    "flow_shift": 3.0,
+}
+
+COSMOS3_EDGE_T2I_PARAMS = {
+    "height": 640,
+    "width": 640,
+    "num_inference_steps": 50,
+    "guidance_scale": 4.0,
+    "flow_shift": 3.0,
+    "guidance_interval": None,
+}
+
+# Model-card validated envelope for Edge; advisory only (the reference
+# runtime accepts a wider range), surfaced as a log line per request.
+COSMOS3_EDGE_ENVELOPE = {
+    "num_frames": (50, 150),
+    "frame_rate": (12.0, 30.0),
+    "max_sequence_length": 4096,
+    "resolutions": frozenset(
+        {(640, 640), (544, 736), (736, 544), (480, 832), (832, 480), (256, 256), (256, 320),
+         (320, 256), (192, 320), (320, 192)}
+    ),
+}
+
+# (family, mode) → generation defaults. Family is the architecture recipe
+# name resolved from the transformer config; mode is the request's output
+# type — never inferred from the checkpoint name (a task-specialized
+# checkpoint can still be asked to run any mode).
+COSMOS3_GENERATION_DEFAULTS: Dict = {
+    ("qwen3", "video"): COSMOS3_720P_PARAMS,
+    ("qwen3", "image"): COSMOS3_T2I_PARAMS,
+    ("nemotron_dense", "video"): COSMOS3_EDGE_VIDEO_PARAMS,
+    ("nemotron_dense", "image"): COSMOS3_EDGE_T2I_PARAMS,
+}
+
+# Families without an entry get no envelope advisory.
+COSMOS3_ENVELOPES: Dict = {
+    "nemotron_dense": COSMOS3_EDGE_ENVELOPE,
+}
+
+
+
+
 COSMOS3_EXTRA_SPECS: Dict[str, ExtraParamSchema] = {
     "use_duration_template": ExtraParamSchema(
         type="bool",