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

docs/source/models/supported-models.md

@@ -171,6 +171,7 @@ For full documentation, see the [Visual Generation](./visual-generation.md) page
 | `Qwen/Qwen-Image-Layered` | Image-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) |
 
 ### Feature Matrix
 

docs/source/models/visual-generation.md

@@ -40,6 +40,7 @@ TensorRT-LLM **VisualGen** provides a unified inference stack for diffusion mode
 | `Qwen/Qwen-Image-Layered` | Image-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) |
 
 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

@@ -13,6 +13,7 @@ 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`.
 
 ## Guardrails
 
@@ -36,6 +37,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 +72,14 @@ 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
+
 # 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" \

requirements-dev.txt

@@ -1,7 +1,8 @@
 -r requirements.txt
-# VisualGen LPIPS goldens were recorded with this version. Install it before
-# pytest collection so already-imported Diffusers modules match the package files.
-diffusers==0.38.0
+# Pin the exact diffusers version VisualGen LPIPS tests run against; install it
+# before pytest collection so already-imported Diffusers modules match the
+# package files. 0.39.0 matches the runtime floor in requirements.txt.
+diffusers==0.39.0
 boto3
 einops
 lpips

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,18 @@
     "guidance_interval": (400.0, 1000.0),
 }
 
+# Fields merged by the executor into every request. Mode-dependent values
+# remain None until infer() selects the request mode; key membership also
+# declares these fields supported during request validation.
+COSMOS3_PIPELINE_DEFAULTS = {
+    **COSMOS3_720P_PARAMS,
+    "height": None,
+    "width": None,
+    "num_inference_steps": None,
+    "guidance_scale": None,
+}
+
+
 COSMOS3_EXTRA_SPECS: Dict[str, ExtraParamSchema] = {
     "use_duration_template": ExtraParamSchema(
         type="bool",