[TRTLLM-14268][feat] Cosmos3 Transfer (control-video conditioning)

examples/visual_gen/models/cosmos3/README.md

@@ -6,6 +6,7 @@ Cosmos3 supports the following generation modes from a single checkpoint:
 - **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)). Validated for Nano / Super only.
+- **Transfer** — control-video conditioning (`edge`/`blur`/`depth`/`seg`/`wsm` hints via `--extra_params`). The control constrains structure frame by frame; the prompt supplies appearance. `edge` and `blur` are auto-computed from `--video_path`; any hint also accepts a precomputed control clip (`{"edge": "control.mp4"}` — the example reads it and sends encoded bytes, the same contract as the `video` reference). Multiple hints compose (each adds a full control-token copy of the video sequence); long videos run chunked (93 frames/chunk, stitched on overlap frames) — but only past the first chunk, so raise `num_frames` above the pipeline default to generate one: it bounds how many frames are decoded from the inputs, and so how long the output can be. A single-hint request picks up that hint's tuned sampling preset — guidance scale, control guidance and flow shift — for any of those the request leaves unset; requests with several hints fall back to the generic video defaults. The active hint names are also appended to the prompt as a one-sentence control-adherence directive; pass `"emphasize_control_in_prompt": false` to suppress it for clean baselines or ablations.
 - **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
@@ -38,6 +39,7 @@ export TRTLLM_DISABLE_COSMOS3_GUARDRAILS=1
 
 - Saving `.mp4` output requires the `ffmpeg` CLI on `PATH` (`apt-get install -y ffmpeg`); without it the encoder falls back to `.avi`.
 - Decoding MP4/AVI reference videos (V2V) happens in the worker processes on NVDEC via PyNvVideoCodec, a declared TensorRT-LLM dependency — nothing extra to install. Tested combinations: H.264 in MP4 and H.264 in AVI; other containers/codecs/profiles depend on the demuxer and the GPU's NVDEC capabilities and are best-effort.
+- Transfer's `edge`/`blur` controls are derived on the GPU from the reference video — nothing extra to install. Precomputed controls (`depth`/`seg`/`wsm`, or a precomputed `edge`/`blur`) are decoded like any other reference video.
 
 ## Deployment configs
 
@@ -90,8 +92,9 @@ python cosmos3.py --model nvidia/Cosmos3-Nano \
 
 # V2V: video-conditioned video (continues the first frames of --video_path).
 # Best results when the prompt describes the input video — e.g. continue a
-# T2V output reusing its original prompt. Output size is fixed (1280x720
-# default); inputs are center-cropped, not aspect-matched.
+# T2V output reusing its original prompt. Output size follows the source's
+# aspect ratio (the closest supported bucket) unless the request sets
+# height/width; the reference is center-cropped to whatever size is chosen.
 python cosmos3.py --model /path/to/Cosmos3-Nano \
     --prompt_file prompts/v2v.json \
     --video_path /path/to/Cosmos3-Nano/assets/example_i2v_output.mp4 \
@@ -124,6 +127,34 @@ python cosmos3.py --model nvidia/Cosmos3-Super-Image2Video-4Step \
     --image_path https://example.com/frame.jpg \
     --output_path output.mp4
 
+# Transfer: control-video conditioning — structure from the control video,
+# appearance from the prompt. edge/blur are computed from --video_path.
+python cosmos3.py --model nvidia/Cosmos3-Nano \
+    --prompt "The same scene rendered as a photorealistic video, sharp detail." \
+    --video_path /path/to/reference.mp4 \
+    --extra_params '{"edge": true}' \
+    --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml
+
+# Transfer with a fully synthetic control (no assets): generate an edge-map
+# video of a bouncing ball, then let the prompt paint it photoreal.
+# Keep synthetic controls edge-style: the blur hint expects the low
+# frequencies of natural video, and flat synthetic color fields degrade
+# generation. Temporal exposure swings (e.g. pulsing global light) do not
+# transfer — express lighting spatially or in the prompt instead.
+python generate_bouncing_ball_control.py --out_dir ./ball_control
+python cosmos3.py --model nvidia/Cosmos3-Nano \
+    --prompt "A photorealistic beach ball with colorful panels bouncing between the walls of an enclosed room, studio lighting." \
+    --extra_params '{"edge": "./ball_control/control.mp4"}' \
+    --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml
+
+# Multi-hint transfer: edge pins the layout, blur pins the palette/lighting.
+# Hints must describe the same underlying video as each other and the prompt.
+python cosmos3.py --model nvidia/Cosmos3-Nano \
+    --prompt "The same scene, ultra sharp, professional photography." \
+    --video_path /path/to/reference.mp4 \
+    --extra_params '{"edge": true, "blur": true}' \
+    --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml
+
 # 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
@@ -134,7 +165,7 @@ python cosmos3.py --model nvidia/Cosmos3-Edge \
     --image_path Cosmos3-Edge/assets/example_i2v_input.jpg \
     --output_path output.mp4
 
-# Inline prompt
+# 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" \
     --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml

examples/visual_gen/models/cosmos3/cosmos3.py

@@ -15,10 +15,10 @@
 # limitations under the License.
 """Cosmos3 Text(+Image/Video)-to-Video(+Audio) generation.
 
-One checkpoint serves T2V, T2I, I2V/TI2V, V2V and T2AV; ``prompts/`` holds a
-prompt file per mode and ``--help`` lists the flags. See ``README.md`` in this
-directory for the checkpoints, guardrail setup, deployment configs, and a
-worked command line per mode.
+One checkpoint serves T2V, T2I, I2V/TI2V, V2V, Transfer and T2AV;
+``prompts/`` holds a prompt file per mode and ``--help`` lists the flags.
+See ``README.md`` in this directory for the checkpoints, guardrail setup,
+deployment configs, and a worked command line per mode.
 """
 
 import argparse
@@ -28,6 +28,7 @@
 from typing import Any, Dict, Optional
 
 from tensorrt_llm import VisualGen, VisualGenArgs
+from tensorrt_llm._torch.visual_gen.models.cosmos3.transfer import TRANSFER_HINT_KEYS
 
 _SCRIPT_DIR = Path(__file__).resolve().parent
 
@@ -45,6 +46,48 @@ def _resolve_path(path: str) -> str:
     return path
 
 
+def _load_transfer_controls(extra_params: dict[str, Any]) -> None:
+    """Read precomputed transfer controls into ``control`` bytes, client-side.
+
+    A hint may name a control file (``{"edge": "ctrl.mp4"}`` or
+    ``{"edge": {"control_path": "ctrl.mp4"}}``); the worker only accepts encoded
+    bytes, so the media is read here.
+    """
+    for key in TRANSFER_HINT_KEYS:
+        hint = extra_params.get(key)
+        if isinstance(hint, str):
+            hint = {"control_path": hint}
+        if not isinstance(hint, dict):
+            continue
+        control_path = hint.pop("control_path", None)
+        if control_path is None:
+            continue
+        if not isinstance(control_path, str) or not control_path.strip():
+            raise ValueError(
+                f"--extra_params {key}.control_path must be a non-empty file path, "
+                f"got {control_path!r}."
+            )
+        hint["control"] = Path(_resolve_path(control_path)).read_bytes()
+        extra_params[key] = hint
+
+
+def _json_object(text: str) -> dict[str, Any]:
+    """Argparse type for a JSON *object*.
+
+    ``json.loads`` alone also accepts arrays, scalars and null, which then
+    either fail deep in the merge or, for ``[]``, succeed while doing nothing.
+    """
+    try:
+        value = json.loads(text)
+    except json.JSONDecodeError as exc:
+        raise argparse.ArgumentTypeError(f"not valid JSON: {exc}") from exc
+    if not isinstance(value, dict):
+        raise argparse.ArgumentTypeError(
+            f"expected a JSON object, got {type(value).__name__}: {text!r}"
+        )
+    return value
+
+
 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))
@@ -117,14 +160,14 @@ def resolve_negative_prompt(
 
 def resolve_prompt_and_options(
     *,
-    prompt: Optional[str],
-    prompt_file: Optional[str],
-    image_path: Optional[str],
+    prompt: str | None,
+    prompt_file: str | None,
+    image_path: str | None,
     enable_audio: bool,
     output_type: str,
-) -> tuple[str, Optional[str], bool, str]:
+) -> tuple[str, str | None, bool, str]:
     """Merge CLI args with optional prompt-file defaults."""
-    prompt_data: Dict[str, Any] = {}
+    prompt_data: dict[str, Any] = {}
     if prompt_file is not None:
         prompt_data = load_prompt_file(prompt_file)
 
@@ -240,6 +283,19 @@ def main():
     parser.add_argument(
         "--output_type", type=str, default="video", help="Output type (video, image)"
     )
+    parser.add_argument(
+        "--extra_params",
+        type=_json_object,
+        default=None,
+        help=(
+            "Model-specific extra params as a JSON object, merged last (overrides "
+            "flag-derived values). Keys are validated against the pipeline's "
+            "extra_param_specs. Transfer example: "
+            '\'{"edge": true, "blur": true, "control_guidance": 1.5}\' with --video_path, '
+            'or \'{"edge": "/path/control.mp4"}\' for a precomputed control (read here and '
+            "sent as encoded bytes)."
+        ),
+    )
 
     # Guardrails
     parser.add_argument(
@@ -282,6 +338,12 @@ def main():
 
     if args.video_path is not None:
         params.extra_params["video"] = Path(args.video_path).read_bytes()
+    if args.extra_params:
+        # Merged last: explicit JSON wins over flag-derived values.
+        params.extra_params.update(args.extra_params)
+    # The pipeline fits the output to the reference's aspect when height/width
+    # are unset, so there is nothing to do client-side.
+    _load_transfer_controls(params.extra_params)
 
     params.negative_prompt = negative_prompt
 

examples/visual_gen/models/cosmos3/generate_bouncing_ball_control.py

@@ -0,0 +1,124 @@
+#!/usr/bin/env python3
+# SPDX-FileCopyrightText: Copyright (c) 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.
+r"""Generate a synthetic edge-map control video: a ball bouncing off walls.
+
+Draws white outlines on black — a room border plus a ball following simple
+elastic-bounce physics — which is exactly what the Cosmos3 transfer ``edge``
+hint expects. No media assets required: the control is 30 lines of math, and
+transfer turns it into a photorealistic video whose subject follows the
+physics frame by frame.
+
+Generate the control, then run transfer with it:
+
+    python generate_bouncing_ball_control.py --out_dir ./ball_control
+
+    python cosmos3.py --model nvidia/Cosmos3-Nano \
+        --visual_gen_args ../../configs/cosmos3-nano-1gpu.yaml \
+        --prompt "A photorealistic beach ball with colorful panels bouncing \
+                  between the walls of an enclosed room, studio lighting." \
+        --extra_params '{"edge": "./ball_control/control.mp4"}' \
+        --output_path cosmos3_bouncing_ball.mp4
+
+Tip: keep synthetic controls edge-style. The ``blur`` hint expects the low
+frequencies of natural video; flat synthetic color fields are far from its
+training distribution and degrade generation quality.
+"""
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+import PIL.Image
+import PIL.ImageDraw
+import torch
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Bouncing-ball edge-control generator")
+    parser.add_argument("--out_dir", default="./ball_control")
+    parser.add_argument("--width", type=int, default=1280)
+    parser.add_argument("--height", type=int, default=720)
+    parser.add_argument("--num_frames", type=int, default=49)
+    parser.add_argument("--fps", type=int, default=24)
+    parser.add_argument("--radius", type=int, default=90)
+    parser.add_argument("--wall_inset", type=int, default=10, help="Room border inset in px")
+    parser.add_argument("--line_width", type=int, default=6)
+    parser.add_argument("--start", type=float, nargs=2, default=(300.0, 250.0))
+    parser.add_argument(
+        "--velocity",
+        type=float,
+        nargs=2,
+        default=(26.0, 19.0),
+        help="px/frame; the defaults bounce a few times over 49 frames",
+    )
+    parser.add_argument(
+        "--save_frames", action="store_true", help="Also write the individual PNG frames"
+    )
+    return parser.parse_args()
+
+
+def draw_frame(args: argparse.Namespace, x: float, y: float) -> PIL.Image.Image:
+    image = PIL.Image.new("RGB", (args.width, args.height), (0, 0, 0))
+    draw = PIL.ImageDraw.Draw(image)
+    inset, r = args.wall_inset, args.radius
+    draw.rectangle(
+        [inset, inset, args.width - inset, args.height - inset],
+        outline=(255, 255, 255),
+        width=args.line_width,
+    )
+    draw.ellipse([x - r, y - r, x + r, y + r], outline=(255, 255, 255), width=args.line_width)
+    return image
+
+
+def main() -> None:
+    args = parse_args()
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    lo_x, hi_x = args.wall_inset + args.radius, args.width - args.wall_inset - args.radius
+    lo_y, hi_y = args.wall_inset + args.radius, args.height - args.wall_inset - args.radius
+    x, y = args.start
+    vx, vy = args.velocity
+
+    frames = []
+    for i in range(args.num_frames):
+        frame = draw_frame(args, x, y)
+        frames.append(frame)
+        if args.save_frames:
+            frame.save(out_dir / f"frame_{i:03d}.png")
+        x, y = x + vx, y + vy
+        if x < lo_x or x > hi_x:
+            vx = -vx
+            x = max(lo_x, min(x, hi_x))
+        if y < lo_y or y > hi_y:
+            vy = -vy
+            y = max(lo_y, min(y, hi_y))
+
+    # Same encoder the pipeline saves its outputs through, so this needs
+    # nothing installed beyond what running a generation already needs.
+    from tensorrt_llm.media.encoding import save_video
+
+    video_path = out_dir / "control.mp4"
+    clip = torch.from_numpy(np.stack([np.asarray(frame) for frame in frames]))
+    save_video(clip, video_path, frame_rate=args.fps)
+
+    print(
+        f"Wrote {video_path}" + (f" and {args.num_frames} PNG frames" if args.save_frames else "")
+    )
+
+
+if __name__ == "__main__":
+    main()

examples/visual_gen/serve/README.md

@@ -316,7 +316,37 @@ Examples:
 - **LTX-2**: `stg_scale`, `stg_blocks`, `modality_scale`, `guidance_rescale`, `output_type`, ...
 - **Wan 2.2 A14B**: `guidance_scale_2`, `boundary_ratio`
 - **Wan 2.1 / Flux**: no model-specific `extra_params` declared
-- **Cosmos3**: `condition_video_latent_indexes`, `condition_video_keep` (V2V conditioning), `flow_shift`, `use_system_prompt`, ...
+- **Cosmos3**: `condition_video_latent_indexes`, `condition_video_keep` (V2V conditioning), `flow_shift`, `use_system_prompt`, and the transfer hints `edge`/`blur`/`depth`/`seg`/`wsm` with `control_guidance`, `control_guidance_interval`, `num_video_frames_per_chunk`, ... (see below)
+
+##### Cosmos3 transfer hints
+
+`extra_params` is JSON, so a control clip travels as a **base64-encoded** MP4/AVI
+string under `<hint>.control`; the server decodes it at the HTTP boundary. Only
+`edge` and `blur` can be auto-computed — pass `true` and supply a `video`
+reference for them to derive from. `depth`/`seg`/`wsm` have no generator, so
+they always need a control clip.
+
+```json
+{
+  "prompt": "a city street at dusk",
+  "extra_params": {
+    "video": "<base64 MP4/AVI>",
+    "edge": {"preset_edge_threshold": "medium"},
+    "blur": {"preset_blur_strength": "medium"},
+    "depth": {"control": "<base64 MP4/AVI>"},
+    "control_guidance": 1.5
+  }
+}
+```
+
+`preset_edge_threshold` and `preset_blur_strength` accept
+`none`/`very_low`/`low`/`medium`/`high`/`very_high` and default to `medium`; a
+bare `true` (or `"<base64>"`) is shorthand for the object form. Individual
+values are validated before the job is queued, so a bad preset or an
+unsupported frame count fails fast; combinations that only make sense together
+— a transfer option with no hint selected, or `edge`/`blur` asked to
+auto-compute with no `video` — are still reported by the worker, as a client
+error, once the request is running.
 
 > **Note:** LTX-2 generates video **with audio**. The `ltx2.yml` config must include
 > `text_encoder_path` pointing to a Gemma3 model (e.g., `google/gemma-3-12b-it`).