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

docs/source/models/supported-models.md

@@ -206,6 +206,7 @@ For full documentation, see the [Visual Generation](./visual-generation.md) page
 | `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

@@ -44,6 +44,7 @@ TensorRT-LLM **VisualGen** provides a unified inference stack for diffusion mode
 | `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/models/cosmos3/README.md

@@ -5,7 +5,7 @@ Cosmos3 supports the following 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.
-- **V2V** — video-conditioned video (`prompts/v2v.json`). Condition on a reference video via `--video_path` (a local MP4/AVI file). Only the first (or last, per `condition_video_keep`) `max(condition_video_latent_indexes) * 4 + 1` input frames condition the output (5 by default); the encoded bytes pass through and each worker decodes just that window on NVDEC (see [Media I/O dependencies](#media-io-dependencies)).
+- **V2V** — video-conditioned video (`prompts/v2v.json`). Condition on a reference video via `--video_path` (a local MP4/AVI file). Only the first (or last, per `condition_video_keep`) `max(condition_video_latent_indexes) * 4 + 1` input frames condition the output (5 by default); the encoded bytes pass through and each worker decodes just that window on NVDEC (see [Media I/O dependencies](#media-io-dependencies)). Validated for Nano / Super only.
 - **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).
 
 ## Checkpoints
@@ -16,6 +16,7 @@ Pass the Hub ID or local path via `--model`:
 - [`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
 
@@ -48,6 +49,25 @@ See `examples/visual_gen/configs/`:
 
 Example prompts live under `prompts/` (mirroring `cosmos3-internal/inputs/omni`).
 
+### Prompt inputs
+
+`--prompt` and `--negative_prompt` each accept **either literal text or a path to a
+prompt file**, chosen by whether the value names an existing file. `--prompt_file`
+and `--negative_prompt_file` accept a path only and fail if the file is missing, so
+use those when a silent fallback to literal text would be a bug (scripts, CI).
+
+A prompt file may hold any of three shapes:
+
+| Shape | Example | Notes |
+|---|---|---|
+| Omni prompt object | `prompts/t2v.json` | `prompt` plus optional `model_mode`, `vision_path`, `enable_audio`, which supply defaults for the matching flags |
+| Structured caption | a checkpoint's `assets/example_i2v_prompt.json` | the object *is* the caption; carries no options |
+| Plain text | any `.txt` | used verbatim |
+
+Structured captions are what the model cards ship and what the checkpoints were
+tuned on; they give noticeably cleaner output than a one-line summary.
+`--negative_prompt` defaults to `cosmos3_negative_prompt.json` in this directory.
+
 ## Usage
 
 ```bash
@@ -104,7 +124,17 @@ python cosmos3.py --model nvidia/Cosmos3-Super-Image2Video-4Step \
     --image_path https://example.com/frame.jpg \
     --output_path output.mp4
 
-# Inline prompt (--prompt or a JSON file path)
+# Cosmos3-Edge image-to-video (480p-native defaults: 832x480 x 121 frames).
+# Reproduces the model-card sample: the checkpoint ships a structured prompt and
+# its own negative prompt alongside the conditioning image. Fetch them with
+#   hf download nvidia/Cosmos3-Edge --local-dir Cosmos3-Edge
+python cosmos3.py --model nvidia/Cosmos3-Edge \
+    --prompt Cosmos3-Edge/assets/example_i2v_prompt.json \
+    --negative_prompt Cosmos3-Edge/assets/negative_prompt.json \
+    --image_path Cosmos3-Edge/assets/example_i2v_input.jpg \
+    --output_path output.mp4
+
+# Inline prompt
 python cosmos3.py --model nvidia/Cosmos3-Nano \
     --prompt "A cute puppy playing with a ball in a park" \
     --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml

examples/visual_gen/models/cosmos3/cosmos3.py

@@ -31,6 +31,9 @@
 
 _SCRIPT_DIR = Path(__file__).resolve().parent
 
+DEFAULT_PROMPT_FILE = "prompts/t2v.json"
+DEFAULT_NEGATIVE_PROMPT_FILE = "cosmos3_negative_prompt.json"
+
 
 def _resolve_path(path: str) -> str:
     candidate = Path(path)
@@ -42,18 +45,76 @@ def _resolve_path(path: str) -> str:
     return path
 
 
-def load_prompt_file(path: str) -> Dict[str, Any]:
-    """Load a Cosmos3 omni prompt JSON (``prompt``, optional ``vision_path``, etc.)."""
+def _is_prompt_file(value: str) -> bool:
+    """Whether a ``--prompt``/``--negative_prompt`` value names an existing file."""
+    return bool(value) and os.path.isfile(_resolve_path(value))
+
+
+def _read_prompt_payload(path: str) -> Any:
+    """Read a prompt file, decoding it as JSON when it parses and as text otherwise."""
     resolved = _resolve_path(path)
+    if not os.path.isfile(resolved):
+        raise ValueError(f"Prompt file {path!r} does not exist (resolved to {resolved!r}).")
     with open(resolved, encoding="utf-8") as f:
-        data = json.load(f)
+        raw = f.read()
+    try:
+        return json.loads(raw)
+    except json.JSONDecodeError:
+        return raw.strip()
+
+
+def load_prompt_file(path: str) -> Dict[str, Any]:
+    """Load a Cosmos3 prompt file.
+
+    Three shapes are accepted: an omni prompt object (``prompt`` plus optional
+    ``vision_path`` / ``model_mode`` / ``enable_audio``), a structured caption
+    object such as the ``assets/*_prompt.json`` files shipped with a checkpoint,
+    or plain text. The latter two carry no options, so they yield ``prompt`` only.
+    """
+    data = _read_prompt_payload(path)
+    if isinstance(data, str):
+        if not data:
+            raise ValueError(f"Prompt file {path!r} is empty.")
+        return {"prompt": data}
     if not isinstance(data, dict):
-        raise ValueError(f"Prompt file must be a JSON object, got {type(data)!r}.")
-    if not data.get("prompt"):
-        raise ValueError(f"Prompt file {resolved!r} is missing a non-empty 'prompt' field.")
+        raise ValueError(
+            f"Prompt file {path!r} must hold a JSON object or text, got {type(data).__name__}."
+        )
+    if "prompt" not in data:
+        if not data:
+            raise ValueError(f"Prompt file {path!r} is an empty JSON object.")
+        return {"prompt": json.dumps(data)}
+    if not data["prompt"]:
+        raise ValueError(f"Prompt file {path!r} is missing a non-empty 'prompt' field.")
     return data
 
 
+def load_negative_prompt_file(path: str) -> str:
+    """Load a negative prompt file (structured JSON object or plain text)."""
+    data = _read_prompt_payload(path)
+    if isinstance(data, dict):
+        return json.dumps(data)
+    if isinstance(data, str):
+        return data
+    raise ValueError(
+        f"Negative prompt file {path!r} must hold a JSON object or text, got {type(data).__name__}."
+    )
+
+
+def resolve_negative_prompt(
+    *,
+    negative_prompt: Optional[str],
+    negative_prompt_file: Optional[str],
+) -> str:
+    """Pick the negative prompt: ``--negative_prompt``, then the file, then the default."""
+    if negative_prompt is not None:
+        # --negative_prompt takes either literal text or a path to a prompt file.
+        if _is_prompt_file(negative_prompt):
+            return load_negative_prompt_file(negative_prompt)
+        return negative_prompt
+    return load_negative_prompt_file(negative_prompt_file or DEFAULT_NEGATIVE_PROMPT_FILE)
+
+
 def resolve_prompt_and_options(
     *,
     prompt: Optional[str],
@@ -67,7 +128,15 @@ def resolve_prompt_and_options(
     if prompt_file is not None:
         prompt_data = load_prompt_file(prompt_file)
 
-    resolved_prompt = prompt
+    inline_prompt: Optional[str] = None
+    if prompt is not None:
+        # --prompt takes either literal text or a path to a prompt file.
+        if _is_prompt_file(prompt):
+            prompt_data = {**prompt_data, **load_prompt_file(prompt)}
+        else:
+            inline_prompt = prompt
+
+    resolved_prompt = inline_prompt
     if resolved_prompt is None:
         resolved_prompt = prompt_data.get("prompt")
     if not resolved_prompt:
@@ -93,7 +162,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",
@@ -106,19 +176,27 @@ def main():
         "--prompt",
         type=str,
         default=None,
-        help="Text prompt for generation (overrides --prompt_file when both are set)",
+        help="Prompt text, or a path to a prompt file (overrides --prompt_file when both are set)",
     )
     parser.add_argument(
         "--prompt_file",
         type=str,
-        default="prompts/t2v.json",
-        help="Path to a JSON prompt file (default: prompts/t2v.json)",
+        default=DEFAULT_PROMPT_FILE,
+        help=f"Path to a prompt file; must exist (default: {DEFAULT_PROMPT_FILE})",
     )
     parser.add_argument(
         "--negative_prompt",
         type=str,
-        default="cosmos3_negative_prompt.json",
-        help="Text prompt or path to JSON file for negative prompt",
+        default=None,
+        help="Negative prompt text, or a path to a negative prompt file "
+        f"(overrides --negative_prompt_file; default: {DEFAULT_NEGATIVE_PROMPT_FILE})",
+    )
+    parser.add_argument(
+        "--negative_prompt_file",
+        type=str,
+        default=None,
+        help=f"Path to a negative prompt file; must exist "
+        f"(default: {DEFAULT_NEGATIVE_PROMPT_FILE})",
     )
     parser.add_argument(
         "--image_path",
@@ -187,15 +265,10 @@ def main():
     if image_path is not None:
         params.image = image_path
 
-    negative_prompt_path = _resolve_path(args.negative_prompt)
-    if args.negative_prompt is not None:
-        if os.path.isfile(negative_prompt_path) and negative_prompt_path.endswith(".json"):
-            with open(negative_prompt_path, encoding="utf-8") as f:
-                negative_prompt = json.load(f)
-        else:
-            negative_prompt = args.negative_prompt
-    else:
-        negative_prompt = None
+    negative_prompt = resolve_negative_prompt(
+        negative_prompt=args.negative_prompt,
+        negative_prompt_file=args.negative_prompt_file,
+    )
 
     if args.disable_duration_template:
         params.extra_params["use_duration_template"] = False
@@ -210,12 +283,7 @@ def main():
     if args.video_path is not None:
         params.extra_params["video"] = Path(args.video_path).read_bytes()
 
-    if negative_prompt is None:
-        params.negative_prompt = None
-    elif isinstance(negative_prompt, str):
-        params.negative_prompt = negative_prompt
-    else:
-        params.negative_prompt = json.dumps(negative_prompt)
+    params.negative_prompt = negative_prompt
 
     output = visual_gen.generate(
         inputs=prompt,

tensorrt_llm/_torch/visual_gen/executor.py

@@ -418,6 +418,12 @@ def _merge_defaults(self, req: DiffusionRequest):
                 ):
                     continue
                 setattr(params, field_name, default_value)
+                # Marks it as a pipeline default rather than caller intent, so
+                # request-dependent defaults stay re-resolvable; assigning the
+                # field re-marks it.
+                # Assumes model_fields_set is the live __pydantic_fields_set__, not a
+                # copy; TestDefaultMarksThroughRealPath fails loudly if that changes.
+                params.model_fields_set.discard(field_name)
 
         # Extra param defaults — fill all declared keys so infer() can use direct access
         specs = self.pipeline.extra_param_specs

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

@@ -111,15 +111,67 @@ def _validate_video_reference(video) -> None:
     "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,
+# 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,
 }