From 34a74cb3cd3d01252df0cd30e786a82acd0eb115 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Thu, 11 Jun 2026 08:11:55 -0700 Subject: [PATCH 01/64] cosmos3 action init Signed-off-by: Shreyas Misra --- examples/visual_gen/models/cosmos3/cosmos3.py | 288 +++++++++ .../prompts/action_forward_dynamics.json | 5 + .../prompts/action_inverse_dynamics.json | 4 + .../models/cosmos3/prompts/action_policy.json | 5 + .../visual_gen/models/cosmos3/action.py | 221 +++++++ .../visual_gen/models/cosmos3/defaults.py | 63 +- .../models/cosmos3/pipeline_cosmos3.py | 548 ++++++++++++++++-- .../models/cosmos3/transformer_cosmos3.py | 250 +++++++- tensorrt_llm/_torch/visual_gen/output.py | 34 +- tensorrt_llm/_torch/visual_gen/pipeline.py | 17 +- tensorrt_llm/visual_gen/output.py | 4 + .../test_cosmos3_transformer_parallel.py | 72 +++ .../_torch/visual_gen/test_cosmos3_action.py | 79 +++ .../visual_gen/test_cosmos3_pipeline.py | 143 ++++- .../visual_gen/test_cosmos3_transformer.py | 138 +++++ 15 files changed, 1795 insertions(+), 76 deletions(-) create mode 100644 examples/visual_gen/models/cosmos3/prompts/action_forward_dynamics.json create mode 100644 examples/visual_gen/models/cosmos3/prompts/action_inverse_dynamics.json create mode 100644 examples/visual_gen/models/cosmos3/prompts/action_policy.json create mode 100644 tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py create mode 100644 tests/unittest/_torch/visual_gen/test_cosmos3_action.py diff --git a/examples/visual_gen/models/cosmos3/cosmos3.py b/examples/visual_gen/models/cosmos3/cosmos3.py index de9e9e5010ad..6186efef5f73 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3.py +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -13,6 +13,114 @@ # 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"""Cosmos3 Text(+Image)-to-Video(+Audio) and action generation. + +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. +- **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). +- **Action** — policy / forward dynamics / inverse dynamics generation + (pass ``--action_mode``). Action and audio generation are mutually exclusive. + +Checkpoints (pass the Hub ID or local path via ``--model``): + +- `nvidia/Cosmos3-Nano `_ +- `nvidia/Cosmos3-Super `_ + +Guardrails are enabled by default (required by the +`NVIDIA Open Model License Agreement +`_). +Install and authenticate as follows:: + + pip install cosmos_guardrail==0.3.0 && pip uninstall opencv-python + +Accept the terms for the guardrail checkpoint at +https://huggingface.co/nvidia/Cosmos-1.0-Guardrail and set a valid ``HF_TOKEN`` +(the checkpoint is downloaded automatically on first run). + +To run without guardrails (you are responsible for safe deployment):: + + export TRTLLM_DISABLE_COSMOS3_GUARDRAILS=1 + +Deployment configs (``examples/visual_gen/configs/``): + +- ``cosmos3-nano-1gpu.yaml`` — 1 GPU +- ``cosmos3-super-4gpu.yaml`` — 4 GPU, CFG + Ulysses + parallel VAE + +Example prompts live under ``prompts/`` (mirroring ``cosmos3-internal/inputs/omni``). + +Usage:: + + # T2V: text-to-video + python cosmos3.py --model nvidia/Cosmos3-Nano \ + --prompt_file prompts/t2v.json \ + --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml + + # I2V/TI2V: image-conditioned video (vision_path is read from the prompt file; + # local path, file://, http(s):// URL, or data: URI are all accepted) + python cosmos3.py --model nvidia/Cosmos3-Nano \ + --prompt_file prompts/i2v.json \ + --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml + + # I2V with an explicit conditioning image (overrides the prompt file) + python cosmos3.py --model nvidia/Cosmos3-Nano \ + --prompt_file prompts/i2v.json \ + --image_path https://example.com/frame.jpg \ + --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml + + # T2AV: text-to-video with synchronized audio + python cosmos3.py --model nvidia/Cosmos3-Nano \ + --prompt_file prompts/t2av.json \ + --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml + + # T2I: text-to-image + python cosmos3.py --model nvidia/Cosmos3-Nano \ + --prompt_file prompts/t2i.json \ + --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml \ + --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" \ + --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml + + # Action — policy (first frame + instruction -> predicted action + rollout video) + python cosmos3.py --model nvidia/Cosmos3-Nano \ + --prompt_file prompts/action_policy.json \ + --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml \ + --action_mode policy \ + --domain_name bridge_orig_lerobot \ + --raw_action_dim 10 \ + --output_path policy_rollout.mp4 \ + --action_output_path policy_action.json + + # Action — forward dynamics (first frame + action trajectory -> rollout video) + python cosmos3.py --model nvidia/Cosmos3-Nano \ + --prompt_file prompts/action_forward_dynamics.json \ + --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml \ + --action_mode forward_dynamics \ + --domain_name av \ + --action_json action_trajectory.json \ + --output_path forward_dynamics.mp4 + + # Action — inverse dynamics (video -> predicted action) + python cosmos3.py --model nvidia/Cosmos3-Nano \ + --prompt_file prompts/action_inverse_dynamics.json \ + --video_path /path/to/frames/ \ + --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml \ + --action_mode inverse_dynamics \ + --domain_name bridge_orig_lerobot \ + --raw_action_dim 10 \ + --output_path inverse_video.mp4 \ + --action_output_path inverse_action.json +""" import argparse import json @@ -21,9 +129,12 @@ from typing import Any, Dict, Optional from tensorrt_llm import VisualGen, VisualGenArgs +from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import COSMOS3_ACTION_PARAMS _SCRIPT_DIR = Path(__file__).resolve().parent +ACTION_MODES = frozenset({"policy", "forward_dynamics", "inverse_dynamics"}) + def _resolve_path(path: str) -> str: candidate = Path(path) @@ -80,6 +191,93 @@ def resolve_prompt_and_options( return resolved_prompt, resolved_image, resolved_enable_audio, resolved_output_type +def _validate_action_args( + args: argparse.Namespace, resolved_image_path: Optional[str] = None +) -> None: + if args.action_mode is None: + return + + # The first frame may come from --image_path or a prompt file's vision_path. + has_first_frame = resolved_image_path is not None or args.video_path is not None + + mode = args.action_mode.strip().lower() + if mode not in ACTION_MODES: + raise SystemExit( + f"Invalid --action_mode {args.action_mode!r}; expected one of {sorted(ACTION_MODES)}." + ) + if args.enable_audio: + raise SystemExit("Cosmos3 does not support joint action and audio generation.") + if args.output_type != "video": + raise SystemExit("Action generation requires --output_type video.") + + if mode == "forward_dynamics": + if args.action_json is None: + raise SystemExit(f"{mode} requires --action_json.") + if not has_first_frame: + raise SystemExit( + f"{mode} requires --image_path, a prompt-file vision_path, or --video_path " + "for the first frame." + ) + elif mode == "policy": + if not has_first_frame: + raise SystemExit( + f"{mode} requires --image_path, a prompt-file vision_path, or --video_path " + "for the first frame." + ) + if args.raw_action_dim is None: + raise SystemExit(f"{mode} requires --raw_action_dim.") + elif mode == "inverse_dynamics": + if args.video_path is None: + raise SystemExit(f"{mode} requires --video_path (frame directory or image list).") + if args.raw_action_dim is None: + raise SystemExit(f"{mode} requires --raw_action_dim.") + + +def _apply_action_generation_params(params, args: argparse.Namespace) -> None: + """Set 480p action defaults on the request before pipeline overrides.""" + chunk = args.action_chunk_size or COSMOS3_ACTION_PARAMS["action_chunk_size"] + params.height = 480 + params.width = 832 + params.num_frames = chunk + 1 + params.num_inference_steps = COSMOS3_ACTION_PARAMS["num_inference_steps"] + params.guidance_scale = COSMOS3_ACTION_PARAMS["guidance_scale"] + params.frame_rate = COSMOS3_ACTION_PARAMS["frame_rate"] + params.extra_params["action_chunk_size"] = chunk + + +def _default_action_output_path(video_path: str) -> str: + stem = Path(video_path) + if stem.suffix: + return str(stem.with_name(f"{stem.stem}_action.json")) + return f"{video_path}_action.json" + + +def _save_action_output(output, path: str) -> None: + if output.action is None: + return + + action = output.action + if action.ndim == 3 and action.shape[0] == 1: + action_data = action[0].tolist() + shape = list(action.shape[1:]) + else: + action_data = action.tolist() + shape = list(action.shape) + + payload = { + "action_mode": output.action_mode, + "domain_id": output.domain_id, + "raw_action_dim": output.raw_action_dim, + "shape": shape, + "dtype": str(action.dtype).replace("torch.", ""), + "data": action_data, + } + out_path = Path(path) + out_path.parent.mkdir(parents=True, exist_ok=True) + with out_path.open("w", encoding="utf-8") as f: + json.dump(payload, f, indent=2) + + def main(): parser = argparse.ArgumentParser(description="Cosmos3 Text(+Image)-to-Video(+Audio) example") parser.add_argument( @@ -139,6 +337,62 @@ def main(): "--use_system_prompt", action="store_true", help="Use system prompt in prompt" ) parser.add_argument("--enable_audio", action="store_true", help="Enable audio generation") + parser.add_argument( + "--action_mode", + type=str, + default=None, + choices=sorted(ACTION_MODES), + help="Action mode: policy, forward_dynamics, or inverse_dynamics", + ) + parser.add_argument( + "--domain_name", + type=str, + default=None, + help="Embodiment domain name (e.g. bridge_orig_lerobot, av, droid_lerobot)", + ) + parser.add_argument( + "--domain_id", + type=int, + default=None, + help="Embodiment domain id (alternative to --domain_name)", + ) + parser.add_argument( + "--raw_action_dim", + type=int, + default=None, + help="Raw action DOF for policy/inverse_dynamics", + ) + parser.add_argument( + "--action_chunk_size", + type=int, + default=None, + help=f"Action tokens to generate (default {COSMOS3_ACTION_PARAMS['action_chunk_size']})", + ) + parser.add_argument( + "--action_json", + type=str, + default=None, + help="JSON file with action trajectory [T, D] for forward_dynamics", + ) + parser.add_argument( + "--video_path", + type=str, + default=None, + help="Frame directory or image path for inverse_dynamics (or first-frame fallback)", + ) + parser.add_argument( + "--action_resolution", + type=int, + default=480, + choices=[256, 480, 704, 720], + help="Resolution bucket for action image sizing", + ) + parser.add_argument( + "--action_output_path", + type=str, + default=None, + help="Path to save predicted action JSON (default: _action.json)", + ) parser.add_argument( "--output_type", type=str, default="video", help="Output type (video, image)" ) @@ -156,6 +410,7 @@ def main(): enable_audio=args.enable_audio, output_type=args.output_type, ) + _validate_action_args(args, resolved_image_path=image_path) # Engine config from shared YAML (optional); model-specific defaults apply otherwise. extra_args = VisualGenArgs.from_yaml(args.visual_gen_args) if args.visual_gen_args else None @@ -167,6 +422,9 @@ def main(): if image_path is not None: params.image = image_path + if args.action_mode is not None: + _apply_action_generation_params(params, args) + 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"): @@ -185,6 +443,23 @@ def main(): params.extra_params["enable_audio"] = enable_audio params.extra_params["use_guardrails"] = not args.disable_guardrails params.extra_params["output_type"] = output_type + params.extra_params["action_resolution"] = args.action_resolution + + if args.action_mode is not None: + params.extra_params["action_mode"] = args.action_mode + if args.domain_name is not None: + params.extra_params["domain_name"] = args.domain_name + if args.domain_id is not None: + params.extra_params["domain_id"] = args.domain_id + if args.raw_action_dim is not None: + params.extra_params["raw_action_dim"] = args.raw_action_dim + if args.action_chunk_size is not None: + params.extra_params["action_chunk_size"] = args.action_chunk_size + if args.action_json is not None: + with open(args.action_json, encoding="utf-8") as f: + params.extra_params["action"] = json.load(f) + if args.video_path is not None: + params.extra_params["video"] = args.video_path if negative_prompt is None: params.negative_prompt = None @@ -200,6 +475,19 @@ def main(): output.save(args.output_path) print(f"Saved: {args.output_path}") + + if args.action_mode is not None: + action_path = args.action_output_path or _default_action_output_path(args.output_path) + _save_action_output(output, action_path) + if output.action is not None: + print(f"Saved action: {action_path}") + print( + f"Action shape: {tuple(output.action.shape)}, " + f"raw_action_dim={output.raw_action_dim}, domain_id={output.domain_id}" + ) + else: + print("Warning: action_mode was set but the output carried no action tensor.") + print(output.metrics) diff --git a/examples/visual_gen/models/cosmos3/prompts/action_forward_dynamics.json b/examples/visual_gen/models/cosmos3/prompts/action_forward_dynamics.json new file mode 100644 index 000000000000..e780ef7b26b5 --- /dev/null +++ b/examples/visual_gen/models/cosmos3/prompts/action_forward_dynamics.json @@ -0,0 +1,5 @@ +{ + "model_mode": "image2video", + "prompt": "Robot manipulation rollout: the right arm extends to the fruit display, picks up a pear, and places it into the bag in the shopping cart.", + "vision_path": "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/vision/robot_153.jpg" +} diff --git a/examples/visual_gen/models/cosmos3/prompts/action_inverse_dynamics.json b/examples/visual_gen/models/cosmos3/prompts/action_inverse_dynamics.json new file mode 100644 index 000000000000..326e74cb77f5 --- /dev/null +++ b/examples/visual_gen/models/cosmos3/prompts/action_inverse_dynamics.json @@ -0,0 +1,4 @@ +{ + "model_mode": "image2video", + "prompt": "Recover the robot action trajectory from this clip of the arm picking up a pear and placing it into the bag." +} diff --git a/examples/visual_gen/models/cosmos3/prompts/action_policy.json b/examples/visual_gen/models/cosmos3/prompts/action_policy.json new file mode 100644 index 000000000000..1eb4e7572ee5 --- /dev/null +++ b/examples/visual_gen/models/cosmos3/prompts/action_policy.json @@ -0,0 +1,5 @@ +{ + "model_mode": "image2video", + "prompt": "Pick up the pear from the fruit display and place it into the plastic bag in the shopping cart.", + "vision_path": "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/vision/robot_153.jpg" +} diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py new file mode 100644 index 000000000000..60d1ea8bdbf6 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py @@ -0,0 +1,221 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Action-token helpers for Cosmos3 UVA/action generation.""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +import torch + +ACTION_MODE_POLICY = "policy" +ACTION_MODE_FORWARD_DYNAMICS = "forward_dynamics" +ACTION_MODE_INVERSE_DYNAMICS = "inverse_dynamics" +ACTION_MODES = { + ACTION_MODE_POLICY, + ACTION_MODE_FORWARD_DYNAMICS, + ACTION_MODE_INVERSE_DYNAMICS, +} + +EMBODIMENT_TO_DOMAIN_ID: dict[str, int] = { + "no_action": 0, + "av": 1, + "camera_pose": 2, + "hand_pose": 3, + "pusht": 4, + "libero": 5, + "umi": 6, + "bridge_orig_lerobot": 7, + "droid_lerobot": 8, + "robomind-franka": 8, + "galbot": 9, + "robomind-franka-dual": 12, + "robomind-ur": 13, + "agibotworld": 15, + "agibot_gear_gripper": 15, + "agibot_gear_gripper_ext": 15, + "fractal": 20, +} + +VIDEO_RES_SIZE_INFO: dict[str, dict[str, tuple[int, int]]] = { + "256": { + "1,1": (256, 256), + "4,3": (320, 256), + "3,4": (256, 320), + "16,9": (320, 192), + "9,16": (192, 320), + }, + "480": { + "1,1": (640, 640), + "4,3": (736, 544), + "3,4": (544, 736), + "16,9": (832, 480), + "9,16": (480, 832), + }, + "704": { + "1,1": (960, 960), + "4,3": (1088, 832), + "3,4": (832, 1088), + "16,9": (1280, 704), + "9,16": (704, 1280), + }, + "720": { + "1,1": (960, 960), + "4,3": (1104, 832), + "3,4": (832, 1104), + "16,9": (1280, 720), + "9,16": (720, 1280), + }, +} + + +def normalize_action_mode(mode: Any) -> str | None: + if mode is None: + return None + normalized = str(mode).strip().lower() + if not normalized: + return None + if normalized not in ACTION_MODES: + raise ValueError( + f"Unsupported Cosmos3 action_mode={mode!r}; expected one of {sorted(ACTION_MODES)}." + ) + return normalized + + +def resolve_domain_id( + *, + domain_id: Any = None, + domain_name: Any = None, + require_explicit: bool = False, +) -> int: + if domain_id is not None: + resolved = int(domain_id) + if resolved < 0: + raise ValueError(f"Cosmos3 domain_id must be non-negative, got {resolved}.") + return resolved + + if domain_name is None or str(domain_name).strip() == "": + if require_explicit: + raise ValueError( + "Cosmos3 action generation requires domain_id or non-empty domain_name." + ) + return 0 + + key = str(domain_name).strip().lower() + if key not in EMBODIMENT_TO_DOMAIN_ID: + raise ValueError( + f"Unknown Cosmos3 action domain_name={domain_name!r}; " + f"expected one of {sorted(EMBODIMENT_TO_DOMAIN_ID)} or pass domain_id directly." + ) + return EMBODIMENT_TO_DOMAIN_ID[key] + + +def action_condition_indexes(mode: str, action_length: int) -> list[int]: + mode = normalize_action_mode(mode) + if mode == ACTION_MODE_FORWARD_DYNAMICS: + return list(range(action_length)) + if mode in {ACTION_MODE_POLICY, ACTION_MODE_INVERSE_DYNAMICS}: + return [] + raise AssertionError(f"Unexpected action mode: {mode!r}") + + +def vision_condition_indexes( + mode: str, video_length: int, temporal_compression_factor: int +) -> list[int]: + mode = normalize_action_mode(mode) + latent_frames = (video_length - 1) // temporal_compression_factor + 1 + if mode in {ACTION_MODE_POLICY, ACTION_MODE_FORWARD_DYNAMICS}: + return [0] + if mode == ACTION_MODE_INVERSE_DYNAMICS: + return list(range(latent_frames)) + raise AssertionError(f"Unexpected action mode: {mode!r}") + + +def action_start_frame_offset(mode: str, action_length: int, video_length: int) -> int: + del mode + if action_length == video_length - 1: + return 1 + if action_length == video_length: + return 0 + raise ValueError( + "Cosmos3 action_chunk_size must equal num_frames - 1 or num_frames; " + f"got action_chunk_size={action_length}, num_frames={video_length}." + ) + + +def build_action_condition_mask( + mode: str, + action_length: int, + *, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + mask = torch.zeros(1, action_length, 1, device=device, dtype=dtype) + for idx in action_condition_indexes(mode, action_length): + mask[:, idx, :] = 1.0 + return mask + + +def build_vision_condition_mask( + mode: str, + video_length: int, + temporal_compression_factor: int, + *, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + latent_frames = (video_length - 1) // temporal_compression_factor + 1 + mask = torch.zeros(1, 1, latent_frames, 1, 1, device=device, dtype=dtype) + for idx in vision_condition_indexes(mode, video_length, temporal_compression_factor): + mask[:, :, idx, :, :] = 1.0 + return mask + + +def pad_action_to_dim(action: torch.Tensor, action_dim: int) -> torch.Tensor: + if action.shape[-1] > action_dim: + raise ValueError( + f"Cosmos3 action dimension {action.shape[-1]} exceeds model action_dim={action_dim}." + ) + if action.shape[-1] == action_dim: + return action + padding = torch.zeros( + *action.shape[:-1], action_dim - action.shape[-1], dtype=action.dtype, device=action.device + ) + return torch.cat([action, padding], dim=-1) + + +def load_action_tensor(action: Any = None) -> torch.Tensor: + if action is None: + raise ValueError( + "Cosmos3 forward_dynamics action mode requires an action tensor of shape [T, D]." + ) + if isinstance(action, torch.Tensor): + tensor = action.detach().to(dtype=torch.float32) + else: + tensor = torch.as_tensor(np.asarray(action), dtype=torch.float32) + if tensor.ndim == 3 and tensor.shape[0] == 1: + tensor = tensor.squeeze(0) + if tensor.ndim != 2: + raise ValueError(f"Cosmos3 action must have shape [T, D], got {tuple(tensor.shape)}.") + return tensor + + +def find_closest_target_size(h: int, w: int, resolution: str | int) -> tuple[int, int]: + key = str(resolution) + if key not in VIDEO_RES_SIZE_INFO: + raise ValueError( + f"Unknown Cosmos3 action resolution={resolution!r}; " + f"expected one of {sorted(VIDEO_RES_SIZE_INFO)}." + ) + input_ratio = h / w + best_size = None + best_diff = float("inf") + for cand_w, cand_h in VIDEO_RES_SIZE_INFO[key].values(): + diff = abs(input_ratio - cand_h / cand_w) + if diff < best_diff: + best_diff = diff + best_size = (cand_w, cand_h) + assert best_size is not None + return best_size diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index f5747544946d..780fc0f397c7 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -35,10 +35,16 @@ "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. +# Fields merged by the executor for every request. Modality-specific sizing and +# sampling defaults (height/width/num_frames/steps/guidance) stay ``None`` on the +# request until ``forward()`` resolves them from T2V/T2I/action context. +COSMOS3_PIPELINE_DEFAULTS = { + "max_sequence_length": COSMOS3_720P_PARAMS["max_sequence_length"], + "frame_rate": COSMOS3_720P_PARAMS["frame_rate"], +} + +# Text-to-image (``output_type="image"``) defaults. Applied when the request +# field is ``None``. COSMOS3_T2I_PARAMS = { "height": 1024, "width": 1024, @@ -48,6 +54,15 @@ "guidance_interval": (400.0, 1000.0), } +COSMOS3_ACTION_PARAMS = { + "action_chunk_size": 16, + "num_frames": 17, + "num_inference_steps": 30, + "guidance_scale": 1.0, + "flow_shift": 5.0, + "frame_rate": 24.0, +} + COSMOS3_EXTRA_SPECS: Dict[str, ExtraParamSchema] = { "use_duration_template": ExtraParamSchema( type="bool", @@ -79,4 +94,44 @@ default="video", description="Output modality: 'video' (T2V/I2V) or 'image' (text-to-image).", ), + "action_mode": ExtraParamSchema( + type="str", + default=None, + description="Action generation mode: policy, forward_dynamics, or inverse_dynamics.", + ), + "domain_name": ExtraParamSchema( + type="str", + default=None, + description="Embodiment domain name for action generation.", + ), + "domain_id": ExtraParamSchema( + type="int", + default=None, + description="Embodiment domain id for action generation.", + ), + "raw_action_dim": ExtraParamSchema( + type="int", + default=None, + description="Raw action DOF count for policy/inverse_dynamics.", + ), + "action_chunk_size": ExtraParamSchema( + type="int", + default=COSMOS3_ACTION_PARAMS["action_chunk_size"], + description="Number of action tokens to generate.", + ), + "action": ExtraParamSchema( + type="list", + default=None, + description="Action trajectory [T, D] for forward_dynamics mode.", + ), + "action_resolution": ExtraParamSchema( + type="int", + default=480, + description="Resolution bucket for action image sizing (256/480/704/720).", + ), + "video": ExtraParamSchema( + type="list", + default=None, + description="Video frames (PIL images or paths) for inverse_dynamics mode.", + ), } diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index 2dac241f2115..b50a36f87515 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -17,8 +17,10 @@ import math import os import time -from typing import List, Optional, Union +from pathlib import Path +from typing import Any, List, Optional, Union +import numpy as np import PIL.Image import torch from diffusers import AutoencoderKLWan, UniPCMultistepScheduler @@ -34,7 +36,25 @@ from tensorrt_llm.inputs.utils import load_image from tensorrt_llm.logger import logger -from .defaults import COSMOS3_720P_PARAMS, COSMOS3_EXTRA_SPECS, COSMOS3_T2I_PARAMS +from .action import ( + ACTION_MODE_FORWARD_DYNAMICS, + ACTION_MODE_INVERSE_DYNAMICS, + action_start_frame_offset, + build_action_condition_mask, + build_vision_condition_mask, + find_closest_target_size, + load_action_tensor, + normalize_action_mode, + pad_action_to_dim, + resolve_domain_id, +) +from .defaults import ( + COSMOS3_720P_PARAMS, + COSMOS3_ACTION_PARAMS, + COSMOS3_EXTRA_SPECS, + COSMOS3_PIPELINE_DEFAULTS, + COSMOS3_T2I_PARAMS, +) from .guardrails import check_video_safety, download_guardrail_checkpoint from .sound_tokenizer import LatentAutoEncoderV2 from .transformer_cosmos3 import Cosmos3VFMTransformer @@ -212,12 +232,56 @@ def default_warmup_num_frames(self): @property def default_generation_params(self): - return dict(COSMOS3_720P_PARAMS) + return dict(COSMOS3_PIPELINE_DEFAULTS) @property def extra_param_specs(self): return dict(COSMOS3_EXTRA_SPECS) + @staticmethod + def _resolve_action_size( + height: Optional[int], + width: Optional[int], + ref_image: PIL.Image.Image, + action_resolution: int, + ) -> tuple[int, int]: + """Fill unset action H/W from the action resolution bucket; honor explicit values.""" + if height is not None and width is not None: + return height, width + target_w, target_h = find_closest_target_size( + ref_image.height, ref_image.width, action_resolution + ) + return ( + height if height is not None else target_h, + width if width is not None else target_w, + ) + + def _action_reference_image( + self, + *, + action_mode: str, + image: Any, + video: Any, + ) -> PIL.Image.Image: + if action_mode == ACTION_MODE_INVERSE_DYNAMICS: + frames = self._normalize_action_video_input(video) + if not frames: + raise ValueError("Cosmos3 action_mode='inverse_dynamics' requires a video input.") + return self._pil_to_rgb(frames[0]) + + if image is None and video is not None: + frames = self._normalize_action_video_input(video) + return self._pil_to_rgb(frames[0]) + if image is None: + raise ValueError(f"Cosmos3 action_mode={action_mode!r} requires an image input.") + if isinstance(image, str): + return PIL.Image.open(image).convert("RGB") + if isinstance(image, PIL.Image.Image): + return image.convert("RGB") + raise TypeError( + f"Cosmos3 action reference image must be PIL.Image or path, got {type(image)!r}." + ) + def _run_warmup(self, height: int, width: int, num_frames: int, steps: int) -> None: with torch.no_grad(): self.forward( @@ -235,53 +299,19 @@ def _run_warmup(self, height: int, width: int, num_frames: int, steps: int) -> N enable_audio=False, ) - @staticmethod - def _resolve_t2i_default(merged_value, video_default, t2i_default): - """Pick the T2I default when the field still carries the merged video default. - - The executor merges a single ``default_generation_params`` dict (the - video params) into the request before ``infer()``, so an unspecified - field arrives equal to its video default. For T2I we substitute the - T2I default in that case while honoring any explicit user override. - """ - return t2i_default if merged_value == video_default else merged_value - def infer(self, req): extra_params = req.params.extra_params or {} output_type = extra_params.get("output_type", "video") - is_t2i = str(output_type).lower() == "image" - - height = req.params.height - width = req.params.width - num_inference_steps = req.params.num_inference_steps - guidance_scale = req.params.guidance_scale - if is_t2i: - height = self._resolve_t2i_default( - height, COSMOS3_720P_PARAMS["height"], COSMOS3_T2I_PARAMS["height"] - ) - width = self._resolve_t2i_default( - width, COSMOS3_720P_PARAMS["width"], COSMOS3_T2I_PARAMS["width"] - ) - num_inference_steps = self._resolve_t2i_default( - num_inference_steps, - COSMOS3_720P_PARAMS["num_inference_steps"], - COSMOS3_T2I_PARAMS["num_inference_steps"], - ) - guidance_scale = self._resolve_t2i_default( - guidance_scale, - COSMOS3_720P_PARAMS["guidance_scale"], - COSMOS3_T2I_PARAMS["guidance_scale"], - ) return self.forward( prompt=req.prompt, negative_prompt=req.params.negative_prompt, image=req.params.image, - height=height, - width=width, + height=req.params.height, + width=req.params.width, num_frames=req.params.num_frames, - num_inference_steps=num_inference_steps, - guidance_scale=guidance_scale, + num_inference_steps=req.params.num_inference_steps, + guidance_scale=req.params.guidance_scale, seed=req.params.seed, max_sequence_length=req.params.max_sequence_length, frame_rate=req.params.frame_rate, @@ -297,6 +327,16 @@ def infer(self, req): use_guardrails=extra_params.get("use_guardrails", True), enable_audio=extra_params.get("enable_audio", False), output_type=output_type, + action_mode=extra_params.get("action_mode"), + domain_name=extra_params.get("domain_name"), + domain_id=extra_params.get("domain_id"), + raw_action_dim=extra_params.get("raw_action_dim"), + action_chunk_size=extra_params.get("action_chunk_size"), + action=extra_params.get("action"), + action_resolution=extra_params.get("action_resolution") + or extra_params.get("image_size") + or 480, + video=extra_params.get("video"), ) def _apply_metadata_templates( @@ -608,6 +648,202 @@ def decode_audio(self, latent: torch.Tensor) -> torch.Tensor: """ return self.audio_tokenizer.decode(latent).float() # [B, audio_channels, N_samples] + @staticmethod + def _normalize_action_video_input(video: Any) -> List[Any]: + """Normalize inverse-dynamics video input to a frame list. + + Accepts a list of PIL images / paths, a single image path, or a + directory of frame images (sorted lexicographically). + """ + if video is None: + return [] + if isinstance(video, list): + if not video: + raise ValueError("Cosmos3 action video input must contain at least one frame.") + return video + if isinstance(video, str): + path = Path(video) + if not path.exists(): + raise ValueError(f"Cosmos3 action video path does not exist: {video}") + if path.is_dir(): + exts = {".png", ".jpg", ".jpeg", ".webp", ".bmp"} + frames = sorted(p for p in path.iterdir() if p.suffix.lower() in exts) + if not frames: + raise ValueError( + f"No image frames found in Cosmos3 action video directory: {video}" + ) + return [str(p) for p in frames] + return [video] + return [video] + + @staticmethod + def _pil_to_rgb(value: Any) -> PIL.Image.Image: + if isinstance(value, str): + return PIL.Image.open(value).convert("RGB") + if isinstance(value, PIL.Image.Image): + return value.convert("RGB") + raise TypeError( + f"Cosmos3 action preprocessing expected PIL image or image path, got {type(value)!r}." + ) + + @staticmethod + def _resize_and_pad_action_image( + image: PIL.Image.Image, target_h: int, target_w: int + ) -> PIL.Image.Image: + scale = min(target_w / image.width, target_h / image.height, 1.0) + resize_w = max(1, int(scale * image.width + 0.5)) + resize_h = max(1, int(scale * image.height + 0.5)) + if (resize_w, resize_h) != image.size: + image = image.resize((resize_w, resize_h), PIL.Image.Resampling.BICUBIC) + + array = np.asarray(image) + pad_h = target_h - resize_h + pad_w = target_w - resize_w + if pad_h < 0 or pad_w < 0: + raise ValueError( + f"Cosmos3 action image resize exceeded target size: resized={(resize_h, resize_w)}, " + f"target={(target_h, target_w)}." + ) + if pad_h == 0 and pad_w == 0: + return image + pad_mode = "reflect" if pad_h < resize_h and pad_w < resize_w else "edge" + padded = np.pad(array, ((0, pad_h), (0, pad_w), (0, 0)), mode=pad_mode) + return PIL.Image.fromarray(padded) + + def _preprocess_action_image( + self, image: PIL.Image.Image, target_h: int, target_w: int + ) -> torch.Tensor: + image = self._resize_and_pad_action_image(image, target_h, target_w) + return self.video_processor.preprocess(image, height=target_h, width=target_w) + + def _preprocess_action_video( + self, frames: List[Any], target_h: int, target_w: int + ) -> torch.Tensor: + if not frames: + raise ValueError("Cosmos3 action video input must contain at least one frame.") + processed = [ + self._preprocess_action_image(self._pil_to_rgb(frame), target_h, target_w).squeeze(0) + for frame in frames + ] + return torch.stack(processed, dim=1).unsqueeze(0).contiguous() + + def _encode_video_tensor(self, video_tensor: torch.Tensor) -> torch.Tensor: + """VAE-encode a preprocessed pixel video [1, 3, T, H, W].""" + if video_tensor.ndim == 4: + video_tensor = video_tensor.unsqueeze(0) + if video_tensor.ndim != 5 or video_tensor.shape[0] != 1 or video_tensor.shape[1] != 3: + raise ValueError( + f"Cosmos3 video tensor must have shape [1, 3, T, H, W], got {tuple(video_tensor.shape)}." + ) + + video = video_tensor.to(device=self.device, dtype=self.vae.dtype) + latent = self.vae.encode(video).latent_dist.mode() + + if hasattr(self.vae.config, "latents_mean") and hasattr(self.vae.config, "latents_std"): + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, -1, 1, 1, 1) + .to(latent.device, latent.dtype) + ) + latents_std = ( + torch.tensor(self.vae.config.latents_std) + .view(1, -1, 1, 1, 1) + .to(latent.device, latent.dtype) + ) + latent = (latent - latents_mean) / latents_std + else: + scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0) + latent = latent * scaling_factor + + return latent.to(self.dtype) + + def _prepare_latents_action_video( + self, + video_tensor: torch.Tensor, + mode: str, + num_frames: int, + generator: torch.Generator, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + C = self.transformer.latent_channel_size + T_lat = (num_frames - 1) // self.vae_scale_factor_temporal + 1 + H_lat = video_tensor.shape[-2] // self.vae_scale_factor_spatial + W_lat = video_tensor.shape[-1] // self.vae_scale_factor_spatial + + noise = randn_tensor( + (1, C, T_lat, H_lat, W_lat), + generator=generator, + device=self.device, + dtype=self.dtype, + ) + cond_latent = self._encode_video_tensor(video_tensor) + if cond_latent.shape[2:] != noise.shape[2:]: + raise ValueError( + "Cosmos3 action video latent shape mismatch: " + f"encoded={tuple(cond_latent.shape)}, expected={tuple(noise.shape)}." + ) + condition_mask = build_vision_condition_mask( + mode, + num_frames, + self.vae_scale_factor_temporal, + device=self.device, + dtype=self.dtype, + ) + latents = condition_mask * cond_latent + (1.0 - condition_mask) * noise + velocity_mask = 1.0 - condition_mask + return latents, velocity_mask, cond_latent + + def _prepare_action_latents( + self, + *, + mode: str, + action_chunk_size: int, + raw_action_dim: Optional[int], + generator: torch.Generator, + action_input: Any = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: + action_dim = int(getattr(self.transformer, "action_dim", 64)) + if mode == ACTION_MODE_FORWARD_DYNAMICS: + action = load_action_tensor(action_input) + if action.shape[0] < action_chunk_size: + pad = action[-1:].repeat(action_chunk_size - action.shape[0], 1) + action = torch.cat([action, pad], dim=0) + elif action.shape[0] > action_chunk_size: + action = action[:action_chunk_size] + if raw_action_dim is None: + raw_action_dim = int(action.shape[-1]) + clean_action = pad_action_to_dim(action, action_dim) + else: + if raw_action_dim is None: + raise ValueError( + "Cosmos3 action_mode='policy' and 'inverse_dynamics' require raw_action_dim." + ) + clean_action = torch.zeros(action_chunk_size, action_dim, dtype=torch.float32) + + raw_action_dim = int(raw_action_dim) + if raw_action_dim <= 0 or raw_action_dim > action_dim: + raise ValueError( + f"Cosmos3 raw_action_dim must be in [1, {action_dim}], got {raw_action_dim}." + ) + + clean_action = clean_action.to(device=self.device, dtype=self.dtype).unsqueeze(0) + condition_mask = build_action_condition_mask( + mode, + action_chunk_size, + device=self.device, + dtype=self.dtype, + ) + noise = randn_tensor( + (1, action_chunk_size, action_dim), + generator=generator, + device=self.device, + dtype=self.dtype, + ) + noise[:, :, raw_action_dim:] = 0 + clean_action[:, :, raw_action_dim:] = 0 + action_latents = condition_mask * clean_action + (1.0 - condition_mask) * noise + action_velocity_mask = 1.0 - condition_mask + return action_latents, action_velocity_mask, clean_action, raw_action_dim + # ========================================================================= # Forward (main generation entry point) # ========================================================================= @@ -620,19 +856,28 @@ def forward( seed: int, negative_prompt: Optional[str] = None, image: Optional[Union[PIL.Image.Image, torch.Tensor, str]] = None, - height: int = COSMOS3_720P_PARAMS["height"], - width: int = COSMOS3_720P_PARAMS["width"], - num_frames: int = COSMOS3_720P_PARAMS["num_frames"], - num_inference_steps: int = COSMOS3_720P_PARAMS["num_inference_steps"], - guidance_scale: float = COSMOS3_720P_PARAMS["guidance_scale"], - max_sequence_length: int = COSMOS3_720P_PARAMS["max_sequence_length"], - frame_rate: float = COSMOS3_720P_PARAMS["frame_rate"], + height: Optional[int] = None, + width: Optional[int] = None, + num_frames: Optional[int] = None, + num_inference_steps: Optional[int] = None, + guidance_scale: Optional[float] = None, + seed: int = 42, + max_sequence_length: Optional[int] = None, + frame_rate: Optional[float] = None, use_duration_template: bool = COSMOS3_EXTRA_SPECS["use_duration_template"].default, use_resolution_template: bool = COSMOS3_EXTRA_SPECS["use_resolution_template"].default, use_system_prompt: bool = COSMOS3_EXTRA_SPECS["use_system_prompt"].default, use_guardrails: bool = COSMOS3_EXTRA_SPECS["use_guardrails"].default, enable_audio: bool = COSMOS3_EXTRA_SPECS["enable_audio"].default, output_type: str = COSMOS3_EXTRA_SPECS["output_type"].default, + action_mode: Optional[str] = None, + domain_name: Optional[str] = None, + domain_id: Optional[int] = None, + raw_action_dim: Optional[int] = None, + action_chunk_size: Optional[int] = None, + action: Any = None, + action_resolution: int = 480, + video: Any = None, ): pipeline_start = time.time() timer = CudaPhaseTimer() @@ -640,6 +885,16 @@ def forward( use_guardrails = use_guardrails and not TRTLLM_DISABLE_COSMOS3_GUARDRAILS + normalized_action_mode = normalize_action_mode(action_mode) + do_action = normalized_action_mode is not None + if do_action and not self.action_gen: + raise ValueError( + "Cosmos3 action generation was requested, but this checkpoint " + "does not enable action_gen." + ) + if do_action and enable_audio: + raise ValueError("Cosmos3 does not support joint action and audio generation.") + # Text-to-image mode: same checkpoint/forward path as T2V, but a single # latent frame, image-flavored prompt templates, flow_shift=3.0, a CFG # guidance interval, and an image (rather than video) output. @@ -650,15 +905,76 @@ def forward( raise ValueError( "Cosmos3 text-to-image (output_type='image') does not accept an image input." ) + if do_action: + raise ValueError("Cosmos3 action generation does not support output_type='image'.") + if enable_audio: + raise ValueError("Cosmos3 audio generation does not support output_type='image'.") num_frames = 1 - enable_audio = False + height = height or COSMOS3_T2I_PARAMS["height"] + width = width or COSMOS3_T2I_PARAMS["width"] + num_inference_steps = num_inference_steps or COSMOS3_T2I_PARAMS["num_inference_steps"] + if guidance_scale is None: + guidance_scale = COSMOS3_T2I_PARAMS["guidance_scale"] guidance_interval = COSMOS3_T2I_PARAMS["guidance_interval"] self._set_flow_shift(COSMOS3_T2I_PARAMS["flow_shift"]) + elif do_action: + action_chunk_size = action_chunk_size or COSMOS3_ACTION_PARAMS["action_chunk_size"] + if num_frames is None: + num_frames = COSMOS3_ACTION_PARAMS["num_frames"] + num_inference_steps = ( + num_inference_steps or COSMOS3_ACTION_PARAMS["num_inference_steps"] + ) + if guidance_scale is None: + guidance_scale = COSMOS3_ACTION_PARAMS["guidance_scale"] + frame_rate = frame_rate or COSMOS3_ACTION_PARAMS["frame_rate"] + self._set_flow_shift(COSMOS3_ACTION_PARAMS["flow_shift"]) + enable_audio = False else: + height = height or COSMOS3_720P_PARAMS["height"] + width = width or COSMOS3_720P_PARAMS["width"] + num_frames = num_frames or COSMOS3_720P_PARAMS["num_frames"] + num_inference_steps = num_inference_steps or COSMOS3_720P_PARAMS["num_inference_steps"] + if guidance_scale is None: + guidance_scale = COSMOS3_720P_PARAMS["guidance_scale"] # Restore the checkpoint flow_shift in case a prior T2I request # rebuilt the scheduler with shift=3.0. self._set_flow_shift(getattr(self, "_engine_init_flow_shift", 1.0)) + max_sequence_length = max_sequence_length or COSMOS3_720P_PARAMS["max_sequence_length"] + if frame_rate is None: + frame_rate = COSMOS3_720P_PARAMS["frame_rate"] + + action_ref_image = None + if do_action: + action_ref_image = self._action_reference_image( + action_mode=normalized_action_mode, + image=image, + video=video, + ) + height, width = self._resolve_action_size( + height, width, action_ref_image, action_resolution + ) + + if self.rank == 0: + logger.info( + "Cosmos3 generation dims: %dx%d (WxH), num_frames=%d, " + "num_inference_steps=%d, guidance_scale=%.2f, frame_rate=%.1f", + width, + height, + num_frames, + num_inference_steps, + guidance_scale, + frame_rate, + ) + if do_action: + logger.info( + "Cosmos3 action dims: action_chunk_size=%d, action_resolution=%s, " + "input_aspect=%.3f", + action_chunk_size, + action_resolution, + action_ref_image.width / action_ref_image.height, + ) + if isinstance(prompt, str): prompt = [prompt] batch_size = len(prompt) @@ -753,7 +1069,67 @@ def forward( ) # 2. Prepare latents - if image is not None: + action_latents = None + action_velocity_mask = None + action_condition_latents = None + action_domain_id = None + action_frame_offset = 1 + resolved_raw_action_dim = raw_action_dim + condition_latents = None + + if do_action: + if action_chunk_size not in {num_frames, num_frames - 1}: + raise ValueError( + "Cosmos3 num_frames must equal action_chunk_size or action_chunk_size + 1." + ) + action_domain_id = resolve_domain_id( + domain_id=domain_id, + domain_name=domain_name, + require_explicit=True, + ) + action_frame_offset = action_start_frame_offset( + normalized_action_mode, action_chunk_size, num_frames + ) + + if normalized_action_mode == ACTION_MODE_INVERSE_DYNAMICS: + video = self._normalize_action_video_input(video) + video_tensor = self._preprocess_action_video(video, height, width) + latents, velocity_mask, condition_latents = self._prepare_latents_action_video( + video_tensor, + normalized_action_mode, + num_frames, + generator, + ) + image_latent = None + else: + image_tensor = self._preprocess_action_image(action_ref_image, height, width) + if image_tensor.ndim == 4: + video_tensor = ( + image_tensor.unsqueeze(2).expand(-1, -1, num_frames, -1, -1).contiguous() + ) + else: + video_tensor = image_tensor + latents, velocity_mask, condition_latents = self._prepare_latents_action_video( + video_tensor, + normalized_action_mode, + num_frames, + generator, + ) + image_latent = None + + ( + action_latents, + action_velocity_mask, + action_condition_latents, + resolved_raw_action_dim, + ) = self._prepare_action_latents( + mode=normalized_action_mode, + action_chunk_size=action_chunk_size, + raw_action_dim=raw_action_dim, + generator=generator, + action_input=action, + ) + elif image is not None: if isinstance(image, str): image = load_image(image, format="pil") @@ -801,6 +1177,12 @@ def forward( self.audio_scheduler.set_timesteps(num_inference_steps, device=self.device) # 4. Build forward_fn for the denoise loop + action_domain_ids_tensor = None + if do_action and action_domain_id is not None: + action_domain_ids_tensor = torch.tensor( + [action_domain_id], dtype=torch.long, device=self.device + ) + def forward_fn( latent_input, extra_stream_latents, @@ -815,6 +1197,16 @@ def forward_fn( rather than through encoder_hidden_states. """ current_audio = extra_stream_latents.get("audio") if extra_stream_latents else None + current_action = extra_stream_latents.get("action") if extra_stream_latents else None + + action_domain_ids = action_domain_ids_tensor + if ( + action_domain_ids is not None + and current_action is not None + and action_domain_ids.shape[0] == 1 + and current_action.shape[0] > 1 + ): + action_domain_ids = action_domain_ids.expand(current_action.shape[0]) result = self.transformer( hidden_states=latent_input, @@ -826,18 +1218,58 @@ def forward_fn( fps=frame_rate, noisy_frame_mask=velocity_mask, audio_latents=current_audio, + action_latents=current_action, + action_domain_ids=action_domain_ids, + action_noisy_mask=action_velocity_mask, + action_start_frame_offset=action_frame_offset, + action_fps=frame_rate, ) video_noise_pred = result.video audio_noise_pred = result.audio + action_noise_pred = result.action if velocity_mask is not None: video_noise_pred = video_noise_pred * velocity_mask + if action_noise_pred is not None: + if action_velocity_mask is not None: + action_noise_pred = action_noise_pred * action_velocity_mask + if ( + resolved_raw_action_dim is not None + and 0 < resolved_raw_action_dim < action_noise_pred.shape[-1] + ): + action_noise_pred = action_noise_pred.clone() + action_noise_pred[..., resolved_raw_action_dim:] = 0 + return video_noise_pred, {"action": action_noise_pred} + if audio_noise_pred is not None: return video_noise_pred, {"audio": audio_noise_pred} return video_noise_pred + def post_step_fn(step_latents, step_extra_stream_latents): + if velocity_mask is not None and condition_latents is not None: + step_latents = ( + velocity_mask * step_latents + (1.0 - velocity_mask) * condition_latents + ) + elif velocity_mask is not None and image_latent is not None: + step_latents = step_latents.clone() + step_latents[:, :, 0:1, :, :] = image_latent.to( + device=step_latents.device, dtype=step_latents.dtype + ) + if ( + action_velocity_mask is not None + and action_condition_latents is not None + and step_extra_stream_latents is not None + and "action" in step_extra_stream_latents + ): + action_key = step_extra_stream_latents["action"] + step_extra_stream_latents["action"] = ( + action_velocity_mask * action_key + + (1.0 - action_velocity_mask) * action_condition_latents + ) + return step_latents, step_extra_stream_latents + # 5. Build CFG tensors — text_ids and text_mask need to be split for CFG # BasePipeline.denoise batches [uncond, cond] when guidance_scale > 1 # We pass text IDs/masks through extra_cfg_tensors so they get split correctly @@ -850,7 +1282,12 @@ def forward_fn( # 6. Denoise timer.mark_denoise_start() - extra_streams = {"audio": (audio_latents, self.audio_scheduler)} if do_audio else None + extra_streams = None + if do_action: + extra_streams = {"action": (action_latents, self.scheduler)} + elif do_audio: + extra_streams = {"audio": (audio_latents, self.audio_scheduler)} + # FUTURE(action+audio): merge both keys; extend forward_fn return dict and post_step_fn. denoise_result = self.denoise( latents=latents, scheduler=self.scheduler, @@ -861,11 +1298,14 @@ def forward_fn( extra_cfg_tensors=extra_cfg_tensors, extra_streams=extra_streams, guidance_interval=guidance_interval, + post_step_fn=post_step_fn if do_action else None, ) if extra_streams is not None: latents, extra_latents = denoise_result audio_latents = extra_latents.get("audio") + if do_action: + action_latents = extra_latents.get("action") else: latents = denoise_result audio_latents = None @@ -911,5 +1351,11 @@ def forward_fn( audio_sample_rate=self.audio_tokenizer.model_config["sampling_rate"] if waveform is not None else None, + action=action_latents[:, :, :resolved_raw_action_dim].float().cpu() + if do_action and action_latents is not None + else None, + raw_action_dim=resolved_raw_action_dim if do_action else None, + action_mode=normalized_action_mode if do_action else None, + domain_id=action_domain_id if do_action else None, ) ) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py index b3df78a06ec1..166229b3eaa1 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -102,6 +102,7 @@ def compute_mrope_position_ids_vision( base_fps: float = 24.0, temporal_compression_factor: int = 4, enable_fps_modulation: bool = False, + start_frame_offset: int = 0, ) -> tuple[torch.Tensor, int | float]: """Generate 3D mRoPE position IDs for vision tokens. @@ -121,15 +122,17 @@ def compute_mrope_position_ids_vision( base_tps = base_fps / temporal_compression_factor frame_indices = torch.arange(grid_t, dtype=torch.float32) t_index = ( - (frame_indices / tps * base_tps + temporal_offset) + ((frame_indices + start_frame_offset) / tps * base_tps + temporal_offset) .view(-1, 1) .expand(-1, grid_h * grid_w) .flatten() ) else: - t_index = torch.arange(grid_t, dtype=torch.long).view(-1, 1).expand( - -1, grid_h * grid_w - ).flatten() + int(temporal_offset) + t_index = ( + torch.arange(grid_t, dtype=torch.long).view(-1, 1).expand(-1, grid_h * grid_w).flatten() + + int(temporal_offset) + + start_frame_offset + ) h_index = ( torch.arange(grid_h, dtype=torch.long).view(1, -1, 1).expand(grid_t, -1, grid_w).flatten() @@ -149,6 +152,75 @@ def compute_mrope_position_ids_vision( return mrope_ids, next_offset +def compute_mrope_position_ids_action( + grid_t: int, + temporal_offset: int | float, + action_fps: float | None, + base_fps: float = 24.0, + base_temporal_compression_factor: int = 4, + enable_fps_modulation: bool = True, + start_frame_offset: int = 1, +) -> tuple[torch.Tensor, int | float]: + """Generate mRoPE IDs for action tokens as a frame-rate (T, 1, 1) grid.""" + return compute_mrope_position_ids_vision( + grid_t=grid_t, + grid_h=1, + grid_w=1, + temporal_offset=temporal_offset, + fps=action_fps, + base_fps=base_fps, + temporal_compression_factor=1, + enable_fps_modulation=enable_fps_modulation, + start_frame_offset=start_frame_offset, + ) + + +class DomainAwareLinear(nn.Module): + """Linear projection with one weight/bias pair per action embodiment domain.""" + + def __init__( + self, + input_size: int, + output_size: int, + num_domains: int, + *, + dtype: torch.dtype = torch.bfloat16, + ) -> None: + super().__init__() + self.input_size = int(input_size) + self.output_size = int(output_size) + self.num_domains = int(num_domains) + self.fc = nn.Embedding(self.num_domains, self.output_size * self.input_size, dtype=dtype) + self.bias = nn.Embedding(self.num_domains, self.output_size, dtype=dtype) + nn.init.xavier_uniform_(self.fc.weight) + nn.init.zeros_(self.bias.weight) + + def forward(self, x: torch.Tensor, domain_id: torch.Tensor) -> torch.Tensor: + if domain_id.ndim == 0: + domain_id = domain_id.unsqueeze(0) + domain_id = domain_id.to(device=x.device, dtype=torch.long).reshape(-1) + if x.shape[0] != domain_id.shape[0]: + raise ValueError( + "Cosmos3 action domain_id batch size must match action batch: " + f"tokens={x.shape[0]}, domain_id={domain_id.shape[0]}." + ) + if torch.any((domain_id < 0) | (domain_id >= self.num_domains)): + raise ValueError( + f"Cosmos3 action domain_id must be in [0, {self.num_domains}), " + f"got {domain_id.tolist()}." + ) + + weight = self.fc(domain_id).view(domain_id.shape[0], self.input_size, self.output_size) + bias = self.bias(domain_id).view(domain_id.shape[0], self.output_size) + if x.ndim == 2: + return torch.bmm(x.unsqueeze(1), weight).squeeze(1) + bias + if x.ndim == 3: + return torch.bmm(x, weight) + bias.unsqueeze(1) + raise ValueError( + f"Cosmos3 DomainAwareLinear expected rank-2 or rank-3 input, got {tuple(x.shape)}." + ) + + class TimestepEmbedder(nn.Module): """ Embeds scalar timesteps into vector representations. @@ -731,6 +803,29 @@ def __init__(self, model_config: DiffusionModelConfig): pretrained_config.temporal_compression_factor_sound ) + if self.action_gen: + action_dim_value = getattr(pretrained_config, "action_dim", None) + if action_dim_value is None: + action_dim_value = getattr(pretrained_config, "max_action_dim", 64) + self.action_dim = int(action_dim_value) + self.num_embodiment_domains = int( + getattr(pretrained_config, "num_embodiment_domains", 32) + ) + dtype = torch.bfloat16 + self.action_proj_in = DomainAwareLinear( + self.action_dim, + self.hidden_size, + self.num_embodiment_domains, + dtype=dtype, + ) + self.action_proj_out = DomainAwareLinear( + self.hidden_size, + self.action_dim, + self.num_embodiment_domains, + dtype=dtype, + ) + self.action_modality_embed = nn.Parameter(torch.zeros(self.hidden_size, dtype=dtype)) + if pretrained_config.position_embedding_type != "unified_3d_mrope": raise ValueError( f"Position embedding type {pretrained_config.position_embedding_type} not supported" @@ -964,6 +1059,59 @@ def unpack_audio_latents(self, hidden_audio: torch.Tensor) -> torch.Tensor: """[B, T_audio, audio_dim] → [B, audio_dim, T_audio].""" return hidden_audio.permute(0, 2, 1) + # ------------------------------------------------------------------------- + # Action helpers + # ------------------------------------------------------------------------- + + def _compute_action_rope_freqs( + self, + T_action: int, + text_mask: torch.Tensor, + action_fps: float, + action_start_frame_offset: int, + device: torch.device, + dtype: torch.dtype, + ) -> Tuple[torch.Tensor, torch.Tensor]: + B = text_mask.shape[0] + text_lengths = text_mask.sum(dim=1).long() + + action_pos_list = [] + for b in range(B): + real_len = int(text_lengths[b].item()) + _, t_offset = compute_mrope_position_ids_text(real_len, temporal_offset=0) + a_pos, _ = compute_mrope_position_ids_action( + T_action, + temporal_offset=t_offset + self.unified_3d_mrope_temporal_modality_margin, + action_fps=action_fps, + base_fps=self.base_fps, + base_temporal_compression_factor=self.temporal_compression_factor, + enable_fps_modulation=self.enable_fps_modulation, + start_frame_offset=action_start_frame_offset, + ) + action_pos_list.append(a_pos) + + action_pos_ids = torch.stack(action_pos_list, dim=1).to(device) + rotary_emb = self.language_model.rotary_emb + _dummy = torch.tensor([], dtype=dtype, device=device) + cos_a, sin_a = rotary_emb(_dummy, position_ids=action_pos_ids) + return cos_a.unsqueeze(2), sin_a.unsqueeze(2) + + def pack_action(self, action_latents: torch.Tensor) -> torch.Tensor: + if action_latents.ndim != 3: + raise ValueError( + f"Cosmos3 action latents must have shape [B, T, D], got {tuple(action_latents.shape)}." + ) + if action_latents.shape[-1] != self.action_dim: + raise ValueError( + f"Cosmos3 action latent dimension mismatch: expected {self.action_dim}, " + f"got {action_latents.shape[-1]}." + ) + return action_latents.contiguous() + + @staticmethod + def unpack_action(tokens: torch.Tensor) -> torch.Tensor: + return tokens + def reset_cache(self): self.cached_kv = None self.cached_freqs_gen = None @@ -979,6 +1127,11 @@ def forward( fps: float | None = None, noisy_frame_mask: torch.Tensor | None = None, audio_latents: Optional[torch.Tensor] = None, + action_latents: Optional[torch.Tensor] = None, + action_domain_ids: Optional[torch.Tensor] = None, + action_noisy_mask: Optional[torch.Tensor] = None, + action_start_frame_offset: int = 1, + action_fps: float | None = None, **kwargs, ) -> "TransformerOutput": """ @@ -1006,13 +1159,26 @@ def forward( Returns: TransformerOutput with video (and image alias) always set. audio is set to the predicted audio velocity when audio_latents is - provided; otherwise None. action is always None for now. + provided; otherwise None. action is set when action_latents is provided. """ del kwargs # Kept for diffusers API compatibility. +<<<<<<< HEAD if timestep is None: raise ValueError("Cosmos3VFMTransformer.forward requires normalized timestep.") if raw_timestep is None: raise ValueError("Cosmos3VFMTransformer.forward requires raw_timestep.") +======= + + if action_latents is not None and audio_latents is not None: + raise ValueError( + "Cosmos3 transformer does not support joint action and audio generation." + ) + if action_latents is not None and not self.action_gen: + raise ValueError( + "Cosmos3 action generation was requested, but this transformer " + "was initialized without action modules." + ) +>>>>>>> 270db1b2f7 (cosmos3 action init) T, H, W = video_shape Hp, Wp, _, _ = self._pad_to_patch_size(H, W) max_real_len = text_mask.sum(dim=1).max().item() @@ -1075,10 +1241,44 @@ def forward( else: self.cached_kv = cached_kv_full - # --- Audio token injection ------------------------------------------------- + # --- Extra modality token injection (mutually exclusive: action OR audio) --- T_vid_tokens = hidden_gen.shape[1] # T * Hp * Wp + T_action = 0 T_audio = 0 - if audio_latents is not None and self.audio_gen: + action_domain_ids_tensor = action_domain_ids + if action_latents is not None and self.action_gen: + # FUTURE(action+audio): concat order is video|action|audio; adjust slices below. + if action_domain_ids_tensor is None: + action_domain_ids_tensor = torch.zeros( + action_latents.shape[0], dtype=torch.long, device=action_latents.device + ) + T_action = action_latents.shape[1] + hidden_action = self.action_proj_in( + self.pack_action(action_latents), action_domain_ids_tensor + ) + hidden_action = hidden_action + self.action_modality_embed.to(hidden_action.dtype) + if action_noisy_mask is None: + hidden_action = hidden_action + time_embed.unsqueeze(1) + else: + hidden_action = hidden_action + time_embed.unsqueeze(1) * action_noisy_mask.to( + hidden_action.dtype + ) + hidden_gen = torch.cat([hidden_gen, hidden_action], dim=1) + effective_action_fps = action_fps if action_fps is not None else (fps or self.base_fps) + cos_a, sin_a = self._compute_action_rope_freqs( + T_action, + text_mask, + float(effective_action_fps), + action_start_frame_offset, + hidden_states.device, + hidden_gen.dtype, + ) + cos_v, sin_v = self.cached_freqs_gen + freqs_gen_combined = ( + torch.cat([cos_v, cos_a], dim=1), + torch.cat([sin_v, sin_a], dim=1), + ) + elif audio_latents is not None and self.audio_gen: T_audio = audio_latents.shape[2] hidden_audio = self.pack_audio_latents(audio_latents).to(hidden_gen.dtype) hidden_audio = self.audio2llm(hidden_audio) + self.audio_modality_embed @@ -1090,7 +1290,6 @@ def forward( hidden_states.device, hidden_gen.dtype, ) - # [B, T_vid+T_audio, hidden_size] hidden_gen = torch.cat([hidden_gen, hidden_audio], dim=1) cos_v, sin_v = self.cached_freqs_gen freqs_gen_combined = ( @@ -1137,16 +1336,27 @@ def forward( # --- Decode video velocity ------------------------------------------------ video_vel = self.unpatchify(self.llm2vae(hidden_gen[:, :T_vid_tokens]), T, H, W) - # --- Decode audio velocity (if requested) --------------------------------- + # --- Decode extra-modality velocity (action XOR audio; follows video) --- + extra_start = T_vid_tokens audio_vel = None if T_audio > 0 and audio_latents is not None and self.audio_gen: - # hidden_gen[:, T_vid_tokens:] → [B, T_audio, hidden_size] - # → llm2audio → [B, T_audio, audio_dim] → unpack → [B, audio_dim, T_audio] audio_vel = self.unpack_audio_latents( - self.llm2audio(hidden_gen[:, T_vid_tokens : T_vid_tokens + T_audio]) + self.llm2audio(hidden_gen[:, extra_start : extra_start + T_audio]) + ) + + action_vel = None + if T_action > 0 and action_latents is not None and self.action_gen: + assert action_domain_ids_tensor is not None + action_vel = self.unpack_action( + self.action_proj_out( + hidden_gen[:, extra_start : extra_start + T_action], + action_domain_ids_tensor, + ) ) - return TransformerOutput(video=video_vel, image=video_vel, audio=audio_vel) + return TransformerOutput( + video=video_vel, image=video_vel, audio=audio_vel, action=action_vel + ) def load_weights(self, weights: dict) -> None: """Load weights with key remapping from Cosmos3-Nano / Diffusers checkpoints. @@ -1158,8 +1368,7 @@ def load_weights(self, weights: dict) -> None: remapped = {} skip_prefixes = ( "lm_head.", - "action_modality_embed", - "action_proj_", + "action_pos_embed.", ) for key, value in weights.items(): @@ -1194,6 +1403,14 @@ def load_weights(self, weights: dict) -> None: remapped[k] = value continue + if k.startswith("action_modality_embed"): + remapped[k] = value + continue + + if k.startswith("action_proj_in.") or k.startswith("action_proj_out."): + remapped[k] = value + continue + if k.startswith("time_embedder.linear"): k = k.replace("time_embedder.linear_1.", "time_embedder.mlp.linear_1.") k = k.replace("time_embedder.linear_2.", "time_embedder.mlp.linear_2.") @@ -1326,6 +1543,9 @@ def post_load_weights(self) -> None: self.llm2audio.to(target_dtype) self.audio_modality_embed.data = self.audio_modality_embed.data.to(target_dtype) + if self.action_gen: + self.action_modality_embed.data = self.action_modality_embed.data.to(target_dtype) + for _, module in self.named_modules(): if isinstance(module, Linear) or isinstance(module, Qwen3VLTextRMSNorm): module.post_load_weights() diff --git a/tensorrt_llm/_torch/visual_gen/output.py b/tensorrt_llm/_torch/visual_gen/output.py index b77956b7b6d8..ece7fd203f3f 100644 --- a/tensorrt_llm/_torch/visual_gen/output.py +++ b/tensorrt_llm/_torch/visual_gen/output.py @@ -37,8 +37,9 @@ class PipelineOutput: """Internal per-pipeline output. Each pipeline ``infer()`` populates the media tensor it produces plus - the metadata it owns (``frame_rate``, ``audio_sample_rate``) and the - three CUDA-event-measured timing phases that decompose ``pipeline.infer()``. + the metadata it owns (``frame_rate``, ``audio_sample_rate``, action fields) + and the three CUDA-event-measured timing phases that decompose + ``pipeline.infer()``. Attributes: image: Generated image as ``torch.Tensor`` shape ``(B, H, W, C)``, @@ -51,12 +52,25 @@ class PipelineOutput: ``(B, channels, T_audio)``, dtype ``float32``. Populated by LTX-2. The leading batch dim is always present, even for single-prompt requests (size 1). + action: Predicted or refined action trajectory as ``torch.Tensor`` + shape ``(B, T_action, D_raw)``, dtype ``float32``. Populated by + Cosmos3 action generation (``policy`` / ``forward_dynamics`` / + ``inverse_dynamics``). Sliced to ``raw_action_dim`` DOF; no VAE + decode. ``None`` when action generation was not requested. frame_rate: Video frame rate in fps. Populated by video pipelines (Wan T2V/I2V emit ``16.0``; LTX-2 emits ``params.frame_rate``). ``None`` for image-only pipelines. audio_sample_rate: Audio sample rate in Hz. Populated by LTX-2 from its audio config (no hard-coded literal). ``None`` for pipelines without audio. + raw_action_dim: Number of action degrees of freedom in the trailing + dimension of ``action``. Populated by Cosmos3 action generation. + ``None`` when ``action`` is ``None``. + action_mode: Requested Cosmos3 action mode (``policy``, + ``forward_dynamics``, or ``inverse_dynamics``). ``None`` when + action generation was not requested. + domain_id: Resolved embodiment domain id for Cosmos3 action + generation. ``None`` when action generation was not requested. pre_denoise: Wall-clock GPU-stream time (seconds) before the denoising loop (text encoding, latent prep, conditioning), measured by CUDA events. ``0.0`` if not measured. @@ -72,8 +86,12 @@ class PipelineOutput: image: Optional[torch.Tensor] = None video: Optional[torch.Tensor] = None audio: Optional[torch.Tensor] = None + action: Optional[torch.Tensor] = None frame_rate: Optional[float] = None audio_sample_rate: Optional[int] = None + raw_action_dim: Optional[int] = None + action_mode: Optional[str] = None + domain_id: Optional[int] = None pre_denoise: float = 0.0 denoise: float = 0.0 post_denoise: float = 0.0 @@ -201,8 +219,12 @@ def to_visual_gen_output(resp: "DiffusionResponse") -> "VisualGenOutput": image=out.image, video=out.video, audio=out.audio, + action=out.action, frame_rate=out.frame_rate, audio_sample_rate=out.audio_sample_rate, + raw_action_dim=out.raw_action_dim, + action_mode=out.action_mode, + domain_id=out.domain_id, metrics=metrics, ) @@ -245,6 +267,10 @@ def split_visual_gen_output(resp: "DiffusionResponse", batch_size: int) -> List[ assert out.audio.shape[0] == batch_size, ( f"audio leading dim {out.audio.shape[0]} != batch_size {batch_size}" ) + if out.action is not None: + assert out.action.shape[0] == batch_size, ( + f"action leading dim {out.action.shape[0]} != batch_size {batch_size}" + ) metrics = VisualGenMetrics( generation=resp.generation, pre_denoise=out.pre_denoise, @@ -259,8 +285,12 @@ def split_visual_gen_output(resp: "DiffusionResponse", batch_size: int) -> List[ image=out.image[i] if out.image is not None else None, video=out.video[i] if out.video is not None else None, audio=out.audio[i] if out.audio is not None else None, + action=out.action[i] if out.action is not None else None, frame_rate=out.frame_rate, audio_sample_rate=out.audio_sample_rate, + raw_action_dim=out.raw_action_dim, + action_mode=out.action_mode, + domain_id=out.domain_id, metrics=metrics, ) ) diff --git a/tensorrt_llm/_torch/visual_gen/pipeline.py b/tensorrt_llm/_torch/visual_gen/pipeline.py index 088548581e4c..f460ab8023df 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline.py @@ -1,4 +1,5 @@ import contextlib +import inspect import itertools import os import time @@ -1096,8 +1097,14 @@ def denoise( guidance_interval: Optional ``(lo, hi)`` scheduler-timestep range in which CFG is active. Outside the interval the effective scale is 1.0 (conditional prediction only); both branches still run. - post_step_fn: Optional callable applied to latents after each scheduler step. - Signature: post_step_fn(latents) -> latents + post_step_fn: Optional callable applied after each scheduler step. + If the callable accepts one argument, it is invoked as + ``post_step_fn(latents) -> latents``. If it accepts two + or more, it is invoked as + ``post_step_fn(latents, extra_stream_latents) -> + (latents, extra_stream_latents)``, where + ``extra_stream_latents`` is the dict of parallel streams + from ``extra_streams`` (e.g. ``{"action": ...}``). Use for constraints that must hold throughout denoising. Returns: @@ -1228,7 +1235,11 @@ def denoise( ) if post_step_fn is not None: - latents = post_step_fn(latents) + sig = inspect.signature(post_step_fn) + if len(sig.parameters) >= 2: + latents, extra_stream_latents = post_step_fn(latents, extra_stream_latents) + else: + latents = post_step_fn(latents) # Logging if self.rank == 0: diff --git a/tensorrt_llm/visual_gen/output.py b/tensorrt_llm/visual_gen/output.py index 0660e661122f..9d1f6e99336a 100644 --- a/tensorrt_llm/visual_gen/output.py +++ b/tensorrt_llm/visual_gen/output.py @@ -96,8 +96,12 @@ class VisualGenOutput: image: Optional[torch.Tensor] = None video: Optional[torch.Tensor] = None audio: Optional[torch.Tensor] = None + action: Optional[torch.Tensor] = None frame_rate: Optional[float] = None audio_sample_rate: Optional[int] = None + raw_action_dim: Optional[int] = None + action_mode: Optional[str] = None + domain_id: Optional[int] = None error: Optional[str] = None metrics: Optional[VisualGenMetrics] = None diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_cosmos3_transformer_parallel.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_cosmos3_transformer_parallel.py index e5b770f5dfa8..f5dab83047d4 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_cosmos3_transformer_parallel.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_cosmos3_transformer_parallel.py @@ -129,6 +129,15 @@ temporal_compression_factor_sound=1, ) +_ACTION_DIM = 64 +_T_ACTION = 4 +_COSMOS3_ACTION_CONFIG = dict( + **_COSMOS3_TEST_CONFIG, + action_gen=True, + action_dim=_ACTION_DIM, + num_embodiment_domains=32, +) + SEED_WEIGHTS = 123 SEED_INPUT = 456 SEED_COND_TEXT = 42 @@ -426,6 +435,33 @@ def _forward_with_audio( return out.video, out.audio +def _forward_with_action( + model: Cosmos3VFMTransformer, device: torch.device, text_seed: int +) -> Tuple[torch.Tensor, torch.Tensor]: + channels = _COSMOS3_TEST_CONFIG["latent_channel"] + hs, ts, text_ids, text_mask, video_shape = _cosmos3_inputs( + device, channels=channels, text_seed=text_seed + ) + torch.manual_seed(SEED_INPUT + 2) + action_latents = ( + torch.randn(hs.shape[0], _T_ACTION, _ACTION_DIM, device=device, dtype=hs.dtype) * 0.1 + ) + domain_ids = torch.tensor([7], dtype=torch.long, device=device) + model.reset_cache() + with torch.inference_mode(): + out = model( + hidden_states=hs, + timestep=ts, + text_ids=text_ids, + text_mask=text_mask, + video_shape=video_shape, + fps=_FPS, + action_latents=action_latents, + action_domain_ids=domain_ids, + ) + return out.video, out.action + + def _build_ref_and_parallel( *, tp_size: int = 1, @@ -562,6 +598,37 @@ def _logic_cosmos3_ulysses_audio_vs_single_gpu(rank, world_size): ) +def _logic_cosmos3_ulysses_action_vs_single_gpu(rank, world_size): + ref_model, ulysses_model, _, device = _build_ref_and_parallel( + ulysses_size=world_size, pretrained_dict=_COSMOS3_ACTION_CONFIG + ) + text_seed = _cfg_text_seed(rank, tp_size=1, ulysses_size=world_size, cfg_size=1) + + ref_video, ref_action = _forward_with_action(ref_model, device, text_seed) + ulysses_video, ulysses_action = _forward_with_action(ulysses_model, device, text_seed) + + if rank == 0: + vdiff = (ulysses_video.float() - ref_video.float()).abs() + adiff = (ulysses_action.float() - ref_action.float()).abs() + print( + f"[ulysses={world_size}+action] " + f"video max_abs_diff={vdiff.max().item():.6e}, " + f"action max_abs_diff={adiff.max().item():.6e}", + flush=True, + ) + + _assert_parity( + ulysses_video, + ref_video, + msg=f"Rank {rank}: Ulysses+action VIDEO differs from single-GPU reference", + ) + _assert_parity( + ulysses_action, + ref_action, + msg=f"Rank {rank}: Ulysses+action ACTION differs from single-GPU reference", + ) + + def _logic_cosmos3_tp_ulysses_vs_single_gpu(rank, world_size): tp_size = 2 ulysses_size = 2 @@ -711,6 +778,11 @@ def test_ulysses2_audio_vs_single_gpu(self): self._skip_if_unavailable() run_test_in_distributed(world_size=2, test_fn=_logic_cosmos3_ulysses_audio_vs_single_gpu) + def test_ulysses2_action_vs_single_gpu(self): + """Ulysses parity with action tokens appended to the GEN sequence.""" + self._skip_if_unavailable() + run_test_in_distributed(world_size=2, test_fn=_logic_cosmos3_ulysses_action_vs_single_gpu) + @pytest.mark.gpu4 def test_tp2_ulysses2_vs_single_gpu(self): self._skip_if_unavailable() diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py new file mode 100644 index 000000000000..b85b1d50b75c --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py @@ -0,0 +1,79 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for Cosmos3 action sizing helpers (no checkpoint / GPU required). + +Run: + pytest tests/unittest/_torch/visual_gen/test_cosmos3_action.py -v +""" + +import PIL.Image +import pytest + +from tensorrt_llm._torch.visual_gen.models.cosmos3.action import ( + VIDEO_RES_SIZE_INFO, + find_closest_target_size, +) +from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import COSMOS3_EXTRA_SPECS +from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniMoTPipeline + +pytestmark = pytest.mark.cosmos3 + + +class TestFindClosestTargetSize: + @pytest.mark.parametrize( + "input_h,input_w,action_resolution,expected", + [ + (480, 832, 480, (832, 480)), + (832, 480, 480, (480, 832)), + (512, 512, 480, (640, 640)), + (704, 1280, 704, (1280, 704)), + (256, 256, 256, (256, 256)), + (720, 1280, 720, (1280, 720)), + ], + ) + def test_picks_closest_aspect_bucket(self, input_h, input_w, action_resolution, expected): + assert find_closest_target_size(input_h, input_w, action_resolution) == expected + + def test_accepts_string_and_int_resolution_keys(self): + ref_h, ref_w = 480, 832 + assert find_closest_target_size(ref_h, ref_w, 480) == find_closest_target_size( + ref_h, ref_w, "480" + ) + + def test_unknown_resolution_raises(self): + with pytest.raises(ValueError, match="Unknown Cosmos3 action resolution"): + find_closest_target_size(480, 832, 1080) + + @pytest.mark.parametrize("action_resolution", sorted(VIDEO_RES_SIZE_INFO)) + def test_all_buckets_have_aspect_entries(self, action_resolution): + assert VIDEO_RES_SIZE_INFO[action_resolution] + + +class TestResolveActionSize: + @staticmethod + def _ref_image(width: int, height: int) -> PIL.Image.Image: + return PIL.Image.new("RGB", (width, height)) + + def test_explicit_height_and_width_are_unchanged(self): + ref = self._ref_image(832, 480) + assert Cosmos3OmniMoTPipeline._resolve_action_size(400, 600, ref, 480) == (400, 600) + + def test_unset_height_and_width_use_action_resolution_bucket(self): + ref = self._ref_image(832, 480) + assert Cosmos3OmniMoTPipeline._resolve_action_size(None, None, ref, 480) == (480, 832) + + def test_partial_height_fills_width_from_bucket(self): + ref = self._ref_image(832, 480) + assert Cosmos3OmniMoTPipeline._resolve_action_size(400, None, ref, 480) == (400, 832) + + def test_partial_width_fills_height_from_bucket(self): + ref = self._ref_image(832, 480) + assert Cosmos3OmniMoTPipeline._resolve_action_size(None, 600, ref, 480) == (480, 600) + + +class TestActionResolutionExtraParam: + def test_extra_param_spec_uses_action_resolution_key(self): + spec = COSMOS3_EXTRA_SPECS["action_resolution"] + assert spec.type == "int" + assert spec.default == 480 diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index c0329b71a060..63aceb9cf51e 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -36,13 +36,16 @@ import pytest import torch -from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import COSMOS3_T2I_PARAMS from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import ( COSMOS3_DEFAULT_RESOLUTION_TEMPLATE, COSMOS3_DURATION_TEMPLATE, COSMOS3_IMAGE_RESOLUTION_TEMPLATE, Cosmos3OmniMoTPipeline, ) +from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( + COSMOS3_ACTION_PARAMS, + COSMOS3_T2I_PARAMS, +) from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader from tensorrt_llm.visual_gen.args import TorchCompileConfig, VisualGenArgs @@ -215,6 +218,24 @@ def _require_audio_pipeline(pipeline) -> None: pytest.skip("Audio tokenizer was not loaded for this pipeline") +def _require_action_pipeline(pipeline) -> None: + if not getattr(pipeline, "action_gen", False): + pytest.skip("Checkpoint does not enable action generation") + + +def _assert_valid_action(action: torch.Tensor, *, raw_action_dim: int, chunk_size: int): + assert action is not None + assert action.dtype == torch.float32 + assert action.dim() == 3, f"Expected (B,T,D), got {action.shape}" + batch, t, d = action.shape + assert batch == 1 + assert t == chunk_size + assert d == raw_action_dim + af = action.float() + assert not torch.isnan(af).any() + assert not torch.isinf(af).any() + + def _make_test_image() -> PIL.Image.Image: image_path = os.environ.get("COSMOS3_TEST_IMAGE") if image_path and os.path.exists(image_path): @@ -433,6 +454,126 @@ def test_audio_smoke(self, cosmos3_pipeline): _assert_valid_audio(result.audio, result.audio_sample_rate) +@pytest.mark.integration +@pytest.mark.cosmos3_action +@pytest.mark.high_cuda_memory +class TestCosmos3Action: + ACTION_HEIGHT = 480 + ACTION_WIDTH = 832 + ACTION_FRAMES = COSMOS3_ACTION_PARAMS["num_frames"] + ACTION_CHUNK = COSMOS3_ACTION_PARAMS["action_chunk_size"] + RAW_ACTION_DIM = 10 + + def test_policy_smoke(self, cosmos3_pipeline): + _require_action_pipeline(cosmos3_pipeline) + image = _make_test_image().resize((self.ACTION_WIDTH, self.ACTION_HEIGHT)) + result = _run_forward( + cosmos3_pipeline, + image=image, + height=self.ACTION_HEIGHT, + width=self.ACTION_WIDTH, + num_frames=self.ACTION_FRAMES, + guidance_scale=COSMOS3_ACTION_PARAMS["guidance_scale"], + action_mode="policy", + domain_name="bridge_orig_lerobot", + raw_action_dim=self.RAW_ACTION_DIM, + action_chunk_size=self.ACTION_CHUNK, + ) + _assert_valid_video( + result.video, + num_frames=self.ACTION_FRAMES, + height=self.ACTION_HEIGHT, + width=self.ACTION_WIDTH, + ) + _assert_valid_action( + result.action, + raw_action_dim=self.RAW_ACTION_DIM, + chunk_size=self.ACTION_CHUNK, + ) + assert result.action_mode == "policy" + assert result.domain_id == 7 + + def test_forward_dynamics_smoke(self, cosmos3_pipeline): + _require_action_pipeline(cosmos3_pipeline) + image = _make_test_image().resize((self.ACTION_WIDTH, self.ACTION_HEIGHT)) + action_traj = [[0.1] * self.RAW_ACTION_DIM for _ in range(self.ACTION_CHUNK)] + result = _run_forward( + cosmos3_pipeline, + image=image, + height=self.ACTION_HEIGHT, + width=self.ACTION_WIDTH, + num_frames=self.ACTION_FRAMES, + guidance_scale=COSMOS3_ACTION_PARAMS["guidance_scale"], + action_mode="forward_dynamics", + domain_name="bridge_orig_lerobot", + action=action_traj, + action_chunk_size=self.ACTION_CHUNK, + ) + _assert_valid_video( + result.video, + num_frames=self.ACTION_FRAMES, + height=self.ACTION_HEIGHT, + width=self.ACTION_WIDTH, + ) + _assert_valid_action( + result.action, + raw_action_dim=self.RAW_ACTION_DIM, + chunk_size=self.ACTION_CHUNK, + ) + + def test_inverse_dynamics_smoke(self, cosmos3_pipeline): + _require_action_pipeline(cosmos3_pipeline) + image = _make_test_image().resize((self.ACTION_WIDTH, self.ACTION_HEIGHT)) + video = [image.copy() for _ in range(NUM_FRAMES)] + result = _run_forward( + cosmos3_pipeline, + image=None, + height=self.ACTION_HEIGHT, + width=self.ACTION_WIDTH, + num_frames=NUM_FRAMES, + guidance_scale=COSMOS3_ACTION_PARAMS["guidance_scale"], + action_mode="inverse_dynamics", + domain_name="bridge_orig_lerobot", + raw_action_dim=self.RAW_ACTION_DIM, + action_chunk_size=NUM_FRAMES, + video=video, + ) + _assert_valid_video( + result.video, + num_frames=NUM_FRAMES, + height=self.ACTION_HEIGHT, + width=self.ACTION_WIDTH, + ) + _assert_valid_action( + result.action, + raw_action_dim=self.RAW_ACTION_DIM, + chunk_size=NUM_FRAMES, + ) + + def test_action_and_audio_rejected(self, cosmos3_pipeline): + _require_action_pipeline(cosmos3_pipeline) + with pytest.raises(ValueError, match="joint action and audio"): + _run_forward( + cosmos3_pipeline, + image=_make_test_image(), + action_mode="policy", + domain_name="bridge_orig_lerobot", + raw_action_dim=self.RAW_ACTION_DIM, + enable_audio=True, + ) + + def test_action_and_t2i_rejected(self, cosmos3_pipeline): + _require_action_pipeline(cosmos3_pipeline) + with pytest.raises(ValueError, match="output_type='image'"): + _run_forward( + cosmos3_pipeline, + output_type="image", + action_mode="policy", + domain_name="bridge_orig_lerobot", + raw_action_dim=self.RAW_ACTION_DIM, + ) + + @pytest.mark.integration @pytest.mark.cosmos3_t2v @pytest.mark.high_cuda_memory diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py index 76e6c1a01a8c..7f02d97ad117 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py @@ -397,6 +397,144 @@ def test_forward_with_audio_multiframe(self, audio_model_config): _assert_finite_output(out.audio, torch.Size([1, model.audio_dim, self.T_AUDIO])) +@pytest.mark.integration +class TestCosmos3Action: + """Action modality — Nano architecture, random weights, action_gen on.""" + + ACTION_DIM = 64 + T_ACTION = 4 + NUM_DOMAINS = 32 + + @pytest.fixture(autouse=True) + def _require_cuda(self): + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + + @pytest.fixture + def action_model_config(self): + checkpoint_dir = _require_checkpoint() + model_config = _load_model_config(checkpoint_dir) + cfg = model_config.pretrained_config + cfg.action_gen = True + cfg.action_dim = self.ACTION_DIM + cfg.num_embodiment_domains = self.NUM_DOMAINS + cfg.sound_gen = False + return model_config + + @pytest.fixture + def cosmos3_model_config_noaction(self): + checkpoint_dir = _require_checkpoint() + model_config = _load_model_config(checkpoint_dir) + model_config.pretrained_config.action_gen = False + model_config.pretrained_config.sound_gen = False + return model_config + + def test_action_model_structure(self, action_model_config): + model = Cosmos3VFMTransformer(model_config=action_model_config) + assert model.action_gen is True + assert model.action_dim == self.ACTION_DIM + assert hasattr(model, "action_proj_in") + assert hasattr(model, "action_proj_out") + assert hasattr(model, "action_modality_embed") + assert model.action_modality_embed.shape == (model.hidden_size,) + + def test_video_only_model_has_no_action_heads(self, cosmos3_model_config_noaction): + model = Cosmos3VFMTransformer(model_config=cosmos3_model_config_noaction) + assert model.action_gen is False + assert not hasattr(model, "action_proj_in") + + @pytest.mark.high_cuda_memory + def test_forward_with_action(self, action_model_config): + cfg = action_model_config.pretrained_config + model = _build_random_weight_model(action_model_config) + hs, ts, text_ids, text_mask, video_shape = _cosmos3_inputs( + DEVICE, channels=cfg.latent_channel + ) + action_latents = torch.randn(1, self.T_ACTION, model.action_dim, device=DEVICE, dtype=DTYPE) + domain_ids = torch.tensor([7], dtype=torch.long, device=DEVICE) + with torch.inference_mode(): + out = model( + hidden_states=hs, + timestep=ts, + text_ids=text_ids, + text_mask=text_mask, + video_shape=video_shape, + fps=24.0, + action_latents=action_latents, + action_domain_ids=domain_ids, + ) + _assert_finite_output(out.video, hs.shape) + assert out.action is not None + _assert_finite_output(out.action, torch.Size([1, self.T_ACTION, model.action_dim])) + + @pytest.mark.high_cuda_memory + def test_forward_without_action_latents_returns_none(self, action_model_config): + cfg = action_model_config.pretrained_config + model = _build_random_weight_model(action_model_config) + hs, ts, text_ids, text_mask, video_shape = _cosmos3_inputs( + DEVICE, channels=cfg.latent_channel + ) + with torch.inference_mode(): + out = model( + hidden_states=hs, + timestep=ts, + text_ids=text_ids, + text_mask=text_mask, + video_shape=video_shape, + ) + _assert_finite_output(out.video, hs.shape) + assert out.action is None + + @pytest.mark.high_cuda_memory + def test_forward_with_action_noisy_mask(self, action_model_config): + cfg = action_model_config.pretrained_config + model = _build_random_weight_model(action_model_config) + hs, ts, text_ids, text_mask, video_shape = _cosmos3_inputs( + DEVICE, channels=cfg.latent_channel, t=2 + ) + action_latents = torch.randn(1, self.T_ACTION, model.action_dim, device=DEVICE, dtype=DTYPE) + noisy_mask = torch.ones(1, self.T_ACTION, 1, device=DEVICE, dtype=DTYPE) + noisy_mask[:, 0, :] = 0.0 + domain_ids = torch.tensor([7], dtype=torch.long, device=DEVICE) + with torch.inference_mode(): + out = model( + hidden_states=hs, + timestep=ts, + text_ids=text_ids, + text_mask=text_mask, + video_shape=video_shape, + fps=24.0, + action_latents=action_latents, + action_domain_ids=domain_ids, + action_noisy_mask=noisy_mask, + ) + _assert_finite_output(out.video, hs.shape) + _assert_finite_output(out.action, torch.Size([1, self.T_ACTION, model.action_dim])) + + @pytest.mark.high_cuda_memory + def test_forward_with_action_multiframe(self, action_model_config): + cfg = action_model_config.pretrained_config + model = _build_random_weight_model(action_model_config) + hs, ts, text_ids, text_mask, video_shape = _cosmos3_inputs( + DEVICE, channels=cfg.latent_channel, t=3 + ) + action_latents = torch.randn(1, self.T_ACTION, model.action_dim, device=DEVICE, dtype=DTYPE) + domain_ids = torch.tensor([7], dtype=torch.long, device=DEVICE) + with torch.inference_mode(): + out = model( + hidden_states=hs, + timestep=ts, + text_ids=text_ids, + text_mask=text_mask, + video_shape=video_shape, + fps=24.0, + action_latents=action_latents, + action_domain_ids=domain_ids, + ) + _assert_finite_output(out.video, hs.shape) + _assert_finite_output(out.action, torch.Size([1, self.T_ACTION, model.action_dim])) + + @pytest.mark.integration class TestCosmos3TransformerCheckpoint: """Load Cosmos3-Nano transformer weights and run a single forward step.""" From b98df2068e02ccde61d84304c6eb1b53c0ff5f80 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Thu, 11 Jun 2026 09:26:53 -0700 Subject: [PATCH 02/64] updates Signed-off-by: Shreyas Misra --- examples/visual_gen/models/cosmos3/cosmos3.py | 58 ++-- .../visual_gen/models/cosmos3/action.py | 223 ++++++++++++++- .../visual_gen/models/cosmos3/defaults.py | 270 +++++++++++++++++- .../models/cosmos3/pipeline_cosmos3.py | 254 +++++----------- .../models/cosmos3/transformer_cosmos3.py | 9 +- .../_torch/visual_gen/test_cosmos3_action.py | 167 ++++++++++- 6 files changed, 768 insertions(+), 213 deletions(-) diff --git a/examples/visual_gen/models/cosmos3/cosmos3.py b/examples/visual_gen/models/cosmos3/cosmos3.py index 6186efef5f73..368399101d5d 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3.py +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -129,7 +129,12 @@ from typing import Any, Dict, Optional from tensorrt_llm import VisualGen, VisualGenArgs -from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import COSMOS3_ACTION_PARAMS +from tensorrt_llm._torch.visual_gen.models.cosmos3.action import VIDEO_RES_SIZE_INFO +from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( + COSMOS3_ACTION_PARAMS, + get_domain_preset, + resolve_domain_action_config, +) _SCRIPT_DIR = Path(__file__).resolve().parent @@ -224,25 +229,46 @@ def _validate_action_args( f"{mode} requires --image_path, a prompt-file vision_path, or --video_path " "for the first frame." ) - if args.raw_action_dim is None: - raise SystemExit(f"{mode} requires --raw_action_dim.") + preset = get_domain_preset(args.domain_name, args.domain_id) + effective_raw_dim = args.raw_action_dim or (preset or {}).get("raw_action_dim") + if effective_raw_dim is None: + raise SystemExit( + f"{mode} requires --raw_action_dim or a known --domain_name with a preset." + ) elif mode == "inverse_dynamics": if args.video_path is None: - raise SystemExit(f"{mode} requires --video_path (frame directory or image list).") - if args.raw_action_dim is None: - raise SystemExit(f"{mode} requires --raw_action_dim.") + raise SystemExit( + f"{mode} requires --video_path (frame directory, .mp4/.avi, or image)." + ) + preset = get_domain_preset(args.domain_name, args.domain_id) + effective_raw_dim = args.raw_action_dim or (preset or {}).get("raw_action_dim") + if effective_raw_dim is None: + raise SystemExit( + f"{mode} requires --raw_action_dim or a known --domain_name with a preset." + ) def _apply_action_generation_params(params, args: argparse.Namespace) -> None: - """Set 480p action defaults on the request before pipeline overrides.""" - chunk = args.action_chunk_size or COSMOS3_ACTION_PARAMS["action_chunk_size"] - params.height = 480 - params.width = 832 - params.num_frames = chunk + 1 + """Set action defaults on the request; domain presets override generic 480p.""" + cfg = resolve_domain_action_config( + domain_name=args.domain_name, + domain_id=args.domain_id, + raw_action_dim=args.raw_action_dim, + action_chunk_size=args.action_chunk_size, + action_resolution=args.action_resolution, + ) + bucket = str(cfg["action_resolution"]) + width, height = VIDEO_RES_SIZE_INFO[bucket]["16,9"] + params.width = width + params.height = height + params.num_frames = cfg["num_frames"] params.num_inference_steps = COSMOS3_ACTION_PARAMS["num_inference_steps"] params.guidance_scale = COSMOS3_ACTION_PARAMS["guidance_scale"] - params.frame_rate = COSMOS3_ACTION_PARAMS["frame_rate"] - params.extra_params["action_chunk_size"] = chunk + params.frame_rate = cfg["frame_rate"] + params.extra_params["action_chunk_size"] = cfg["action_chunk_size"] + if cfg["raw_action_dim"] is not None: + params.extra_params["raw_action_dim"] = cfg["raw_action_dim"] + params.extra_params["action_resolution"] = cfg["action_resolution"] def _default_action_output_path(video_path: str) -> str: @@ -378,14 +404,14 @@ def main(): "--video_path", type=str, default=None, - help="Frame directory or image path for inverse_dynamics (or first-frame fallback)", + help="Frame directory, .mp4/.avi video, or image path for inverse_dynamics", ) parser.add_argument( "--action_resolution", type=int, - default=480, + default=None, choices=[256, 480, 704, 720], - help="Resolution bucket for action image sizing", + help=("Resolution bucket for action image sizing. Defaults to the domain preset or 480."), ) parser.add_argument( "--action_output_path", diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py index 60d1ea8bdbf6..b66dda5b6ac2 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py @@ -5,10 +5,13 @@ from __future__ import annotations -from typing import Any +from pathlib import Path +from typing import Any, List, Optional import numpy as np +import PIL.Image import torch +from diffusers.utils.torch_utils import randn_tensor ACTION_MODE_POLICY = "policy" ACTION_MODE_FORWARD_DYNAMICS = "forward_dynamics" @@ -71,6 +74,26 @@ } +COSMOS3_ACTION_RESOLUTIONS = tuple(int(key) for key in sorted(VIDEO_RES_SIZE_INFO, key=int)) + + +def normalize_action_resolution(resolution: Any) -> int: + if resolution is None: + raise ValueError("Cosmos3 action_resolution is required for action generation.") + try: + bucket = int(resolution) + except (TypeError, ValueError) as exc: + raise ValueError( + f"Cosmos3 action_resolution must be an int bucket, got {resolution!r}." + ) from exc + if bucket not in COSMOS3_ACTION_RESOLUTIONS: + raise ValueError( + f"Unknown Cosmos3 action_resolution={bucket}; " + f"expected one of {COSMOS3_ACTION_RESOLUTIONS}." + ) + return bucket + + def normalize_action_mode(mode: Any) -> str | None: if mode is None: return None @@ -219,3 +242,201 @@ def find_closest_target_size(h: int, w: int, resolution: str | int) -> tuple[int best_size = (cand_w, cand_h) assert best_size is not None return best_size + + +ACTION_IMAGE_EXTENSIONS = frozenset({".png", ".jpg", ".jpeg", ".webp", ".bmp"}) +ACTION_VIDEO_EXTENSIONS = frozenset({".mp4", ".avi"}) + + +def pil_to_rgb(value: Any) -> PIL.Image.Image: + if isinstance(value, str): + return PIL.Image.open(value).convert("RGB") + if isinstance(value, PIL.Image.Image): + return value.convert("RGB") + raise TypeError( + f"Cosmos3 action preprocessing expected PIL image or image path, got {type(value)!r}." + ) + + +def decode_action_video_file(path: Path, max_frames: Optional[int] = None) -> List[PIL.Image.Image]: + import torchvision.io as io + + frames, _, _ = io.read_video(str(path), pts_unit="sec") + if frames.numel() == 0: + raise ValueError(f"Cosmos3 action video file contains no frames: {path}") + if max_frames is not None: + frames = frames[:max_frames] + return [PIL.Image.fromarray(frames[i].numpy()) for i in range(frames.shape[0])] + + +def normalize_action_video_path(path: Path, max_frames: Optional[int] = None) -> List[Any]: + if not path.exists(): + raise ValueError(f"Cosmos3 action video path does not exist: {path}") + if path.is_dir(): + frames = sorted(p for p in path.iterdir() if p.suffix.lower() in ACTION_IMAGE_EXTENSIONS) + if not frames: + raise ValueError(f"No image frames found in Cosmos3 action video directory: {path}") + frame_paths = [str(p) for p in frames] + if max_frames is not None: + frame_paths = frame_paths[:max_frames] + return frame_paths + + suffix = path.suffix.lower() + if suffix in ACTION_IMAGE_EXTENSIONS: + return [str(path)] + if suffix in ACTION_VIDEO_EXTENSIONS: + return decode_action_video_file(path, max_frames=max_frames) + raise ValueError( + "Cosmos3 action video path must be a frame directory, an image file " + f"{sorted(ACTION_IMAGE_EXTENSIONS)}, or a video file " + f"{sorted(ACTION_VIDEO_EXTENSIONS)}; got {path}" + ) + + +def normalize_action_video_input(video: Any, max_frames: Optional[int] = None) -> List[Any]: + """Normalize action video input to a frame list. + + Accepts a list of PIL images / paths, a single image or video file path, + or a directory of frame images (sorted lexicographically). + """ + if video is None: + return [] + if isinstance(video, list): + if not video: + raise ValueError("Cosmos3 action video input must contain at least one frame.") + if max_frames is not None: + return video[:max_frames] + return video + if isinstance(video, (str, Path)): + return normalize_action_video_path(Path(video), max_frames=max_frames) + return [video] + + +def resolve_action_size( + height: Optional[int], + width: Optional[int], + ref_image: PIL.Image.Image, + action_resolution: int, +) -> tuple[int, int]: + """Fill unset action H/W from the action resolution bucket; honor explicit values.""" + if height is not None and width is not None: + return height, width + target_w, target_h = find_closest_target_size( + ref_image.height, ref_image.width, action_resolution + ) + return ( + height if height is not None else target_h, + width if width is not None else target_w, + ) + + +def action_reference_image( + *, + action_mode: str, + image: Any, + video: Any, +) -> PIL.Image.Image: + """Resolve the reference frame used for action sizing and conditioning.""" + if action_mode == ACTION_MODE_INVERSE_DYNAMICS: + source = video if video is not None else image + frames = normalize_action_video_input(source, max_frames=1) + if not frames: + raise ValueError("Cosmos3 action_mode='inverse_dynamics' requires a video input.") + return pil_to_rgb(frames[0]) + + source = image if image is not None else video + if source is None: + raise ValueError(f"Cosmos3 action_mode={action_mode!r} requires an image or video input.") + if isinstance(source, PIL.Image.Image): + return source.convert("RGB") + if isinstance(source, str): + path = Path(source) + if path.is_file() and path.suffix.lower() in ACTION_IMAGE_EXTENSIONS: + return PIL.Image.open(source).convert("RGB") + frames = normalize_action_video_input(source, max_frames=1) + if not frames: + raise ValueError( + f"Cosmos3 action_mode={action_mode!r} requires an image or video input." + ) + return pil_to_rgb(frames[0]) + raise TypeError( + f"Cosmos3 action reference image must be PIL.Image or path, got {type(source)!r}." + ) + + +def resize_and_pad_action_image( + image: PIL.Image.Image, target_h: int, target_w: int +) -> PIL.Image.Image: + scale = min(target_w / image.width, target_h / image.height, 1.0) + resize_w = max(1, int(scale * image.width + 0.5)) + resize_h = max(1, int(scale * image.height + 0.5)) + if (resize_w, resize_h) != image.size: + image = image.resize((resize_w, resize_h), PIL.Image.Resampling.BICUBIC) + + array = np.asarray(image) + pad_h = target_h - resize_h + pad_w = target_w - resize_w + if pad_h < 0 or pad_w < 0: + raise ValueError( + f"Cosmos3 action image resize exceeded target size: resized={(resize_h, resize_w)}, " + f"target={(target_h, target_w)}." + ) + if pad_h == 0 and pad_w == 0: + return image + pad_mode = "reflect" if pad_h < resize_h and pad_w < resize_w else "edge" + padded = np.pad(array, ((0, pad_h), (0, pad_w), (0, 0)), mode=pad_mode) + return PIL.Image.fromarray(padded) + + +def prepare_action_latents( + *, + mode: str, + action_chunk_size: int, + raw_action_dim: Optional[int], + action_dim: int, + generator: torch.Generator, + device: torch.device, + dtype: torch.dtype, + action_input: Any = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: + if mode == ACTION_MODE_FORWARD_DYNAMICS: + action = load_action_tensor(action_input) + if action.shape[0] < action_chunk_size: + pad = action[-1:].repeat(action_chunk_size - action.shape[0], 1) + action = torch.cat([action, pad], dim=0) + elif action.shape[0] > action_chunk_size: + action = action[:action_chunk_size] + if raw_action_dim is None: + raw_action_dim = int(action.shape[-1]) + clean_action = pad_action_to_dim(action, action_dim) + else: + if raw_action_dim is None: + raise ValueError( + "Cosmos3 action_mode='policy' and 'inverse_dynamics' require raw_action_dim." + ) + clean_action = torch.zeros(action_chunk_size, action_dim, dtype=torch.float32) + + raw_action_dim = int(raw_action_dim) + if raw_action_dim <= 0 or raw_action_dim > action_dim: + raise ValueError( + f"Cosmos3 raw_action_dim must be in [1, {action_dim}], got {raw_action_dim}." + ) + + clean_action = clean_action.to(device=device, dtype=dtype).unsqueeze(0) + condition_mask = build_action_condition_mask( + mode, + action_chunk_size, + device=device, + dtype=dtype, + ) + noise = randn_tensor( + (1, action_chunk_size, action_dim), + generator=generator, + device=device, + dtype=dtype, + ) + noise[:, :, raw_action_dim:] = 0 + clean_action[:, :, raw_action_dim:] = 0 + action_latents = condition_mask * clean_action + (1.0 - condition_mask) * noise + action_velocity_mask = 1.0 - condition_mask + return action_latents, action_velocity_mask, clean_action, raw_action_dim diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index 780fc0f397c7..60a3a496a7a9 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -15,10 +15,24 @@ """Per-model default generation parameters for Cosmos3 pipelines. Shared by the Cosmos3 OmniMoT text-to-video and image-to-video generation paths. + +Action generation +----------------- +``COSMOS3_DOMAIN_PRESETS`` lists training-aligned defaults per embodiment. When +``domain_name`` (or a uniquely mapped ``domain_id``) is set, the pipeline fills +omitted ``raw_action_dim``, ``action_chunk_size``, ``num_frames``, +``action_resolution``, and ``frame_rate`` from the preset and logs a warning if +explicit values differ. See Cosmos3 omni ``action_*.json`` inputs for reference +configs (bridge, av, droid, libero, etc.). """ -from typing import Dict +from typing import Any, Dict, List, Optional, TypedDict +from tensorrt_llm._torch.visual_gen.models.cosmos3.action import ( + COSMOS3_ACTION_RESOLUTIONS, + EMBODIMENT_TO_DOMAIN_ID, + normalize_action_resolution, +) from tensorrt_llm._torch.visual_gen.pipeline import ExtraParamSchema # --------------------------------------------------------------------------- @@ -35,10 +49,16 @@ "frame_rate": 24.0, } -# Fields merged by the executor for every request. Modality-specific sizing and -# sampling defaults (height/width/num_frames/steps/guidance) stay ``None`` on the -# request until ``forward()`` resolves them from T2V/T2I/action context. +# Fields merged by the executor for every request. Modality-specific values +# (height/width/num_frames/steps/guidance) are declared with ``None`` so the +# executor accepts explicit overrides; ``forward()`` resolves unset fields from +# T2V/T2I/action context. COSMOS3_PIPELINE_DEFAULTS = { + "height": None, + "width": None, + "num_frames": None, + "num_inference_steps": None, + "guidance_scale": None, "max_sequence_length": COSMOS3_720P_PARAMS["max_sequence_length"], "frame_rate": COSMOS3_720P_PARAMS["frame_rate"], } @@ -63,6 +83,219 @@ "frame_rate": 24.0, } + +class Cosmos3DomainPreset(TypedDict, total=False): + """Recommended action-generation settings for a trained embodiment.""" + + raw_action_dim: int + action_chunk_size: int + num_frames: int + action_resolution: int + frame_rate: float + + +# Training-aligned defaults. Values mirror Cosmos3 omni action JSON examples where available. +COSMOS3_DOMAIN_PRESETS: Dict[str, Cosmos3DomainPreset] = { + # WidowX bridge; 7-DOF arm + gripper in 10-D state. + "bridge_orig_lerobot": { + "raw_action_dim": 10, + "action_chunk_size": 16, + "num_frames": 17, + "action_resolution": 480, + "frame_rate": 5.0, + }, + # Autonomous-vehicle steering/throttle; longer action horizon. + "av": { + "raw_action_dim": 9, + "action_chunk_size": 60, + "num_frames": 61, + "action_resolution": 480, + "frame_rate": 10.0, + }, + # 6-DoF camera pose + shutter; matches AV-style horizon. + "camera_pose": { + "raw_action_dim": 9, + "action_chunk_size": 60, + "num_frames": 61, + "action_resolution": 480, + "frame_rate": 30.0, + }, + # Franka single-arm tabletop; same domain_id as robomind-franka. + "droid_lerobot": { + "raw_action_dim": 10, + "action_chunk_size": 16, + "num_frames": 17, + "action_resolution": 480, + "frame_rate": 15.0, + }, + # LIBERO sim single-arm; lower action resolution bucket. + "libero": { + "raw_action_dim": 10, + "action_chunk_size": 16, + "num_frames": 17, + "action_resolution": 256, + "frame_rate": 10.0, + }, + # MANO hand pose; high-DOF hand articulation. + "hand_pose": { + "raw_action_dim": 57, + "action_chunk_size": 16, + "num_frames": 17, + "action_resolution": 480, + "frame_rate": 24.0, + }, + # AgiBot humanoid; shared domain_id with agibot_gear_gripper*. + "agibotworld": { + "raw_action_dim": 29, + "action_chunk_size": 16, + "num_frames": 17, + "action_resolution": 480, + "frame_rate": 10.0, + }, + # Google Robot (RT-1 / fractal) single-arm. + "fractal": { + "raw_action_dim": 10, + "action_chunk_size": 16, + "num_frames": 17, + "action_resolution": 480, + "frame_rate": 5.0, + }, + # 2-D planar push task. + "pusht": { + "raw_action_dim": 2, + "action_chunk_size": 16, + "num_frames": 17, + "action_resolution": 256, + "frame_rate": 10.0, + }, + # UMI handheld gripper setup. + "umi": { + "raw_action_dim": 10, + "action_chunk_size": 16, + "num_frames": 17, + "action_resolution": 480, + "frame_rate": 10.0, + }, +} + +# Map alias domain_name keys to a canonical preset entry. +COSMOS3_DOMAIN_PRESET_ALIASES: Dict[str, str] = { + "robomind-franka": "droid_lerobot", + "robomind-franka-dual": "droid_lerobot", + "robomind-ur": "droid_lerobot", + "agibot_gear_gripper": "agibotworld", + "agibot_gear_gripper_ext": "agibotworld", + "galbot": "agibotworld", +} + + +def canonical_domain_preset_key( + domain_name: Any = None, + domain_id: Any = None, +) -> Optional[str]: + if domain_name is not None and str(domain_name).strip(): + key = str(domain_name).strip().lower() + key = COSMOS3_DOMAIN_PRESET_ALIASES.get(key, key) + if key in COSMOS3_DOMAIN_PRESETS: + return key + return None + + if domain_id is None: + return None + + resolved_id = int(domain_id) + if resolved_id == 0: + return None + + candidates: List[str] = [] + for name, mapped_id in EMBODIMENT_TO_DOMAIN_ID.items(): + if mapped_id != resolved_id: + continue + canon = COSMOS3_DOMAIN_PRESET_ALIASES.get(name, name) + if canon in COSMOS3_DOMAIN_PRESETS and canon not in candidates: + candidates.append(canon) + + if len(candidates) == 1: + return candidates[0] + return None + + +def get_domain_preset( + domain_name: Any = None, + domain_id: Any = None, +) -> Optional[Cosmos3DomainPreset]: + key = canonical_domain_preset_key(domain_name, domain_id) + if key is None: + return None + return COSMOS3_DOMAIN_PRESETS[key] + + +def resolve_domain_action_config( + *, + domain_name: Any = None, + domain_id: Any = None, + raw_action_dim: Optional[int] = None, + action_chunk_size: Optional[int] = None, + action_resolution: Optional[int] = None, + frame_rate: Optional[float] = None, + num_frames: Optional[int] = None, +) -> Dict[str, Any]: + """Merge user action params with domain presets and generic fallbacks.""" + preset_key = canonical_domain_preset_key(domain_name, domain_id) + preset = COSMOS3_DOMAIN_PRESETS.get(preset_key) if preset_key else None + warnings: List[str] = [] + + def _resolve_field( + field: str, + current: Any, + *, + fallback: Any = None, + ) -> Any: + recommended = preset.get(field) if preset else None + if current is not None: + if recommended is not None and current != recommended: + warnings.append( + f"Cosmos3 {field}={current} differs from recommended " + f"{recommended} for domain {preset_key!r}." + ) + return current + if recommended is not None: + return recommended + return fallback + + resolved_raw_action_dim = _resolve_field("raw_action_dim", raw_action_dim) + resolved_chunk = _resolve_field( + "action_chunk_size", + action_chunk_size, + fallback=COSMOS3_ACTION_PARAMS["action_chunk_size"], + ) + resolved_resolution = normalize_action_resolution( + _resolve_field( + "action_resolution", + action_resolution, + fallback=480, + ) + ) + resolved_frame_rate = _resolve_field( + "frame_rate", + frame_rate, + fallback=COSMOS3_ACTION_PARAMS["frame_rate"], + ) + resolved_num_frames = _resolve_field("num_frames", num_frames) + if resolved_num_frames is None: + resolved_num_frames = int(resolved_chunk) + 1 + + return { + "raw_action_dim": resolved_raw_action_dim, + "action_chunk_size": int(resolved_chunk), + "action_resolution": resolved_resolution, + "frame_rate": float(resolved_frame_rate), + "num_frames": int(resolved_num_frames), + "preset_key": preset_key, + "warnings": warnings, + } + + COSMOS3_EXTRA_SPECS: Dict[str, ExtraParamSchema] = { "use_duration_template": ExtraParamSchema( type="bool", @@ -102,7 +335,11 @@ "domain_name": ExtraParamSchema( type="str", default=None, - description="Embodiment domain name for action generation.", + description=( + "Embodiment domain name for action generation (e.g. bridge_orig_lerobot, av). " + "When set, omitted raw_action_dim/action_chunk_size/action_resolution/frame_rate " + "are filled from COSMOS3_DOMAIN_PRESETS; mismatches are logged as warnings." + ), ), "domain_id": ExtraParamSchema( type="int", @@ -112,12 +349,18 @@ "raw_action_dim": ExtraParamSchema( type="int", default=None, - description="Raw action DOF count for policy/inverse_dynamics.", + description=( + "Raw action DOF for policy/inverse_dynamics (e.g. 10 bridge, 9 av, 29 agibot). " + "Inferred from domain_name preset when omitted." + ), ), "action_chunk_size": ExtraParamSchema( type="int", default=COSMOS3_ACTION_PARAMS["action_chunk_size"], - description="Number of action tokens to generate.", + description=( + "Number of action tokens to generate (16 for most robots, 60 for av/camera_pose). " + "Inferred from domain_name preset when omitted." + ), ), "action": ExtraParamSchema( type="list", @@ -127,11 +370,18 @@ "action_resolution": ExtraParamSchema( type="int", default=480, - description="Resolution bucket for action image sizing (256/480/704/720).", + description=( + "Resolution bucket for action image sizing. Must be one of " + f"{list(COSMOS3_ACTION_RESOLUTIONS)}. Inferred from domain_name preset when omitted." + ), + range=(min(COSMOS3_ACTION_RESOLUTIONS), max(COSMOS3_ACTION_RESOLUTIONS)), ), "video": ExtraParamSchema( - type="list", + type="path_or_list", default=None, - description="Video frames (PIL images or paths) for inverse_dynamics mode.", + description=( + "Video for inverse_dynamics: .mp4/.avi file, frame directory, " + "image path, or list of PIL images / frame paths." + ), ), } diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index b50a36f87515..8b09f7cf4720 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -17,10 +17,8 @@ import math import os import time -from pathlib import Path from typing import Any, List, Optional, Union -import numpy as np import PIL.Image import torch from diffusers import AutoencoderKLWan, UniPCMultistepScheduler @@ -37,15 +35,16 @@ from tensorrt_llm.logger import logger from .action import ( - ACTION_MODE_FORWARD_DYNAMICS, ACTION_MODE_INVERSE_DYNAMICS, + action_reference_image, action_start_frame_offset, - build_action_condition_mask, build_vision_condition_mask, - find_closest_target_size, - load_action_tensor, normalize_action_mode, - pad_action_to_dim, + normalize_action_video_input, + pil_to_rgb, + prepare_action_latents, + resize_and_pad_action_image, + resolve_action_size, resolve_domain_id, ) from .defaults import ( @@ -54,6 +53,7 @@ COSMOS3_EXTRA_SPECS, COSMOS3_PIPELINE_DEFAULTS, COSMOS3_T2I_PARAMS, + resolve_domain_action_config, ) from .guardrails import check_video_safety, download_guardrail_checkpoint from .sound_tokenizer import LatentAutoEncoderV2 @@ -176,6 +176,10 @@ def load_standard_components( # (UniPC mutates internal correction buffers on every .step() call). self.audio_scheduler = UniPCMultistepScheduler.from_config(self.scheduler.config) + if self.action_gen: + # Action uses its own scheduler for the same reason as audio. + self.action_scheduler = UniPCMultistepScheduler.from_config(self.scheduler.config) + # Re-check the env var in case it was changed after initialization like in unit tests. guardrails_disabled = os.environ.get("TRTLLM_DISABLE_COSMOS3_GUARDRAILS", "0") == "1" global TRTLLM_DISABLE_COSMOS3_GUARDRAILS @@ -220,6 +224,14 @@ def _set_flow_shift(self, target_shift: float) -> None: self.scheduler = UniPCMultistepScheduler.from_config( self._base_scheduler_config, flow_shift=target ) + if self.audio_gen: + self.audio_scheduler = UniPCMultistepScheduler.from_config( + self._base_scheduler_config, flow_shift=target + ) + if self.action_gen: + self.action_scheduler = UniPCMultistepScheduler.from_config( + self._base_scheduler_config, flow_shift=target + ) self._current_flow_shift = target @property @@ -238,50 +250,6 @@ def default_generation_params(self): def extra_param_specs(self): return dict(COSMOS3_EXTRA_SPECS) - @staticmethod - def _resolve_action_size( - height: Optional[int], - width: Optional[int], - ref_image: PIL.Image.Image, - action_resolution: int, - ) -> tuple[int, int]: - """Fill unset action H/W from the action resolution bucket; honor explicit values.""" - if height is not None and width is not None: - return height, width - target_w, target_h = find_closest_target_size( - ref_image.height, ref_image.width, action_resolution - ) - return ( - height if height is not None else target_h, - width if width is not None else target_w, - ) - - def _action_reference_image( - self, - *, - action_mode: str, - image: Any, - video: Any, - ) -> PIL.Image.Image: - if action_mode == ACTION_MODE_INVERSE_DYNAMICS: - frames = self._normalize_action_video_input(video) - if not frames: - raise ValueError("Cosmos3 action_mode='inverse_dynamics' requires a video input.") - return self._pil_to_rgb(frames[0]) - - if image is None and video is not None: - frames = self._normalize_action_video_input(video) - return self._pil_to_rgb(frames[0]) - if image is None: - raise ValueError(f"Cosmos3 action_mode={action_mode!r} requires an image input.") - if isinstance(image, str): - return PIL.Image.open(image).convert("RGB") - if isinstance(image, PIL.Image.Image): - return image.convert("RGB") - raise TypeError( - f"Cosmos3 action reference image must be PIL.Image or path, got {type(image)!r}." - ) - def _run_warmup(self, height: int, width: int, num_frames: int, steps: int) -> None: with torch.no_grad(): self.forward( @@ -334,8 +302,7 @@ def infer(self, req): action_chunk_size=extra_params.get("action_chunk_size"), action=extra_params.get("action"), action_resolution=extra_params.get("action_resolution") - or extra_params.get("image_size") - or 480, + or extra_params.get("image_size"), video=extra_params.get("video"), ) @@ -492,7 +459,9 @@ def _prepare_latents(self, height, width, num_frames, generator): ) return randn_tensor(shape, generator=generator, device=self.device, dtype=self.dtype) - # -- I2V latent preparation ----------------------------------------------- + # ========================================================================= + # I2V latent preparation + # ========================================================================= def _encode_conditioning_video( self, @@ -648,72 +617,14 @@ def decode_audio(self, latent: torch.Tensor) -> torch.Tensor: """ return self.audio_tokenizer.decode(latent).float() # [B, audio_channels, N_samples] - @staticmethod - def _normalize_action_video_input(video: Any) -> List[Any]: - """Normalize inverse-dynamics video input to a frame list. - - Accepts a list of PIL images / paths, a single image path, or a - directory of frame images (sorted lexicographically). - """ - if video is None: - return [] - if isinstance(video, list): - if not video: - raise ValueError("Cosmos3 action video input must contain at least one frame.") - return video - if isinstance(video, str): - path = Path(video) - if not path.exists(): - raise ValueError(f"Cosmos3 action video path does not exist: {video}") - if path.is_dir(): - exts = {".png", ".jpg", ".jpeg", ".webp", ".bmp"} - frames = sorted(p for p in path.iterdir() if p.suffix.lower() in exts) - if not frames: - raise ValueError( - f"No image frames found in Cosmos3 action video directory: {video}" - ) - return [str(p) for p in frames] - return [video] - return [video] - - @staticmethod - def _pil_to_rgb(value: Any) -> PIL.Image.Image: - if isinstance(value, str): - return PIL.Image.open(value).convert("RGB") - if isinstance(value, PIL.Image.Image): - return value.convert("RGB") - raise TypeError( - f"Cosmos3 action preprocessing expected PIL image or image path, got {type(value)!r}." - ) - - @staticmethod - def _resize_and_pad_action_image( - image: PIL.Image.Image, target_h: int, target_w: int - ) -> PIL.Image.Image: - scale = min(target_w / image.width, target_h / image.height, 1.0) - resize_w = max(1, int(scale * image.width + 0.5)) - resize_h = max(1, int(scale * image.height + 0.5)) - if (resize_w, resize_h) != image.size: - image = image.resize((resize_w, resize_h), PIL.Image.Resampling.BICUBIC) - - array = np.asarray(image) - pad_h = target_h - resize_h - pad_w = target_w - resize_w - if pad_h < 0 or pad_w < 0: - raise ValueError( - f"Cosmos3 action image resize exceeded target size: resized={(resize_h, resize_w)}, " - f"target={(target_h, target_w)}." - ) - if pad_h == 0 and pad_w == 0: - return image - pad_mode = "reflect" if pad_h < resize_h and pad_w < resize_w else "edge" - padded = np.pad(array, ((0, pad_h), (0, pad_w), (0, 0)), mode=pad_mode) - return PIL.Image.fromarray(padded) + # ========================================================================= + # Action generation + # ========================================================================= def _preprocess_action_image( self, image: PIL.Image.Image, target_h: int, target_w: int ) -> torch.Tensor: - image = self._resize_and_pad_action_image(image, target_h, target_w) + image = resize_and_pad_action_image(image, target_h, target_w) return self.video_processor.preprocess(image, height=target_h, width=target_w) def _preprocess_action_video( @@ -722,7 +633,7 @@ def _preprocess_action_video( if not frames: raise ValueError("Cosmos3 action video input must contain at least one frame.") processed = [ - self._preprocess_action_image(self._pil_to_rgb(frame), target_h, target_w).squeeze(0) + self._preprocess_action_image(pil_to_rgb(frame), target_h, target_w).squeeze(0) for frame in frames ] return torch.stack(processed, dim=1).unsqueeze(0).contiguous() @@ -801,48 +712,16 @@ def _prepare_action_latents( generator: torch.Generator, action_input: Any = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: - action_dim = int(getattr(self.transformer, "action_dim", 64)) - if mode == ACTION_MODE_FORWARD_DYNAMICS: - action = load_action_tensor(action_input) - if action.shape[0] < action_chunk_size: - pad = action[-1:].repeat(action_chunk_size - action.shape[0], 1) - action = torch.cat([action, pad], dim=0) - elif action.shape[0] > action_chunk_size: - action = action[:action_chunk_size] - if raw_action_dim is None: - raw_action_dim = int(action.shape[-1]) - clean_action = pad_action_to_dim(action, action_dim) - else: - if raw_action_dim is None: - raise ValueError( - "Cosmos3 action_mode='policy' and 'inverse_dynamics' require raw_action_dim." - ) - clean_action = torch.zeros(action_chunk_size, action_dim, dtype=torch.float32) - - raw_action_dim = int(raw_action_dim) - if raw_action_dim <= 0 or raw_action_dim > action_dim: - raise ValueError( - f"Cosmos3 raw_action_dim must be in [1, {action_dim}], got {raw_action_dim}." - ) - - clean_action = clean_action.to(device=self.device, dtype=self.dtype).unsqueeze(0) - condition_mask = build_action_condition_mask( - mode, - action_chunk_size, - device=self.device, - dtype=self.dtype, - ) - noise = randn_tensor( - (1, action_chunk_size, action_dim), + return prepare_action_latents( + mode=mode, + action_chunk_size=action_chunk_size, + raw_action_dim=raw_action_dim, + action_dim=int(getattr(self.transformer, "action_dim", 64)), generator=generator, device=self.device, dtype=self.dtype, + action_input=action_input, ) - noise[:, :, raw_action_dim:] = 0 - clean_action[:, :, raw_action_dim:] = 0 - action_latents = condition_mask * clean_action + (1.0 - condition_mask) * noise - action_velocity_mask = 1.0 - condition_mask - return action_latents, action_velocity_mask, clean_action, raw_action_dim # ========================================================================= # Forward (main generation entry point) @@ -876,7 +755,7 @@ def forward( raw_action_dim: Optional[int] = None, action_chunk_size: Optional[int] = None, action: Any = None, - action_resolution: int = 480, + action_resolution: Optional[int] = None, video: Any = None, ): pipeline_start = time.time() @@ -918,15 +797,38 @@ def forward( guidance_interval = COSMOS3_T2I_PARAMS["guidance_interval"] self._set_flow_shift(COSMOS3_T2I_PARAMS["flow_shift"]) elif do_action: - action_chunk_size = action_chunk_size or COSMOS3_ACTION_PARAMS["action_chunk_size"] - if num_frames is None: - num_frames = COSMOS3_ACTION_PARAMS["num_frames"] + action_cfg = resolve_domain_action_config( + domain_name=domain_name, + domain_id=domain_id, + raw_action_dim=raw_action_dim, + action_chunk_size=action_chunk_size, + action_resolution=action_resolution, + frame_rate=frame_rate, + num_frames=num_frames, + ) + if self.rank == 0: + for warning in action_cfg["warnings"]: + logger.warning(warning) + if action_cfg["preset_key"] is not None: + logger.info( + f"Cosmos3 action domain preset {action_cfg['preset_key']!r}: " + f"raw_action_dim={action_cfg['raw_action_dim']}, " + f"action_chunk_size={action_cfg['action_chunk_size']}, " + f"action_resolution={action_cfg['action_resolution']}, " + f"frame_rate={action_cfg['frame_rate']:.1f}, " + f"num_frames={action_cfg['num_frames']}" + ) + + raw_action_dim = action_cfg["raw_action_dim"] + action_chunk_size = action_cfg["action_chunk_size"] + action_resolution = action_cfg["action_resolution"] + num_frames = action_cfg["num_frames"] + frame_rate = action_cfg["frame_rate"] num_inference_steps = ( num_inference_steps or COSMOS3_ACTION_PARAMS["num_inference_steps"] ) if guidance_scale is None: guidance_scale = COSMOS3_ACTION_PARAMS["guidance_scale"] - frame_rate = frame_rate or COSMOS3_ACTION_PARAMS["frame_rate"] self._set_flow_shift(COSMOS3_ACTION_PARAMS["flow_shift"]) enable_audio = False else: @@ -946,33 +848,24 @@ def forward( action_ref_image = None if do_action: - action_ref_image = self._action_reference_image( + action_ref_image = action_reference_image( action_mode=normalized_action_mode, image=image, video=video, ) - height, width = self._resolve_action_size( - height, width, action_ref_image, action_resolution - ) + height, width = resolve_action_size(height, width, action_ref_image, action_resolution) if self.rank == 0: logger.info( - "Cosmos3 generation dims: %dx%d (WxH), num_frames=%d, " - "num_inference_steps=%d, guidance_scale=%.2f, frame_rate=%.1f", - width, - height, - num_frames, - num_inference_steps, - guidance_scale, - frame_rate, + f"Cosmos3 generation dims: {width}x{height} (WxH), num_frames={num_frames}, " + f"num_inference_steps={num_inference_steps}, guidance_scale={guidance_scale:.2f}, " + f"frame_rate={frame_rate:.1f}" ) if do_action: logger.info( - "Cosmos3 action dims: action_chunk_size=%d, action_resolution=%s, " - "input_aspect=%.3f", - action_chunk_size, - action_resolution, - action_ref_image.width / action_ref_image.height, + f"Cosmos3 action dims: action_chunk_size={action_chunk_size}, " + f"action_resolution={action_resolution}, " + f"input_aspect={action_ref_image.width / action_ref_image.height:.3f}" ) if isinstance(prompt, str): @@ -1092,7 +985,8 @@ def forward( ) if normalized_action_mode == ACTION_MODE_INVERSE_DYNAMICS: - video = self._normalize_action_video_input(video) + inverse_video = video if video is not None else image + video = normalize_action_video_input(inverse_video, max_frames=num_frames) video_tensor = self._preprocess_action_video(video, height, width) latents, velocity_mask, condition_latents = self._prepare_latents_action_video( video_tensor, @@ -1158,6 +1052,8 @@ def forward( # 3. Set up scheduler self.scheduler.set_timesteps(num_inference_steps, device=self.device) + if do_action: + self.action_scheduler.set_timesteps(num_inference_steps, device=self.device) # 3b. Audio noise init — latent length matches diffusers Cosmos3OmniPipeline.prepare_latents. do_audio = enable_audio and self.audio_gen and hasattr(self, "audio_tokenizer") @@ -1284,7 +1180,7 @@ def post_step_fn(step_latents, step_extra_stream_latents): timer.mark_denoise_start() extra_streams = None if do_action: - extra_streams = {"action": (action_latents, self.scheduler)} + extra_streams = {"action": (action_latents, self.action_scheduler)} elif do_audio: extra_streams = {"audio": (audio_latents, self.audio_scheduler)} # FUTURE(action+audio): merge both keys; extend forward_fn return dict and post_step_fn. diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py index 166229b3eaa1..0840ee763bab 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -190,10 +190,13 @@ def __init__( self.input_size = int(input_size) self.output_size = int(output_size) self.num_domains = int(num_domains) + self.dtype = dtype self.fc = nn.Embedding(self.num_domains, self.output_size * self.input_size, dtype=dtype) self.bias = nn.Embedding(self.num_domains, self.output_size, dtype=dtype) - nn.init.xavier_uniform_(self.fc.weight) - nn.init.zeros_(self.bias.weight) + + def post_load_weights(self) -> None: + self.fc.to(self.dtype) + self.bias.to(self.dtype) def forward(self, x: torch.Tensor, domain_id: torch.Tensor) -> torch.Tensor: if domain_id.ndim == 0: @@ -1545,6 +1548,8 @@ def post_load_weights(self) -> None: if self.action_gen: self.action_modality_embed.data = self.action_modality_embed.data.to(target_dtype) + self.action_proj_in.post_load_weights() + self.action_proj_out.post_load_weights() for _, module in self.named_modules(): if isinstance(module, Linear) or isinstance(module, Qwen3VLTextRMSNorm): diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py index b85b1d50b75c..673534c9fbab 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py @@ -7,15 +7,18 @@ pytest tests/unittest/_torch/visual_gen/test_cosmos3_action.py -v """ +import numpy as np import PIL.Image import pytest from tensorrt_llm._torch.visual_gen.models.cosmos3.action import ( VIDEO_RES_SIZE_INFO, + action_reference_image, find_closest_target_size, + normalize_action_video_input, + resolve_action_size, ) from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import COSMOS3_EXTRA_SPECS -from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniMoTPipeline pytestmark = pytest.mark.cosmos3 @@ -57,19 +60,19 @@ def _ref_image(width: int, height: int) -> PIL.Image.Image: def test_explicit_height_and_width_are_unchanged(self): ref = self._ref_image(832, 480) - assert Cosmos3OmniMoTPipeline._resolve_action_size(400, 600, ref, 480) == (400, 600) + assert resolve_action_size(400, 600, ref, 480) == (400, 600) def test_unset_height_and_width_use_action_resolution_bucket(self): ref = self._ref_image(832, 480) - assert Cosmos3OmniMoTPipeline._resolve_action_size(None, None, ref, 480) == (480, 832) + assert resolve_action_size(None, None, ref, 480) == (480, 832) def test_partial_height_fills_width_from_bucket(self): ref = self._ref_image(832, 480) - assert Cosmos3OmniMoTPipeline._resolve_action_size(400, None, ref, 480) == (400, 832) + assert resolve_action_size(400, None, ref, 480) == (400, 832) def test_partial_width_fills_height_from_bucket(self): ref = self._ref_image(832, 480) - assert Cosmos3OmniMoTPipeline._resolve_action_size(None, 600, ref, 480) == (480, 600) + assert resolve_action_size(None, 600, ref, 480) == (480, 600) class TestActionResolutionExtraParam: @@ -77,3 +80,157 @@ def test_extra_param_spec_uses_action_resolution_key(self): spec = COSMOS3_EXTRA_SPECS["action_resolution"] assert spec.type == "int" assert spec.default == 480 + + +class TestDomainActionPresets: + def test_bridge_preset_fills_missing_fields(self): + from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( + resolve_domain_action_config, + ) + + cfg = resolve_domain_action_config(domain_name="bridge_orig_lerobot") + assert cfg["raw_action_dim"] == 10 + assert cfg["action_chunk_size"] == 16 + assert cfg["num_frames"] == 17 + assert cfg["action_resolution"] == 480 + assert cfg["frame_rate"] == 5.0 + assert cfg["warnings"] == [] + + def test_av_preset_uses_longer_chunk(self): + from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( + resolve_domain_action_config, + ) + + cfg = resolve_domain_action_config(domain_name="av") + assert cfg["action_chunk_size"] == 60 + assert cfg["num_frames"] == 61 + assert cfg["raw_action_dim"] == 9 + + def test_mismatch_emits_warning(self): + from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( + resolve_domain_action_config, + ) + + cfg = resolve_domain_action_config( + domain_name="bridge_orig_lerobot", + raw_action_dim=9, + ) + assert cfg["raw_action_dim"] == 9 + assert len(cfg["warnings"]) == 1 + assert "raw_action_dim=9" in cfg["warnings"][0] + + def test_alias_maps_to_canonical_preset(self): + from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import get_domain_preset + + preset = get_domain_preset("robomind-franka") + assert preset is not None + assert preset["raw_action_dim"] == 10 + + def test_unknown_resolution_raises(self): + from tensorrt_llm._torch.visual_gen.models.cosmos3.action import normalize_action_resolution + + with pytest.raises(ValueError, match="Unknown Cosmos3 action_resolution"): + normalize_action_resolution(1080) + + +class TestActionReferenceImage: + def test_forward_dynamics_accepts_mp4_on_image_path(self, tmp_path, monkeypatch): + video_path = tmp_path / "clip.mp4" + video_path.write_bytes(b"fake") + expected = PIL.Image.new("RGB", (4, 2), "red") + + def _fake_read_video(path, pts_unit="sec"): + import torch + + tensor = torch.from_numpy(np.array(expected)).unsqueeze(0) + return tensor, None, {} + + monkeypatch.setattr("torchvision.io.read_video", _fake_read_video) + ref = action_reference_image( + action_mode="forward_dynamics", + image=str(video_path), + video=None, + ) + assert ref.size == expected.size + assert ref.getpixel((0, 0)) == (255, 0, 0) + + def test_policy_prefers_image_path_over_video(self, tmp_path): + image_path = tmp_path / "frame.png" + PIL.Image.new("RGB", (3, 3), "blue").save(image_path) + ref = action_reference_image( + action_mode="policy", + image=str(image_path), + video=str(tmp_path / "unused.mp4"), + ) + assert ref.getpixel((0, 0)) == (0, 0, 255) + + +class TestNormalizeActionVideoInput: + def test_image_path_returns_singleton_list(self, tmp_path): + image_path = tmp_path / "frame.png" + PIL.Image.new("RGB", (8, 4), "red").save(image_path) + assert normalize_action_video_input(str(image_path)) == [str(image_path)] + + def test_frame_directory_returns_sorted_paths(self, tmp_path): + (tmp_path / "b.png").write_bytes(b"") + (tmp_path / "a.png").write_bytes(b"") + (tmp_path / "skip.txt").write_text("x") + assert normalize_action_video_input(str(tmp_path)) == [ + str(tmp_path / "a.png"), + str(tmp_path / "b.png"), + ] + + def test_unsupported_file_extension_raises(self, tmp_path): + bad_path = tmp_path / "clip.mov" + bad_path.write_bytes(b"fake") + with pytest.raises(ValueError, match="must be a frame directory"): + normalize_action_video_input(str(bad_path)) + + def test_decode_mp4_returns_pil_frames(self, tmp_path, monkeypatch): + video_path = tmp_path / "clip.mp4" + video_path.write_bytes(b"fake") + expected = [ + PIL.Image.new("RGB", (2, 2), "red"), + PIL.Image.new("RGB", (2, 2), "blue"), + ] + + def _fake_read_video(path, pts_unit="sec"): + assert path == str(video_path) + assert pts_unit == "sec" + import torch + + tensor = torch.stack( + [torch.from_numpy(np.array(image)) for image in expected], + dim=0, + ) + return tensor, None, {} + + monkeypatch.setattr( + "torchvision.io.read_video", + _fake_read_video, + ) + frames = normalize_action_video_input(str(video_path)) + assert len(frames) == 2 + assert all(isinstance(frame, PIL.Image.Image) for frame in frames) + assert frames[0].getpixel((0, 0)) == (255, 0, 0) + + def test_decode_respects_max_frames(self, tmp_path, monkeypatch): + video_path = tmp_path / "clip.avi" + video_path.write_bytes(b"fake") + images = [PIL.Image.new("RGB", (1, 1), color) for color in ("red", "green", "blue")] + + def _fake_read_video(path, pts_unit="sec"): + import torch + + tensor = torch.stack( + [torch.from_numpy(np.array(image)) for image in images], + dim=0, + ) + return tensor, None, {} + + monkeypatch.setattr("torchvision.io.read_video", _fake_read_video) + frames = normalize_action_video_input( + str(video_path), + max_frames=2, + ) + assert len(frames) == 2 From 9841567ac5b8f72a71541985aa48b14d2ea3028e Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Thu, 11 Jun 2026 09:32:46 -0700 Subject: [PATCH 03/64] add action fps optional input Signed-off-by: Shreyas Misra --- examples/visual_gen/models/cosmos3/cosmos3.py | 8 ++++++++ .../visual_gen/models/cosmos3/defaults.py | 12 ++++++++++++ .../models/cosmos3/pipeline_cosmos3.py | 11 ++++++++++- .../_torch/visual_gen/test_cosmos3_action.py | 18 ++++++++++++++++++ 4 files changed, 48 insertions(+), 1 deletion(-) diff --git a/examples/visual_gen/models/cosmos3/cosmos3.py b/examples/visual_gen/models/cosmos3/cosmos3.py index 368399101d5d..01f804fc83a6 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3.py +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -413,6 +413,12 @@ def main(): choices=[256, 480, 704, 720], help=("Resolution bucket for action image sizing. Defaults to the domain preset or 480."), ) + parser.add_argument( + "--action_fps", + type=float, + default=None, + help="Action-token temporal rate for mRoPE (Hz). Defaults to frame_rate.", + ) parser.add_argument( "--action_output_path", type=str, @@ -481,6 +487,8 @@ def main(): params.extra_params["raw_action_dim"] = args.raw_action_dim if args.action_chunk_size is not None: params.extra_params["action_chunk_size"] = args.action_chunk_size + if args.action_fps is not None: + params.extra_params["action_fps"] = args.action_fps if args.action_json is not None: with open(args.action_json, encoding="utf-8") as f: params.extra_params["action"] = json.load(f) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index 60a3a496a7a9..b01e98334276 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -238,6 +238,7 @@ def resolve_domain_action_config( action_chunk_size: Optional[int] = None, action_resolution: Optional[int] = None, frame_rate: Optional[float] = None, + action_fps: Optional[float] = None, num_frames: Optional[int] = None, ) -> Dict[str, Any]: """Merge user action params with domain presets and generic fallbacks.""" @@ -284,12 +285,16 @@ def _resolve_field( resolved_num_frames = _resolve_field("num_frames", num_frames) if resolved_num_frames is None: resolved_num_frames = int(resolved_chunk) + 1 + resolved_action_fps = ( + float(action_fps) if action_fps is not None else float(resolved_frame_rate) + ) return { "raw_action_dim": resolved_raw_action_dim, "action_chunk_size": int(resolved_chunk), "action_resolution": resolved_resolution, "frame_rate": float(resolved_frame_rate), + "action_fps": resolved_action_fps, "num_frames": int(resolved_num_frames), "preset_key": preset_key, "warnings": warnings, @@ -376,6 +381,13 @@ def _resolve_field( ), range=(min(COSMOS3_ACTION_RESOLUTIONS), max(COSMOS3_ACTION_RESOLUTIONS)), ), + "action_fps": ExtraParamSchema( + type="float", + default=None, + description=( + "Action-token temporal rate for mRoPE (Hz). Defaults to frame_rate when omitted." + ), + ), "video": ExtraParamSchema( type="path_or_list", default=None, diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index 8b09f7cf4720..ada83ed31324 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -303,6 +303,7 @@ def infer(self, req): action=extra_params.get("action"), action_resolution=extra_params.get("action_resolution") or extra_params.get("image_size"), + action_fps=extra_params.get("action_fps"), video=extra_params.get("video"), ) @@ -756,6 +757,7 @@ def forward( action_chunk_size: Optional[int] = None, action: Any = None, action_resolution: Optional[int] = None, + action_fps: Optional[float] = None, video: Any = None, ): pipeline_start = time.time() @@ -779,6 +781,7 @@ def forward( # guidance interval, and an image (rather than video) output. is_t2i = str(output_type).lower() == "image" guidance_interval = None + resolved_action_fps: Optional[float] = None if is_t2i: if image is not None: raise ValueError( @@ -804,6 +807,7 @@ def forward( action_chunk_size=action_chunk_size, action_resolution=action_resolution, frame_rate=frame_rate, + action_fps=action_fps, num_frames=num_frames, ) if self.rank == 0: @@ -816,6 +820,7 @@ def forward( f"action_chunk_size={action_cfg['action_chunk_size']}, " f"action_resolution={action_cfg['action_resolution']}, " f"frame_rate={action_cfg['frame_rate']:.1f}, " + f"action_fps={action_cfg['action_fps']:.1f}, " f"num_frames={action_cfg['num_frames']}" ) @@ -824,6 +829,7 @@ def forward( action_resolution = action_cfg["action_resolution"] num_frames = action_cfg["num_frames"] frame_rate = action_cfg["frame_rate"] + resolved_action_fps = action_cfg["action_fps"] num_inference_steps = ( num_inference_steps or COSMOS3_ACTION_PARAMS["num_inference_steps"] ) @@ -845,6 +851,8 @@ def forward( max_sequence_length = max_sequence_length or COSMOS3_720P_PARAMS["max_sequence_length"] if frame_rate is None: frame_rate = COSMOS3_720P_PARAMS["frame_rate"] + if resolved_action_fps is None: + resolved_action_fps = frame_rate action_ref_image = None if do_action: @@ -865,6 +873,7 @@ def forward( logger.info( f"Cosmos3 action dims: action_chunk_size={action_chunk_size}, " f"action_resolution={action_resolution}, " + f"action_fps={resolved_action_fps:.1f}, " f"input_aspect={action_ref_image.width / action_ref_image.height:.3f}" ) @@ -1118,7 +1127,7 @@ def forward_fn( action_domain_ids=action_domain_ids, action_noisy_mask=action_velocity_mask, action_start_frame_offset=action_frame_offset, - action_fps=frame_rate, + action_fps=resolved_action_fps, ) video_noise_pred = result.video diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py index 673534c9fbab..9a64ebdf4048 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py @@ -119,6 +119,24 @@ def test_mismatch_emits_warning(self): assert len(cfg["warnings"]) == 1 assert "raw_action_dim=9" in cfg["warnings"][0] + def test_action_fps_defaults_to_frame_rate(self): + from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( + resolve_domain_action_config, + ) + + cfg = resolve_domain_action_config(domain_name="av") + assert cfg["frame_rate"] == 10.0 + assert cfg["action_fps"] == 10.0 + + def test_explicit_action_fps_overrides_default(self): + from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( + resolve_domain_action_config, + ) + + cfg = resolve_domain_action_config(domain_name="av", action_fps=5.0, frame_rate=24.0) + assert cfg["frame_rate"] == 24.0 + assert cfg["action_fps"] == 5.0 + def test_alias_maps_to_canonical_preset(self): from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import get_domain_preset From 2cc00f236b91b2282cbf61b5c93bfd928eef3b9f Mon Sep 17 00:00:00 2001 From: Bartosz Stefaniak Date: Wed, 24 Jun 2026 14:57:19 +0000 Subject: [PATCH 04/64] V2V for cosmos3 Signed-off-by: Bartosz Stefaniak --- .../visual_gen/models/cosmos3/action.py | 78 +---- .../visual_gen/models/cosmos3/defaults.py | 25 +- .../models/cosmos3/pipeline_cosmos3.py | 280 +++++++++++++++-- .../_torch/visual_gen/models/cosmos3/utils.py | 78 +++++ .../_torch/visual_gen/test_cosmos3_action.py | 14 +- .../visual_gen/test_cosmos3_pipeline.py | 288 +++++++++++++++++- 6 files changed, 651 insertions(+), 112 deletions(-) create mode 100644 tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py index b66dda5b6ac2..2d4a1c448988 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py @@ -6,13 +6,15 @@ from __future__ import annotations from pathlib import Path -from typing import Any, List, Optional +from typing import Any, Optional import numpy as np import PIL.Image import torch from diffusers.utils.torch_utils import randn_tensor +from .utils import IMAGE_EXTENSIONS, normalize_video_input, pil_to_rgb + ACTION_MODE_POLICY = "policy" ACTION_MODE_FORWARD_DYNAMICS = "forward_dynamics" ACTION_MODE_INVERSE_DYNAMICS = "inverse_dynamics" @@ -244,74 +246,6 @@ def find_closest_target_size(h: int, w: int, resolution: str | int) -> tuple[int return best_size -ACTION_IMAGE_EXTENSIONS = frozenset({".png", ".jpg", ".jpeg", ".webp", ".bmp"}) -ACTION_VIDEO_EXTENSIONS = frozenset({".mp4", ".avi"}) - - -def pil_to_rgb(value: Any) -> PIL.Image.Image: - if isinstance(value, str): - return PIL.Image.open(value).convert("RGB") - if isinstance(value, PIL.Image.Image): - return value.convert("RGB") - raise TypeError( - f"Cosmos3 action preprocessing expected PIL image or image path, got {type(value)!r}." - ) - - -def decode_action_video_file(path: Path, max_frames: Optional[int] = None) -> List[PIL.Image.Image]: - import torchvision.io as io - - frames, _, _ = io.read_video(str(path), pts_unit="sec") - if frames.numel() == 0: - raise ValueError(f"Cosmos3 action video file contains no frames: {path}") - if max_frames is not None: - frames = frames[:max_frames] - return [PIL.Image.fromarray(frames[i].numpy()) for i in range(frames.shape[0])] - - -def normalize_action_video_path(path: Path, max_frames: Optional[int] = None) -> List[Any]: - if not path.exists(): - raise ValueError(f"Cosmos3 action video path does not exist: {path}") - if path.is_dir(): - frames = sorted(p for p in path.iterdir() if p.suffix.lower() in ACTION_IMAGE_EXTENSIONS) - if not frames: - raise ValueError(f"No image frames found in Cosmos3 action video directory: {path}") - frame_paths = [str(p) for p in frames] - if max_frames is not None: - frame_paths = frame_paths[:max_frames] - return frame_paths - - suffix = path.suffix.lower() - if suffix in ACTION_IMAGE_EXTENSIONS: - return [str(path)] - if suffix in ACTION_VIDEO_EXTENSIONS: - return decode_action_video_file(path, max_frames=max_frames) - raise ValueError( - "Cosmos3 action video path must be a frame directory, an image file " - f"{sorted(ACTION_IMAGE_EXTENSIONS)}, or a video file " - f"{sorted(ACTION_VIDEO_EXTENSIONS)}; got {path}" - ) - - -def normalize_action_video_input(video: Any, max_frames: Optional[int] = None) -> List[Any]: - """Normalize action video input to a frame list. - - Accepts a list of PIL images / paths, a single image or video file path, - or a directory of frame images (sorted lexicographically). - """ - if video is None: - return [] - if isinstance(video, list): - if not video: - raise ValueError("Cosmos3 action video input must contain at least one frame.") - if max_frames is not None: - return video[:max_frames] - return video - if isinstance(video, (str, Path)): - return normalize_action_video_path(Path(video), max_frames=max_frames) - return [video] - - def resolve_action_size( height: Optional[int], width: Optional[int], @@ -339,7 +273,7 @@ def action_reference_image( """Resolve the reference frame used for action sizing and conditioning.""" if action_mode == ACTION_MODE_INVERSE_DYNAMICS: source = video if video is not None else image - frames = normalize_action_video_input(source, max_frames=1) + frames = normalize_video_input(source, max_frames=1) if not frames: raise ValueError("Cosmos3 action_mode='inverse_dynamics' requires a video input.") return pil_to_rgb(frames[0]) @@ -351,9 +285,9 @@ def action_reference_image( return source.convert("RGB") if isinstance(source, str): path = Path(source) - if path.is_file() and path.suffix.lower() in ACTION_IMAGE_EXTENSIONS: + if path.is_file() and path.suffix.lower() in IMAGE_EXTENSIONS: return PIL.Image.open(source).convert("RGB") - frames = normalize_action_video_input(source, max_frames=1) + frames = normalize_video_input(source, max_frames=1) if not frames: raise ValueError( f"Cosmos3 action_mode={action_mode!r} requires an image or video input." diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index b01e98334276..ae1d4cdfc2bc 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -49,6 +49,9 @@ "frame_rate": 24.0, } +COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION = (0, 1) +COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP = "first" + # Fields merged by the executor for every request. Modality-specific values # (height/width/num_frames/steps/guidance) are declared with ``None`` so the # executor accepts explicit overrides; ``forward()`` resolves unset fields from @@ -314,7 +317,7 @@ def _resolve_field( ), "use_system_prompt": ExtraParamSchema( type="bool", - default=False, + default=None, description="Whether to use the system prompt.", ), "use_guardrails": ExtraParamSchema( @@ -388,12 +391,28 @@ def _resolve_field( "Action-token temporal rate for mRoPE (Hz). Defaults to frame_rate when omitted." ), ), + "condition_frame_indexes_vision": ExtraParamSchema( + type="list", + default=list(COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION), + description="Latent frame indexes to keep fixed for video conditioning.", + ), + "condition_video_keep": ExtraParamSchema( + type="str", + default=COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP, + description="Which side of the input video to use for conditioning: first or last.", + ), + "flow_shift": ExtraParamSchema( + type="float", + default=None, + description="Optional scheduler flow shift override. Uses the Cosmos3 mode default when omitted.", + ), "video": ExtraParamSchema( type="path_or_list", default=None, description=( - "Video for inverse_dynamics: .mp4/.avi file, frame directory, " - "image path, or list of PIL images / frame paths." + "Video input for video-to-video generation or inverse_dynamics: " + ".mp4/.avi file, frame directory, image path, or list of PIL " + "images / frame paths." ), ), } diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index ada83ed31324..810b15fa8bd1 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -17,7 +17,7 @@ import math import os import time -from typing import Any, List, Optional, Union +from typing import Any, Iterable, List, Optional, Union import PIL.Image import torch @@ -40,8 +40,6 @@ action_start_frame_offset, build_vision_condition_mask, normalize_action_mode, - normalize_action_video_input, - pil_to_rgb, prepare_action_latents, resize_and_pad_action_image, resolve_action_size, @@ -50,6 +48,8 @@ from .defaults import ( COSMOS3_720P_PARAMS, COSMOS3_ACTION_PARAMS, + COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION, + COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP, COSMOS3_EXTRA_SPECS, COSMOS3_PIPELINE_DEFAULTS, COSMOS3_T2I_PARAMS, @@ -58,13 +58,15 @@ from .guardrails import check_video_safety, download_guardrail_checkpoint from .sound_tokenizer import LatentAutoEncoderV2 from .transformer_cosmos3 import Cosmos3VFMTransformer +from .utils import normalize_video_input, pil_to_rgb COSMOS3_DEFAULT_NEGATIVE_PROMPT = "" +# NOTE: Intentional typo in "give" instead of "given" to match training setup. COSMOS3_DEFAULT_SYSTEM_PROMPT = ( - "You are a helpful assistant who will generate videos from a given prompt." + "You are a helpful assistant who will generate videos from a give prompt." ) COSMOS3_T2I_SYSTEM_PROMPT = ( - "You are a helpful assistant who will generate images from a given prompt." + "You are a helpful assistant who will generate images from a give prompt." ) COSMOS3_DURATION_TEMPLATE = "The video is {duration:.1f} seconds long and is of {fps:.0f} FPS." COSMOS3_DEFAULT_RESOLUTION_TEMPLATE = "This video is of {height}x{width} resolution." @@ -73,6 +75,43 @@ TRTLLM_DISABLE_COSMOS3_GUARDRAILS = os.environ.get("TRTLLM_DISABLE_COSMOS3_GUARDRAILS", "0") == "1" +def _normalize_condition_frame_indexes_vision( + indexes: Iterable[int] | int | str | None, +) -> tuple[int, ...]: + if indexes is None: + return COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION + if isinstance(indexes, int): + normalized = (indexes,) + elif isinstance(indexes, str): + parts = [part.strip() for part in indexes.split(",") if part.strip()] + normalized = tuple(int(part) for part in parts) + else: + normalized = tuple(int(index) for index in indexes) + + if not normalized: + raise ValueError("Cosmos3 condition_frame_indexes_vision must not be empty.") + if any(index < 0 for index in normalized): + raise ValueError( + "Cosmos3 condition_frame_indexes_vision must be non-negative, " + f"got {normalized}." + ) + return normalized + + +def _condition_pixel_frame_count( + condition_frame_indexes_vision: Iterable[int], + temporal_compression: int, +) -> int: + return max(condition_frame_indexes_vision) * int(temporal_compression) + 1 + + +def _normalize_condition_video_keep(keep: str | None) -> str: + normalized = str(keep or COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP).strip().lower() + if normalized not in {"first", "last"}: + raise ValueError("Cosmos3 condition_video_keep must be either first or last.") + return normalized + + @register_pipeline( "Cosmos3OmniMoTPipeline", hf_ids=[ @@ -165,12 +204,16 @@ def load_standard_components( ) # Snapshot the checkpoint scheduler config so the scheduler can be # rebuilt at request time when a mode-specific ``flow_shift`` is - # needed (T2I uses shift=3.0; T2V/I2V keep the checkpoint default). + # needed. self._base_scheduler_config = self.scheduler.config self._engine_init_flow_shift = float( getattr(self.scheduler.config, "flow_shift", 1.0) or 1.0 ) self._current_flow_shift = self._engine_init_flow_shift + self._base_scheduler_use_karras_sigmas = self._scheduler_use_karras_sigmas( + self.scheduler.config + ) + self._current_scheduler_use_karras_sigmas = self._base_scheduler_use_karras_sigmas if self.audio_gen: # Separate instance so video and audio scheduler states don't collide # (UniPC mutates internal correction buffers on every .step() call). @@ -209,30 +252,52 @@ def load_standard_components( self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) - def _set_flow_shift(self, target_shift: float) -> None: - """Rebuild the UniPC scheduler with ``flow_shift=target_shift`` if needed. + @staticmethod + def _scheduler_use_karras_sigmas(config: Any) -> Optional[bool]: + value = getattr(config, "use_karras_sigmas", None) + return None if value is None else bool(value) - T2I uses ``flow_shift=3.0`` while T2V/I2V use the checkpoint default. - ``self._current_flow_shift`` is tracked explicitly so a prior T2I rebuild - does not leak into a subsequent video request. + def _set_flow_shift( + self, target_shift: float, *, use_karras_sigmas: Optional[bool] = None + ) -> None: + """Rebuild the UniPC scheduler when request scheduler defaults change. + + The effective flow-shift changes when switching between mode defaults + (T2I=3.0, action=5.0, V2V=10.0, T2V/I2V=checkpoint default) or when a + request provides ``flow_shift``. V2V also forces Karras sigmas off. """ if not hasattr(self, "_base_scheduler_config"): return target = float(target_shift) - if target == float(self._current_flow_shift): + target_use_karras_sigmas = ( + self._base_scheduler_use_karras_sigmas + if use_karras_sigmas is None + else bool(use_karras_sigmas) + ) + if ( + target == float(self._current_flow_shift) + and target_use_karras_sigmas == self._current_scheduler_use_karras_sigmas + ): return + + scheduler_kwargs = {"flow_shift": target} + if use_karras_sigmas is not None: + scheduler_kwargs["use_karras_sigmas"] = bool(use_karras_sigmas) self.scheduler = UniPCMultistepScheduler.from_config( - self._base_scheduler_config, flow_shift=target + self._base_scheduler_config, **scheduler_kwargs ) if self.audio_gen: self.audio_scheduler = UniPCMultistepScheduler.from_config( - self._base_scheduler_config, flow_shift=target + self._base_scheduler_config, **scheduler_kwargs ) if self.action_gen: self.action_scheduler = UniPCMultistepScheduler.from_config( - self._base_scheduler_config, flow_shift=target + self._base_scheduler_config, **scheduler_kwargs ) self._current_flow_shift = target + self._current_scheduler_use_karras_sigmas = self._scheduler_use_karras_sigmas( + self.scheduler.config + ) @property def default_warmup_resolutions(self): @@ -291,7 +356,9 @@ def infer(self, req): "use_resolution_template", COSMOS3_EXTRA_SPECS["use_resolution_template"].default, ), - use_system_prompt=extra_params.get("use_system_prompt", False), + use_system_prompt=extra_params.get( + "use_system_prompt", COSMOS3_EXTRA_SPECS["use_system_prompt"].default + ), use_guardrails=extra_params.get("use_guardrails", True), enable_audio=extra_params.get("enable_audio", False), output_type=output_type, @@ -305,6 +372,9 @@ def infer(self, req): or extra_params.get("image_size"), action_fps=extra_params.get("action_fps"), video=extra_params.get("video"), + condition_frame_indexes_vision=extra_params.get("condition_frame_indexes_vision"), + condition_video_keep=extra_params.get("condition_video_keep"), + flow_shift=extra_params.get("flow_shift"), ) def _apply_metadata_templates( @@ -639,6 +709,21 @@ def _preprocess_action_video( ] return torch.stack(processed, dim=1).unsqueeze(0).contiguous() + def _preprocess_condition_video( + self, frames: List[Any], target_h: int, target_w: int + ) -> torch.Tensor: + if not frames: + raise ValueError("Cosmos3 condition video input must contain at least one frame.") + processed = [ + self.video_processor.preprocess( + self._resize_and_center_crop_image(pil_to_rgb(frame), target_h, target_w), + height=target_h, + width=target_w, + ).squeeze(0) + for frame in frames + ] + return torch.stack(processed, dim=1).unsqueeze(0).contiguous() + def _encode_video_tensor(self, video_tensor: torch.Tensor) -> torch.Tensor: """VAE-encode a preprocessed pixel video [1, 3, T, H, W].""" if video_tensor.ndim == 4: @@ -724,6 +809,79 @@ def _prepare_action_latents( action_input=action_input, ) + # ========================================================================= + # Video to video + # ========================================================================= + + def _prepare_latents_v2v( + self, + video_tensor: torch.Tensor, + num_frames: int, + generator: torch.Generator, + condition_frame_indexes_vision: Iterable[int] | int | str | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Prepare V2V latents with explicit clean conditioned latent frames.""" + if video_tensor.ndim == 4: + video_tensor = video_tensor.unsqueeze(0) + if video_tensor.ndim != 5 or video_tensor.shape[0] != 1 or video_tensor.shape[1] != 3: + raise ValueError( + "Cosmos3 video tensor must have shape [1, 3, T, H, W], " + f"got {tuple(video_tensor.shape)}." + ) + if video_tensor.shape[2] < 1: + raise ValueError("Cosmos3 V2V video tensor must contain at least one frame.") + + C = self.transformer.latent_channel_size + T_lat = (num_frames - 1) // self.vae_scale_factor_temporal + 1 + H_lat = video_tensor.shape[-2] // self.vae_scale_factor_spatial + W_lat = video_tensor.shape[-1] // self.vae_scale_factor_spatial + indexes = _normalize_condition_frame_indexes_vision(condition_frame_indexes_vision) + out_of_range = [index for index in indexes if index >= T_lat] + if out_of_range: + raise ValueError( + "Cosmos3 condition_frame_indexes_vision contains indexes outside the latent video: " + f"indexes={indexes}, latent_frames={T_lat}." + ) + + noise = randn_tensor( + (1, C, T_lat, H_lat, W_lat), + generator=generator, + device=self.device, + dtype=self.dtype, + ) + + condition_pixel_frames = _condition_pixel_frame_count( + indexes, self.vae_scale_factor_temporal + ) + condition_video = video_tensor[:, :, :condition_pixel_frames] + if condition_video.shape[2] < condition_pixel_frames: + pad = condition_video[:, :, -1:].repeat( + 1, 1, condition_pixel_frames - condition_video.shape[2], 1, 1 + ) + condition_video = torch.cat([condition_video, pad], dim=2) + + cond_latent = self._encode_video_tensor(condition_video) + expected_prefix = (1, C, max(indexes) + 1, H_lat, W_lat) + if ( + cond_latent.shape[0] != expected_prefix[0] + or cond_latent.shape[1] != expected_prefix[1] + or cond_latent.shape[2] < expected_prefix[2] + or cond_latent.shape[3:] != expected_prefix[3:] + ): + raise ValueError( + "Cosmos3 V2V condition latent shape mismatch: " + f"encoded={tuple(cond_latent.shape)}, expected at least {expected_prefix}." + ) + + condition_mask = torch.zeros(1, 1, T_lat, 1, 1, device=self.device, dtype=self.dtype) + condition_latents = torch.zeros_like(noise) + for index in indexes: + condition_mask[:, :, index, :, :] = 1.0 + condition_latents[:, :, index : index + 1] = cond_latent[:, :, index : index + 1] + latents = condition_mask * condition_latents + (1.0 - condition_mask) * noise + velocity_mask = 1.0 - condition_mask + return latents, velocity_mask, condition_latents + # ========================================================================= # Forward (main generation entry point) # ========================================================================= @@ -733,7 +891,6 @@ def _prepare_action_latents( def forward( self, prompt: Union[str, List[str]], - seed: int, negative_prompt: Optional[str] = None, image: Optional[Union[PIL.Image.Image, torch.Tensor, str]] = None, height: Optional[int] = None, @@ -746,7 +903,7 @@ def forward( frame_rate: Optional[float] = None, use_duration_template: bool = COSMOS3_EXTRA_SPECS["use_duration_template"].default, use_resolution_template: bool = COSMOS3_EXTRA_SPECS["use_resolution_template"].default, - use_system_prompt: bool = COSMOS3_EXTRA_SPECS["use_system_prompt"].default, + use_system_prompt: Optional[bool] = COSMOS3_EXTRA_SPECS["use_system_prompt"].default, use_guardrails: bool = COSMOS3_EXTRA_SPECS["use_guardrails"].default, enable_audio: bool = COSMOS3_EXTRA_SPECS["enable_audio"].default, output_type: str = COSMOS3_EXTRA_SPECS["output_type"].default, @@ -759,6 +916,9 @@ def forward( action_resolution: Optional[int] = None, action_fps: Optional[float] = None, video: Any = None, + condition_frame_indexes_vision: Any = None, + condition_video_keep: Any = None, + flow_shift: Optional[float] = None, ): pipeline_start = time.time() timer = CudaPhaseTimer() @@ -780,6 +940,20 @@ def forward( # latent frame, image-flavored prompt templates, flow_shift=3.0, a CFG # guidance interval, and an image (rather than video) output. is_t2i = str(output_type).lower() == "image" + if not do_action and image is not None and video is not None: + raise ValueError( + "Cosmos3 non-action generation supports text-only, text + image, " + "or text + video input, but not both image and video." + ) + if is_t2i and video is not None: + raise ValueError( + "Cosmos3 video-to-video generation is supported only for video outputs." + ) + is_v2v = video is not None and not is_t2i and not do_action + if use_system_prompt is None: + use_system_prompt = is_v2v + else: + use_system_prompt = bool(use_system_prompt) guidance_interval = None resolved_action_fps: Optional[float] = None if is_t2i: @@ -798,7 +972,9 @@ def forward( if guidance_scale is None: guidance_scale = COSMOS3_T2I_PARAMS["guidance_scale"] guidance_interval = COSMOS3_T2I_PARAMS["guidance_interval"] - self._set_flow_shift(COSMOS3_T2I_PARAMS["flow_shift"]) + self._set_flow_shift( + flow_shift if flow_shift is not None else COSMOS3_T2I_PARAMS["flow_shift"] + ) elif do_action: action_cfg = resolve_domain_action_config( domain_name=domain_name, @@ -835,7 +1011,9 @@ def forward( ) if guidance_scale is None: guidance_scale = COSMOS3_ACTION_PARAMS["guidance_scale"] - self._set_flow_shift(COSMOS3_ACTION_PARAMS["flow_shift"]) + self._set_flow_shift( + flow_shift if flow_shift is not None else COSMOS3_ACTION_PARAMS["flow_shift"] + ) enable_audio = False else: height = height or COSMOS3_720P_PARAMS["height"] @@ -844,9 +1022,19 @@ def forward( num_inference_steps = num_inference_steps or COSMOS3_720P_PARAMS["num_inference_steps"] if guidance_scale is None: guidance_scale = COSMOS3_720P_PARAMS["guidance_scale"] - # Restore the checkpoint flow_shift in case a prior T2I request - # rebuilt the scheduler with shift=3.0. - self._set_flow_shift(getattr(self, "_engine_init_flow_shift", 1.0)) + if is_v2v: + self._set_flow_shift( + flow_shift if flow_shift is not None else 10.0, + use_karras_sigmas=False, + ) + else: + # Restore the checkpoint flow_shift in case a prior T2I/V2V + # request rebuilt the scheduler with a mode-specific shift. + self._set_flow_shift( + flow_shift + if flow_shift is not None + else getattr(self, "_engine_init_flow_shift", 1.0) + ) max_sequence_length = max_sequence_length or COSMOS3_720P_PARAMS["max_sequence_length"] if frame_rate is None: @@ -978,6 +1166,8 @@ def forward( action_frame_offset = 1 resolved_raw_action_dim = raw_action_dim condition_latents = None + image_latent = None + velocity_mask = None if do_action: if action_chunk_size not in {num_frames, num_frames - 1}: @@ -995,7 +1185,7 @@ def forward( if normalized_action_mode == ACTION_MODE_INVERSE_DYNAMICS: inverse_video = video if video is not None else image - video = normalize_action_video_input(inverse_video, max_frames=num_frames) + video = normalize_video_input(inverse_video, max_frames=num_frames) video_tensor = self._preprocess_action_video(video, height, width) latents, velocity_mask, condition_latents = self._prepare_latents_action_video( video_tensor, @@ -1003,7 +1193,6 @@ def forward( num_frames, generator, ) - image_latent = None else: image_tensor = self._preprocess_action_image(action_ref_image, height, width) if image_tensor.ndim == 4: @@ -1018,7 +1207,6 @@ def forward( num_frames, generator, ) - image_latent = None ( action_latents, @@ -1048,10 +1236,41 @@ def forward( latents, velocity_mask, image_latent = self._prepare_latents_i2v( image, height=height, width=width, num_frames=num_frames, generator=generator ) + elif video is not None: + condition_frame_indexes_vision = _normalize_condition_frame_indexes_vision( + condition_frame_indexes_vision + ) + condition_video_keep = _normalize_condition_video_keep(condition_video_keep) + condition_pixel_frames = min( + _condition_pixel_frame_count( + condition_frame_indexes_vision, self.vae_scale_factor_temporal + ), + num_frames, + ) + video = normalize_video_input( + video, + max_frames=None if condition_video_keep == "last" else condition_pixel_frames, + ) + video = ( + video[-condition_pixel_frames:] + if condition_video_keep == "last" + else video[:condition_pixel_frames] + ) + video = self._preprocess_condition_video(video, height, width) + + if self.rank == 0: + logger.info( + f"Cosmos3 V2V conditioning: frames={video.shape[2]}, " + f"latent_indexes={condition_frame_indexes_vision}" + ) + latents, velocity_mask, condition_latents = self._prepare_latents_v2v( + video, + num_frames=num_frames, + generator=generator, + condition_frame_indexes_vision=condition_frame_indexes_vision, + ) else: latents = self._prepare_latents(height, width, num_frames, generator) - velocity_mask = None - image_latent = None # Compute video shape in latent space T_latent = latents.shape[2] @@ -1192,6 +1411,9 @@ def post_step_fn(step_latents, step_extra_stream_latents): extra_streams = {"action": (action_latents, self.action_scheduler)} elif do_audio: extra_streams = {"audio": (audio_latents, self.audio_scheduler)} + should_pin_condition_latents = ( + do_action or condition_latents is not None or image_latent is not None + ) # FUTURE(action+audio): merge both keys; extend forward_fn return dict and post_step_fn. denoise_result = self.denoise( latents=latents, @@ -1203,7 +1425,7 @@ def post_step_fn(step_latents, step_extra_stream_latents): extra_cfg_tensors=extra_cfg_tensors, extra_streams=extra_streams, guidance_interval=guidance_interval, - post_step_fn=post_step_fn if do_action else None, + post_step_fn=post_step_fn if should_pin_condition_latents else None, ) if extra_streams is not None: diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py new file mode 100644 index 000000000000..b18d90284ce2 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared Cosmos3 media helpers.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, List, Optional + +import PIL.Image + +IMAGE_EXTENSIONS = frozenset({".png", ".jpg", ".jpeg", ".webp", ".bmp"}) +VIDEO_EXTENSIONS = frozenset({".mp4", ".avi"}) + + +def pil_to_rgb(value: Any) -> PIL.Image.Image: + if isinstance(value, str): + return PIL.Image.open(value).convert("RGB") + if isinstance(value, PIL.Image.Image): + return value.convert("RGB") + raise TypeError( + f"Cosmos3 preprocessing expected PIL image or image path, got {type(value)!r}." + ) + + +def decode_video_file(path: Path, max_frames: Optional[int] = None) -> List[PIL.Image.Image]: + import torchvision.io as io + + frames, _, _ = io.read_video(str(path), pts_unit="sec") + if frames.numel() == 0: + raise ValueError(f"Cosmos3 video file contains no frames: {path}") + if max_frames is not None: + frames = frames[:max_frames] + return [PIL.Image.fromarray(frames[i].numpy()) for i in range(frames.shape[0])] + + +def normalize_video_input_path(path: Path, max_frames: Optional[int] = None) -> List[Any]: + if not path.exists(): + raise ValueError(f"Cosmos3 video path does not exist: {path}") + if path.is_dir(): + frames = sorted(p for p in path.iterdir() if p.suffix.lower() in IMAGE_EXTENSIONS) + if not frames: + raise ValueError(f"No image frames found in Cosmos3 video directory: {path}") + frame_paths = [str(p) for p in frames] + if max_frames is not None: + frame_paths = frame_paths[:max_frames] + return frame_paths + + suffix = path.suffix.lower() + if suffix in IMAGE_EXTENSIONS: + return [str(path)] + if suffix in VIDEO_EXTENSIONS: + return decode_video_file(path, max_frames=max_frames) + raise ValueError( + "Cosmos3 video path must be a frame directory, an image file " + f"{sorted(IMAGE_EXTENSIONS)}, or a video file " + f"{sorted(VIDEO_EXTENSIONS)}; got {path}" + ) + + +def normalize_video_input(video: Any, max_frames: Optional[int] = None) -> List[Any]: + """Normalize video input to a frame list. + + Accepts a list of PIL images / paths, a single image or video file path, + or a directory of frame images (sorted lexicographically). + """ + if video is None: + return [] + if isinstance(video, list): + if not video: + raise ValueError("Cosmos3 video input must contain at least one frame.") + if max_frames is not None: + return video[:max_frames] + return video + if isinstance(video, (str, Path)): + return normalize_video_input_path(Path(video), max_frames=max_frames) + return [video] diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py index 9a64ebdf4048..e2e8b032bd19 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py @@ -15,10 +15,10 @@ VIDEO_RES_SIZE_INFO, action_reference_image, find_closest_target_size, - normalize_action_video_input, resolve_action_size, ) from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import COSMOS3_EXTRA_SPECS +from tensorrt_llm._torch.visual_gen.models.cosmos3.utils import normalize_video_input pytestmark = pytest.mark.cosmos3 @@ -183,17 +183,17 @@ def test_policy_prefers_image_path_over_video(self, tmp_path): assert ref.getpixel((0, 0)) == (0, 0, 255) -class TestNormalizeActionVideoInput: +class TestNormalizeVideoInput: def test_image_path_returns_singleton_list(self, tmp_path): image_path = tmp_path / "frame.png" PIL.Image.new("RGB", (8, 4), "red").save(image_path) - assert normalize_action_video_input(str(image_path)) == [str(image_path)] + assert normalize_video_input(str(image_path)) == [str(image_path)] def test_frame_directory_returns_sorted_paths(self, tmp_path): (tmp_path / "b.png").write_bytes(b"") (tmp_path / "a.png").write_bytes(b"") (tmp_path / "skip.txt").write_text("x") - assert normalize_action_video_input(str(tmp_path)) == [ + assert normalize_video_input(str(tmp_path)) == [ str(tmp_path / "a.png"), str(tmp_path / "b.png"), ] @@ -202,7 +202,7 @@ def test_unsupported_file_extension_raises(self, tmp_path): bad_path = tmp_path / "clip.mov" bad_path.write_bytes(b"fake") with pytest.raises(ValueError, match="must be a frame directory"): - normalize_action_video_input(str(bad_path)) + normalize_video_input(str(bad_path)) def test_decode_mp4_returns_pil_frames(self, tmp_path, monkeypatch): video_path = tmp_path / "clip.mp4" @@ -227,7 +227,7 @@ def _fake_read_video(path, pts_unit="sec"): "torchvision.io.read_video", _fake_read_video, ) - frames = normalize_action_video_input(str(video_path)) + frames = normalize_video_input(str(video_path)) assert len(frames) == 2 assert all(isinstance(frame, PIL.Image.Image) for frame in frames) assert frames[0].getpixel((0, 0)) == (255, 0, 0) @@ -247,7 +247,7 @@ def _fake_read_video(path, pts_unit="sec"): return tensor, None, {} monkeypatch.setattr("torchvision.io.read_video", _fake_read_video) - frames = normalize_action_video_input( + frames = normalize_video_input( str(video_path), max_frames=2, ) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index 63aceb9cf51e..e7280979126c 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -28,6 +28,7 @@ import json import os from pathlib import Path +from types import SimpleNamespace os.environ["TLLM_DISABLE_MPI"] = "1" os.environ["TRTLLM_DISABLE_COSMOS3_GUARDRAILS"] = "1" @@ -38,12 +39,19 @@ from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import ( COSMOS3_DEFAULT_RESOLUTION_TEMPLATE, + COSMOS3_DEFAULT_SYSTEM_PROMPT, COSMOS3_DURATION_TEMPLATE, COSMOS3_IMAGE_RESOLUTION_TEMPLATE, Cosmos3OmniMoTPipeline, + _condition_pixel_frame_count, + _normalize_condition_frame_indexes_vision, + _normalize_condition_video_keep, ) from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( COSMOS3_ACTION_PARAMS, + COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION, + COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP, + COSMOS3_EXTRA_SPECS, COSMOS3_T2I_PARAMS, ) from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader @@ -236,6 +244,33 @@ def _assert_valid_action(action: torch.Tensor, *, raw_action_dim: int, chunk_siz assert not torch.isinf(af).any() +def _scheduler_use_karras_sigmas(scheduler) -> bool | None: + value = getattr(scheduler.config, "use_karras_sigmas", None) + return None if value is None else bool(value) + + +def _assert_scheduler_config( + pipeline, + *, + flow_shift: float, + use_karras_sigmas: bool | None, +): + assert float(getattr(pipeline.scheduler.config, "flow_shift")) == pytest.approx( + float(flow_shift) + ) + assert float(pipeline._current_flow_shift) == pytest.approx(float(flow_shift)) + assert pipeline._current_scheduler_use_karras_sigmas == use_karras_sigmas + assert _scheduler_use_karras_sigmas(pipeline.scheduler) == use_karras_sigmas + + +def _assert_default_video_scheduler_config(pipeline): + _assert_scheduler_config( + pipeline, + flow_shift=pipeline._engine_init_flow_shift, + use_karras_sigmas=pipeline._base_scheduler_use_karras_sigmas, + ) + + def _make_test_image() -> PIL.Image.Image: image_path = os.environ.get("COSMOS3_TEST_IMAGE") if image_path and os.path.exists(image_path): @@ -243,6 +278,16 @@ def _make_test_image() -> PIL.Image.Image: return PIL.Image.new("RGB", (WIDTH, HEIGHT), color=(64, 128, 192)) +def _make_test_video( + num_frames: int = NUM_FRAMES, + *, + width: int = WIDTH, + height: int = HEIGHT, +) -> list[PIL.Image.Image]: + image = _make_test_image().resize((width, height)) + return [image.copy() for _ in range(num_frames)] + + @pytest.fixture def cosmos3_format_pipeline(): """Minimal pipeline for prompt formatting helpers (no checkpoint).""" @@ -319,6 +364,70 @@ def test_json_array_falls_back_to_append(self, cosmos3_format_pipeline): assert "720x1280" in result +class _CapturingTokenizer: + eos_token_id = 99 + pad_token_id = 0 + + def __init__(self): + self.conversations = [] + + def apply_chat_template( + self, + conversations, + tokenize=True, + add_generation_prompt=True, + return_dict=False, + ): + assert tokenize is True + assert add_generation_prompt is True + assert return_dict is False + self.conversations.append(conversations) + return [1, 2, 3] + + def convert_tokens_to_ids(self, token): + assert token == "<|vision_start|>" + return 98 + + +class TestTokenizePrompt: + def test_system_prompt_included_when_enabled(self, cosmos3_format_pipeline): + tokenizer = _CapturingTokenizer() + cosmos3_format_pipeline.tokenizer = tokenizer + cosmos3_format_pipeline.transformer = SimpleNamespace(device=torch.device("cpu")) + + input_ids, attention_mask = cosmos3_format_pipeline._tokenize_prompt( + "Describe motion.", + max_sequence_length=8, + use_system_prompt=True, + system_prompt="System text.", + ) + + assert tokenizer.conversations == [ + [ + {"role": "system", "content": "System text."}, + {"role": "user", "content": "Describe motion."}, + ] + ] + assert input_ids.tolist() == [[1, 2, 3, 99, 98, 0, 0, 0]] + assert attention_mask.tolist() == [[1, 1, 1, 1, 1, 0, 0, 0]] + + def test_system_prompt_omitted_when_disabled(self, cosmos3_format_pipeline): + tokenizer = _CapturingTokenizer() + cosmos3_format_pipeline.tokenizer = tokenizer + cosmos3_format_pipeline.transformer = SimpleNamespace(device=torch.device("cpu")) + + cosmos3_format_pipeline._tokenize_prompt( + "Describe motion.", + max_sequence_length=8, + use_system_prompt=False, + system_prompt="System text.", + ) + + assert tokenizer.conversations == [ + [{"role": "user", "content": "Describe motion."}] + ] + + class TestFormatPromptWithMetadataJson: def test_injects_metadata_fields(self, cosmos3_format_pipeline): prompt = json.dumps({"prompt": "A foundry pour", "subjects": []}) @@ -412,6 +521,7 @@ def test_t2v_smoke(self, cosmos3_pipeline): result = _run_forward(cosmos3_pipeline, image=None, num_frames=NUM_FRAMES) _assert_valid_video(result.video, num_frames=NUM_FRAMES) assert result.frame_rate == FRAME_RATE + _assert_default_video_scheduler_config(cosmos3_pipeline) @pytest.mark.integration @@ -423,6 +533,161 @@ def test_i2v_smoke(self, cosmos3_pipeline): result = _run_forward(cosmos3_pipeline, image=image, num_frames=NUM_FRAMES) _assert_valid_video(result.video, num_frames=NUM_FRAMES) assert result.frame_rate == FRAME_RATE + _assert_default_video_scheduler_config(cosmos3_pipeline) + + +class TestCosmos3V2VExtraParams: + def test_condition_defaults_are_declared(self): + assert COSMOS3_EXTRA_SPECS["condition_frame_indexes_vision"].default == list( + COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION + ) + assert ( + COSMOS3_EXTRA_SPECS["condition_video_keep"].default + == COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP + ) + + def test_flow_shift_default_is_request_optional(self): + spec = COSMOS3_EXTRA_SPECS["flow_shift"] + assert spec.type == "float" + assert spec.default is None + + def test_video_spec_declares_path_or_list_input(self): + spec = COSMOS3_EXTRA_SPECS["video"] + assert spec.type == "path_or_list" + assert spec.default is None + + +class TestCosmos3V2VConditioningParams: + @pytest.mark.parametrize( + "value,expected", + [ + (None, (0, 1)), + (0, (0,)), + ([0, 2], (0, 2)), + ((1, 3), (1, 3)), + ("0, 2", (0, 2)), + ], + ) + def test_normalize_condition_frame_indexes_vision(self, value, expected): + assert _normalize_condition_frame_indexes_vision(value) == expected + + @pytest.mark.parametrize("value", [[], "", [-1], "0, -1", [0, -2]]) + def test_invalid_condition_frame_indexes_vision_raise(self, value): + with pytest.raises(ValueError): + _normalize_condition_frame_indexes_vision(value) + + @pytest.mark.parametrize( + "indexes,expected", + [ + ((0,), 1), + ((0, 1), 5), + ((2,), 9), + ], + ) + def test_condition_pixel_frame_count(self, indexes, expected): + assert _condition_pixel_frame_count(indexes, temporal_compression=4) == expected + + @pytest.mark.parametrize( + "value,expected", + [ + (None, "first"), + ("first", "first"), + ("FIRST", "first"), + (" last ", "last"), + ], + ) + def test_normalize_condition_video_keep(self, value, expected): + assert _normalize_condition_video_keep(value) == expected + + def test_invalid_condition_video_keep_raises(self): + with pytest.raises(ValueError, match="first or last"): + _normalize_condition_video_keep("middle") + + +@pytest.mark.integration +@pytest.mark.cosmos3_v2v +@pytest.mark.high_cuda_memory +class TestCosmos3V2V: + def test_v2v_smoke(self, cosmos3_pipeline): + video = _make_test_video(NUM_FRAMES) + result = _run_forward( + cosmos3_pipeline, + image=None, + video=video, + num_frames=NUM_FRAMES, + condition_frame_indexes_vision=[0, 1], + condition_video_keep="first", + ) + _assert_valid_video(result.video, num_frames=NUM_FRAMES) + assert result.frame_rate == FRAME_RATE + _assert_scheduler_config( + cosmos3_pipeline, + flow_shift=10.0, + use_karras_sigmas=False, + ) + + def test_v2v_flow_shift_override_request_path(self): + pipeline = Cosmos3OmniMoTPipeline.__new__(Cosmos3OmniMoTPipeline) + pipeline.transformer = SimpleNamespace(device=torch.device("cpu")) + pipeline.action_gen = False + pipeline.audio_gen = False + calls = [] + token_calls = [] + + class StopAfterTokenize(Exception): + pass + + def fake_set_flow_shift(target, *, use_karras_sigmas=None): + calls.append((target, use_karras_sigmas)) + + def fake_tokenize_prompt(text, max_sequence_length, use_system_prompt, system_prompt=None): + token_calls.append((text, max_sequence_length, use_system_prompt, system_prompt)) + raise StopAfterTokenize + + pipeline._set_flow_shift = fake_set_flow_shift + pipeline._tokenize_prompt = fake_tokenize_prompt + + with pytest.raises(StopAfterTokenize): + pipeline.forward( + prompt="continue", + video=_make_test_video(5, width=16, height=16), + height=16, + width=16, + num_frames=5, + num_inference_steps=1, + guidance_scale=1.0, + seed=1, + max_sequence_length=8, + frame_rate=8.0, + use_duration_template=False, + use_resolution_template=False, + use_system_prompt=None, + use_guardrails=False, + flow_shift=7.0, + ) + + assert calls == [(7.0, False)] + assert token_calls[0][2] is True + assert token_calls[0][3] == COSMOS3_DEFAULT_SYSTEM_PROMPT + + def test_image_and_video_rejected(self, cosmos3_pipeline): + with pytest.raises(ValueError, match="not both image and video"): + _run_forward( + cosmos3_pipeline, + image=_make_test_image(), + video=_make_test_video(5), + ) + + def test_t2i_and_video_rejected(self, cosmos3_pipeline): + with pytest.raises(ValueError, match="supported only for video outputs"): + _run_forward( + cosmos3_pipeline, + image=None, + video=_make_test_video(5), + output_type="image", + height=T2I_HEIGHT, + width=T2I_WIDTH, + ) @pytest.mark.integration @@ -440,6 +705,11 @@ def test_t2i_smoke(self, cosmos3_pipeline): ) assert result.video is None _assert_valid_image(result.image, height=T2I_HEIGHT, width=T2I_WIDTH) + _assert_scheduler_config( + cosmos3_pipeline, + flow_shift=COSMOS3_T2I_PARAMS["flow_shift"], + use_karras_sigmas=cosmos3_pipeline._base_scheduler_use_karras_sigmas, + ) @pytest.mark.integration @@ -492,6 +762,11 @@ def test_policy_smoke(self, cosmos3_pipeline): ) assert result.action_mode == "policy" assert result.domain_id == 7 + _assert_scheduler_config( + cosmos3_pipeline, + flow_shift=COSMOS3_ACTION_PARAMS["flow_shift"], + use_karras_sigmas=cosmos3_pipeline._base_scheduler_use_karras_sigmas, + ) def test_forward_dynamics_smoke(self, cosmos3_pipeline): _require_action_pipeline(cosmos3_pipeline) @@ -520,6 +795,11 @@ def test_forward_dynamics_smoke(self, cosmos3_pipeline): raw_action_dim=self.RAW_ACTION_DIM, chunk_size=self.ACTION_CHUNK, ) + _assert_scheduler_config( + cosmos3_pipeline, + flow_shift=COSMOS3_ACTION_PARAMS["flow_shift"], + use_karras_sigmas=cosmos3_pipeline._base_scheduler_use_karras_sigmas, + ) def test_inverse_dynamics_smoke(self, cosmos3_pipeline): _require_action_pipeline(cosmos3_pipeline) @@ -549,6 +829,11 @@ def test_inverse_dynamics_smoke(self, cosmos3_pipeline): raw_action_dim=self.RAW_ACTION_DIM, chunk_size=NUM_FRAMES, ) + _assert_scheduler_config( + cosmos3_pipeline, + flow_shift=COSMOS3_ACTION_PARAMS["flow_shift"], + use_karras_sigmas=cosmos3_pipeline._base_scheduler_use_karras_sigmas, + ) def test_action_and_audio_rejected(self, cosmos3_pipeline): _require_action_pipeline(cosmos3_pipeline) @@ -584,8 +869,9 @@ class TestCosmos3PromptTemplates: (True, True, True), (False, False, False), (False, False, True), + (False, False, None), ], - ids=["all-on", "all-off", "system-prompt-only"], + ids=["all-on", "all-off", "system-prompt-only", "system-prompt-default"], ) def test_template_variants( self, From 91285132108a06f20cd18d4dd6bea8ef870c7ac7 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 8 Jul 2026 14:45:56 -0700 Subject: [PATCH 05/64] Add content-based classification for video input_reference Previously `input_reference` was always treated as an image and stored as `_reference.png`, routing unconditionally to `params.image`. Introduce `_reference_is_image` and `_reference_is_video` helpers that probe decoded content (PIL and PyAV respectively) rather than relying on file extension or content-type. Classification order is image-first: FFmpeg demuxes still images as single-frame video streams, so the PIL probe must gate the PyAV probe. On a positive video classification the reference is stored as `_reference.mp4` and routed to `params.extra_params["video"]` instead of `params.image`. Undecodable content removes the temporary `.part` file and raises `ValueError`. The field docstring is updated to document the content-based dispatch contract. Signed-off-by: Igor Shovkun --- tensorrt_llm/serve/openai_protocol.py | 8 +- tensorrt_llm/serve/visual_gen_utils.py | 64 ++++++++++++++- .../visual_gen/test_visual_gen_utils.py | 79 +++++++++++++++++++ 3 files changed, 146 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index 4f497787e143..529821302eaf 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -1686,7 +1686,13 @@ class VideoGenerationRequest(OpenAIBaseModel): description="Random seed for reproducibility.") input_reference: Optional[Union[str, UploadFile]] = Field( default=None, - description="Optional image reference that guides generation.", + description=( + "Optional image or video reference that guides generation. " + "Content is classified by decoding, not by extension or " + "content-type: images (anything PIL reads) condition " + "image-to-video; videos (anything PyAV reads) condition " + "video-to-video on models that support it. JSON requests " + "carry base64 bytes; multipart requests upload the file."), ) # Resolution diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index 3094bf66cb04..0659d513f140 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -4,6 +4,8 @@ import shutil from typing import Any, Dict, List, Optional +from PIL import Image, UnidentifiedImageError + from tensorrt_llm.logger import logger from tensorrt_llm.serve.openai_protocol import ImageGenerationRequest, VideoGenerationRequest from tensorrt_llm.visual_gen import VisualGen, VisualGenParams @@ -84,6 +86,40 @@ def _merge_extra_params( params.extra_params = None +def _reference_is_image(path: str) -> bool: + """True when ``path`` holds image content (anything PIL can open). + + Capability-based: the supported set is whatever the decoder accepts, + with no enumerated format table. The probe parses only the header — + pixel decode stays in the worker. + """ + try: + with Image.open(path): + return True + except UnidentifiedImageError: + return False + + +def _reference_is_video(path: str) -> bool: + """True when ``path`` holds video content (a PyAV-openable video stream). + + Total predicate: False for images, audio, and undecodable content + alike. Images must be excluded explicitly because FFmpeg demuxes a + still image as a valid single-frame video stream, so an av probe + alone would claim every PNG/JPEG. A missing ``av`` package raises + ``ImportError`` — that is a deployment problem, not a content + verdict. + """ + if _reference_is_image(path): + return False + import av + try: + with av.open(path) as container: + return bool(container.streams.video) + except av.FFmpegError: + return False + + def parse_visual_gen_params( request: ImageGenerationRequest | VideoGenerationRequest, id: str, @@ -156,14 +192,34 @@ def parse_visual_gen_params( if request.input_reference is not None: if media_storage_path is None: raise ValueError("media_storage_path is required when input_reference is provided") - ref_path = os.path.join(media_storage_path, f"{id}_reference.png") + tmp_path = os.path.join(media_storage_path, f"{id}_reference.part") if isinstance(request.input_reference, str): - with open(ref_path, "wb") as f: + with open(tmp_path, "wb") as f: f.write(base64.b64decode(request.input_reference)) else: - with open(ref_path, "wb") as f: + with open(tmp_path, "wb") as f: shutil.copyfileobj(request.input_reference.file, f) - params.image = ref_path + # image, video, or reject. + if _reference_is_image(tmp_path): + is_video = False + elif _reference_is_video(tmp_path): + is_video = True + else: + os.remove(tmp_path) + raise ValueError( + "input_reference content is neither a decodable image " + "nor a decodable video." + ) + ref_path = os.path.join( + media_storage_path, f"{id}_reference{'.mp4' if is_video else '.png'}" + ) + os.replace(tmp_path, ref_path) + if is_video: + if params.extra_params is None: + params.extra_params = {} + params.extra_params["video"] = ref_path + else: + params.image = ref_path _warn_if_set_with_no_semantic(request, getattr(generator, "model", None)) _merge_extra_params(params, request.extra_params, generator.extra_param_specs) diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index 4b4836392a1f..e60ce22df50c 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -15,7 +15,9 @@ from io import BytesIO from typing import Any, Dict, Optional +import numpy as np import pytest +from fastapi import UploadFile from PIL import Image from tensorrt_llm.serve.openai_protocol import ImageGenerationRequest, VideoGenerationRequest @@ -289,6 +291,83 @@ def test_missing_media_storage_path_raises(self): with pytest.raises(ValueError, match="media_storage_path"): parse_visual_gen_params(request, "vid-2", generator, media_storage_path=None) + @staticmethod + def _mp4_bytes() -> bytes: + """Encode a 2-frame 16x16 mpeg4-in-mp4 clip in memory. + + ``mpeg4`` is a built-in FFmpeg encoder, so this works on any + PyAV wheel (h264 may be absent from LGPL builds). + """ + av = pytest.importorskip("av") + buf = BytesIO() + with av.open(buf, "w", format="mp4") as container: + stream = container.add_stream("mpeg4", rate=4) + stream.width = 16 + stream.height = 16 + stream.pix_fmt = "yuv420p" + for _ in range(2): + frame = av.VideoFrame.from_ndarray( + np.zeros((16, 16, 3), dtype=np.uint8), format="rgb24" + ) + container.mux(stream.encode(frame)) + container.mux(stream.encode()) + return buf.getvalue() + + def test_multipart_video_reference_routes_to_extra_params(self, tmp_path): + generator = _StubVisualGen() + upload = UploadFile(file=BytesIO(self._mp4_bytes()), filename="clip.mp4") + request = VideoGenerationRequest(prompt="x", input_reference=upload) + params = parse_visual_gen_params( + request, "vid-3", generator, media_storage_path=str(tmp_path) + ) + # Video content routes to extra_params["video"], not params.image, + # and the written suffix drives the worker's decode dispatch. + assert params.image is None + assert params.extra_params is not None + assert str(params.extra_params["video"]).endswith("vid-3_reference.mp4") + assert (tmp_path / "vid-3_reference.mp4").exists() + + def test_base64_video_reference_routes_to_extra_params(self, tmp_path): + # Classification is content-based, so the JSON/base64 path can + # carry video even though it has no content-type or filename. + generator = _StubVisualGen() + b64 = base64.b64encode(self._mp4_bytes()).decode() + request = VideoGenerationRequest(prompt="x", input_reference=b64) + params = parse_visual_gen_params( + request, "vid-4", generator, media_storage_path=str(tmp_path) + ) + assert params.image is None + assert str(params.extra_params["video"]).endswith("vid-4_reference.mp4") + + def test_multipart_image_reference_routes_to_image(self, tmp_path): + # JPEG upload: content sniffing classifies it as an image even + # though the stored name is the cosmetic ``.png`` (PIL identifies + # by content, not suffix). + generator = _StubVisualGen() + img = Image.new("RGB", (4, 4), (10, 20, 30)) + buf = BytesIO() + img.save(buf, format="JPEG") + buf.seek(0) + upload = UploadFile(file=buf, filename="ref.jpg") + request = VideoGenerationRequest(prompt="x", input_reference=upload) + params = parse_visual_gen_params( + request, "vid-5", generator, media_storage_path=str(tmp_path) + ) + assert params.extra_params is None + assert str(params.image).endswith("vid-5_reference.png") + + def test_undecodable_reference_raises_and_cleans_up(self, tmp_path): + pytest.importorskip("av") + generator = _StubVisualGen() + b64 = base64.b64encode(b"neither an image nor a video").decode() + request = VideoGenerationRequest(prompt="x", input_reference=b64) + with pytest.raises(ValueError, match="neither a decodable image"): + parse_visual_gen_params( + request, "vid-6", generator, media_storage_path=str(tmp_path) + ) + # The temporary materialization is removed on rejection. + assert list(tmp_path.iterdir()) == [] + # ============================================================================= # _merge_extra_params — the merge truth table From 05ed1361a7f9eccd2a684a73c7db3a93ffe4f824 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 8 Jul 2026 14:45:58 -0700 Subject: [PATCH 06/64] Add Cosmos3 V2V docs, tests, and serve endpoint coverage - Document V2V mode in README and cosmos3.py docstring alongside existing T2V/T2I/I2V/T2AV modes - Update `--video_path` help text to reflect V2V as the primary use case - Add `condition_video_keep="last"` smoke test and audio+V2V combined smoke test to the pipeline unit tests - Add serve endpoint tests: multipart video reference routes to `extra_params["video"]` (V2V), undecodable reference returns HTTP 400 - Document Cosmos3 `extra_params` knobs and V2V multipart curl example in the serve README; clarify `input_reference` classification behavior Signed-off-by: Igor Shovkun --- examples/visual_gen/models/cosmos3/README.md | 12 +++- examples/visual_gen/models/cosmos3/cosmos3.py | 14 ++++- examples/visual_gen/serve/README.md | 15 ++++- .../visual_gen/test_cosmos3_pipeline.py | 51 ++++++++++++++++ .../visual_gen/test_trtllm_serve_endpoints.py | 58 +++++++++++++++++++ 5 files changed, 147 insertions(+), 3 deletions(-) diff --git a/examples/visual_gen/models/cosmos3/README.md b/examples/visual_gen/models/cosmos3/README.md index 69be21fe4880..f3ddcf5910f3 100644 --- a/examples/visual_gen/models/cosmos3/README.md +++ b/examples/visual_gen/models/cosmos3/README.md @@ -1,10 +1,11 @@ # Cosmos3 Text(+Image)-to-Video(+Audio) generation -Cosmos3 supports four generation modes from a single checkpoint: +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 frame directory, `.mp4`/`.avi` file, or single image; passing it without `--action_mode` selects V2V). Only the first (or last, per `condition_video_keep`) `max(condition_frame_indexes_vision) * 4 + 1` input frames condition the output (5 by default); `.mp4`/`.avi` decode requires `pip install av`. - **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 @@ -59,6 +60,15 @@ python cosmos3.py --model nvidia/Cosmos3-Nano \ --image_path https://example.com/frame.jpg \ --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml +# 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. +python cosmos3.py --model /path/to/Cosmos3-Nano \ + --prompt_file prompts/v2v.json \ + --video_path /path/to/Cosmos3-Nano/assets/example_i2v_output.mp4 \ + --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml + # T2AV: text-to-video with synchronized audio python cosmos3.py --model nvidia/Cosmos3-Nano \ --prompt_file prompts/t2av.json \ diff --git a/examples/visual_gen/models/cosmos3/cosmos3.py b/examples/visual_gen/models/cosmos3/cosmos3.py index 01f804fc83a6..6795c1495eeb 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3.py +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -23,6 +23,11 @@ - **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 the + first (or last, per ``condition_video_keep``) frames of a reference video + via ``--video_path`` (a local frame directory, ``.mp4``/``.avi`` file, or + single image; ``.mp4``/``.avi`` decode requires the ``av`` package). + Passing ``--video_path`` without ``--action_mode`` selects V2V. - **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). @@ -75,6 +80,12 @@ --image_path https://example.com/frame.jpg \ --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml + # V2V: video-conditioned video (continues the first frames of --video_path) + python cosmos3.py --model nvidia/Cosmos3-Nano \ + --prompt_file prompts/v2v.json \ + --video_path /path/to/reference.mp4 \ + --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml + # T2AV: text-to-video with synchronized audio python cosmos3.py --model nvidia/Cosmos3-Nano \ --prompt_file prompts/t2av.json \ @@ -404,7 +415,8 @@ def main(): "--video_path", type=str, default=None, - help="Frame directory, .mp4/.avi video, or image path for inverse_dynamics", + help="Reference video for V2V (or inverse_dynamics with --action_mode): " + "a local frame directory, .mp4/.avi file, or image path", ) parser.add_argument( "--action_resolution", diff --git a/examples/visual_gen/serve/README.md b/examples/visual_gen/serve/README.md index 3c3d9d979831..63fae4b5d93e 100644 --- a/examples/visual_gen/serve/README.md +++ b/examples/visual_gen/serve/README.md @@ -286,7 +286,7 @@ You can customize these by: - `frame_rate` (canonical) or `fps` (alias): frames per second - `num_frames`: when set, wins over the `seconds * frame_rate` derivation - `seed`, `num_inference_steps`, `guidance_scale`, `max_sequence_length`, `negative_prompt`: per-request denoise controls -- `input_reference`: Reference image (TI2V mode); accepted as base64-encoded string in JSON or as a file in multipart form-data +- `input_reference`: Reference image (I2V/TI2V) or video (V2V), classified by decoding the content — filename and content type are ignored; accepted as base64-encoded string in JSON or as a file in multipart form-data. Undecodable content returns HTTP 400. Video decode requires the `av` (PyAV) package on the server. - `extra_params`: model-specific overflow (see below) - `response_format`: `"b64_json"` or `"url"` - `format`: Generation content encoding. Video encoders: `"mp4"`, `"avi"`, `"auto"`. Tensor formats: `"safetensors"`, `"pt"` (carries video + audio + scalar metadata in one payload for LTX-2). @@ -315,6 +315,7 @@ 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_frame_indexes_vision`, `condition_video_keep` (V2V conditioning), `flow_shift`, `use_system_prompt`, ... > **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`). @@ -357,6 +358,18 @@ curl -X POST "http://localhost:8000/v1/videos" \ -F "guidance_scale=5.0" ``` +### Video-to-Video (Multipart with File Upload, Cosmos3) +```bash +# The reference is classified by content: image -> I2V, video -> V2V. +# V2V conditioning knobs ride in extra_params (values below are the defaults). +curl -X POST "http://localhost:8000/v1/videos" \ + -F "prompt=Continue the same scene with smooth natural motion and consistent subjects." \ + -F "input_reference=@./media/reference.mp4" \ + -F "num_frames=189" \ + -F "fps=24" \ + -F 'extra_params={"condition_frame_indexes_vision": [0, 1], "condition_video_keep": "first"}' +``` + ### Check Video Status ```bash curl -X GET "http://localhost:8000/v1/videos/{video_id}" diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index e7280979126c..9964fa57f2b2 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -626,6 +626,38 @@ def test_v2v_smoke(self, cosmos3_pipeline): use_karras_sigmas=False, ) + def test_v2v_keep_last_smoke(self, cosmos3_pipeline): + """condition_video_keep="last" pins the tail of the input, not the head. + + The input is longer than the conditioning window and color-coded + (dark head, bright tail), so this exercises the full-decode + + tail-slice path and asserts behavior: frame 0 of the output is a + pinned VAE round-trip of the bright tail frames. + """ + dark = PIL.Image.new("RGB", (WIDTH, HEIGHT), (40, 40, 40)) + bright = PIL.Image.new("RGB", (WIDTH, HEIGHT), (230, 230, 230)) + # 5 = max(condition_frame_indexes_vision) * 4 + 1 conditioning frames. + video = [dark.copy() for _ in range(NUM_FRAMES)] + [bright.copy() for _ in range(5)] + result = _run_forward( + cosmos3_pipeline, + image=None, + video=video, + num_frames=NUM_FRAMES, + condition_frame_indexes_vision=[0, 1], + condition_video_keep="last", + ) + _assert_valid_video(result.video, num_frames=NUM_FRAMES) + first_frame_mean = result.video[0, 0].float().mean().item() + assert first_frame_mean > 135, ( + f"keep='last' must condition on the bright tail frames; frame-0 mean " + f"{first_frame_mean:.1f} matches the dark head instead" + ) + _assert_scheduler_config( + cosmos3_pipeline, + flow_shift=10.0, + use_karras_sigmas=False, + ) + def test_v2v_flow_shift_override_request_path(self): pipeline = Cosmos3OmniMoTPipeline.__new__(Cosmos3OmniMoTPipeline) pipeline.transformer = SimpleNamespace(device=torch.device("cpu")) @@ -723,6 +755,25 @@ def test_audio_smoke(self, cosmos3_pipeline): assert result.frame_rate == FRAME_RATE _assert_valid_audio(result.audio, result.audio_sample_rate) + def test_v2v_audio_smoke(self, cosmos3_pipeline): + """Audio + V2V combined — allowed in both implementations (no guard), + previously untested in either.""" + _require_audio_pipeline(cosmos3_pipeline) + result = _run_forward( + cosmos3_pipeline, + enable_audio=True, + video=_make_test_video(NUM_FRAMES), + condition_frame_indexes_vision=[0, 1], + condition_video_keep="first", + ) + _assert_valid_video(result.video, num_frames=NUM_FRAMES) + _assert_valid_audio(result.audio, result.audio_sample_rate) + _assert_scheduler_config( + cosmos3_pipeline, + flow_shift=10.0, + use_karras_sigmas=False, + ) + @pytest.mark.integration @pytest.mark.cosmos3_action diff --git a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py index 105f9ad90307..97e8e15b5c30 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py @@ -842,6 +842,64 @@ def test_sync_video_generation_multipart_with_reference(self, video_client, tmp_ assert params.image.endswith("_reference.png") assert os.path.exists(params.image) + def test_sync_video_generation_multipart_with_video_reference(self, video_client, tmp_path): + """A video ``input_reference`` routes to ``extra_params["video"]`` (V2V). + + The reference is classified by decoding its content, so the clip is + synthesized in-test with PyAV — no video asset ships with the repo. + """ + av = pytest.importorskip("av") + np = pytest.importorskip("numpy") + ref_path = tmp_path / "ref.mp4" + with av.open(str(ref_path), "w") as container: + # mpeg4 is a built-in FFmpeg encoder (h264 may be absent from + # LGPL PyAV wheels). + stream = container.add_stream("mpeg4", rate=4) + stream.width = 16 + stream.height = 16 + stream.pix_fmt = "yuv420p" + for _ in range(2): + frame = av.VideoFrame.from_ndarray( + np.zeros((16, 16, 3), dtype=np.uint8), format="rgb24" + ) + container.mux(stream.encode(frame)) + container.mux(stream.encode()) + + with open(ref_path, "rb") as f: + resp = video_client.post( + "/v1/videos/generations", + data={ + "prompt": "Continue the same scene", + "size": "64x64", + "seconds": "1.0", + "fps": "8", + }, + files={"input_reference": ("ref.mp4", f, "video/mp4")}, + ) + assert resp.status_code == 200 + assert len(resp.content) > 0 + + # Video content must NOT land on params.image; it rides + # extra_params["video"] (the same pipeline entry the offline + # example's --video_path uses), stored with a .mp4 suffix. + params = video_client.mock_gen.last_params + assert params.image is None + assert isinstance(params.extra_params, dict) + video_ref = params.extra_params["video"] + assert video_ref.endswith("_reference.mp4") + assert os.path.exists(video_ref) + + def test_sync_video_generation_undecodable_reference_400(self, video_client): + """Content neither PIL nor PyAV can decode is rejected at the boundary.""" + pytest.importorskip("av") + resp = video_client.post( + "/v1/videos/generations", + data={"prompt": "x"}, + files={"input_reference": ("doc.txt", BytesIO(b"not media"), "text/plain")}, + ) + assert resp.status_code == 400 + assert "neither a decodable image" in resp.text + def test_sync_video_failure(self, failing_client): resp = failing_client.post( "/v1/videos/generations", From 433feb9b93c2b4bb5031c8e0d2f321e8dbfba395 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 8 Jul 2026 14:46:00 -0700 Subject: [PATCH 07/64] Add Cosmos3 V2V prompt file and document media I/O deps - Add missing `prompts/v2v.json` example for video-to-video mode - Add `av` (PyAV) to `requirements.txt` for video decode support - Add "Media I/O dependencies" section to README covering ffmpeg (mp4 output) and av (V2V reference video decode) - Update V2V bullet to cross-reference the new section instead of inline pip-install note Signed-off-by: Igor Shovkun --- examples/visual_gen/models/cosmos3/README.md | 7 ++++++- examples/visual_gen/models/cosmos3/prompts/v2v.json | 4 ++++ requirements.txt | 3 +++ 3 files changed, 13 insertions(+), 1 deletion(-) create mode 100644 examples/visual_gen/models/cosmos3/prompts/v2v.json diff --git a/examples/visual_gen/models/cosmos3/README.md b/examples/visual_gen/models/cosmos3/README.md index f3ddcf5910f3..77784e8b164f 100644 --- a/examples/visual_gen/models/cosmos3/README.md +++ b/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 frame directory, `.mp4`/`.avi` file, or single image; passing it without `--action_mode` selects V2V). Only the first (or last, per `condition_video_keep`) `max(condition_frame_indexes_vision) * 4 + 1` input frames condition the output (5 by default); `.mp4`/`.avi` decode requires `pip install av`. +- **V2V** — video-conditioned video (`prompts/v2v.json`). Condition on a reference video via `--video_path` (a local frame directory, `.mp4`/`.avi` file, or single image; passing it without `--action_mode` selects V2V). Only the first (or last, per `condition_video_keep`) `max(condition_frame_indexes_vision) * 4 + 1` input frames condition the output (5 by default); `.mp4`/`.avi` decode uses the `av` package (see [Media I/O dependencies](#media-io-dependencies)). - **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 @@ -31,6 +31,11 @@ To run without guardrails (you are responsible for safe deployment): export TRTLLM_DISABLE_COSMOS3_GUARDRAILS=1 ``` +## Media I/O dependencies + +- 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, inverse dynamics) uses the `av` (PyAV) package, installed with TensorRT-LLM's requirements. + ## Deployment configs See `examples/visual_gen/configs/`: diff --git a/examples/visual_gen/models/cosmos3/prompts/v2v.json b/examples/visual_gen/models/cosmos3/prompts/v2v.json new file mode 100644 index 000000000000..9c42fdcf67fa --- /dev/null +++ b/examples/visual_gen/models/cosmos3/prompts/v2v.json @@ -0,0 +1,4 @@ +{ + "model_mode": "video2video", + "prompt": "Continue the same scene with smooth natural motion and consistent subjects." +} diff --git a/requirements.txt b/requirements.txt index d735182a8fe8..a7f38440bfff 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,9 @@ --extra-index-url https://download.pytorch.org/whl/cu130 -c constraints.txt accelerate>=1.7.0 +# Video reference decode (torchvision.io.read_video backend) and serve-side +# image-vs-video content classification for visual_gen (Cosmos3 V2V). +av build colored cuda-python>=13 From ae5b03d0851ec2fbd1cd2ae9b81cd31b3f857d3c Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 8 Jul 2026 16:40:20 -0700 Subject: [PATCH 08/64] Resolve merge conflict and apply formatting fixes - Remove leftover conflict markers from transformer_cosmos3.py - Consolidate multi-line f-strings and error messages onto single lines - Reorder imports alphabetically in test_cosmos3_pipeline.py - Add blank line after `import av` per formatting conventions Signed-off-by: Igor Shovkun --- .../models/cosmos3/pipeline_cosmos3.py | 3 +-- .../models/cosmos3/transformer_cosmos3.py | 5 +---- .../_torch/visual_gen/models/cosmos3/utils.py | 4 +--- tensorrt_llm/serve/visual_gen_utils.py | 4 ++-- .../_torch/visual_gen/test_cosmos3_pipeline.py | 18 ++++++++---------- .../_torch/visual_gen/test_visual_gen_utils.py | 4 +--- 6 files changed, 14 insertions(+), 24 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index 810b15fa8bd1..e7192b12db05 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -92,8 +92,7 @@ def _normalize_condition_frame_indexes_vision( raise ValueError("Cosmos3 condition_frame_indexes_vision must not be empty.") if any(index < 0 for index in normalized): raise ValueError( - "Cosmos3 condition_frame_indexes_vision must be non-negative, " - f"got {normalized}." + f"Cosmos3 condition_frame_indexes_vision must be non-negative, got {normalized}." ) return normalized diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py index 0840ee763bab..7740efc1a24b 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -1165,13 +1165,11 @@ def forward( provided; otherwise None. action is set when action_latents is provided. """ del kwargs # Kept for diffusers API compatibility. -<<<<<<< HEAD if timestep is None: raise ValueError("Cosmos3VFMTransformer.forward requires normalized timestep.") if raw_timestep is None: raise ValueError("Cosmos3VFMTransformer.forward requires raw_timestep.") -======= - + if action_latents is not None and audio_latents is not None: raise ValueError( "Cosmos3 transformer does not support joint action and audio generation." @@ -1181,7 +1179,6 @@ def forward( "Cosmos3 action generation was requested, but this transformer " "was initialized without action modules." ) ->>>>>>> 270db1b2f7 (cosmos3 action init) T, H, W = video_shape Hp, Wp, _, _ = self._pad_to_patch_size(H, W) max_real_len = text_mask.sum(dim=1).max().item() diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py index b18d90284ce2..75dcec8a3ebe 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py @@ -19,9 +19,7 @@ def pil_to_rgb(value: Any) -> PIL.Image.Image: return PIL.Image.open(value).convert("RGB") if isinstance(value, PIL.Image.Image): return value.convert("RGB") - raise TypeError( - f"Cosmos3 preprocessing expected PIL image or image path, got {type(value)!r}." - ) + raise TypeError(f"Cosmos3 preprocessing expected PIL image or image path, got {type(value)!r}.") def decode_video_file(path: Path, max_frames: Optional[int] = None) -> List[PIL.Image.Image]: diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index 0659d513f140..589b6b9d27c1 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -113,6 +113,7 @@ def _reference_is_video(path: str) -> bool: if _reference_is_image(path): return False import av + try: with av.open(path) as container: return bool(container.streams.video) @@ -207,8 +208,7 @@ def parse_visual_gen_params( else: os.remove(tmp_path) raise ValueError( - "input_reference content is neither a decodable image " - "nor a decodable video." + "input_reference content is neither a decodable image nor a decodable video." ) ref_path = os.path.join( media_storage_path, f"{id}_reference{'.mp4' if is_video else '.png'}" diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index 9964fa57f2b2..d992b275ef6a 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -37,6 +37,13 @@ import pytest import torch +from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( + COSMOS3_ACTION_PARAMS, + COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION, + COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP, + COSMOS3_EXTRA_SPECS, + COSMOS3_T2I_PARAMS, +) from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import ( COSMOS3_DEFAULT_RESOLUTION_TEMPLATE, COSMOS3_DEFAULT_SYSTEM_PROMPT, @@ -47,13 +54,6 @@ _normalize_condition_frame_indexes_vision, _normalize_condition_video_keep, ) -from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( - COSMOS3_ACTION_PARAMS, - COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION, - COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP, - COSMOS3_EXTRA_SPECS, - COSMOS3_T2I_PARAMS, -) from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader from tensorrt_llm.visual_gen.args import TorchCompileConfig, VisualGenArgs @@ -423,9 +423,7 @@ def test_system_prompt_omitted_when_disabled(self, cosmos3_format_pipeline): system_prompt="System text.", ) - assert tokenizer.conversations == [ - [{"role": "user", "content": "Describe motion."}] - ] + assert tokenizer.conversations == [[{"role": "user", "content": "Describe motion."}]] class TestFormatPromptWithMetadataJson: diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index e60ce22df50c..952854161935 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -362,9 +362,7 @@ def test_undecodable_reference_raises_and_cleans_up(self, tmp_path): b64 = base64.b64encode(b"neither an image nor a video").decode() request = VideoGenerationRequest(prompt="x", input_reference=b64) with pytest.raises(ValueError, match="neither a decodable image"): - parse_visual_gen_params( - request, "vid-6", generator, media_storage_path=str(tmp_path) - ) + parse_visual_gen_params(request, "vid-6", generator, media_storage_path=str(tmp_path)) # The temporary materialization is removed on rejection. assert list(tmp_path.iterdir()) == [] From 14d45558a4846020c94685a2a53decce342ac49c Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 8 Jul 2026 22:27:44 -0700 Subject: [PATCH 09/64] Add URI support for reference image in Cosmos3 action Local path handling remains unchanged; only HTTP(S), data:, and file: URIs are routed through `load_image` to match the existing I2V image branch behavior. Signed-off-by: Igor Shovkun --- .../_torch/visual_gen/models/cosmos3/action.py | 7 +++++++ .../_torch/visual_gen/test_cosmos3_action.py | 17 +++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py index 2d4a1c448988..0f65df684718 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py @@ -7,12 +7,15 @@ from pathlib import Path from typing import Any, Optional +from urllib.parse import urlparse import numpy as np import PIL.Image import torch from diffusers.utils.torch_utils import randn_tensor +from tensorrt_llm.inputs.utils import load_image + from .utils import IMAGE_EXTENSIONS, normalize_video_input, pil_to_rgb ACTION_MODE_POLICY = "policy" @@ -284,6 +287,10 @@ def action_reference_image( if isinstance(source, PIL.Image.Image): return source.convert("RGB") if isinstance(source, str): + if urlparse(source).scheme in ("http", "https", "data", "file"): + # Same URI-aware loader as the I2V image branch; plain local + # paths keep the richer path handling below (frame dirs, videos). + return load_image(source, format="pil").convert("RGB") path = Path(source) if path.is_file() and path.suffix.lower() in IMAGE_EXTENSIONS: return PIL.Image.open(source).convert("RGB") diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py index e2e8b032bd19..ef9856e3056a 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py @@ -182,6 +182,23 @@ def test_policy_prefers_image_path_over_video(self, tmp_path): ) assert ref.getpixel((0, 0)) == (0, 0, 255) + def test_policy_accepts_data_uri_image(self): + # URI references must go through the same loader as the I2V image + # branch instead of being treated as local filesystem paths. + import base64 + from io import BytesIO + + buf = BytesIO() + PIL.Image.new("RGB", (3, 3), "green").save(buf, format="PNG") + data_uri = "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode() + ref = action_reference_image( + action_mode="policy", + image=data_uri, + video=None, + ) + assert ref.size == (3, 3) + assert ref.getpixel((0, 0)) == (0, 128, 0) + class TestNormalizeVideoInput: def test_image_path_returns_singleton_list(self, tmp_path): From 4e7c4039d6a41a3077056ceee0c66bf3697dfab8 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 9 Jul 2026 07:59:07 -0700 Subject: [PATCH 10/64] Clean up temp file on input_reference parse failure Wrap the input_reference materialization block in a try/except that removes the `.part` temp file before re-raising. Previously, any failure after the file was opened (bad base64, broken upload stream, unrecognized media type) would leave the partial file on disk. Also add explicit validation for malformed base64 input, raising a descriptive ValueError instead of propagating the raw decode error. Signed-off-by: Igor Shovkun --- tensorrt_llm/serve/visual_gen_utils.py | 49 ++++++++++++------- .../visual_gen/test_visual_gen_utils.py | 24 +++++++++ 2 files changed, 54 insertions(+), 19 deletions(-) diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index 589b6b9d27c1..29bc4e7c2969 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -194,26 +194,37 @@ def parse_visual_gen_params( if media_storage_path is None: raise ValueError("media_storage_path is required when input_reference is provided") tmp_path = os.path.join(media_storage_path, f"{id}_reference.part") - if isinstance(request.input_reference, str): - with open(tmp_path, "wb") as f: - f.write(base64.b64decode(request.input_reference)) - else: - with open(tmp_path, "wb") as f: - shutil.copyfileobj(request.input_reference.file, f) - # image, video, or reject. - if _reference_is_image(tmp_path): - is_video = False - elif _reference_is_video(tmp_path): - is_video = True - else: - os.remove(tmp_path) - raise ValueError( - "input_reference content is neither a decodable image nor a decodable video." + try: + if isinstance(request.input_reference, str): + try: + payload = base64.b64decode(request.input_reference) + except ValueError as exc: + raise ValueError("input_reference is not valid base64 data.") from exc + with open(tmp_path, "wb") as f: + f.write(payload) + else: + with open(tmp_path, "wb") as f: + shutil.copyfileobj(request.input_reference.file, f) + # image, video, or reject. + if _reference_is_image(tmp_path): + is_video = False + elif _reference_is_video(tmp_path): + is_video = True + else: + raise ValueError( + "input_reference content is neither a decodable image nor a decodable video." + ) + ref_path = os.path.join( + media_storage_path, f"{id}_reference{'.mp4' if is_video else '.png'}" ) - ref_path = os.path.join( - media_storage_path, f"{id}_reference{'.mp4' if is_video else '.png'}" - ) - os.replace(tmp_path, ref_path) + os.replace(tmp_path, ref_path) + except Exception: + # Cleanup-and-reraise, not handling: every failure path — + # validation errors (400) and I/O or dependency errors (500) + # alike — must not leak the temporary materialization. + if os.path.exists(tmp_path): + os.remove(tmp_path) + raise if is_video: if params.extra_params is None: params.extra_params = {} diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index 952854161935..00c41baec4c5 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -366,6 +366,30 @@ def test_undecodable_reference_raises_and_cleans_up(self, tmp_path): # The temporary materialization is removed on rejection. assert list(tmp_path.iterdir()) == [] + def test_malformed_base64_reference_raises_and_cleans_up(self, tmp_path): + generator = _StubVisualGen() + # "ABC" survives the lenient alphabet filter but has an invalid + # length, so b64decode raises. + request = VideoGenerationRequest(prompt="x", input_reference="ABC") + with pytest.raises(ValueError, match="not valid base64"): + parse_visual_gen_params(request, "vid-7", generator, media_storage_path=str(tmp_path)) + assert list(tmp_path.iterdir()) == [] + + def test_upload_stream_failure_cleans_up_tmp(self, tmp_path): + generator = _StubVisualGen() + + class _BrokenStream: + def read(self, *args, **kwargs): + raise OSError("client went away") + + upload = UploadFile(file=_BrokenStream(), filename="clip.mp4") + request = VideoGenerationRequest(prompt="x", input_reference=upload) + # I/O failures keep their server-error semantics (no 400 masking) … + with pytest.raises(OSError, match="client went away"): + parse_visual_gen_params(request, "vid-8", generator, media_storage_path=str(tmp_path)) + # … but the partial materialization must not leak. + assert list(tmp_path.iterdir()) == [] + # ============================================================================= # _merge_extra_params — the merge truth table From f48a358b46d353550335ce55fecf13bc57a01e94 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 9 Jul 2026 08:18:15 -0700 Subject: [PATCH 11/64] Add docstrings to Cosmos3 action helper functions Signed-off-by: Igor Shovkun --- .../visual_gen/models/cosmos3/action.py | 72 +++++++++++++++++++ 1 file changed, 72 insertions(+) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py index 0f65df684718..b3bbf4c63941 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py @@ -83,6 +83,20 @@ def normalize_action_resolution(resolution: Any) -> int: + """Validate and coerce the requested action resolution bucket. + + Args: + resolution: Bucket selector, int-like (e.g. ``480``, ``"720"``). Must + match one of the buckets declared in ``VIDEO_RES_SIZE_INFO`` + (``COSMOS3_ACTION_RESOLUTIONS``). + + Returns: + The validated bucket as ``int``. + + Raises: + ValueError: If ``resolution`` is ``None``, not int-coercible, or not a + known bucket. + """ if resolution is None: raise ValueError("Cosmos3 action_resolution is required for action generation.") try: @@ -100,6 +114,19 @@ def normalize_action_resolution(resolution: Any) -> int: def normalize_action_mode(mode: Any) -> str | None: + """Canonicalize the requested action mode. + + Args: + mode: Raw request value. Strings are stripped and lower-cased; + ``None`` or an empty string means action generation is disabled. + + Returns: + One of ``ACTION_MODES`` (``"policy"``, ``"forward_dynamics"``, + ``"inverse_dynamics"``), or ``None`` when no action mode is requested. + + Raises: + ValueError: If ``mode`` is non-empty but not a supported action mode. + """ if mode is None: return None normalized = str(mode).strip().lower() @@ -340,6 +367,51 @@ def prepare_action_latents( dtype: torch.dtype, action_input: Any = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: + """Build the initial action-token latents for denoising. + + The action analog of the video-side latent preparation: tokens that are + *given* are pinned to their clean values via the condition mask, tokens + that are *predicted* start as Gaussian noise. For ``forward_dynamics`` + the trajectory in ``action_input`` is the given conditioning; for + ``policy`` and ``inverse_dynamics`` every action token is generated. + + Args: + mode: Canonical action mode (see :func:`normalize_action_mode`). + action_chunk_size: Number of action timesteps (tokens) in the + sequence. A ``forward_dynamics`` input trajectory is right-padded + by repeating its last row, or truncated, to this length. + raw_action_dim: The embodiment's true degrees of freedom — the number + of *meaningful* leading channels of each action vector (e.g. 9 + for ``av``, 10 for ``bridge_orig_lerobot``). Channels beyond it + are zeroed in both the clean values and the noise. Optional for + ``forward_dynamics`` (inferred from ``action_input``'s last dim); + required for ``policy`` / ``inverse_dynamics``. + action_dim: Model-side padded action width (the checkpoint's + ``max_action_dim``, 64 for Cosmos3-Nano). Trajectories are + zero-padded from ``raw_action_dim`` up to this width. + generator: RNG for the noise draw (seed reproducibility). + device: Device for the returned tensors. + dtype: Dtype for the returned tensors. + action_input: ``forward_dynamics`` only — the given trajectory of + shape ``[T, raw_action_dim]`` (tensor, array, nested list, or + anything :func:`load_action_tensor` accepts). Ignored otherwise. + + Returns: + A ``(action_latents, action_velocity_mask, clean_action, + raw_action_dim)`` tuple: + + * ``action_latents``: ``[1, action_chunk_size, action_dim]`` — clean + values at conditioned positions, noise elsewhere. + * ``action_velocity_mask``: ``1 - condition_mask``; 1 where the + scheduler should integrate the predicted velocity. + * ``clean_action``: ``[1, action_chunk_size, action_dim]`` padded + clean trajectory (all zeros in the generative modes). + * ``raw_action_dim``: the resolved value as ``int``. + + Raises: + ValueError: If ``raw_action_dim`` is missing in a generative mode or + falls outside ``[1, action_dim]``. + """ if mode == ACTION_MODE_FORWARD_DYNAMICS: action = load_action_tensor(action_input) if action.shape[0] < action_chunk_size: From f2129425cca522adcf5d48eb75447ad0242c2751 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 9 Jul 2026 09:37:11 -0700 Subject: [PATCH 12/64] Enforce path_or_list extra-param type validation Signed-off-by: Igor Shovkun --- tensorrt_llm/visual_gen/params.py | 3 +++ .../visual_gen/test_visual_gen_params.py | 24 +++++++++++++++++++ 2 files changed, 27 insertions(+) diff --git a/tensorrt_llm/visual_gen/params.py b/tensorrt_llm/visual_gen/params.py index 9e71a25d0505..3b00d4080ee3 100644 --- a/tensorrt_llm/visual_gen/params.py +++ b/tensorrt_llm/visual_gen/params.py @@ -12,6 +12,7 @@ # 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. +import os from typing import Any, Dict, List, Optional, Union from pydantic import Field @@ -93,6 +94,8 @@ class VisualGenParams(StrictBaseModel): "bool": (bool,), "str": (str,), "list": (list,), + # Cosmos3 V2V `video` reference: a media path or an in-memory frame list. + "path_or_list": (str, os.PathLike, list), } # Generation config fields that pipelines declare defaults for. If a user diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py index e69c81d21f1b..fb497bf5b928 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py @@ -760,6 +760,30 @@ def test_valid_extra_params_accepted(self): req = self._make_request(extra_params={"stg_scale": 0.5}) self._merge_and_validate(executor, req) # should not raise + def test_path_or_list_extra_param_type_enforced(self): + """Cosmos3 V2V's `video` extra param declares type "path_or_list"; the + type map must know it, otherwise the type check is silently skipped + and a bad value reaches the pipeline.""" + from tensorrt_llm._torch.visual_gen.pipeline import ExtraParamSchema + from tensorrt_llm.visual_gen.params import VisualGenParams, validate_visual_gen_params + + specs = {"video": ExtraParamSchema(type="path_or_list", default=None)} + + # A media path and an in-memory frame list are both valid forms. + for good in ("/tmp/clip.mp4", ["frame0.png", "frame1.png"]): + validate_visual_gen_params( + VisualGenParams(extra_params={"video": good}), + declared_defaults={}, + extra_param_specs=specs, + ) + + with pytest.raises(ValueError, match="expected type 'path_or_list'"): + validate_visual_gen_params( + VisualGenParams(extra_params={"video": 42}), + declared_defaults={}, + extra_param_specs=specs, + ) + # --- unsupported universal fields --- def test_num_frames_on_image_pipeline_raises(self): From 300cd9b15ff2cb14a7ddf02160e4a82ee29fd24e Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Fri, 10 Jul 2026 12:55:54 -0700 Subject: [PATCH 13/64] Remove av from requirements; document manual install Signed-off-by: Igor Shovkun --- examples/visual_gen/models/cosmos3/README.md | 2 +- examples/visual_gen/serve/README.md | 2 +- requirements.txt | 3 --- 3 files changed, 2 insertions(+), 5 deletions(-) diff --git a/examples/visual_gen/models/cosmos3/README.md b/examples/visual_gen/models/cosmos3/README.md index 77784e8b164f..25ae32d357b3 100644 --- a/examples/visual_gen/models/cosmos3/README.md +++ b/examples/visual_gen/models/cosmos3/README.md @@ -34,7 +34,7 @@ export TRTLLM_DISABLE_COSMOS3_GUARDRAILS=1 ## Media I/O dependencies - 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, inverse dynamics) uses the `av` (PyAV) package, installed with TensorRT-LLM's requirements. +- Decoding `.mp4`/`.avi` reference videos (V2V, inverse dynamics) uses the `av` (PyAV) package. It is **not** bundled with TensorRT-LLM — install it yourself: `pip install av`. Frame directories and single-image references work without it. ## Deployment configs diff --git a/examples/visual_gen/serve/README.md b/examples/visual_gen/serve/README.md index 63fae4b5d93e..b267a1923205 100644 --- a/examples/visual_gen/serve/README.md +++ b/examples/visual_gen/serve/README.md @@ -286,7 +286,7 @@ You can customize these by: - `frame_rate` (canonical) or `fps` (alias): frames per second - `num_frames`: when set, wins over the `seconds * frame_rate` derivation - `seed`, `num_inference_steps`, `guidance_scale`, `max_sequence_length`, `negative_prompt`: per-request denoise controls -- `input_reference`: Reference image (I2V/TI2V) or video (V2V), classified by decoding the content — filename and content type are ignored; accepted as base64-encoded string in JSON or as a file in multipart form-data. Undecodable content returns HTTP 400. Video decode requires the `av` (PyAV) package on the server. +- `input_reference`: Reference image (I2V/TI2V) or video (V2V), classified by decoding the content — filename and content type are ignored; accepted as base64-encoded string in JSON or as a file in multipart form-data. Undecodable content returns HTTP 400. Video decode requires the `av` (PyAV) package on the server (not bundled — `pip install av`). - `extra_params`: model-specific overflow (see below) - `response_format`: `"b64_json"` or `"url"` - `format`: Generation content encoding. Video encoders: `"mp4"`, `"avi"`, `"auto"`. Tensor formats: `"safetensors"`, `"pt"` (carries video + audio + scalar metadata in one payload for LTX-2). diff --git a/requirements.txt b/requirements.txt index a7f38440bfff..d735182a8fe8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,9 +1,6 @@ --extra-index-url https://download.pytorch.org/whl/cu130 -c constraints.txt accelerate>=1.7.0 -# Video reference decode (torchvision.io.read_video backend) and serve-side -# image-vs-video content classification for visual_gen (Cosmos3 V2V). -av build colored cuda-python>=13 From afff945f5861832f18375ce8949ad4cd0a129e21 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Fri, 10 Jul 2026 13:07:08 -0700 Subject: [PATCH 14/64] Update output dataclass tests for action fields Signed-off-by: Igor Shovkun --- tests/unittest/visual_gen/test_output.py | 57 +++++++++++++++++++++++- 1 file changed, 55 insertions(+), 2 deletions(-) diff --git a/tests/unittest/visual_gen/test_output.py b/tests/unittest/visual_gen/test_output.py index e27f9c33e3b5..54097e64558e 100644 --- a/tests/unittest/visual_gen/test_output.py +++ b/tests/unittest/visual_gen/test_output.py @@ -32,8 +32,12 @@ def test_visual_gen_output_is_dataclass(): "image", "video", "audio", + "action", "frame_rate", "audio_sample_rate", + "raw_action_dim", + "action_mode", + "domain_id", "error", "metrics", } @@ -62,6 +66,10 @@ def test_minimal_construction_defaults(): assert out.audio is None assert out.frame_rate is None assert out.audio_sample_rate is None + assert out.action is None + assert out.raw_action_dim is None + assert out.action_mode is None + assert out.domain_id is None assert out.error is None assert out.metrics is None @@ -753,15 +761,19 @@ def test_encoding_not_top_level_reexport(): # --------------------------------------------------------------------------- -def test_pipeline_output_has_eight_fields(): - """PipelineOutput has the eight expected fields.""" +def test_pipeline_output_field_set(): + """PipelineOutput exposes exactly the expected fields.""" field_names = {f.name for f in fields(PipelineOutput)} assert field_names == { "image", "video", "audio", + "action", "frame_rate", "audio_sample_rate", + "raw_action_dim", + "action_mode", + "domain_id", "pre_denoise", "denoise", "post_denoise", @@ -774,6 +786,10 @@ def test_pipeline_output_default_construction(): assert p.image is None assert p.video is None assert p.audio is None + assert p.action is None + assert p.raw_action_dim is None + assert p.action_mode is None + assert p.domain_id is None assert p.frame_rate is None assert p.audio_sample_rate is None assert p.pre_denoise == 0.0 @@ -781,6 +797,43 @@ def test_pipeline_output_default_construction(): assert p.post_denoise == 0.0 +def test_from_response_propagates_action_fields(): + """to_visual_gen_output carries the action tensor and its metadata.""" + action = torch.zeros(1, 16, 10, dtype=torch.float32) + pipeline_out = PipelineOutput( + action=action, + raw_action_dim=10, + action_mode="policy", + domain_id=7, + ) + resp = DiffusionResponse(request_id=5, output=pipeline_out, generation=12.0) + out = to_visual_gen_output(resp) + assert out.action is action + assert out.raw_action_dim == 10 + assert out.action_mode == "policy" + assert out.domain_id == 7 + + +def test_batch_split_slices_action(): + """split_visual_gen_output slices batched actions per item and shares scalar metadata.""" + action = torch.stack([torch.full((16, 10), float(v), dtype=torch.float32) for v in (1, 2, 3)]) + pipeline_out = PipelineOutput( + action=action, + raw_action_dim=10, + action_mode="inverse_dynamics", + domain_id=1, + ) + resp = DiffusionResponse(request_id=9, output=pipeline_out, generation=30.0) + outs = split_visual_gen_output(resp, batch_size=3) + assert len(outs) == 3 + for i, out in enumerate(outs): + assert out.action.shape == (16, 10) + assert float(out.action[0, 0]) == float(i + 1) + assert out.raw_action_dim == 10 + assert out.action_mode == "inverse_dynamics" + assert out.domain_id == 1 + + def test_media_output_unimportable(): """``MediaOutput`` is not importable from any path.""" with pytest.raises(ImportError): From 8ad23864d741d2ed23afdea263df6cb45aa69b27 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Fri, 10 Jul 2026 20:29:48 -0700 Subject: [PATCH 15/64] Pass raw_timestep in Cosmos3 action transformer tests Signed-off-by: Igor Shovkun --- .../_torch/visual_gen/test_cosmos3_transformer.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py index 7f02d97ad117..c5d563c7fc9e 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py @@ -455,7 +455,8 @@ def test_forward_with_action(self, action_model_config): with torch.inference_mode(): out = model( hidden_states=hs, - timestep=ts, + timestep=ts / _NUM_TRAIN_TIMESTEPS, + raw_timestep=ts, text_ids=text_ids, text_mask=text_mask, video_shape=video_shape, @@ -477,7 +478,8 @@ def test_forward_without_action_latents_returns_none(self, action_model_config): with torch.inference_mode(): out = model( hidden_states=hs, - timestep=ts, + timestep=ts / _NUM_TRAIN_TIMESTEPS, + raw_timestep=ts, text_ids=text_ids, text_mask=text_mask, video_shape=video_shape, @@ -499,7 +501,8 @@ def test_forward_with_action_noisy_mask(self, action_model_config): with torch.inference_mode(): out = model( hidden_states=hs, - timestep=ts, + timestep=ts / _NUM_TRAIN_TIMESTEPS, + raw_timestep=ts, text_ids=text_ids, text_mask=text_mask, video_shape=video_shape, @@ -523,7 +526,8 @@ def test_forward_with_action_multiframe(self, action_model_config): with torch.inference_mode(): out = model( hidden_states=hs, - timestep=ts, + timestep=ts / _NUM_TRAIN_TIMESTEPS, + raw_timestep=ts, text_ids=text_ids, text_mask=text_mask, video_shape=video_shape, From 2716347b74958893f27181d83b22502acc89199a Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 13 Jul 2026 12:53:09 -0700 Subject: [PATCH 16/64] Remove Cosmos3 action generation from the V2V PR Action generation (policy / forward_dynamics / inverse_dynamics), embodiment domain presets, action mRoPE, the DomainAwareLinear projections, and all action output fields, CLI flags, prompts, and tests are removed. V2V, audio, and the shared video encode/decode infrastructure they were built on are retained. Action generation will land in a separate PR. Signed-off-by: Igor Shovkun --- examples/visual_gen/models/cosmos3/cosmos3.py | 244 +--------- .../prompts/action_forward_dynamics.json | 5 - .../prompts/action_inverse_dynamics.json | 4 - .../models/cosmos3/prompts/action_policy.json | 5 - .../visual_gen/models/cosmos3/action.py | 455 ------------------ .../visual_gen/models/cosmos3/defaults.py | 302 +----------- .../models/cosmos3/pipeline_cosmos3.py | 320 +----------- .../models/cosmos3/transformer_cosmos3.py | 240 +-------- tensorrt_llm/_torch/visual_gen/output.py | 31 +- tensorrt_llm/visual_gen/output.py | 4 - .../test_cosmos3_transformer_parallel.py | 72 --- .../_torch/visual_gen/test_cosmos3_action.py | 271 ----------- .../visual_gen/test_cosmos3_pipeline.py | 155 ------ .../visual_gen/test_cosmos3_transformer.py | 142 ------ tests/unittest/visual_gen/test_output.py | 57 +-- 15 files changed, 22 insertions(+), 2285 deletions(-) delete mode 100644 examples/visual_gen/models/cosmos3/prompts/action_forward_dynamics.json delete mode 100644 examples/visual_gen/models/cosmos3/prompts/action_inverse_dynamics.json delete mode 100644 examples/visual_gen/models/cosmos3/prompts/action_policy.json delete mode 100644 tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py delete mode 100644 tests/unittest/_torch/visual_gen/test_cosmos3_action.py diff --git a/examples/visual_gen/models/cosmos3/cosmos3.py b/examples/visual_gen/models/cosmos3/cosmos3.py index 6795c1495eeb..2e35925a5b6c 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3.py +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -13,7 +13,7 @@ # 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"""Cosmos3 Text(+Image)-to-Video(+Audio) and action generation. +r"""Cosmos3 Text(+Image)-to-Video(+Audio) generation. Cosmos3 supports the following generation modes from a single checkpoint: @@ -27,12 +27,9 @@ first (or last, per ``condition_video_keep``) frames of a reference video via ``--video_path`` (a local frame directory, ``.mp4``/``.avi`` file, or single image; ``.mp4``/``.avi`` decode requires the ``av`` package). - Passing ``--video_path`` without ``--action_mode`` selects V2V. - **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). -- **Action** — policy / forward dynamics / inverse dynamics generation - (pass ``--action_mode``). Action and audio generation are mutually exclusive. Checkpoints (pass the Hub ID or local path via ``--model``): @@ -101,36 +98,6 @@ 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 - - # Action — policy (first frame + instruction -> predicted action + rollout video) - python cosmos3.py --model nvidia/Cosmos3-Nano \ - --prompt_file prompts/action_policy.json \ - --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml \ - --action_mode policy \ - --domain_name bridge_orig_lerobot \ - --raw_action_dim 10 \ - --output_path policy_rollout.mp4 \ - --action_output_path policy_action.json - - # Action — forward dynamics (first frame + action trajectory -> rollout video) - python cosmos3.py --model nvidia/Cosmos3-Nano \ - --prompt_file prompts/action_forward_dynamics.json \ - --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml \ - --action_mode forward_dynamics \ - --domain_name av \ - --action_json action_trajectory.json \ - --output_path forward_dynamics.mp4 - - # Action — inverse dynamics (video -> predicted action) - python cosmos3.py --model nvidia/Cosmos3-Nano \ - --prompt_file prompts/action_inverse_dynamics.json \ - --video_path /path/to/frames/ \ - --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml \ - --action_mode inverse_dynamics \ - --domain_name bridge_orig_lerobot \ - --raw_action_dim 10 \ - --output_path inverse_video.mp4 \ - --action_output_path inverse_action.json """ import argparse @@ -140,17 +107,9 @@ from typing import Any, Dict, Optional from tensorrt_llm import VisualGen, VisualGenArgs -from tensorrt_llm._torch.visual_gen.models.cosmos3.action import VIDEO_RES_SIZE_INFO -from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( - COSMOS3_ACTION_PARAMS, - get_domain_preset, - resolve_domain_action_config, -) _SCRIPT_DIR = Path(__file__).resolve().parent -ACTION_MODES = frozenset({"policy", "forward_dynamics", "inverse_dynamics"}) - def _resolve_path(path: str) -> str: candidate = Path(path) @@ -207,114 +166,6 @@ def resolve_prompt_and_options( return resolved_prompt, resolved_image, resolved_enable_audio, resolved_output_type -def _validate_action_args( - args: argparse.Namespace, resolved_image_path: Optional[str] = None -) -> None: - if args.action_mode is None: - return - - # The first frame may come from --image_path or a prompt file's vision_path. - has_first_frame = resolved_image_path is not None or args.video_path is not None - - mode = args.action_mode.strip().lower() - if mode not in ACTION_MODES: - raise SystemExit( - f"Invalid --action_mode {args.action_mode!r}; expected one of {sorted(ACTION_MODES)}." - ) - if args.enable_audio: - raise SystemExit("Cosmos3 does not support joint action and audio generation.") - if args.output_type != "video": - raise SystemExit("Action generation requires --output_type video.") - - if mode == "forward_dynamics": - if args.action_json is None: - raise SystemExit(f"{mode} requires --action_json.") - if not has_first_frame: - raise SystemExit( - f"{mode} requires --image_path, a prompt-file vision_path, or --video_path " - "for the first frame." - ) - elif mode == "policy": - if not has_first_frame: - raise SystemExit( - f"{mode} requires --image_path, a prompt-file vision_path, or --video_path " - "for the first frame." - ) - preset = get_domain_preset(args.domain_name, args.domain_id) - effective_raw_dim = args.raw_action_dim or (preset or {}).get("raw_action_dim") - if effective_raw_dim is None: - raise SystemExit( - f"{mode} requires --raw_action_dim or a known --domain_name with a preset." - ) - elif mode == "inverse_dynamics": - if args.video_path is None: - raise SystemExit( - f"{mode} requires --video_path (frame directory, .mp4/.avi, or image)." - ) - preset = get_domain_preset(args.domain_name, args.domain_id) - effective_raw_dim = args.raw_action_dim or (preset or {}).get("raw_action_dim") - if effective_raw_dim is None: - raise SystemExit( - f"{mode} requires --raw_action_dim or a known --domain_name with a preset." - ) - - -def _apply_action_generation_params(params, args: argparse.Namespace) -> None: - """Set action defaults on the request; domain presets override generic 480p.""" - cfg = resolve_domain_action_config( - domain_name=args.domain_name, - domain_id=args.domain_id, - raw_action_dim=args.raw_action_dim, - action_chunk_size=args.action_chunk_size, - action_resolution=args.action_resolution, - ) - bucket = str(cfg["action_resolution"]) - width, height = VIDEO_RES_SIZE_INFO[bucket]["16,9"] - params.width = width - params.height = height - params.num_frames = cfg["num_frames"] - params.num_inference_steps = COSMOS3_ACTION_PARAMS["num_inference_steps"] - params.guidance_scale = COSMOS3_ACTION_PARAMS["guidance_scale"] - params.frame_rate = cfg["frame_rate"] - params.extra_params["action_chunk_size"] = cfg["action_chunk_size"] - if cfg["raw_action_dim"] is not None: - params.extra_params["raw_action_dim"] = cfg["raw_action_dim"] - params.extra_params["action_resolution"] = cfg["action_resolution"] - - -def _default_action_output_path(video_path: str) -> str: - stem = Path(video_path) - if stem.suffix: - return str(stem.with_name(f"{stem.stem}_action.json")) - return f"{video_path}_action.json" - - -def _save_action_output(output, path: str) -> None: - if output.action is None: - return - - action = output.action - if action.ndim == 3 and action.shape[0] == 1: - action_data = action[0].tolist() - shape = list(action.shape[1:]) - else: - action_data = action.tolist() - shape = list(action.shape) - - payload = { - "action_mode": output.action_mode, - "domain_id": output.domain_id, - "raw_action_dim": output.raw_action_dim, - "shape": shape, - "dtype": str(action.dtype).replace("torch.", ""), - "data": action_data, - } - out_path = Path(path) - out_path.parent.mkdir(parents=True, exist_ok=True) - with out_path.open("w", encoding="utf-8") as f: - json.dump(payload, f, indent=2) - - def main(): parser = argparse.ArgumentParser(description="Cosmos3 Text(+Image)-to-Video(+Audio) example") parser.add_argument( @@ -374,68 +225,11 @@ def main(): "--use_system_prompt", action="store_true", help="Use system prompt in prompt" ) parser.add_argument("--enable_audio", action="store_true", help="Enable audio generation") - parser.add_argument( - "--action_mode", - type=str, - default=None, - choices=sorted(ACTION_MODES), - help="Action mode: policy, forward_dynamics, or inverse_dynamics", - ) - parser.add_argument( - "--domain_name", - type=str, - default=None, - help="Embodiment domain name (e.g. bridge_orig_lerobot, av, droid_lerobot)", - ) - parser.add_argument( - "--domain_id", - type=int, - default=None, - help="Embodiment domain id (alternative to --domain_name)", - ) - parser.add_argument( - "--raw_action_dim", - type=int, - default=None, - help="Raw action DOF for policy/inverse_dynamics", - ) - parser.add_argument( - "--action_chunk_size", - type=int, - default=None, - help=f"Action tokens to generate (default {COSMOS3_ACTION_PARAMS['action_chunk_size']})", - ) - parser.add_argument( - "--action_json", - type=str, - default=None, - help="JSON file with action trajectory [T, D] for forward_dynamics", - ) parser.add_argument( "--video_path", type=str, default=None, - help="Reference video for V2V (or inverse_dynamics with --action_mode): " - "a local frame directory, .mp4/.avi file, or image path", - ) - parser.add_argument( - "--action_resolution", - type=int, - default=None, - choices=[256, 480, 704, 720], - help=("Resolution bucket for action image sizing. Defaults to the domain preset or 480."), - ) - parser.add_argument( - "--action_fps", - type=float, - default=None, - help="Action-token temporal rate for mRoPE (Hz). Defaults to frame_rate.", - ) - parser.add_argument( - "--action_output_path", - type=str, - default=None, - help="Path to save predicted action JSON (default: _action.json)", + help="Reference video for V2V: a local frame directory, .mp4/.avi file, or image path", ) parser.add_argument( "--output_type", type=str, default="video", help="Output type (video, image)" @@ -454,7 +248,6 @@ def main(): enable_audio=args.enable_audio, output_type=args.output_type, ) - _validate_action_args(args, resolved_image_path=image_path) # Engine config from shared YAML (optional); model-specific defaults apply otherwise. extra_args = VisualGenArgs.from_yaml(args.visual_gen_args) if args.visual_gen_args else None @@ -466,9 +259,6 @@ def main(): if image_path is not None: params.image = image_path - if args.action_mode is not None: - _apply_action_generation_params(params, args) - 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"): @@ -487,23 +277,7 @@ def main(): params.extra_params["enable_audio"] = enable_audio params.extra_params["use_guardrails"] = not args.disable_guardrails params.extra_params["output_type"] = output_type - params.extra_params["action_resolution"] = args.action_resolution - - if args.action_mode is not None: - params.extra_params["action_mode"] = args.action_mode - if args.domain_name is not None: - params.extra_params["domain_name"] = args.domain_name - if args.domain_id is not None: - params.extra_params["domain_id"] = args.domain_id - if args.raw_action_dim is not None: - params.extra_params["raw_action_dim"] = args.raw_action_dim - if args.action_chunk_size is not None: - params.extra_params["action_chunk_size"] = args.action_chunk_size - if args.action_fps is not None: - params.extra_params["action_fps"] = args.action_fps - if args.action_json is not None: - with open(args.action_json, encoding="utf-8") as f: - params.extra_params["action"] = json.load(f) + if args.video_path is not None: params.extra_params["video"] = args.video_path @@ -522,18 +296,6 @@ def main(): output.save(args.output_path) print(f"Saved: {args.output_path}") - if args.action_mode is not None: - action_path = args.action_output_path or _default_action_output_path(args.output_path) - _save_action_output(output, action_path) - if output.action is not None: - print(f"Saved action: {action_path}") - print( - f"Action shape: {tuple(output.action.shape)}, " - f"raw_action_dim={output.raw_action_dim}, domain_id={output.domain_id}" - ) - else: - print("Warning: action_mode was set but the output carried no action tensor.") - print(output.metrics) diff --git a/examples/visual_gen/models/cosmos3/prompts/action_forward_dynamics.json b/examples/visual_gen/models/cosmos3/prompts/action_forward_dynamics.json deleted file mode 100644 index e780ef7b26b5..000000000000 --- a/examples/visual_gen/models/cosmos3/prompts/action_forward_dynamics.json +++ /dev/null @@ -1,5 +0,0 @@ -{ - "model_mode": "image2video", - "prompt": "Robot manipulation rollout: the right arm extends to the fruit display, picks up a pear, and places it into the bag in the shopping cart.", - "vision_path": "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/vision/robot_153.jpg" -} diff --git a/examples/visual_gen/models/cosmos3/prompts/action_inverse_dynamics.json b/examples/visual_gen/models/cosmos3/prompts/action_inverse_dynamics.json deleted file mode 100644 index 326e74cb77f5..000000000000 --- a/examples/visual_gen/models/cosmos3/prompts/action_inverse_dynamics.json +++ /dev/null @@ -1,4 +0,0 @@ -{ - "model_mode": "image2video", - "prompt": "Recover the robot action trajectory from this clip of the arm picking up a pear and placing it into the bag." -} diff --git a/examples/visual_gen/models/cosmos3/prompts/action_policy.json b/examples/visual_gen/models/cosmos3/prompts/action_policy.json deleted file mode 100644 index 1eb4e7572ee5..000000000000 --- a/examples/visual_gen/models/cosmos3/prompts/action_policy.json +++ /dev/null @@ -1,5 +0,0 @@ -{ - "model_mode": "image2video", - "prompt": "Pick up the pear from the fruit display and place it into the plastic bag in the shopping cart.", - "vision_path": "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/vision/robot_153.jpg" -} diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py deleted file mode 100644 index b3bbf4c63941..000000000000 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py +++ /dev/null @@ -1,455 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Action-token helpers for Cosmos3 UVA/action generation.""" - -from __future__ import annotations - -from pathlib import Path -from typing import Any, Optional -from urllib.parse import urlparse - -import numpy as np -import PIL.Image -import torch -from diffusers.utils.torch_utils import randn_tensor - -from tensorrt_llm.inputs.utils import load_image - -from .utils import IMAGE_EXTENSIONS, normalize_video_input, pil_to_rgb - -ACTION_MODE_POLICY = "policy" -ACTION_MODE_FORWARD_DYNAMICS = "forward_dynamics" -ACTION_MODE_INVERSE_DYNAMICS = "inverse_dynamics" -ACTION_MODES = { - ACTION_MODE_POLICY, - ACTION_MODE_FORWARD_DYNAMICS, - ACTION_MODE_INVERSE_DYNAMICS, -} - -EMBODIMENT_TO_DOMAIN_ID: dict[str, int] = { - "no_action": 0, - "av": 1, - "camera_pose": 2, - "hand_pose": 3, - "pusht": 4, - "libero": 5, - "umi": 6, - "bridge_orig_lerobot": 7, - "droid_lerobot": 8, - "robomind-franka": 8, - "galbot": 9, - "robomind-franka-dual": 12, - "robomind-ur": 13, - "agibotworld": 15, - "agibot_gear_gripper": 15, - "agibot_gear_gripper_ext": 15, - "fractal": 20, -} - -VIDEO_RES_SIZE_INFO: dict[str, dict[str, tuple[int, int]]] = { - "256": { - "1,1": (256, 256), - "4,3": (320, 256), - "3,4": (256, 320), - "16,9": (320, 192), - "9,16": (192, 320), - }, - "480": { - "1,1": (640, 640), - "4,3": (736, 544), - "3,4": (544, 736), - "16,9": (832, 480), - "9,16": (480, 832), - }, - "704": { - "1,1": (960, 960), - "4,3": (1088, 832), - "3,4": (832, 1088), - "16,9": (1280, 704), - "9,16": (704, 1280), - }, - "720": { - "1,1": (960, 960), - "4,3": (1104, 832), - "3,4": (832, 1104), - "16,9": (1280, 720), - "9,16": (720, 1280), - }, -} - - -COSMOS3_ACTION_RESOLUTIONS = tuple(int(key) for key in sorted(VIDEO_RES_SIZE_INFO, key=int)) - - -def normalize_action_resolution(resolution: Any) -> int: - """Validate and coerce the requested action resolution bucket. - - Args: - resolution: Bucket selector, int-like (e.g. ``480``, ``"720"``). Must - match one of the buckets declared in ``VIDEO_RES_SIZE_INFO`` - (``COSMOS3_ACTION_RESOLUTIONS``). - - Returns: - The validated bucket as ``int``. - - Raises: - ValueError: If ``resolution`` is ``None``, not int-coercible, or not a - known bucket. - """ - if resolution is None: - raise ValueError("Cosmos3 action_resolution is required for action generation.") - try: - bucket = int(resolution) - except (TypeError, ValueError) as exc: - raise ValueError( - f"Cosmos3 action_resolution must be an int bucket, got {resolution!r}." - ) from exc - if bucket not in COSMOS3_ACTION_RESOLUTIONS: - raise ValueError( - f"Unknown Cosmos3 action_resolution={bucket}; " - f"expected one of {COSMOS3_ACTION_RESOLUTIONS}." - ) - return bucket - - -def normalize_action_mode(mode: Any) -> str | None: - """Canonicalize the requested action mode. - - Args: - mode: Raw request value. Strings are stripped and lower-cased; - ``None`` or an empty string means action generation is disabled. - - Returns: - One of ``ACTION_MODES`` (``"policy"``, ``"forward_dynamics"``, - ``"inverse_dynamics"``), or ``None`` when no action mode is requested. - - Raises: - ValueError: If ``mode`` is non-empty but not a supported action mode. - """ - if mode is None: - return None - normalized = str(mode).strip().lower() - if not normalized: - return None - if normalized not in ACTION_MODES: - raise ValueError( - f"Unsupported Cosmos3 action_mode={mode!r}; expected one of {sorted(ACTION_MODES)}." - ) - return normalized - - -def resolve_domain_id( - *, - domain_id: Any = None, - domain_name: Any = None, - require_explicit: bool = False, -) -> int: - if domain_id is not None: - resolved = int(domain_id) - if resolved < 0: - raise ValueError(f"Cosmos3 domain_id must be non-negative, got {resolved}.") - return resolved - - if domain_name is None or str(domain_name).strip() == "": - if require_explicit: - raise ValueError( - "Cosmos3 action generation requires domain_id or non-empty domain_name." - ) - return 0 - - key = str(domain_name).strip().lower() - if key not in EMBODIMENT_TO_DOMAIN_ID: - raise ValueError( - f"Unknown Cosmos3 action domain_name={domain_name!r}; " - f"expected one of {sorted(EMBODIMENT_TO_DOMAIN_ID)} or pass domain_id directly." - ) - return EMBODIMENT_TO_DOMAIN_ID[key] - - -def action_condition_indexes(mode: str, action_length: int) -> list[int]: - mode = normalize_action_mode(mode) - if mode == ACTION_MODE_FORWARD_DYNAMICS: - return list(range(action_length)) - if mode in {ACTION_MODE_POLICY, ACTION_MODE_INVERSE_DYNAMICS}: - return [] - raise AssertionError(f"Unexpected action mode: {mode!r}") - - -def vision_condition_indexes( - mode: str, video_length: int, temporal_compression_factor: int -) -> list[int]: - mode = normalize_action_mode(mode) - latent_frames = (video_length - 1) // temporal_compression_factor + 1 - if mode in {ACTION_MODE_POLICY, ACTION_MODE_FORWARD_DYNAMICS}: - return [0] - if mode == ACTION_MODE_INVERSE_DYNAMICS: - return list(range(latent_frames)) - raise AssertionError(f"Unexpected action mode: {mode!r}") - - -def action_start_frame_offset(mode: str, action_length: int, video_length: int) -> int: - del mode - if action_length == video_length - 1: - return 1 - if action_length == video_length: - return 0 - raise ValueError( - "Cosmos3 action_chunk_size must equal num_frames - 1 or num_frames; " - f"got action_chunk_size={action_length}, num_frames={video_length}." - ) - - -def build_action_condition_mask( - mode: str, - action_length: int, - *, - device: torch.device, - dtype: torch.dtype, -) -> torch.Tensor: - mask = torch.zeros(1, action_length, 1, device=device, dtype=dtype) - for idx in action_condition_indexes(mode, action_length): - mask[:, idx, :] = 1.0 - return mask - - -def build_vision_condition_mask( - mode: str, - video_length: int, - temporal_compression_factor: int, - *, - device: torch.device, - dtype: torch.dtype, -) -> torch.Tensor: - latent_frames = (video_length - 1) // temporal_compression_factor + 1 - mask = torch.zeros(1, 1, latent_frames, 1, 1, device=device, dtype=dtype) - for idx in vision_condition_indexes(mode, video_length, temporal_compression_factor): - mask[:, :, idx, :, :] = 1.0 - return mask - - -def pad_action_to_dim(action: torch.Tensor, action_dim: int) -> torch.Tensor: - if action.shape[-1] > action_dim: - raise ValueError( - f"Cosmos3 action dimension {action.shape[-1]} exceeds model action_dim={action_dim}." - ) - if action.shape[-1] == action_dim: - return action - padding = torch.zeros( - *action.shape[:-1], action_dim - action.shape[-1], dtype=action.dtype, device=action.device - ) - return torch.cat([action, padding], dim=-1) - - -def load_action_tensor(action: Any = None) -> torch.Tensor: - if action is None: - raise ValueError( - "Cosmos3 forward_dynamics action mode requires an action tensor of shape [T, D]." - ) - if isinstance(action, torch.Tensor): - tensor = action.detach().to(dtype=torch.float32) - else: - tensor = torch.as_tensor(np.asarray(action), dtype=torch.float32) - if tensor.ndim == 3 and tensor.shape[0] == 1: - tensor = tensor.squeeze(0) - if tensor.ndim != 2: - raise ValueError(f"Cosmos3 action must have shape [T, D], got {tuple(tensor.shape)}.") - return tensor - - -def find_closest_target_size(h: int, w: int, resolution: str | int) -> tuple[int, int]: - key = str(resolution) - if key not in VIDEO_RES_SIZE_INFO: - raise ValueError( - f"Unknown Cosmos3 action resolution={resolution!r}; " - f"expected one of {sorted(VIDEO_RES_SIZE_INFO)}." - ) - input_ratio = h / w - best_size = None - best_diff = float("inf") - for cand_w, cand_h in VIDEO_RES_SIZE_INFO[key].values(): - diff = abs(input_ratio - cand_h / cand_w) - if diff < best_diff: - best_diff = diff - best_size = (cand_w, cand_h) - assert best_size is not None - return best_size - - -def resolve_action_size( - height: Optional[int], - width: Optional[int], - ref_image: PIL.Image.Image, - action_resolution: int, -) -> tuple[int, int]: - """Fill unset action H/W from the action resolution bucket; honor explicit values.""" - if height is not None and width is not None: - return height, width - target_w, target_h = find_closest_target_size( - ref_image.height, ref_image.width, action_resolution - ) - return ( - height if height is not None else target_h, - width if width is not None else target_w, - ) - - -def action_reference_image( - *, - action_mode: str, - image: Any, - video: Any, -) -> PIL.Image.Image: - """Resolve the reference frame used for action sizing and conditioning.""" - if action_mode == ACTION_MODE_INVERSE_DYNAMICS: - source = video if video is not None else image - frames = normalize_video_input(source, max_frames=1) - if not frames: - raise ValueError("Cosmos3 action_mode='inverse_dynamics' requires a video input.") - return pil_to_rgb(frames[0]) - - source = image if image is not None else video - if source is None: - raise ValueError(f"Cosmos3 action_mode={action_mode!r} requires an image or video input.") - if isinstance(source, PIL.Image.Image): - return source.convert("RGB") - if isinstance(source, str): - if urlparse(source).scheme in ("http", "https", "data", "file"): - # Same URI-aware loader as the I2V image branch; plain local - # paths keep the richer path handling below (frame dirs, videos). - return load_image(source, format="pil").convert("RGB") - path = Path(source) - if path.is_file() and path.suffix.lower() in IMAGE_EXTENSIONS: - return PIL.Image.open(source).convert("RGB") - frames = normalize_video_input(source, max_frames=1) - if not frames: - raise ValueError( - f"Cosmos3 action_mode={action_mode!r} requires an image or video input." - ) - return pil_to_rgb(frames[0]) - raise TypeError( - f"Cosmos3 action reference image must be PIL.Image or path, got {type(source)!r}." - ) - - -def resize_and_pad_action_image( - image: PIL.Image.Image, target_h: int, target_w: int -) -> PIL.Image.Image: - scale = min(target_w / image.width, target_h / image.height, 1.0) - resize_w = max(1, int(scale * image.width + 0.5)) - resize_h = max(1, int(scale * image.height + 0.5)) - if (resize_w, resize_h) != image.size: - image = image.resize((resize_w, resize_h), PIL.Image.Resampling.BICUBIC) - - array = np.asarray(image) - pad_h = target_h - resize_h - pad_w = target_w - resize_w - if pad_h < 0 or pad_w < 0: - raise ValueError( - f"Cosmos3 action image resize exceeded target size: resized={(resize_h, resize_w)}, " - f"target={(target_h, target_w)}." - ) - if pad_h == 0 and pad_w == 0: - return image - pad_mode = "reflect" if pad_h < resize_h and pad_w < resize_w else "edge" - padded = np.pad(array, ((0, pad_h), (0, pad_w), (0, 0)), mode=pad_mode) - return PIL.Image.fromarray(padded) - - -def prepare_action_latents( - *, - mode: str, - action_chunk_size: int, - raw_action_dim: Optional[int], - action_dim: int, - generator: torch.Generator, - device: torch.device, - dtype: torch.dtype, - action_input: Any = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: - """Build the initial action-token latents for denoising. - - The action analog of the video-side latent preparation: tokens that are - *given* are pinned to their clean values via the condition mask, tokens - that are *predicted* start as Gaussian noise. For ``forward_dynamics`` - the trajectory in ``action_input`` is the given conditioning; for - ``policy`` and ``inverse_dynamics`` every action token is generated. - - Args: - mode: Canonical action mode (see :func:`normalize_action_mode`). - action_chunk_size: Number of action timesteps (tokens) in the - sequence. A ``forward_dynamics`` input trajectory is right-padded - by repeating its last row, or truncated, to this length. - raw_action_dim: The embodiment's true degrees of freedom — the number - of *meaningful* leading channels of each action vector (e.g. 9 - for ``av``, 10 for ``bridge_orig_lerobot``). Channels beyond it - are zeroed in both the clean values and the noise. Optional for - ``forward_dynamics`` (inferred from ``action_input``'s last dim); - required for ``policy`` / ``inverse_dynamics``. - action_dim: Model-side padded action width (the checkpoint's - ``max_action_dim``, 64 for Cosmos3-Nano). Trajectories are - zero-padded from ``raw_action_dim`` up to this width. - generator: RNG for the noise draw (seed reproducibility). - device: Device for the returned tensors. - dtype: Dtype for the returned tensors. - action_input: ``forward_dynamics`` only — the given trajectory of - shape ``[T, raw_action_dim]`` (tensor, array, nested list, or - anything :func:`load_action_tensor` accepts). Ignored otherwise. - - Returns: - A ``(action_latents, action_velocity_mask, clean_action, - raw_action_dim)`` tuple: - - * ``action_latents``: ``[1, action_chunk_size, action_dim]`` — clean - values at conditioned positions, noise elsewhere. - * ``action_velocity_mask``: ``1 - condition_mask``; 1 where the - scheduler should integrate the predicted velocity. - * ``clean_action``: ``[1, action_chunk_size, action_dim]`` padded - clean trajectory (all zeros in the generative modes). - * ``raw_action_dim``: the resolved value as ``int``. - - Raises: - ValueError: If ``raw_action_dim`` is missing in a generative mode or - falls outside ``[1, action_dim]``. - """ - if mode == ACTION_MODE_FORWARD_DYNAMICS: - action = load_action_tensor(action_input) - if action.shape[0] < action_chunk_size: - pad = action[-1:].repeat(action_chunk_size - action.shape[0], 1) - action = torch.cat([action, pad], dim=0) - elif action.shape[0] > action_chunk_size: - action = action[:action_chunk_size] - if raw_action_dim is None: - raw_action_dim = int(action.shape[-1]) - clean_action = pad_action_to_dim(action, action_dim) - else: - if raw_action_dim is None: - raise ValueError( - "Cosmos3 action_mode='policy' and 'inverse_dynamics' require raw_action_dim." - ) - clean_action = torch.zeros(action_chunk_size, action_dim, dtype=torch.float32) - - raw_action_dim = int(raw_action_dim) - if raw_action_dim <= 0 or raw_action_dim > action_dim: - raise ValueError( - f"Cosmos3 raw_action_dim must be in [1, {action_dim}], got {raw_action_dim}." - ) - - clean_action = clean_action.to(device=device, dtype=dtype).unsqueeze(0) - condition_mask = build_action_condition_mask( - mode, - action_chunk_size, - device=device, - dtype=dtype, - ) - noise = randn_tensor( - (1, action_chunk_size, action_dim), - generator=generator, - device=device, - dtype=dtype, - ) - noise[:, :, raw_action_dim:] = 0 - clean_action[:, :, raw_action_dim:] = 0 - action_latents = condition_mask * clean_action + (1.0 - condition_mask) * noise - action_velocity_mask = 1.0 - condition_mask - return action_latents, action_velocity_mask, clean_action, raw_action_dim diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index ae1d4cdfc2bc..3928817951b9 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -15,24 +15,10 @@ """Per-model default generation parameters for Cosmos3 pipelines. Shared by the Cosmos3 OmniMoT text-to-video and image-to-video generation paths. - -Action generation ------------------ -``COSMOS3_DOMAIN_PRESETS`` lists training-aligned defaults per embodiment. When -``domain_name`` (or a uniquely mapped ``domain_id``) is set, the pipeline fills -omitted ``raw_action_dim``, ``action_chunk_size``, ``num_frames``, -``action_resolution``, and ``frame_rate`` from the preset and logs a warning if -explicit values differ. See Cosmos3 omni ``action_*.json`` inputs for reference -configs (bridge, av, droid, libero, etc.). """ -from typing import Any, Dict, List, Optional, TypedDict +from typing import Dict -from tensorrt_llm._torch.visual_gen.models.cosmos3.action import ( - COSMOS3_ACTION_RESOLUTIONS, - EMBODIMENT_TO_DOMAIN_ID, - normalize_action_resolution, -) from tensorrt_llm._torch.visual_gen.pipeline import ExtraParamSchema # --------------------------------------------------------------------------- @@ -55,7 +41,7 @@ # Fields merged by the executor for every request. Modality-specific values # (height/width/num_frames/steps/guidance) are declared with ``None`` so the # executor accepts explicit overrides; ``forward()`` resolves unset fields from -# T2V/T2I/action context. +# T2V/T2I context. COSMOS3_PIPELINE_DEFAULTS = { "height": None, "width": None, @@ -77,232 +63,6 @@ "guidance_interval": (400.0, 1000.0), } -COSMOS3_ACTION_PARAMS = { - "action_chunk_size": 16, - "num_frames": 17, - "num_inference_steps": 30, - "guidance_scale": 1.0, - "flow_shift": 5.0, - "frame_rate": 24.0, -} - - -class Cosmos3DomainPreset(TypedDict, total=False): - """Recommended action-generation settings for a trained embodiment.""" - - raw_action_dim: int - action_chunk_size: int - num_frames: int - action_resolution: int - frame_rate: float - - -# Training-aligned defaults. Values mirror Cosmos3 omni action JSON examples where available. -COSMOS3_DOMAIN_PRESETS: Dict[str, Cosmos3DomainPreset] = { - # WidowX bridge; 7-DOF arm + gripper in 10-D state. - "bridge_orig_lerobot": { - "raw_action_dim": 10, - "action_chunk_size": 16, - "num_frames": 17, - "action_resolution": 480, - "frame_rate": 5.0, - }, - # Autonomous-vehicle steering/throttle; longer action horizon. - "av": { - "raw_action_dim": 9, - "action_chunk_size": 60, - "num_frames": 61, - "action_resolution": 480, - "frame_rate": 10.0, - }, - # 6-DoF camera pose + shutter; matches AV-style horizon. - "camera_pose": { - "raw_action_dim": 9, - "action_chunk_size": 60, - "num_frames": 61, - "action_resolution": 480, - "frame_rate": 30.0, - }, - # Franka single-arm tabletop; same domain_id as robomind-franka. - "droid_lerobot": { - "raw_action_dim": 10, - "action_chunk_size": 16, - "num_frames": 17, - "action_resolution": 480, - "frame_rate": 15.0, - }, - # LIBERO sim single-arm; lower action resolution bucket. - "libero": { - "raw_action_dim": 10, - "action_chunk_size": 16, - "num_frames": 17, - "action_resolution": 256, - "frame_rate": 10.0, - }, - # MANO hand pose; high-DOF hand articulation. - "hand_pose": { - "raw_action_dim": 57, - "action_chunk_size": 16, - "num_frames": 17, - "action_resolution": 480, - "frame_rate": 24.0, - }, - # AgiBot humanoid; shared domain_id with agibot_gear_gripper*. - "agibotworld": { - "raw_action_dim": 29, - "action_chunk_size": 16, - "num_frames": 17, - "action_resolution": 480, - "frame_rate": 10.0, - }, - # Google Robot (RT-1 / fractal) single-arm. - "fractal": { - "raw_action_dim": 10, - "action_chunk_size": 16, - "num_frames": 17, - "action_resolution": 480, - "frame_rate": 5.0, - }, - # 2-D planar push task. - "pusht": { - "raw_action_dim": 2, - "action_chunk_size": 16, - "num_frames": 17, - "action_resolution": 256, - "frame_rate": 10.0, - }, - # UMI handheld gripper setup. - "umi": { - "raw_action_dim": 10, - "action_chunk_size": 16, - "num_frames": 17, - "action_resolution": 480, - "frame_rate": 10.0, - }, -} - -# Map alias domain_name keys to a canonical preset entry. -COSMOS3_DOMAIN_PRESET_ALIASES: Dict[str, str] = { - "robomind-franka": "droid_lerobot", - "robomind-franka-dual": "droid_lerobot", - "robomind-ur": "droid_lerobot", - "agibot_gear_gripper": "agibotworld", - "agibot_gear_gripper_ext": "agibotworld", - "galbot": "agibotworld", -} - - -def canonical_domain_preset_key( - domain_name: Any = None, - domain_id: Any = None, -) -> Optional[str]: - if domain_name is not None and str(domain_name).strip(): - key = str(domain_name).strip().lower() - key = COSMOS3_DOMAIN_PRESET_ALIASES.get(key, key) - if key in COSMOS3_DOMAIN_PRESETS: - return key - return None - - if domain_id is None: - return None - - resolved_id = int(domain_id) - if resolved_id == 0: - return None - - candidates: List[str] = [] - for name, mapped_id in EMBODIMENT_TO_DOMAIN_ID.items(): - if mapped_id != resolved_id: - continue - canon = COSMOS3_DOMAIN_PRESET_ALIASES.get(name, name) - if canon in COSMOS3_DOMAIN_PRESETS and canon not in candidates: - candidates.append(canon) - - if len(candidates) == 1: - return candidates[0] - return None - - -def get_domain_preset( - domain_name: Any = None, - domain_id: Any = None, -) -> Optional[Cosmos3DomainPreset]: - key = canonical_domain_preset_key(domain_name, domain_id) - if key is None: - return None - return COSMOS3_DOMAIN_PRESETS[key] - - -def resolve_domain_action_config( - *, - domain_name: Any = None, - domain_id: Any = None, - raw_action_dim: Optional[int] = None, - action_chunk_size: Optional[int] = None, - action_resolution: Optional[int] = None, - frame_rate: Optional[float] = None, - action_fps: Optional[float] = None, - num_frames: Optional[int] = None, -) -> Dict[str, Any]: - """Merge user action params with domain presets and generic fallbacks.""" - preset_key = canonical_domain_preset_key(domain_name, domain_id) - preset = COSMOS3_DOMAIN_PRESETS.get(preset_key) if preset_key else None - warnings: List[str] = [] - - def _resolve_field( - field: str, - current: Any, - *, - fallback: Any = None, - ) -> Any: - recommended = preset.get(field) if preset else None - if current is not None: - if recommended is not None and current != recommended: - warnings.append( - f"Cosmos3 {field}={current} differs from recommended " - f"{recommended} for domain {preset_key!r}." - ) - return current - if recommended is not None: - return recommended - return fallback - - resolved_raw_action_dim = _resolve_field("raw_action_dim", raw_action_dim) - resolved_chunk = _resolve_field( - "action_chunk_size", - action_chunk_size, - fallback=COSMOS3_ACTION_PARAMS["action_chunk_size"], - ) - resolved_resolution = normalize_action_resolution( - _resolve_field( - "action_resolution", - action_resolution, - fallback=480, - ) - ) - resolved_frame_rate = _resolve_field( - "frame_rate", - frame_rate, - fallback=COSMOS3_ACTION_PARAMS["frame_rate"], - ) - resolved_num_frames = _resolve_field("num_frames", num_frames) - if resolved_num_frames is None: - resolved_num_frames = int(resolved_chunk) + 1 - resolved_action_fps = ( - float(action_fps) if action_fps is not None else float(resolved_frame_rate) - ) - - return { - "raw_action_dim": resolved_raw_action_dim, - "action_chunk_size": int(resolved_chunk), - "action_resolution": resolved_resolution, - "frame_rate": float(resolved_frame_rate), - "action_fps": resolved_action_fps, - "num_frames": int(resolved_num_frames), - "preset_key": preset_key, - "warnings": warnings, - } - COSMOS3_EXTRA_SPECS: Dict[str, ExtraParamSchema] = { "use_duration_template": ExtraParamSchema( @@ -335,62 +95,6 @@ def _resolve_field( default="video", description="Output modality: 'video' (T2V/I2V) or 'image' (text-to-image).", ), - "action_mode": ExtraParamSchema( - type="str", - default=None, - description="Action generation mode: policy, forward_dynamics, or inverse_dynamics.", - ), - "domain_name": ExtraParamSchema( - type="str", - default=None, - description=( - "Embodiment domain name for action generation (e.g. bridge_orig_lerobot, av). " - "When set, omitted raw_action_dim/action_chunk_size/action_resolution/frame_rate " - "are filled from COSMOS3_DOMAIN_PRESETS; mismatches are logged as warnings." - ), - ), - "domain_id": ExtraParamSchema( - type="int", - default=None, - description="Embodiment domain id for action generation.", - ), - "raw_action_dim": ExtraParamSchema( - type="int", - default=None, - description=( - "Raw action DOF for policy/inverse_dynamics (e.g. 10 bridge, 9 av, 29 agibot). " - "Inferred from domain_name preset when omitted." - ), - ), - "action_chunk_size": ExtraParamSchema( - type="int", - default=COSMOS3_ACTION_PARAMS["action_chunk_size"], - description=( - "Number of action tokens to generate (16 for most robots, 60 for av/camera_pose). " - "Inferred from domain_name preset when omitted." - ), - ), - "action": ExtraParamSchema( - type="list", - default=None, - description="Action trajectory [T, D] for forward_dynamics mode.", - ), - "action_resolution": ExtraParamSchema( - type="int", - default=480, - description=( - "Resolution bucket for action image sizing. Must be one of " - f"{list(COSMOS3_ACTION_RESOLUTIONS)}. Inferred from domain_name preset when omitted." - ), - range=(min(COSMOS3_ACTION_RESOLUTIONS), max(COSMOS3_ACTION_RESOLUTIONS)), - ), - "action_fps": ExtraParamSchema( - type="float", - default=None, - description=( - "Action-token temporal rate for mRoPE (Hz). Defaults to frame_rate when omitted." - ), - ), "condition_frame_indexes_vision": ExtraParamSchema( type="list", default=list(COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION), @@ -410,7 +114,7 @@ def _resolve_field( type="path_or_list", default=None, description=( - "Video input for video-to-video generation or inverse_dynamics: " + "Video input for video-to-video generation: " ".mp4/.avi file, frame directory, image path, or list of PIL " "images / frame paths." ), diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index e7192b12db05..cd982b837c2d 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -34,26 +34,13 @@ from tensorrt_llm.inputs.utils import load_image from tensorrt_llm.logger import logger -from .action import ( - ACTION_MODE_INVERSE_DYNAMICS, - action_reference_image, - action_start_frame_offset, - build_vision_condition_mask, - normalize_action_mode, - prepare_action_latents, - resize_and_pad_action_image, - resolve_action_size, - resolve_domain_id, -) from .defaults import ( COSMOS3_720P_PARAMS, - COSMOS3_ACTION_PARAMS, COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION, COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP, COSMOS3_EXTRA_SPECS, COSMOS3_PIPELINE_DEFAULTS, COSMOS3_T2I_PARAMS, - resolve_domain_action_config, ) from .guardrails import check_video_safety, download_guardrail_checkpoint from .sound_tokenizer import LatentAutoEncoderV2 @@ -125,7 +112,6 @@ class Cosmos3OmniMoTPipeline(BasePipeline): def __init__(self, pipeline_config): primary_pretrained_config = pipeline_config.primary_pretrained_config self.audio_gen = False - self.action_gen = False if getattr( primary_pretrained_config, "audio_gen", @@ -134,10 +120,6 @@ def __init__(self, pipeline_config): logger.info("Initializing Cosmos3OmniMoTPipeline with audio generation.") self.audio_gen = True - if getattr(primary_pretrained_config, "action_gen", False): - logger.info("Initializing Cosmos3OmniMoTPipeline with action generation.") - self.action_gen = True - super().__init__(pipeline_config) def _init_transformer(self) -> None: @@ -218,10 +200,6 @@ def load_standard_components( # (UniPC mutates internal correction buffers on every .step() call). self.audio_scheduler = UniPCMultistepScheduler.from_config(self.scheduler.config) - if self.action_gen: - # Action uses its own scheduler for the same reason as audio. - self.action_scheduler = UniPCMultistepScheduler.from_config(self.scheduler.config) - # Re-check the env var in case it was changed after initialization like in unit tests. guardrails_disabled = os.environ.get("TRTLLM_DISABLE_COSMOS3_GUARDRAILS", "0") == "1" global TRTLLM_DISABLE_COSMOS3_GUARDRAILS @@ -262,7 +240,7 @@ def _set_flow_shift( """Rebuild the UniPC scheduler when request scheduler defaults change. The effective flow-shift changes when switching between mode defaults - (T2I=3.0, action=5.0, V2V=10.0, T2V/I2V=checkpoint default) or when a + (T2I=3.0, V2V=10.0, T2V/I2V=checkpoint default) or when a request provides ``flow_shift``. V2V also forces Karras sigmas off. """ if not hasattr(self, "_base_scheduler_config"): @@ -289,10 +267,6 @@ def _set_flow_shift( self.audio_scheduler = UniPCMultistepScheduler.from_config( self._base_scheduler_config, **scheduler_kwargs ) - if self.action_gen: - self.action_scheduler = UniPCMultistepScheduler.from_config( - self._base_scheduler_config, **scheduler_kwargs - ) self._current_flow_shift = target self._current_scheduler_use_karras_sigmas = self._scheduler_use_karras_sigmas( self.scheduler.config @@ -361,15 +335,6 @@ def infer(self, req): use_guardrails=extra_params.get("use_guardrails", True), enable_audio=extra_params.get("enable_audio", False), output_type=output_type, - action_mode=extra_params.get("action_mode"), - domain_name=extra_params.get("domain_name"), - domain_id=extra_params.get("domain_id"), - raw_action_dim=extra_params.get("raw_action_dim"), - action_chunk_size=extra_params.get("action_chunk_size"), - action=extra_params.get("action"), - action_resolution=extra_params.get("action_resolution") - or extra_params.get("image_size"), - action_fps=extra_params.get("action_fps"), video=extra_params.get("video"), condition_frame_indexes_vision=extra_params.get("condition_frame_indexes_vision"), condition_video_keep=extra_params.get("condition_video_keep"), @@ -687,27 +652,6 @@ def decode_audio(self, latent: torch.Tensor) -> torch.Tensor: """ return self.audio_tokenizer.decode(latent).float() # [B, audio_channels, N_samples] - # ========================================================================= - # Action generation - # ========================================================================= - - def _preprocess_action_image( - self, image: PIL.Image.Image, target_h: int, target_w: int - ) -> torch.Tensor: - image = resize_and_pad_action_image(image, target_h, target_w) - return self.video_processor.preprocess(image, height=target_h, width=target_w) - - def _preprocess_action_video( - self, frames: List[Any], target_h: int, target_w: int - ) -> torch.Tensor: - if not frames: - raise ValueError("Cosmos3 action video input must contain at least one frame.") - processed = [ - self._preprocess_action_image(pil_to_rgb(frame), target_h, target_w).squeeze(0) - for frame in frames - ] - return torch.stack(processed, dim=1).unsqueeze(0).contiguous() - def _preprocess_condition_video( self, frames: List[Any], target_h: int, target_w: int ) -> torch.Tensor: @@ -753,61 +697,6 @@ def _encode_video_tensor(self, video_tensor: torch.Tensor) -> torch.Tensor: return latent.to(self.dtype) - def _prepare_latents_action_video( - self, - video_tensor: torch.Tensor, - mode: str, - num_frames: int, - generator: torch.Generator, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - C = self.transformer.latent_channel_size - T_lat = (num_frames - 1) // self.vae_scale_factor_temporal + 1 - H_lat = video_tensor.shape[-2] // self.vae_scale_factor_spatial - W_lat = video_tensor.shape[-1] // self.vae_scale_factor_spatial - - noise = randn_tensor( - (1, C, T_lat, H_lat, W_lat), - generator=generator, - device=self.device, - dtype=self.dtype, - ) - cond_latent = self._encode_video_tensor(video_tensor) - if cond_latent.shape[2:] != noise.shape[2:]: - raise ValueError( - "Cosmos3 action video latent shape mismatch: " - f"encoded={tuple(cond_latent.shape)}, expected={tuple(noise.shape)}." - ) - condition_mask = build_vision_condition_mask( - mode, - num_frames, - self.vae_scale_factor_temporal, - device=self.device, - dtype=self.dtype, - ) - latents = condition_mask * cond_latent + (1.0 - condition_mask) * noise - velocity_mask = 1.0 - condition_mask - return latents, velocity_mask, cond_latent - - def _prepare_action_latents( - self, - *, - mode: str, - action_chunk_size: int, - raw_action_dim: Optional[int], - generator: torch.Generator, - action_input: Any = None, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: - return prepare_action_latents( - mode=mode, - action_chunk_size=action_chunk_size, - raw_action_dim=raw_action_dim, - action_dim=int(getattr(self.transformer, "action_dim", 64)), - generator=generator, - device=self.device, - dtype=self.dtype, - action_input=action_input, - ) - # ========================================================================= # Video to video # ========================================================================= @@ -906,14 +795,6 @@ def forward( use_guardrails: bool = COSMOS3_EXTRA_SPECS["use_guardrails"].default, enable_audio: bool = COSMOS3_EXTRA_SPECS["enable_audio"].default, output_type: str = COSMOS3_EXTRA_SPECS["output_type"].default, - action_mode: Optional[str] = None, - domain_name: Optional[str] = None, - domain_id: Optional[int] = None, - raw_action_dim: Optional[int] = None, - action_chunk_size: Optional[int] = None, - action: Any = None, - action_resolution: Optional[int] = None, - action_fps: Optional[float] = None, video: Any = None, condition_frame_indexes_vision: Any = None, condition_video_keep: Any = None, @@ -925,43 +806,30 @@ def forward( use_guardrails = use_guardrails and not TRTLLM_DISABLE_COSMOS3_GUARDRAILS - normalized_action_mode = normalize_action_mode(action_mode) - do_action = normalized_action_mode is not None - if do_action and not self.action_gen: - raise ValueError( - "Cosmos3 action generation was requested, but this checkpoint " - "does not enable action_gen." - ) - if do_action and enable_audio: - raise ValueError("Cosmos3 does not support joint action and audio generation.") - # Text-to-image mode: same checkpoint/forward path as T2V, but a single # latent frame, image-flavored prompt templates, flow_shift=3.0, a CFG # guidance interval, and an image (rather than video) output. is_t2i = str(output_type).lower() == "image" - if not do_action and image is not None and video is not None: + if image is not None and video is not None: raise ValueError( - "Cosmos3 non-action generation supports text-only, text + image, " + "Cosmos3 generation supports text-only, text + image, " "or text + video input, but not both image and video." ) if is_t2i and video is not None: raise ValueError( "Cosmos3 video-to-video generation is supported only for video outputs." ) - is_v2v = video is not None and not is_t2i and not do_action + is_v2v = video is not None and not is_t2i if use_system_prompt is None: use_system_prompt = is_v2v else: use_system_prompt = bool(use_system_prompt) guidance_interval = None - resolved_action_fps: Optional[float] = None if is_t2i: if image is not None: raise ValueError( "Cosmos3 text-to-image (output_type='image') does not accept an image input." ) - if do_action: - raise ValueError("Cosmos3 action generation does not support output_type='image'.") if enable_audio: raise ValueError("Cosmos3 audio generation does not support output_type='image'.") num_frames = 1 @@ -974,46 +842,6 @@ def forward( self._set_flow_shift( flow_shift if flow_shift is not None else COSMOS3_T2I_PARAMS["flow_shift"] ) - elif do_action: - action_cfg = resolve_domain_action_config( - domain_name=domain_name, - domain_id=domain_id, - raw_action_dim=raw_action_dim, - action_chunk_size=action_chunk_size, - action_resolution=action_resolution, - frame_rate=frame_rate, - action_fps=action_fps, - num_frames=num_frames, - ) - if self.rank == 0: - for warning in action_cfg["warnings"]: - logger.warning(warning) - if action_cfg["preset_key"] is not None: - logger.info( - f"Cosmos3 action domain preset {action_cfg['preset_key']!r}: " - f"raw_action_dim={action_cfg['raw_action_dim']}, " - f"action_chunk_size={action_cfg['action_chunk_size']}, " - f"action_resolution={action_cfg['action_resolution']}, " - f"frame_rate={action_cfg['frame_rate']:.1f}, " - f"action_fps={action_cfg['action_fps']:.1f}, " - f"num_frames={action_cfg['num_frames']}" - ) - - raw_action_dim = action_cfg["raw_action_dim"] - action_chunk_size = action_cfg["action_chunk_size"] - action_resolution = action_cfg["action_resolution"] - num_frames = action_cfg["num_frames"] - frame_rate = action_cfg["frame_rate"] - resolved_action_fps = action_cfg["action_fps"] - num_inference_steps = ( - num_inference_steps or COSMOS3_ACTION_PARAMS["num_inference_steps"] - ) - if guidance_scale is None: - guidance_scale = COSMOS3_ACTION_PARAMS["guidance_scale"] - self._set_flow_shift( - flow_shift if flow_shift is not None else COSMOS3_ACTION_PARAMS["flow_shift"] - ) - enable_audio = False else: height = height or COSMOS3_720P_PARAMS["height"] width = width or COSMOS3_720P_PARAMS["width"] @@ -1038,17 +866,6 @@ def forward( max_sequence_length = max_sequence_length or COSMOS3_720P_PARAMS["max_sequence_length"] if frame_rate is None: frame_rate = COSMOS3_720P_PARAMS["frame_rate"] - if resolved_action_fps is None: - resolved_action_fps = frame_rate - - action_ref_image = None - if do_action: - action_ref_image = action_reference_image( - action_mode=normalized_action_mode, - image=image, - video=video, - ) - height, width = resolve_action_size(height, width, action_ref_image, action_resolution) if self.rank == 0: logger.info( @@ -1056,13 +873,6 @@ def forward( f"num_inference_steps={num_inference_steps}, guidance_scale={guidance_scale:.2f}, " f"frame_rate={frame_rate:.1f}" ) - if do_action: - logger.info( - f"Cosmos3 action dims: action_chunk_size={action_chunk_size}, " - f"action_resolution={action_resolution}, " - f"action_fps={resolved_action_fps:.1f}, " - f"input_aspect={action_ref_image.width / action_ref_image.height:.3f}" - ) if isinstance(prompt, str): prompt = [prompt] @@ -1158,68 +968,11 @@ def forward( ) # 2. Prepare latents - action_latents = None - action_velocity_mask = None - action_condition_latents = None - action_domain_id = None - action_frame_offset = 1 - resolved_raw_action_dim = raw_action_dim condition_latents = None image_latent = None velocity_mask = None - if do_action: - if action_chunk_size not in {num_frames, num_frames - 1}: - raise ValueError( - "Cosmos3 num_frames must equal action_chunk_size or action_chunk_size + 1." - ) - action_domain_id = resolve_domain_id( - domain_id=domain_id, - domain_name=domain_name, - require_explicit=True, - ) - action_frame_offset = action_start_frame_offset( - normalized_action_mode, action_chunk_size, num_frames - ) - - if normalized_action_mode == ACTION_MODE_INVERSE_DYNAMICS: - inverse_video = video if video is not None else image - video = normalize_video_input(inverse_video, max_frames=num_frames) - video_tensor = self._preprocess_action_video(video, height, width) - latents, velocity_mask, condition_latents = self._prepare_latents_action_video( - video_tensor, - normalized_action_mode, - num_frames, - generator, - ) - else: - image_tensor = self._preprocess_action_image(action_ref_image, height, width) - if image_tensor.ndim == 4: - video_tensor = ( - image_tensor.unsqueeze(2).expand(-1, -1, num_frames, -1, -1).contiguous() - ) - else: - video_tensor = image_tensor - latents, velocity_mask, condition_latents = self._prepare_latents_action_video( - video_tensor, - normalized_action_mode, - num_frames, - generator, - ) - - ( - action_latents, - action_velocity_mask, - action_condition_latents, - resolved_raw_action_dim, - ) = self._prepare_action_latents( - mode=normalized_action_mode, - action_chunk_size=action_chunk_size, - raw_action_dim=raw_action_dim, - generator=generator, - action_input=action, - ) - elif image is not None: + if image is not None: if isinstance(image, str): image = load_image(image, format="pil") @@ -1279,8 +1032,6 @@ def forward( # 3. Set up scheduler self.scheduler.set_timesteps(num_inference_steps, device=self.device) - if do_action: - self.action_scheduler.set_timesteps(num_inference_steps, device=self.device) # 3b. Audio noise init — latent length matches diffusers Cosmos3OmniPipeline.prepare_latents. do_audio = enable_audio and self.audio_gen and hasattr(self, "audio_tokenizer") @@ -1300,12 +1051,6 @@ def forward( self.audio_scheduler.set_timesteps(num_inference_steps, device=self.device) # 4. Build forward_fn for the denoise loop - action_domain_ids_tensor = None - if do_action and action_domain_id is not None: - action_domain_ids_tensor = torch.tensor( - [action_domain_id], dtype=torch.long, device=self.device - ) - def forward_fn( latent_input, extra_stream_latents, @@ -1320,16 +1065,6 @@ def forward_fn( rather than through encoder_hidden_states. """ current_audio = extra_stream_latents.get("audio") if extra_stream_latents else None - current_action = extra_stream_latents.get("action") if extra_stream_latents else None - - action_domain_ids = action_domain_ids_tensor - if ( - action_domain_ids is not None - and current_action is not None - and action_domain_ids.shape[0] == 1 - and current_action.shape[0] > 1 - ): - action_domain_ids = action_domain_ids.expand(current_action.shape[0]) result = self.transformer( hidden_states=latent_input, @@ -1341,31 +1076,14 @@ def forward_fn( fps=frame_rate, noisy_frame_mask=velocity_mask, audio_latents=current_audio, - action_latents=current_action, - action_domain_ids=action_domain_ids, - action_noisy_mask=action_velocity_mask, - action_start_frame_offset=action_frame_offset, - action_fps=resolved_action_fps, ) video_noise_pred = result.video audio_noise_pred = result.audio - action_noise_pred = result.action if velocity_mask is not None: video_noise_pred = video_noise_pred * velocity_mask - if action_noise_pred is not None: - if action_velocity_mask is not None: - action_noise_pred = action_noise_pred * action_velocity_mask - if ( - resolved_raw_action_dim is not None - and 0 < resolved_raw_action_dim < action_noise_pred.shape[-1] - ): - action_noise_pred = action_noise_pred.clone() - action_noise_pred[..., resolved_raw_action_dim:] = 0 - return video_noise_pred, {"action": action_noise_pred} - if audio_noise_pred is not None: return video_noise_pred, {"audio": audio_noise_pred} return video_noise_pred @@ -1380,17 +1098,6 @@ def post_step_fn(step_latents, step_extra_stream_latents): step_latents[:, :, 0:1, :, :] = image_latent.to( device=step_latents.device, dtype=step_latents.dtype ) - if ( - action_velocity_mask is not None - and action_condition_latents is not None - and step_extra_stream_latents is not None - and "action" in step_extra_stream_latents - ): - action_key = step_extra_stream_latents["action"] - step_extra_stream_latents["action"] = ( - action_velocity_mask * action_key - + (1.0 - action_velocity_mask) * action_condition_latents - ) return step_latents, step_extra_stream_latents # 5. Build CFG tensors — text_ids and text_mask need to be split for CFG @@ -1406,14 +1113,9 @@ def post_step_fn(step_latents, step_extra_stream_latents): # 6. Denoise timer.mark_denoise_start() extra_streams = None - if do_action: - extra_streams = {"action": (action_latents, self.action_scheduler)} - elif do_audio: + if do_audio: extra_streams = {"audio": (audio_latents, self.audio_scheduler)} - should_pin_condition_latents = ( - do_action or condition_latents is not None or image_latent is not None - ) - # FUTURE(action+audio): merge both keys; extend forward_fn return dict and post_step_fn. + should_pin_condition_latents = condition_latents is not None or image_latent is not None denoise_result = self.denoise( latents=latents, scheduler=self.scheduler, @@ -1430,8 +1132,6 @@ def post_step_fn(step_latents, step_extra_stream_latents): if extra_streams is not None: latents, extra_latents = denoise_result audio_latents = extra_latents.get("audio") - if do_action: - action_latents = extra_latents.get("action") else: latents = denoise_result audio_latents = None @@ -1477,11 +1177,5 @@ def post_step_fn(step_latents, step_extra_stream_latents): audio_sample_rate=self.audio_tokenizer.model_config["sampling_rate"] if waveform is not None else None, - action=action_latents[:, :, :resolved_raw_action_dim].float().cpu() - if do_action and action_latents is not None - else None, - raw_action_dim=resolved_raw_action_dim if do_action else None, - action_mode=normalized_action_mode if do_action else None, - domain_id=action_domain_id if do_action else None, ) ) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py index 7740efc1a24b..8deedeed044a 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -72,9 +72,6 @@ class TransformerOutput: audio: Optional[torch.Tensor] = None """[B, audio_dim, T_audio] audio velocity prediction, or None.""" - action: Optional[torch.Tensor] = None - """[B, T_action, action_dim] action velocity prediction, or None.""" - def compute_mrope_position_ids_text( num_tokens: int, @@ -152,78 +149,6 @@ def compute_mrope_position_ids_vision( return mrope_ids, next_offset -def compute_mrope_position_ids_action( - grid_t: int, - temporal_offset: int | float, - action_fps: float | None, - base_fps: float = 24.0, - base_temporal_compression_factor: int = 4, - enable_fps_modulation: bool = True, - start_frame_offset: int = 1, -) -> tuple[torch.Tensor, int | float]: - """Generate mRoPE IDs for action tokens as a frame-rate (T, 1, 1) grid.""" - return compute_mrope_position_ids_vision( - grid_t=grid_t, - grid_h=1, - grid_w=1, - temporal_offset=temporal_offset, - fps=action_fps, - base_fps=base_fps, - temporal_compression_factor=1, - enable_fps_modulation=enable_fps_modulation, - start_frame_offset=start_frame_offset, - ) - - -class DomainAwareLinear(nn.Module): - """Linear projection with one weight/bias pair per action embodiment domain.""" - - def __init__( - self, - input_size: int, - output_size: int, - num_domains: int, - *, - dtype: torch.dtype = torch.bfloat16, - ) -> None: - super().__init__() - self.input_size = int(input_size) - self.output_size = int(output_size) - self.num_domains = int(num_domains) - self.dtype = dtype - self.fc = nn.Embedding(self.num_domains, self.output_size * self.input_size, dtype=dtype) - self.bias = nn.Embedding(self.num_domains, self.output_size, dtype=dtype) - - def post_load_weights(self) -> None: - self.fc.to(self.dtype) - self.bias.to(self.dtype) - - def forward(self, x: torch.Tensor, domain_id: torch.Tensor) -> torch.Tensor: - if domain_id.ndim == 0: - domain_id = domain_id.unsqueeze(0) - domain_id = domain_id.to(device=x.device, dtype=torch.long).reshape(-1) - if x.shape[0] != domain_id.shape[0]: - raise ValueError( - "Cosmos3 action domain_id batch size must match action batch: " - f"tokens={x.shape[0]}, domain_id={domain_id.shape[0]}." - ) - if torch.any((domain_id < 0) | (domain_id >= self.num_domains)): - raise ValueError( - f"Cosmos3 action domain_id must be in [0, {self.num_domains}), " - f"got {domain_id.tolist()}." - ) - - weight = self.fc(domain_id).view(domain_id.shape[0], self.input_size, self.output_size) - bias = self.bias(domain_id).view(domain_id.shape[0], self.output_size) - if x.ndim == 2: - return torch.bmm(x.unsqueeze(1), weight).squeeze(1) + bias - if x.ndim == 3: - return torch.bmm(x, weight) + bias.unsqueeze(1) - raise ValueError( - f"Cosmos3 DomainAwareLinear expected rank-2 or rank-3 input, got {tuple(x.shape)}." - ) - - class TimestepEmbedder(nn.Module): """ Embeds scalar timesteps into vector representations. @@ -779,7 +704,6 @@ def __init__(self, model_config: DiffusionModelConfig): super().__init__(model_config) pretrained_config = model_config.pretrained_config self.audio_gen = getattr(pretrained_config, "sound_gen", False) - self.action_gen = getattr(pretrained_config, "action_gen", False) self.hidden_size = pretrained_config.hidden_size self.num_hidden_layers = pretrained_config.num_hidden_layers @@ -806,29 +730,6 @@ def __init__(self, model_config: DiffusionModelConfig): pretrained_config.temporal_compression_factor_sound ) - if self.action_gen: - action_dim_value = getattr(pretrained_config, "action_dim", None) - if action_dim_value is None: - action_dim_value = getattr(pretrained_config, "max_action_dim", 64) - self.action_dim = int(action_dim_value) - self.num_embodiment_domains = int( - getattr(pretrained_config, "num_embodiment_domains", 32) - ) - dtype = torch.bfloat16 - self.action_proj_in = DomainAwareLinear( - self.action_dim, - self.hidden_size, - self.num_embodiment_domains, - dtype=dtype, - ) - self.action_proj_out = DomainAwareLinear( - self.hidden_size, - self.action_dim, - self.num_embodiment_domains, - dtype=dtype, - ) - self.action_modality_embed = nn.Parameter(torch.zeros(self.hidden_size, dtype=dtype)) - if pretrained_config.position_embedding_type != "unified_3d_mrope": raise ValueError( f"Position embedding type {pretrained_config.position_embedding_type} not supported" @@ -1062,59 +963,6 @@ def unpack_audio_latents(self, hidden_audio: torch.Tensor) -> torch.Tensor: """[B, T_audio, audio_dim] → [B, audio_dim, T_audio].""" return hidden_audio.permute(0, 2, 1) - # ------------------------------------------------------------------------- - # Action helpers - # ------------------------------------------------------------------------- - - def _compute_action_rope_freqs( - self, - T_action: int, - text_mask: torch.Tensor, - action_fps: float, - action_start_frame_offset: int, - device: torch.device, - dtype: torch.dtype, - ) -> Tuple[torch.Tensor, torch.Tensor]: - B = text_mask.shape[0] - text_lengths = text_mask.sum(dim=1).long() - - action_pos_list = [] - for b in range(B): - real_len = int(text_lengths[b].item()) - _, t_offset = compute_mrope_position_ids_text(real_len, temporal_offset=0) - a_pos, _ = compute_mrope_position_ids_action( - T_action, - temporal_offset=t_offset + self.unified_3d_mrope_temporal_modality_margin, - action_fps=action_fps, - base_fps=self.base_fps, - base_temporal_compression_factor=self.temporal_compression_factor, - enable_fps_modulation=self.enable_fps_modulation, - start_frame_offset=action_start_frame_offset, - ) - action_pos_list.append(a_pos) - - action_pos_ids = torch.stack(action_pos_list, dim=1).to(device) - rotary_emb = self.language_model.rotary_emb - _dummy = torch.tensor([], dtype=dtype, device=device) - cos_a, sin_a = rotary_emb(_dummy, position_ids=action_pos_ids) - return cos_a.unsqueeze(2), sin_a.unsqueeze(2) - - def pack_action(self, action_latents: torch.Tensor) -> torch.Tensor: - if action_latents.ndim != 3: - raise ValueError( - f"Cosmos3 action latents must have shape [B, T, D], got {tuple(action_latents.shape)}." - ) - if action_latents.shape[-1] != self.action_dim: - raise ValueError( - f"Cosmos3 action latent dimension mismatch: expected {self.action_dim}, " - f"got {action_latents.shape[-1]}." - ) - return action_latents.contiguous() - - @staticmethod - def unpack_action(tokens: torch.Tensor) -> torch.Tensor: - return tokens - def reset_cache(self): self.cached_kv = None self.cached_freqs_gen = None @@ -1130,11 +978,6 @@ def forward( fps: float | None = None, noisy_frame_mask: torch.Tensor | None = None, audio_latents: Optional[torch.Tensor] = None, - action_latents: Optional[torch.Tensor] = None, - action_domain_ids: Optional[torch.Tensor] = None, - action_noisy_mask: Optional[torch.Tensor] = None, - action_start_frame_offset: int = 1, - action_fps: float | None = None, **kwargs, ) -> "TransformerOutput": """ @@ -1162,7 +1005,7 @@ def forward( Returns: TransformerOutput with video (and image alias) always set. audio is set to the predicted audio velocity when audio_latents is - provided; otherwise None. action is set when action_latents is provided. + provided; otherwise None. """ del kwargs # Kept for diffusers API compatibility. if timestep is None: @@ -1170,15 +1013,6 @@ def forward( if raw_timestep is None: raise ValueError("Cosmos3VFMTransformer.forward requires raw_timestep.") - if action_latents is not None and audio_latents is not None: - raise ValueError( - "Cosmos3 transformer does not support joint action and audio generation." - ) - if action_latents is not None and not self.action_gen: - raise ValueError( - "Cosmos3 action generation was requested, but this transformer " - "was initialized without action modules." - ) T, H, W = video_shape Hp, Wp, _, _ = self._pad_to_patch_size(H, W) max_real_len = text_mask.sum(dim=1).max().item() @@ -1241,44 +1075,10 @@ def forward( else: self.cached_kv = cached_kv_full - # --- Extra modality token injection (mutually exclusive: action OR audio) --- + # --- Extra modality token injection (audio) --- T_vid_tokens = hidden_gen.shape[1] # T * Hp * Wp - T_action = 0 T_audio = 0 - action_domain_ids_tensor = action_domain_ids - if action_latents is not None and self.action_gen: - # FUTURE(action+audio): concat order is video|action|audio; adjust slices below. - if action_domain_ids_tensor is None: - action_domain_ids_tensor = torch.zeros( - action_latents.shape[0], dtype=torch.long, device=action_latents.device - ) - T_action = action_latents.shape[1] - hidden_action = self.action_proj_in( - self.pack_action(action_latents), action_domain_ids_tensor - ) - hidden_action = hidden_action + self.action_modality_embed.to(hidden_action.dtype) - if action_noisy_mask is None: - hidden_action = hidden_action + time_embed.unsqueeze(1) - else: - hidden_action = hidden_action + time_embed.unsqueeze(1) * action_noisy_mask.to( - hidden_action.dtype - ) - hidden_gen = torch.cat([hidden_gen, hidden_action], dim=1) - effective_action_fps = action_fps if action_fps is not None else (fps or self.base_fps) - cos_a, sin_a = self._compute_action_rope_freqs( - T_action, - text_mask, - float(effective_action_fps), - action_start_frame_offset, - hidden_states.device, - hidden_gen.dtype, - ) - cos_v, sin_v = self.cached_freqs_gen - freqs_gen_combined = ( - torch.cat([cos_v, cos_a], dim=1), - torch.cat([sin_v, sin_a], dim=1), - ) - elif audio_latents is not None and self.audio_gen: + if audio_latents is not None and self.audio_gen: T_audio = audio_latents.shape[2] hidden_audio = self.pack_audio_latents(audio_latents).to(hidden_gen.dtype) hidden_audio = self.audio2llm(hidden_audio) + self.audio_modality_embed @@ -1336,7 +1136,7 @@ def forward( # --- Decode video velocity ------------------------------------------------ video_vel = self.unpatchify(self.llm2vae(hidden_gen[:, :T_vid_tokens]), T, H, W) - # --- Decode extra-modality velocity (action XOR audio; follows video) --- + # --- Decode extra-modality velocity (audio; follows video) --- extra_start = T_vid_tokens audio_vel = None if T_audio > 0 and audio_latents is not None and self.audio_gen: @@ -1344,19 +1144,7 @@ def forward( self.llm2audio(hidden_gen[:, extra_start : extra_start + T_audio]) ) - action_vel = None - if T_action > 0 and action_latents is not None and self.action_gen: - assert action_domain_ids_tensor is not None - action_vel = self.unpack_action( - self.action_proj_out( - hidden_gen[:, extra_start : extra_start + T_action], - action_domain_ids_tensor, - ) - ) - - return TransformerOutput( - video=video_vel, image=video_vel, audio=audio_vel, action=action_vel - ) + return TransformerOutput(video=video_vel, image=video_vel, audio=audio_vel) def load_weights(self, weights: dict) -> None: """Load weights with key remapping from Cosmos3-Nano / Diffusers checkpoints. @@ -1366,10 +1154,7 @@ def load_weights(self, weights: dict) -> None: Maps UND vs GEN blocks into this module's layout (causal self-attn vs cross-attn + MLPs). """ remapped = {} - skip_prefixes = ( - "lm_head.", - "action_pos_embed.", - ) + skip_prefixes = ("lm_head.",) for key, value in weights.items(): k = key @@ -1403,14 +1188,6 @@ def load_weights(self, weights: dict) -> None: remapped[k] = value continue - if k.startswith("action_modality_embed"): - remapped[k] = value - continue - - if k.startswith("action_proj_in.") or k.startswith("action_proj_out."): - remapped[k] = value - continue - if k.startswith("time_embedder.linear"): k = k.replace("time_embedder.linear_1.", "time_embedder.mlp.linear_1.") k = k.replace("time_embedder.linear_2.", "time_embedder.mlp.linear_2.") @@ -1543,11 +1320,6 @@ def post_load_weights(self) -> None: self.llm2audio.to(target_dtype) self.audio_modality_embed.data = self.audio_modality_embed.data.to(target_dtype) - if self.action_gen: - self.action_modality_embed.data = self.action_modality_embed.data.to(target_dtype) - self.action_proj_in.post_load_weights() - self.action_proj_out.post_load_weights() - for _, module in self.named_modules(): if isinstance(module, Linear) or isinstance(module, Qwen3VLTextRMSNorm): module.post_load_weights() diff --git a/tensorrt_llm/_torch/visual_gen/output.py b/tensorrt_llm/_torch/visual_gen/output.py index ece7fd203f3f..64963edc76be 100644 --- a/tensorrt_llm/_torch/visual_gen/output.py +++ b/tensorrt_llm/_torch/visual_gen/output.py @@ -37,7 +37,7 @@ class PipelineOutput: """Internal per-pipeline output. Each pipeline ``infer()`` populates the media tensor it produces plus - the metadata it owns (``frame_rate``, ``audio_sample_rate``, action fields) + the metadata it owns (``frame_rate``, ``audio_sample_rate``) and the three CUDA-event-measured timing phases that decompose ``pipeline.infer()``. @@ -52,25 +52,12 @@ class PipelineOutput: ``(B, channels, T_audio)``, dtype ``float32``. Populated by LTX-2. The leading batch dim is always present, even for single-prompt requests (size 1). - action: Predicted or refined action trajectory as ``torch.Tensor`` - shape ``(B, T_action, D_raw)``, dtype ``float32``. Populated by - Cosmos3 action generation (``policy`` / ``forward_dynamics`` / - ``inverse_dynamics``). Sliced to ``raw_action_dim`` DOF; no VAE - decode. ``None`` when action generation was not requested. frame_rate: Video frame rate in fps. Populated by video pipelines (Wan T2V/I2V emit ``16.0``; LTX-2 emits ``params.frame_rate``). ``None`` for image-only pipelines. audio_sample_rate: Audio sample rate in Hz. Populated by LTX-2 from its audio config (no hard-coded literal). ``None`` for pipelines without audio. - raw_action_dim: Number of action degrees of freedom in the trailing - dimension of ``action``. Populated by Cosmos3 action generation. - ``None`` when ``action`` is ``None``. - action_mode: Requested Cosmos3 action mode (``policy``, - ``forward_dynamics``, or ``inverse_dynamics``). ``None`` when - action generation was not requested. - domain_id: Resolved embodiment domain id for Cosmos3 action - generation. ``None`` when action generation was not requested. pre_denoise: Wall-clock GPU-stream time (seconds) before the denoising loop (text encoding, latent prep, conditioning), measured by CUDA events. ``0.0`` if not measured. @@ -86,12 +73,8 @@ class PipelineOutput: image: Optional[torch.Tensor] = None video: Optional[torch.Tensor] = None audio: Optional[torch.Tensor] = None - action: Optional[torch.Tensor] = None frame_rate: Optional[float] = None audio_sample_rate: Optional[int] = None - raw_action_dim: Optional[int] = None - action_mode: Optional[str] = None - domain_id: Optional[int] = None pre_denoise: float = 0.0 denoise: float = 0.0 post_denoise: float = 0.0 @@ -219,12 +202,8 @@ def to_visual_gen_output(resp: "DiffusionResponse") -> "VisualGenOutput": image=out.image, video=out.video, audio=out.audio, - action=out.action, frame_rate=out.frame_rate, audio_sample_rate=out.audio_sample_rate, - raw_action_dim=out.raw_action_dim, - action_mode=out.action_mode, - domain_id=out.domain_id, metrics=metrics, ) @@ -267,10 +246,6 @@ def split_visual_gen_output(resp: "DiffusionResponse", batch_size: int) -> List[ assert out.audio.shape[0] == batch_size, ( f"audio leading dim {out.audio.shape[0]} != batch_size {batch_size}" ) - if out.action is not None: - assert out.action.shape[0] == batch_size, ( - f"action leading dim {out.action.shape[0]} != batch_size {batch_size}" - ) metrics = VisualGenMetrics( generation=resp.generation, pre_denoise=out.pre_denoise, @@ -285,12 +260,8 @@ def split_visual_gen_output(resp: "DiffusionResponse", batch_size: int) -> List[ image=out.image[i] if out.image is not None else None, video=out.video[i] if out.video is not None else None, audio=out.audio[i] if out.audio is not None else None, - action=out.action[i] if out.action is not None else None, frame_rate=out.frame_rate, audio_sample_rate=out.audio_sample_rate, - raw_action_dim=out.raw_action_dim, - action_mode=out.action_mode, - domain_id=out.domain_id, metrics=metrics, ) ) diff --git a/tensorrt_llm/visual_gen/output.py b/tensorrt_llm/visual_gen/output.py index 9d1f6e99336a..0660e661122f 100644 --- a/tensorrt_llm/visual_gen/output.py +++ b/tensorrt_llm/visual_gen/output.py @@ -96,12 +96,8 @@ class VisualGenOutput: image: Optional[torch.Tensor] = None video: Optional[torch.Tensor] = None audio: Optional[torch.Tensor] = None - action: Optional[torch.Tensor] = None frame_rate: Optional[float] = None audio_sample_rate: Optional[int] = None - raw_action_dim: Optional[int] = None - action_mode: Optional[str] = None - domain_id: Optional[int] = None error: Optional[str] = None metrics: Optional[VisualGenMetrics] = None diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_cosmos3_transformer_parallel.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_cosmos3_transformer_parallel.py index f5dab83047d4..e5b770f5dfa8 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_cosmos3_transformer_parallel.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_cosmos3_transformer_parallel.py @@ -129,15 +129,6 @@ temporal_compression_factor_sound=1, ) -_ACTION_DIM = 64 -_T_ACTION = 4 -_COSMOS3_ACTION_CONFIG = dict( - **_COSMOS3_TEST_CONFIG, - action_gen=True, - action_dim=_ACTION_DIM, - num_embodiment_domains=32, -) - SEED_WEIGHTS = 123 SEED_INPUT = 456 SEED_COND_TEXT = 42 @@ -435,33 +426,6 @@ def _forward_with_audio( return out.video, out.audio -def _forward_with_action( - model: Cosmos3VFMTransformer, device: torch.device, text_seed: int -) -> Tuple[torch.Tensor, torch.Tensor]: - channels = _COSMOS3_TEST_CONFIG["latent_channel"] - hs, ts, text_ids, text_mask, video_shape = _cosmos3_inputs( - device, channels=channels, text_seed=text_seed - ) - torch.manual_seed(SEED_INPUT + 2) - action_latents = ( - torch.randn(hs.shape[0], _T_ACTION, _ACTION_DIM, device=device, dtype=hs.dtype) * 0.1 - ) - domain_ids = torch.tensor([7], dtype=torch.long, device=device) - model.reset_cache() - with torch.inference_mode(): - out = model( - hidden_states=hs, - timestep=ts, - text_ids=text_ids, - text_mask=text_mask, - video_shape=video_shape, - fps=_FPS, - action_latents=action_latents, - action_domain_ids=domain_ids, - ) - return out.video, out.action - - def _build_ref_and_parallel( *, tp_size: int = 1, @@ -598,37 +562,6 @@ def _logic_cosmos3_ulysses_audio_vs_single_gpu(rank, world_size): ) -def _logic_cosmos3_ulysses_action_vs_single_gpu(rank, world_size): - ref_model, ulysses_model, _, device = _build_ref_and_parallel( - ulysses_size=world_size, pretrained_dict=_COSMOS3_ACTION_CONFIG - ) - text_seed = _cfg_text_seed(rank, tp_size=1, ulysses_size=world_size, cfg_size=1) - - ref_video, ref_action = _forward_with_action(ref_model, device, text_seed) - ulysses_video, ulysses_action = _forward_with_action(ulysses_model, device, text_seed) - - if rank == 0: - vdiff = (ulysses_video.float() - ref_video.float()).abs() - adiff = (ulysses_action.float() - ref_action.float()).abs() - print( - f"[ulysses={world_size}+action] " - f"video max_abs_diff={vdiff.max().item():.6e}, " - f"action max_abs_diff={adiff.max().item():.6e}", - flush=True, - ) - - _assert_parity( - ulysses_video, - ref_video, - msg=f"Rank {rank}: Ulysses+action VIDEO differs from single-GPU reference", - ) - _assert_parity( - ulysses_action, - ref_action, - msg=f"Rank {rank}: Ulysses+action ACTION differs from single-GPU reference", - ) - - def _logic_cosmos3_tp_ulysses_vs_single_gpu(rank, world_size): tp_size = 2 ulysses_size = 2 @@ -778,11 +711,6 @@ def test_ulysses2_audio_vs_single_gpu(self): self._skip_if_unavailable() run_test_in_distributed(world_size=2, test_fn=_logic_cosmos3_ulysses_audio_vs_single_gpu) - def test_ulysses2_action_vs_single_gpu(self): - """Ulysses parity with action tokens appended to the GEN sequence.""" - self._skip_if_unavailable() - run_test_in_distributed(world_size=2, test_fn=_logic_cosmos3_ulysses_action_vs_single_gpu) - @pytest.mark.gpu4 def test_tp2_ulysses2_vs_single_gpu(self): self._skip_if_unavailable() diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py deleted file mode 100644 index ef9856e3056a..000000000000 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py +++ /dev/null @@ -1,271 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Unit tests for Cosmos3 action sizing helpers (no checkpoint / GPU required). - -Run: - pytest tests/unittest/_torch/visual_gen/test_cosmos3_action.py -v -""" - -import numpy as np -import PIL.Image -import pytest - -from tensorrt_llm._torch.visual_gen.models.cosmos3.action import ( - VIDEO_RES_SIZE_INFO, - action_reference_image, - find_closest_target_size, - resolve_action_size, -) -from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import COSMOS3_EXTRA_SPECS -from tensorrt_llm._torch.visual_gen.models.cosmos3.utils import normalize_video_input - -pytestmark = pytest.mark.cosmos3 - - -class TestFindClosestTargetSize: - @pytest.mark.parametrize( - "input_h,input_w,action_resolution,expected", - [ - (480, 832, 480, (832, 480)), - (832, 480, 480, (480, 832)), - (512, 512, 480, (640, 640)), - (704, 1280, 704, (1280, 704)), - (256, 256, 256, (256, 256)), - (720, 1280, 720, (1280, 720)), - ], - ) - def test_picks_closest_aspect_bucket(self, input_h, input_w, action_resolution, expected): - assert find_closest_target_size(input_h, input_w, action_resolution) == expected - - def test_accepts_string_and_int_resolution_keys(self): - ref_h, ref_w = 480, 832 - assert find_closest_target_size(ref_h, ref_w, 480) == find_closest_target_size( - ref_h, ref_w, "480" - ) - - def test_unknown_resolution_raises(self): - with pytest.raises(ValueError, match="Unknown Cosmos3 action resolution"): - find_closest_target_size(480, 832, 1080) - - @pytest.mark.parametrize("action_resolution", sorted(VIDEO_RES_SIZE_INFO)) - def test_all_buckets_have_aspect_entries(self, action_resolution): - assert VIDEO_RES_SIZE_INFO[action_resolution] - - -class TestResolveActionSize: - @staticmethod - def _ref_image(width: int, height: int) -> PIL.Image.Image: - return PIL.Image.new("RGB", (width, height)) - - def test_explicit_height_and_width_are_unchanged(self): - ref = self._ref_image(832, 480) - assert resolve_action_size(400, 600, ref, 480) == (400, 600) - - def test_unset_height_and_width_use_action_resolution_bucket(self): - ref = self._ref_image(832, 480) - assert resolve_action_size(None, None, ref, 480) == (480, 832) - - def test_partial_height_fills_width_from_bucket(self): - ref = self._ref_image(832, 480) - assert resolve_action_size(400, None, ref, 480) == (400, 832) - - def test_partial_width_fills_height_from_bucket(self): - ref = self._ref_image(832, 480) - assert resolve_action_size(None, 600, ref, 480) == (480, 600) - - -class TestActionResolutionExtraParam: - def test_extra_param_spec_uses_action_resolution_key(self): - spec = COSMOS3_EXTRA_SPECS["action_resolution"] - assert spec.type == "int" - assert spec.default == 480 - - -class TestDomainActionPresets: - def test_bridge_preset_fills_missing_fields(self): - from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( - resolve_domain_action_config, - ) - - cfg = resolve_domain_action_config(domain_name="bridge_orig_lerobot") - assert cfg["raw_action_dim"] == 10 - assert cfg["action_chunk_size"] == 16 - assert cfg["num_frames"] == 17 - assert cfg["action_resolution"] == 480 - assert cfg["frame_rate"] == 5.0 - assert cfg["warnings"] == [] - - def test_av_preset_uses_longer_chunk(self): - from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( - resolve_domain_action_config, - ) - - cfg = resolve_domain_action_config(domain_name="av") - assert cfg["action_chunk_size"] == 60 - assert cfg["num_frames"] == 61 - assert cfg["raw_action_dim"] == 9 - - def test_mismatch_emits_warning(self): - from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( - resolve_domain_action_config, - ) - - cfg = resolve_domain_action_config( - domain_name="bridge_orig_lerobot", - raw_action_dim=9, - ) - assert cfg["raw_action_dim"] == 9 - assert len(cfg["warnings"]) == 1 - assert "raw_action_dim=9" in cfg["warnings"][0] - - def test_action_fps_defaults_to_frame_rate(self): - from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( - resolve_domain_action_config, - ) - - cfg = resolve_domain_action_config(domain_name="av") - assert cfg["frame_rate"] == 10.0 - assert cfg["action_fps"] == 10.0 - - def test_explicit_action_fps_overrides_default(self): - from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( - resolve_domain_action_config, - ) - - cfg = resolve_domain_action_config(domain_name="av", action_fps=5.0, frame_rate=24.0) - assert cfg["frame_rate"] == 24.0 - assert cfg["action_fps"] == 5.0 - - def test_alias_maps_to_canonical_preset(self): - from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import get_domain_preset - - preset = get_domain_preset("robomind-franka") - assert preset is not None - assert preset["raw_action_dim"] == 10 - - def test_unknown_resolution_raises(self): - from tensorrt_llm._torch.visual_gen.models.cosmos3.action import normalize_action_resolution - - with pytest.raises(ValueError, match="Unknown Cosmos3 action_resolution"): - normalize_action_resolution(1080) - - -class TestActionReferenceImage: - def test_forward_dynamics_accepts_mp4_on_image_path(self, tmp_path, monkeypatch): - video_path = tmp_path / "clip.mp4" - video_path.write_bytes(b"fake") - expected = PIL.Image.new("RGB", (4, 2), "red") - - def _fake_read_video(path, pts_unit="sec"): - import torch - - tensor = torch.from_numpy(np.array(expected)).unsqueeze(0) - return tensor, None, {} - - monkeypatch.setattr("torchvision.io.read_video", _fake_read_video) - ref = action_reference_image( - action_mode="forward_dynamics", - image=str(video_path), - video=None, - ) - assert ref.size == expected.size - assert ref.getpixel((0, 0)) == (255, 0, 0) - - def test_policy_prefers_image_path_over_video(self, tmp_path): - image_path = tmp_path / "frame.png" - PIL.Image.new("RGB", (3, 3), "blue").save(image_path) - ref = action_reference_image( - action_mode="policy", - image=str(image_path), - video=str(tmp_path / "unused.mp4"), - ) - assert ref.getpixel((0, 0)) == (0, 0, 255) - - def test_policy_accepts_data_uri_image(self): - # URI references must go through the same loader as the I2V image - # branch instead of being treated as local filesystem paths. - import base64 - from io import BytesIO - - buf = BytesIO() - PIL.Image.new("RGB", (3, 3), "green").save(buf, format="PNG") - data_uri = "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode() - ref = action_reference_image( - action_mode="policy", - image=data_uri, - video=None, - ) - assert ref.size == (3, 3) - assert ref.getpixel((0, 0)) == (0, 128, 0) - - -class TestNormalizeVideoInput: - def test_image_path_returns_singleton_list(self, tmp_path): - image_path = tmp_path / "frame.png" - PIL.Image.new("RGB", (8, 4), "red").save(image_path) - assert normalize_video_input(str(image_path)) == [str(image_path)] - - def test_frame_directory_returns_sorted_paths(self, tmp_path): - (tmp_path / "b.png").write_bytes(b"") - (tmp_path / "a.png").write_bytes(b"") - (tmp_path / "skip.txt").write_text("x") - assert normalize_video_input(str(tmp_path)) == [ - str(tmp_path / "a.png"), - str(tmp_path / "b.png"), - ] - - def test_unsupported_file_extension_raises(self, tmp_path): - bad_path = tmp_path / "clip.mov" - bad_path.write_bytes(b"fake") - with pytest.raises(ValueError, match="must be a frame directory"): - normalize_video_input(str(bad_path)) - - def test_decode_mp4_returns_pil_frames(self, tmp_path, monkeypatch): - video_path = tmp_path / "clip.mp4" - video_path.write_bytes(b"fake") - expected = [ - PIL.Image.new("RGB", (2, 2), "red"), - PIL.Image.new("RGB", (2, 2), "blue"), - ] - - def _fake_read_video(path, pts_unit="sec"): - assert path == str(video_path) - assert pts_unit == "sec" - import torch - - tensor = torch.stack( - [torch.from_numpy(np.array(image)) for image in expected], - dim=0, - ) - return tensor, None, {} - - monkeypatch.setattr( - "torchvision.io.read_video", - _fake_read_video, - ) - frames = normalize_video_input(str(video_path)) - assert len(frames) == 2 - assert all(isinstance(frame, PIL.Image.Image) for frame in frames) - assert frames[0].getpixel((0, 0)) == (255, 0, 0) - - def test_decode_respects_max_frames(self, tmp_path, monkeypatch): - video_path = tmp_path / "clip.avi" - video_path.write_bytes(b"fake") - images = [PIL.Image.new("RGB", (1, 1), color) for color in ("red", "green", "blue")] - - def _fake_read_video(path, pts_unit="sec"): - import torch - - tensor = torch.stack( - [torch.from_numpy(np.array(image)) for image in images], - dim=0, - ) - return tensor, None, {} - - monkeypatch.setattr("torchvision.io.read_video", _fake_read_video) - frames = normalize_video_input( - str(video_path), - max_frames=2, - ) - assert len(frames) == 2 diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index d992b275ef6a..56626844a0d1 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -38,7 +38,6 @@ import torch from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( - COSMOS3_ACTION_PARAMS, COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION, COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP, COSMOS3_EXTRA_SPECS, @@ -226,24 +225,6 @@ def _require_audio_pipeline(pipeline) -> None: pytest.skip("Audio tokenizer was not loaded for this pipeline") -def _require_action_pipeline(pipeline) -> None: - if not getattr(pipeline, "action_gen", False): - pytest.skip("Checkpoint does not enable action generation") - - -def _assert_valid_action(action: torch.Tensor, *, raw_action_dim: int, chunk_size: int): - assert action is not None - assert action.dtype == torch.float32 - assert action.dim() == 3, f"Expected (B,T,D), got {action.shape}" - batch, t, d = action.shape - assert batch == 1 - assert t == chunk_size - assert d == raw_action_dim - af = action.float() - assert not torch.isnan(af).any() - assert not torch.isinf(af).any() - - def _scheduler_use_karras_sigmas(scheduler) -> bool | None: value = getattr(scheduler.config, "use_karras_sigmas", None) return None if value is None else bool(value) @@ -659,7 +640,6 @@ def test_v2v_keep_last_smoke(self, cosmos3_pipeline): def test_v2v_flow_shift_override_request_path(self): pipeline = Cosmos3OmniMoTPipeline.__new__(Cosmos3OmniMoTPipeline) pipeline.transformer = SimpleNamespace(device=torch.device("cpu")) - pipeline.action_gen = False pipeline.audio_gen = False calls = [] token_calls = [] @@ -773,141 +753,6 @@ def test_v2v_audio_smoke(self, cosmos3_pipeline): ) -@pytest.mark.integration -@pytest.mark.cosmos3_action -@pytest.mark.high_cuda_memory -class TestCosmos3Action: - ACTION_HEIGHT = 480 - ACTION_WIDTH = 832 - ACTION_FRAMES = COSMOS3_ACTION_PARAMS["num_frames"] - ACTION_CHUNK = COSMOS3_ACTION_PARAMS["action_chunk_size"] - RAW_ACTION_DIM = 10 - - def test_policy_smoke(self, cosmos3_pipeline): - _require_action_pipeline(cosmos3_pipeline) - image = _make_test_image().resize((self.ACTION_WIDTH, self.ACTION_HEIGHT)) - result = _run_forward( - cosmos3_pipeline, - image=image, - height=self.ACTION_HEIGHT, - width=self.ACTION_WIDTH, - num_frames=self.ACTION_FRAMES, - guidance_scale=COSMOS3_ACTION_PARAMS["guidance_scale"], - action_mode="policy", - domain_name="bridge_orig_lerobot", - raw_action_dim=self.RAW_ACTION_DIM, - action_chunk_size=self.ACTION_CHUNK, - ) - _assert_valid_video( - result.video, - num_frames=self.ACTION_FRAMES, - height=self.ACTION_HEIGHT, - width=self.ACTION_WIDTH, - ) - _assert_valid_action( - result.action, - raw_action_dim=self.RAW_ACTION_DIM, - chunk_size=self.ACTION_CHUNK, - ) - assert result.action_mode == "policy" - assert result.domain_id == 7 - _assert_scheduler_config( - cosmos3_pipeline, - flow_shift=COSMOS3_ACTION_PARAMS["flow_shift"], - use_karras_sigmas=cosmos3_pipeline._base_scheduler_use_karras_sigmas, - ) - - def test_forward_dynamics_smoke(self, cosmos3_pipeline): - _require_action_pipeline(cosmos3_pipeline) - image = _make_test_image().resize((self.ACTION_WIDTH, self.ACTION_HEIGHT)) - action_traj = [[0.1] * self.RAW_ACTION_DIM for _ in range(self.ACTION_CHUNK)] - result = _run_forward( - cosmos3_pipeline, - image=image, - height=self.ACTION_HEIGHT, - width=self.ACTION_WIDTH, - num_frames=self.ACTION_FRAMES, - guidance_scale=COSMOS3_ACTION_PARAMS["guidance_scale"], - action_mode="forward_dynamics", - domain_name="bridge_orig_lerobot", - action=action_traj, - action_chunk_size=self.ACTION_CHUNK, - ) - _assert_valid_video( - result.video, - num_frames=self.ACTION_FRAMES, - height=self.ACTION_HEIGHT, - width=self.ACTION_WIDTH, - ) - _assert_valid_action( - result.action, - raw_action_dim=self.RAW_ACTION_DIM, - chunk_size=self.ACTION_CHUNK, - ) - _assert_scheduler_config( - cosmos3_pipeline, - flow_shift=COSMOS3_ACTION_PARAMS["flow_shift"], - use_karras_sigmas=cosmos3_pipeline._base_scheduler_use_karras_sigmas, - ) - - def test_inverse_dynamics_smoke(self, cosmos3_pipeline): - _require_action_pipeline(cosmos3_pipeline) - image = _make_test_image().resize((self.ACTION_WIDTH, self.ACTION_HEIGHT)) - video = [image.copy() for _ in range(NUM_FRAMES)] - result = _run_forward( - cosmos3_pipeline, - image=None, - height=self.ACTION_HEIGHT, - width=self.ACTION_WIDTH, - num_frames=NUM_FRAMES, - guidance_scale=COSMOS3_ACTION_PARAMS["guidance_scale"], - action_mode="inverse_dynamics", - domain_name="bridge_orig_lerobot", - raw_action_dim=self.RAW_ACTION_DIM, - action_chunk_size=NUM_FRAMES, - video=video, - ) - _assert_valid_video( - result.video, - num_frames=NUM_FRAMES, - height=self.ACTION_HEIGHT, - width=self.ACTION_WIDTH, - ) - _assert_valid_action( - result.action, - raw_action_dim=self.RAW_ACTION_DIM, - chunk_size=NUM_FRAMES, - ) - _assert_scheduler_config( - cosmos3_pipeline, - flow_shift=COSMOS3_ACTION_PARAMS["flow_shift"], - use_karras_sigmas=cosmos3_pipeline._base_scheduler_use_karras_sigmas, - ) - - def test_action_and_audio_rejected(self, cosmos3_pipeline): - _require_action_pipeline(cosmos3_pipeline) - with pytest.raises(ValueError, match="joint action and audio"): - _run_forward( - cosmos3_pipeline, - image=_make_test_image(), - action_mode="policy", - domain_name="bridge_orig_lerobot", - raw_action_dim=self.RAW_ACTION_DIM, - enable_audio=True, - ) - - def test_action_and_t2i_rejected(self, cosmos3_pipeline): - _require_action_pipeline(cosmos3_pipeline) - with pytest.raises(ValueError, match="output_type='image'"): - _run_forward( - cosmos3_pipeline, - output_type="image", - action_mode="policy", - domain_name="bridge_orig_lerobot", - raw_action_dim=self.RAW_ACTION_DIM, - ) - - @pytest.mark.integration @pytest.mark.cosmos3_t2v @pytest.mark.high_cuda_memory diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py index c5d563c7fc9e..76e6c1a01a8c 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py @@ -397,148 +397,6 @@ def test_forward_with_audio_multiframe(self, audio_model_config): _assert_finite_output(out.audio, torch.Size([1, model.audio_dim, self.T_AUDIO])) -@pytest.mark.integration -class TestCosmos3Action: - """Action modality — Nano architecture, random weights, action_gen on.""" - - ACTION_DIM = 64 - T_ACTION = 4 - NUM_DOMAINS = 32 - - @pytest.fixture(autouse=True) - def _require_cuda(self): - if not torch.cuda.is_available(): - pytest.skip("CUDA not available") - - @pytest.fixture - def action_model_config(self): - checkpoint_dir = _require_checkpoint() - model_config = _load_model_config(checkpoint_dir) - cfg = model_config.pretrained_config - cfg.action_gen = True - cfg.action_dim = self.ACTION_DIM - cfg.num_embodiment_domains = self.NUM_DOMAINS - cfg.sound_gen = False - return model_config - - @pytest.fixture - def cosmos3_model_config_noaction(self): - checkpoint_dir = _require_checkpoint() - model_config = _load_model_config(checkpoint_dir) - model_config.pretrained_config.action_gen = False - model_config.pretrained_config.sound_gen = False - return model_config - - def test_action_model_structure(self, action_model_config): - model = Cosmos3VFMTransformer(model_config=action_model_config) - assert model.action_gen is True - assert model.action_dim == self.ACTION_DIM - assert hasattr(model, "action_proj_in") - assert hasattr(model, "action_proj_out") - assert hasattr(model, "action_modality_embed") - assert model.action_modality_embed.shape == (model.hidden_size,) - - def test_video_only_model_has_no_action_heads(self, cosmos3_model_config_noaction): - model = Cosmos3VFMTransformer(model_config=cosmos3_model_config_noaction) - assert model.action_gen is False - assert not hasattr(model, "action_proj_in") - - @pytest.mark.high_cuda_memory - def test_forward_with_action(self, action_model_config): - cfg = action_model_config.pretrained_config - model = _build_random_weight_model(action_model_config) - hs, ts, text_ids, text_mask, video_shape = _cosmos3_inputs( - DEVICE, channels=cfg.latent_channel - ) - action_latents = torch.randn(1, self.T_ACTION, model.action_dim, device=DEVICE, dtype=DTYPE) - domain_ids = torch.tensor([7], dtype=torch.long, device=DEVICE) - with torch.inference_mode(): - out = model( - hidden_states=hs, - timestep=ts / _NUM_TRAIN_TIMESTEPS, - raw_timestep=ts, - text_ids=text_ids, - text_mask=text_mask, - video_shape=video_shape, - fps=24.0, - action_latents=action_latents, - action_domain_ids=domain_ids, - ) - _assert_finite_output(out.video, hs.shape) - assert out.action is not None - _assert_finite_output(out.action, torch.Size([1, self.T_ACTION, model.action_dim])) - - @pytest.mark.high_cuda_memory - def test_forward_without_action_latents_returns_none(self, action_model_config): - cfg = action_model_config.pretrained_config - model = _build_random_weight_model(action_model_config) - hs, ts, text_ids, text_mask, video_shape = _cosmos3_inputs( - DEVICE, channels=cfg.latent_channel - ) - with torch.inference_mode(): - out = model( - hidden_states=hs, - timestep=ts / _NUM_TRAIN_TIMESTEPS, - raw_timestep=ts, - text_ids=text_ids, - text_mask=text_mask, - video_shape=video_shape, - ) - _assert_finite_output(out.video, hs.shape) - assert out.action is None - - @pytest.mark.high_cuda_memory - def test_forward_with_action_noisy_mask(self, action_model_config): - cfg = action_model_config.pretrained_config - model = _build_random_weight_model(action_model_config) - hs, ts, text_ids, text_mask, video_shape = _cosmos3_inputs( - DEVICE, channels=cfg.latent_channel, t=2 - ) - action_latents = torch.randn(1, self.T_ACTION, model.action_dim, device=DEVICE, dtype=DTYPE) - noisy_mask = torch.ones(1, self.T_ACTION, 1, device=DEVICE, dtype=DTYPE) - noisy_mask[:, 0, :] = 0.0 - domain_ids = torch.tensor([7], dtype=torch.long, device=DEVICE) - with torch.inference_mode(): - out = model( - hidden_states=hs, - timestep=ts / _NUM_TRAIN_TIMESTEPS, - raw_timestep=ts, - text_ids=text_ids, - text_mask=text_mask, - video_shape=video_shape, - fps=24.0, - action_latents=action_latents, - action_domain_ids=domain_ids, - action_noisy_mask=noisy_mask, - ) - _assert_finite_output(out.video, hs.shape) - _assert_finite_output(out.action, torch.Size([1, self.T_ACTION, model.action_dim])) - - @pytest.mark.high_cuda_memory - def test_forward_with_action_multiframe(self, action_model_config): - cfg = action_model_config.pretrained_config - model = _build_random_weight_model(action_model_config) - hs, ts, text_ids, text_mask, video_shape = _cosmos3_inputs( - DEVICE, channels=cfg.latent_channel, t=3 - ) - action_latents = torch.randn(1, self.T_ACTION, model.action_dim, device=DEVICE, dtype=DTYPE) - domain_ids = torch.tensor([7], dtype=torch.long, device=DEVICE) - with torch.inference_mode(): - out = model( - hidden_states=hs, - timestep=ts / _NUM_TRAIN_TIMESTEPS, - raw_timestep=ts, - text_ids=text_ids, - text_mask=text_mask, - video_shape=video_shape, - fps=24.0, - action_latents=action_latents, - action_domain_ids=domain_ids, - ) - _assert_finite_output(out.video, hs.shape) - _assert_finite_output(out.action, torch.Size([1, self.T_ACTION, model.action_dim])) - - @pytest.mark.integration class TestCosmos3TransformerCheckpoint: """Load Cosmos3-Nano transformer weights and run a single forward step.""" diff --git a/tests/unittest/visual_gen/test_output.py b/tests/unittest/visual_gen/test_output.py index 54097e64558e..e27f9c33e3b5 100644 --- a/tests/unittest/visual_gen/test_output.py +++ b/tests/unittest/visual_gen/test_output.py @@ -32,12 +32,8 @@ def test_visual_gen_output_is_dataclass(): "image", "video", "audio", - "action", "frame_rate", "audio_sample_rate", - "raw_action_dim", - "action_mode", - "domain_id", "error", "metrics", } @@ -66,10 +62,6 @@ def test_minimal_construction_defaults(): assert out.audio is None assert out.frame_rate is None assert out.audio_sample_rate is None - assert out.action is None - assert out.raw_action_dim is None - assert out.action_mode is None - assert out.domain_id is None assert out.error is None assert out.metrics is None @@ -761,19 +753,15 @@ def test_encoding_not_top_level_reexport(): # --------------------------------------------------------------------------- -def test_pipeline_output_field_set(): - """PipelineOutput exposes exactly the expected fields.""" +def test_pipeline_output_has_eight_fields(): + """PipelineOutput has the eight expected fields.""" field_names = {f.name for f in fields(PipelineOutput)} assert field_names == { "image", "video", "audio", - "action", "frame_rate", "audio_sample_rate", - "raw_action_dim", - "action_mode", - "domain_id", "pre_denoise", "denoise", "post_denoise", @@ -786,10 +774,6 @@ def test_pipeline_output_default_construction(): assert p.image is None assert p.video is None assert p.audio is None - assert p.action is None - assert p.raw_action_dim is None - assert p.action_mode is None - assert p.domain_id is None assert p.frame_rate is None assert p.audio_sample_rate is None assert p.pre_denoise == 0.0 @@ -797,43 +781,6 @@ def test_pipeline_output_default_construction(): assert p.post_denoise == 0.0 -def test_from_response_propagates_action_fields(): - """to_visual_gen_output carries the action tensor and its metadata.""" - action = torch.zeros(1, 16, 10, dtype=torch.float32) - pipeline_out = PipelineOutput( - action=action, - raw_action_dim=10, - action_mode="policy", - domain_id=7, - ) - resp = DiffusionResponse(request_id=5, output=pipeline_out, generation=12.0) - out = to_visual_gen_output(resp) - assert out.action is action - assert out.raw_action_dim == 10 - assert out.action_mode == "policy" - assert out.domain_id == 7 - - -def test_batch_split_slices_action(): - """split_visual_gen_output slices batched actions per item and shares scalar metadata.""" - action = torch.stack([torch.full((16, 10), float(v), dtype=torch.float32) for v in (1, 2, 3)]) - pipeline_out = PipelineOutput( - action=action, - raw_action_dim=10, - action_mode="inverse_dynamics", - domain_id=1, - ) - resp = DiffusionResponse(request_id=9, output=pipeline_out, generation=30.0) - outs = split_visual_gen_output(resp, batch_size=3) - assert len(outs) == 3 - for i, out in enumerate(outs): - assert out.action.shape == (16, 10) - assert float(out.action[0, 0]) == float(i + 1) - assert out.raw_action_dim == 10 - assert out.action_mode == "inverse_dynamics" - assert out.domain_id == 1 - - def test_media_output_unimportable(): """``MediaOutput`` is not importable from any path.""" with pytest.raises(ImportError): From cbd5419e2c98ddcae0957e610adddb1fe705441c Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 13 Jul 2026 21:39:15 -0700 Subject: [PATCH 17/64] Replace PyAV with ffprobe for video stream detection Signed-off-by: Igor Shovkun --- tensorrt_llm/serve/visual_gen_utils.py | 45 ++++++++++++++++++-------- 1 file changed, 32 insertions(+), 13 deletions(-) diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index 29bc4e7c2969..ca0845dbcabd 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -1,7 +1,9 @@ import asyncio import base64 +import json import os import shutil +import subprocess from typing import Any, Dict, List, Optional from PIL import Image, UnidentifiedImageError @@ -101,24 +103,41 @@ def _reference_is_image(path: str) -> bool: def _reference_is_video(path: str) -> bool: - """True when ``path`` holds video content (a PyAV-openable video stream). - - Total predicate: False for images, audio, and undecodable content - alike. Images must be excluded explicitly because FFmpeg demuxes a - still image as a valid single-frame video stream, so an av probe - alone would claim every PNG/JPEG. A missing ``av`` package raises - ``ImportError`` — that is a deployment problem, not a content - verdict. + """True when ``path`` holds video content (a stream ffprobe recognizes). + + Total predicate: False for images, audio, and undecodable content alike. + Images must be excluded explicitly because FFmpeg demuxes a still image as + a single-frame video stream, so the ffprobe check must be gated by the PIL + probe. ``ffprobe`` (part of the ffmpeg CLI, a documented system dependency) + absent → the probe returns False rather than raising. """ if _reference_is_image(path): return False - import av - try: - with av.open(path) as container: - return bool(container.streams.video) - except av.FFmpegError: + result = subprocess.run( + [ + "ffprobe", + "-v", + "error", + "-select_streams", + "v", + "-show_entries", + "stream=codec_type", + "-of", + "json", + path, + ], + capture_output=True, + text=True, + check=True, + ) + except (subprocess.CalledProcessError, FileNotFoundError): + return False + try: + streams = json.loads(result.stdout).get("streams", []) + except json.JSONDecodeError: return False + return any(s.get("codec_type") == "video" for s in streams) def parse_visual_gen_params( From 82f351f897ee3dde5dab0666d0976a6cb6531447 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 13 Jul 2026 21:48:55 -0700 Subject: [PATCH 18/64] Remove single-image support mention from V2V video_path docs Signed-off-by: Igor Shovkun --- examples/visual_gen/models/cosmos3/README.md | 2 +- examples/visual_gen/models/cosmos3/cosmos3.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/visual_gen/models/cosmos3/README.md b/examples/visual_gen/models/cosmos3/README.md index 25ae32d357b3..c6d310fedc7b 100644 --- a/examples/visual_gen/models/cosmos3/README.md +++ b/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 frame directory, `.mp4`/`.avi` file, or single image; passing it without `--action_mode` selects V2V). Only the first (or last, per `condition_video_keep`) `max(condition_frame_indexes_vision) * 4 + 1` input frames condition the output (5 by default); `.mp4`/`.avi` decode uses the `av` package (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 frame directory or `.mp4`/`.avi` file). Only the first (or last, per `condition_video_keep`) `max(condition_frame_indexes_vision) * 4 + 1` input frames condition the output (5 by default); `.mp4`/`.avi` decode uses the `av` package (see [Media I/O dependencies](#media-io-dependencies)). - **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 diff --git a/examples/visual_gen/models/cosmos3/cosmos3.py b/examples/visual_gen/models/cosmos3/cosmos3.py index 2e35925a5b6c..995736f533ed 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3.py +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -229,7 +229,7 @@ def main(): "--video_path", type=str, default=None, - help="Reference video for V2V: a local frame directory, .mp4/.avi file, or image path", + help="Reference video for V2V: a local frame directory or .mp4/.avi file", ) parser.add_argument( "--output_type", type=str, default="video", help="Output type (video, image)" From 29d62ec74961efaa486852c97572e1822aad7e85 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 13 Jul 2026 22:03:30 -0700 Subject: [PATCH 19/64] Refactor image-to-latent encoding to reuse `_encode_video_tensor` Signed-off-by: Igor Shovkun --- .../models/cosmos3/pipeline_cosmos3.py | 23 +------------------ 1 file changed, 1 insertion(+), 22 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index cd982b837c2d..3b4eb1d4d1ef 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -528,28 +528,7 @@ def _encode_conditioning_video( # Build pixel-space video: repeat the conditioning image across all frames # image_tensor: [1, 3, H, W] -> [1, 3, 1, H, W] -> [1, 3, num_frames, H, W] video = image_tensor.unsqueeze(2).expand(-1, -1, num_frames, -1, -1).contiguous() - video = video.to(device=self.device, dtype=self.vae.dtype) - - latent = self.vae.encode(video).latent_dist.mode() - - # Normalize (inverse of _decode_latents denormalization) - if hasattr(self.vae.config, "latents_mean") and hasattr(self.vae.config, "latents_std"): - latents_mean = ( - torch.tensor(self.vae.config.latents_mean) - .view(1, -1, 1, 1, 1) - .to(latent.device, latent.dtype) - ) - latents_std = ( - torch.tensor(self.vae.config.latents_std) - .view(1, -1, 1, 1, 1) - .to(latent.device, latent.dtype) - ) - latent = (latent - latents_mean) / latents_std - else: - scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0) - latent = latent * scaling_factor - - return latent.to(self.dtype) + return self._encode_video_tensor(video) def _prepare_latents_i2v( self, From a92f8bb525824b0c4885733bffb8e25b77bdbaae Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 13 Jul 2026 22:09:04 -0700 Subject: [PATCH 20/64] Remove unused height/width params from conditioning video encoder Signed-off-by: Igor Shovkun --- .../_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index 3b4eb1d4d1ef..c62364590e05 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -502,8 +502,6 @@ def _encode_conditioning_video( self, image_tensor: torch.Tensor, num_frames: int, - height: int, - width: int, ) -> torch.Tensor: """VAE-encode a conditioning image as a full-length video. @@ -518,8 +516,6 @@ def _encode_conditioning_video( Args: image_tensor: [1, 3, H, W] in [-1, 1] num_frames: total pixel frames for the video - height: pixel height - width: pixel width Returns: [1, C, T_latent, H_latent, W_latent] normalized latent of the @@ -572,8 +568,6 @@ def _prepare_latents_i2v( cond_latent = self._encode_conditioning_video( image_tensor, num_frames, - height, - width, ) # [1, C, T_lat, H_lat, W_lat] # Keep only frame 0 for conditioning; replace rest with noise From 4a0702d0319b0bcc63a633a51b5a7ef0de065467 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 13 Jul 2026 22:14:46 -0700 Subject: [PATCH 21/64] Fix `post_step_fn` signature in Cosmos3 pipeline Remove unused `step_extra_stream_latents` parameter from `post_step_fn` and update its return value accordingly. Signed-off-by: Igor Shovkun --- .../_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index c62364590e05..ec3d1516e1f9 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -1061,7 +1061,7 @@ def forward_fn( return video_noise_pred, {"audio": audio_noise_pred} return video_noise_pred - def post_step_fn(step_latents, step_extra_stream_latents): + def post_step_fn(step_latents): if velocity_mask is not None and condition_latents is not None: step_latents = ( velocity_mask * step_latents + (1.0 - velocity_mask) * condition_latents @@ -1071,7 +1071,7 @@ def post_step_fn(step_latents, step_extra_stream_latents): step_latents[:, :, 0:1, :, :] = image_latent.to( device=step_latents.device, dtype=step_latents.dtype ) - return step_latents, step_extra_stream_latents + return step_latents # 5. Build CFG tensors — text_ids and text_mask need to be split for CFG # BasePipeline.denoise batches [uncond, cond] when guidance_scale > 1 From e698f6d9b2ac55e65f48da3165f7c1744c0d5b15 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Tue, 14 Jul 2026 09:10:25 -0700 Subject: [PATCH 22/64] Simplify `post_step_fn` to single-argument interface Remove the two-argument `post_step_fn(latents, extra_stream_latents)` overload and the `inspect.signature` dispatch logic. The callable now always takes and returns `latents` only. Signed-off-by: Igor Shovkun --- tensorrt_llm/_torch/visual_gen/pipeline.py | 17 +++-------------- 1 file changed, 3 insertions(+), 14 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/pipeline.py b/tensorrt_llm/_torch/visual_gen/pipeline.py index f460ab8023df..4428de277928 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline.py @@ -1,5 +1,4 @@ import contextlib -import inspect import itertools import os import time @@ -1097,14 +1096,8 @@ def denoise( guidance_interval: Optional ``(lo, hi)`` scheduler-timestep range in which CFG is active. Outside the interval the effective scale is 1.0 (conditional prediction only); both branches still run. - post_step_fn: Optional callable applied after each scheduler step. - If the callable accepts one argument, it is invoked as - ``post_step_fn(latents) -> latents``. If it accepts two - or more, it is invoked as - ``post_step_fn(latents, extra_stream_latents) -> - (latents, extra_stream_latents)``, where - ``extra_stream_latents`` is the dict of parallel streams - from ``extra_streams`` (e.g. ``{"action": ...}``). + post_step_fn: Optional callable applied after each scheduler step, + invoked as ``post_step_fn(latents) -> latents``. Use for constraints that must hold throughout denoising. Returns: @@ -1235,11 +1228,7 @@ def denoise( ) if post_step_fn is not None: - sig = inspect.signature(post_step_fn) - if len(sig.parameters) >= 2: - latents, extra_stream_latents = post_step_fn(latents, extra_stream_latents) - else: - latents = post_step_fn(latents) + latents = post_step_fn(latents) # Logging if self.rank == 0: From ebe02d8f8a4259939fbce0af0ca2430a5cc67040 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Tue, 14 Jul 2026 09:22:16 -0700 Subject: [PATCH 23/64] Skip action module weights during Cosmos3 checkpoint load Signed-off-by: Igor Shovkun --- .../visual_gen/models/cosmos3/transformer_cosmos3.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py index 8deedeed044a..baa397d66538 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -1154,7 +1154,13 @@ def load_weights(self, weights: dict) -> None: Maps UND vs GEN blocks into this module's layout (causal self-attn vs cross-attn + MLPs). """ remapped = {} - skip_prefixes = ("lm_head.",) + # The Cosmos3 checkpoint ships action modules (action_gen=true), but this + # transformer no longer builds them — skip their weights explicitly so the + # load stays quiet instead of warning on each as an unknown key. + skip_prefixes = ( + "lm_head.", + "action_", + ) for key, value in weights.items(): k = key From 2209843d1749dd735029bcee510bba3079ad4abc Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Tue, 14 Jul 2026 12:51:16 -0700 Subject: [PATCH 24/64] Replace ffprobe with PyAV for video reference detection; make pyav opt-in explicitly Signed-off-by: Igor Shovkun --- examples/visual_gen/models/cosmos3/README.md | 2 +- tensorrt_llm/serve/visual_gen_utils.py | 54 +++++++++---------- .../visual_gen/test_trtllm_serve_endpoints.py | 18 ++++++- 3 files changed, 41 insertions(+), 33 deletions(-) diff --git a/examples/visual_gen/models/cosmos3/README.md b/examples/visual_gen/models/cosmos3/README.md index c6d310fedc7b..d7c034e27707 100644 --- a/examples/visual_gen/models/cosmos3/README.md +++ b/examples/visual_gen/models/cosmos3/README.md @@ -34,7 +34,7 @@ export TRTLLM_DISABLE_COSMOS3_GUARDRAILS=1 ## Media I/O dependencies - 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, inverse dynamics) uses the `av` (PyAV) package. It is **not** bundled with TensorRT-LLM — install it yourself: `pip install av`. Frame directories and single-image references work without it. +- Decoding `.mp4`/`.avi` reference videos (V2V) uses the `av` (PyAV) package. It is **not** bundled with TensorRT-LLM — install it yourself: `pip install av`. Frame directories work without it. ## Deployment configs diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index ca0845dbcabd..25520fe42aff 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -1,9 +1,7 @@ import asyncio import base64 -import json import os import shutil -import subprocess from typing import Any, Dict, List, Optional from PIL import Image, UnidentifiedImageError @@ -103,41 +101,37 @@ def _reference_is_image(path: str) -> bool: def _reference_is_video(path: str) -> bool: - """True when ``path`` holds video content (a stream ffprobe recognizes). - - Total predicate: False for images, audio, and undecodable content alike. - Images must be excluded explicitly because FFmpeg demuxes a still image as - a single-frame video stream, so the ffprobe check must be gated by the PIL - probe. ``ffprobe`` (part of the ffmpeg CLI, a documented system dependency) - absent → the probe returns False rather than raising. + """True when ``path`` holds video content (a PyAV-openable video stream). + + Total predicate over *content*: False for images, audio, and undecodable + content alike. Images must be excluded explicitly because FFmpeg demuxes a + still image as a valid single-frame video stream, so an av probe alone would + claim every PNG/JPEG. + + PyAV is gated behind ``TRTLLM_ENABLE_PYAV=1`` (the same opt-in the audio + extraction path uses) and is also required to decode ``.mp4``/``.avi`` + references through ``torchvision``. A disabled gate or a missing package is + a deployment problem, not a content verdict, so it raises rather than + masquerading as "not video". """ if _reference_is_image(path): return False + if os.environ.get("TRTLLM_ENABLE_PYAV", "0") != "1": + raise RuntimeError( + "PyAV is required to detect and decode video references. " + "Set the environment variable TRTLLM_ENABLE_PYAV=1 to enable it." + ) try: - result = subprocess.run( - [ - "ffprobe", - "-v", - "error", - "-select_streams", - "v", - "-show_entries", - "stream=codec_type", - "-of", - "json", - path, - ], - capture_output=True, - text=True, - check=True, + import av + except ImportError: + raise ImportError( + "PyAV is required to detect and decode video references but is not installed." ) - except (subprocess.CalledProcessError, FileNotFoundError): - return False try: - streams = json.loads(result.stdout).get("streams", []) - except json.JSONDecodeError: + with av.open(path) as container: + return bool(container.streams.video) + except av.FFmpegError: return False - return any(s.get("codec_type") == "video" for s in streams) def parse_visual_gen_params( diff --git a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py index 97e8e15b5c30..61c3337a7ae7 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py @@ -60,6 +60,20 @@ def _assert_llm_envelope( assert message_contains in body["message"], body["message"] +def _require_pyav_optin(): + """Skip unless PyAV is installed and opted into via ``TRTLLM_ENABLE_PYAV=1``. + + These tests drive ``_reference_is_video``, which gates PyAV behind the same + opt-in ``media_io`` uses. ``av`` ships in no CI image, so this both + ``importorskip``-s the package and skips when the gate is disabled; it + returns the ``av`` module for tests that build fixtures with it. + """ + av = pytest.importorskip("av") + if os.environ.get("TRTLLM_ENABLE_PYAV", "0") != "1": + pytest.skip("requires the PyAV opt-in (TRTLLM_ENABLE_PYAV=1)") + return av + + def _make_dummy_image_tensor(height: int = 64, width: int = 64) -> torch.Tensor: """Create a small dummy uint8 image tensor (H, W, C).""" return torch.randint(0, 256, (height, width, 3), dtype=torch.uint8) @@ -848,7 +862,7 @@ def test_sync_video_generation_multipart_with_video_reference(self, video_client The reference is classified by decoding its content, so the clip is synthesized in-test with PyAV — no video asset ships with the repo. """ - av = pytest.importorskip("av") + av = _require_pyav_optin() np = pytest.importorskip("numpy") ref_path = tmp_path / "ref.mp4" with av.open(str(ref_path), "w") as container: @@ -891,7 +905,7 @@ def test_sync_video_generation_multipart_with_video_reference(self, video_client def test_sync_video_generation_undecodable_reference_400(self, video_client): """Content neither PIL nor PyAV can decode is rejected at the boundary.""" - pytest.importorskip("av") + _require_pyav_optin() resp = video_client.post( "/v1/videos/generations", data={"prompt": "x"}, From 79274ea00755b3a371d5e38b57921878f2236e3e Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Tue, 14 Jul 2026 14:27:35 -0700 Subject: [PATCH 25/64] Replace PyAV with OpenCV for Cosmos3 video decoding Switch the V2V video reference decoder from `av` (PyAV) to OpenCV, aligning it with the shared multimodal video path. This removes the `TRTLLM_ENABLE_PYAV=1` gate, drops the `av` dependency, and updates docs, tests, and the serve-side `_reference_is_video` probe accordingly. Signed-off-by: Igor Shovkun --- examples/visual_gen/models/cosmos3/README.md | 4 +- examples/visual_gen/models/cosmos3/cosmos3.py | 2 +- .../_torch/visual_gen/models/cosmos3/utils.py | 26 ++++++++--- tensorrt_llm/serve/visual_gen_utils.py | 34 +++++--------- .../visual_gen/test_trtllm_serve_endpoints.py | 44 ++++++++----------- .../visual_gen/test_visual_gen_utils.py | 34 +++++++------- 6 files changed, 67 insertions(+), 77 deletions(-) diff --git a/examples/visual_gen/models/cosmos3/README.md b/examples/visual_gen/models/cosmos3/README.md index d7c034e27707..aff88d35d04b 100644 --- a/examples/visual_gen/models/cosmos3/README.md +++ b/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 frame directory or `.mp4`/`.avi` file). Only the first (or last, per `condition_video_keep`) `max(condition_frame_indexes_vision) * 4 + 1` input frames condition the output (5 by default); `.mp4`/`.avi` decode uses the `av` package (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 frame directory or `.mp4`/`.avi` file). Only the first (or last, per `condition_video_keep`) `max(condition_frame_indexes_vision) * 4 + 1` input frames condition the output (5 by default); `.mp4`/`.avi` decode uses OpenCV (see [Media I/O dependencies](#media-io-dependencies)). - **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 @@ -34,7 +34,7 @@ export TRTLLM_DISABLE_COSMOS3_GUARDRAILS=1 ## Media I/O dependencies - 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) uses the `av` (PyAV) package. It is **not** bundled with TensorRT-LLM — install it yourself: `pip install av`. Frame directories work without it. +- Decoding `.mp4`/`.avi` reference videos (V2V) uses OpenCV — the same optional decoder as the multimodal video path. It is **not** bundled with TensorRT-LLM — install it yourself: `pip install opencv-python-headless`. Frame directories work without it. ## Deployment configs diff --git a/examples/visual_gen/models/cosmos3/cosmos3.py b/examples/visual_gen/models/cosmos3/cosmos3.py index 995736f533ed..d1ea4a0448b0 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3.py +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -26,7 +26,7 @@ - **V2V** — video-conditioned video (``prompts/v2v.json``). Condition on the first (or last, per ``condition_video_keep``) frames of a reference video via ``--video_path`` (a local frame directory, ``.mp4``/``.avi`` file, or - single image; ``.mp4``/``.avi`` decode requires the ``av`` package). + single image; ``.mp4``/``.avi`` decode requires OpenCV). - **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). diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py index 75dcec8a3ebe..ee02ec6cf127 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py @@ -23,14 +23,26 @@ def pil_to_rgb(value: Any) -> PIL.Image.Image: def decode_video_file(path: Path, max_frames: Optional[int] = None) -> List[PIL.Image.Image]: - import torchvision.io as io - - frames, _, _ = io.read_video(str(path), pts_unit="sec") - if frames.numel() == 0: + # OpenCV is the shared multimodal video decoder; ``_get_cv2`` raises a clear + # ``pip install opencv-python-headless`` hint when it is not installed. + from tensorrt_llm.inputs.media_io import _get_cv2 + + cv2 = _get_cv2() + capture = cv2.VideoCapture(str(path)) + try: + if not capture.isOpened(): + raise ValueError(f"Cosmos3 could not open video file: {path}") + frames: List[PIL.Image.Image] = [] + while max_frames is None or len(frames) < max_frames: + ok, frame = capture.read() + if not ok: + break + frames.append(PIL.Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))) + finally: + capture.release() + if not frames: raise ValueError(f"Cosmos3 video file contains no frames: {path}") - if max_frames is not None: - frames = frames[:max_frames] - return [PIL.Image.fromarray(frames[i].numpy()) for i in range(frames.shape[0])] + return frames def normalize_video_input_path(path: Path, max_frames: Optional[int] = None) -> List[Any]: diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index 25520fe42aff..20d27990cc0e 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -6,6 +6,7 @@ from PIL import Image, UnidentifiedImageError +from tensorrt_llm.inputs.media_io import _get_cv2 from tensorrt_llm.logger import logger from tensorrt_llm.serve.openai_protocol import ImageGenerationRequest, VideoGenerationRequest from tensorrt_llm.visual_gen import VisualGen, VisualGenParams @@ -101,37 +102,22 @@ def _reference_is_image(path: str) -> bool: def _reference_is_video(path: str) -> bool: - """True when ``path`` holds video content (a PyAV-openable video stream). + """True when ``path`` holds video content (an OpenCV-decodable video stream). Total predicate over *content*: False for images, audio, and undecodable content alike. Images must be excluded explicitly because FFmpeg demuxes a - still image as a valid single-frame video stream, so an av probe alone would - claim every PNG/JPEG. - - PyAV is gated behind ``TRTLLM_ENABLE_PYAV=1`` (the same opt-in the audio - extraction path uses) and is also required to decode ``.mp4``/``.avi`` - references through ``torchvision``. A disabled gate or a missing package is - a deployment problem, not a content verdict, so it raises rather than - masquerading as "not video". + still image as a valid single-frame video stream, so a video probe alone + would claim every PNG/JPEG. OpenCV is the optional decoder shared with the + multimodal video path; ``_get_cv2`` raises a clear install hint if missing. """ if _reference_is_image(path): return False - if os.environ.get("TRTLLM_ENABLE_PYAV", "0") != "1": - raise RuntimeError( - "PyAV is required to detect and decode video references. " - "Set the environment variable TRTLLM_ENABLE_PYAV=1 to enable it." - ) - try: - import av - except ImportError: - raise ImportError( - "PyAV is required to detect and decode video references but is not installed." - ) + cv2 = _get_cv2() + capture = cv2.VideoCapture(path) try: - with av.open(path) as container: - return bool(container.streams.video) - except av.FFmpegError: - return False + return bool(capture.isOpened() and capture.read()[0]) + finally: + capture.release() def parse_visual_gen_params( diff --git a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py index 61c3337a7ae7..2f0e21bbebbe 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py @@ -60,18 +60,15 @@ def _assert_llm_envelope( assert message_contains in body["message"], body["message"] -def _require_pyav_optin(): - """Skip unless PyAV is installed and opted into via ``TRTLLM_ENABLE_PYAV=1``. +def _require_opencv(): + """Skip unless OpenCV is installed (the shared optional video decoder). - These tests drive ``_reference_is_video``, which gates PyAV behind the same - opt-in ``media_io`` uses. ``av`` ships in no CI image, so this both - ``importorskip``-s the package and skips when the gate is disabled; it - returns the ``av`` module for tests that build fixtures with it. + These tests drive ``_reference_is_video``, which decodes references via + OpenCV (the same optional dep the multimodal video path uses). ``cv2`` ships + in no CI image, so this ``importorskip``-s it and returns the module for + tests that synthesize a clip with ``cv2.VideoWriter``. """ - av = pytest.importorskip("av") - if os.environ.get("TRTLLM_ENABLE_PYAV", "0") != "1": - pytest.skip("requires the PyAV opt-in (TRTLLM_ENABLE_PYAV=1)") - return av + return pytest.importorskip("cv2") def _make_dummy_image_tensor(height: int = 64, width: int = 64) -> torch.Tensor: @@ -860,24 +857,19 @@ def test_sync_video_generation_multipart_with_video_reference(self, video_client """A video ``input_reference`` routes to ``extra_params["video"]`` (V2V). The reference is classified by decoding its content, so the clip is - synthesized in-test with PyAV — no video asset ships with the repo. + synthesized in-test with OpenCV — no video asset ships with the repo. """ - av = _require_pyav_optin() + cv2 = _require_opencv() np = pytest.importorskip("numpy") ref_path = tmp_path / "ref.mp4" - with av.open(str(ref_path), "w") as container: - # mpeg4 is a built-in FFmpeg encoder (h264 may be absent from - # LGPL PyAV wheels). - stream = container.add_stream("mpeg4", rate=4) - stream.width = 16 - stream.height = 16 - stream.pix_fmt = "yuv420p" + # mp4v is a built-in FFmpeg mpeg4 encoder present in the opencv wheel. + writer = cv2.VideoWriter(str(ref_path), cv2.VideoWriter_fourcc(*"mp4v"), 4.0, (16, 16)) + try: for _ in range(2): - frame = av.VideoFrame.from_ndarray( - np.zeros((16, 16, 3), dtype=np.uint8), format="rgb24" - ) - container.mux(stream.encode(frame)) - container.mux(stream.encode()) + writer.write(np.zeros((16, 16, 3), dtype=np.uint8)) + finally: + writer.release() + assert ref_path.exists() and ref_path.stat().st_size > 0 with open(ref_path, "rb") as f: resp = video_client.post( @@ -904,8 +896,8 @@ def test_sync_video_generation_multipart_with_video_reference(self, video_client assert os.path.exists(video_ref) def test_sync_video_generation_undecodable_reference_400(self, video_client): - """Content neither PIL nor PyAV can decode is rejected at the boundary.""" - _require_pyav_optin() + """Content neither PIL nor OpenCV can decode is rejected at the boundary.""" + _require_opencv() resp = video_client.post( "/v1/videos/generations", data={"prompt": "x"}, diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index 00c41baec4c5..2d26f6a13fa3 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -12,6 +12,8 @@ from __future__ import annotations import base64 +import os +import tempfile from io import BytesIO from typing import Any, Dict, Optional @@ -293,25 +295,23 @@ def test_missing_media_storage_path_raises(self): @staticmethod def _mp4_bytes() -> bytes: - """Encode a 2-frame 16x16 mpeg4-in-mp4 clip in memory. + """Encode a 2-frame 16x16 mp4v-in-mp4 clip and return its bytes. - ``mpeg4`` is a built-in FFmpeg encoder, so this works on any - PyAV wheel (h264 may be absent from LGPL builds). + ``mp4v`` is a built-in FFmpeg mpeg4 encoder present in the opencv wheel. + OpenCV writes only to a path, so encode to a tempfile and read it back. """ - av = pytest.importorskip("av") - buf = BytesIO() - with av.open(buf, "w", format="mp4") as container: - stream = container.add_stream("mpeg4", rate=4) - stream.width = 16 - stream.height = 16 - stream.pix_fmt = "yuv420p" + cv2 = pytest.importorskip("cv2") + with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp: + path = tmp.name + try: + writer = cv2.VideoWriter(path, cv2.VideoWriter_fourcc(*"mp4v"), 4.0, (16, 16)) for _ in range(2): - frame = av.VideoFrame.from_ndarray( - np.zeros((16, 16, 3), dtype=np.uint8), format="rgb24" - ) - container.mux(stream.encode(frame)) - container.mux(stream.encode()) - return buf.getvalue() + writer.write(np.zeros((16, 16, 3), dtype=np.uint8)) + writer.release() + with open(path, "rb") as f: + return f.read() + finally: + os.remove(path) def test_multipart_video_reference_routes_to_extra_params(self, tmp_path): generator = _StubVisualGen() @@ -357,7 +357,7 @@ def test_multipart_image_reference_routes_to_image(self, tmp_path): assert str(params.image).endswith("vid-5_reference.png") def test_undecodable_reference_raises_and_cleans_up(self, tmp_path): - pytest.importorskip("av") + pytest.importorskip("cv2") generator = _StubVisualGen() b64 = base64.b64encode(b"neither an image nor a video").decode() request = VideoGenerationRequest(prompt="x", input_reference=b64) From 8a524ed624004503fe78328ca15396569e44a6f1 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 15 Jul 2026 09:10:00 -0700 Subject: [PATCH 26/64] Drop type-suffix from stored reference files Reference files are now stored without a `.png`/`.mp4` suffix. Content classification uses `is_image_file`/`is_video_file` (PIL and OpenCV header probes) instead of file extension, matching what the serve path already produces when writing raw bytes. The two helpers are lifted from `visual_gen_utils.py` into `media_io.py` for reuse by the Cosmos3 `normalize_video_input_path` utility. Signed-off-by: Igor Shovkun --- .../_torch/visual_gen/models/cosmos3/utils.py | 17 +++--- tensorrt_llm/inputs/media_io.py | 33 +++++++++++- tensorrt_llm/serve/openai_protocol.py | 2 +- tensorrt_llm/serve/visual_gen_utils.py | 47 ++--------------- .../visual_gen/test_cosmos3_pipeline.py | 52 +++++++++++++++++++ .../visual_gen/test_trtllm_serve_endpoints.py | 8 +-- .../visual_gen/test_visual_gen_utils.py | 21 ++++---- 7 files changed, 115 insertions(+), 65 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py index ee02ec6cf127..ba75faee32a0 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py @@ -11,7 +11,6 @@ import PIL.Image IMAGE_EXTENSIONS = frozenset({".png", ".jpg", ".jpeg", ".webp", ".bmp"}) -VIDEO_EXTENSIONS = frozenset({".mp4", ".avi"}) def pil_to_rgb(value: Any) -> PIL.Image.Image: @@ -57,15 +56,19 @@ def normalize_video_input_path(path: Path, max_frames: Optional[int] = None) -> frame_paths = frame_paths[:max_frames] return frame_paths - suffix = path.suffix.lower() - if suffix in IMAGE_EXTENSIONS: + # Classify a single file by content, not by suffix: a decodable still is + # one conditioning frame; a decodable video is expanded to its frames. + # Extensions are unreliable — the serve path stores references with no + # type-suffix at all — so the container decides, not the name. + from tensorrt_llm.inputs.media_io import is_image_file, is_video_file + + if is_image_file(path): return [str(path)] - if suffix in VIDEO_EXTENSIONS: + if is_video_file(path): return decode_video_file(path, max_frames=max_frames) raise ValueError( - "Cosmos3 video path must be a frame directory, an image file " - f"{sorted(IMAGE_EXTENSIONS)}, or a video file " - f"{sorted(VIDEO_EXTENSIONS)}; got {path}" + f"Cosmos3 reference must be a frame directory, a decodable image, " + f"or a decodable video; got {path}" ) diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index a618f3f0a540..e2856737ca15 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -38,7 +38,7 @@ import torch from blake3 import blake3 from packaging.version import Version -from PIL import Image +from PIL import Image, UnidentifiedImageError from tensorrt_llm.inputs.multimodal_data import AudioData, VideoData from tensorrt_llm.logger import logger @@ -346,6 +346,37 @@ def _get_cv2(): return cv2 +def is_image_file(path) -> bool: + """True when ``path`` holds still-image content (anything PIL opens). + + Header-only probe: identifies the container without decoding pixels. + Lets callers classify a reference by content instead of by file suffix. + """ + try: + with Image.open(path): + return True + except UnidentifiedImageError: + return False + + +def is_video_file(path) -> bool: + """True when ``path`` holds a decodable video stream (OpenCV-openable). + + Total predicate over content: False for images, audio, and undecodable + data alike. Stills are excluded explicitly — FFmpeg demuxes a single + image as a one-frame video stream, so a bare video probe would accept + every PNG/JPEG. ``_get_cv2`` raises a clear install hint if cv2 is absent. + """ + if is_image_file(path): + return False + cv2 = _get_cv2() + capture = cv2.VideoCapture(str(path)) + try: + return bool(capture.isOpened() and capture.read()[0]) + finally: + capture.release() + + def _select_cv2_stream_buffered_backend() -> Optional[int]: """Return a VideoCapture backend that can read from a Python `BytesIO`. diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index 529821302eaf..300c6a53e375 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -1690,7 +1690,7 @@ class VideoGenerationRequest(OpenAIBaseModel): "Optional image or video reference that guides generation. " "Content is classified by decoding, not by extension or " "content-type: images (anything PIL reads) condition " - "image-to-video; videos (anything PyAV reads) condition " + "image-to-video; videos (anything OpenCV reads) condition " "video-to-video on models that support it. JSON requests " "carry base64 bytes; multipart requests upload the file."), ) diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index 20d27990cc0e..d86bf8975536 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -4,9 +4,7 @@ import shutil from typing import Any, Dict, List, Optional -from PIL import Image, UnidentifiedImageError - -from tensorrt_llm.inputs.media_io import _get_cv2 +from tensorrt_llm.inputs.media_io import is_image_file, is_video_file from tensorrt_llm.logger import logger from tensorrt_llm.serve.openai_protocol import ImageGenerationRequest, VideoGenerationRequest from tensorrt_llm.visual_gen import VisualGen, VisualGenParams @@ -87,39 +85,6 @@ def _merge_extra_params( params.extra_params = None -def _reference_is_image(path: str) -> bool: - """True when ``path`` holds image content (anything PIL can open). - - Capability-based: the supported set is whatever the decoder accepts, - with no enumerated format table. The probe parses only the header — - pixel decode stays in the worker. - """ - try: - with Image.open(path): - return True - except UnidentifiedImageError: - return False - - -def _reference_is_video(path: str) -> bool: - """True when ``path`` holds video content (an OpenCV-decodable video stream). - - Total predicate over *content*: False for images, audio, and undecodable - content alike. Images must be excluded explicitly because FFmpeg demuxes a - still image as a valid single-frame video stream, so a video probe alone - would claim every PNG/JPEG. OpenCV is the optional decoder shared with the - multimodal video path; ``_get_cv2`` raises a clear install hint if missing. - """ - if _reference_is_image(path): - return False - cv2 = _get_cv2() - capture = cv2.VideoCapture(path) - try: - return bool(capture.isOpened() and capture.read()[0]) - finally: - capture.release() - - def parse_visual_gen_params( request: ImageGenerationRequest | VideoGenerationRequest, id: str, @@ -204,18 +169,16 @@ def parse_visual_gen_params( else: with open(tmp_path, "wb") as f: shutil.copyfileobj(request.input_reference.file, f) - # image, video, or reject. - if _reference_is_image(tmp_path): + + if is_image_file(tmp_path): is_video = False - elif _reference_is_video(tmp_path): + elif is_video_file(tmp_path): is_video = True else: raise ValueError( "input_reference content is neither a decodable image nor a decodable video." ) - ref_path = os.path.join( - media_storage_path, f"{id}_reference{'.mp4' if is_video else '.png'}" - ) + ref_path = os.path.join(media_storage_path, f"{id}_reference") os.replace(tmp_path, ref_path) except Exception: # Cleanup-and-reraise, not handling: every failure path — diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index 56626844a0d1..f66f08fec5c9 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -583,6 +583,58 @@ def test_invalid_condition_video_keep_raises(self): _normalize_condition_video_keep("middle") +class TestNormalizeVideoInputContentDispatch: + """``normalize_video_input_path`` classifies files by content, not suffix. + + The serve path stores references with no type-suffix, so decode dispatch + must key on the container, not the filename. CPU-only; the clip is + synthesized with OpenCV (``cv2`` ships in no CI image → importorskip). + """ + + @staticmethod + def _write_mp4(path, num_frames: int = 3) -> None: + cv2 = pytest.importorskip("cv2") + np = pytest.importorskip("numpy") + writer = cv2.VideoWriter(str(path), cv2.VideoWriter_fourcc(*"mp4v"), 4.0, (16, 16)) + try: + for _ in range(num_frames): + writer.write(np.zeros((16, 16, 3), dtype=np.uint8)) + finally: + writer.release() + assert path.exists() and path.stat().st_size > 0 + + def test_video_without_extension_decodes_to_frames(self, tmp_path): + pytest.importorskip("cv2") + from tensorrt_llm._torch.visual_gen.models.cosmos3.utils import normalize_video_input_path + + # OpenCV's *writer* selects the muxer by extension, so encode to a + # ``.mp4`` path, then rename to a suffix-less name — exactly what the + # serve path produces (raw bytes written to ``{id}_reference``). + encoded = tmp_path / "clip.mp4" + self._write_mp4(encoded) + ref = encoded.rename(tmp_path / "reference") + frames = normalize_video_input_path(ref) + # Took the video-decode path (returns PIL frames), not the single-still + # path (which would return the bare ``[str(path)]``). + assert frames and all(isinstance(f, PIL.Image.Image) for f in frames) + + def test_image_without_extension_is_single_frame(self, tmp_path): + from tensorrt_llm._torch.visual_gen.models.cosmos3.utils import normalize_video_input_path + + ref = tmp_path / "reference" # no ``.png`` suffix + PIL.Image.new("RGB", (8, 8), (1, 2, 3)).save(ref, format="PNG") + assert normalize_video_input_path(ref) == [str(ref)] + + def test_undecodable_file_raises(self, tmp_path): + pytest.importorskip("cv2") + from tensorrt_llm._torch.visual_gen.models.cosmos3.utils import normalize_video_input_path + + ref = tmp_path / "reference" + ref.write_bytes(b"not media") + with pytest.raises(ValueError, match="decodable"): + normalize_video_input_path(ref) + + @pytest.mark.integration @pytest.mark.cosmos3_v2v @pytest.mark.high_cuda_memory diff --git a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py index 2f0e21bbebbe..cc266a9d556a 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py @@ -63,7 +63,7 @@ def _assert_llm_envelope( def _require_opencv(): """Skip unless OpenCV is installed (the shared optional video decoder). - These tests drive ``_reference_is_video``, which decodes references via + These tests drive ``is_video_file``, which decodes references via OpenCV (the same optional dep the multimodal video path uses). ``cv2`` ships in no CI image, so this ``importorskip``-s it and returns the module for tests that synthesize a clip with ``cv2.VideoWriter``. @@ -850,7 +850,7 @@ def test_sync_video_generation_multipart_with_reference(self, video_client, tmp_ # through as params.image (a filesystem path). params = video_client.mock_gen.last_params assert isinstance(params.image, str) - assert params.image.endswith("_reference.png") + assert params.image.endswith("_reference") assert os.path.exists(params.image) def test_sync_video_generation_multipart_with_video_reference(self, video_client, tmp_path): @@ -887,12 +887,12 @@ def test_sync_video_generation_multipart_with_video_reference(self, video_client # Video content must NOT land on params.image; it rides # extra_params["video"] (the same pipeline entry the offline - # example's --video_path uses), stored with a .mp4 suffix. + # example's --video_path uses), stored without a type-suffix. params = video_client.mock_gen.last_params assert params.image is None assert isinstance(params.extra_params, dict) video_ref = params.extra_params["video"] - assert video_ref.endswith("_reference.mp4") + assert video_ref.endswith("_reference") assert os.path.exists(video_ref) def test_sync_video_generation_undecodable_reference_400(self, video_client): diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index 2d26f6a13fa3..1da8a2e373e0 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -277,7 +277,7 @@ def test_base64_reference_written_to_disk(self, tmp_path): request, "vid-1", generator, media_storage_path=str(tmp_path) ) assert params.image is not None - assert str(params.image).endswith("vid-1_reference.png") + assert str(params.image).endswith("vid-1_reference") # The decoded image is identical to what we passed in. with open(params.image, "rb") as f: decoded = Image.open(f).convert("RGB") @@ -320,12 +320,13 @@ def test_multipart_video_reference_routes_to_extra_params(self, tmp_path): params = parse_visual_gen_params( request, "vid-3", generator, media_storage_path=str(tmp_path) ) - # Video content routes to extra_params["video"], not params.image, - # and the written suffix drives the worker's decode dispatch. + # Video content routes to extra_params["video"], not params.image. + # The stored file has no type-suffix — the worker classifies by + # content, so a fabricated extension would assert a type, not hint one. assert params.image is None assert params.extra_params is not None - assert str(params.extra_params["video"]).endswith("vid-3_reference.mp4") - assert (tmp_path / "vid-3_reference.mp4").exists() + assert str(params.extra_params["video"]).endswith("vid-3_reference") + assert (tmp_path / "vid-3_reference").exists() def test_base64_video_reference_routes_to_extra_params(self, tmp_path): # Classification is content-based, so the JSON/base64 path can @@ -337,12 +338,12 @@ def test_base64_video_reference_routes_to_extra_params(self, tmp_path): request, "vid-4", generator, media_storage_path=str(tmp_path) ) assert params.image is None - assert str(params.extra_params["video"]).endswith("vid-4_reference.mp4") + assert str(params.extra_params["video"]).endswith("vid-4_reference") def test_multipart_image_reference_routes_to_image(self, tmp_path): - # JPEG upload: content sniffing classifies it as an image even - # though the stored name is the cosmetic ``.png`` (PIL identifies - # by content, not suffix). + # JPEG upload: content sniffing classifies it as an image and routes + # to params.image. The stored file has no type-suffix (PIL identifies + # by content, not name). generator = _StubVisualGen() img = Image.new("RGB", (4, 4), (10, 20, 30)) buf = BytesIO() @@ -354,7 +355,7 @@ def test_multipart_image_reference_routes_to_image(self, tmp_path): request, "vid-5", generator, media_storage_path=str(tmp_path) ) assert params.extra_params is None - assert str(params.image).endswith("vid-5_reference.png") + assert str(params.image).endswith("vid-5_reference") def test_undecodable_reference_raises_and_cleans_up(self, tmp_path): pytest.importorskip("cv2") From 29f64a2f704b21c2a4eda266323e12156e956381 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 16 Jul 2026 15:13:40 -0700 Subject: [PATCH 27/64] Add .zed and .plans/ to .gitignore Signed-off-by: Igor Shovkun --- .gitignore | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.gitignore b/.gitignore index 9a9332335ca5..512373ae8788 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,7 @@ __pycache__/ .mypy_cache/ .vscode +.zed .cursor *.engine *.engine.config @@ -121,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/ From 7e20e1e8166e31bfeb3f6eb3bb09d5ec58514639 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 16 Jul 2026 15:23:30 -0700 Subject: [PATCH 28/64] Refactor Cosmos3 V2V to use `multi_modal_data` for video reference Replace the `extra_params["video"]` path-based approach with decoded `VideoData` under `multi_modal_data["video"]`, aligning V2V with the framework's multimodal convention. The worker now receives pre-decoded frames and VAE-encodes them directly. Also rename `condition_frame_indexes_vision` to `condition_video_latent_indexes` to clarify that the indexes refer to output latent frames, not source frame selection. Signed-off-by: Igor Shovkun --- examples/visual_gen/models/cosmos3/README.md | 2 +- examples/visual_gen/models/cosmos3/cosmos3.py | 3 +- examples/visual_gen/serve/README.md | 4 +- .../visual_gen/models/cosmos3/defaults.py | 22 ++--- .../models/cosmos3/pipeline_cosmos3.py | 47 ++++++---- .../_torch/visual_gen/models/cosmos3/utils.py | 39 ++++----- tensorrt_llm/inputs/media_io.py | 27 ++++++ tensorrt_llm/serve/visual_gen_utils.py | 13 ++- tensorrt_llm/visual_gen/params.py | 11 ++- .../visual_gen/test_cosmos3_pipeline.py | 87 +++++++++++++++---- .../visual_gen/test_trtllm_serve_endpoints.py | 18 ++-- .../visual_gen/test_visual_gen_params.py | 24 ----- .../visual_gen/test_visual_gen_utils.py | 22 +++-- 13 files changed, 199 insertions(+), 120 deletions(-) diff --git a/examples/visual_gen/models/cosmos3/README.md b/examples/visual_gen/models/cosmos3/README.md index aff88d35d04b..60be9a991970 100644 --- a/examples/visual_gen/models/cosmos3/README.md +++ b/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 frame directory or `.mp4`/`.avi` file). Only the first (or last, per `condition_video_keep`) `max(condition_frame_indexes_vision) * 4 + 1` input frames condition the output (5 by default); `.mp4`/`.avi` decode uses OpenCV (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 frame directory or `.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); `.mp4`/`.avi` decode uses OpenCV (see [Media I/O dependencies](#media-io-dependencies)). - **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 diff --git a/examples/visual_gen/models/cosmos3/cosmos3.py b/examples/visual_gen/models/cosmos3/cosmos3.py index d1ea4a0448b0..844d17d17741 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3.py +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -107,6 +107,7 @@ from typing import Any, Dict, Optional from tensorrt_llm import VisualGen, VisualGenArgs +from tensorrt_llm._torch.visual_gen.models.cosmos3.utils import load_reference_video _SCRIPT_DIR = Path(__file__).resolve().parent @@ -279,7 +280,7 @@ def main(): params.extra_params["output_type"] = output_type if args.video_path is not None: - params.extra_params["video"] = args.video_path + params.multi_modal_data = {"video": load_reference_video(args.video_path)} if negative_prompt is None: params.negative_prompt = None diff --git a/examples/visual_gen/serve/README.md b/examples/visual_gen/serve/README.md index b267a1923205..4cf5310ef6c6 100644 --- a/examples/visual_gen/serve/README.md +++ b/examples/visual_gen/serve/README.md @@ -315,7 +315,7 @@ 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_frame_indexes_vision`, `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`, ... > **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`). @@ -367,7 +367,7 @@ curl -X POST "http://localhost:8000/v1/videos" \ -F "input_reference=@./media/reference.mp4" \ -F "num_frames=189" \ -F "fps=24" \ - -F 'extra_params={"condition_frame_indexes_vision": [0, 1], "condition_video_keep": "first"}' + -F 'extra_params={"condition_video_latent_indexes": [0, 1], "condition_video_keep": "first"}' ``` ### Check Video Status diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index 3928817951b9..311844e0f3c1 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -35,7 +35,7 @@ "frame_rate": 24.0, } -COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION = (0, 1) +COSMOS3_DEFAULT_CONDITION_VIDEO_LATENT_INDEXES = (0, 1) COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP = "first" # Fields merged by the executor for every request. Modality-specific values @@ -95,10 +95,15 @@ default="video", description="Output modality: 'video' (T2V/I2V) or 'image' (text-to-image).", ), - "condition_frame_indexes_vision": ExtraParamSchema( + "condition_video_latent_indexes": ExtraParamSchema( type="list", - default=list(COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION), - description="Latent frame indexes to keep fixed for video conditioning.", + default=list(COSMOS3_DEFAULT_CONDITION_VIDEO_LATENT_INDEXES), + description=( + "Latent frame indexes OF THE OUTPUT video to pin to the encoded " + "reference (not source-frame selection). Each latent frame spans 4 " + "pixel frames, so the worker consumes the first (or last, per " + "condition_video_keep) max(indexes)*4+1 reference frames." + ), ), "condition_video_keep": ExtraParamSchema( type="str", @@ -110,13 +115,4 @@ default=None, description="Optional scheduler flow shift override. Uses the Cosmos3 mode default when omitted.", ), - "video": ExtraParamSchema( - type="path_or_list", - default=None, - description=( - "Video input for video-to-video generation: " - ".mp4/.avi file, frame directory, image path, or list of PIL " - "images / frame paths." - ), - ), } diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index ec3d1516e1f9..f05f88f86fd4 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -36,8 +36,8 @@ from .defaults import ( COSMOS3_720P_PARAMS, - COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION, COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP, + COSMOS3_DEFAULT_CONDITION_VIDEO_LATENT_INDEXES, COSMOS3_EXTRA_SPECS, COSMOS3_PIPELINE_DEFAULTS, COSMOS3_T2I_PARAMS, @@ -62,11 +62,11 @@ TRTLLM_DISABLE_COSMOS3_GUARDRAILS = os.environ.get("TRTLLM_DISABLE_COSMOS3_GUARDRAILS", "0") == "1" -def _normalize_condition_frame_indexes_vision( +def _normalize_condition_video_latent_indexes( indexes: Iterable[int] | int | str | None, ) -> tuple[int, ...]: if indexes is None: - return COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION + return COSMOS3_DEFAULT_CONDITION_VIDEO_LATENT_INDEXES if isinstance(indexes, int): normalized = (indexes,) elif isinstance(indexes, str): @@ -76,19 +76,19 @@ def _normalize_condition_frame_indexes_vision( normalized = tuple(int(index) for index in indexes) if not normalized: - raise ValueError("Cosmos3 condition_frame_indexes_vision must not be empty.") + raise ValueError("Cosmos3 condition_video_latent_indexes must not be empty.") if any(index < 0 for index in normalized): raise ValueError( - f"Cosmos3 condition_frame_indexes_vision must be non-negative, got {normalized}." + f"Cosmos3 condition_video_latent_indexes must be non-negative, got {normalized}." ) return normalized def _condition_pixel_frame_count( - condition_frame_indexes_vision: Iterable[int], + condition_video_latent_indexes: Iterable[int], temporal_compression: int, ) -> int: - return max(condition_frame_indexes_vision) * int(temporal_compression) + 1 + return max(condition_video_latent_indexes) * int(temporal_compression) + 1 def _normalize_condition_video_keep(keep: str | None) -> str: @@ -309,6 +309,14 @@ def infer(self, req): extra_params = req.params.extra_params or {} output_type = extra_params.get("output_type", "video") + # The V2V reference rides in ``multi_modal_data["video"]`` as a + # ``VideoData`` (framework convention); the worker crops + VAE-encodes its + # frames. Both producers (offline and serve) build it, so there is no + # legacy-path fallback. + mm_data = req.params.multi_modal_data or {} + video_data = mm_data.get("video") + video = video_data.frames if video_data is not None else None + return self.forward( prompt=req.prompt, negative_prompt=req.params.negative_prompt, @@ -335,8 +343,8 @@ def infer(self, req): use_guardrails=extra_params.get("use_guardrails", True), enable_audio=extra_params.get("enable_audio", False), output_type=output_type, - video=extra_params.get("video"), - condition_frame_indexes_vision=extra_params.get("condition_frame_indexes_vision"), + video=video, + condition_video_latent_indexes=extra_params.get("condition_video_latent_indexes"), condition_video_keep=extra_params.get("condition_video_keep"), flow_shift=extra_params.get("flow_shift"), ) @@ -679,7 +687,7 @@ def _prepare_latents_v2v( video_tensor: torch.Tensor, num_frames: int, generator: torch.Generator, - condition_frame_indexes_vision: Iterable[int] | int | str | None = None, + condition_video_latent_indexes: Iterable[int] | int | str | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Prepare V2V latents with explicit clean conditioned latent frames.""" if video_tensor.ndim == 4: @@ -696,11 +704,11 @@ def _prepare_latents_v2v( T_lat = (num_frames - 1) // self.vae_scale_factor_temporal + 1 H_lat = video_tensor.shape[-2] // self.vae_scale_factor_spatial W_lat = video_tensor.shape[-1] // self.vae_scale_factor_spatial - indexes = _normalize_condition_frame_indexes_vision(condition_frame_indexes_vision) + indexes = _normalize_condition_video_latent_indexes(condition_video_latent_indexes) out_of_range = [index for index in indexes if index >= T_lat] if out_of_range: raise ValueError( - "Cosmos3 condition_frame_indexes_vision contains indexes outside the latent video: " + "Cosmos3 condition_video_latent_indexes contains indexes outside the latent video: " f"indexes={indexes}, latent_frames={T_lat}." ) @@ -769,7 +777,7 @@ def forward( enable_audio: bool = COSMOS3_EXTRA_SPECS["enable_audio"].default, output_type: str = COSMOS3_EXTRA_SPECS["output_type"].default, video: Any = None, - condition_frame_indexes_vision: Any = None, + condition_video_latent_indexes: Any = None, condition_video_keep: Any = None, flow_shift: Optional[float] = None, ): @@ -962,16 +970,19 @@ def forward( image, height=height, width=width, num_frames=num_frames, generator=generator ) elif video is not None: - condition_frame_indexes_vision = _normalize_condition_frame_indexes_vision( - condition_frame_indexes_vision + condition_video_latent_indexes = _normalize_condition_video_latent_indexes( + condition_video_latent_indexes ) condition_video_keep = _normalize_condition_video_keep(condition_video_keep) condition_pixel_frames = min( _condition_pixel_frame_count( - condition_frame_indexes_vision, self.vae_scale_factor_temporal + condition_video_latent_indexes, self.vae_scale_factor_temporal ), num_frames, ) + # ``video`` is the reference frames (a list, from the producer's + # ``VideoData``). The worker crops the first/last conditioning window + # uniformly — producers never crop, so behavior never depends on them. video = normalize_video_input( video, max_frames=None if condition_video_keep == "last" else condition_pixel_frames, @@ -986,13 +997,13 @@ def forward( if self.rank == 0: logger.info( f"Cosmos3 V2V conditioning: frames={video.shape[2]}, " - f"latent_indexes={condition_frame_indexes_vision}" + f"latent_indexes={condition_video_latent_indexes}" ) latents, velocity_mask, condition_latents = self._prepare_latents_v2v( video, num_frames=num_frames, generator=generator, - condition_frame_indexes_vision=condition_frame_indexes_vision, + condition_video_latent_indexes=condition_video_latent_indexes, ) else: latents = self._prepare_latents(height, width, num_frames, generator) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py index ba75faee32a0..40c7a4e7171a 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py @@ -22,26 +22,9 @@ def pil_to_rgb(value: Any) -> PIL.Image.Image: def decode_video_file(path: Path, max_frames: Optional[int] = None) -> List[PIL.Image.Image]: - # OpenCV is the shared multimodal video decoder; ``_get_cv2`` raises a clear - # ``pip install opencv-python-headless`` hint when it is not installed. - from tensorrt_llm.inputs.media_io import _get_cv2 - - cv2 = _get_cv2() - capture = cv2.VideoCapture(str(path)) - try: - if not capture.isOpened(): - raise ValueError(f"Cosmos3 could not open video file: {path}") - frames: List[PIL.Image.Image] = [] - while max_frames is None or len(frames) < max_frames: - ok, frame = capture.read() - if not ok: - break - frames.append(PIL.Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))) - finally: - capture.release() - if not frames: - raise ValueError(f"Cosmos3 video file contains no frames: {path}") - return frames + from tensorrt_llm.inputs.media_io import decode_video_frames + + return decode_video_frames(path, max_frames=max_frames) def normalize_video_input_path(path: Path, max_frames: Optional[int] = None) -> List[Any]: @@ -89,3 +72,19 @@ def normalize_video_input(video: Any, max_frames: Optional[int] = None) -> List[ if isinstance(video, (str, Path)): return normalize_video_input_path(Path(video), max_frames=max_frames) return [video] + + +def load_reference_video(src: Any): + """Decode a reference into the framework's ``VideoData`` (all frames). + + Offline entry point: decodes any supported reference into ``VideoData`` for + ``multi_modal_data["video"]``. It does **not** crop and carries no temporal + metadata — the Cosmos3 worker crops the first/last conditioning window, and + the temporal placement is set by ``condition_video_latent_indexes`` + + ``frame_rate`` (request params), not the reference. The worker reads + ``.frames`` and VAE-encodes them; it never decodes media itself. + """ + from tensorrt_llm.inputs.multimodal_data import VideoData + + frames = normalize_video_input(src, max_frames=None) + return VideoData(frames=[pil_to_rgb(frame) for frame in frames], metadata={}) diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index e2856737ca15..b59c46bf970c 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -377,6 +377,33 @@ def is_video_file(path) -> bool: capture.release() +def decode_video_frames(path, max_frames: Optional[int] = None) -> List["Image.Image"]: + """Decode a video file into its frames as PIL images, in order, no sampling. + + Unlike :func:`load_video` — which samples ``num_frames`` evenly for + VLM-style inputs — this preserves every frame sequentially from the start; + ``max_frames`` bounds the decode when only a prefix is needed. + Reference-conditioning consumers (e.g. video-to-video pipelines) pick + their own frame window downstream. + """ + cv2 = _get_cv2() + capture = cv2.VideoCapture(str(path)) + try: + if not capture.isOpened(): + raise ValueError(f"Could not open video file: {path}") + frames = [] + while max_frames is None or len(frames) < max_frames: + ok, frame = capture.read() + if not ok: + break + frames.append(Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))) + finally: + capture.release() + if not frames: + raise ValueError(f"Video file contains no frames: {path}") + return frames + + def _select_cv2_stream_buffered_backend() -> Optional[int]: """Return a VideoCapture backend that can read from a Python `BytesIO`. diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index d86bf8975536..a9589539e158 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -4,7 +4,8 @@ import shutil from typing import Any, Dict, List, Optional -from tensorrt_llm.inputs.media_io import is_image_file, is_video_file +from tensorrt_llm.inputs.media_io import decode_video_frames, is_image_file, is_video_file +from tensorrt_llm.inputs.multimodal_data import VideoData from tensorrt_llm.logger import logger from tensorrt_llm.serve.openai_protocol import ImageGenerationRequest, VideoGenerationRequest from tensorrt_llm.visual_gen import VisualGen, VisualGenParams @@ -188,9 +189,13 @@ def parse_visual_gen_params( os.remove(tmp_path) raise if is_video: - if params.extra_params is None: - params.extra_params = {} - params.extra_params["video"] = ref_path + # V2V reference: decode into ``VideoData`` under + # ``multi_modal_data["video"]`` — the one intake video-capable + # pipelines read. The decode is model-agnostic (all frames, no + # crop); each pipeline picks its own conditioning window. + params.multi_modal_data = { + "video": VideoData(frames=decode_video_frames(ref_path), metadata={}) + } else: params.image = ref_path diff --git a/tensorrt_llm/visual_gen/params.py b/tensorrt_llm/visual_gen/params.py index 3b00d4080ee3..88251fd2056d 100644 --- a/tensorrt_llm/visual_gen/params.py +++ b/tensorrt_llm/visual_gen/params.py @@ -12,7 +12,6 @@ # 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. -import os from typing import Any, Dict, List, Optional, Union from pydantic import Field @@ -73,6 +72,14 @@ class VisualGenParams(StrictBaseModel): image: Optional[Union[str, bytes, List[Union[str, bytes]]]] = Field( default=None, description="Reference image(s) for I2V/I2I." ) + # Framework multimodal convention: modality -> data. Cosmos3 V2V carries its + # reference here as ``{"video": VideoData(frames, metadata)}``; the worker + # reads the frames and VAE-encodes them (see ``inputs.multimodal_data``). + multi_modal_data: Optional[Dict[str, Any]] = Field( + default=None, + description="Multimodal conditioning inputs keyed by modality " + "(e.g. {'video': VideoData}; see inputs.multimodal_data).", + ) # Per-prompt multiplier num_images_per_prompt: int = Field(default=1, description="Number of images per prompt.") @@ -94,8 +101,6 @@ class VisualGenParams(StrictBaseModel): "bool": (bool,), "str": (str,), "list": (list,), - # Cosmos3 V2V `video` reference: a media path or an in-memory frame list. - "path_or_list": (str, os.PathLike, list), } # Generation config fields that pipelines declare defaults for. If a user diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index f66f08fec5c9..7bd77d13dee9 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -38,8 +38,8 @@ import torch from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( - COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION, COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP, + COSMOS3_DEFAULT_CONDITION_VIDEO_LATENT_INDEXES, COSMOS3_EXTRA_SPECS, COSMOS3_T2I_PARAMS, ) @@ -50,8 +50,8 @@ COSMOS3_IMAGE_RESOLUTION_TEMPLATE, Cosmos3OmniMoTPipeline, _condition_pixel_frame_count, - _normalize_condition_frame_indexes_vision, _normalize_condition_video_keep, + _normalize_condition_video_latent_indexes, ) from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader from tensorrt_llm.visual_gen.args import TorchCompileConfig, VisualGenArgs @@ -517,8 +517,8 @@ def test_i2v_smoke(self, cosmos3_pipeline): class TestCosmos3V2VExtraParams: def test_condition_defaults_are_declared(self): - assert COSMOS3_EXTRA_SPECS["condition_frame_indexes_vision"].default == list( - COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION + assert COSMOS3_EXTRA_SPECS["condition_video_latent_indexes"].default == list( + COSMOS3_DEFAULT_CONDITION_VIDEO_LATENT_INDEXES ) assert ( COSMOS3_EXTRA_SPECS["condition_video_keep"].default @@ -530,11 +530,6 @@ def test_flow_shift_default_is_request_optional(self): assert spec.type == "float" assert spec.default is None - def test_video_spec_declares_path_or_list_input(self): - spec = COSMOS3_EXTRA_SPECS["video"] - assert spec.type == "path_or_list" - assert spec.default is None - class TestCosmos3V2VConditioningParams: @pytest.mark.parametrize( @@ -547,13 +542,13 @@ class TestCosmos3V2VConditioningParams: ("0, 2", (0, 2)), ], ) - def test_normalize_condition_frame_indexes_vision(self, value, expected): - assert _normalize_condition_frame_indexes_vision(value) == expected + def test_normalize_condition_video_latent_indexes(self, value, expected): + assert _normalize_condition_video_latent_indexes(value) == expected @pytest.mark.parametrize("value", [[], "", [-1], "0, -1", [0, -2]]) - def test_invalid_condition_frame_indexes_vision_raise(self, value): + def test_invalid_condition_video_latent_indexes_raise(self, value): with pytest.raises(ValueError): - _normalize_condition_frame_indexes_vision(value) + _normalize_condition_video_latent_indexes(value) @pytest.mark.parametrize( "indexes,expected", @@ -635,6 +630,46 @@ def test_undecodable_file_raises(self, tmp_path): normalize_video_input_path(ref) +class TestLoadReferenceVideo: + """``load_reference_video`` decodes any reference into the framework's + ``VideoData`` (all frames) for ``multi_modal_data["video"]``. It does not + crop — the worker does — so the producer stays model-agnostic. CPU-only.""" + + def test_from_pil_list_all_frames(self): + from tensorrt_llm._torch.visual_gen.models.cosmos3.utils import load_reference_video + from tensorrt_llm.inputs.multimodal_data import VideoData + + frames = [PIL.Image.new("RGB", (8, 8), (i, i, i)) for i in range(6)] + vd = load_reference_video(frames) + assert isinstance(vd, VideoData) + assert len(vd.frames) == 6 # all frames, no crop + assert all(isinstance(f, PIL.Image.Image) for f in vd.frames) + # No temporal metadata: placement comes from request params + # (condition_video_latent_indexes, frame_rate), not the reference. + assert vd.metadata == {} + + def test_from_video_file(self, tmp_path): + pytest.importorskip("cv2") + from tensorrt_llm._torch.visual_gen.models.cosmos3.utils import load_reference_video + + enc = tmp_path / "clip.mp4" + TestNormalizeVideoInputContentDispatch._write_mp4(enc, num_frames=8) + assert len(load_reference_video(enc).frames) == 8 # all frames + + def test_from_directory(self, tmp_path): + from tensorrt_llm._torch.visual_gen.models.cosmos3.utils import load_reference_video + + for i in range(5): + PIL.Image.new("RGB", (8, 8), (i, i, i)).save(tmp_path / f"{i:03d}.png") + assert len(load_reference_video(tmp_path).frames) == 5 + + def test_missing_path_raises(self, tmp_path): + from tensorrt_llm._torch.visual_gen.models.cosmos3.utils import load_reference_video + + with pytest.raises(ValueError, match="does not exist"): + load_reference_video(tmp_path / "nope") + + @pytest.mark.integration @pytest.mark.cosmos3_v2v @pytest.mark.high_cuda_memory @@ -646,7 +681,7 @@ def test_v2v_smoke(self, cosmos3_pipeline): image=None, video=video, num_frames=NUM_FRAMES, - condition_frame_indexes_vision=[0, 1], + condition_video_latent_indexes=[0, 1], condition_video_keep="first", ) _assert_valid_video(result.video, num_frames=NUM_FRAMES) @@ -657,6 +692,24 @@ def test_v2v_smoke(self, cosmos3_pipeline): use_karras_sigmas=False, ) + def test_v2v_multimodal_reference_smoke(self, cosmos3_pipeline): + """The V2V reference arrives as ``VideoData`` (built by + ``load_reference_video``, all frames) under ``multi_modal_data["video"]``; + the worker crops the conditioning window and VAE-encodes. Mirrors what the + offline example feeds the pipeline.""" + from tensorrt_llm._torch.visual_gen.models.cosmos3.utils import load_reference_video + + video_data = load_reference_video(_make_test_video(NUM_FRAMES)) + result = _run_forward( + cosmos3_pipeline, + image=None, + video=video_data.frames, + num_frames=NUM_FRAMES, + condition_video_latent_indexes=[0, 1], + condition_video_keep="first", + ) + _assert_valid_video(result.video, num_frames=NUM_FRAMES) + def test_v2v_keep_last_smoke(self, cosmos3_pipeline): """condition_video_keep="last" pins the tail of the input, not the head. @@ -667,14 +720,14 @@ def test_v2v_keep_last_smoke(self, cosmos3_pipeline): """ dark = PIL.Image.new("RGB", (WIDTH, HEIGHT), (40, 40, 40)) bright = PIL.Image.new("RGB", (WIDTH, HEIGHT), (230, 230, 230)) - # 5 = max(condition_frame_indexes_vision) * 4 + 1 conditioning frames. + # 5 = max(condition_video_latent_indexes) * 4 + 1 conditioning frames. video = [dark.copy() for _ in range(NUM_FRAMES)] + [bright.copy() for _ in range(5)] result = _run_forward( cosmos3_pipeline, image=None, video=video, num_frames=NUM_FRAMES, - condition_frame_indexes_vision=[0, 1], + condition_video_latent_indexes=[0, 1], condition_video_keep="last", ) _assert_valid_video(result.video, num_frames=NUM_FRAMES) @@ -793,7 +846,7 @@ def test_v2v_audio_smoke(self, cosmos3_pipeline): cosmos3_pipeline, enable_audio=True, video=_make_test_video(NUM_FRAMES), - condition_frame_indexes_vision=[0, 1], + condition_video_latent_indexes=[0, 1], condition_video_keep="first", ) _assert_valid_video(result.video, num_frames=NUM_FRAMES) diff --git a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py index cc266a9d556a..661e4bcce091 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py @@ -854,7 +854,8 @@ def test_sync_video_generation_multipart_with_reference(self, video_client, tmp_ assert os.path.exists(params.image) def test_sync_video_generation_multipart_with_video_reference(self, video_client, tmp_path): - """A video ``input_reference`` routes to ``extra_params["video"]`` (V2V). + """A video ``input_reference`` is decoded into ``VideoData`` under + ``multi_modal_data["video"]`` (V2V). The reference is classified by decoding its content, so the clip is synthesized in-test with OpenCV — no video asset ships with the repo. @@ -885,15 +886,16 @@ def test_sync_video_generation_multipart_with_video_reference(self, video_client assert resp.status_code == 200 assert len(resp.content) > 0 - # Video content must NOT land on params.image; it rides - # extra_params["video"] (the same pipeline entry the offline - # example's --video_path uses), stored without a type-suffix. + # Video content must NOT land on params.image; it's decoded into + # VideoData under multi_modal_data["video"] (framework convention; the + # same pipeline entry the offline example's --video_path uses). + from tensorrt_llm.inputs.multimodal_data import VideoData + params = video_client.mock_gen.last_params assert params.image is None - assert isinstance(params.extra_params, dict) - video_ref = params.extra_params["video"] - assert video_ref.endswith("_reference") - assert os.path.exists(video_ref) + video_data = params.multi_modal_data["video"] + assert isinstance(video_data, VideoData) + assert len(video_data.frames) >= 1 def test_sync_video_generation_undecodable_reference_400(self, video_client): """Content neither PIL nor OpenCV can decode is rejected at the boundary.""" diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py index fb497bf5b928..e69c81d21f1b 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py @@ -760,30 +760,6 @@ def test_valid_extra_params_accepted(self): req = self._make_request(extra_params={"stg_scale": 0.5}) self._merge_and_validate(executor, req) # should not raise - def test_path_or_list_extra_param_type_enforced(self): - """Cosmos3 V2V's `video` extra param declares type "path_or_list"; the - type map must know it, otherwise the type check is silently skipped - and a bad value reaches the pipeline.""" - from tensorrt_llm._torch.visual_gen.pipeline import ExtraParamSchema - from tensorrt_llm.visual_gen.params import VisualGenParams, validate_visual_gen_params - - specs = {"video": ExtraParamSchema(type="path_or_list", default=None)} - - # A media path and an in-memory frame list are both valid forms. - for good in ("/tmp/clip.mp4", ["frame0.png", "frame1.png"]): - validate_visual_gen_params( - VisualGenParams(extra_params={"video": good}), - declared_defaults={}, - extra_param_specs=specs, - ) - - with pytest.raises(ValueError, match="expected type 'path_or_list'"): - validate_visual_gen_params( - VisualGenParams(extra_params={"video": 42}), - declared_defaults={}, - extra_param_specs=specs, - ) - # --- unsupported universal fields --- def test_num_frames_on_image_pipeline_raises(self): diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index 1da8a2e373e0..79248a14ad6d 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -313,24 +313,28 @@ def _mp4_bytes() -> bytes: finally: os.remove(path) - def test_multipart_video_reference_routes_to_extra_params(self, tmp_path): + def test_multipart_video_reference_routes_to_multimodal(self, tmp_path): + from tensorrt_llm.inputs.multimodal_data import VideoData + generator = _StubVisualGen() upload = UploadFile(file=BytesIO(self._mp4_bytes()), filename="clip.mp4") request = VideoGenerationRequest(prompt="x", input_reference=upload) params = parse_visual_gen_params( request, "vid-3", generator, media_storage_path=str(tmp_path) ) - # Video content routes to extra_params["video"], not params.image. - # The stored file has no type-suffix — the worker classifies by - # content, so a fabricated extension would assert a type, not hint one. + # Video content is decoded into VideoData under multi_modal_data["video"] + # (framework convention), not params.image. The worker crops + VAE-encodes. assert params.image is None - assert params.extra_params is not None - assert str(params.extra_params["video"]).endswith("vid-3_reference") - assert (tmp_path / "vid-3_reference").exists() + assert params.multi_modal_data is not None + video_data = params.multi_modal_data["video"] + assert isinstance(video_data, VideoData) + assert len(video_data.frames) >= 1 - def test_base64_video_reference_routes_to_extra_params(self, tmp_path): + def test_base64_video_reference_routes_to_multimodal(self, tmp_path): # Classification is content-based, so the JSON/base64 path can # carry video even though it has no content-type or filename. + from tensorrt_llm.inputs.multimodal_data import VideoData + generator = _StubVisualGen() b64 = base64.b64encode(self._mp4_bytes()).decode() request = VideoGenerationRequest(prompt="x", input_reference=b64) @@ -338,7 +342,7 @@ def test_base64_video_reference_routes_to_extra_params(self, tmp_path): request, "vid-4", generator, media_storage_path=str(tmp_path) ) assert params.image is None - assert str(params.extra_params["video"]).endswith("vid-4_reference") + assert isinstance(params.multi_modal_data["video"], VideoData) def test_multipart_image_reference_routes_to_image(self, tmp_path): # JPEG upload: content sniffing classifies it as an image and routes From a5843aa4b18daf3fbfd120c18cc0c805d3047b08 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 20 Jul 2026 10:24:34 -0700 Subject: [PATCH 29/64] Fix numpy import: use top-level import instead of importorskip Signed-off-by: Igor Shovkun --- .../_torch/visual_gen/test_cosmos3_pipeline.py | 5 +++-- .../visual_gen/test_trtllm_serve_endpoints.py | 13 ++++++++----- 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index 7bd77d13dee9..ef0eed637979 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -33,6 +33,7 @@ os.environ["TLLM_DISABLE_MPI"] = "1" os.environ["TRTLLM_DISABLE_COSMOS3_GUARDRAILS"] = "1" +import numpy as np import PIL.Image import pytest import torch @@ -583,13 +584,13 @@ class TestNormalizeVideoInputContentDispatch: The serve path stores references with no type-suffix, so decode dispatch must key on the container, not the filename. CPU-only; the clip is - synthesized with OpenCV (``cv2`` ships in no CI image → importorskip). + synthesized with OpenCV (installed by CI test stages via + ``jenkins/L0_Test.groovy``; the importorskip only spares bare local envs). """ @staticmethod def _write_mp4(path, num_frames: int = 3) -> None: cv2 = pytest.importorskip("cv2") - np = pytest.importorskip("numpy") writer = cv2.VideoWriter(str(path), cv2.VideoWriter_fourcc(*"mp4v"), 4.0, (16, 16)) try: for _ in range(num_frames): diff --git a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py index 661e4bcce091..136e4f8b27dd 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py @@ -20,6 +20,7 @@ from typing import Optional from unittest.mock import patch +import numpy as np import pytest import torch from fastapi.testclient import TestClient @@ -63,10 +64,13 @@ def _assert_llm_envelope( def _require_opencv(): """Skip unless OpenCV is installed (the shared optional video decoder). - These tests drive ``is_video_file``, which decodes references via - OpenCV (the same optional dep the multimodal video path uses). ``cv2`` ships - in no CI image, so this ``importorskip``-s it and returns the module for - tests that synthesize a clip with ``cv2.VideoWriter``. + These tests drive ``is_video_file``, which decodes references via OpenCV + (the same optional dep the multimodal video path uses). CI installs + ``opencv-python-headless`` in every test stage (``jenkins/L0_Test.groovy``), + so these tests always run there; the skip only spares bare local + environments, where cv2 stays optional (kept out of requirements by the + dependency policy). Returns the module for tests that synthesize a clip + with ``cv2.VideoWriter``. """ return pytest.importorskip("cv2") @@ -861,7 +865,6 @@ def test_sync_video_generation_multipart_with_video_reference(self, video_client synthesized in-test with OpenCV — no video asset ships with the repo. """ cv2 = _require_opencv() - np = pytest.importorskip("numpy") ref_path = tmp_path / "ref.mp4" # mp4v is a built-in FFmpeg mpeg4 encoder present in the opencv wheel. writer = cv2.VideoWriter(str(ref_path), cv2.VideoWriter_fourcc(*"mp4v"), 4.0, (16, 16)) From 6eecba3475978480528081b44c66ac7d878500e1 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 20 Jul 2026 12:56:12 -0700 Subject: [PATCH 30/64] Decode video references in memory, skip disk materialization Add `is_image_bytes` and `decode_video_frames_from_bytes` helpers to `media_io.py` so that video references in `parse_visual_gen_params` are classified and decoded entirely in memory. Image references still write one file (the worker-readable path contract), but no temporary file is created for validation or for video content. Refactor `decode_video_frames` to share frame-draining logic via the new `_read_frames_in_order` helper. Video references no longer require `media_storage_path`, and rejected content never touches disk. Signed-off-by: Igor Shovkun --- tensorrt_llm/inputs/media_io.py | 61 +++++++++++++-- tensorrt_llm/serve/visual_gen_utils.py | 74 +++++++++---------- .../visual_gen/test_visual_gen_utils.py | 65 +++++++++++++++- 3 files changed, 153 insertions(+), 47 deletions(-) diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index b59c46bf970c..319457ec2144 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -377,6 +377,30 @@ def is_video_file(path) -> bool: capture.release() +def is_image_bytes(data) -> bool: + """True when ``data`` holds still-image content (anything PIL opens). + + In-memory counterpart of :func:`is_image_file` — header-only probe, no + pixel decode, no filesystem. + """ + try: + with Image.open(BytesIO(data)): + return True + except UnidentifiedImageError: + return False + + +def _read_frames_in_order(cv2, capture, max_frames: Optional[int]) -> List["Image.Image"]: + """Drain an opened ``VideoCapture`` into PIL frames, in order, no sampling.""" + frames = [] + while max_frames is None or len(frames) < max_frames: + ok, frame = capture.read() + if not ok: + break + frames.append(Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))) + return frames + + def decode_video_frames(path, max_frames: Optional[int] = None) -> List["Image.Image"]: """Decode a video file into its frames as PIL images, in order, no sampling. @@ -391,12 +415,7 @@ def decode_video_frames(path, max_frames: Optional[int] = None) -> List["Image.I try: if not capture.isOpened(): raise ValueError(f"Could not open video file: {path}") - frames = [] - while max_frames is None or len(frames) < max_frames: - ok, frame = capture.read() - if not ok: - break - frames.append(Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))) + frames = _read_frames_in_order(cv2, capture, max_frames) finally: capture.release() if not frames: @@ -404,6 +423,36 @@ def decode_video_frames(path, max_frames: Optional[int] = None) -> List["Image.I return frames +def decode_video_frames_from_bytes(data, max_frames: Optional[int] = None) -> List["Image.Image"]: + """Decode raw video bytes into PIL frames, in order, no sampling. + + Fully in-memory when this OpenCV build has a stream-buffered backend + (:func:`_select_cv2_stream_buffered_backend`); otherwise the bytes spill + to an auto-deleted tempfile and take the :func:`decode_video_frames` path. + Raises ``ValueError`` when the bytes are not a decodable video. + """ + cv2 = _get_cv2() + backend = _select_cv2_stream_buffered_backend() + if backend is None: + with tempfile.NamedTemporaryFile() as spill: + spill.write(data) + spill.flush() + return decode_video_frames(spill.name, max_frames=max_frames) + + # cv2 keeps a non-owning view into the buffer; hold it until release(). + buffer = BytesIO(bytes(data)) + capture = cv2.VideoCapture(buffer, backend, []) + try: + if not capture.isOpened(): + raise ValueError(f"Could not open video from <{len(data)} bytes>.") + frames = _read_frames_in_order(cv2, capture, max_frames) + finally: + capture.release() + if not frames: + raise ValueError(f"Video bytes contain no frames (<{len(data)} bytes>).") + return frames + + def _select_cv2_stream_buffered_backend() -> Optional[int]: """Return a VideoCapture backend that can read from a Python `BytesIO`. diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index a9589539e158..029f9052ec63 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -1,10 +1,9 @@ import asyncio import base64 import os -import shutil from typing import Any, Dict, List, Optional -from tensorrt_llm.inputs.media_io import decode_video_frames, is_image_file, is_video_file +from tensorrt_llm.inputs.media_io import decode_video_frames_from_bytes, is_image_bytes from tensorrt_llm.inputs.multimodal_data import VideoData from tensorrt_llm.logger import logger from tensorrt_llm.serve.openai_protocol import ImageGenerationRequest, VideoGenerationRequest @@ -156,48 +155,43 @@ def parse_visual_gen_params( ) params.num_frames = derived if request.input_reference is not None: - if media_storage_path is None: - raise ValueError("media_storage_path is required when input_reference is provided") - tmp_path = os.path.join(media_storage_path, f"{id}_reference.part") - try: - if isinstance(request.input_reference, str): - try: - payload = base64.b64decode(request.input_reference) - except ValueError as exc: - raise ValueError("input_reference is not valid base64 data.") from exc - with open(tmp_path, "wb") as f: - f.write(payload) - else: - with open(tmp_path, "wb") as f: - shutil.copyfileobj(request.input_reference.file, f) - - if is_image_file(tmp_path): - is_video = False - elif is_video_file(tmp_path): - is_video = True - else: + if isinstance(request.input_reference, str): + try: + payload = base64.b64decode(request.input_reference) + except ValueError as exc: + raise ValueError("input_reference is not valid base64 data.") from exc + else: + payload = request.input_reference.file.read() + + # Classify by decoding the bytes in memory — nothing touches disk + # until the modality is known and validated. + if is_image_bytes(payload): + # I2V: the stored image file is the cross-model contract — + # every I2V pipeline reads ``params.image`` as a path. One + # write, straight to the final name; the id is unique per + # request, so no reader can observe it early. + if media_storage_path is None: raise ValueError( - "input_reference content is neither a decodable image nor a decodable video." + "media_storage_path is required when input_reference is an image" ) ref_path = os.path.join(media_storage_path, f"{id}_reference") - os.replace(tmp_path, ref_path) - except Exception: - # Cleanup-and-reraise, not handling: every failure path — - # validation errors (400) and I/O or dependency errors (500) - # alike — must not leak the temporary materialization. - if os.path.exists(tmp_path): - os.remove(tmp_path) - raise - if is_video: - # V2V reference: decode into ``VideoData`` under - # ``multi_modal_data["video"]`` — the one intake video-capable - # pipelines read. The decode is model-agnostic (all frames, no - # crop); each pipeline picks its own conditioning window. - params.multi_modal_data = { - "video": VideoData(frames=decode_video_frames(ref_path), metadata={}) - } - else: + with open(ref_path, "wb") as f: + f.write(payload) params.image = ref_path + else: + # V2V: decode in memory into ``VideoData`` under + # ``multi_modal_data["video"]`` — the one intake video-capable + # pipelines read; no file is materialized. The decode is + # model-agnostic (all frames, no crop); each pipeline picks + # its own conditioning window. + try: + frames = decode_video_frames_from_bytes(payload) + except ValueError as exc: + raise ValueError( + "input_reference content is neither a decodable image " + "nor a decodable video." + ) from exc + params.multi_modal_data = {"video": VideoData(frames=frames, metadata={})} _warn_if_set_with_no_semantic(request, getattr(generator, "model", None)) _merge_extra_params(params, request.extra_params, generator.extra_param_specs) diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index 79248a14ad6d..01a706cec2de 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -329,6 +329,19 @@ def test_multipart_video_reference_routes_to_multimodal(self, tmp_path): video_data = params.multi_modal_data["video"] assert isinstance(video_data, VideoData) assert len(video_data.frames) >= 1 + # Video references are decoded in memory — nothing lands in media storage. + assert list(tmp_path.iterdir()) == [] + + def test_video_reference_needs_no_media_storage(self): + # The decode is in-memory, so V2V works without a storage path at all + # (only image references persist a file for the worker to read). + from tensorrt_llm.inputs.multimodal_data import VideoData + + generator = _StubVisualGen() + b64 = base64.b64encode(self._mp4_bytes()).decode() + request = VideoGenerationRequest(prompt="x", input_reference=b64) + params = parse_visual_gen_params(request, "vid-9", generator, media_storage_path=None) + assert isinstance(params.multi_modal_data["video"], VideoData) def test_base64_video_reference_routes_to_multimodal(self, tmp_path): # Classification is content-based, so the JSON/base64 path can @@ -368,7 +381,7 @@ def test_undecodable_reference_raises_and_cleans_up(self, tmp_path): request = VideoGenerationRequest(prompt="x", input_reference=b64) with pytest.raises(ValueError, match="neither a decodable image"): parse_visual_gen_params(request, "vid-6", generator, media_storage_path=str(tmp_path)) - # The temporary materialization is removed on rejection. + # Classification runs on the bytes; rejected content never touches disk. assert list(tmp_path.iterdir()) == [] def test_malformed_base64_reference_raises_and_cleans_up(self, tmp_path): @@ -396,6 +409,56 @@ def read(self, *args, **kwargs): assert list(tmp_path.iterdir()) == [] +class TestMediaBytesProbes: + """The in-memory probe/decode primitives the serve boundary runs on.""" + + def test_is_image_bytes(self): + from tensorrt_llm.inputs.media_io import is_image_bytes + + buf = BytesIO() + Image.new("RGB", (4, 4), (1, 2, 3)).save(buf, format="PNG") + assert is_image_bytes(buf.getvalue()) + assert not is_image_bytes(b"definitely not an image") + # Video bytes are not an image (mp4 has no PIL-openable header). + assert not is_image_bytes(TestInputReferenceMaterialization._mp4_bytes()) + + def test_decode_video_frames_from_bytes(self): + pytest.importorskip("cv2") + from tensorrt_llm.inputs.media_io import decode_video_frames_from_bytes + + frames = decode_video_frames_from_bytes(TestInputReferenceMaterialization._mp4_bytes()) + assert len(frames) == 2 + assert all(isinstance(f, Image.Image) for f in frames) + + def test_decode_video_frames_from_bytes_max_frames(self): + pytest.importorskip("cv2") + from tensorrt_llm.inputs.media_io import decode_video_frames_from_bytes + + frames = decode_video_frames_from_bytes( + TestInputReferenceMaterialization._mp4_bytes(), max_frames=1 + ) + assert len(frames) == 1 + + def test_decode_video_frames_from_bytes_rejects_garbage(self): + pytest.importorskip("cv2") + from tensorrt_llm.inputs.media_io import decode_video_frames_from_bytes + + with pytest.raises(ValueError): + decode_video_frames_from_bytes(b"not a video at all") + + def test_tempfile_fallback_without_stream_backend(self, monkeypatch): + # Old OpenCV builds have no stream-buffered backend; the bytes spill + # to an auto-deleted tempfile and decode through the path route. + pytest.importorskip("cv2") + from tensorrt_llm.inputs import media_io + + monkeypatch.setattr(media_io, "_select_cv2_stream_buffered_backend", lambda: None) + frames = media_io.decode_video_frames_from_bytes( + TestInputReferenceMaterialization._mp4_bytes() + ) + assert len(frames) == 2 + + # ============================================================================= # _merge_extra_params — the merge truth table # ============================================================================= From cf054b90c59264a87e670638f4c5aed9698bb303 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 20 Jul 2026 13:35:52 -0700 Subject: [PATCH 31/64] Rename media probe functions to reflect decodability guarantee Signed-off-by: Igor Shovkun --- .../_torch/visual_gen/models/cosmos3/utils.py | 6 ++--- tensorrt_llm/inputs/media_io.py | 26 +++++++++++++++---- tensorrt_llm/serve/visual_gen_utils.py | 4 +-- .../visual_gen/test_trtllm_serve_endpoints.py | 5 ++-- .../visual_gen/test_visual_gen_utils.py | 10 +++---- 5 files changed, 34 insertions(+), 17 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py index 40c7a4e7171a..5656ae493ea7 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py @@ -43,11 +43,11 @@ def normalize_video_input_path(path: Path, max_frames: Optional[int] = None) -> # one conditioning frame; a decodable video is expanded to its frames. # Extensions are unreliable — the serve path stores references with no # type-suffix at all — so the container decides, not the name. - from tensorrt_llm.inputs.media_io import is_image_file, is_video_file + from tensorrt_llm.inputs.media_io import is_decodable_image_file, is_decodable_video_file - if is_image_file(path): + if is_decodable_image_file(path): return [str(path)] - if is_video_file(path): + if is_decodable_video_file(path): return decode_video_file(path, max_frames=max_frames) raise ValueError( f"Cosmos3 reference must be a frame directory, a decodable image, " diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index 319457ec2144..0bf68c3d72c2 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -346,7 +346,23 @@ def _get_cv2(): return cv2 -def is_image_file(path) -> bool: +# --- Content-classification probes ------------------------------------------ +# These classify media by asking the decoder itself, deliberately: +# * File suffixes are unreliable — client-controlled, often absent (the serve +# handles raw uploaded bytes with no filename at all), and never proof that +# the content matches the name. +# * Sniffing signatures ("magic numbers") would need an extra dependency +# (libmagic, as behind Unix `file`) or a hand-rolled signature table — and +# would still only name the container, not prove that this build's decoder +# can actually open it (codec support varies per PIL/OpenCV build). +# * A classifier that can disagree with the decoder is a deferred failure: +# content that passes the check but fails to decode later, deeper in the +# pipeline. Probing PIL/OpenCV directly makes "classified as X" and +# "decodes as X" the same statement by construction. +# Hence the names: these predicates promise decodability, not format identity. + + +def is_decodable_image_file(path) -> bool: """True when ``path`` holds still-image content (anything PIL opens). Header-only probe: identifies the container without decoding pixels. @@ -359,7 +375,7 @@ def is_image_file(path) -> bool: return False -def is_video_file(path) -> bool: +def is_decodable_video_file(path) -> bool: """True when ``path`` holds a decodable video stream (OpenCV-openable). Total predicate over content: False for images, audio, and undecodable @@ -367,7 +383,7 @@ def is_video_file(path) -> bool: image as a one-frame video stream, so a bare video probe would accept every PNG/JPEG. ``_get_cv2`` raises a clear install hint if cv2 is absent. """ - if is_image_file(path): + if is_decodable_image_file(path): return False cv2 = _get_cv2() capture = cv2.VideoCapture(str(path)) @@ -377,10 +393,10 @@ def is_video_file(path) -> bool: capture.release() -def is_image_bytes(data) -> bool: +def is_decodable_image_bytes(data) -> bool: """True when ``data`` holds still-image content (anything PIL opens). - In-memory counterpart of :func:`is_image_file` — header-only probe, no + In-memory counterpart of :func:`is_decodable_image_file` — header-only probe, no pixel decode, no filesystem. """ try: diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index 029f9052ec63..5857106b50f6 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -3,7 +3,7 @@ import os from typing import Any, Dict, List, Optional -from tensorrt_llm.inputs.media_io import decode_video_frames_from_bytes, is_image_bytes +from tensorrt_llm.inputs.media_io import decode_video_frames_from_bytes, is_decodable_image_bytes from tensorrt_llm.inputs.multimodal_data import VideoData from tensorrt_llm.logger import logger from tensorrt_llm.serve.openai_protocol import ImageGenerationRequest, VideoGenerationRequest @@ -165,7 +165,7 @@ def parse_visual_gen_params( # Classify by decoding the bytes in memory — nothing touches disk # until the modality is known and validated. - if is_image_bytes(payload): + if is_decodable_image_bytes(payload): # I2V: the stored image file is the cross-model contract — # every I2V pipeline reads ``params.image`` as a path. One # write, straight to the final name; the id is unique per diff --git a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py index 136e4f8b27dd..43097fcd9150 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py @@ -64,8 +64,9 @@ def _assert_llm_envelope( def _require_opencv(): """Skip unless OpenCV is installed (the shared optional video decoder). - These tests drive ``is_video_file``, which decodes references via OpenCV - (the same optional dep the multimodal video path uses). CI installs + These tests drive the serve's in-memory reference classification/decode + (``is_decodable_image_bytes`` / ``decode_video_frames_from_bytes``), which decodes + via OpenCV (the same optional dep the multimodal video path uses). CI installs ``opencv-python-headless`` in every test stage (``jenkins/L0_Test.groovy``), so these tests always run there; the skip only spares bare local environments, where cv2 stays optional (kept out of requirements by the diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index 01a706cec2de..df097ca85c35 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -412,15 +412,15 @@ def read(self, *args, **kwargs): class TestMediaBytesProbes: """The in-memory probe/decode primitives the serve boundary runs on.""" - def test_is_image_bytes(self): - from tensorrt_llm.inputs.media_io import is_image_bytes + def test_is_decodable_image_bytes(self): + from tensorrt_llm.inputs.media_io import is_decodable_image_bytes buf = BytesIO() Image.new("RGB", (4, 4), (1, 2, 3)).save(buf, format="PNG") - assert is_image_bytes(buf.getvalue()) - assert not is_image_bytes(b"definitely not an image") + assert is_decodable_image_bytes(buf.getvalue()) + assert not is_decodable_image_bytes(b"definitely not an image") # Video bytes are not an image (mp4 has no PIL-openable header). - assert not is_image_bytes(TestInputReferenceMaterialization._mp4_bytes()) + assert not is_decodable_image_bytes(TestInputReferenceMaterialization._mp4_bytes()) def test_decode_video_frames_from_bytes(self): pytest.importorskip("cv2") From a699afe3b587f33fb8a28ceeecb745f446f430c2 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 20 Jul 2026 14:39:14 -0700 Subject: [PATCH 32/64] Migrate Cosmos3 V2V reference from VideoData to uint8 tensor Replace the `multi_modal_data["video"]` / `VideoData` pathway with a plain `extra_params["video"]` uint8 `[T, H, W, C]` tensor contract. - Add `load_video_frames_tensor` and `frames_to_tensor` to `tensorrt_llm/inputs/media_io.py` as the canonical client-side helper - Register `"video"` as a typed `ExtraParamSchema` in Cosmos3 defaults - Update the serve layer (`visual_gen_utils.py`) to decode video references into the tensor form instead of `VideoData` - Remove `normalize_video_input`, `load_reference_video`, and related path helpers from `cosmos3/utils.py`; remove `normalize_video_input` import from the pipeline - Simplify `_normalize_condition_video_latent_indexes` to accept only iterables (drop `int` / `str` overloads) - Drop `multi_modal_data` field from `VisualGenParams` - Update offline example and all affected tests Signed-off-by: Igor Shovkun --- examples/visual_gen/models/cosmos3/cosmos3.py | 6 +- examples/visual_gen/serve/README.md | 2 +- .../visual_gen/models/cosmos3/defaults.py | 12 ++ .../models/cosmos3/pipeline_cosmos3.py | 48 ++---- .../_torch/visual_gen/models/cosmos3/utils.py | 74 +------- tensorrt_llm/inputs/media_io.py | 51 ++++++ tensorrt_llm/serve/visual_gen_utils.py | 23 ++- tensorrt_llm/visual_gen/params.py | 11 +- .../visual_gen/test_cosmos3_pipeline.py | 158 ++++++++---------- .../visual_gen/test_trtllm_serve_endpoints.py | 19 +-- .../visual_gen/test_visual_gen_utils.py | 27 +-- 11 files changed, 190 insertions(+), 241 deletions(-) diff --git a/examples/visual_gen/models/cosmos3/cosmos3.py b/examples/visual_gen/models/cosmos3/cosmos3.py index 844d17d17741..083a83b307cc 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3.py +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -107,7 +107,7 @@ from typing import Any, Dict, Optional from tensorrt_llm import VisualGen, VisualGenArgs -from tensorrt_llm._torch.visual_gen.models.cosmos3.utils import load_reference_video +from tensorrt_llm.inputs.media_io import load_video_frames_tensor _SCRIPT_DIR = Path(__file__).resolve().parent @@ -280,7 +280,9 @@ def main(): params.extra_params["output_type"] = output_type if args.video_path is not None: - params.multi_modal_data = {"video": load_reference_video(args.video_path)} + # Decode client-side into the uint8 [T, H, W, C] tensor contract; the + # worker keeps the conditioning window and VAE-encodes. + params.extra_params["video"] = load_video_frames_tensor(args.video_path) if negative_prompt is None: params.negative_prompt = None diff --git a/examples/visual_gen/serve/README.md b/examples/visual_gen/serve/README.md index 4cf5310ef6c6..4bc9840b7f10 100644 --- a/examples/visual_gen/serve/README.md +++ b/examples/visual_gen/serve/README.md @@ -286,7 +286,7 @@ You can customize these by: - `frame_rate` (canonical) or `fps` (alias): frames per second - `num_frames`: when set, wins over the `seconds * frame_rate` derivation - `seed`, `num_inference_steps`, `guidance_scale`, `max_sequence_length`, `negative_prompt`: per-request denoise controls -- `input_reference`: Reference image (I2V/TI2V) or video (V2V), classified by decoding the content — filename and content type are ignored; accepted as base64-encoded string in JSON or as a file in multipart form-data. Undecodable content returns HTTP 400. Video decode requires the `av` (PyAV) package on the server (not bundled — `pip install av`). +- `input_reference`: Reference image (I2V/TI2V) or video (V2V), classified by decoding the content — filename and content type are ignored; accepted as base64-encoded string in JSON or as a file in multipart form-data. Undecodable content returns HTTP 400. Video decode requires OpenCV on the server (not bundled — `pip install opencv-python-headless`). - `extra_params`: model-specific overflow (see below) - `response_format`: `"b64_json"` or `"url"` - `format`: Generation content encoding. Video encoders: `"mp4"`, `"avi"`, `"auto"`. Tensor formats: `"safetensors"`, `"pt"` (carries video + audio + scalar metadata in one payload for LTX-2). diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index 311844e0f3c1..e8544730b94e 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -115,4 +115,16 @@ default=None, description="Optional scheduler flow shift override. Uses the Cosmos3 mode default when omitted.", ), + "video": ExtraParamSchema( + type="tensor", + default=None, + description=( + "V2V reference: decoded video frames as a uint8 [T, H, W, C] RGB " + "torch.Tensor (build one from a file with " + "tensorrt_llm.inputs.media_io.load_video_frames_tensor). The worker " + "keeps the first/last conditioning window per " + "condition_video_latent_indexes / condition_video_keep and " + "VAE-encodes it; media is always decoded by the producer." + ), + ), } diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index f05f88f86fd4..1204e497c55a 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -45,7 +45,7 @@ from .guardrails import check_video_safety, download_guardrail_checkpoint from .sound_tokenizer import LatentAutoEncoderV2 from .transformer_cosmos3 import Cosmos3VFMTransformer -from .utils import normalize_video_input, pil_to_rgb +from .utils import pil_to_rgb COSMOS3_DEFAULT_NEGATIVE_PROMPT = "" # NOTE: Intentional typo in "give" instead of "given" to match training setup. @@ -63,17 +63,11 @@ def _normalize_condition_video_latent_indexes( - indexes: Iterable[int] | int | str | None, + indexes: Iterable[int] | None, ) -> tuple[int, ...]: if indexes is None: return COSMOS3_DEFAULT_CONDITION_VIDEO_LATENT_INDEXES - if isinstance(indexes, int): - normalized = (indexes,) - elif isinstance(indexes, str): - parts = [part.strip() for part in indexes.split(",") if part.strip()] - normalized = tuple(int(part) for part in parts) - else: - normalized = tuple(int(index) for index in indexes) + normalized = tuple(int(index) for index in indexes) if not normalized: raise ValueError("Cosmos3 condition_video_latent_indexes must not be empty.") @@ -308,14 +302,7 @@ def _run_warmup(self, height: int, width: int, num_frames: int, steps: int) -> N def infer(self, req): extra_params = req.params.extra_params or {} output_type = extra_params.get("output_type", "video") - - # The V2V reference rides in ``multi_modal_data["video"]`` as a - # ``VideoData`` (framework convention); the worker crops + VAE-encodes its - # frames. Both producers (offline and serve) build it, so there is no - # legacy-path fallback. - mm_data = req.params.multi_modal_data or {} - video_data = mm_data.get("video") - video = video_data.frames if video_data is not None else None + video = extra_params.get("video") # Tensor[T, H, W, C, dtype=uint8] return self.forward( prompt=req.prompt, @@ -687,7 +674,7 @@ def _prepare_latents_v2v( video_tensor: torch.Tensor, num_frames: int, generator: torch.Generator, - condition_video_latent_indexes: Iterable[int] | int | str | None = None, + condition_video_latent_indexes: Iterable[int] | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Prepare V2V latents with explicit clean conditioned latent frames.""" if video_tensor.ndim == 4: @@ -776,9 +763,9 @@ def forward( use_guardrails: bool = COSMOS3_EXTRA_SPECS["use_guardrails"].default, enable_audio: bool = COSMOS3_EXTRA_SPECS["enable_audio"].default, output_type: str = COSMOS3_EXTRA_SPECS["output_type"].default, - video: Any = None, - condition_video_latent_indexes: Any = None, - condition_video_keep: Any = None, + video: torch.Tensor | None = None, # uint8 [T, H, W, C, dtype=uint8] + condition_video_latent_indexes: Iterable[int] | None = None, + condition_video_keep: str | None = None, flow_shift: Optional[float] = None, ): pipeline_start = time.time() @@ -980,19 +967,20 @@ def forward( ), num_frames, ) - # ``video`` is the reference frames (a list, from the producer's - # ``VideoData``). The worker crops the first/last conditioning window - # uniformly — producers never crop, so behavior never depends on them. - video = normalize_video_input( - video, - max_frames=None if condition_video_keep == "last" else condition_pixel_frames, - ) - video = ( + if not isinstance(video, torch.Tensor) or video.ndim != 4: + raise ValueError( + "Cosmos3 V2V reference must be a uint8 [T, H, W, C] tensor " + f"(the 'video' extra-param contract), got {type(video).__name__}" + f"{' of shape ' + str(tuple(video.shape)) if isinstance(video, torch.Tensor) else ''}." + ) + # Crop the first/last conditioning window uniformly + window = ( video[-condition_pixel_frames:] if condition_video_keep == "last" else video[:condition_pixel_frames] ) - video = self._preprocess_condition_video(video, height, width) + frames = [PIL.Image.fromarray(frame.cpu().numpy()) for frame in window] + video = self._preprocess_condition_video(frames, height, width) if self.rank == 0: logger.info( diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py index 5656ae493ea7..a45e5e03f73f 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py @@ -5,13 +5,10 @@ from __future__ import annotations -from pathlib import Path -from typing import Any, List, Optional +from typing import Any import PIL.Image -IMAGE_EXTENSIONS = frozenset({".png", ".jpg", ".jpeg", ".webp", ".bmp"}) - def pil_to_rgb(value: Any) -> PIL.Image.Image: if isinstance(value, str): @@ -19,72 +16,3 @@ def pil_to_rgb(value: Any) -> PIL.Image.Image: if isinstance(value, PIL.Image.Image): return value.convert("RGB") raise TypeError(f"Cosmos3 preprocessing expected PIL image or image path, got {type(value)!r}.") - - -def decode_video_file(path: Path, max_frames: Optional[int] = None) -> List[PIL.Image.Image]: - from tensorrt_llm.inputs.media_io import decode_video_frames - - return decode_video_frames(path, max_frames=max_frames) - - -def normalize_video_input_path(path: Path, max_frames: Optional[int] = None) -> List[Any]: - if not path.exists(): - raise ValueError(f"Cosmos3 video path does not exist: {path}") - if path.is_dir(): - frames = sorted(p for p in path.iterdir() if p.suffix.lower() in IMAGE_EXTENSIONS) - if not frames: - raise ValueError(f"No image frames found in Cosmos3 video directory: {path}") - frame_paths = [str(p) for p in frames] - if max_frames is not None: - frame_paths = frame_paths[:max_frames] - return frame_paths - - # Classify a single file by content, not by suffix: a decodable still is - # one conditioning frame; a decodable video is expanded to its frames. - # Extensions are unreliable — the serve path stores references with no - # type-suffix at all — so the container decides, not the name. - from tensorrt_llm.inputs.media_io import is_decodable_image_file, is_decodable_video_file - - if is_decodable_image_file(path): - return [str(path)] - if is_decodable_video_file(path): - return decode_video_file(path, max_frames=max_frames) - raise ValueError( - f"Cosmos3 reference must be a frame directory, a decodable image, " - f"or a decodable video; got {path}" - ) - - -def normalize_video_input(video: Any, max_frames: Optional[int] = None) -> List[Any]: - """Normalize video input to a frame list. - - Accepts a list of PIL images / paths, a single image or video file path, - or a directory of frame images (sorted lexicographically). - """ - if video is None: - return [] - if isinstance(video, list): - if not video: - raise ValueError("Cosmos3 video input must contain at least one frame.") - if max_frames is not None: - return video[:max_frames] - return video - if isinstance(video, (str, Path)): - return normalize_video_input_path(Path(video), max_frames=max_frames) - return [video] - - -def load_reference_video(src: Any): - """Decode a reference into the framework's ``VideoData`` (all frames). - - Offline entry point: decodes any supported reference into ``VideoData`` for - ``multi_modal_data["video"]``. It does **not** crop and carries no temporal - metadata — the Cosmos3 worker crops the first/last conditioning window, and - the temporal placement is set by ``condition_video_latent_indexes`` + - ``frame_rate`` (request params), not the reference. The worker reads - ``.frames`` and VAE-encodes them; it never decodes media itself. - """ - from tensorrt_llm.inputs.multimodal_data import VideoData - - frames = normalize_video_input(src, max_frames=None) - return VideoData(frames=[pil_to_rgb(frame) for frame in frames], metadata={}) diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index 0bf68c3d72c2..1a160b93ce7f 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -469,6 +469,57 @@ def decode_video_frames_from_bytes(data, max_frames: Optional[int] = None) -> Li return frames +# Frame-image suffixes recognized when expanding a frame directory (see the +# selection rationale in ``load_video_frames_tensor``). +_IMAGE_SUFFIXES = frozenset({".png", ".jpg", ".jpeg", ".webp", ".bmp"}) + + +def frames_to_tensor(frames: List["Image.Image"]) -> torch.Tensor: + """Stack PIL frames into a uint8 ``[T, H, W, C]`` RGB CPU tensor. + + The plain-tensor form video references travel in (e.g. + ``extra_params["video"]`` for video-to-video pipelines). + """ + if not frames: + raise ValueError("Cannot build a video tensor from an empty frame list.") + return torch.from_numpy(np.stack([np.asarray(f.convert("RGB")) for f in frames])) + + +def load_video_frames_tensor(source, max_frames: Optional[int] = None) -> torch.Tensor: + """Load a video reference from disk as a uint8 ``[T, H, W, C]`` RGB tensor. + + Public helper for building video-reference tensors client-side. Accepts a + video file, a single still image (one frame), or a directory of frame + images (sorted lexicographically). Dispatch is by content, not suffix + (see the content-classification probes above). + """ + path = Path(source) + if not path.exists(): + raise ValueError(f"Video reference path does not exist: {path}") + if path.is_dir(): + # Directories are selected by suffix, deliberately unlike single files: + # a frame directory is user-curated (names are the interface), suffix + # selection costs no file opens (a decodability probe is an open per + # entry — painful for thousands of frames on network filesystems), and + # it fails the right way — a selected frame that doesn't decode raises + # below, whereas a probe would silently drop corrupt frames and produce + # a video with holes. + frame_paths = sorted(p for p in path.iterdir() if p.suffix.lower() in _IMAGE_SUFFIXES) + if not frame_paths: + raise ValueError(f"No image frames found in directory: {path}") + if max_frames is not None: + frame_paths = frame_paths[:max_frames] + return frames_to_tensor([Image.open(p) for p in frame_paths]) + if is_decodable_image_file(path): + return frames_to_tensor([Image.open(path)]) + if is_decodable_video_file(path): + return frames_to_tensor(decode_video_frames(path, max_frames=max_frames)) + raise ValueError( + f"Video reference must be a decodable video, a decodable image, or a " + f"directory of frame images; got undecodable {path}" + ) + + def _select_cv2_stream_buffered_backend() -> Optional[int]: """Return a VideoCapture backend that can read from a Python `BytesIO`. diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index 5857106b50f6..1837fd69813b 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -3,8 +3,11 @@ import os from typing import Any, Dict, List, Optional -from tensorrt_llm.inputs.media_io import decode_video_frames_from_bytes, is_decodable_image_bytes -from tensorrt_llm.inputs.multimodal_data import VideoData +from tensorrt_llm.inputs.media_io import ( + decode_video_frames_from_bytes, + frames_to_tensor, + is_decodable_image_bytes, +) from tensorrt_llm.logger import logger from tensorrt_llm.serve.openai_protocol import ImageGenerationRequest, VideoGenerationRequest from tensorrt_llm.visual_gen import VisualGen, VisualGenParams @@ -166,10 +169,8 @@ def parse_visual_gen_params( # Classify by decoding the bytes in memory — nothing touches disk # until the modality is known and validated. if is_decodable_image_bytes(payload): - # I2V: the stored image file is the cross-model contract — - # every I2V pipeline reads ``params.image`` as a path. One - # write, straight to the final name; the id is unique per - # request, so no reader can observe it early. + # I2V: the stored image file is the cross-model contract. + # every I2V pipeline reads ``params.image`` as a path. if media_storage_path is None: raise ValueError( "media_storage_path is required when input_reference is an image" @@ -179,11 +180,7 @@ def parse_visual_gen_params( f.write(payload) params.image = ref_path else: - # V2V: decode in memory into ``VideoData`` under - # ``multi_modal_data["video"]`` — the one intake video-capable - # pipelines read; no file is materialized. The decode is - # model-agnostic (all frames, no crop); each pipeline picks - # its own conditioning window. + # V2V: decode in memory into a uint8 [T, H, W, C] tensor try: frames = decode_video_frames_from_bytes(payload) except ValueError as exc: @@ -191,7 +188,9 @@ def parse_visual_gen_params( "input_reference content is neither a decodable image " "nor a decodable video." ) from exc - params.multi_modal_data = {"video": VideoData(frames=frames, metadata={})} + if params.extra_params is None: + params.extra_params = {} + params.extra_params["video"] = frames_to_tensor(frames) _warn_if_set_with_no_semantic(request, getattr(generator, "model", None)) _merge_extra_params(params, request.extra_params, generator.extra_param_specs) diff --git a/tensorrt_llm/visual_gen/params.py b/tensorrt_llm/visual_gen/params.py index 88251fd2056d..2756ec79158d 100644 --- a/tensorrt_llm/visual_gen/params.py +++ b/tensorrt_llm/visual_gen/params.py @@ -14,6 +14,7 @@ # limitations under the License. from typing import Any, Dict, List, Optional, Union +import torch from pydantic import Field from tensorrt_llm.llmapi.utils import StrictBaseModel, set_api_status @@ -72,15 +73,6 @@ class VisualGenParams(StrictBaseModel): image: Optional[Union[str, bytes, List[Union[str, bytes]]]] = Field( default=None, description="Reference image(s) for I2V/I2I." ) - # Framework multimodal convention: modality -> data. Cosmos3 V2V carries its - # reference here as ``{"video": VideoData(frames, metadata)}``; the worker - # reads the frames and VAE-encodes them (see ``inputs.multimodal_data``). - multi_modal_data: Optional[Dict[str, Any]] = Field( - default=None, - description="Multimodal conditioning inputs keyed by modality " - "(e.g. {'video': VideoData}; see inputs.multimodal_data).", - ) - # Per-prompt multiplier num_images_per_prompt: int = Field(default=1, description="Number of images per prompt.") @@ -101,6 +93,7 @@ class VisualGenParams(StrictBaseModel): "bool": (bool,), "str": (str,), "list": (list,), + "tensor": (torch.Tensor,), # Decoded media payloads (e.g. a V2V reference as a uint8 [T, H, W, C] } # Generation config fields that pipelines declare defaults for. If a user diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index ef0eed637979..3b155a23b105 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -55,6 +55,7 @@ _normalize_condition_video_latent_indexes, ) from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader +from tensorrt_llm.inputs.media_io import frames_to_tensor from tensorrt_llm.visual_gen.args import TorchCompileConfig, VisualGenArgs pytestmark = pytest.mark.cosmos3 @@ -537,16 +538,14 @@ class TestCosmos3V2VConditioningParams: "value,expected", [ (None, (0, 1)), - (0, (0,)), ([0, 2], (0, 2)), ((1, 3), (1, 3)), - ("0, 2", (0, 2)), ], ) def test_normalize_condition_video_latent_indexes(self, value, expected): assert _normalize_condition_video_latent_indexes(value) == expected - @pytest.mark.parametrize("value", [[], "", [-1], "0, -1", [0, -2]]) + @pytest.mark.parametrize("value", [[], [-1], [0, -2]]) def test_invalid_condition_video_latent_indexes_raise(self, value): with pytest.raises(ValueError): _normalize_condition_video_latent_indexes(value) @@ -579,96 +578,70 @@ def test_invalid_condition_video_keep_raises(self): _normalize_condition_video_keep("middle") -class TestNormalizeVideoInputContentDispatch: - """``normalize_video_input_path`` classifies files by content, not suffix. +def _write_mp4(path, num_frames: int = 3) -> None: + """Synthesize a tiny mp4v clip (no video asset ships with the repo).""" + cv2 = pytest.importorskip("cv2") + writer = cv2.VideoWriter(str(path), cv2.VideoWriter_fourcc(*"mp4v"), 4.0, (16, 16)) + try: + for _ in range(num_frames): + writer.write(np.zeros((16, 16, 3), dtype=np.uint8)) + finally: + writer.release() + assert path.exists() and path.stat().st_size > 0 - The serve path stores references with no type-suffix, so decode dispatch - must key on the container, not the filename. CPU-only; the clip is - synthesized with OpenCV (installed by CI test stages via - ``jenkins/L0_Test.groovy``; the importorskip only spares bare local envs). - """ - @staticmethod - def _write_mp4(path, num_frames: int = 3) -> None: - cv2 = pytest.importorskip("cv2") - writer = cv2.VideoWriter(str(path), cv2.VideoWriter_fourcc(*"mp4v"), 4.0, (16, 16)) - try: - for _ in range(num_frames): - writer.write(np.zeros((16, 16, 3), dtype=np.uint8)) - finally: - writer.release() - assert path.exists() and path.stat().st_size > 0 +class TestLoadVideoFramesTensor: + """``media_io.load_video_frames_tensor`` builds the uint8 [T, H, W, C] + tensor the ``video`` extra param carries (all frames, no crop — the worker + keeps the conditioning window). Public helper, used by the example and + available to API clients. CPU-only.""" - def test_video_without_extension_decodes_to_frames(self, tmp_path): + def test_from_video_file_all_frames(self, tmp_path): pytest.importorskip("cv2") - from tensorrt_llm._torch.visual_gen.models.cosmos3.utils import normalize_video_input_path - - # OpenCV's *writer* selects the muxer by extension, so encode to a - # ``.mp4`` path, then rename to a suffix-less name — exactly what the - # serve path produces (raw bytes written to ``{id}_reference``). - encoded = tmp_path / "clip.mp4" - self._write_mp4(encoded) - ref = encoded.rename(tmp_path / "reference") - frames = normalize_video_input_path(ref) - # Took the video-decode path (returns PIL frames), not the single-still - # path (which would return the bare ``[str(path)]``). - assert frames and all(isinstance(f, PIL.Image.Image) for f in frames) - - def test_image_without_extension_is_single_frame(self, tmp_path): - from tensorrt_llm._torch.visual_gen.models.cosmos3.utils import normalize_video_input_path - - ref = tmp_path / "reference" # no ``.png`` suffix - PIL.Image.new("RGB", (8, 8), (1, 2, 3)).save(ref, format="PNG") - assert normalize_video_input_path(ref) == [str(ref)] - - def test_undecodable_file_raises(self, tmp_path): - pytest.importorskip("cv2") - from tensorrt_llm._torch.visual_gen.models.cosmos3.utils import normalize_video_input_path - - ref = tmp_path / "reference" - ref.write_bytes(b"not media") - with pytest.raises(ValueError, match="decodable"): - normalize_video_input_path(ref) + from tensorrt_llm.inputs.media_io import load_video_frames_tensor + enc = tmp_path / "clip.mp4" + _write_mp4(enc, num_frames=8) + video = load_video_frames_tensor(enc) + assert video.dtype == torch.uint8 + assert video.ndim == 4 and video.shape[0] == 8 and video.shape[-1] == 3 -class TestLoadReferenceVideo: - """``load_reference_video`` decodes any reference into the framework's - ``VideoData`` (all frames) for ``multi_modal_data["video"]``. It does not - crop — the worker does — so the producer stays model-agnostic. CPU-only.""" - - def test_from_pil_list_all_frames(self): - from tensorrt_llm._torch.visual_gen.models.cosmos3.utils import load_reference_video - from tensorrt_llm.inputs.multimodal_data import VideoData - - frames = [PIL.Image.new("RGB", (8, 8), (i, i, i)) for i in range(6)] - vd = load_reference_video(frames) - assert isinstance(vd, VideoData) - assert len(vd.frames) == 6 # all frames, no crop - assert all(isinstance(f, PIL.Image.Image) for f in vd.frames) - # No temporal metadata: placement comes from request params - # (condition_video_latent_indexes, frame_rate), not the reference. - assert vd.metadata == {} + def test_from_image_file_is_single_frame(self, tmp_path): + from tensorrt_llm.inputs.media_io import load_video_frames_tensor - def test_from_video_file(self, tmp_path): - pytest.importorskip("cv2") - from tensorrt_llm._torch.visual_gen.models.cosmos3.utils import load_reference_video - - enc = tmp_path / "clip.mp4" - TestNormalizeVideoInputContentDispatch._write_mp4(enc, num_frames=8) - assert len(load_reference_video(enc).frames) == 8 # all frames + p = tmp_path / "img.png" + PIL.Image.new("RGB", (8, 8), (7, 7, 7)).save(p) + assert load_video_frames_tensor(p).shape == (1, 8, 8, 3) def test_from_directory(self, tmp_path): - from tensorrt_llm._torch.visual_gen.models.cosmos3.utils import load_reference_video + from tensorrt_llm.inputs.media_io import load_video_frames_tensor for i in range(5): PIL.Image.new("RGB", (8, 8), (i, i, i)).save(tmp_path / f"{i:03d}.png") - assert len(load_reference_video(tmp_path).frames) == 5 + video = load_video_frames_tensor(tmp_path) + assert video.shape == (5, 8, 8, 3) + # Sorted lexicographically: frame k is the solid (k, k, k) image. + assert int(video[0].float().mean()) == 0 and int(video[4].float().mean()) == 4 + + def test_directory_selects_by_suffix_and_fails_loud_on_corrupt(self, tmp_path): + # Directories are user-curated: non-frame entries are ignored by name, + # but a selected frame that doesn't decode raises — corrupt frames are + # never silently dropped into a video with holes. + from tensorrt_llm.inputs.media_io import load_video_frames_tensor + + PIL.Image.new("RGB", (8, 8), (1, 1, 1)).save(tmp_path / "000.png") + (tmp_path / "notes.txt").write_text("not a frame") # ignored by suffix + assert load_video_frames_tensor(tmp_path).shape[0] == 1 + + (tmp_path / "001.png").write_bytes(b"corrupt") + with pytest.raises(PIL.UnidentifiedImageError): + load_video_frames_tensor(tmp_path) def test_missing_path_raises(self, tmp_path): - from tensorrt_llm._torch.visual_gen.models.cosmos3.utils import load_reference_video + from tensorrt_llm.inputs.media_io import load_video_frames_tensor with pytest.raises(ValueError, match="does not exist"): - load_reference_video(tmp_path / "nope") + load_video_frames_tensor(tmp_path / "nope") @pytest.mark.integration @@ -676,7 +649,7 @@ def test_missing_path_raises(self, tmp_path): @pytest.mark.high_cuda_memory class TestCosmos3V2V: def test_v2v_smoke(self, cosmos3_pipeline): - video = _make_test_video(NUM_FRAMES) + video = frames_to_tensor(_make_test_video(NUM_FRAMES)) result = _run_forward( cosmos3_pipeline, image=None, @@ -693,18 +666,19 @@ def test_v2v_smoke(self, cosmos3_pipeline): use_karras_sigmas=False, ) - def test_v2v_multimodal_reference_smoke(self, cosmos3_pipeline): - """The V2V reference arrives as ``VideoData`` (built by - ``load_reference_video``, all frames) under ``multi_modal_data["video"]``; - the worker crops the conditioning window and VAE-encodes. Mirrors what the - offline example feeds the pipeline.""" - from tensorrt_llm._torch.visual_gen.models.cosmos3.utils import load_reference_video + def test_v2v_tensor_reference_smoke(self, cosmos3_pipeline): + """The V2V reference arrives as a decoded uint8 [T, H, W, C] tensor + (the ``video`` extra-param contract); the worker crops the conditioning + window and VAE-encodes. Mirrors what the offline example and serve feed + the pipeline.""" + from tensorrt_llm.inputs.media_io import frames_to_tensor - video_data = load_reference_video(_make_test_video(NUM_FRAMES)) + video = frames_to_tensor(_make_test_video(NUM_FRAMES)) + assert video.dtype == torch.uint8 and video.ndim == 4 result = _run_forward( cosmos3_pipeline, image=None, - video=video_data.frames, + video=video, num_frames=NUM_FRAMES, condition_video_latent_indexes=[0, 1], condition_video_keep="first", @@ -722,7 +696,9 @@ def test_v2v_keep_last_smoke(self, cosmos3_pipeline): dark = PIL.Image.new("RGB", (WIDTH, HEIGHT), (40, 40, 40)) bright = PIL.Image.new("RGB", (WIDTH, HEIGHT), (230, 230, 230)) # 5 = max(condition_video_latent_indexes) * 4 + 1 conditioning frames. - video = [dark.copy() for _ in range(NUM_FRAMES)] + [bright.copy() for _ in range(5)] + video = frames_to_tensor( + [dark.copy() for _ in range(NUM_FRAMES)] + [bright.copy() for _ in range(5)] + ) result = _run_forward( cosmos3_pipeline, image=None, @@ -766,7 +742,7 @@ def fake_tokenize_prompt(text, max_sequence_length, use_system_prompt, system_pr with pytest.raises(StopAfterTokenize): pipeline.forward( prompt="continue", - video=_make_test_video(5, width=16, height=16), + video=frames_to_tensor(_make_test_video(5, width=16, height=16)), height=16, width=16, num_frames=5, @@ -791,7 +767,7 @@ def test_image_and_video_rejected(self, cosmos3_pipeline): _run_forward( cosmos3_pipeline, image=_make_test_image(), - video=_make_test_video(5), + video=frames_to_tensor(_make_test_video(5)), ) def test_t2i_and_video_rejected(self, cosmos3_pipeline): @@ -799,7 +775,7 @@ def test_t2i_and_video_rejected(self, cosmos3_pipeline): _run_forward( cosmos3_pipeline, image=None, - video=_make_test_video(5), + video=frames_to_tensor(_make_test_video(5)), output_type="image", height=T2I_HEIGHT, width=T2I_WIDTH, @@ -846,7 +822,7 @@ def test_v2v_audio_smoke(self, cosmos3_pipeline): result = _run_forward( cosmos3_pipeline, enable_audio=True, - video=_make_test_video(NUM_FRAMES), + video=frames_to_tensor(_make_test_video(NUM_FRAMES)), condition_video_latent_indexes=[0, 1], condition_video_keep="first", ) diff --git a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py index 43097fcd9150..1295d24c3a6c 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py @@ -859,8 +859,8 @@ def test_sync_video_generation_multipart_with_reference(self, video_client, tmp_ assert os.path.exists(params.image) def test_sync_video_generation_multipart_with_video_reference(self, video_client, tmp_path): - """A video ``input_reference`` is decoded into ``VideoData`` under - ``multi_modal_data["video"]`` (V2V). + """A video ``input_reference`` is decoded into a uint8 [T, H, W, C] + tensor on the model-specific ``video`` extra param (V2V). The reference is classified by decoding its content, so the clip is synthesized in-test with OpenCV — no video asset ships with the repo. @@ -890,16 +890,15 @@ def test_sync_video_generation_multipart_with_video_reference(self, video_client assert resp.status_code == 200 assert len(resp.content) > 0 - # Video content must NOT land on params.image; it's decoded into - # VideoData under multi_modal_data["video"] (framework convention; the - # same pipeline entry the offline example's --video_path uses). - from tensorrt_llm.inputs.multimodal_data import VideoData - + # Video content must NOT land on params.image; it's decoded into a + # uint8 [T, H, W, C] tensor on the model-specific ``video`` extra param + # (the same intake the offline example's --video_path uses). params = video_client.mock_gen.last_params assert params.image is None - video_data = params.multi_modal_data["video"] - assert isinstance(video_data, VideoData) - assert len(video_data.frames) >= 1 + video = params.extra_params["video"] + assert isinstance(video, torch.Tensor) + assert video.dtype == torch.uint8 + assert video.ndim == 4 and video.shape[-1] == 3 def test_sync_video_generation_undecodable_reference_400(self, video_client): """Content neither PIL nor OpenCV can decode is rejected at the boundary.""" diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index df097ca85c35..4ad2f99c3b2c 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -313,8 +313,8 @@ def _mp4_bytes() -> bytes: finally: os.remove(path) - def test_multipart_video_reference_routes_to_multimodal(self, tmp_path): - from tensorrt_llm.inputs.multimodal_data import VideoData + def test_multipart_video_reference_routes_to_extra_params_tensor(self, tmp_path): + import torch generator = _StubVisualGen() upload = UploadFile(file=BytesIO(self._mp4_bytes()), filename="clip.mp4") @@ -322,31 +322,32 @@ def test_multipart_video_reference_routes_to_multimodal(self, tmp_path): params = parse_visual_gen_params( request, "vid-3", generator, media_storage_path=str(tmp_path) ) - # Video content is decoded into VideoData under multi_modal_data["video"] - # (framework convention), not params.image. The worker crops + VAE-encodes. + # Video content is decoded into a uint8 [T, H, W, C] tensor on the + # model-specific ``video`` extra param, not params.image. The worker + # crops the conditioning window + VAE-encodes. assert params.image is None - assert params.multi_modal_data is not None - video_data = params.multi_modal_data["video"] - assert isinstance(video_data, VideoData) - assert len(video_data.frames) >= 1 + video = params.extra_params["video"] + assert isinstance(video, torch.Tensor) + assert video.dtype == torch.uint8 + assert video.ndim == 4 and video.shape[0] == 2 and video.shape[-1] == 3 # Video references are decoded in memory — nothing lands in media storage. assert list(tmp_path.iterdir()) == [] def test_video_reference_needs_no_media_storage(self): # The decode is in-memory, so V2V works without a storage path at all # (only image references persist a file for the worker to read). - from tensorrt_llm.inputs.multimodal_data import VideoData + import torch generator = _StubVisualGen() b64 = base64.b64encode(self._mp4_bytes()).decode() request = VideoGenerationRequest(prompt="x", input_reference=b64) params = parse_visual_gen_params(request, "vid-9", generator, media_storage_path=None) - assert isinstance(params.multi_modal_data["video"], VideoData) + assert isinstance(params.extra_params["video"], torch.Tensor) - def test_base64_video_reference_routes_to_multimodal(self, tmp_path): + def test_base64_video_reference_routes_to_extra_params_tensor(self, tmp_path): # Classification is content-based, so the JSON/base64 path can # carry video even though it has no content-type or filename. - from tensorrt_llm.inputs.multimodal_data import VideoData + import torch generator = _StubVisualGen() b64 = base64.b64encode(self._mp4_bytes()).decode() @@ -355,7 +356,7 @@ def test_base64_video_reference_routes_to_multimodal(self, tmp_path): request, "vid-4", generator, media_storage_path=str(tmp_path) ) assert params.image is None - assert isinstance(params.multi_modal_data["video"], VideoData) + assert isinstance(params.extra_params["video"], torch.Tensor) def test_multipart_image_reference_routes_to_image(self, tmp_path): # JPEG upload: content sniffing classifies it as an image and routes From 3f9fe5004dfb73ec9b2a1a6fd11510473e3bf99a Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 20 Jul 2026 14:41:00 -0700 Subject: [PATCH 33/64] Clean comment Signed-off-by: Igor Shovkun --- tensorrt_llm/serve/visual_gen_utils.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index 1837fd69813b..e3bbfefbca45 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -166,8 +166,6 @@ def parse_visual_gen_params( else: payload = request.input_reference.file.read() - # Classify by decoding the bytes in memory — nothing touches disk - # until the modality is known and validated. if is_decodable_image_bytes(payload): # I2V: the stored image file is the cross-model contract. # every I2V pipeline reads ``params.image`` as a path. From 01195643424e990c954e539f8ca50d7fba2bdfb9 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 20 Jul 2026 15:22:56 -0700 Subject: [PATCH 34/64] Add per-param validators to ExtraParamSchema for preflight 400s Wire optional `validator` callables into `ExtraParamSchema` and `validate_visual_gen_params` so invalid conditioning values (e.g. negative latent indexes, bad `condition_video_keep`, malformed video tensor shape/dtype) are rejected at the API boundary with a 400 error instead of failing deep inside the worker. Move the Cosmos3 normalizer/validator functions from `pipeline_cosmos3.py` into `defaults.py` so they are module-level and survive pickling when specs are exchanged over ZMQ during worker startup. Signed-off-by: Igor Shovkun --- .../visual_gen/models/cosmos3/defaults.py | 48 ++++++++++++++++++- .../models/cosmos3/pipeline_cosmos3.py | 29 ++--------- tensorrt_llm/_torch/visual_gen/pipeline.py | 4 ++ tensorrt_llm/visual_gen/params.py | 12 ++++- .../visual_gen/test_visual_gen_params.py | 42 ++++++++++++++++ 5 files changed, 107 insertions(+), 28 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index e8544730b94e..f790cd86be94 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -17,7 +17,9 @@ Shared by the Cosmos3 OmniMoT text-to-video and image-to-video generation paths. """ -from typing import Dict +from typing import Dict, Iterable + +import torch from tensorrt_llm._torch.visual_gen.pipeline import ExtraParamSchema @@ -38,6 +40,47 @@ COSMOS3_DEFAULT_CONDITION_VIDEO_LATENT_INDEXES = (0, 1) COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP = "first" + +# --------------------------------------------------------------------------- +# Conditioning-value normalizers / validators. Declared as the ``validator`` +# of the matching extra-param specs below, so invalid values 400 at preflight; +# the pipeline reuses them at run time to normalize the same inputs. +# --------------------------------------------------------------------------- + + +def _normalize_condition_video_latent_indexes( + indexes: Iterable[int] | None, +) -> tuple[int, ...]: + if indexes is None: + return COSMOS3_DEFAULT_CONDITION_VIDEO_LATENT_INDEXES + normalized = tuple(int(index) for index in indexes) + + if not normalized: + raise ValueError("Cosmos3 condition_video_latent_indexes must not be empty.") + if any(index < 0 for index in normalized): + raise ValueError( + f"Cosmos3 condition_video_latent_indexes must be non-negative, got {normalized}." + ) + return normalized + + +def _normalize_condition_video_keep(keep: str | None) -> str: + normalized = str(keep or COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP).strip().lower() + if normalized not in {"first", "last"}: + raise ValueError("Cosmos3 condition_video_keep must be either first or last.") + return normalized + + +def _validate_video_reference_tensor(video: torch.Tensor) -> None: + if video.ndim != 4 or video.shape[-1] != 3: + raise ValueError( + f"Cosmos3 video reference must be a uint8 [T, H, W, C] RGB tensor, " + f"got shape {tuple(video.shape)}." + ) + if video.dtype != torch.uint8: + raise ValueError(f"Cosmos3 video reference must have dtype uint8, got {video.dtype}.") + + # Fields merged by the executor for every request. Modality-specific values # (height/width/num_frames/steps/guidance) are declared with ``None`` so the # executor accepts explicit overrides; ``forward()`` resolves unset fields from @@ -104,11 +147,13 @@ "pixel frames, so the worker consumes the first (or last, per " "condition_video_keep) max(indexes)*4+1 reference frames." ), + validator=_normalize_condition_video_latent_indexes, ), "condition_video_keep": ExtraParamSchema( type="str", default=COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP, description="Which side of the input video to use for conditioning: first or last.", + validator=_normalize_condition_video_keep, ), "flow_shift": ExtraParamSchema( type="float", @@ -126,5 +171,6 @@ "condition_video_latent_indexes / condition_video_keep and " "VAE-encodes it; media is always decoded by the producer." ), + validator=_validate_video_reference_tensor, ), } diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index 1204e497c55a..2f1c2b3ed476 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -36,11 +36,11 @@ from .defaults import ( COSMOS3_720P_PARAMS, - COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP, - COSMOS3_DEFAULT_CONDITION_VIDEO_LATENT_INDEXES, COSMOS3_EXTRA_SPECS, COSMOS3_PIPELINE_DEFAULTS, COSMOS3_T2I_PARAMS, + _normalize_condition_video_keep, + _normalize_condition_video_latent_indexes, ) from .guardrails import check_video_safety, download_guardrail_checkpoint from .sound_tokenizer import LatentAutoEncoderV2 @@ -62,22 +62,6 @@ TRTLLM_DISABLE_COSMOS3_GUARDRAILS = os.environ.get("TRTLLM_DISABLE_COSMOS3_GUARDRAILS", "0") == "1" -def _normalize_condition_video_latent_indexes( - indexes: Iterable[int] | None, -) -> tuple[int, ...]: - if indexes is None: - return COSMOS3_DEFAULT_CONDITION_VIDEO_LATENT_INDEXES - normalized = tuple(int(index) for index in indexes) - - if not normalized: - raise ValueError("Cosmos3 condition_video_latent_indexes must not be empty.") - if any(index < 0 for index in normalized): - raise ValueError( - f"Cosmos3 condition_video_latent_indexes must be non-negative, got {normalized}." - ) - return normalized - - def _condition_pixel_frame_count( condition_video_latent_indexes: Iterable[int], temporal_compression: int, @@ -85,13 +69,6 @@ def _condition_pixel_frame_count( return max(condition_video_latent_indexes) * int(temporal_compression) + 1 -def _normalize_condition_video_keep(keep: str | None) -> str: - normalized = str(keep or COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP).strip().lower() - if normalized not in {"first", "last"}: - raise ValueError("Cosmos3 condition_video_keep must be either first or last.") - return normalized - - @register_pipeline( "Cosmos3OmniMoTPipeline", hf_ids=[ @@ -763,7 +740,7 @@ def forward( use_guardrails: bool = COSMOS3_EXTRA_SPECS["use_guardrails"].default, enable_audio: bool = COSMOS3_EXTRA_SPECS["enable_audio"].default, output_type: str = COSMOS3_EXTRA_SPECS["output_type"].default, - video: torch.Tensor | None = None, # uint8 [T, H, W, C, dtype=uint8] + video: torch.Tensor | None = None, # [T, H, W, C, dtype=uint8] condition_video_latent_indexes: Iterable[int] | None = None, condition_video_keep: str | None = None, flow_shift: Optional[float] = None, diff --git a/tensorrt_llm/_torch/visual_gen/pipeline.py b/tensorrt_llm/_torch/visual_gen/pipeline.py index 4428de277928..b972b1c506ea 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline.py @@ -36,6 +36,10 @@ class ExtraParamSchema(StrictBaseModel): range: Optional[tuple] = Field( default=None, description="Optional (min, max) range for numeric params." ) + validator: Optional[Callable[[Any], Any]] = Field( + default=None, + description="Optional value validator; raises ValueError on invalid values.", + ) def _parse_profile_range(): diff --git a/tensorrt_llm/visual_gen/params.py b/tensorrt_llm/visual_gen/params.py index 2756ec79158d..ef41b060c560 100644 --- a/tensorrt_llm/visual_gen/params.py +++ b/tensorrt_llm/visual_gen/params.py @@ -93,7 +93,7 @@ class VisualGenParams(StrictBaseModel): "bool": (bool,), "str": (str,), "list": (list,), - "tensor": (torch.Tensor,), # Decoded media payloads (e.g. a V2V reference as a uint8 [T, H, W, C] + "tensor": (torch.Tensor,), } # Generation config fields that pipelines declare defaults for. If a user @@ -175,6 +175,16 @@ def validate_visual_gen_params( f"got {type(value).__name__}: {value!r}" ) continue # skip range check if type is wrong + # Validator (enums, bounds, tensor shapes) declared on + # the spec so deterministic client errors 400 at preflight + # instead of failing deep in the worker. + validator = getattr(spec, "validator", None) + if validator is not None: + try: + validator(value) + except ValueError as exc: + messages.append(f"extra_params['{key}']: {exc}") + continue # Range check (numeric only) if spec.range is not None and isinstance(value, (int, float)): lo, hi = spec.range diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py index e69c81d21f1b..7fb589126696 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py @@ -760,6 +760,48 @@ def test_valid_extra_params_accepted(self): req = self._make_request(extra_params={"stg_scale": 0.5}) self._merge_and_validate(executor, req) # should not raise + def test_spec_validator_runs_at_preflight(self): + """Per-param validators turn deterministic client errors into 400s at + the boundary instead of worker-side failures (Cosmos3 conditioning).""" + import torch + + from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import COSMOS3_EXTRA_SPECS + from tensorrt_llm.visual_gen.params import VisualGenParams, validate_visual_gen_params + + def _validate(extras): + validate_visual_gen_params( + VisualGenParams(extra_params=extras), + declared_defaults={}, + extra_param_specs=COSMOS3_EXTRA_SPECS, + ) + + # Valid values pass. + _validate({"condition_video_latent_indexes": [0, 1], "condition_video_keep": "last"}) + _validate({"video": torch.zeros(3, 4, 4, 3, dtype=torch.uint8)}) + + with pytest.raises(ValueError, match="non-negative"): + _validate({"condition_video_latent_indexes": [0, -1]}) + with pytest.raises(ValueError, match="must not be empty"): + _validate({"condition_video_latent_indexes": []}) + with pytest.raises(ValueError, match="first or last"): + _validate({"condition_video_keep": "middle"}) + with pytest.raises(ValueError, match=r"\[T, H, W, C\]"): + _validate({"video": torch.zeros(4, 4, 3, dtype=torch.uint8)}) # 3-D + with pytest.raises(ValueError, match="uint8"): + _validate({"video": torch.zeros(3, 4, 4, 3, dtype=torch.float32)}) + + def test_spec_validators_survive_pickling(self): + """Specs travel worker -> coordinator in the READY handshake (pickled + over ZMQ); validators must be module-level functions so they serialize + by reference — a lambda/closure here would crash worker startup.""" + import pickle + + from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import COSMOS3_EXTRA_SPECS + + specs = pickle.loads(pickle.dumps(COSMOS3_EXTRA_SPECS)) + with pytest.raises(ValueError, match="first or last"): + specs["condition_video_keep"].validator("middle") + # --- unsupported universal fields --- def test_num_frames_on_image_pipeline_raises(self): From a3183a9d4f8bec4419ea746fddbb4ea5f9890455 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 20 Jul 2026 15:50:10 -0700 Subject: [PATCH 35/64] Make `--use_system_prompt` optional with model-driven default Signed-off-by: Igor Shovkun --- examples/visual_gen/models/cosmos3/cosmos3.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/examples/visual_gen/models/cosmos3/cosmos3.py b/examples/visual_gen/models/cosmos3/cosmos3.py index 083a83b307cc..13afc4081f81 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3.py +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -223,7 +223,12 @@ 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="Force the system prompt on (--use_system_prompt) or off " + "(--no-use_system_prompt). Default: the model decides by mode " + "(on for V2V, off otherwise).", ) parser.add_argument("--enable_audio", action="store_true", help="Enable audio generation") parser.add_argument( @@ -274,7 +279,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 From 9eaa5efff54f1e396fdb9dbfeb83159ebce27d58 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 20 Jul 2026 21:49:01 -0700 Subject: [PATCH 36/64] Add V2V transport reducer to crop reference video before serialization Before this change, a full V2V reference (e.g. 189-frame 720p ~520 MiB) was deep-copied, pickled over ZMQ, and broadcast to every rank. The coordinator now crops the tensor to the conditioning window (~5 frames, ~14 MiB) before the copy/serialize path via a new `reducer` field on `ExtraParamSchema`. Key changes: - Add `reducer: Callable` field to `ExtraParamSchema` for semantics-preserving pre-serialization transforms - Implement `_crop_video_frames` reducer in Cosmos3 defaults; wire it onto the `video` extra-param spec - Add `reduce_visual_gen_params` in Signed-off-by: Igor Shovkun --- .../visual_gen/models/cosmos3/defaults.py | 45 +++- .../models/cosmos3/pipeline_cosmos3.py | 17 +- tensorrt_llm/_torch/visual_gen/pipeline.py | 9 + tensorrt_llm/inputs/media_io.py | 141 +++++++++++- tensorrt_llm/serve/openai_protocol.py | 4 +- tensorrt_llm/serve/visual_gen_utils.py | 18 +- tensorrt_llm/visual_gen/params.py | 33 +++ tensorrt_llm/visual_gen/visual_gen.py | 10 +- .../visual_gen/test_cosmos3_pipeline.py | 24 +- .../visual_gen/test_trtllm_serve_endpoints.py | 25 +++ .../visual_gen/test_visual_gen_params.py | 62 +++++- .../visual_gen/test_visual_gen_utils.py | 208 +++++++++++++++++- 12 files changed, 546 insertions(+), 50 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index f790cd86be94..9b61c3ea5b78 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -71,6 +71,36 @@ def _normalize_condition_video_keep(keep: str | None) -> str: return normalized +def _crop_video_frames(video, extra_params) -> torch.Tensor: + """Crop the V2V reference to the conditioning window before transport. + + Runs once in the coordinator (spec ``reducer``) before the request is + deep-copied, pickled over ZMQ, and broadcast per rank: a full 189-frame + 720p reference is ~520 MiB while the default conditioning window is 5 + frames (~14 MiB). Semantics-preserving — the worker's own first/last crop + is idempotent, so reduced and unreduced tensors generate identically. + Anything invalid is returned unchanged for the validators to reject. + """ + if not isinstance(video, torch.Tensor) or video.ndim != 4: + return video + try: + indexes = _normalize_condition_video_latent_indexes( + extra_params.get("condition_video_latent_indexes") + ) + keep = _normalize_condition_video_keep(extra_params.get("condition_video_keep")) + except (TypeError, ValueError): + return video + # 4 = Cosmos3 VAE temporal compression; if a future VAE changes it, the + # worker pads/crops the window itself, so a mismatch degrades gracefully. + window = max(indexes) * 4 + 1 + if video.shape[0] <= window: + return video + sliced = video[-window:] if keep == "last" else video[:window] + # A slice is a view over the full storage and would pickle all of it; + # clone so the transport payload owns only the window. + return sliced.clone() + + def _validate_video_reference_tensor(video: torch.Tensor) -> None: if video.ndim != 4 or video.shape[-1] != 3: raise ValueError( @@ -79,6 +109,11 @@ def _validate_video_reference_tensor(video: torch.Tensor) -> None: ) if video.dtype != torch.uint8: raise ValueError(f"Cosmos3 video reference must have dtype uint8, got {video.dtype}.") + if video.device.type != "cpu": + raise ValueError( + f"Cosmos3 video reference must be a CPU tensor, got device '{video.device}' " + "(it is pickled to the workers; keep decoded references on the host)." + ) # Fields merged by the executor for every request. Modality-specific values @@ -166,11 +201,13 @@ def _validate_video_reference_tensor(video: torch.Tensor) -> None: description=( "V2V reference: decoded video frames as a uint8 [T, H, W, C] RGB " "torch.Tensor (build one from a file with " - "tensorrt_llm.inputs.media_io.load_video_frames_tensor). The worker " - "keeps the first/last conditioning window per " - "condition_video_latent_indexes / condition_video_keep and " - "VAE-encodes it; media is always decoded by the producer." + "tensorrt_llm.inputs.media_io.load_video_frames_tensor). The " + "coordinator crops it to the conditioning window per " + "condition_video_latent_indexes / condition_video_keep before " + "dispatch; the worker VAE-encodes it. Media is always decoded " + "by the producer." ), validator=_validate_video_reference_tensor, + reducer=_crop_video_frames, ), } diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index 2f1c2b3ed476..d29793ab23c9 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -39,7 +39,6 @@ COSMOS3_EXTRA_SPECS, COSMOS3_PIPELINE_DEFAULTS, COSMOS3_T2I_PARAMS, - _normalize_condition_video_keep, _normalize_condition_video_latent_indexes, ) from .guardrails import check_video_safety, download_guardrail_checkpoint @@ -937,26 +936,14 @@ def forward( condition_video_latent_indexes = _normalize_condition_video_latent_indexes( condition_video_latent_indexes ) - condition_video_keep = _normalize_condition_video_keep(condition_video_keep) - condition_pixel_frames = min( - _condition_pixel_frame_count( - condition_video_latent_indexes, self.vae_scale_factor_temporal - ), - num_frames, - ) if not isinstance(video, torch.Tensor) or video.ndim != 4: raise ValueError( "Cosmos3 V2V reference must be a uint8 [T, H, W, C] tensor " f"(the 'video' extra-param contract), got {type(video).__name__}" f"{' of shape ' + str(tuple(video.shape)) if isinstance(video, torch.Tensor) else ''}." ) - # Crop the first/last conditioning window uniformly - window = ( - video[-condition_pixel_frames:] - if condition_video_keep == "last" - else video[:condition_pixel_frames] - ) - frames = [PIL.Image.fromarray(frame.cpu().numpy()) for frame in window] + # video is already the conditioning window: the coordinator crops it + frames = [PIL.Image.fromarray(frame.cpu().numpy()) for frame in video] video = self._preprocess_condition_video(frames, height, width) if self.rank == 0: diff --git a/tensorrt_llm/_torch/visual_gen/pipeline.py b/tensorrt_llm/_torch/visual_gen/pipeline.py index b972b1c506ea..feabc7452df0 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline.py @@ -40,6 +40,15 @@ class ExtraParamSchema(StrictBaseModel): default=None, description="Optional value validator; raises ValueError on invalid values.", ) + # Like ``validator``, must be a module-level function (specs are pickled to + # the coordinator in the READY handshake). + reducer: Optional[Callable[[Any, Dict[str, Any]], Any]] = Field( + default=None, + description="Optional transport reducer, run once in the coordinator " + "before the request is copied/serialized: (value, extra_params) -> " + "reduced value. Must be semantics-preserving (the worker treats " + "reduced and unreduced values identically).", + ) def _parse_profile_range(): diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index 1a160b93ce7f..1db4fbc06321 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -394,15 +394,20 @@ def is_decodable_video_file(path) -> bool: def is_decodable_image_bytes(data) -> bool: - """True when ``data`` holds still-image content (anything PIL opens). + """True when ``data`` holds still-image content PIL can fully decode. - In-memory counterpart of :func:`is_decodable_image_file` — header-only probe, no - pixel decode, no filesystem. + In-memory counterpart of :func:`is_decodable_image_file`, but strict: + ``Image.open`` is lazy, so this also decodes the pixels (``load``) — + a truncated file passes a header-only probe and would then 500 at the + worker's load instead of 400ing at the boundary. """ try: - with Image.open(BytesIO(data)): + with Image.open(BytesIO(data)) as image: + image.load() return True - except UnidentifiedImageError: + except OSError: + # ``UnidentifiedImageError`` (bad header) subclasses ``OSError``; + # truncated files raise plain ``OSError`` from ``load``. return False @@ -474,6 +479,130 @@ def decode_video_frames_from_bytes(data, max_frames: Optional[int] = None) -> Li _IMAGE_SUFFIXES = frozenset({".png", ".jpg", ".jpeg", ".webp", ".bmp"}) +# Longest video, in frames, a client may request as *output*; used by the +# serve's ``num_frames`` cap (``openai_protocol``). +MAX_VIDEO_FRAMES = 7200 + +# Hard budget for *decoded* video bytes when reading a reference. Bounds both +# the preallocated buffer and total accumulation, whatever the resolution — +# a frame count alone is no guard (7200 frames is ~18.5 GiB at 720p). The +# canonical 189-frame 720p reference is ~0.5 GiB, so this allows 2x headroom. +MAX_DECODED_VIDEO_BYTES = 1 << 30 # 1 GiB + + +class DecodedVideoTooLargeError(ValueError): + """Decoded reference exceeds ``MAX_DECODED_VIDEO_BYTES``. + + A ``ValueError`` subclass so boundary handlers still map it to a client + error (400), while letting its actionable message pass through instead of + being folded into generic "undecodable content" handling. + """ + + +def _decode_capture_to_tensor( + cv2, capture, max_frames: Optional[int], src_repr: str +) -> torch.Tensor: + """Drain an opened ``VideoCapture`` straight into a uint8 [T, H, W, C] tensor. + + Streams frames into one preallocated array, so with accurate container + metadata peak memory is a single copy of the video — roughly half of the + decode-to-PIL-list-then-``np.stack`` route. Unknown or misreported lengths + take the spill paths below, whose ``stack``/``concatenate`` transients can + reach ~2-3x the decoded size — still bounded, since everything is capped + by ``MAX_DECODED_VIDEO_BYTES`` + measured against real frame sizes: the preallocation (the declared frame + count is container metadata, not evidence) and the total decoded + accumulation (streams that exceed the budget raise instead of growing + without bound). Misreported lengths degrade gracefully: extra frames spill + to a side list, an over-declared buffer is trimmed, and an unknown length + falls back to list-and-stack. + """ + if not capture.isOpened(): + raise ValueError(f"Could not open video from {src_repr}.") + declared = int(capture.get(cv2.CAP_PROP_FRAME_COUNT) or 0) + + buffer = None + overflow = [] + count = 0 + decoded_bytes = 0 + while max_frames is None or count < max_frames: + ok, frame = capture.read() + if not ok: + break + rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + decoded_bytes += rgb.nbytes + if decoded_bytes > MAX_DECODED_VIDEO_BYTES: + raise DecodedVideoTooLargeError( + f"Video from {src_repr} exceeds the decoded-size budget of " + f"{MAX_DECODED_VIDEO_BYTES >> 20} MiB at frame {count}; trim or " + "downscale the reference." + ) + if buffer is None and declared > 0: + # Size the preallocation from the declared count, clamped to the + # byte budget using the actual frame size (an over-declared buffer + # is virtual until written; the trim below drops the excess). + capacity = min(declared, MAX_DECODED_VIDEO_BYTES // max(rgb.nbytes, 1)) + if max_frames is not None: + capacity = min(capacity, max_frames) + if capacity > 0: + buffer = np.empty((capacity, *rgb.shape), dtype=np.uint8) + if buffer is not None and count < buffer.shape[0]: + buffer[count] = rgb + else: + overflow.append(rgb) + count += 1 + + if count == 0: + raise ValueError(f"Video contains no frames ({src_repr}).") + if buffer is None: + return torch.from_numpy(np.stack(overflow)) + filled = min(count, buffer.shape[0]) + if overflow: + return torch.from_numpy(np.concatenate([buffer[:filled], np.stack(overflow)])) + if filled < buffer.shape[0]: + # Over-declared container: trim without keeping the oversized buffer. + return torch.from_numpy(buffer[:filled].copy()) + return torch.from_numpy(buffer) + + +def decode_video_tensor(path, max_frames: Optional[int] = None) -> torch.Tensor: + """Decode a video file into a uint8 ``[T, H, W, C]`` RGB tensor. + + Tensor-native counterpart of :func:`decode_video_frames` — streams into a + single buffer instead of materializing PIL frames first. + """ + cv2 = _get_cv2() + capture = cv2.VideoCapture(str(path)) + try: + return _decode_capture_to_tensor(cv2, capture, max_frames, f"'{path}'") + finally: + capture.release() + + +def decode_video_tensor_from_bytes(data, max_frames: Optional[int] = None) -> torch.Tensor: + """Decode raw video bytes into a uint8 ``[T, H, W, C]`` RGB tensor. + + In-memory when this OpenCV build has a stream-buffered backend; otherwise + the bytes spill to an auto-deleted tempfile. Raises ``ValueError`` when the + bytes are not a decodable video. + """ + cv2 = _get_cv2() + backend = _select_cv2_stream_buffered_backend() + if backend is None: + with tempfile.NamedTemporaryFile() as spill: + spill.write(data) + spill.flush() + return decode_video_tensor(spill.name, max_frames=max_frames) + + # cv2 keeps a non-owning view into the buffer; hold it until release(). + buffer = BytesIO(bytes(data)) + capture = cv2.VideoCapture(buffer, backend, []) + try: + return _decode_capture_to_tensor(cv2, capture, max_frames, f"<{len(data)} bytes>") + finally: + capture.release() + + def frames_to_tensor(frames: List["Image.Image"]) -> torch.Tensor: """Stack PIL frames into a uint8 ``[T, H, W, C]`` RGB CPU tensor. @@ -513,7 +642,7 @@ def load_video_frames_tensor(source, max_frames: Optional[int] = None) -> torch. if is_decodable_image_file(path): return frames_to_tensor([Image.open(path)]) if is_decodable_video_file(path): - return frames_to_tensor(decode_video_frames(path, max_frames=max_frames)) + return decode_video_tensor(path, max_frames=max_frames) raise ValueError( f"Video reference must be a decodable video, a decodable image, or a " f"directory of frame images; got undecodable {path}" diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index 300c6a53e375..bedfe219eafb 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -37,7 +37,7 @@ from typing_extensions import Annotated, Required, TypeAlias, TypedDict from tensorrt_llm.executor.request import LoRARequest -from tensorrt_llm.inputs.media_io import MediaModality +from tensorrt_llm.inputs.media_io import MAX_VIDEO_FRAMES, MediaModality from tensorrt_llm.llmapi import ConversationParams as LlmConversationParams from tensorrt_llm.llmapi import DisaggregatedParams as LlmDisaggregatedParams from tensorrt_llm.llmapi import (DisaggScheduleStyle, GuidedDecodingParams, @@ -1713,7 +1713,7 @@ class VideoGenerationRequest(OpenAIBaseModel): # The numbers are generous (a minute of video at 120 fps) so common # workloads pass; clients that need larger budgets can lift the cap # at deployment time. - num_frames: Optional[int] = Field(default=None, gt=0, le=7200) + num_frames: Optional[int] = Field(default=None, gt=0, le=MAX_VIDEO_FRAMES) seconds: Optional[float] = Field(default=None, gt=0, le=60.0) frame_rate: Optional[float] = Field(default=None, alias="fps", diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index e3bbfefbca45..a779ddb2fb3b 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -4,13 +4,14 @@ from typing import Any, Dict, List, Optional from tensorrt_llm.inputs.media_io import ( - decode_video_frames_from_bytes, - frames_to_tensor, + DecodedVideoTooLargeError, + decode_video_tensor_from_bytes, is_decodable_image_bytes, ) from tensorrt_llm.logger import logger from tensorrt_llm.serve.openai_protocol import ImageGenerationRequest, VideoGenerationRequest from tensorrt_llm.visual_gen import VisualGen, VisualGenParams +from tensorrt_llm.visual_gen.params import reduce_visual_gen_params # Per-field warnings for OpenAI-shaped knobs that the engine has no # semantic for. Each entry maps the request attribute to the message @@ -180,7 +181,10 @@ def parse_visual_gen_params( else: # V2V: decode in memory into a uint8 [T, H, W, C] tensor try: - frames = decode_video_frames_from_bytes(payload) + video = decode_video_tensor_from_bytes(payload) + except DecodedVideoTooLargeError: + # Still a 400, but with the actionable size message intact. + raise except ValueError as exc: raise ValueError( "input_reference content is neither a decodable image " @@ -188,12 +192,16 @@ def parse_visual_gen_params( ) from exc if params.extra_params is None: params.extra_params = {} - params.extra_params["video"] = frames_to_tensor(frames) + params.extra_params["video"] = video _warn_if_set_with_no_semantic(request, getattr(generator, "model", None)) _merge_extra_params(params, request.extra_params, generator.extra_param_specs) - return params + # Apply spec-declared transport reducers here as well (generate_async + # reduces non-mutatively, so without this the serve-owned params — held by + # the sync/async routes for the job's whole lifetime — would retain the + # full decoded reference, e.g. ~500 MiB per queued V2V request). + return reduce_visual_gen_params(params, generator.extra_param_specs) class AsyncDictStore: diff --git a/tensorrt_llm/visual_gen/params.py b/tensorrt_llm/visual_gen/params.py index ef41b060c560..cdb85a107a78 100644 --- a/tensorrt_llm/visual_gen/params.py +++ b/tensorrt_llm/visual_gen/params.py @@ -113,6 +113,39 @@ class VisualGenParams(StrictBaseModel): ) +def reduce_visual_gen_params( + params: VisualGenParams, + extra_param_specs: Dict[str, Any], +) -> VisualGenParams: + """Apply spec-declared transport reducers to ``extra_params`` values. + + Runs once in the coordinator, before the request is deep-copied and + serialized, so oversized payloads (e.g. a full V2V reference when only + the conditioning window is consumed) shrink before they hit the copy / + ZMQ / per-rank broadcast path. Reducers are semantics-preserving by + contract — the worker behaves identically with or without them. + + Never mutates ``params``: returns it unchanged when nothing reduces, + otherwise a shallow copy carrying a new ``extra_params`` dict. + """ + if not params.extra_params: + return params + reduced: Dict[str, Any] = {} + for key, value in params.extra_params.items(): + spec = extra_param_specs.get(key) + reducer = getattr(spec, "reducer", None) if spec is not None else None + if reducer is None or value is None: + continue + new_value = reducer(value, params.extra_params) + if new_value is not value: + reduced[key] = new_value + if not reduced: + return params + out = params.model_copy() + out.extra_params = {**params.extra_params, **reduced} + return out + + def validate_visual_gen_params( params: VisualGenParams, *, diff --git a/tensorrt_llm/visual_gen/visual_gen.py b/tensorrt_llm/visual_gen/visual_gen.py index ce83f862326f..c740ae7d140c 100644 --- a/tensorrt_llm/visual_gen/visual_gen.py +++ b/tensorrt_llm/visual_gen/visual_gen.py @@ -32,7 +32,11 @@ from tensorrt_llm._torch.visual_gen.pipeline_registry import PIPELINE_REGISTRY, AutoPipeline from tensorrt_llm.visual_gen.args import VisualGenArgs from tensorrt_llm.visual_gen.output import VisualGenOutput -from tensorrt_llm.visual_gen.params import VisualGenParams, validate_visual_gen_params +from tensorrt_llm.visual_gen.params import ( + VisualGenParams, + reduce_visual_gen_params, + validate_visual_gen_params, +) __all__ = [ "VisualGen", @@ -387,6 +391,10 @@ def generate_async( # from the READY signal) and skip validation — there's nothing # user-supplied to validate against. if params is not None: + # Shrink oversized transport payloads (spec-declared reducers) + # before the deep copy, so the copy/pickle/broadcast chain only + # ever carries the reduced values. Non-mutating for the caller. + params = reduce_visual_gen_params(params, self.executor.extra_param_specs) resolved_params = params.model_copy(deep=True) # Raising in the caller's process means ``ValueError`` reaches # the user as a natural Python exception; the worker only has diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index 3b155a23b105..554ece9958a5 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -43,6 +43,8 @@ COSMOS3_DEFAULT_CONDITION_VIDEO_LATENT_INDEXES, COSMOS3_EXTRA_SPECS, COSMOS3_T2I_PARAMS, + _crop_video_frames, + _normalize_condition_video_keep, ) from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import ( COSMOS3_DEFAULT_RESOLUTION_TEMPLATE, @@ -51,7 +53,6 @@ COSMOS3_IMAGE_RESOLUTION_TEMPLATE, Cosmos3OmniMoTPipeline, _condition_pixel_frame_count, - _normalize_condition_video_keep, _normalize_condition_video_latent_indexes, ) from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader @@ -656,7 +657,6 @@ def test_v2v_smoke(self, cosmos3_pipeline): video=video, num_frames=NUM_FRAMES, condition_video_latent_indexes=[0, 1], - condition_video_keep="first", ) _assert_valid_video(result.video, num_frames=NUM_FRAMES) assert result.frame_rate == FRAME_RATE @@ -668,9 +668,9 @@ def test_v2v_smoke(self, cosmos3_pipeline): def test_v2v_tensor_reference_smoke(self, cosmos3_pipeline): """The V2V reference arrives as a decoded uint8 [T, H, W, C] tensor - (the ``video`` extra-param contract); the worker crops the conditioning - window and VAE-encodes. Mirrors what the offline example and serve feed - the pipeline.""" + (the ``video`` extra-param contract, cropped by the coordinator's + reducer in real requests); the worker VAE-encodes it, capping/padding + to the latent window in ``_prepare_latents_v2v``.""" from tensorrt_llm.inputs.media_io import frames_to_tensor video = frames_to_tensor(_make_test_video(NUM_FRAMES)) @@ -681,17 +681,17 @@ def test_v2v_tensor_reference_smoke(self, cosmos3_pipeline): video=video, num_frames=NUM_FRAMES, condition_video_latent_indexes=[0, 1], - condition_video_keep="first", ) _assert_valid_video(result.video, num_frames=NUM_FRAMES) def test_v2v_keep_last_smoke(self, cosmos3_pipeline): """condition_video_keep="last" pins the tail of the input, not the head. - The input is longer than the conditioning window and color-coded - (dark head, bright tail), so this exercises the full-decode + - tail-slice path and asserts behavior: frame 0 of the output is a - pinned VAE round-trip of the bright tail frames. + ``keep`` is consumed by the coordinator-side reducer + (``_crop_video_frames``), so this test composes reducer + forward the + way a real request flows. The input is longer than the conditioning + window and color-coded (dark head, bright tail); frame 0 of the output + must be a pinned VAE round-trip of the bright tail frames. """ dark = PIL.Image.new("RGB", (WIDTH, HEIGHT), (40, 40, 40)) bright = PIL.Image.new("RGB", (WIDTH, HEIGHT), (230, 230, 230)) @@ -699,13 +699,14 @@ def test_v2v_keep_last_smoke(self, cosmos3_pipeline): video = frames_to_tensor( [dark.copy() for _ in range(NUM_FRAMES)] + [bright.copy() for _ in range(5)] ) + video = _crop_video_frames(video, {"condition_video_keep": "last"}) + assert video.shape[0] == 5 result = _run_forward( cosmos3_pipeline, image=None, video=video, num_frames=NUM_FRAMES, condition_video_latent_indexes=[0, 1], - condition_video_keep="last", ) _assert_valid_video(result.video, num_frames=NUM_FRAMES) first_frame_mean = result.video[0, 0].float().mean().item() @@ -824,7 +825,6 @@ def test_v2v_audio_smoke(self, cosmos3_pipeline): enable_audio=True, video=frames_to_tensor(_make_test_video(NUM_FRAMES)), condition_video_latent_indexes=[0, 1], - condition_video_keep="first", ) _assert_valid_video(result.video, num_frames=NUM_FRAMES) _assert_valid_audio(result.audio, result.audio_sample_rate) diff --git a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py index 1295d24c3a6c..a1558797a060 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py @@ -911,6 +911,31 @@ def test_sync_video_generation_undecodable_reference_400(self, video_client): assert resp.status_code == 400 assert "neither a decodable image" in resp.text + def test_sync_video_oversized_reference_400_with_message( + self, video_client, tmp_path, monkeypatch + ): + """A reference over the decoded-byte budget gets an HTTP 400 whose body + carries the actionable size message (not the generic undecodable one).""" + cv2 = _require_opencv() + from tensorrt_llm.inputs import media_io + + monkeypatch.setattr(media_io, "MAX_DECODED_VIDEO_BYTES", 100) + ref_path = tmp_path / "ref.mp4" + writer = cv2.VideoWriter(str(ref_path), cv2.VideoWriter_fourcc(*"mp4v"), 4.0, (16, 16)) + try: + for _ in range(4): + writer.write(np.zeros((16, 16, 3), dtype=np.uint8)) + finally: + writer.release() + with open(ref_path, "rb") as f: + resp = video_client.post( + "/v1/videos/generations", + data={"prompt": "x"}, + files={"input_reference": ("ref.mp4", f, "video/mp4")}, + ) + assert resp.status_code == 400 + assert "decoded-size budget" in resp.text + def test_sync_video_failure(self, failing_client): resp = failing_client.post( "/v1/videos/generations", diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py index 7fb589126696..9a7ba0373b06 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py @@ -789,18 +789,76 @@ def _validate(extras): _validate({"video": torch.zeros(4, 4, 3, dtype=torch.uint8)}) # 3-D with pytest.raises(ValueError, match="uint8"): _validate({"video": torch.zeros(3, 4, 4, 3, dtype=torch.float32)}) + with pytest.raises(ValueError, match="CPU tensor"): + _validate({"video": torch.zeros(3, 4, 4, 3, dtype=torch.uint8, device="meta")}) def test_spec_validators_survive_pickling(self): """Specs travel worker -> coordinator in the READY handshake (pickled - over ZMQ); validators must be module-level functions so they serialize - by reference — a lambda/closure here would crash worker startup.""" + over ZMQ); validators/reducers must be module-level functions so they + serialize by reference — a lambda/closure here would crash worker + startup.""" import pickle + import torch + from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import COSMOS3_EXTRA_SPECS specs = pickle.loads(pickle.dumps(COSMOS3_EXTRA_SPECS)) with pytest.raises(ValueError, match="first or last"): specs["condition_video_keep"].validator("middle") + reduced = specs["video"].reducer(torch.zeros(20, 4, 4, 3, dtype=torch.uint8), {}) + assert reduced.shape[0] == 5 + + def test_video_transport_reducer_crops_to_conditioning_window(self): + """The coordinator-side reducer ships only the conditioning window — + never the full clip — and the payload owns its storage (a bare slice + would pickle the entire original tensor).""" + import torch + + from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import _crop_video_frames + + # Frame k is solid value k, so window position is observable. + full = torch.arange(189, dtype=torch.uint8).view(189, 1, 1, 1).expand(189, 4, 4, 3) + full = full.contiguous() + + first = _crop_video_frames(full, {}) + assert first.shape[0] == 5 # default indexes (0, 1) -> 1*4+1 + assert int(first[0, 0, 0, 0]) == 0 and int(first[-1, 0, 0, 0]) == 4 + # Owns its storage: pickling must carry the window, not the clip. + assert first.untyped_storage().size() < full.untyped_storage().size() + + last = _crop_video_frames(full, {"condition_video_keep": "last"}) + assert last.shape[0] == 5 and int(last[-1, 0, 0, 0]) == 188 + + wider = _crop_video_frames(full, {"condition_video_latent_indexes": [0, 2]}) + assert wider.shape[0] == 9 # 2*4+1 + + # Short-enough inputs and invalid context pass through unchanged. + short = torch.zeros(3, 4, 4, 3, dtype=torch.uint8) + assert _crop_video_frames(short, {}) is short + assert _crop_video_frames(full, {"condition_video_latent_indexes": [-1]}) is full + assert _crop_video_frames("not a tensor", {}) == "not a tensor" + + def test_reduce_visual_gen_params_is_non_mutating(self): + """generate_async reduces before the deep copy; the caller's params + object and tensor must be untouched.""" + import torch + + from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import COSMOS3_EXTRA_SPECS + from tensorrt_llm.visual_gen.params import VisualGenParams, reduce_visual_gen_params + + full = torch.zeros(189, 4, 4, 3, dtype=torch.uint8) + params = VisualGenParams(extra_params={"video": full, "flow_shift": 10.0}) + out = reduce_visual_gen_params(params, COSMOS3_EXTRA_SPECS) + + assert out is not params + assert params.extra_params["video"] is full # caller untouched + assert out.extra_params["video"].shape[0] == 5 + assert out.extra_params["flow_shift"] == 10.0 # non-reduced keys intact + + # Nothing to reduce -> same object back, zero copies. + plain = VisualGenParams(extra_params={"flow_shift": 10.0}) + assert reduce_visual_gen_params(plain, COSMOS3_EXTRA_SPECS) is plain # --- unsupported universal fields --- diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index 4ad2f99c3b2c..b67ebc14d006 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -294,8 +294,8 @@ def test_missing_media_storage_path_raises(self): parse_visual_gen_params(request, "vid-2", generator, media_storage_path=None) @staticmethod - def _mp4_bytes() -> bytes: - """Encode a 2-frame 16x16 mp4v-in-mp4 clip and return its bytes. + def _mp4_bytes(num_frames: int = 2) -> bytes: + """Encode a 16x16 mp4v-in-mp4 clip and return its bytes. ``mp4v`` is a built-in FFmpeg mpeg4 encoder present in the opencv wheel. OpenCV writes only to a path, so encode to a tempfile and read it back. @@ -305,7 +305,7 @@ def _mp4_bytes() -> bytes: path = tmp.name try: writer = cv2.VideoWriter(path, cv2.VideoWriter_fourcc(*"mp4v"), 4.0, (16, 16)) - for _ in range(2): + for _ in range(num_frames): writer.write(np.zeros((16, 16, 3), dtype=np.uint8)) writer.release() with open(path, "rb") as f: @@ -358,6 +358,34 @@ def test_base64_video_reference_routes_to_extra_params_tensor(self, tmp_path): assert params.image is None assert isinstance(params.extra_params["video"], torch.Tensor) + def test_video_reference_reduced_before_routes_hold_params(self): + """``parse_visual_gen_params`` applies the spec reducers itself: the + sync/async routes hold the returned params for the whole job lifetime, + and ``generate_async`` reduces non-mutatively — without reduction at + parse, the serve would retain the full decoded clip per queued + request.""" + from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import COSMOS3_EXTRA_SPECS + + generator = _StubVisualGen(extra_param_specs=COSMOS3_EXTRA_SPECS) + b64 = base64.b64encode(self._mp4_bytes(num_frames=8)).decode() + request = VideoGenerationRequest(prompt="x", input_reference=b64) + params = parse_visual_gen_params(request, "vid-10", generator, media_storage_path=None) + # Cropped to the default conditioning window (5) at parse, not later. + assert params.extra_params["video"].shape[0] == 5 + + def test_budget_error_message_survives_parse(self, monkeypatch): + # Helper-level: DecodedVideoTooLargeError passes through the generic + # "undecodable" handler with its message intact. The HTTP 400 itself + # is asserted in test_trtllm_serve_endpoints.py. + from tensorrt_llm.inputs import media_io + + monkeypatch.setattr(media_io, "MAX_DECODED_VIDEO_BYTES", 100) + generator = _StubVisualGen() + b64 = base64.b64encode(self._mp4_bytes(num_frames=4)).decode() + request = VideoGenerationRequest(prompt="x", input_reference=b64) + with pytest.raises(ValueError, match="decoded-size budget"): + parse_visual_gen_params(request, "vid-11", generator, media_storage_path=None) + def test_multipart_image_reference_routes_to_image(self, tmp_path): # JPEG upload: content sniffing classifies it as an image and routes # to params.image. The stored file has no type-suffix (PIL identifies @@ -447,6 +475,65 @@ def test_decode_video_frames_from_bytes_rejects_garbage(self): with pytest.raises(ValueError): decode_video_frames_from_bytes(b"not a video at all") + def test_truncated_image_bytes_are_not_decodable(self): + # A truncated PNG still opens (the header parses) but cannot decode + # its pixels; the probe must reject it so the boundary 400s instead + # of the worker 500ing at load time. + rng_pixels = np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8) + buf = BytesIO() + Image.fromarray(rng_pixels).save(buf, format="PNG") + whole = buf.getvalue() + truncated = whole[: len(whole) // 2] + Image.open(BytesIO(truncated)) # sanity: header-only open succeeds + + from tensorrt_llm.inputs.media_io import is_decodable_image_bytes + + assert is_decodable_image_bytes(whole) + assert not is_decodable_image_bytes(truncated) + + def test_truncated_image_reference_rejected_at_parse(self): + # End of the chain: a truncated image upload is rejected as a client + # error at the boundary (never routed into the worker). + pytest.importorskip("cv2") + rng_pixels = np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8) + buf = BytesIO() + Image.fromarray(rng_pixels).save(buf, format="PNG") + truncated = buf.getvalue()[: len(buf.getvalue()) // 2] + + generator = _StubVisualGen() + request = VideoGenerationRequest( + prompt="x", input_reference=base64.b64encode(truncated).decode() + ) + with pytest.raises(ValueError, match="neither a decodable"): + parse_visual_gen_params(request, "vid-12", generator, media_storage_path=None) + + def test_decode_video_tensor_matches_pil_route(self): + # The streaming decoder (single preallocated buffer — the low-peak + # path the serve uses) must produce byte-identical output to the + # PIL-frames route. + pytest.importorskip("cv2") + import torch + + from tensorrt_llm.inputs.media_io import ( + decode_video_frames_from_bytes, + decode_video_tensor_from_bytes, + frames_to_tensor, + ) + + data = TestInputReferenceMaterialization._mp4_bytes() + streamed = decode_video_tensor_from_bytes(data) + via_pil = frames_to_tensor(decode_video_frames_from_bytes(data)) + assert streamed.dtype == torch.uint8 and streamed.ndim == 4 + assert torch.equal(streamed, via_pil) + assert torch.equal(decode_video_tensor_from_bytes(data, max_frames=1), via_pil[:1]) + + def test_decode_video_tensor_rejects_garbage(self): + pytest.importorskip("cv2") + from tensorrt_llm.inputs.media_io import decode_video_tensor_from_bytes + + with pytest.raises(ValueError): + decode_video_tensor_from_bytes(b"not a video at all") + def test_tempfile_fallback_without_stream_backend(self, monkeypatch): # Old OpenCV builds have no stream-buffered backend; the bytes spill # to an auto-deleted tempfile and decode through the path route. @@ -511,3 +598,118 @@ def test_empty_extras_dict_normalizes_to_none(self): params = self._make_params() _merge_extra_params(params, request_extras=None, extra_param_specs={}) assert params.extra_params is None + + +class _FakeCapture: + """Stands in for ``cv2.VideoCapture`` to exercise declared-count handling.""" + + def __init__(self, frames, declared): + self._frames = list(frames) + self._pos = 0 + self._declared = declared + + def isOpened(self): + return True + + def get(self, prop): + return self._declared + + def read(self): + if self._pos < len(self._frames): + frame = self._frames[self._pos] + self._pos += 1 + return True, frame + return False, None + + +class _FakeCv2: + CAP_PROP_FRAME_COUNT = 7 + COLOR_BGR2RGB = 4 + + @staticmethod + def cvtColor(frame, code): + return frame + + +class TestDecodeCaptureGuards: + """``_decode_capture_to_tensor`` against containers that misreport length. + + The declared frame count is metadata, not evidence — the decoder must + stream correctly whether it is accurate, unknown, under-, over-, or + absurdly reported.""" + + def _decode(self, num_frames, declared, max_frames=None): + import torch + + from tensorrt_llm.inputs.media_io import _decode_capture_to_tensor + + frames = [np.full((4, 4, 3), i, dtype=np.uint8) for i in range(num_frames)] + out = _decode_capture_to_tensor( + _FakeCv2, _FakeCapture(frames, declared), max_frames, "test" + ) + assert out.dtype == torch.uint8 + return out + + def test_accurate_declaration(self): + out = self._decode(3, declared=3) + assert out.shape == (3, 4, 4, 3) + assert int(out[2, 0, 0, 0]) == 2 # frame order preserved + + def test_unknown_declaration_falls_back(self): + assert self._decode(3, declared=0).shape == (3, 4, 4, 3) + assert self._decode(3, declared=-1).shape == (3, 4, 4, 3) + + def test_underreported_declaration_keeps_overflow(self): + out = self._decode(5, declared=2) + assert out.shape == (5, 4, 4, 3) + assert int(out[4, 0, 0, 0]) == 4 + + def test_overreported_declaration_trims_storage(self): + out = self._decode(3, declared=10) + assert out.shape == (3, 4, 4, 3) + # The oversized buffer is not retained behind the result. + assert out.untyped_storage().size() == out.numel() + + def test_absurd_declaration_allocation_is_byte_budgeted(self, monkeypatch): + # Realistic 720p frames + an absurd declared count: the preallocation + # request itself must stay within MAX_DECODED_VIDEO_BYTES (a frame + # cap alone would still be ~18.5 GiB at 720p). Spy on np.empty to + # assert the requested size, not just the result. + import math + + from tensorrt_llm.inputs import media_io + + requested = [] + real_empty = media_io.np.empty + + def spy(shape, dtype=None): + requested.append((tuple(shape), dtype)) + return real_empty(shape, dtype=dtype) + + monkeypatch.setattr(media_io.np, "empty", spy) + # Lower the budget so the (real) allocation the spy delegates to stays + # small; the assertion is about the *requested* size honoring it. + monkeypatch.setattr(media_io, "MAX_DECODED_VIDEO_BYTES", 32 << 20) + frames = [np.zeros((720, 1280, 3), dtype=np.uint8) for _ in range(3)] + out = media_io._decode_capture_to_tensor( + _FakeCv2, _FakeCapture(frames, declared=10**9), None, "test" + ) + assert out.shape == (3, 720, 1280, 3) + (shape, _dtype) = requested[0] + assert math.prod(shape) <= media_io.MAX_DECODED_VIDEO_BYTES + + def test_decoded_byte_budget_rejects_oversized_streams(self, monkeypatch): + # Total accumulation is bounded too — a stream that exceeds the budget + # raises instead of growing without bound (unknown-length containers + # included, where no buffer is preallocated at all). + from tensorrt_llm.inputs import media_io + + monkeypatch.setattr(media_io, "MAX_DECODED_VIDEO_BYTES", 100) + frames = [np.zeros((4, 4, 3), dtype=np.uint8) for _ in range(5)] # 48 B each + with pytest.raises(ValueError, match="decoded-size budget"): + media_io._decode_capture_to_tensor( + _FakeCv2, _FakeCapture(frames, declared=0), None, "test" + ) + + def test_max_frames_bounds_decode(self): + assert self._decode(5, declared=5, max_frames=2).shape[0] == 2 From 750598bd95e9c7b833e8180c2229faaaa5d8a9df Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 20 Jul 2026 21:59:08 -0700 Subject: [PATCH 37/64] Fix Cosmos3 V2V post-step latent anchoring and pin condition - Simplify `post_step_fn` to unconditionally apply velocity mask blending (V2V path); remove the I2V per-step write-back branch since I2V anchoring is handled separately outside the loop - Fix `should_pin_condition_latents` guard: require both `condition_latents` and `velocity_mask` to be non-None, preventing spurious pinning when only one is set Signed-off-by: Igor Shovkun --- .../models/cosmos3/pipeline_cosmos3.py | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index d29793ab23c9..854fe849d93e 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -1025,16 +1025,12 @@ def forward_fn( return video_noise_pred def post_step_fn(step_latents): - if velocity_mask is not None and condition_latents is not None: - step_latents = ( - velocity_mask * step_latents + (1.0 - velocity_mask) * condition_latents - ) - elif velocity_mask is not None and image_latent is not None: - step_latents = step_latents.clone() - step_latents[:, :, 0:1, :, :] = image_latent.to( - device=step_latents.device, dtype=step_latents.dtype - ) - return step_latents + # V2V only: re-impose the clean condition latents after every + # scheduler step. I2V deliberately keeps its pre-existing behavior + # (velocity mask during the loop, one write-back after it) so this + # PR does not alter I2V denoising; per-step anchoring for + # stochastic distilled schedulers belongs to the distilled work. + return velocity_mask * step_latents + (1.0 - velocity_mask) * condition_latents # 5. Build CFG tensors — text_ids and text_mask need to be split for CFG # BasePipeline.denoise batches [uncond, cond] when guidance_scale > 1 @@ -1051,7 +1047,7 @@ def post_step_fn(step_latents): extra_streams = None if do_audio: extra_streams = {"audio": (audio_latents, self.audio_scheduler)} - should_pin_condition_latents = condition_latents is not None or image_latent is not None + should_pin_condition_latents = condition_latents is not None and velocity_mask is not None denoise_result = self.denoise( latents=latents, scheduler=self.scheduler, From 6925f0f78dfe640a4217c1c8d803a28b4f09e8b6 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 20 Jul 2026 22:14:25 -0700 Subject: [PATCH 38/64] Tighten Cosmos3 extra-param validation and error handling - Reject non-integer and non-integral-float values in `condition_video_latent_indexes` instead of silently truncating - Add `_validate_output_type` validator for `output_type` extra param - Catch `TypeError` alongside `ValueError` in `validate_visual_gen_params` so wrong-shaped values surface as client errors (400) not server faults - Add unit tests covering all new validation paths Signed-off-by: Igor Shovkun --- .../visual_gen/models/cosmos3/defaults.py | 25 ++++++++++++++++-- tensorrt_llm/visual_gen/params.py | 4 ++- .../visual_gen/test_visual_gen_params.py | 26 +++++++++++++++++++ 3 files changed, 52 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index 9b61c3ea5b78..102ad4f27a09 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -53,7 +53,22 @@ def _normalize_condition_video_latent_indexes( ) -> tuple[int, ...]: if indexes is None: return COSMOS3_DEFAULT_CONDITION_VIDEO_LATENT_INDEXES - normalized = tuple(int(index) for index in indexes) + values = [] + for index in indexes: + # Strict: reject non-integers instead of silently truncating (1.9 -> 1) + # or TypeError-ing on None. Integral floats (JSON emitters) coerce. + if isinstance(index, bool) or not isinstance(index, (int, float)): + raise ValueError( + f"Cosmos3 condition_video_latent_indexes must be integers, got {index!r}." + ) + if isinstance(index, float): + if not index.is_integer(): + raise ValueError( + f"Cosmos3 condition_video_latent_indexes must be integers, got {index!r}." + ) + index = int(index) + values.append(index) + normalized = tuple(values) if not normalized: raise ValueError("Cosmos3 condition_video_latent_indexes must not be empty.") @@ -101,6 +116,11 @@ def _crop_video_frames(video, extra_params) -> torch.Tensor: return sliced.clone() +def _validate_output_type(output_type: str) -> None: + if output_type not in ("video", "image"): + raise ValueError(f"Cosmos3 output_type must be 'video' or 'image', got {output_type!r}.") + + def _validate_video_reference_tensor(video: torch.Tensor) -> None: if video.ndim != 4 or video.shape[-1] != 3: raise ValueError( @@ -169,9 +189,10 @@ def _validate_video_reference_tensor(video: torch.Tensor) -> None: description="Whether to enable audio generation.", ), "output_type": ExtraParamSchema( - type="Literal['video', 'image']", + type="str", default="video", description="Output modality: 'video' (T2V/I2V) or 'image' (text-to-image).", + validator=_validate_output_type, ), "condition_video_latent_indexes": ExtraParamSchema( type="list", diff --git a/tensorrt_llm/visual_gen/params.py b/tensorrt_llm/visual_gen/params.py index cdb85a107a78..bb1663239612 100644 --- a/tensorrt_llm/visual_gen/params.py +++ b/tensorrt_llm/visual_gen/params.py @@ -215,7 +215,9 @@ def validate_visual_gen_params( if validator is not None: try: validator(value) - except ValueError as exc: + except (TypeError, ValueError) as exc: + # TypeError included: a validator tripping on a wrong-shaped + # value is still a client error, not a server fault. messages.append(f"extra_params['{key}']: {exc}") continue # Range check (numeric only) diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py index 9a7ba0373b06..e155c017ce22 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py @@ -785,6 +785,15 @@ def _validate(extras): _validate({"condition_video_latent_indexes": []}) with pytest.raises(ValueError, match="first or last"): _validate({"condition_video_keep": "middle"}) + # No silent float truncation; None elements are a 400, not a TypeError. + with pytest.raises(ValueError, match="must be integers"): + _validate({"condition_video_latent_indexes": [1.9]}) + with pytest.raises(ValueError, match="must be integers"): + _validate({"condition_video_latent_indexes": [None]}) + _validate({"condition_video_latent_indexes": [0, 1.0]}) # integral floats OK + with pytest.raises(ValueError, match="output_type"): + _validate({"output_type": "gif"}) + _validate({"output_type": "image"}) with pytest.raises(ValueError, match=r"\[T, H, W, C\]"): _validate({"video": torch.zeros(4, 4, 3, dtype=torch.uint8)}) # 3-D with pytest.raises(ValueError, match="uint8"): @@ -792,6 +801,23 @@ def _validate(extras): with pytest.raises(ValueError, match="CPU tensor"): _validate({"video": torch.zeros(3, 4, 4, 3, dtype=torch.uint8, device="meta")}) + def test_validator_type_errors_become_client_errors(self): + """A validator raising TypeError (wrong-shaped value it didn't guard) + still folds into the 400 message list instead of escaping as a 500.""" + from tensorrt_llm._torch.visual_gen.pipeline import ExtraParamSchema + from tensorrt_llm.visual_gen.params import VisualGenParams, validate_visual_gen_params + + def touchy(value): + len(value) # TypeError on ints + + specs = {"knob": ExtraParamSchema(type="int", default=None, validator=touchy)} + with pytest.raises(ValueError, match="extra_params\\['knob'\\]"): + validate_visual_gen_params( + VisualGenParams(extra_params={"knob": 3}), + declared_defaults={}, + extra_param_specs=specs, + ) + def test_spec_validators_survive_pickling(self): """Specs travel worker -> coordinator in the READY handshake (pickled over ZMQ); validators/reducers must be module-level functions so they From 87a7b22ee5f2ad59f5f19cb5819e130bbc44a2d6 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Tue, 21 Jul 2026 10:33:15 -0700 Subject: [PATCH 39/64] Add Cosmos3-Nano V2V LPIPS integration test Signed-off-by: Igor Shovkun --- .../examples/visual_gen/test_visual_gen.py | 68 ++++++++++++++++++- .../test_lists/test-db/l0_b200.yml | 1 + 2 files changed, 67 insertions(+), 2 deletions(-) diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen.py b/tests/integration/defs/examples/visual_gen/test_visual_gen.py index 88d538570597..6d086d516009 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen.py @@ -30,6 +30,7 @@ import zipfile from typing import Any +import numpy as np import pytest import torch import torch._inductor.config as inductor_config @@ -126,6 +127,10 @@ COSMOS3_LPIPS_WIDTH = 1280 COSMOS3_LPIPS_T2V_NUM_FRAMES = 189 COSMOS3_LPIPS_T2I_NUM_FRAMES = 1 +# 9 frames = 3 latent frames: latents (0, 1) are pinned to the V2V reference, +# latent 2 (pixel frames 5-8) is generated. Frame 8 is the golden-compared frame. +COSMOS3_LPIPS_V2V_NUM_FRAMES = 9 +COSMOS3_LPIPS_V2V_FREE_FRAME_INDEX = 8 COSMOS3_LPIPS_NUM_INFERENCE_STEPS = 35 COSMOS3_LPIPS_GUIDANCE_SCALE = 6.0 COSMOS3_LPIPS_SEED = 42 @@ -978,12 +983,13 @@ def _generate_qwen_image_layered_lpips_image(model_path, input_path, output_path save_image(generated_image, output_path) -def _run_cosmos3_lpips_pipeline(num_frames): +def _run_cosmos3_lpips_pipeline(num_frames, video=None): """Run the Cosmos3-Nano pipeline (default setting, VANILLA attn, compile-off). Returns the generated video tensor ``(B, T, H, W, C)`` (T == ``num_frames``), or ``None`` if generation produced no video. ``num_frames=1`` yields the - single-frame text-to-image path. + single-frame text-to-image path; passing ``video`` (a uint8 [T, H, W, C] + conditioning-window tensor) yields the video-to-video path. """ # Cosmos3 re-reads the guardrail flag in __init__; set it before the pipeline loads. guardrails_env_key = "TRTLLM_DISABLE_COSMOS3_GUARDRAILS" @@ -1020,6 +1026,7 @@ def _run_cosmos3_lpips_pipeline(num_frames): guidance_scale=COSMOS3_LPIPS_GUIDANCE_SCALE, frame_rate=COSMOS3_LPIPS_FRAME_RATE, use_guardrails=False, + video=video, ) if result is None or result.video is None: return None @@ -1041,6 +1048,43 @@ def _generate_cosmos3_lpips_video(output_path): _save_lpips_video_mp4(video, output_path, frame_rate=COSMOS3_LPIPS_FRAME_RATE) +def _synthesize_cosmos3_v2v_lpips_reference(): + """Deterministic 5-frame 720p conditioning window. + + Five frames is the default window (``max(condition_video_latent_indexes) + * 4 + 1``). Built directly as a tensor — no codec round-trip — so it is + bit-identical on every machine and OpenCV build; a moving block gives the + conditioning a real structure signal. + """ + frames = [] + for i in range(5): + frame = np.full((COSMOS3_LPIPS_HEIGHT, COSMOS3_LPIPS_WIDTH, 3), 30, dtype=np.uint8) + x = 100 + i * 40 + frame[200:520, x : x + 200] = (200, 120, 40) + frames.append(frame) + return torch.from_numpy(np.stack(frames)) + + +def _generate_cosmos3_v2v_lpips_frame(output_path): + """Generate the Cosmos3-Nano video-to-video LPIPS sample (free frame only). + + The golden stores a single frame, not a video: frames 0-4 are pinned to the + reference's VAE round-trip (that contract is asserted semantically by the + V2V unit smokes), so the only content unique to this gate is what the model + generates *under* conditioning — the free latent, pixel frames 5-8. Frame 8 + is compared. 9 frames keeps the run to seconds while still exercising the + pinned/free latent boundary. + """ + from tensorrt_llm.media.encoding import save_image + + video = _run_cosmos3_lpips_pipeline( + COSMOS3_LPIPS_V2V_NUM_FRAMES, video=_synthesize_cosmos3_v2v_lpips_reference() + ) + assert video is not None, "Cosmos3-Nano V2V LPIPS run produced no video" + # video is (B, T, H, W, C); take the free frame -> (H, W, C) for save_image. + save_image(video[0, COSMOS3_LPIPS_V2V_FREE_FRAME_INDEX], output_path) + + def _generate_cosmos3_lpips_image(output_path): """Generate the Cosmos3-Nano text-to-image LPIPS sample (single frame).""" from tensorrt_llm.media.encoding import save_image @@ -1288,6 +1332,26 @@ def test_cosmos3_nano_t2v_lpips_against_golden(_visual_gen_deps, tmp_path): _assert_lpips_below_threshold(score, COSMOS3_LPIPS_THRESHOLD) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_cosmos3_nano_v2v_lpips_against_golden(_visual_gen_deps, tmp_path): + generated_path = tmp_path / "cosmos3_nano_v2v_generated_frame.png" + golden_path = _golden_media_path( + tmp_path, + "cosmos3_nano_v2v_lpips_golden_frame.png", + "Cosmos3-Nano V2V LPIPS golden frame", + ) + _generate_cosmos3_v2v_lpips_frame(generated_path) + score = _run_lpips_eval( + tmp_path, + "cosmos3_nano_v2v", + "image", + COSMOS3_LPIPS_PROMPT, + golden_path, + generated_path, + ) + _assert_lpips_below_threshold(score, COSMOS3_LPIPS_THRESHOLD) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") def test_cosmos3_nano_t2i_lpips_against_golden(tmp_path): generated_path = tmp_path / "cosmos3_nano_t2i_generated.png" diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 8f0e48d83458..bf158c5f530e 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -341,6 +341,7 @@ l0_b200: - examples/visual_gen/test_visual_gen.py::test_qwen_image_layered_lpips_against_golden TIMEOUT (10) - examples/visual_gen/test_visual_gen.py::test_cosmos3_nano_t2i_lpips_against_golden TIMEOUT (10) - examples/visual_gen/test_visual_gen.py::test_cosmos3_nano_t2v_lpips_against_golden TIMEOUT (15) + - examples/visual_gen/test_visual_gen.py::test_cosmos3_nano_v2v_lpips_against_golden TIMEOUT (10) - visual_gen/test_visual_gen_benchmark.py::test_offline_benchmark - visual_gen/test_visual_gen_benchmark.py::test_online_benchmark[openai-videos] # ---- moved to post-merge (MoE CI optimization) ---- From e1cc4016236cfdf0fb8f59d14f83307311175385 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Tue, 21 Jul 2026 11:16:26 -0700 Subject: [PATCH 40/64] Document and test supported input_reference formats - Make `is_decodable_image_file` strict (calls `image.load()`) so truncated files are rejected at classification, matching the bytes probe behavior - Add `TestMediaFileProbes` with a truncated-PNG test covering the file-path probe - Add AVI/MJPEG as a documented second container/codec pair: new `_avi_bytes` helper, `test_multipart_avi_reference_routes_to_video`, and `test_avi_mjpeg_bytes_decode` - Update `VideoGenerationRequest.input_reference` docstring and `examples/visual_gen/serve/README.md` to enumerate tested formats Signed-off-by: Igor Shovkun --- examples/visual_gen/serve/README.md | 1 + tensorrt_llm/inputs/media_io.py | 25 +++-- tensorrt_llm/serve/openai_protocol.py | 18 +++- .../visual_gen/test_visual_gen_utils.py | 96 +++++++++++++++++-- 4 files changed, 119 insertions(+), 21 deletions(-) diff --git a/examples/visual_gen/serve/README.md b/examples/visual_gen/serve/README.md index 4bc9840b7f10..38abcb633237 100644 --- a/examples/visual_gen/serve/README.md +++ b/examples/visual_gen/serve/README.md @@ -287,6 +287,7 @@ You can customize these by: - `num_frames`: when set, wins over the `seconds * frame_rate` derivation - `seed`, `num_inference_steps`, `guidance_scale`, `max_sequence_length`, `negative_prompt`: per-request denoise controls - `input_reference`: Reference image (I2V/TI2V) or video (V2V), classified by decoding the content — filename and content type are ignored; accepted as base64-encoded string in JSON or as a file in multipart form-data. Undecodable content returns HTTP 400. Video decode requires OpenCV on the server (not bundled — `pip install opencv-python-headless`). + - **Supported formats**: reference images are tested with PNG (`image/png`) and JPEG (`image/jpeg`). Reference videos are tested with MPEG-4 Part 2 video in MP4 (`video/mp4`) and Motion JPEG in AVI (`video/x-msvideo`). Filename and MIME metadata are not used for routing. Pillow must fully load an image; OpenCV must open a video and return decodable frames. Other formats depend on the installed decoder backend and are not guaranteed. Video references exceeding 1 GiB after RGB decoding are rejected. - `extra_params`: model-specific overflow (see below) - `response_format`: `"b64_json"` or `"url"` - `format`: Generation content encoding. Video encoders: `"mp4"`, `"avi"`, `"auto"`. Tensor formats: `"safetensors"`, `"pt"` (carries video + audio + scalar metadata in one payload for LTX-2). diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index 1db4fbc06321..739c6cd02c81 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -38,7 +38,7 @@ import torch from blake3 import blake3 from packaging.version import Version -from PIL import Image, UnidentifiedImageError +from PIL import Image from tensorrt_llm.inputs.multimodal_data import AudioData, VideoData from tensorrt_llm.logger import logger @@ -363,15 +363,20 @@ def _get_cv2(): def is_decodable_image_file(path) -> bool: - """True when ``path`` holds still-image content (anything PIL opens). + """True when ``path`` holds still-image content PIL can fully decode. - Header-only probe: identifies the container without decoding pixels. - Lets callers classify a reference by content instead of by file suffix. + Strict, like :func:`is_decodable_image_bytes`: ``Image.open`` is lazy, + so this also decodes the pixels (``load``) — a truncated file passes a + header-only probe and then fails at the actual load, far from the cause. + Non-image content is still rejected cheaply at the header parse. """ try: - with Image.open(path): + with Image.open(path) as image: + image.load() return True - except UnidentifiedImageError: + except OSError: + # ``UnidentifiedImageError`` (bad header) subclasses ``OSError``; + # truncated files raise plain ``OSError`` from ``load``. return False @@ -396,10 +401,10 @@ def is_decodable_video_file(path) -> bool: def is_decodable_image_bytes(data) -> bool: """True when ``data`` holds still-image content PIL can fully decode. - In-memory counterpart of :func:`is_decodable_image_file`, but strict: - ``Image.open`` is lazy, so this also decodes the pixels (``load``) — - a truncated file passes a header-only probe and would then 500 at the - worker's load instead of 400ing at the boundary. + In-memory counterpart of :func:`is_decodable_image_file`, equally + strict: ``Image.open`` is lazy, so this also decodes the pixels + (``load``) — a truncated file passes a header-only probe and would + then 500 at the worker's load instead of 400ing at the boundary. """ try: with Image.open(BytesIO(data)) as image: diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index bedfe219eafb..aa44551a592a 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -1688,11 +1688,19 @@ class VideoGenerationRequest(OpenAIBaseModel): default=None, description=( "Optional image or video reference that guides generation. " - "Content is classified by decoding, not by extension or " - "content-type: images (anything PIL reads) condition " - "image-to-video; videos (anything OpenCV reads) condition " - "video-to-video on models that support it. JSON requests " - "carry base64 bytes; multipart requests upload the file."), + "Content is classified by decoding it — filename and MIME " + "metadata are not used for routing (JSON requests carry bare " + "base64 with no such metadata): Pillow must fully load an " + "image; OpenCV must open a video and return decodable " + "frames. Images condition image-to-video; videos condition " + "video-to-video on models that support it. Tested formats: " + "PNG (image/png) and JPEG (image/jpeg) images; MPEG-4 Part 2 " + "video in MP4 (video/mp4) and Motion JPEG in AVI " + "(video/x-msvideo). Other containers, codecs, profiles, and " + "pixel formats depend on the installed decoder backend and " + "are not guaranteed. Video references exceeding 1 GiB after " + "RGB decoding are rejected. JSON requests carry base64 " + "bytes; multipart requests upload the file."), ) # Resolution diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index b67ebc14d006..7ea50014dab6 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -305,14 +305,57 @@ def _mp4_bytes(num_frames: int = 2) -> bytes: path = tmp.name try: writer = cv2.VideoWriter(path, cv2.VideoWriter_fourcc(*"mp4v"), 4.0, (16, 16)) - for _ in range(num_frames): - writer.write(np.zeros((16, 16, 3), dtype=np.uint8)) - writer.release() + try: + assert writer.isOpened(), "cv2 VideoWriter failed to open (mp4v in MP4)" + for _ in range(num_frames): + writer.write(np.zeros((16, 16, 3), dtype=np.uint8)) + finally: + writer.release() + with open(path, "rb") as f: + return f.read() + finally: + os.remove(path) + + @staticmethod + def _avi_bytes(num_frames: int = 2) -> bytes: + """Encode a 16x16 Motion-JPEG-in-AVI clip and return its bytes. + + The second container/codec pair in the documented support contract; + ``MJPG`` is built into the opencv wheel like ``mp4v``. + """ + cv2 = pytest.importorskip("cv2") + with tempfile.NamedTemporaryFile(suffix=".avi", delete=False) as tmp: + path = tmp.name + try: + writer = cv2.VideoWriter(path, cv2.VideoWriter_fourcc(*"MJPG"), 4.0, (16, 16)) + try: + assert writer.isOpened(), "cv2 VideoWriter failed to open (MJPG in AVI)" + for _ in range(num_frames): + writer.write(np.zeros((16, 16, 3), dtype=np.uint8)) + finally: + writer.release() with open(path, "rb") as f: return f.read() finally: os.remove(path) + def test_multipart_avi_reference_routes_to_video(self, tmp_path): + # The AVI/MJPEG contract pair must survive the real boundary, not + # just the decode primitive: classified as video by content, routed + # to the ``video`` extra param. + import torch + + generator = _StubVisualGen() + upload = UploadFile(file=BytesIO(self._avi_bytes()), filename="clip.avi") + request = VideoGenerationRequest(prompt="x", input_reference=upload) + params = parse_visual_gen_params( + request, "vid-avi", generator, media_storage_path=str(tmp_path) + ) + assert params.image is None + video = params.extra_params["video"] + assert isinstance(video, torch.Tensor) + assert tuple(video.shape) == (2, 16, 16, 3) + def test_multipart_video_reference_routes_to_extra_params_tensor(self, tmp_path): import torch @@ -438,15 +481,38 @@ def read(self, *args, **kwargs): assert list(tmp_path.iterdir()) == [] +class TestMediaFileProbes: + """File-path probes backing the offline producer path (``media_io``).""" + + def test_truncated_image_file_is_not_decodable(self, tmp_path): + # Same strictness as the bytes probe: a truncated PNG parses its + # header but must not classify as a decodable image — and the video + # probe must not rescue it as a one-frame video either. + pytest.importorskip("cv2") + from tensorrt_llm.inputs.media_io import is_decodable_image_file, is_decodable_video_file + + whole = tmp_path / "whole.png" + Image.fromarray(np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8)).save(whole) + truncated = tmp_path / "truncated.png" + truncated.write_bytes(whole.read_bytes()[: whole.stat().st_size // 2]) + + assert is_decodable_image_file(whole) + assert not is_decodable_image_file(truncated) + assert not is_decodable_video_file(truncated) + + class TestMediaBytesProbes: """The in-memory probe/decode primitives the serve boundary runs on.""" def test_is_decodable_image_bytes(self): from tensorrt_llm.inputs.media_io import is_decodable_image_bytes - buf = BytesIO() - Image.new("RGB", (4, 4), (1, 2, 3)).save(buf, format="PNG") - assert is_decodable_image_bytes(buf.getvalue()) + # PNG and JPEG are the two image formats in the documented support + # contract; both must probe as decodable. + for fmt in ("PNG", "JPEG"): + buf = BytesIO() + Image.new("RGB", (4, 4), (1, 2, 3)).save(buf, format=fmt) + assert is_decodable_image_bytes(buf.getvalue()), fmt assert not is_decodable_image_bytes(b"definitely not an image") # Video bytes are not an image (mp4 has no PIL-openable header). assert not is_decodable_image_bytes(TestInputReferenceMaterialization._mp4_bytes()) @@ -475,6 +541,24 @@ def test_decode_video_frames_from_bytes_rejects_garbage(self): with pytest.raises(ValueError): decode_video_frames_from_bytes(b"not a video at all") + def test_avi_mjpeg_bytes_decode(self): + # Motion-JPEG-in-AVI — the second container/codec pair in the + # documented support contract: not an image, decodes as video. + pytest.importorskip("cv2") + import torch + + from tensorrt_llm.inputs.media_io import ( + decode_video_tensor_from_bytes, + is_decodable_image_bytes, + ) + + payload = TestInputReferenceMaterialization._avi_bytes() + assert not is_decodable_image_bytes(payload) + video = decode_video_tensor_from_bytes(payload) + assert isinstance(video, torch.Tensor) + assert video.dtype == torch.uint8 + assert tuple(video.shape) == (2, 16, 16, 3) + def test_truncated_image_bytes_are_not_decodable(self): # A truncated PNG still opens (the header parses) but cannot decode # its pixels; the probe must reject it so the boundary 400s instead From b9168e5515f065bb328769aca5e982ac853fa35c Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 23 Jul 2026 21:35:59 -0700 Subject: [PATCH 41/64] Add worker-side NVDEC reference decoding via PyNvVideoCodec Each worker rank demuxes encoded video bytes from memory and decodes the V2V conditioning window on NVDEC: forward-only feed, preallocated ring at the request's target resolution (keep=first stops early, keep=last rings to EOS), emitted-frame work limit (TRTLLM_MAX_REFERENCE_DECODE_FRAMES), local-tap Lanczos resize with bounded PIL parity, client/capacity error classes, and an all-rank CPU/gloo status protocol so a rank-local decode failure cannot hang healthy ranks in model collectives. PyNvVideoCodec becomes a declared, pinned dependency (compliance approval recorded, covering the bundled FFmpeg demux libraries); the import stays function-local so importing tensorrt_llm never loads the driver-linked module. Checked-in H.264 fixtures (MP4 + AVI, forced B-frames, per-frame index patterns, provenance in README) drive the direct tests. Signed-off-by: Igor Shovkun --- requirements.txt | 1 + .../_torch/visual_gen/media_decode.py | 327 ++++++++++++++++++ .../_torch/visual_gen/test_data/README.md | 59 ++++ .../test_data/cosmos3_v2v_ref_9f_bframes.avi | Bin 0 -> 7842 bytes .../test_data/cosmos3_v2v_ref_9f_bframes.mp4 | Bin 0 -> 2786 bytes .../_torch/visual_gen/test_media_decode.py | 272 +++++++++++++++ 6 files changed, 659 insertions(+) create mode 100644 tensorrt_llm/_torch/visual_gen/media_decode.py create mode 100644 tests/unittest/_torch/visual_gen/test_data/README.md create mode 100644 tests/unittest/_torch/visual_gen/test_data/cosmos3_v2v_ref_9f_bframes.avi create mode 100644 tests/unittest/_torch/visual_gen/test_data/cosmos3_v2v_ref_9f_bframes.mp4 create mode 100644 tests/unittest/_torch/visual_gen/test_media_decode.py diff --git a/requirements.txt b/requirements.txt index d735182a8fe8..1f480852c623 100644 --- a/requirements.txt +++ b/requirements.txt @@ -37,6 +37,7 @@ pydantic-settings[yaml] msgspec omegaconf pillow +PyNvVideoCodec~=2.1.0 optimum # evaluate needs datasets>=2.0.0 which triggers datasets>3.1.0 which is not stable: https://github.com/huggingface/datasets/issues/7467 datasets==3.1.0 diff --git a/tensorrt_llm/_torch/visual_gen/media_decode.py b/tensorrt_llm/_torch/visual_gen/media_decode.py new file mode 100644 index 000000000000..a9f9e810424b --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/media_decode.py @@ -0,0 +1,327 @@ +# 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. +"""Worker-side video-reference decoding on NVDEC (PyNvVideoCodec). + +Encoded reference bytes arrive from the coordinator; each worker rank demuxes +them from memory, decodes on NVDEC, and retains only the conditioning window, +resized to the request's output resolution — so retained memory is bounded by +the request's own output shape and at any instant only one source-resolution +frame is alive. + +PyNvVideoCodec is imported function-locally: ``import tensorrt_llm`` on a +CPU-only host must never load the driver-linked extension. +""" + +import functools +import math +import os + +import torch + + +class MediaDecodeError(ValueError): + """Client-class failure: the reference itself is unusable. + + Corrupt/undecodable content, unsupported codec, zero frames, or the + decode-work limit. Maps to HTTP 400 / a plain ``ValueError`` at the + public Python boundary. + """ + + +class VisualGenCapacityError(RuntimeError): + """Capacity-class failure: a valid request does not fit this deployment. + + CUDA/NVDEC allocation or session-init failure. Maps to HTTP 503; never + a client error — the input is not malformed. + """ + + +def classify_worker_error(exc: BaseException) -> str | None: + """Failure class for the response channel: "client", "capacity", or None. + + Only the two dedicated classes are mapped — a bare ``ValueError`` from a + model bug must stay an unclassified runtime failure, not become a 400. + """ + if isinstance(exc, MediaDecodeError): + return "client" + if isinstance(exc, (VisualGenCapacityError, torch.cuda.OutOfMemoryError)): + return "capacity" + return None + + +def synchronize_media_prepare_status(exc: Exception | None) -> None: + """All-rank convergence point between media prepare and model collectives. + + Every rank decodes/prepares its media independently; a rank that failed + while others proceed into the transformer's collectives would hang the + job. All ranks call this with their local outcome; if any failed, the + lowest failing rank's error class + message is broadcast, the failing + rank(s) re-raise their own exception, and every healthy rank raises a + reconstructed equivalent in lockstep. Runs on CPU tensors so + the hybrid (``cpu:gloo``) process group carries it even when the failure + was CUDA/NVDEC initialization. Converges *caught* failures only — a fatal + process or context death is beyond its reach. + """ + import torch.distributed as dist + + if not (dist.is_available() and dist.is_initialized()) or dist.get_world_size() == 1: + if exc is not None: + raise exc + return + + healthy_sentinel = 2**31 - 1 + rank = dist.get_rank() + flag = torch.tensor([rank if exc is not None else healthy_sentinel], dtype=torch.int64) + dist.all_reduce(flag, op=dist.ReduceOp.MIN) + failing_rank = int(flag.item()) + if failing_rank == healthy_sentinel: + return + + payload = [None] + if rank == failing_rank: + payload = [(classify_worker_error(exc), str(exc))] + dist.broadcast_object_list(payload, src=failing_rank) + + if exc is not None: + raise exc + kind, message = payload[0] + message = f"[rank {failing_rank}] {message}" + if kind == "client": + raise MediaDecodeError(message) + if kind == "capacity": + raise VisualGenCapacityError(message) + raise RuntimeError(message) + + +# Decode-work default: 5 min @ 24 fps, far above the canonical ~8 s reference. +# Deliberately its own constant — the serve's output cap (``MAX_VIDEO_FRAMES``) +# is a different policy that happens to share the value today. +DEFAULT_MAX_REFERENCE_DECODE_FRAMES = 7200 + + +def max_reference_decode_frames() -> int | None: + """Decode-work limit for a video reference, or ``None`` when disabled. + + Bounds serial worker occupancy for forward-only ``keep="last"`` decoding + (encoded size cannot: hours of low-bitrate video are small on disk). The + default sits far above the canonical ~8 s reference; trusted deployments + may raise it or disable it with ``TRTLLM_MAX_REFERENCE_DECODE_FRAMES=0``. + """ + raw = os.environ.get("TRTLLM_MAX_REFERENCE_DECODE_FRAMES") + if raw is None: + return DEFAULT_MAX_REFERENCE_DECODE_FRAMES + limit = int(raw) + return None if limit <= 0 else limit + + +@functools.lru_cache(maxsize=32) +def _lanczos_taps( + in_size: int, out_size: int, device_str: str, a: int = 3 +) -> tuple[torch.Tensor, torch.Tensor]: + """Local-support Lanczos-a taps: ``([out, K] weights, [out, K] indices)``. + + PIL semantics (the reference vllm-omni preprocess resizes with + ``PIL.Image.Resampling.LANCZOS``): the kernel is stretched by the + downscale ratio, taps span ``a * scale`` source pixels around each output + center, out-of-range taps get zero weight, and each row normalizes. + float32, like PIL's internal filter precision (PIL then quantizes + coefficients — parity is bounded, not bit-exact). + + ``K = ceil(2 * a * max(in/out, 1)) + 1`` is the filter's true support, so + applying these by gather + weighted sum costs ``O(out * K)`` per row + instead of the ``O(out * in)`` of a dense resampling matrix. Cached per + (in, out, device): every frame of a clip shares sizes, and the cached + tensors are kilobytes. + """ + device = torch.device(device_str) + ratio = in_size / out_size + scale = max(ratio, 1.0) + support = a * scale + centers = (torch.arange(out_size, device=device, dtype=torch.float32) + 0.5) * ratio + first = torch.floor(centers - support) + num_taps = int(math.ceil(2 * support)) + 1 + taps = first.unsqueeze(1) + torch.arange(num_taps, device=device, dtype=torch.float32) + x = (taps + 0.5 - centers.unsqueeze(1)) / scale + weights = torch.sinc(x) * torch.sinc(x / a) + weights = torch.where(x.abs() < a, weights, torch.zeros_like(weights)) + valid = (taps >= 0) & (taps < in_size) + weights = weights * valid + weights = weights / weights.sum(dim=1, keepdim=True) + return weights, taps.clamp(0, in_size - 1).long() + + +def _resample_last_dim(x: torch.Tensor, weights: torch.Tensor, taps: torch.Tensor) -> torch.Tensor: + """Resample the last dim ``[..., in] -> [..., out]`` by gather + weighted sum.""" + return (x[..., taps] * weights).sum(-1) + + +def resize_center_crop_uint8(frames: torch.Tensor, target_h: int, target_w: int) -> torch.Tensor: + """Resize + center-crop uint8 ``[T, H, W, C]`` frames to the target size. + + Applied to the worker-decoded reference frames before retention. + Semantics mirror the reference implementation's PIL path (cover-scale by + ``max(target/source)``, ceil-rounded resize with Lanczos-3, center crop), + implemented as separable local-tap resampling: per output pixel only the + filter's ``K`` support taps are gathered and summed. + """ + t, h, w, c = frames.shape + if (h, w) == (target_h, target_w): + return frames + ratio = max(target_w / w, target_h / h) + resize_w = int(math.ceil(ratio * w)) + resize_h = int(math.ceil(ratio * h)) + + # PIL resamples in two passes (horizontal, then vertical) and stores the + # intermediate as uint8 — clamping away Lanczos overshoot between passes. + # Same order and intermediate quantization here, so parity with the PIL + # reference path stays within coefficient-rounding noise. + x = frames.permute(0, 3, 1, 2).to(torch.float32) # [T, C, H, W] + if resize_w != w: + weights, taps = _lanczos_taps(w, resize_w, str(frames.device)) + x = _resample_last_dim(x, weights, taps) + x = x.round_().clamp_(0, 255) + if resize_h != h: + weights, taps = _lanczos_taps(h, resize_h, str(frames.device)) + x = _resample_last_dim(x.transpose(-1, -2), weights, taps).transpose(-1, -2) + + left = max((resize_w - target_w) // 2, 0) + top = max((resize_h - target_h) // 2, 0) + x = x[:, :, top : top + target_h, left : left + target_w] + return x.round_().clamp_(0, 255).to(torch.uint8).permute(0, 2, 3, 1).contiguous() + + +def decode_video_reference_window( + data: bytes, + *, + window: int, + keep: str, + target_h: int, + target_w: int, + device: torch.device, +) -> torch.Tensor: + """Decode encoded reference bytes into the conditioning window on device. + + Returns a uint8 ``[T, target_h, target_w, 3]`` tensor, ``T <= window``: + the first ``window`` frames (``keep="first"``, decode stops early) or the + last ``window`` (``keep="last"``, sequential decode to EOS through a + preallocated ring — the memory-buffer demuxer is a forward-only feeder, + seeking is not assumed). Shorter-than-window clips return what exists; + the pipeline right-pads. Frames are resized to the target resolution + before retention, so a high-resolution source never dominates memory. + """ + try: + import PyNvVideoCodec as nvc + except ImportError as exc: + raise ImportError( + "PyNvVideoCodec is required for video-reference decoding; " + "install the declared dependency (pip install PyNvVideoCodec)." + ) from exc + + frame_cap = max_reference_decode_frames() + if frame_cap is not None and keep == "first" and window > frame_cap: + raise MediaDecodeError( + f"Conditioning window of {window} frames exceeds the reference " + f"decode limit of {frame_cap} (TRTLLM_MAX_REFERENCE_DECODE_FRAMES)." + ) + + position = 0 + + def _read(buf: bytearray) -> int: + nonlocal position + chunk = data[position : position + len(buf)] + buf[: len(chunk)] = chunk + position += len(chunk) + return len(chunk) + + demuxer = None + decoder = None + try: + try: + # CPU-side FFmpeg demux: failure here means the bytes are not a + # readable stream — a content problem, not a capacity one. + demuxer = nvc.CreateDemuxer(_read) + except nvc.PyNvVCException as exc: + raise MediaDecodeError( + f"Video reference could not be demuxed (corrupt or not a " + f"supported container): {exc}" + ) from exc + try: + decoder = nvc.CreateDecoder( + gpuid=device.index or 0, + codec=demuxer.GetNvCodecId(), + usedevicememory=True, + outputColorType=nvc.OutputColorType.RGB, + ) + except nvc.PyNvVCException as exc: + # Init failure on a demuxable stream is genuinely ambiguous — an + # unsupported codec/profile (client-fixable by re-encoding) and a + # driver/session failure (deployment fault) raise the same + # exception type with no inspectable code. Make neither + # categorical claim: stay unclassified (500), with a message + # naming both possibilities. + raise RuntimeError( + f"NVDEC decoder initialization failed for this stream — the " + f"codec/profile may be unsupported on this GPU, or the " + f"decoder session could not be created: {exc}" + ) from exc + + ring = torch.empty(window, target_h, target_w, 3, dtype=torch.uint8, device=device) + count = 0 + try: + for packet in demuxer: + for frame in decoder.Decode(packet): + if frame_cap is not None and count >= frame_cap: + raise MediaDecodeError( + f"Video reference exceeds the decode limit of " + f"{frame_cap} frames " + f"(TRTLLM_MAX_REFERENCE_DECODE_FRAMES); trim the " + f"clip{' or use condition_video_keep=first' if keep == 'last' else ''}." + ) + decoded = torch.from_dlpack(frame) + # Ownership copy off the NVDEC surface (recycled by the + # decoder) and resize-before-retain in one step. + ring[count % window].copy_( + resize_center_crop_uint8(decoded.unsqueeze(0), target_h, target_w)[0] + ) + count += 1 + if keep == "first" and count >= window: + break + if keep == "first" and count >= window: + break + except torch.cuda.OutOfMemoryError as exc: + raise VisualGenCapacityError( + f"Out of device memory while decoding the video reference " + f"({window} frames @ {target_w}x{target_h} retained): {exc}" + ) from exc + except nvc.PyNvVCException as exc: + raise MediaDecodeError( + f"Video reference failed to decode (corrupt or unsupported " + f"stream for this deployment's decoder): {exc}" + ) from exc + + if count == 0: + raise MediaDecodeError( + "Video reference contains no decodable frames; the payload " + "may be corrupt or use an unsupported codec." + ) + if count <= window: + return ring[:count] + start = count % window + if start == 0: + return ring + return torch.cat([ring[start:], ring[:start]]) + finally: + del decoder + del demuxer diff --git a/tests/unittest/_torch/visual_gen/test_data/README.md b/tests/unittest/_torch/visual_gen/test_data/README.md new file mode 100644 index 000000000000..369e5c0c5189 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_data/README.md @@ -0,0 +1,59 @@ + + +# VisualGen test media fixtures + +## `cosmos3_v2v_ref_9f_bframes.mp4` + +9-frame 64×64 H.264-in-MP4 V2V reference fixture (2,786 bytes), encoded once +offline with ffmpeg/libx264 so decode tests exercise **cross-encoder** interop +(x264 encodes, NVDEC decodes; the H.264 spec makes decoded YUV +bit-exact for conformant decoders; the YUV->RGB conversion is this stack's +(PyNvVideoCodec), so decoded RGB is stable for this decode path). B-frames are forced — with only 9 frames x264 +would otherwise skip them — and verified present (`I B B P I B B P I`). + +Each frame encodes its own display index three ways, so tests recover ordering +from content alone (catching B-frame reorder bugs): +- red channel: solid ramp, `R = 20 + 25 * i` +- green channel: horizontal bar at rows `[7*i, 7*i + 7)` +- blue channel: vertical bar at columns `[7*i, 7*i + 7)` + +Regeneration (exact provenance; ffmpeg 6.1.1 / libx264): + +```python +import numpy as np +from PIL import Image + +for i in range(9): + frame = np.zeros((64, 64, 3), dtype=np.uint8) + frame[:, :, 0] = 20 + i * 25 + frame[7 * i : 7 * i + 7, :, 1] = 255 + frame[:, 7 * i : 7 * i + 7, 2] = 255 + Image.fromarray(frame).save(f"frame_{i:02d}.png") +``` + +```bash +ffmpeg -y -framerate 24 -i frame_%02d.png \ + -c:v libx264 -pix_fmt yuv420p -g 4 -bf 2 \ + -x264-params b_adapt=0:scenecut=0 \ + -movflags +faststart cosmos3_v2v_ref_9f_bframes.mp4 + +# verify B-frames survived: +ffprobe -v error -select_streams v:0 -show_entries frame=pict_type \ + -of csv=p=0 cosmos3_v2v_ref_9f_bframes.mp4 +``` + +## `cosmos3_v2v_ref_9f_bframes.avi` + +The **same 9 frames** as the MP4 fixture, re-muxed as H.264-in-AVI (7,842 +bytes) so the second supported container is exercised through the real decode +path. Same `frame_%02d.png` source as above; only the container differs: + +```bash +ffmpeg -y -framerate 24 -i frame_%02d.png \ + -c:v libx264 -pix_fmt yuv420p -g 4 -bf 2 \ + -x264-params b_adapt=0:scenecut=0 \ + cosmos3_v2v_ref_9f_bframes.avi +``` diff --git a/tests/unittest/_torch/visual_gen/test_data/cosmos3_v2v_ref_9f_bframes.avi b/tests/unittest/_torch/visual_gen/test_data/cosmos3_v2v_ref_9f_bframes.avi new file mode 100644 index 0000000000000000000000000000000000000000..33cc8c4ff4178439d98a008dfaf80e6a2c17ebd7 GIT binary patch literal 7842 zcmeHMdsq`!79SP_d?C72v*O|`S`n2nlR$ViEfM4;4~vMW3Vq>nI zkznwI_uB>f+Tn~K9M5XNZT6K#E4c`f#OK$#(2rn50DVuhMGzAx7M)<=9oVa80|C5) zM{XNnW%h*u!CNJ2N$f%a=-X~Yj$umjEfp(kMZk)H6#**(Rs^gFSP`%yU`4=+fE58N z0#*d92>i|wfcXJCka%E}S`E0=1Hbd}{;%P|Jk3;~5E&a57r8hr!3ogOM3zc|32})K z$1ntm<$xesCX0rdRG6U!#U@pWjH$RDIx7!BY+w2z7RPCDgBYJ(t<6a|30i#|h$*>fP1CkZvm_I7U zDLgD@u&#H#9ksyBnCn?&+7xdy$uF~ z5agFBMnmX?44di=dl9O6tp?~YdY)!l2Nhj2Rky~Uq%4uj4#9|R(sE#%FLDE`+hjGw4 zo&vdjpakR!O(siZ%SlG7CwRb<03x0xXdQ3?6j=i5s8~Wv!TuC-eKw$JrI{v_gdW-` ze z=AI#(EXAn-&yw>04i%WOq?BY;s8#_Wg(CutR3sK+Xa*dI6cb7SsfS+8Pr@jb1ONfY zQ+i-c>%shhwE(8i2h4Z|n0h#oNG@2b)W$pnH>FN$`EbnWFN;noD?%1da=-cRYJul; zrwd z!T7dY;@dAXu0K00QFX68v6VLs9d5W1bI8B=r{}9Ke9YRd4olFV8b&U0WV#11J60+#zK4usZ9Q=X*>=GZTpEAW`fFReEfw>IV%EP3)* zRo>c*w;hE}H(thPP!Uck|EQ?KLMlAnQi8gtq&42iom{kfS*$2HcHfs@`ONwY0%H@6 z#AdT!kv6|%XgldQaCiNc&jf-d87UUNkx=`fTEh=)AniuG=fSY~^vr z%J!wv2XnOA^Q$v5iS*WzsC~tqF?Ndv?Q-LNulX#=3{ANjSv2(Z!?SJj z4bCmZr(H8wWDQN#<@N^TByKuZxuxdWk7a+mbGN3UU3|SDt!GYJZ8`SMjfSJVL4Rf7ve)71S_MZ$OAN9U%aRjY0plDo$FPW`i2 z#$l=`c+A(iC zb^Zmn<Lr6-DqhBl1|1d#tcfsM7@gFTLb1*RP8$PO@{hc$?#*?>$`(7s7 z7k=`3#qQFIFt)9=u+dN5R=K&UaXI6|?43f(I<%LXi(BijO@9zIVVxvEB=N3rz47|A zWP^Ff#{Fv5g`0+lu3j!mTsC!L{I)_H`^ubm?l)!yJ;nV>PpE2Ln>~-zmmfX9Y>cyf zUCANe1*NHnH%cYHT~zGm+S=Twc-b#F#mKH~Scye;7gItK^WkKrr;GX???*+iB}47u6uT zU(of%`;XhE57SB`{uziC?j-Q@)cp4pzUv_m2QeGa@LaR-O@N2zh=o^zZveP=E%y5X k4|k-cy?hjcz?Es?BkU0bt^*6-33xaQ7JCu+=78w^2VgdZ)&Kwi literal 0 HcmV?d00001 diff --git a/tests/unittest/_torch/visual_gen/test_data/cosmos3_v2v_ref_9f_bframes.mp4 b/tests/unittest/_torch/visual_gen/test_data/cosmos3_v2v_ref_9f_bframes.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..4bae887a195857e11b2ce1b5c60e52e3d14ccb20 GIT binary patch literal 2786 zcmaJB3se(Vb`TW&;b&2^LUoEC7L|}mApDvZ1PdyN^@Cv1Fquq9Ad?9*0}0q7D%LOS zQhx|4t%}x)qFogMsVLM%Sky%+m3r_)<)~Fr@Iyt&-WRY--93BHx$oY0|L@)V9)@AK zMo2a4c$UGi0T`lyQZj`i7^4?~VYXQe%cfu$re{*LYM}kV_CAK$m@pey`$_+}0mUDb z7#7&y_P_Z7P)D)Ak%^#+6%wsEZTp$^MT5NW!aI5XXWl??LWp6LX|=!ub)ogVV3k-a zMRD!=Y98RpsC6Xha7O*^*(D$zjV-RSdU|Npdaf@9mZDSBZ+S$cV9$n;2DP3>Ha1O+ z&Y%ItDTT2H?O*m}b+67v&C$pP0n!gk=Ja?kZ7m6IMQfD-TNIp zbiMO8<*>k_fbXeEt1~JRgC1$1M}!=DgeJ~G!_hEY_8tsaL8XTW4YXJaZq9m8-g+af zc*l&N@Een!w4p}g9yrBvG)@4^$q0!rE+c5dPvZxQFfkZ-M$DfT>fwb4MFoSUnx;Sz z%o2~HYFeJW2YBxaev#7qz}oS_w3+KltbD&Tx@BcBSc zuqrvVf)Go<2^O5unQ65WMI^vO$&rS5T0zKhO3SefsRUO77dTq4*YUvlntj!j01TDH zC?rqyM=~_Zr&1bIf!C?6JgFv)$Va7A=|~ro@I3>e-pdVMBFSmH<5 zelle6*SV+Fg~4-2yWX0zQsh3-adKnCgW0E^M>$5%nC*9DaKr7^KXqq~8*_F+cO`Z9%z{^fS+$UDtlR&GM|fM%58~f?D(Wkz$)+(U{nLJZsvs z_PXIm7sXC&yyiUNU)YKB7v^L?bsD;3MZ;;vYh7#4G=(RzETk=c%lX>RYKlMl?t;ge zGjGV`AUA30@b--DvpQ#vJ*fMkF|>5`wxI)!MPCI4m%Wg+Uw$D(zNv89bG{~|r070o zE9`XZ*s#g<$n}K|5zVns*Os8YLFTpRdgPjxw3AtarTcK@)rdpBdB40^aq&~mZso^| zjHi3BOHLzS1Z~)NEzPg1GGx=F=bH-VUr|)A${#YXTb1p%%v-fTq)ZX_JV7*V(BkPL zx5>hi3QJyH)u!ylPruM)uD*1~LF{<*)q*5C%rWjCg~gKKLiY)A0WISas%~bC&0QHg zPZ~IH&)460O?@9q65H%mCCtq4-b(p&-dlS$rb$)uDSx-?rIVv~UrX0TM%h*8 zcKAjdtTRk40gpy{{c3EOrRSEk)9FX2Chn|Pv|G*54CkSvYG(xV=7tx=kA6<;T2&BJ zE}^46wq;KG=2+*B{jb`JM(r0>kG8y0Sx$FmUw5)S68UCNc~H&#!2#jtqF#n)W{!2< zTGVW-j5HTDED1l5&M+5NCM7T6x0IS)(xg$kto%89@|q&-=6Bg;EbzYWwKzE>?%M3! z?r$HRt5>cYSxbK2Jb8Ia_jp4_n_qg=hT}z>OP~K#@YlQdN-G;=H?k92rzMmflpMQR zc~meN4=nHUo+I%~e?0T4wBgl*mQTvj7rniAZZK@t-DL+ zg>IX|No78lUQYj~DKFaHg+K8+IBpdb_5~W8EX%1|FUvZ=`+J73r~Qcl%>Z%ZCVA*X z&4UqJ@8(!W_}iBa|JD>77vbPjhjlz!B)T{2aP9D=a|`TE?1!eMC4c#TBvvm_cLLjA z#SF^%?9K9>`wBni>g#f_39v0m)%ak`*J=5kG5UB&f@59f?nBlnZaJqw+0 zzBwabho074G%ul+p6+W`a-(9$k6N%L$7WDb`j`h*DFM$2pM57aWzIG3M=A~;y%0NO zq;gIEA@5oH;=fF#@9#3_xvK7Zum5t@?vKuEsuXL`!P3#IceC_d*LBR;+O#ohXLCUQ z0oR!ASG-00c8^Qwv|XE|D|`HSL->=(lMPSjOb94z8T+WQ@^#F^6MS)X%+<5Ql0|m1 z?538EK_xzgyDl@G&8-)Ggs{`~?p*EN%|EutQ}PdaibUP7tJhYpx|6VmcG7g=dqvG( LefXq)V$Z(;V^@_M literal 0 HcmV?d00001 diff --git a/tests/unittest/_torch/visual_gen/test_media_decode.py b/tests/unittest/_torch/visual_gen/test_media_decode.py new file mode 100644 index 000000000000..e5cda550d66c --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_media_decode.py @@ -0,0 +1,272 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Direct tests for :mod:`tensorrt_llm._torch.visual_gen.media_decode`. + +The decode tests run on the checked-in H.264 fixtures (see +``test_data/README.md`` for provenance). Each fixture frame encodes its own +display index three ways — red-channel ramp, green horizontal bar, blue +vertical bar — so ordering, channel layout, and surface ownership are all +observable from content alone. + +GPU (NVDEC) tests skip only on CUDA-less platforms; a missing PyNvVideoCodec +on a supported platform fails them (it is a declared dependency). +""" + +import math +import subprocess +import sys +import time +from pathlib import Path + +import numpy as np +import pytest +import torch +from PIL import Image + +from tensorrt_llm._torch.visual_gen.media_decode import ( + MediaDecodeError, + _lanczos_taps, + decode_video_reference_window, + max_reference_decode_frames, + resize_center_crop_uint8, + synchronize_media_prepare_status, +) + +_TEST_DATA = Path(__file__).parent / "test_data" +_MP4 = _TEST_DATA / "cosmos3_v2v_ref_9f_bframes.mp4" +_AVI = _TEST_DATA / "cosmos3_v2v_ref_9f_bframes.avi" + + +def _frame_indices(frames: torch.Tensor) -> list[int]: + """Recover each frame's display index from the red-channel ramp.""" + return [round((f[:, :, 0].float().mean().item() - 20) / 25) for f in frames] + + +class TestResizeCenterCrop: + """CPU-runnable checks of the shared Lanczos resize/crop front.""" + + def _pil_reference(self, frames_u8: torch.Tensor, target_h: int, target_w: int): + """The reference implementation's geometry: PIL LANCZOS cover-scale + + center crop (vllm-omni ``_preprocess_condition_image``).""" + out = [] + for frame in frames_u8.numpy(): + h, w = frame.shape[:2] + scale = max(target_w / w, target_h / h) + resize_w = int(math.ceil(scale * w)) + resize_h = int(math.ceil(scale * h)) + img = Image.fromarray(frame).resize((resize_w, resize_h), Image.Resampling.LANCZOS) + left = (resize_w - target_w) // 2 + top = (resize_h - target_h) // 2 + out.append(np.asarray(img.crop((left, top, left + target_w, top + target_h)))) + return torch.from_numpy(np.stack(out)) + + def test_parity_with_pil_lanczos(self): + rng = np.random.default_rng(7) + frames = torch.from_numpy(rng.integers(0, 256, (3, 64, 96, 3), dtype=np.uint8)) + for target_h, target_w in ((32, 48), (48, 32), (128, 96)): + ours = resize_center_crop_uint8(frames, target_h, target_w) + ref = self._pil_reference(frames, target_h, target_w) + assert ours.shape == ref.shape == (3, target_h, target_w, 3) + diff = (ours.int() - ref.int()).abs() + # PIL quantizes filter coefficients to fixed point; we keep + # float. Bounded, not bit-exact. + assert diff.max().item() <= 2, f"{target_h}x{target_w}: max diff {diff.max()}" + + def test_native_resolution_is_identity(self): + frames = torch.zeros(2, 32, 32, 3, dtype=torch.uint8) + assert resize_center_crop_uint8(frames, 32, 32) is frames + + def test_taps_are_local_support_and_cached(self): + # The taps depend only on (in, out) sizes — one build per clip, not + # one per frame — and their width is the filter's true support + # (K = ceil(2 * a * scale) + 1), NOT the full input row: this is what + # makes the resample O(out * K) instead of O(out * in). + _lanczos_taps.cache_clear() + weights, taps = _lanczos_taps(1920, 1280, "cpu") + assert taps.shape == weights.shape + assert taps.shape[1] <= 2 * math.ceil(3 * (1920 / 1280)) + 1 # K = 10 << 1920 + frames = torch.zeros(4, 40, 40, 3, dtype=torch.uint8) + resize_center_crop_uint8(frames, 20, 20) + info = _lanczos_taps.cache_info() + resize_center_crop_uint8(frames, 20, 20) + assert _lanczos_taps.cache_info().hits >= info.hits + 2 + + +class TestImportIsolation: + def test_import_tensorrt_llm_does_not_load_pynvvideocodec(self): + """PyNvVideoCodec is driver-linked; ``import tensorrt_llm`` (and the + decode module itself) must not load it — only an actual decode may.""" + code = ( + "import sys; import tensorrt_llm; " + "import tensorrt_llm._torch.visual_gen.media_decode; " + "assert 'PyNvVideoCodec' not in sys.modules, " + "'driver-linked PyNvVideoCodec loaded at import time'" + ) + subprocess.run([sys.executable, "-c", code], check=True, timeout=600) + + +def _status_protocol_rank(rank: int, world_size: int, init_file: str, results_dir: str): + """Spawn target: run the convergence protocol with rank 1 failing.""" + import os + + # Containers often have hostnames that don't resolve to a usable + # interface; pin gloo to loopback or its rendezvous hangs. + os.environ.setdefault("GLOO_SOCKET_IFNAME", "lo") + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + + import torch.distributed as dist + + from tensorrt_llm._torch.visual_gen.media_decode import ( + MediaDecodeError, + synchronize_media_prepare_status, + ) + + dist.init_process_group( + "gloo", init_method=f"file://{init_file}", rank=rank, world_size=world_size + ) + try: + local_error = MediaDecodeError("rank-local decode failure") if rank == 1 else None + try: + synchronize_media_prepare_status(local_error) + outcome = "no-error" + except MediaDecodeError as exc: + outcome = f"client:{exc}" + except Exception as exc: # pragma: no cover - diagnostic path + outcome = f"unexpected:{type(exc).__name__}:{exc}" + finally: + dist.destroy_process_group() + Path(results_dir, f"rank{rank}.txt").write_text(outcome) + + +class TestPrepareStatusProtocol: + def test_local_failure_is_reraised_without_group(self): + err = MediaDecodeError("boom") + with pytest.raises(MediaDecodeError, match="boom"): + synchronize_media_prepare_status(err) + + def test_success_passes_through_without_group(self): + synchronize_media_prepare_status(None) + + def test_two_rank_convergence_over_gloo(self, tmp_path): + """Rank 1 fails decode, rank 0 is healthy: both must exit the + protocol with the SAME client error — rank 1 re-raising its own, + rank 0 raising the reconstructed equivalent — instead of rank 0 + proceeding into (and hanging in) model collectives.""" + import torch.multiprocessing as mp + + init_file = tmp_path / "gloo_init" + ctx = mp.spawn( + _status_protocol_rank, + args=(2, str(init_file), str(tmp_path)), + nprocs=2, + join=False, + ) + # ``ProcessContext.join`` returns False whenever ANY child has + # exited but others still run — it is designed to be polled. + deadline = time.monotonic() + 240 + converged = False + while time.monotonic() < deadline: + if ctx.join(timeout=deadline - time.monotonic()): + converged = True + break + if not converged: + for proc in ctx.processes: + if proc.is_alive(): + proc.terminate() + assert converged, "status protocol did not converge (hang)" + healthy = (tmp_path / "rank0.txt").read_text() + failing = (tmp_path / "rank1.txt").read_text() + assert failing == "client:rank-local decode failure" + assert healthy == "client:[rank 1] rank-local decode failure" + + +class TestDecodeFrameLimit: + def test_default_and_env_override(self, monkeypatch): + monkeypatch.delenv("TRTLLM_MAX_REFERENCE_DECODE_FRAMES", raising=False) + assert max_reference_decode_frames() == 7200 + monkeypatch.setenv("TRTLLM_MAX_REFERENCE_DECODE_FRAMES", "12") + assert max_reference_decode_frames() == 12 + monkeypatch.setenv("TRTLLM_MAX_REFERENCE_DECODE_FRAMES", "0") + assert max_reference_decode_frames() is None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +class TestDecodeVideoReferenceWindow: + _DEVICE = torch.device("cuda:0") + + def _decode(self, data: bytes, **kwargs): + defaults = dict(window=5, keep="first", target_h=64, target_w=64, device=self._DEVICE) + defaults.update(kwargs) + return decode_video_reference_window(data, **defaults) + + @pytest.mark.parametrize("fixture", [_MP4, _AVI], ids=["mp4", "avi"]) + def test_keep_first_display_order(self, fixture): + # Display order despite forced B-frames, in both supported containers. + window = self._decode(fixture.read_bytes()) + assert window.shape == (5, 64, 64, 3) and window.dtype == torch.uint8 + assert window.device.type == "cuda" + assert _frame_indices(window) == [0, 1, 2, 3, 4] + + @pytest.mark.parametrize("fixture", [_MP4, _AVI], ids=["mp4", "avi"]) + def test_keep_last_ring_reorder(self, fixture): + window = self._decode(fixture.read_bytes(), keep="last") + assert _frame_indices(window) == [4, 5, 6, 7, 8] + + def test_rgb_channel_layout(self): + # Frame i carries a green horizontal bar at rows [7i, 7i+7) and a + # blue vertical bar at cols [7i, 7i+7): asserts the NVDEC output is + # RGB (not BGR) and spatially unflipped. 4:2:0 chroma blurs edges, + # so thresholds are generous. + window = self._decode(_MP4.read_bytes(), window=9, keep="first") + for i in (2, 6): + frame = window[i].float() + band = slice(7 * i, 7 * i + 7) + assert frame[band, :, 1].mean() > 150 # green bar rows + assert frame[:, band, 2].mean() > 150 # blue bar cols + assert frame[:, :, 1].mean() < frame[band, :, 1].mean() - 60 + + def test_surface_ownership_across_full_decode(self): + # If the ring held DLPack views instead of owned copies, NVDEC's + # surface recycling would overwrite earlier frames during the later + # decodes; per-frame content proves each retained frame is intact + # after the decoder finished the whole stream. + window = self._decode(_MP4.read_bytes(), window=9, keep="first") + assert _frame_indices(window) == list(range(9)) + + def test_target_resolution_resize(self): + window = self._decode(_MP4.read_bytes(), window=3, target_h=96, target_w=128) + assert window.shape == (3, 96, 128, 3) + + def test_window_longer_than_clip_returns_all(self): + window = self._decode(_MP4.read_bytes(), window=20) + assert _frame_indices(window) == list(range(9)) + + def test_corrupt_bytes_with_valid_magic_is_client_error(self): + payload = b"\x00\x00\x00\x18ftypmp42" + b"\x00" * 64 + with pytest.raises(MediaDecodeError): + self._decode(payload) + + def test_frame_limit_trips_on_emitted_frames(self, monkeypatch): + monkeypatch.setenv("TRTLLM_MAX_REFERENCE_DECODE_FRAMES", "4") + with pytest.raises(MediaDecodeError, match="decode limit"): + self._decode(_MP4.read_bytes(), keep="last") + + def test_frame_limit_disabled(self, monkeypatch): + monkeypatch.setenv("TRTLLM_MAX_REFERENCE_DECODE_FRAMES", "0") + window = self._decode(_MP4.read_bytes(), window=20, keep="last") + assert window.shape[0] == 9 + + def test_resize_perf_representative(self): + # Representative evidence for the local-tap resample: a 1080p frame + # to 720p-cover must be in the low-millisecond range. The bound is + # ~100x actual so it never flakes; a dense O(out*in) implementation + # (~11 GMAC/frame) would still trip it under contention. + frame = torch.randint(0, 256, (1, 1080, 1920, 3), dtype=torch.uint8, device=self._DEVICE) + resize_center_crop_uint8(frame, 720, 1280) # warmup + tap-cache build + torch.cuda.synchronize() + start = time.perf_counter() + for _ in range(10): + resize_center_crop_uint8(frame, 720, 1280) + torch.cuda.synchronize() + per_frame = (time.perf_counter() - start) / 10 + assert per_frame < 0.25, f"resize took {per_frame * 1e3:.1f} ms/frame" From a97f31b0da38f45b19742645814fafa9cb0c38dd Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 23 Jul 2026 21:36:17 -0700 Subject: [PATCH 42/64] Route video references by container signature, not decode probing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The serve boundary now classifies input_reference by magic bytes (PNG, JPEG, MP4 ftyp, AVI RIFF) — routing only, never acceptance: images still require a full PIL decode at the boundary, while video payloads pass through as untouched encoded bytes for the workers' NVDEC to decode. This removes every OpenCV import and probe from the VisualGen path (cv2 cannot be declared in requirements) along with the boundary video decoders and the decoded-size budget. Base64 decodes strictly; payload size is deliberately not part of the request-validity contract (body limits are deployment policy). Format contract documented accordingly: PNG/JPEG images; H.264 in MP4/AVI decode-tested; the rest best-effort per the GPU's decoder capabilities. Signed-off-by: Igor Shovkun --- examples/visual_gen/serve/README.md | 4 +- tensorrt_llm/inputs/media_io.py | 311 ++------------ tensorrt_llm/serve/openai_protocol.py | 27 +- tensorrt_llm/serve/visual_gen_utils.py | 75 ++-- .../visual_gen/test_trtllm_serve_endpoints.py | 81 +--- .../visual_gen/test_visual_gen_utils.py | 384 +++--------------- 6 files changed, 163 insertions(+), 719 deletions(-) diff --git a/examples/visual_gen/serve/README.md b/examples/visual_gen/serve/README.md index 38abcb633237..7eff00f35d81 100644 --- a/examples/visual_gen/serve/README.md +++ b/examples/visual_gen/serve/README.md @@ -286,8 +286,8 @@ You can customize these by: - `frame_rate` (canonical) or `fps` (alias): frames per second - `num_frames`: when set, wins over the `seconds * frame_rate` derivation - `seed`, `num_inference_steps`, `guidance_scale`, `max_sequence_length`, `negative_prompt`: per-request denoise controls -- `input_reference`: Reference image (I2V/TI2V) or video (V2V), classified by decoding the content — filename and content type are ignored; accepted as base64-encoded string in JSON or as a file in multipart form-data. Undecodable content returns HTTP 400. Video decode requires OpenCV on the server (not bundled — `pip install opencv-python-headless`). - - **Supported formats**: reference images are tested with PNG (`image/png`) and JPEG (`image/jpeg`). Reference videos are tested with MPEG-4 Part 2 video in MP4 (`video/mp4`) and Motion JPEG in AVI (`video/x-msvideo`). Filename and MIME metadata are not used for routing. Pillow must fully load an image; OpenCV must open a video and return decodable frames. Other formats depend on the installed decoder backend and are not guaranteed. Video references exceeding 1 GiB after RGB decoding are rejected. +- `input_reference`: Reference image (I2V/TI2V) or video (V2V), routed by container signature — filename and content type are ignored; accepted as base64-encoded string in JSON or as a file in multipart form-data. Images must fully decode at the boundary; video bytes pass through and decode on the workers' NVDEC (PyNvVideoCodec, a declared dependency). Unrecognized or undecodable content returns HTTP 400. + - **Supported formats**: reference images are tested with PNG (`image/png`) and JPEG (`image/jpeg`); they must fully decode with Pillow at the boundary. Reference videos are routed by container signature (MP4 `video/mp4`, AVI `video/x-msvideo`) and decoded on the workers' NVDEC — tested codec is H.264 in both containers; other codecs/profiles depend on the GPU's decoder capabilities and are best-effort. Filename and MIME metadata are never used for routing. Corrupt or undecodable content returns 400; a valid reference the deployment cannot fit returns 503. Reference decoding is bounded to 7200 frames by default (override with `TRTLLM_MAX_REFERENCE_DECODE_FRAMES`, `0` disables). - `extra_params`: model-specific overflow (see below) - `response_format`: `"b64_json"` or `"url"` - `format`: Generation content encoding. Video encoders: `"mp4"`, `"avi"`, `"auto"`. Tensor formats: `"safetensors"`, `"pt"` (carries video + audio + scalar metadata in one payload for LTX-2). diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index 739c6cd02c81..df6c01f8dedb 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -346,65 +346,38 @@ def _get_cv2(): return cv2 -# --- Content-classification probes ------------------------------------------ -# These classify media by asking the decoder itself, deliberately: -# * File suffixes are unreliable — client-controlled, often absent (the serve -# handles raw uploaded bytes with no filename at all), and never proof that -# the content matches the name. -# * Sniffing signatures ("magic numbers") would need an extra dependency -# (libmagic, as behind Unix `file`) or a hand-rolled signature table — and -# would still only name the container, not prove that this build's decoder -# can actually open it (codec support varies per PIL/OpenCV build). -# * A classifier that can disagree with the decoder is a deferred failure: -# content that passes the check but fails to decode later, deeper in the -# pipeline. Probing PIL/OpenCV directly makes "classified as X" and -# "decodes as X" the same statement by construction. -# Hence the names: these predicates promise decodability, not format identity. - - -def is_decodable_image_file(path) -> bool: - """True when ``path`` holds still-image content PIL can fully decode. - - Strict, like :func:`is_decodable_image_bytes`: ``Image.open`` is lazy, - so this also decodes the pixels (``load``) — a truncated file passes a - header-only probe and then fails at the actual load, far from the cause. - Non-image content is still rejected cheaply at the header parse. - """ - try: - with Image.open(path) as image: - image.load() - return True - except OSError: - # ``UnidentifiedImageError`` (bad header) subclasses ``OSError``; - # truncated files raise plain ``OSError`` from ``load``. - return False +# --- Reference-payload classification ---------------------------------------- +# The serve boundary routes a metadata-free reference payload (JSON base64 +# carries no filename or MIME type; multipart metadata is client-typed) to the +# image or video slot by its container signature. Routing only, never +# acceptance: an image is accepted by a full PIL decode at the boundary, a +# video by the worker's decoder — so corrupt content behind a valid signature +# still fails cleanly as a client error, and the signature can never disagree +# with what actually decodes. -def is_decodable_video_file(path) -> bool: - """True when ``path`` holds a decodable video stream (OpenCV-openable). +def sniff_media_kind(data) -> Optional[str]: + """Classify a reference payload by container signature. - Total predicate over content: False for images, audio, and undecodable - data alike. Stills are excluded explicitly — FFmpeg demuxes a single - image as a one-frame video stream, so a bare video probe would accept - every PNG/JPEG. ``_get_cv2`` raises a clear install hint if cv2 is absent. + Returns ``"image"`` (PNG/JPEG), ``"video"`` (ISO-BMFF/MP4 family or AVI), + or ``None`` for anything unrecognized. """ - if is_decodable_image_file(path): - return False - cv2 = _get_cv2() - capture = cv2.VideoCapture(str(path)) - try: - return bool(capture.isOpened() and capture.read()[0]) - finally: - capture.release() + header = bytes(data[:12]) + if header.startswith(b"\x89PNG\r\n\x1a\n") or header.startswith(b"\xff\xd8\xff"): + return "image" + if header[4:8] == b"ftyp": + return "video" + if header.startswith(b"RIFF") and header[8:12] == b"AVI ": + return "video" + return None def is_decodable_image_bytes(data) -> bool: """True when ``data`` holds still-image content PIL can fully decode. - In-memory counterpart of :func:`is_decodable_image_file`, equally - strict: ``Image.open`` is lazy, so this also decodes the pixels - (``load``) — a truncated file passes a header-only probe and would - then 500 at the worker's load instead of 400ing at the boundary. + Strict on purpose: ``Image.open`` is lazy, so this also decodes the + pixels (``load``) — a truncated file passes a header-only probe and + would then 500 at the worker's load instead of 400ing at the boundary. """ try: with Image.open(BytesIO(data)) as image: @@ -416,243 +389,11 @@ def is_decodable_image_bytes(data) -> bool: return False -def _read_frames_in_order(cv2, capture, max_frames: Optional[int]) -> List["Image.Image"]: - """Drain an opened ``VideoCapture`` into PIL frames, in order, no sampling.""" - frames = [] - while max_frames is None or len(frames) < max_frames: - ok, frame = capture.read() - if not ok: - break - frames.append(Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))) - return frames - - -def decode_video_frames(path, max_frames: Optional[int] = None) -> List["Image.Image"]: - """Decode a video file into its frames as PIL images, in order, no sampling. - - Unlike :func:`load_video` — which samples ``num_frames`` evenly for - VLM-style inputs — this preserves every frame sequentially from the start; - ``max_frames`` bounds the decode when only a prefix is needed. - Reference-conditioning consumers (e.g. video-to-video pipelines) pick - their own frame window downstream. - """ - cv2 = _get_cv2() - capture = cv2.VideoCapture(str(path)) - try: - if not capture.isOpened(): - raise ValueError(f"Could not open video file: {path}") - frames = _read_frames_in_order(cv2, capture, max_frames) - finally: - capture.release() - if not frames: - raise ValueError(f"Video file contains no frames: {path}") - return frames - - -def decode_video_frames_from_bytes(data, max_frames: Optional[int] = None) -> List["Image.Image"]: - """Decode raw video bytes into PIL frames, in order, no sampling. - - Fully in-memory when this OpenCV build has a stream-buffered backend - (:func:`_select_cv2_stream_buffered_backend`); otherwise the bytes spill - to an auto-deleted tempfile and take the :func:`decode_video_frames` path. - Raises ``ValueError`` when the bytes are not a decodable video. - """ - cv2 = _get_cv2() - backend = _select_cv2_stream_buffered_backend() - if backend is None: - with tempfile.NamedTemporaryFile() as spill: - spill.write(data) - spill.flush() - return decode_video_frames(spill.name, max_frames=max_frames) - - # cv2 keeps a non-owning view into the buffer; hold it until release(). - buffer = BytesIO(bytes(data)) - capture = cv2.VideoCapture(buffer, backend, []) - try: - if not capture.isOpened(): - raise ValueError(f"Could not open video from <{len(data)} bytes>.") - frames = _read_frames_in_order(cv2, capture, max_frames) - finally: - capture.release() - if not frames: - raise ValueError(f"Video bytes contain no frames (<{len(data)} bytes>).") - return frames - - -# Frame-image suffixes recognized when expanding a frame directory (see the -# selection rationale in ``load_video_frames_tensor``). -_IMAGE_SUFFIXES = frozenset({".png", ".jpg", ".jpeg", ".webp", ".bmp"}) - - -# Longest video, in frames, a client may request as *output*; used by the -# serve's ``num_frames`` cap (``openai_protocol``). +# Longest video, in frames, a client may request as *output* (the serve's +# ``num_frames`` cap in ``openai_protocol``) and the most reference frames a +# worker will decode from a video reference before raising. MAX_VIDEO_FRAMES = 7200 -# Hard budget for *decoded* video bytes when reading a reference. Bounds both -# the preallocated buffer and total accumulation, whatever the resolution — -# a frame count alone is no guard (7200 frames is ~18.5 GiB at 720p). The -# canonical 189-frame 720p reference is ~0.5 GiB, so this allows 2x headroom. -MAX_DECODED_VIDEO_BYTES = 1 << 30 # 1 GiB - - -class DecodedVideoTooLargeError(ValueError): - """Decoded reference exceeds ``MAX_DECODED_VIDEO_BYTES``. - - A ``ValueError`` subclass so boundary handlers still map it to a client - error (400), while letting its actionable message pass through instead of - being folded into generic "undecodable content" handling. - """ - - -def _decode_capture_to_tensor( - cv2, capture, max_frames: Optional[int], src_repr: str -) -> torch.Tensor: - """Drain an opened ``VideoCapture`` straight into a uint8 [T, H, W, C] tensor. - - Streams frames into one preallocated array, so with accurate container - metadata peak memory is a single copy of the video — roughly half of the - decode-to-PIL-list-then-``np.stack`` route. Unknown or misreported lengths - take the spill paths below, whose ``stack``/``concatenate`` transients can - reach ~2-3x the decoded size — still bounded, since everything is capped - by ``MAX_DECODED_VIDEO_BYTES`` - measured against real frame sizes: the preallocation (the declared frame - count is container metadata, not evidence) and the total decoded - accumulation (streams that exceed the budget raise instead of growing - without bound). Misreported lengths degrade gracefully: extra frames spill - to a side list, an over-declared buffer is trimmed, and an unknown length - falls back to list-and-stack. - """ - if not capture.isOpened(): - raise ValueError(f"Could not open video from {src_repr}.") - declared = int(capture.get(cv2.CAP_PROP_FRAME_COUNT) or 0) - - buffer = None - overflow = [] - count = 0 - decoded_bytes = 0 - while max_frames is None or count < max_frames: - ok, frame = capture.read() - if not ok: - break - rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) - decoded_bytes += rgb.nbytes - if decoded_bytes > MAX_DECODED_VIDEO_BYTES: - raise DecodedVideoTooLargeError( - f"Video from {src_repr} exceeds the decoded-size budget of " - f"{MAX_DECODED_VIDEO_BYTES >> 20} MiB at frame {count}; trim or " - "downscale the reference." - ) - if buffer is None and declared > 0: - # Size the preallocation from the declared count, clamped to the - # byte budget using the actual frame size (an over-declared buffer - # is virtual until written; the trim below drops the excess). - capacity = min(declared, MAX_DECODED_VIDEO_BYTES // max(rgb.nbytes, 1)) - if max_frames is not None: - capacity = min(capacity, max_frames) - if capacity > 0: - buffer = np.empty((capacity, *rgb.shape), dtype=np.uint8) - if buffer is not None and count < buffer.shape[0]: - buffer[count] = rgb - else: - overflow.append(rgb) - count += 1 - - if count == 0: - raise ValueError(f"Video contains no frames ({src_repr}).") - if buffer is None: - return torch.from_numpy(np.stack(overflow)) - filled = min(count, buffer.shape[0]) - if overflow: - return torch.from_numpy(np.concatenate([buffer[:filled], np.stack(overflow)])) - if filled < buffer.shape[0]: - # Over-declared container: trim without keeping the oversized buffer. - return torch.from_numpy(buffer[:filled].copy()) - return torch.from_numpy(buffer) - - -def decode_video_tensor(path, max_frames: Optional[int] = None) -> torch.Tensor: - """Decode a video file into a uint8 ``[T, H, W, C]`` RGB tensor. - - Tensor-native counterpart of :func:`decode_video_frames` — streams into a - single buffer instead of materializing PIL frames first. - """ - cv2 = _get_cv2() - capture = cv2.VideoCapture(str(path)) - try: - return _decode_capture_to_tensor(cv2, capture, max_frames, f"'{path}'") - finally: - capture.release() - - -def decode_video_tensor_from_bytes(data, max_frames: Optional[int] = None) -> torch.Tensor: - """Decode raw video bytes into a uint8 ``[T, H, W, C]`` RGB tensor. - - In-memory when this OpenCV build has a stream-buffered backend; otherwise - the bytes spill to an auto-deleted tempfile. Raises ``ValueError`` when the - bytes are not a decodable video. - """ - cv2 = _get_cv2() - backend = _select_cv2_stream_buffered_backend() - if backend is None: - with tempfile.NamedTemporaryFile() as spill: - spill.write(data) - spill.flush() - return decode_video_tensor(spill.name, max_frames=max_frames) - - # cv2 keeps a non-owning view into the buffer; hold it until release(). - buffer = BytesIO(bytes(data)) - capture = cv2.VideoCapture(buffer, backend, []) - try: - return _decode_capture_to_tensor(cv2, capture, max_frames, f"<{len(data)} bytes>") - finally: - capture.release() - - -def frames_to_tensor(frames: List["Image.Image"]) -> torch.Tensor: - """Stack PIL frames into a uint8 ``[T, H, W, C]`` RGB CPU tensor. - - The plain-tensor form video references travel in (e.g. - ``extra_params["video"]`` for video-to-video pipelines). - """ - if not frames: - raise ValueError("Cannot build a video tensor from an empty frame list.") - return torch.from_numpy(np.stack([np.asarray(f.convert("RGB")) for f in frames])) - - -def load_video_frames_tensor(source, max_frames: Optional[int] = None) -> torch.Tensor: - """Load a video reference from disk as a uint8 ``[T, H, W, C]`` RGB tensor. - - Public helper for building video-reference tensors client-side. Accepts a - video file, a single still image (one frame), or a directory of frame - images (sorted lexicographically). Dispatch is by content, not suffix - (see the content-classification probes above). - """ - path = Path(source) - if not path.exists(): - raise ValueError(f"Video reference path does not exist: {path}") - if path.is_dir(): - # Directories are selected by suffix, deliberately unlike single files: - # a frame directory is user-curated (names are the interface), suffix - # selection costs no file opens (a decodability probe is an open per - # entry — painful for thousands of frames on network filesystems), and - # it fails the right way — a selected frame that doesn't decode raises - # below, whereas a probe would silently drop corrupt frames and produce - # a video with holes. - frame_paths = sorted(p for p in path.iterdir() if p.suffix.lower() in _IMAGE_SUFFIXES) - if not frame_paths: - raise ValueError(f"No image frames found in directory: {path}") - if max_frames is not None: - frame_paths = frame_paths[:max_frames] - return frames_to_tensor([Image.open(p) for p in frame_paths]) - if is_decodable_image_file(path): - return frames_to_tensor([Image.open(path)]) - if is_decodable_video_file(path): - return decode_video_tensor(path, max_frames=max_frames) - raise ValueError( - f"Video reference must be a decodable video, a decodable image, or a " - f"directory of frame images; got undecodable {path}" - ) - def _select_cv2_stream_buffered_backend() -> Optional[int]: """Return a VideoCapture backend that can read from a Python `BytesIO`. diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index aa44551a592a..feee9a192692 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -1688,19 +1688,20 @@ class VideoGenerationRequest(OpenAIBaseModel): default=None, description=( "Optional image or video reference that guides generation. " - "Content is classified by decoding it — filename and MIME " - "metadata are not used for routing (JSON requests carry bare " - "base64 with no such metadata): Pillow must fully load an " - "image; OpenCV must open a video and return decodable " - "frames. Images condition image-to-video; videos condition " - "video-to-video on models that support it. Tested formats: " - "PNG (image/png) and JPEG (image/jpeg) images; MPEG-4 Part 2 " - "video in MP4 (video/mp4) and Motion JPEG in AVI " - "(video/x-msvideo). Other containers, codecs, profiles, and " - "pixel formats depend on the installed decoder backend and " - "are not guaranteed. Video references exceeding 1 GiB after " - "RGB decoding are rejected. JSON requests carry base64 " - "bytes; multipart requests upload the file."), + "Content is routed by its container signature — filename and " + "MIME metadata are ignored (JSON requests carry bare base64 " + "with no such metadata). Images (PNG image/png, JPEG " + "image/jpeg) must fully decode with Pillow at the boundary " + "and condition image-to-video. Video containers (MP4 " + "video/mp4, AVI video/x-msvideo) pass through encoded and " + "are decoded on the workers' NVDEC; tested codec is H.264 in " + "both containers, other codecs/profiles depend on the GPU's " + "decoder capabilities and are best-effort. Undecodable or " + "corrupt content fails as a client error (400); a valid " + "reference the deployment cannot fit fails as capacity " + "(503). Reference decoding is bounded to 7200 frames by " + "default (TRTLLM_MAX_REFERENCE_DECODE_FRAMES). JSON requests " + "carry base64 bytes; multipart requests upload the file."), ) # Resolution diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index a779ddb2fb3b..da3803485073 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -3,15 +3,10 @@ import os from typing import Any, Dict, List, Optional -from tensorrt_llm.inputs.media_io import ( - DecodedVideoTooLargeError, - decode_video_tensor_from_bytes, - is_decodable_image_bytes, -) +from tensorrt_llm.inputs.media_io import is_decodable_image_bytes, sniff_media_kind from tensorrt_llm.logger import logger from tensorrt_llm.serve.openai_protocol import ImageGenerationRequest, VideoGenerationRequest from tensorrt_llm.visual_gen import VisualGen, VisualGenParams -from tensorrt_llm.visual_gen.params import reduce_visual_gen_params # Per-field warnings for OpenAI-shaped knobs that the engine has no # semantic for. Each entry maps the request attribute to the message @@ -89,6 +84,23 @@ def _merge_extra_params( params.extra_params = None +def _read_reference_payload(reference) -> bytes: + """Read the ``input_reference`` payload (base64 JSON or multipart file). + + Payload size is deliberately not checked here: encoded size is not part + of the request-validity contract, and body limits belong to the + proxy/ASGI deployment layer (HTTP 413). Base64 decodes strictly so + malformed encodings — not sizes — are rejected. + """ + if isinstance(reference, str): + try: + return base64.b64decode(reference, validate=True) + except ValueError as exc: + # binascii.Error subclasses ValueError. + raise ValueError("input_reference is not valid base64 data.") from exc + return reference.file.read() + + def parse_visual_gen_params( request: ImageGenerationRequest | VideoGenerationRequest, id: str, @@ -159,15 +171,18 @@ def parse_visual_gen_params( ) params.num_frames = derived if request.input_reference is not None: - if isinstance(request.input_reference, str): - try: - payload = base64.b64decode(request.input_reference) - except ValueError as exc: - raise ValueError("input_reference is not valid base64 data.") from exc - else: - payload = request.input_reference.file.read() - - if is_decodable_image_bytes(payload): + payload = _read_reference_payload(request.input_reference) + kind = sniff_media_kind(payload) + if kind == "image": + # Signature routes; the full decode is still the acceptance + # check, so a truncated PNG 400s here instead of 500ing at + # the worker's load. + if not is_decodable_image_bytes(payload): + raise ValueError( + "input_reference has an image container signature but " + "does not fully decode; the file may be truncated or " + "corrupt." + ) # I2V: the stored image file is the cross-model contract. # every I2V pipeline reads ``params.image`` as a path. if media_storage_path is None: @@ -178,30 +193,24 @@ def parse_visual_gen_params( with open(ref_path, "wb") as f: f.write(payload) params.image = ref_path - else: - # V2V: decode in memory into a uint8 [T, H, W, C] tensor - try: - video = decode_video_tensor_from_bytes(payload) - except DecodedVideoTooLargeError: - # Still a 400, but with the actionable size message intact. - raise - except ValueError as exc: - raise ValueError( - "input_reference content is neither a decodable image " - "nor a decodable video." - ) from exc + elif kind == "video": + # V2V: encoded bytes pass through untouched; the worker + # demuxes and NVDEC-decodes them (acceptance happens there, + # so corrupt content behind a valid signature still fails as + # a client error). if params.extra_params is None: params.extra_params = {} - params.extra_params["video"] = video + params.extra_params["video"] = payload + else: + raise ValueError( + "input_reference is not a recognized media container; " + "supported inputs are PNG/JPEG images and MP4/AVI video." + ) _warn_if_set_with_no_semantic(request, getattr(generator, "model", None)) _merge_extra_params(params, request.extra_params, generator.extra_param_specs) - # Apply spec-declared transport reducers here as well (generate_async - # reduces non-mutatively, so without this the serve-owned params — held by - # the sync/async routes for the job's whole lifetime — would retain the - # full decoded reference, e.g. ~500 MiB per queued V2V request). - return reduce_visual_gen_params(params, generator.extra_param_specs) + return params class AsyncDictStore: diff --git a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py index a1558797a060..a905b4ca4084 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py @@ -17,10 +17,10 @@ import base64 import os from io import BytesIO +from pathlib import Path from typing import Optional from unittest.mock import patch -import numpy as np import pytest import torch from fastapi.testclient import TestClient @@ -61,19 +61,7 @@ def _assert_llm_envelope( assert message_contains in body["message"], body["message"] -def _require_opencv(): - """Skip unless OpenCV is installed (the shared optional video decoder). - - These tests drive the serve's in-memory reference classification/decode - (``is_decodable_image_bytes`` / ``decode_video_frames_from_bytes``), which decodes - via OpenCV (the same optional dep the multimodal video path uses). CI installs - ``opencv-python-headless`` in every test stage (``jenkins/L0_Test.groovy``), - so these tests always run there; the skip only spares bare local - environments, where cv2 stays optional (kept out of requirements by the - dependency policy). Returns the module for tests that synthesize a clip - with ``cv2.VideoWriter``. - """ - return pytest.importorskip("cv2") +_V2V_FIXTURE_MP4 = Path(__file__).parent / "test_data" / "cosmos3_v2v_ref_9f_bframes.mp4" def _make_dummy_image_tensor(height: int = 64, width: int = 64) -> torch.Tensor: @@ -858,25 +846,16 @@ def test_sync_video_generation_multipart_with_reference(self, video_client, tmp_ assert params.image.endswith("_reference") assert os.path.exists(params.image) - def test_sync_video_generation_multipart_with_video_reference(self, video_client, tmp_path): - """A video ``input_reference`` is decoded into a uint8 [T, H, W, C] - tensor on the model-specific ``video`` extra param (V2V). + def test_sync_video_generation_multipart_with_video_reference(self, video_client): + """A video ``input_reference`` rides through as the encoded payload on + the model-specific ``video`` extra param (V2V), byte-identical — the + serve never decodes video; the worker demuxes/NVDEC-decodes it. - The reference is classified by decoding its content, so the clip is - synthesized in-test with OpenCV — no video asset ships with the repo. + Routed by container signature, so a checked-in H.264/MP4 fixture drives + the boundary directly. """ - cv2 = _require_opencv() - ref_path = tmp_path / "ref.mp4" - # mp4v is a built-in FFmpeg mpeg4 encoder present in the opencv wheel. - writer = cv2.VideoWriter(str(ref_path), cv2.VideoWriter_fourcc(*"mp4v"), 4.0, (16, 16)) - try: - for _ in range(2): - writer.write(np.zeros((16, 16, 3), dtype=np.uint8)) - finally: - writer.release() - assert ref_path.exists() and ref_path.stat().st_size > 0 - - with open(ref_path, "rb") as f: + payload = _V2V_FIXTURE_MP4.read_bytes() + with open(_V2V_FIXTURE_MP4, "rb") as f: resp = video_client.post( "/v1/videos/generations", data={ @@ -890,51 +869,23 @@ def test_sync_video_generation_multipart_with_video_reference(self, video_client assert resp.status_code == 200 assert len(resp.content) > 0 - # Video content must NOT land on params.image; it's decoded into a - # uint8 [T, H, W, C] tensor on the model-specific ``video`` extra param + # Video content must NOT land on params.image; it rides the + # model-specific ``video`` extra param as the untouched encoded bytes # (the same intake the offline example's --video_path uses). params = video_client.mock_gen.last_params assert params.image is None - video = params.extra_params["video"] - assert isinstance(video, torch.Tensor) - assert video.dtype == torch.uint8 - assert video.ndim == 4 and video.shape[-1] == 3 + assert params.extra_params["video"] == payload def test_sync_video_generation_undecodable_reference_400(self, video_client): - """Content neither PIL nor OpenCV can decode is rejected at the boundary.""" - _require_opencv() + """Content matching no image or video container signature is rejected + at the boundary.""" resp = video_client.post( "/v1/videos/generations", data={"prompt": "x"}, files={"input_reference": ("doc.txt", BytesIO(b"not media"), "text/plain")}, ) assert resp.status_code == 400 - assert "neither a decodable image" in resp.text - - def test_sync_video_oversized_reference_400_with_message( - self, video_client, tmp_path, monkeypatch - ): - """A reference over the decoded-byte budget gets an HTTP 400 whose body - carries the actionable size message (not the generic undecodable one).""" - cv2 = _require_opencv() - from tensorrt_llm.inputs import media_io - - monkeypatch.setattr(media_io, "MAX_DECODED_VIDEO_BYTES", 100) - ref_path = tmp_path / "ref.mp4" - writer = cv2.VideoWriter(str(ref_path), cv2.VideoWriter_fourcc(*"mp4v"), 4.0, (16, 16)) - try: - for _ in range(4): - writer.write(np.zeros((16, 16, 3), dtype=np.uint8)) - finally: - writer.release() - with open(ref_path, "rb") as f: - resp = video_client.post( - "/v1/videos/generations", - data={"prompt": "x"}, - files={"input_reference": ("ref.mp4", f, "video/mp4")}, - ) - assert resp.status_code == 400 - assert "decoded-size budget" in resp.text + assert "not a recognized media container" in resp.text def test_sync_video_failure(self, failing_client): resp = failing_client.post( diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index 7ea50014dab6..7a4e0d83d98d 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -12,9 +12,8 @@ from __future__ import annotations import base64 -import os -import tempfile from io import BytesIO +from pathlib import Path from typing import Any, Dict, Optional import numpy as np @@ -293,141 +292,85 @@ def test_missing_media_storage_path_raises(self): with pytest.raises(ValueError, match="media_storage_path"): parse_visual_gen_params(request, "vid-2", generator, media_storage_path=None) + _TEST_DATA = Path(__file__).parent / "test_data" + @staticmethod - def _mp4_bytes(num_frames: int = 2) -> bytes: - """Encode a 16x16 mp4v-in-mp4 clip and return its bytes. - - ``mp4v`` is a built-in FFmpeg mpeg4 encoder present in the opencv wheel. - OpenCV writes only to a path, so encode to a tempfile and read it back. - """ - cv2 = pytest.importorskip("cv2") - with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp: - path = tmp.name - try: - writer = cv2.VideoWriter(path, cv2.VideoWriter_fourcc(*"mp4v"), 4.0, (16, 16)) - try: - assert writer.isOpened(), "cv2 VideoWriter failed to open (mp4v in MP4)" - for _ in range(num_frames): - writer.write(np.zeros((16, 16, 3), dtype=np.uint8)) - finally: - writer.release() - with open(path, "rb") as f: - return f.read() - finally: - os.remove(path) + def _mp4_bytes() -> bytes: + """9-frame H.264-in-MP4 fixture (provenance: test_data/README.md).""" + return ( + TestInputReferenceMaterialization._TEST_DATA / "cosmos3_v2v_ref_9f_bframes.mp4" + ).read_bytes() @staticmethod - def _avi_bytes(num_frames: int = 2) -> bytes: - """Encode a 16x16 Motion-JPEG-in-AVI clip and return its bytes. - - The second container/codec pair in the documented support contract; - ``MJPG`` is built into the opencv wheel like ``mp4v``. - """ - cv2 = pytest.importorskip("cv2") - with tempfile.NamedTemporaryFile(suffix=".avi", delete=False) as tmp: - path = tmp.name - try: - writer = cv2.VideoWriter(path, cv2.VideoWriter_fourcc(*"MJPG"), 4.0, (16, 16)) - try: - assert writer.isOpened(), "cv2 VideoWriter failed to open (MJPG in AVI)" - for _ in range(num_frames): - writer.write(np.zeros((16, 16, 3), dtype=np.uint8)) - finally: - writer.release() - with open(path, "rb") as f: - return f.read() - finally: - os.remove(path) + def _avi_bytes() -> bytes: + """Same 9 frames as H.264-in-AVI (provenance: test_data/README.md).""" + return ( + TestInputReferenceMaterialization._TEST_DATA / "cosmos3_v2v_ref_9f_bframes.avi" + ).read_bytes() def test_multipart_avi_reference_routes_to_video(self, tmp_path): - # The AVI/MJPEG contract pair must survive the real boundary, not - # just the decode primitive: classified as video by content, routed - # to the ``video`` extra param. - import torch - + # The AVI container signature must survive the real boundary: routed + # to the ``video`` extra param as untouched encoded bytes. generator = _StubVisualGen() - upload = UploadFile(file=BytesIO(self._avi_bytes()), filename="clip.avi") + payload = self._avi_bytes() + upload = UploadFile(file=BytesIO(payload), filename="clip.avi") request = VideoGenerationRequest(prompt="x", input_reference=upload) params = parse_visual_gen_params( request, "vid-avi", generator, media_storage_path=str(tmp_path) ) assert params.image is None - video = params.extra_params["video"] - assert isinstance(video, torch.Tensor) - assert tuple(video.shape) == (2, 16, 16, 3) - - def test_multipart_video_reference_routes_to_extra_params_tensor(self, tmp_path): - import torch + assert params.extra_params["video"] == payload + def test_multipart_video_reference_routes_to_extra_params_bytes(self, tmp_path): generator = _StubVisualGen() - upload = UploadFile(file=BytesIO(self._mp4_bytes()), filename="clip.mp4") + payload = self._mp4_bytes() + upload = UploadFile(file=BytesIO(payload), filename="clip.mp4") request = VideoGenerationRequest(prompt="x", input_reference=upload) params = parse_visual_gen_params( request, "vid-3", generator, media_storage_path=str(tmp_path) ) - # Video content is decoded into a uint8 [T, H, W, C] tensor on the - # model-specific ``video`` extra param, not params.image. The worker - # crops the conditioning window + VAE-encodes. + # Video content rides the model-specific ``video`` extra param as the + # encoded payload, byte-identical — the boundary never decodes video; + # the worker demuxes/NVDEC-decodes the conditioning window. assert params.image is None - video = params.extra_params["video"] - assert isinstance(video, torch.Tensor) - assert video.dtype == torch.uint8 - assert video.ndim == 4 and video.shape[0] == 2 and video.shape[-1] == 3 - # Video references are decoded in memory — nothing lands in media storage. + assert params.extra_params["video"] == payload + # Nothing lands in media storage for video references. assert list(tmp_path.iterdir()) == [] def test_video_reference_needs_no_media_storage(self): - # The decode is in-memory, so V2V works without a storage path at all - # (only image references persist a file for the worker to read). - import torch - + # Video bytes pass through in memory, so V2V works without a storage + # path at all (only image references persist a file for the worker). generator = _StubVisualGen() b64 = base64.b64encode(self._mp4_bytes()).decode() request = VideoGenerationRequest(prompt="x", input_reference=b64) params = parse_visual_gen_params(request, "vid-9", generator, media_storage_path=None) - assert isinstance(params.extra_params["video"], torch.Tensor) - - def test_base64_video_reference_routes_to_extra_params_tensor(self, tmp_path): - # Classification is content-based, so the JSON/base64 path can - # carry video even though it has no content-type or filename. - import torch + assert params.extra_params["video"] == self._mp4_bytes() + def test_base64_video_reference_routes_to_extra_params_bytes(self, tmp_path): + # Routing is signature-based, so the JSON/base64 path can carry video + # even though it has no content-type or filename. generator = _StubVisualGen() - b64 = base64.b64encode(self._mp4_bytes()).decode() + payload = self._mp4_bytes() + b64 = base64.b64encode(payload).decode() request = VideoGenerationRequest(prompt="x", input_reference=b64) params = parse_visual_gen_params( request, "vid-4", generator, media_storage_path=str(tmp_path) ) assert params.image is None - assert isinstance(params.extra_params["video"], torch.Tensor) - - def test_video_reference_reduced_before_routes_hold_params(self): - """``parse_visual_gen_params`` applies the spec reducers itself: the - sync/async routes hold the returned params for the whole job lifetime, - and ``generate_async`` reduces non-mutatively — without reduction at - parse, the serve would retain the full decoded clip per queued - request.""" + assert params.extra_params["video"] == payload + + def test_video_reference_bytes_survive_real_specs(self): + """With the real cosmos3 specs loaded, parsing leaves the encoded + payload byte-identical in ``extra_params['video']`` — the boundary + never transforms video content; the worker decodes the window.""" from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import COSMOS3_EXTRA_SPECS generator = _StubVisualGen(extra_param_specs=COSMOS3_EXTRA_SPECS) - b64 = base64.b64encode(self._mp4_bytes(num_frames=8)).decode() + payload = self._mp4_bytes() + b64 = base64.b64encode(payload).decode() request = VideoGenerationRequest(prompt="x", input_reference=b64) params = parse_visual_gen_params(request, "vid-10", generator, media_storage_path=None) - # Cropped to the default conditioning window (5) at parse, not later. - assert params.extra_params["video"].shape[0] == 5 - - def test_budget_error_message_survives_parse(self, monkeypatch): - # Helper-level: DecodedVideoTooLargeError passes through the generic - # "undecodable" handler with its message intact. The HTTP 400 itself - # is asserted in test_trtllm_serve_endpoints.py. - from tensorrt_llm.inputs import media_io - - monkeypatch.setattr(media_io, "MAX_DECODED_VIDEO_BYTES", 100) - generator = _StubVisualGen() - b64 = base64.b64encode(self._mp4_bytes(num_frames=4)).decode() - request = VideoGenerationRequest(prompt="x", input_reference=b64) - with pytest.raises(ValueError, match="decoded-size budget"): - parse_visual_gen_params(request, "vid-11", generator, media_storage_path=None) + assert params.extra_params["video"] == payload def test_multipart_image_reference_routes_to_image(self, tmp_path): # JPEG upload: content sniffing classifies it as an image and routes @@ -447,11 +390,10 @@ def test_multipart_image_reference_routes_to_image(self, tmp_path): assert str(params.image).endswith("vid-5_reference") def test_undecodable_reference_raises_and_cleans_up(self, tmp_path): - pytest.importorskip("cv2") generator = _StubVisualGen() b64 = base64.b64encode(b"neither an image nor a video").decode() request = VideoGenerationRequest(prompt="x", input_reference=b64) - with pytest.raises(ValueError, match="neither a decodable image"): + with pytest.raises(ValueError, match="not a recognized media container"): parse_visual_gen_params(request, "vid-6", generator, media_storage_path=str(tmp_path)) # Classification runs on the bytes; rejected content never touches disk. assert list(tmp_path.iterdir()) == [] @@ -481,26 +423,6 @@ def read(self, *args, **kwargs): assert list(tmp_path.iterdir()) == [] -class TestMediaFileProbes: - """File-path probes backing the offline producer path (``media_io``).""" - - def test_truncated_image_file_is_not_decodable(self, tmp_path): - # Same strictness as the bytes probe: a truncated PNG parses its - # header but must not classify as a decodable image — and the video - # probe must not rescue it as a one-frame video either. - pytest.importorskip("cv2") - from tensorrt_llm.inputs.media_io import is_decodable_image_file, is_decodable_video_file - - whole = tmp_path / "whole.png" - Image.fromarray(np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8)).save(whole) - truncated = tmp_path / "truncated.png" - truncated.write_bytes(whole.read_bytes()[: whole.stat().st_size // 2]) - - assert is_decodable_image_file(whole) - assert not is_decodable_image_file(truncated) - assert not is_decodable_video_file(truncated) - - class TestMediaBytesProbes: """The in-memory probe/decode primitives the serve boundary runs on.""" @@ -517,47 +439,21 @@ def test_is_decodable_image_bytes(self): # Video bytes are not an image (mp4 has no PIL-openable header). assert not is_decodable_image_bytes(TestInputReferenceMaterialization._mp4_bytes()) - def test_decode_video_frames_from_bytes(self): - pytest.importorskip("cv2") - from tensorrt_llm.inputs.media_io import decode_video_frames_from_bytes - - frames = decode_video_frames_from_bytes(TestInputReferenceMaterialization._mp4_bytes()) - assert len(frames) == 2 - assert all(isinstance(f, Image.Image) for f in frames) - - def test_decode_video_frames_from_bytes_max_frames(self): - pytest.importorskip("cv2") - from tensorrt_llm.inputs.media_io import decode_video_frames_from_bytes - - frames = decode_video_frames_from_bytes( - TestInputReferenceMaterialization._mp4_bytes(), max_frames=1 - ) - assert len(frames) == 1 - - def test_decode_video_frames_from_bytes_rejects_garbage(self): - pytest.importorskip("cv2") - from tensorrt_llm.inputs.media_io import decode_video_frames_from_bytes - - with pytest.raises(ValueError): - decode_video_frames_from_bytes(b"not a video at all") - - def test_avi_mjpeg_bytes_decode(self): - # Motion-JPEG-in-AVI — the second container/codec pair in the - # documented support contract: not an image, decodes as video. - pytest.importorskip("cv2") - import torch - - from tensorrt_llm.inputs.media_io import ( - decode_video_tensor_from_bytes, - is_decodable_image_bytes, - ) - - payload = TestInputReferenceMaterialization._avi_bytes() - assert not is_decodable_image_bytes(payload) - video = decode_video_tensor_from_bytes(payload) - assert isinstance(video, torch.Tensor) - assert video.dtype == torch.uint8 - assert tuple(video.shape) == (2, 16, 16, 3) + def test_sniff_media_kind(self): + from tensorrt_llm.inputs.media_io import sniff_media_kind + + png = BytesIO() + Image.new("RGB", (2, 2)).save(png, format="PNG") + jpg = BytesIO() + Image.new("RGB", (2, 2)).save(jpg, format="JPEG") + assert sniff_media_kind(png.getvalue()) == "image" + assert sniff_media_kind(jpg.getvalue()) == "image" + assert sniff_media_kind(TestInputReferenceMaterialization._mp4_bytes()) == "video" + assert sniff_media_kind(TestInputReferenceMaterialization._avi_bytes()) == "video" + assert sniff_media_kind(b"plain text, not media") is None + assert sniff_media_kind(b"") is None + # RIFF alone is not AVI (e.g. WAV audio is RIFF too). + assert sniff_media_kind(b"RIFF\x00\x00\x00\x00WAVEfmt ") is None def test_truncated_image_bytes_are_not_decodable(self): # A truncated PNG still opens (the header parses) but cannot decode @@ -576,9 +472,9 @@ def test_truncated_image_bytes_are_not_decodable(self): assert not is_decodable_image_bytes(truncated) def test_truncated_image_reference_rejected_at_parse(self): - # End of the chain: a truncated image upload is rejected as a client - # error at the boundary (never routed into the worker). - pytest.importorskip("cv2") + # End of the chain: a truncated image upload sniffs as an image but + # fails the strict decode -> client error at the boundary (never + # routed into the worker). rng_pixels = np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8) buf = BytesIO() Image.fromarray(rng_pixels).save(buf, format="PNG") @@ -588,48 +484,9 @@ def test_truncated_image_reference_rejected_at_parse(self): request = VideoGenerationRequest( prompt="x", input_reference=base64.b64encode(truncated).decode() ) - with pytest.raises(ValueError, match="neither a decodable"): + with pytest.raises(ValueError, match="does not fully decode"): parse_visual_gen_params(request, "vid-12", generator, media_storage_path=None) - def test_decode_video_tensor_matches_pil_route(self): - # The streaming decoder (single preallocated buffer — the low-peak - # path the serve uses) must produce byte-identical output to the - # PIL-frames route. - pytest.importorskip("cv2") - import torch - - from tensorrt_llm.inputs.media_io import ( - decode_video_frames_from_bytes, - decode_video_tensor_from_bytes, - frames_to_tensor, - ) - - data = TestInputReferenceMaterialization._mp4_bytes() - streamed = decode_video_tensor_from_bytes(data) - via_pil = frames_to_tensor(decode_video_frames_from_bytes(data)) - assert streamed.dtype == torch.uint8 and streamed.ndim == 4 - assert torch.equal(streamed, via_pil) - assert torch.equal(decode_video_tensor_from_bytes(data, max_frames=1), via_pil[:1]) - - def test_decode_video_tensor_rejects_garbage(self): - pytest.importorskip("cv2") - from tensorrt_llm.inputs.media_io import decode_video_tensor_from_bytes - - with pytest.raises(ValueError): - decode_video_tensor_from_bytes(b"not a video at all") - - def test_tempfile_fallback_without_stream_backend(self, monkeypatch): - # Old OpenCV builds have no stream-buffered backend; the bytes spill - # to an auto-deleted tempfile and decode through the path route. - pytest.importorskip("cv2") - from tensorrt_llm.inputs import media_io - - monkeypatch.setattr(media_io, "_select_cv2_stream_buffered_backend", lambda: None) - frames = media_io.decode_video_frames_from_bytes( - TestInputReferenceMaterialization._mp4_bytes() - ) - assert len(frames) == 2 - # ============================================================================= # _merge_extra_params — the merge truth table @@ -682,118 +539,3 @@ def test_empty_extras_dict_normalizes_to_none(self): params = self._make_params() _merge_extra_params(params, request_extras=None, extra_param_specs={}) assert params.extra_params is None - - -class _FakeCapture: - """Stands in for ``cv2.VideoCapture`` to exercise declared-count handling.""" - - def __init__(self, frames, declared): - self._frames = list(frames) - self._pos = 0 - self._declared = declared - - def isOpened(self): - return True - - def get(self, prop): - return self._declared - - def read(self): - if self._pos < len(self._frames): - frame = self._frames[self._pos] - self._pos += 1 - return True, frame - return False, None - - -class _FakeCv2: - CAP_PROP_FRAME_COUNT = 7 - COLOR_BGR2RGB = 4 - - @staticmethod - def cvtColor(frame, code): - return frame - - -class TestDecodeCaptureGuards: - """``_decode_capture_to_tensor`` against containers that misreport length. - - The declared frame count is metadata, not evidence — the decoder must - stream correctly whether it is accurate, unknown, under-, over-, or - absurdly reported.""" - - def _decode(self, num_frames, declared, max_frames=None): - import torch - - from tensorrt_llm.inputs.media_io import _decode_capture_to_tensor - - frames = [np.full((4, 4, 3), i, dtype=np.uint8) for i in range(num_frames)] - out = _decode_capture_to_tensor( - _FakeCv2, _FakeCapture(frames, declared), max_frames, "test" - ) - assert out.dtype == torch.uint8 - return out - - def test_accurate_declaration(self): - out = self._decode(3, declared=3) - assert out.shape == (3, 4, 4, 3) - assert int(out[2, 0, 0, 0]) == 2 # frame order preserved - - def test_unknown_declaration_falls_back(self): - assert self._decode(3, declared=0).shape == (3, 4, 4, 3) - assert self._decode(3, declared=-1).shape == (3, 4, 4, 3) - - def test_underreported_declaration_keeps_overflow(self): - out = self._decode(5, declared=2) - assert out.shape == (5, 4, 4, 3) - assert int(out[4, 0, 0, 0]) == 4 - - def test_overreported_declaration_trims_storage(self): - out = self._decode(3, declared=10) - assert out.shape == (3, 4, 4, 3) - # The oversized buffer is not retained behind the result. - assert out.untyped_storage().size() == out.numel() - - def test_absurd_declaration_allocation_is_byte_budgeted(self, monkeypatch): - # Realistic 720p frames + an absurd declared count: the preallocation - # request itself must stay within MAX_DECODED_VIDEO_BYTES (a frame - # cap alone would still be ~18.5 GiB at 720p). Spy on np.empty to - # assert the requested size, not just the result. - import math - - from tensorrt_llm.inputs import media_io - - requested = [] - real_empty = media_io.np.empty - - def spy(shape, dtype=None): - requested.append((tuple(shape), dtype)) - return real_empty(shape, dtype=dtype) - - monkeypatch.setattr(media_io.np, "empty", spy) - # Lower the budget so the (real) allocation the spy delegates to stays - # small; the assertion is about the *requested* size honoring it. - monkeypatch.setattr(media_io, "MAX_DECODED_VIDEO_BYTES", 32 << 20) - frames = [np.zeros((720, 1280, 3), dtype=np.uint8) for _ in range(3)] - out = media_io._decode_capture_to_tensor( - _FakeCv2, _FakeCapture(frames, declared=10**9), None, "test" - ) - assert out.shape == (3, 720, 1280, 3) - (shape, _dtype) = requested[0] - assert math.prod(shape) <= media_io.MAX_DECODED_VIDEO_BYTES - - def test_decoded_byte_budget_rejects_oversized_streams(self, monkeypatch): - # Total accumulation is bounded too — a stream that exceeds the budget - # raises instead of growing without bound (unknown-length containers - # included, where no buffer is preallocated at all). - from tensorrt_llm.inputs import media_io - - monkeypatch.setattr(media_io, "MAX_DECODED_VIDEO_BYTES", 100) - frames = [np.zeros((4, 4, 3), dtype=np.uint8) for _ in range(5)] # 48 B each - with pytest.raises(ValueError, match="decoded-size budget"): - media_io._decode_capture_to_tensor( - _FakeCv2, _FakeCapture(frames, declared=0), None, "test" - ) - - def test_max_frames_bounds_decode(self): - assert self._decode(5, declared=5, max_frames=2).shape[0] == 2 From 7cb9ff388cc89925005f9920b3bfe5def186d6e6 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 23 Jul 2026 21:36:37 -0700 Subject: [PATCH 43/64] Classify worker failures as client or capacity across all surfaces MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DiffusionResponse carries an error_type so worker-side failures keep their class over the wire: client errors (unusable reference content, conditioning bounds, decode-work limit) surface as HTTP 400 and as ValueError from VisualGen.generate() — uniform with coordinator preflight — while capacity failures (allocation/OOM for a valid request) surface as HTTP 503 and RuntimeError. Ambiguous NVDEC decoder-init failures and everything unclassified stay 500. No new public exception class. Signed-off-by: Igor Shovkun --- tensorrt_llm/_torch/visual_gen/executor.py | 13 +++++++- tensorrt_llm/serve/openai_video_routes.py | 10 ++++++ tensorrt_llm/visual_gen/visual_gen.py | 37 ++++++++++++---------- 3 files changed, 43 insertions(+), 17 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/executor.py b/tensorrt_llm/_torch/visual_gen/executor.py index 95a170120795..1d4580bdd6ee 100644 --- a/tensorrt_llm/_torch/visual_gen/executor.py +++ b/tensorrt_llm/_torch/visual_gen/executor.py @@ -15,6 +15,7 @@ import torch.multiprocessing as mp import zmq +from tensorrt_llm._torch.visual_gen.media_decode import classify_worker_error from tensorrt_llm._torch.visual_gen.output import PipelineOutput from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader from tensorrt_llm.executor.ipc import ZeroMqQueue @@ -253,6 +254,11 @@ class DiffusionResponse: model-specific fields populated. Set to ``None`` on the error path; on the READY signal it carries a ``dict`` instead. error_msg: Error message if generation failed. + error_type: Failure class when ``error_msg`` is set: ``"client"`` + (unusable request content → 400 / ``ValueError``), ``"capacity"`` + (valid request does not fit the deployment → 503 / + ``RuntimeError``), or ``None`` for unclassified runtime failures + (500). generation: Wall-clock time the executor measured around the engine's inference call (host ``time.perf_counter()``), in seconds. Default ``0.0`` so the dataclass round-trips through @@ -263,6 +269,7 @@ class DiffusionResponse: request_id: int output: Optional[PipelineOutput] = None error_msg: Optional[str] = None + error_type: Optional[str] = None generation: float = 0.0 @@ -452,7 +459,11 @@ def process_request(self, req: DiffusionRequest): logger.error(traceback.format_exc()) if self.rank == 0: self.response_queue.put( - DiffusionResponse(request_id=req.request_id, error_msg=str(e)) + DiffusionResponse( + request_id=req.request_id, + error_msg=str(e), + error_type=classify_worker_error(e), + ) ) diff --git a/tensorrt_llm/serve/openai_video_routes.py b/tensorrt_llm/serve/openai_video_routes.py index 4c17ba5d5ad4..49689ff1f946 100644 --- a/tensorrt_llm/serve/openai_video_routes.py +++ b/tensorrt_llm/serve/openai_video_routes.py @@ -24,6 +24,7 @@ from fastapi.responses import FileResponse, JSONResponse, Response from pydantic import ValidationError +from tensorrt_llm._torch.visual_gen.media_decode import VisualGenCapacityError from tensorrt_llm.logger import logger from tensorrt_llm.media.encoding import resolve_video_format from tensorrt_llm.media.tensor_payload import is_tensor_format @@ -124,6 +125,15 @@ async def openai_video_generation_sync(self, raw_request: Request) -> Response: except ValueError as exc: logger.error(f"Video request error: {exc}") return self.create_error_response(str(exc), status_code=HTTPStatus.BAD_REQUEST) + except VisualGenCapacityError as exc: + # Valid request that does not fit this deployment (decode / + # allocation capacity) — a server condition, not client error. + logger.error(f"Video request capacity error: {exc}") + return self.create_error_response( + str(exc), + err_type="ServiceUnavailableError", + status_code=HTTPStatus.SERVICE_UNAVAILABLE, + ) if output.video is None: return self.create_error_response( diff --git a/tensorrt_llm/visual_gen/visual_gen.py b/tensorrt_llm/visual_gen/visual_gen.py index c740ae7d140c..7528f5068c70 100644 --- a/tensorrt_llm/visual_gen/visual_gen.py +++ b/tensorrt_llm/visual_gen/visual_gen.py @@ -32,11 +32,7 @@ from tensorrt_llm._torch.visual_gen.pipeline_registry import PIPELINE_REGISTRY, AutoPipeline from tensorrt_llm.visual_gen.args import VisualGenArgs from tensorrt_llm.visual_gen.output import VisualGenOutput -from tensorrt_llm.visual_gen.params import ( - VisualGenParams, - reduce_visual_gen_params, - validate_visual_gen_params, -) +from tensorrt_llm.visual_gen.params import VisualGenParams, validate_visual_gen_params __all__ = [ "VisualGen", @@ -160,20 +156,33 @@ def cancel(self): # ----- internals ----- def _build_resolved(self, response: "DiffusionResponse"): + # Failure class travels on the result object, not on the public + # ``VisualGenOutput`` — no new public field for an error taxonomy. + self._error_type = getattr(response, "error_type", None) if self._batch_size is None: return to_visual_gen_output(response) return split_visual_gen_output(response, self._batch_size) def _resolved_value(self): - # For single prompts, surface engine-side failure as - # ``RuntimeError``. Request-parameter validation is enforced - # synchronously at :meth:`VisualGen.generate_async` entry, so - # anything reaching this point is by definition a runtime - # failure from ``pipeline.infer()``. For batches, return the - # list as-is so callers iterate per-item ``error``. + # For single prompts, surface engine-side failure typed by class: + # worker-classified client errors (unusable reference content, + # conditioning bounds) raise ``ValueError`` — uniform with the + # synchronous parameter validation at ``generate_async`` entry — + # and capacity failures raise ``RuntimeError`` + # (``VisualGenCapacityError``), like any unclassified runtime + # failure. For batches, return the list as-is so callers iterate + # per-item ``error``. if self._batch_size is None and isinstance(self._resolved, VisualGenOutput): if self._resolved.error is not None: - raise RuntimeError(f"Generation failed: {self._resolved.error}") + message = f"Generation failed: {self._resolved.error}" + error_type = getattr(self, "_error_type", None) + if error_type == "client": + raise ValueError(message) + if error_type == "capacity": + from tensorrt_llm._torch.visual_gen.media_decode import VisualGenCapacityError + + raise VisualGenCapacityError(message) + raise RuntimeError(message) return self._resolved @@ -391,10 +400,6 @@ def generate_async( # from the READY signal) and skip validation — there's nothing # user-supplied to validate against. if params is not None: - # Shrink oversized transport payloads (spec-declared reducers) - # before the deep copy, so the copy/pickle/broadcast chain only - # ever carries the reduced values. Non-mutating for the caller. - params = reduce_visual_gen_params(params, self.executor.extra_param_specs) resolved_params = params.model_copy(deep=True) # Raising in the caller's process means ``ValueError`` reaches # the user as a natural Python exception; the worker only has From 51c5f53139db18db01537bde060e9138b211f683 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 23 Jul 2026 21:36:38 -0700 Subject: [PATCH 44/64] Make encoded bytes the sole Cosmos3 V2V reference contract MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit extra_params['video'] now carries encoded MP4/AVI bytes only, decoded per rank inside forward(): conditioning indexes are bound-checked against the output latent length before any window math, the decoded frames are freed as soon as the VAE has encoded them, and per-rank prepare outcomes converge through the status protocol before any model collective. The interim tensor form and everything that existed to ship it are removed: the tensor/tensor_or_bytes type-map entries, the coordinator crop reducer, the ExtraParamSchema.reducer hook and reduce_visual_gen_params plumbing, and cosmos3/utils.py (orphaned). Smokes drive the production bytes path via the checked-in fixture — keep=last is now exercised through the real NVDEC decode. Signed-off-by: Igor Shovkun --- .../visual_gen/models/cosmos3/defaults.py | 68 ++------- .../models/cosmos3/pipeline_cosmos3.py | 119 ++++++++++----- .../_torch/visual_gen/models/cosmos3/utils.py | 18 --- tensorrt_llm/_torch/visual_gen/pipeline.py | 13 +- tensorrt_llm/visual_gen/params.py | 36 +---- .../visual_gen/test_cosmos3_pipeline.py | 141 +++--------------- .../visual_gen/test_visual_gen_params.py | 79 ++-------- 7 files changed, 134 insertions(+), 340 deletions(-) delete mode 100644 tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index 102ad4f27a09..b38dd972e743 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -19,9 +19,8 @@ from typing import Dict, Iterable -import torch - from tensorrt_llm._torch.visual_gen.pipeline import ExtraParamSchema +from tensorrt_llm.inputs.media_io import sniff_media_kind # --------------------------------------------------------------------------- # Constant tables @@ -86,53 +85,19 @@ def _normalize_condition_video_keep(keep: str | None) -> str: return normalized -def _crop_video_frames(video, extra_params) -> torch.Tensor: - """Crop the V2V reference to the conditioning window before transport. - - Runs once in the coordinator (spec ``reducer``) before the request is - deep-copied, pickled over ZMQ, and broadcast per rank: a full 189-frame - 720p reference is ~520 MiB while the default conditioning window is 5 - frames (~14 MiB). Semantics-preserving — the worker's own first/last crop - is idempotent, so reduced and unreduced tensors generate identically. - Anything invalid is returned unchanged for the validators to reject. - """ - if not isinstance(video, torch.Tensor) or video.ndim != 4: - return video - try: - indexes = _normalize_condition_video_latent_indexes( - extra_params.get("condition_video_latent_indexes") - ) - keep = _normalize_condition_video_keep(extra_params.get("condition_video_keep")) - except (TypeError, ValueError): - return video - # 4 = Cosmos3 VAE temporal compression; if a future VAE changes it, the - # worker pads/crops the window itself, so a mismatch degrades gracefully. - window = max(indexes) * 4 + 1 - if video.shape[0] <= window: - return video - sliced = video[-window:] if keep == "last" else video[:window] - # A slice is a view over the full storage and would pickle all of it; - # clone so the transport payload owns only the window. - return sliced.clone() - - def _validate_output_type(output_type: str) -> None: if output_type not in ("video", "image"): raise ValueError(f"Cosmos3 output_type must be 'video' or 'image', got {output_type!r}.") -def _validate_video_reference_tensor(video: torch.Tensor) -> None: - if video.ndim != 4 or video.shape[-1] != 3: - raise ValueError( - f"Cosmos3 video reference must be a uint8 [T, H, W, C] RGB tensor, " - f"got shape {tuple(video.shape)}." - ) - if video.dtype != torch.uint8: - raise ValueError(f"Cosmos3 video reference must have dtype uint8, got {video.dtype}.") - if video.device.type != "cpu": +def _validate_video_reference(video) -> None: + """Preflight for the ``video`` extra param: encoded MP4/AVI bytes.""" + if not video: + raise ValueError("Cosmos3 video reference bytes are empty.") + if sniff_media_kind(video) != "video": raise ValueError( - f"Cosmos3 video reference must be a CPU tensor, got device '{video.device}' " - "(it is pickled to the workers; keep decoded references on the host)." + "Cosmos3 video reference bytes are not a recognized video " + "container (supported: MP4/AVI)." ) @@ -217,18 +182,15 @@ def _validate_video_reference_tensor(video: torch.Tensor) -> None: description="Optional scheduler flow shift override. Uses the Cosmos3 mode default when omitted.", ), "video": ExtraParamSchema( - type="tensor", + type="bytes", default=None, description=( - "V2V reference: decoded video frames as a uint8 [T, H, W, C] RGB " - "torch.Tensor (build one from a file with " - "tensorrt_llm.inputs.media_io.load_video_frames_tensor). The " - "coordinator crops it to the conditioning window per " - "condition_video_latent_indexes / condition_video_keep before " - "dispatch; the worker VAE-encodes it. Media is always decoded " - "by the producer." + "V2V reference: encoded MP4/AVI bytes (e.g. " + "Path(video).read_bytes()). Each worker rank demuxes them from " + "memory and NVDEC-decodes only the conditioning window per " + "condition_video_latent_indexes / condition_video_keep, resized " + "to the output resolution, then VAE-encodes it." ), - validator=_validate_video_reference_tensor, - reducer=_crop_video_frames, + validator=_validate_video_reference, ), } diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index 854fe849d93e..9050e77c8a61 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -26,6 +26,11 @@ from diffusers.video_processor import VideoProcessor from transformers import Qwen2Tokenizer +from tensorrt_llm._torch.visual_gen.media_decode import ( + MediaDecodeError, + decode_video_reference_window, + synchronize_media_prepare_status, +) from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline from tensorrt_llm._torch.visual_gen.pipeline_registry import PipelineComponent, register_pipeline @@ -39,12 +44,12 @@ COSMOS3_EXTRA_SPECS, COSMOS3_PIPELINE_DEFAULTS, COSMOS3_T2I_PARAMS, + _normalize_condition_video_keep, _normalize_condition_video_latent_indexes, ) from .guardrails import check_video_safety, download_guardrail_checkpoint from .sound_tokenizer import LatentAutoEncoderV2 from .transformer_cosmos3 import Cosmos3VFMTransformer -from .utils import pil_to_rgb COSMOS3_DEFAULT_NEGATIVE_PROMPT = "" # NOTE: Intentional typo in "give" instead of "given" to match training setup. @@ -278,7 +283,7 @@ def _run_warmup(self, height: int, width: int, num_frames: int, steps: int) -> N def infer(self, req): extra_params = req.params.extra_params or {} output_type = extra_params.get("output_type", "video") - video = extra_params.get("video") # Tensor[T, H, W, C, dtype=uint8] + video = extra_params.get("video") # encoded MP4/AVI bytes (the extra-param contract) return self.forward( prompt=req.prompt, @@ -596,20 +601,17 @@ def decode_audio(self, latent: torch.Tensor) -> torch.Tensor: """ return self.audio_tokenizer.decode(latent).float() # [B, audio_channels, N_samples] - def _preprocess_condition_video( - self, frames: List[Any], target_h: int, target_w: int - ) -> torch.Tensor: - if not frames: - raise ValueError("Cosmos3 condition video input must contain at least one frame.") - processed = [ - self.video_processor.preprocess( - self._resize_and_center_crop_image(pil_to_rgb(frame), target_h, target_w), - height=target_h, - width=target_w, - ).squeeze(0) - for frame in frames - ] - return torch.stack(processed, dim=1).unsqueeze(0).contiguous() + def _condition_frames_to_video_tensor(self, frames: torch.Tensor) -> torch.Tensor: + """Normalize uint8 ``[T, H, W, C]`` device frames to ``[1, 3, T, H, W]``. + + Same value mapping as ``VideoProcessor.preprocess`` (``[0, 255]`` → + ``[-1, 1]``), applied to the target-resolution frames the worker + decode (``decode_video_reference_window``) retains. + """ + if frames.shape[0] < 1: + raise MediaDecodeError("Cosmos3 condition video must contain at least one frame.") + x = frames.to(torch.float32).div_(255.0).mul_(2.0).sub_(1.0) + return x.permute(3, 0, 1, 2).unsqueeze(0).contiguous() def _encode_video_tensor(self, video_tensor: torch.Tensor) -> torch.Tensor: """VAE-encode a preprocessed pixel video [1, 3, T, H, W].""" @@ -670,7 +672,9 @@ def _prepare_latents_v2v( indexes = _normalize_condition_video_latent_indexes(condition_video_latent_indexes) out_of_range = [index for index in indexes if index >= T_lat] if out_of_range: - raise ValueError( + # Mode-aware bound (num_frames may be a mode-deferred default, so + # this cannot run at coordinator preflight); client error class. + raise MediaDecodeError( "Cosmos3 condition_video_latent_indexes contains indexes outside the latent video: " f"indexes={indexes}, latent_frames={T_lat}." ) @@ -739,7 +743,7 @@ def forward( use_guardrails: bool = COSMOS3_EXTRA_SPECS["use_guardrails"].default, enable_audio: bool = COSMOS3_EXTRA_SPECS["enable_audio"].default, output_type: str = COSMOS3_EXTRA_SPECS["output_type"].default, - video: torch.Tensor | None = None, # [T, H, W, C, dtype=uint8] + video: bytes | None = None, # encoded MP4/AVI reference (V2V) condition_video_latent_indexes: Iterable[int] | None = None, condition_video_keep: str | None = None, flow_shift: Optional[float] = None, @@ -933,30 +937,65 @@ def forward( image, height=height, width=width, num_frames=num_frames, generator=generator ) elif video is not None: - condition_video_latent_indexes = _normalize_condition_video_latent_indexes( - condition_video_latent_indexes - ) - if not isinstance(video, torch.Tensor) or video.ndim != 4: - raise ValueError( - "Cosmos3 V2V reference must be a uint8 [T, H, W, C] tensor " - f"(the 'video' extra-param contract), got {type(video).__name__}" - f"{' of shape ' + str(tuple(video.shape)) if isinstance(video, torch.Tensor) else ''}." + prepare_error: Optional[Exception] = None + try: + condition_video_latent_indexes = _normalize_condition_video_latent_indexes( + condition_video_latent_indexes ) - # video is already the conditioning window: the coordinator crops it - frames = [PIL.Image.fromarray(frame.cpu().numpy()) for frame in video] - video = self._preprocess_condition_video(frames, height, width) - - if self.rank == 0: - logger.info( - f"Cosmos3 V2V conditioning: frames={video.shape[2]}, " - f"latent_indexes={condition_video_latent_indexes}" + # Bound-check the indexes against the OUTPUT latent length + # before any window math: an out-of-range index would + # otherwise size the decode ring (keep="last" decodes to EOF + # through it) from a request that is deterministically + # invalid. + num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 + out_of_range = [i for i in condition_video_latent_indexes if i >= num_latent_frames] + if out_of_range: + raise MediaDecodeError( + f"Cosmos3 condition_video_latent_indexes {out_of_range} are out " + f"of range for a {num_frames}-frame output " + f"({num_latent_frames} latent frames)." + ) + if isinstance(video, bytes): + window = _condition_pixel_frame_count( + condition_video_latent_indexes, self.vae_scale_factor_temporal + ) + frames_u8 = decode_video_reference_window( + video, + window=window, + keep=_normalize_condition_video_keep(condition_video_keep), + target_h=height, + target_w=width, + device=self.device, + ) + else: + raise MediaDecodeError( + "Cosmos3 V2V reference must be encoded MP4/AVI bytes " + f"(the 'video' extra-param contract), got " + f"{type(video).__name__}." + ) + condition_pixels = self._condition_frames_to_video_tensor(frames_u8) + del frames_u8 + + if self.rank == 0: + logger.info( + f"Cosmos3 V2V conditioning: frames={condition_pixels.shape[2]}, " + f"latent_indexes={condition_video_latent_indexes}" + ) + latents, velocity_mask, condition_latents = self._prepare_latents_v2v( + condition_pixels, + num_frames=num_frames, + generator=generator, + condition_video_latent_indexes=condition_video_latent_indexes, ) - latents, velocity_mask, condition_latents = self._prepare_latents_v2v( - video, - num_frames=num_frames, - generator=generator, - condition_video_latent_indexes=condition_video_latent_indexes, - ) + # The VAE-encoded condition latents are all the denoise loop + # needs; drop the decoded pixels before the long generation. + del condition_pixels + except Exception as exc: + prepare_error = exc + # Per-rank decode/prepare can fail non-uniformly (NVDEC init, + # corrupt stream, allocation); converge all ranks on one outcome + # before any model collective so healthy ranks cannot hang. + synchronize_media_prepare_status(prepare_error) else: latents = self._prepare_latents(height, width, num_frames, generator) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py deleted file mode 100644 index a45e5e03f73f..000000000000 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py +++ /dev/null @@ -1,18 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Shared Cosmos3 media helpers.""" - -from __future__ import annotations - -from typing import Any - -import PIL.Image - - -def pil_to_rgb(value: Any) -> PIL.Image.Image: - if isinstance(value, str): - return PIL.Image.open(value).convert("RGB") - if isinstance(value, PIL.Image.Image): - return value.convert("RGB") - raise TypeError(f"Cosmos3 preprocessing expected PIL image or image path, got {type(value)!r}.") diff --git a/tensorrt_llm/_torch/visual_gen/pipeline.py b/tensorrt_llm/_torch/visual_gen/pipeline.py index feabc7452df0..5eac636f3285 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline.py @@ -38,16 +38,9 @@ class ExtraParamSchema(StrictBaseModel): ) validator: Optional[Callable[[Any], Any]] = Field( default=None, - description="Optional value validator; raises ValueError on invalid values.", - ) - # Like ``validator``, must be a module-level function (specs are pickled to - # the coordinator in the READY handshake). - reducer: Optional[Callable[[Any, Dict[str, Any]], Any]] = Field( - default=None, - description="Optional transport reducer, run once in the coordinator " - "before the request is copied/serialized: (value, extra_params) -> " - "reduced value. Must be semantics-preserving (the worker treats " - "reduced and unreduced values identically).", + description="Optional value validator; raises ValueError on invalid " + "values. Must be a module-level function (specs are pickled to the " + "coordinator in the READY handshake).", ) diff --git a/tensorrt_llm/visual_gen/params.py b/tensorrt_llm/visual_gen/params.py index bb1663239612..874698dfc425 100644 --- a/tensorrt_llm/visual_gen/params.py +++ b/tensorrt_llm/visual_gen/params.py @@ -14,7 +14,6 @@ # limitations under the License. from typing import Any, Dict, List, Optional, Union -import torch from pydantic import Field from tensorrt_llm.llmapi.utils import StrictBaseModel, set_api_status @@ -93,7 +92,7 @@ class VisualGenParams(StrictBaseModel): "bool": (bool,), "str": (str,), "list": (list,), - "tensor": (torch.Tensor,), + "bytes": (bytes,), } # Generation config fields that pipelines declare defaults for. If a user @@ -113,39 +112,6 @@ class VisualGenParams(StrictBaseModel): ) -def reduce_visual_gen_params( - params: VisualGenParams, - extra_param_specs: Dict[str, Any], -) -> VisualGenParams: - """Apply spec-declared transport reducers to ``extra_params`` values. - - Runs once in the coordinator, before the request is deep-copied and - serialized, so oversized payloads (e.g. a full V2V reference when only - the conditioning window is consumed) shrink before they hit the copy / - ZMQ / per-rank broadcast path. Reducers are semantics-preserving by - contract — the worker behaves identically with or without them. - - Never mutates ``params``: returns it unchanged when nothing reduces, - otherwise a shallow copy carrying a new ``extra_params`` dict. - """ - if not params.extra_params: - return params - reduced: Dict[str, Any] = {} - for key, value in params.extra_params.items(): - spec = extra_param_specs.get(key) - reducer = getattr(spec, "reducer", None) if spec is not None else None - if reducer is None or value is None: - continue - new_value = reducer(value, params.extra_params) - if new_value is not value: - reduced[key] = new_value - if not reduced: - return params - out = params.model_copy() - out.extra_params = {**params.extra_params, **reduced} - return out - - def validate_visual_gen_params( params: VisualGenParams, *, diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index 554ece9958a5..f03f28738481 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -33,7 +33,6 @@ os.environ["TLLM_DISABLE_MPI"] = "1" os.environ["TRTLLM_DISABLE_COSMOS3_GUARDRAILS"] = "1" -import numpy as np import PIL.Image import pytest import torch @@ -43,7 +42,6 @@ COSMOS3_DEFAULT_CONDITION_VIDEO_LATENT_INDEXES, COSMOS3_EXTRA_SPECS, COSMOS3_T2I_PARAMS, - _crop_video_frames, _normalize_condition_video_keep, ) from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import ( @@ -56,7 +54,6 @@ _normalize_condition_video_latent_indexes, ) from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader -from tensorrt_llm.inputs.media_io import frames_to_tensor from tensorrt_llm.visual_gen.args import TorchCompileConfig, VisualGenArgs pytestmark = pytest.mark.cosmos3 @@ -262,16 +259,6 @@ def _make_test_image() -> PIL.Image.Image: return PIL.Image.new("RGB", (WIDTH, HEIGHT), color=(64, 128, 192)) -def _make_test_video( - num_frames: int = NUM_FRAMES, - *, - width: int = WIDTH, - height: int = HEIGHT, -) -> list[PIL.Image.Image]: - image = _make_test_image().resize((width, height)) - return [image.copy() for _ in range(num_frames)] - - @pytest.fixture def cosmos3_format_pipeline(): """Minimal pipeline for prompt formatting helpers (no checkpoint).""" @@ -579,70 +566,7 @@ def test_invalid_condition_video_keep_raises(self): _normalize_condition_video_keep("middle") -def _write_mp4(path, num_frames: int = 3) -> None: - """Synthesize a tiny mp4v clip (no video asset ships with the repo).""" - cv2 = pytest.importorskip("cv2") - writer = cv2.VideoWriter(str(path), cv2.VideoWriter_fourcc(*"mp4v"), 4.0, (16, 16)) - try: - for _ in range(num_frames): - writer.write(np.zeros((16, 16, 3), dtype=np.uint8)) - finally: - writer.release() - assert path.exists() and path.stat().st_size > 0 - - -class TestLoadVideoFramesTensor: - """``media_io.load_video_frames_tensor`` builds the uint8 [T, H, W, C] - tensor the ``video`` extra param carries (all frames, no crop — the worker - keeps the conditioning window). Public helper, used by the example and - available to API clients. CPU-only.""" - - def test_from_video_file_all_frames(self, tmp_path): - pytest.importorskip("cv2") - from tensorrt_llm.inputs.media_io import load_video_frames_tensor - - enc = tmp_path / "clip.mp4" - _write_mp4(enc, num_frames=8) - video = load_video_frames_tensor(enc) - assert video.dtype == torch.uint8 - assert video.ndim == 4 and video.shape[0] == 8 and video.shape[-1] == 3 - - def test_from_image_file_is_single_frame(self, tmp_path): - from tensorrt_llm.inputs.media_io import load_video_frames_tensor - - p = tmp_path / "img.png" - PIL.Image.new("RGB", (8, 8), (7, 7, 7)).save(p) - assert load_video_frames_tensor(p).shape == (1, 8, 8, 3) - - def test_from_directory(self, tmp_path): - from tensorrt_llm.inputs.media_io import load_video_frames_tensor - - for i in range(5): - PIL.Image.new("RGB", (8, 8), (i, i, i)).save(tmp_path / f"{i:03d}.png") - video = load_video_frames_tensor(tmp_path) - assert video.shape == (5, 8, 8, 3) - # Sorted lexicographically: frame k is the solid (k, k, k) image. - assert int(video[0].float().mean()) == 0 and int(video[4].float().mean()) == 4 - - def test_directory_selects_by_suffix_and_fails_loud_on_corrupt(self, tmp_path): - # Directories are user-curated: non-frame entries are ignored by name, - # but a selected frame that doesn't decode raises — corrupt frames are - # never silently dropped into a video with holes. - from tensorrt_llm.inputs.media_io import load_video_frames_tensor - - PIL.Image.new("RGB", (8, 8), (1, 1, 1)).save(tmp_path / "000.png") - (tmp_path / "notes.txt").write_text("not a frame") # ignored by suffix - assert load_video_frames_tensor(tmp_path).shape[0] == 1 - - (tmp_path / "001.png").write_bytes(b"corrupt") - with pytest.raises(PIL.UnidentifiedImageError): - load_video_frames_tensor(tmp_path) - - def test_missing_path_raises(self, tmp_path): - from tensorrt_llm.inputs.media_io import load_video_frames_tensor - - with pytest.raises(ValueError, match="does not exist"): - load_video_frames_tensor(tmp_path / "nope") +_V2V_FIXTURE_MP4 = Path(__file__).parent / "test_data" / "cosmos3_v2v_ref_9f_bframes.mp4" @pytest.mark.integration @@ -650,11 +574,14 @@ def test_missing_path_raises(self, tmp_path): @pytest.mark.high_cuda_memory class TestCosmos3V2V: def test_v2v_smoke(self, cosmos3_pipeline): - video = frames_to_tensor(_make_test_video(NUM_FRAMES)) + """The production V2V path end to end: encoded MP4 bytes (the only + ``video`` form) — each rank demuxes from memory, NVDEC-decodes the + conditioning window, resizes to the output resolution, VAE-encodes, + and generates with the V2V scheduler policy.""" result = _run_forward( cosmos3_pipeline, image=None, - video=video, + video=_V2V_FIXTURE_MP4.read_bytes(), num_frames=NUM_FRAMES, condition_video_latent_indexes=[0, 1], ) @@ -666,53 +593,29 @@ def test_v2v_smoke(self, cosmos3_pipeline): use_karras_sigmas=False, ) - def test_v2v_tensor_reference_smoke(self, cosmos3_pipeline): - """The V2V reference arrives as a decoded uint8 [T, H, W, C] tensor - (the ``video`` extra-param contract, cropped by the coordinator's - reducer in real requests); the worker VAE-encodes it, capping/padding - to the latent window in ``_prepare_latents_v2v``.""" - from tensorrt_llm.inputs.media_io import frames_to_tensor - - video = frames_to_tensor(_make_test_video(NUM_FRAMES)) - assert video.dtype == torch.uint8 and video.ndim == 4 - result = _run_forward( - cosmos3_pipeline, - image=None, - video=video, - num_frames=NUM_FRAMES, - condition_video_latent_indexes=[0, 1], - ) - _assert_valid_video(result.video, num_frames=NUM_FRAMES) - def test_v2v_keep_last_smoke(self, cosmos3_pipeline): """condition_video_keep="last" pins the tail of the input, not the head. - ``keep`` is consumed by the coordinator-side reducer - (``_crop_video_frames``), so this test composes reducer + forward the - way a real request flows. The input is longer than the conditioning - window and color-coded (dark head, bright tail); frame 0 of the output - must be a pinned VAE round-trip of the bright tail frames. + Drives the real bytes path end to end: ``keep`` is applied inside the + worker's NVDEC decode (ring buffer over the demuxed stream), exactly + as a request flows. The fixture's red channel encodes the frame index + (R = 20 + 25*i over 9 frames); with keep="last" the conditioning + window is frames 4-8, so output frame 0 must be a pinned VAE + round-trip of fixture frame 4 (R=120) — not fixture frame 0 (R=20). """ - dark = PIL.Image.new("RGB", (WIDTH, HEIGHT), (40, 40, 40)) - bright = PIL.Image.new("RGB", (WIDTH, HEIGHT), (230, 230, 230)) - # 5 = max(condition_video_latent_indexes) * 4 + 1 conditioning frames. - video = frames_to_tensor( - [dark.copy() for _ in range(NUM_FRAMES)] + [bright.copy() for _ in range(5)] - ) - video = _crop_video_frames(video, {"condition_video_keep": "last"}) - assert video.shape[0] == 5 result = _run_forward( cosmos3_pipeline, image=None, - video=video, + video=_V2V_FIXTURE_MP4.read_bytes(), num_frames=NUM_FRAMES, condition_video_latent_indexes=[0, 1], + condition_video_keep="last", ) _assert_valid_video(result.video, num_frames=NUM_FRAMES) - first_frame_mean = result.video[0, 0].float().mean().item() - assert first_frame_mean > 135, ( - f"keep='last' must condition on the bright tail frames; frame-0 mean " - f"{first_frame_mean:.1f} matches the dark head instead" + red_mean = result.video[0, 0, :, :, 0].float().mean().item() + assert red_mean > 70, ( + f"keep='last' must condition on the tail frames (R=120..220); " + f"output frame-0 red mean {red_mean:.1f} matches the head (R=20) instead" ) _assert_scheduler_config( cosmos3_pipeline, @@ -743,7 +646,7 @@ def fake_tokenize_prompt(text, max_sequence_length, use_system_prompt, system_pr with pytest.raises(StopAfterTokenize): pipeline.forward( prompt="continue", - video=frames_to_tensor(_make_test_video(5, width=16, height=16)), + video=_V2V_FIXTURE_MP4.read_bytes(), height=16, width=16, num_frames=5, @@ -768,7 +671,7 @@ def test_image_and_video_rejected(self, cosmos3_pipeline): _run_forward( cosmos3_pipeline, image=_make_test_image(), - video=frames_to_tensor(_make_test_video(5)), + video=_V2V_FIXTURE_MP4.read_bytes(), ) def test_t2i_and_video_rejected(self, cosmos3_pipeline): @@ -776,7 +679,7 @@ def test_t2i_and_video_rejected(self, cosmos3_pipeline): _run_forward( cosmos3_pipeline, image=None, - video=frames_to_tensor(_make_test_video(5)), + video=_V2V_FIXTURE_MP4.read_bytes(), output_type="image", height=T2I_HEIGHT, width=T2I_WIDTH, @@ -823,7 +726,7 @@ def test_v2v_audio_smoke(self, cosmos3_pipeline): result = _run_forward( cosmos3_pipeline, enable_audio=True, - video=frames_to_tensor(_make_test_video(NUM_FRAMES)), + video=_V2V_FIXTURE_MP4.read_bytes(), condition_video_latent_indexes=[0, 1], ) _assert_valid_video(result.video, num_frames=NUM_FRAMES) diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py index e155c017ce22..abcb784e4084 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py @@ -777,7 +777,16 @@ def _validate(extras): # Valid values pass. _validate({"condition_video_latent_indexes": [0, 1], "condition_video_keep": "last"}) - _validate({"video": torch.zeros(3, 4, 4, 3, dtype=torch.uint8)}) + # ``video`` carries encoded MP4/AVI bytes: a video signature passes, + # empty / non-video bytes are client errors, and anything that is not + # bytes (e.g. a decoded tensor) fails the type check. + _validate({"video": b"\x00\x00\x00\x18ftypmp42\x00\x00\x00\x00"}) + with pytest.raises(ValueError, match="empty"): + _validate({"video": b""}) + with pytest.raises(ValueError, match="not a recognized video container"): + _validate({"video": b"\x89PNG\r\n\x1a\n and not a video"}) + with pytest.raises(ValueError, match="expected type 'bytes'"): + _validate({"video": torch.zeros(3, 4, 4, 3, dtype=torch.uint8)}) with pytest.raises(ValueError, match="non-negative"): _validate({"condition_video_latent_indexes": [0, -1]}) @@ -794,12 +803,6 @@ def _validate(extras): with pytest.raises(ValueError, match="output_type"): _validate({"output_type": "gif"}) _validate({"output_type": "image"}) - with pytest.raises(ValueError, match=r"\[T, H, W, C\]"): - _validate({"video": torch.zeros(4, 4, 3, dtype=torch.uint8)}) # 3-D - with pytest.raises(ValueError, match="uint8"): - _validate({"video": torch.zeros(3, 4, 4, 3, dtype=torch.float32)}) - with pytest.raises(ValueError, match="CPU tensor"): - _validate({"video": torch.zeros(3, 4, 4, 3, dtype=torch.uint8, device="meta")}) def test_validator_type_errors_become_client_errors(self): """A validator raising TypeError (wrong-shaped value it didn't guard) @@ -820,71 +823,17 @@ def touchy(value): def test_spec_validators_survive_pickling(self): """Specs travel worker -> coordinator in the READY handshake (pickled - over ZMQ); validators/reducers must be module-level functions so they - serialize by reference — a lambda/closure here would crash worker - startup.""" + over ZMQ); validators must be module-level functions so they serialize + by reference — a lambda/closure here would crash worker startup.""" import pickle - import torch - from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import COSMOS3_EXTRA_SPECS specs = pickle.loads(pickle.dumps(COSMOS3_EXTRA_SPECS)) with pytest.raises(ValueError, match="first or last"): specs["condition_video_keep"].validator("middle") - reduced = specs["video"].reducer(torch.zeros(20, 4, 4, 3, dtype=torch.uint8), {}) - assert reduced.shape[0] == 5 - - def test_video_transport_reducer_crops_to_conditioning_window(self): - """The coordinator-side reducer ships only the conditioning window — - never the full clip — and the payload owns its storage (a bare slice - would pickle the entire original tensor).""" - import torch - - from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import _crop_video_frames - - # Frame k is solid value k, so window position is observable. - full = torch.arange(189, dtype=torch.uint8).view(189, 1, 1, 1).expand(189, 4, 4, 3) - full = full.contiguous() - - first = _crop_video_frames(full, {}) - assert first.shape[0] == 5 # default indexes (0, 1) -> 1*4+1 - assert int(first[0, 0, 0, 0]) == 0 and int(first[-1, 0, 0, 0]) == 4 - # Owns its storage: pickling must carry the window, not the clip. - assert first.untyped_storage().size() < full.untyped_storage().size() - - last = _crop_video_frames(full, {"condition_video_keep": "last"}) - assert last.shape[0] == 5 and int(last[-1, 0, 0, 0]) == 188 - - wider = _crop_video_frames(full, {"condition_video_latent_indexes": [0, 2]}) - assert wider.shape[0] == 9 # 2*4+1 - - # Short-enough inputs and invalid context pass through unchanged. - short = torch.zeros(3, 4, 4, 3, dtype=torch.uint8) - assert _crop_video_frames(short, {}) is short - assert _crop_video_frames(full, {"condition_video_latent_indexes": [-1]}) is full - assert _crop_video_frames("not a tensor", {}) == "not a tensor" - - def test_reduce_visual_gen_params_is_non_mutating(self): - """generate_async reduces before the deep copy; the caller's params - object and tensor must be untouched.""" - import torch - - from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import COSMOS3_EXTRA_SPECS - from tensorrt_llm.visual_gen.params import VisualGenParams, reduce_visual_gen_params - - full = torch.zeros(189, 4, 4, 3, dtype=torch.uint8) - params = VisualGenParams(extra_params={"video": full, "flow_shift": 10.0}) - out = reduce_visual_gen_params(params, COSMOS3_EXTRA_SPECS) - - assert out is not params - assert params.extra_params["video"] is full # caller untouched - assert out.extra_params["video"].shape[0] == 5 - assert out.extra_params["flow_shift"] == 10.0 # non-reduced keys intact - - # Nothing to reduce -> same object back, zero copies. - plain = VisualGenParams(extra_params={"flow_shift": 10.0}) - assert reduce_visual_gen_params(plain, COSMOS3_EXTRA_SPECS) is plain + with pytest.raises(ValueError, match="not a recognized video container"): + specs["video"].validator(b"garbage bytes") # --- unsupported universal fields --- From 71f00e97b50cdcfad26efddc54602190955b0abc Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 23 Jul 2026 21:36:53 -0700 Subject: [PATCH 45/64] Pass encoded video bytes from the Cosmos3 offline example --video_path reads the file's bytes straight into extra_params['video']; each worker decodes the conditioning window on NVDEC. Docs updated for the NVDEC path (no OpenCV install step, MP4/AVI files). Signed-off-by: Igor Shovkun --- examples/visual_gen/models/cosmos3/README.md | 4 ++-- examples/visual_gen/models/cosmos3/cosmos3.py | 11 ++++------- 2 files changed, 6 insertions(+), 9 deletions(-) diff --git a/examples/visual_gen/models/cosmos3/README.md b/examples/visual_gen/models/cosmos3/README.md index 60be9a991970..a7c5f86c4f7f 100644 --- a/examples/visual_gen/models/cosmos3/README.md +++ b/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 frame directory or `.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); `.mp4`/`.avi` decode uses OpenCV (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)). - **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 @@ -34,7 +34,7 @@ export TRTLLM_DISABLE_COSMOS3_GUARDRAILS=1 ## Media I/O dependencies - 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) uses OpenCV — the same optional decoder as the multimodal video path. It is **not** bundled with TensorRT-LLM — install it yourself: `pip install opencv-python-headless`. Frame directories work without it. +- 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. ## Deployment configs diff --git a/examples/visual_gen/models/cosmos3/cosmos3.py b/examples/visual_gen/models/cosmos3/cosmos3.py index 13afc4081f81..f0bb6b2bc8aa 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3.py +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -25,8 +25,8 @@ ``file://`` / ``http(s)://`` URL, or a ``data:`` URI. - **V2V** — video-conditioned video (``prompts/v2v.json``). Condition on the first (or last, per ``condition_video_keep``) frames of a reference video - via ``--video_path`` (a local frame directory, ``.mp4``/``.avi`` file, or - single image; ``.mp4``/``.avi`` decode requires OpenCV). + via ``--video_path`` (a local MP4/AVI file; its encoded bytes pass through + and each worker decodes the conditioning window on NVDEC). - **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). @@ -107,7 +107,6 @@ from typing import Any, Dict, Optional from tensorrt_llm import VisualGen, VisualGenArgs -from tensorrt_llm.inputs.media_io import load_video_frames_tensor _SCRIPT_DIR = Path(__file__).resolve().parent @@ -235,7 +234,7 @@ def main(): "--video_path", type=str, default=None, - help="Reference video for V2V: a local frame directory or .mp4/.avi file", + help="Reference video for V2V: a local MP4/AVI file (decoded on worker NVDEC)", ) parser.add_argument( "--output_type", type=str, default="video", help="Output type (video, image)" @@ -286,9 +285,7 @@ def main(): params.extra_params["output_type"] = output_type if args.video_path is not None: - # Decode client-side into the uint8 [T, H, W, C] tensor contract; the - # worker keeps the conditioning window and VAE-encodes. - params.extra_params["video"] = load_video_frames_tensor(args.video_path) + params.extra_params["video"] = Path(args.video_path).read_bytes() if negative_prompt is None: params.negative_prompt = None From ed441d4109b1df094b0e93e88640a38c506692c7 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 23 Jul 2026 21:36:55 -0700 Subject: [PATCH 46/64] Feed the V2V LPIPS gate from a checked-in H.264 fixture MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The synthesized-tensor reference died with the tensor intake; the gate now conditions on a 2.7 KB checked-in H.264 clip (same moving-block content, provenance in test_data/README.md — decoded YUV is normative per spec, and the YUV->RGB conversion is this stack's, so the decode is stable for this path). Golden frame regenerated once for the codec-lossy conditioning pixels. Registers test_media_decode.py in l0_b200. Signed-off-by: Igor Shovkun --- .../visual_gen_lpips_golden_media.zip | 4 +-- .../examples/visual_gen/test_data/README.md | 34 ++++++++++++++++++ .../test_data/cosmos3_v2v_lpips_reference.mp4 | Bin 0 -> 2729 bytes .../examples/visual_gen/test_visual_gen.py | 34 +++++++++--------- .../test_lists/test-db/l0_b200.yml | 1 + 5 files changed, 53 insertions(+), 20 deletions(-) create mode 100644 tests/integration/defs/examples/visual_gen/test_data/README.md create mode 100644 tests/integration/defs/examples/visual_gen/test_data/cosmos3_v2v_lpips_reference.mp4 diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zip b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zip index 1b85d11a6f67..305a209e3c62 100644 --- a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zip +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zip @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:de95c0e23c15209a3f00f6955433917a156d057473fe6c5054450f9db38243ab -size 17636502 +oid sha256:d891a0a4e25d3a46d543c71a90fa83b39488d898ad484715c2a13f8f070e66b4 +size 18921612 diff --git a/tests/integration/defs/examples/visual_gen/test_data/README.md b/tests/integration/defs/examples/visual_gen/test_data/README.md new file mode 100644 index 000000000000..cf96a72fce52 --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/test_data/README.md @@ -0,0 +1,34 @@ + + +# VisualGen integration-test input fixtures + +## `cosmos3_v2v_lpips_reference.mp4` + +5-frame 720p H.264/MP4 conditioning reference for the Cosmos3-Nano V2V LPIPS +gate (2,729 bytes). Deterministic content: gray (30, 30, 30) background with a +(200, 120, 40) block moving 40 px/frame. the H.264 spec makes decoded YUV bit-exact +for conformant decoders; the YUV->RGB conversion is this stack's +(PyNvVideoCodec), so decoded RGB is stable for this decode path. + +Regeneration (exact provenance; ffmpeg 6.1.1 / libx264): + +```python +import numpy as np +from PIL import Image + +for i in range(5): + frame = np.full((720, 1280, 3), 30, dtype=np.uint8) + x = 100 + i * 40 + frame[200:520, x : x + 200] = (200, 120, 40) + Image.fromarray(frame).save(f"frame_{i:02d}.png") +``` + +```bash +ffmpeg -y -framerate 24 -i frame_%02d.png \ + -c:v libx264 -pix_fmt yuv420p -g 4 -bf 2 \ + -x264-params b_adapt=0:scenecut=0 \ + -movflags +faststart cosmos3_v2v_lpips_reference.mp4 +``` diff --git a/tests/integration/defs/examples/visual_gen/test_data/cosmos3_v2v_lpips_reference.mp4 b/tests/integration/defs/examples/visual_gen/test_data/cosmos3_v2v_lpips_reference.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..776cb846bf07a4a139c51cf937a492c144f14d46 GIT binary patch literal 2729 zcma(S2~-s4`P*GVBt(qjfd`I)F^^-J*+m5Erpuy+#Jf*?iLtJ;Gqdb=XLp8~;aJqT z1e1_dV>KG^3Pi0Y9z0X8)&?bpq_JAjs4=MTiAOCaR!gNGPjUPGyAY|b@3sHi|Nrjq z|E_-)04S&Cm))x51<(V5xJ_&c$9Uxo2Ef31C8-1eL~n_UN7!Q-+K@N}Iv~F%>puu& zKSf)gX#O`$k8v#56xM^nd73A5r=f>px3^P0cww|#(k*z3wjw_QxBg^EAxYph6|v!h zsA?fjI8yZ6t1D)GgqL@-D5Jdm)49{gj)sD!kY$>S7nLp_P~zr=Zk)h5NK5&wj~4}^ zqucFu`<#fxm3Tu=dxn|Bhh@WgMIbt45qV6pB2r;oC##x6M0~fZsof1BI(4LJ2+8W{ zGsK7>Y3CzKHd(rXQ&B&YHkuhDLz^hk?I<&4W`u>Ib4W0VFcY=%bUcgi1Mj;;q}H zJyaUSNs2%*h$?1=&Z5kWz+^kKQIKy$h4g7NC%u%GLFLZLLrq@bP>?6deo=60l!>Oz zX(rlaMy5;CM($Z2R8CXeVeMU(sOeeP(U7A-!9ZA;Ql26T{IM%^(7KXy1w3sPg zaEKD;u`oGw4o$H>R`jcag)TGa&}9r`ro4j1C6rNWu>-O3l&t#E6i>UtTNopaMtD$O zcbUN3iAM%4>6+t7>*x^|)oC5q3}{+a;${6*NnliPBsOTBsHi4!R_HSQk?EZl;-1MK_1RbTeqm zr-Uxy++J2Afw+B|polCgA>SbuE3DtnNnV-Nkmj&P_`!1f&;=?gEYWc)tXIGbDTiJ5 zBktxyILos#v2oZPZdOGEXGaJnRw}p)U77%sIfh#5+LpPIpDuU`lTH)+JrHW<>XWIpAMv5R7X3!^bR13jBt#kvtzh^BBHLg94VC~L1w$3Yw!+ka zekjC!kiNEZc`Gpee4crx3n!Z18R?0;a=)QIw&_lWxH<9os=jgGtg;QN)%^otY!s9Z zp84*F^t>O2jHr13&i1@nS8a6wiOO?(->3C69?`Jzof@#&7Q^7zn;?4MHHi4?4(L~x zgKj~C6%HP11UBdbI(h>nU-}&4*UW)I53fLswH~4iREXU3D#V@$m=m){nAb~5{=aUi z-CmH?I%QPW>B2qL_mT!rDELc7wm5rXj(invB>(4WD8z3-bMCrC(RTQ8 zmNldwHU5jfW3%r@SME6d+s5x|too}ff?vl3o3-{qbJwolmb>X3d)7I-dPc(>|M5j* zE~b|=(Wi3jXL(~i)&888X=VK@hkkxVzass+npw*>wcabVfD?RUqleC9#Z z?$saX#SEPKcKfxeH%8s-44hi`EXNri6ui7D?Pt7g;Ankz>g42I&tm!ph={?I*5-bc zTlUtHuKbIv#kJDL$UiUZf9sFMbNB6^8n-`qSRy ztK8{4_hp*pdWr`*2OJSXK=ZF54(L;K0CslAtcJ9_x9mmPlZ;x_BZIueYb6N7W< zK63oZcDO4zvG>FJvBwU4d}z;#dH!p`&WiCF<7Q1fw0UP+=DF(Mohfz<_XdRromHDG48;Hv*+7)joUliQ#Ik7bxHjfvGZ-lhNa6F_bGog;1cnz|2F(W z-h>!@AvX_2jN8ySv--Kzug(i@cBvmo5i-n8seznz zjj(Wj!mIxpyJ%_E&i>zYRNrc8Ucg?oWp=KdQZw-8^ui?<@fr (H, W, C) for save_image. diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index bf158c5f530e..b34305248d07 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -193,6 +193,7 @@ l0_b200: # ------------- Prefix-aware scheduling E2E tests --------------- - kv_cache/test_prefix_aware_scheduling.py::TestServePrefixAwareScheduling::test_multi_round_qa_shared_prefix_smoke # ------------- Visual Gen tests --------------- + - unittest/_torch/visual_gen/test_media_decode.py - unittest/_torch/visual_gen/test_visual_gen_args.py - unittest/_torch/visual_gen/test_visual_gen_params.py - unittest/_torch/visual_gen/test_visual_gen_utils.py From a188655851e41a7df472ca44e79b3af8253665cb Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 27 Jul 2026 15:24:32 -0700 Subject: [PATCH 47/64] Document VisualGenResult failure classes and test them The class and aresult() docstrings promised RuntimeError for any single-prompt failure, but _resolved_value() raises by failure class: ValueError for client errors, VisualGenCapacityError for capacity, and RuntimeError only when unclassified. ValueError is not a RuntimeError subclass, so a caller following the docs would miss client failures entirely. Adds regression tests over all three classes plus the batch path (which keeps Option B semantics and never raises). They assert the exact type: VisualGenCapacityError subclasses RuntimeError, so a plain pytest.raises(RuntimeError) would pass for a capacity failure and leave the 400-vs-503 distinction the routes depend on untested. Signed-off-by: Igor Shovkun --- tensorrt_llm/visual_gen/visual_gen.py | 13 ++++-- tests/unittest/visual_gen/test_output.py | 55 ++++++++++++++++++++++++ 2 files changed, 65 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/visual_gen/visual_gen.py b/tensorrt_llm/visual_gen/visual_gen.py index 7528f5068c70..ecfc0a98f893 100644 --- a/tensorrt_llm/visual_gen/visual_gen.py +++ b/tensorrt_llm/visual_gen/visual_gen.py @@ -51,7 +51,12 @@ class VisualGenResult: A single instance backs both single-prompt and batch-prompt requests: - Single prompt: ``await handle`` resolves to a :class:`VisualGenOutput`. - Underlying-request failure raises :class:`RuntimeError`. + Underlying-request failure raises by failure class: + :class:`ValueError` for client errors (unusable request content — an + undecodable media reference, out-of-range conditioning), + ``VisualGenCapacityError`` for capacity failures (a valid request that + does not fit this deployment), and :class:`RuntimeError` for anything + unclassified. - Batch prompt: ``await handle`` resolves to ``List[VisualGenOutput]``. Per-item or whole-batch failure never raises; failed items carry ``error != None`` (Option B semantics). @@ -90,8 +95,10 @@ def __await__(self): async def aresult(self, timeout: Optional[float] = None): """Wait for the underlying request and return the resolved value. - For single-prompt requests, returns a :class:`VisualGenOutput`. Raises - :class:`RuntimeError` on underlying-request failure. + For single-prompt requests, returns a :class:`VisualGenOutput`. + Underlying-request failure raises :class:`ValueError` (client), + ``VisualGenCapacityError`` (capacity), or :class:`RuntimeError` + (unclassified) — see the class docstring. For batch-prompt requests, returns ``List[VisualGenOutput]``. Never raises; failed items carry ``error != None``. diff --git a/tests/unittest/visual_gen/test_output.py b/tests/unittest/visual_gen/test_output.py index e27f9c33e3b5..4e096ff153d1 100644 --- a/tests/unittest/visual_gen/test_output.py +++ b/tests/unittest/visual_gen/test_output.py @@ -785,3 +785,58 @@ def test_media_output_unimportable(): """``MediaOutput`` is not importable from any path.""" with pytest.raises(ImportError): from tensorrt_llm._torch.visual_gen.output import MediaOutput # noqa: F401 + + +# --------------------------------------------------------------------------- +# Single-prompt failure classes reach the caller as distinct exception types +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "error_type, expected", + [ + ("client", ValueError), + ("capacity", "capacity"), # resolved below; avoids a module-level import + (None, RuntimeError), + ], + ids=["client", "capacity", "unclassified"], +) +def test_single_prompt_failure_raises_by_error_class(error_type, expected): + """The worker's failure class survives to the public API. + + ``VisualGenCapacityError`` subclasses ``RuntimeError``, so these assert the + *exact* type — otherwise the unclassified case would pass for a capacity + failure and the distinction the routes rely on (400 vs 503) would be + untested. + """ + from tensorrt_llm._torch.visual_gen.media_decode import VisualGenCapacityError + from tensorrt_llm.visual_gen.visual_gen import VisualGenResult + + if expected == "capacity": + expected = VisualGenCapacityError + + resp = DiffusionResponse(request_id=20, error_msg="boom", error_type=error_type) + fx = _FakeExecutor(resp) + try: + handle = VisualGenResult(request_id=20, executor=fx, batch_size=None) + with pytest.raises(expected) as excinfo: + handle.result(timeout=5.0) + assert type(excinfo.value) is expected + assert "boom" in str(excinfo.value) + finally: + fx.stop() + + +def test_batch_failure_never_raises_regardless_of_error_class(): + """Batch handles keep Option B semantics: per-item ``error``, no raise.""" + from tensorrt_llm.visual_gen.visual_gen import VisualGenResult + + resp = DiffusionResponse(request_id=21, error_msg="boom", error_type="client") + fx = _FakeExecutor(resp) + try: + handle = VisualGenResult(request_id=21, executor=fx, batch_size=2) + outs = handle.result(timeout=5.0) + assert len(outs) == 2 + assert all(o.error is not None for o in outs) + finally: + fx.stop() From 5f8883b2b4205f0581ba219f78d0f8691988779a Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 27 Jul 2026 15:33:34 -0700 Subject: [PATCH 48/64] Parameterize the batch failure test over all error classes The test asserted Option B semantics 'regardless of error class' while only exercising the client class; run it over client, capacity, and unclassified so the name matches what it covers. Signed-off-by: Igor Shovkun --- tests/unittest/visual_gen/test_output.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/tests/unittest/visual_gen/test_output.py b/tests/unittest/visual_gen/test_output.py index 4e096ff153d1..002813f1f495 100644 --- a/tests/unittest/visual_gen/test_output.py +++ b/tests/unittest/visual_gen/test_output.py @@ -827,11 +827,20 @@ def test_single_prompt_failure_raises_by_error_class(error_type, expected): fx.stop() -def test_batch_failure_never_raises_regardless_of_error_class(): - """Batch handles keep Option B semantics: per-item ``error``, no raise.""" +@pytest.mark.parametrize( + "error_type", + ["client", "capacity", None], + ids=["client", "capacity", "unclassified"], +) +def test_batch_failure_never_raises_regardless_of_error_class(error_type): + """Batch handles keep Option B semantics for every failure class. + + Per-item ``error`` is set and nothing raises, unlike the single-prompt + path which raises by class. + """ from tensorrt_llm.visual_gen.visual_gen import VisualGenResult - resp = DiffusionResponse(request_id=21, error_msg="boom", error_type="client") + resp = DiffusionResponse(request_id=21, error_msg="boom", error_type=error_type) fx = _FakeExecutor(resp) try: handle = VisualGenResult(request_id=21, executor=fx, batch_size=2) From d3f1e7fec3d8036d3bccf8699f181c4b511ae07a Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Tue, 28 Jul 2026 21:54:10 -0700 Subject: [PATCH 49/64] Route ISO-BMFF still images away from the video decoder `ftyp` marks the ISO-BMFF family, not video: HEIF/AVIF photos share it with MP4, so a HEIC upload (the iOS camera default) was routed to the video slot and failed at NVDEC with a misleading demux error. Check the major brand and the compatible brands against the registered HEIF/AVIF/AVC-image set, and reject those at the boundary asking for PNG/JPEG. Compatible brands matter because the AVIF spec requires avif/avis there rather than as major brand. Parsing stays inside the leading ftyp box and is bounds-checked: the 64-bit largesize escape and absurd sizes degrade to the major brand instead of reading brands at shifted offsets, and a box truncated before the major brand is now unrecognized rather than a video candidate. The set is deliberately not an exhaustive image registry; unfamiliar ISO-BMFF still images remain video candidates and are rejected by the decoder, which is safe. Signed-off-by: Igor Shovkun --- examples/visual_gen/serve/README.md | 2 +- tensorrt_llm/inputs/media_io.py | 86 ++++++++++++++++++- tensorrt_llm/serve/openai_protocol.py | 3 + tensorrt_llm/serve/visual_gen_utils.py | 12 ++- .../visual_gen/test_visual_gen_utils.py | 59 +++++++++++++ 5 files changed, 158 insertions(+), 4 deletions(-) diff --git a/examples/visual_gen/serve/README.md b/examples/visual_gen/serve/README.md index 7eff00f35d81..03cd358da68c 100644 --- a/examples/visual_gen/serve/README.md +++ b/examples/visual_gen/serve/README.md @@ -287,7 +287,7 @@ You can customize these by: - `num_frames`: when set, wins over the `seconds * frame_rate` derivation - `seed`, `num_inference_steps`, `guidance_scale`, `max_sequence_length`, `negative_prompt`: per-request denoise controls - `input_reference`: Reference image (I2V/TI2V) or video (V2V), routed by container signature — filename and content type are ignored; accepted as base64-encoded string in JSON or as a file in multipart form-data. Images must fully decode at the boundary; video bytes pass through and decode on the workers' NVDEC (PyNvVideoCodec, a declared dependency). Unrecognized or undecodable content returns HTTP 400. - - **Supported formats**: reference images are tested with PNG (`image/png`) and JPEG (`image/jpeg`); they must fully decode with Pillow at the boundary. Reference videos are routed by container signature (MP4 `video/mp4`, AVI `video/x-msvideo`) and decoded on the workers' NVDEC — tested codec is H.264 in both containers; other codecs/profiles depend on the GPU's decoder capabilities and are best-effort. Filename and MIME metadata are never used for routing. Corrupt or undecodable content returns 400; a valid reference the deployment cannot fit returns 503. Reference decoding is bounded to 7200 frames by default (override with `TRTLLM_MAX_REFERENCE_DECODE_FRAMES`, `0` disables). + - **Supported formats**: reference images are tested with PNG (`image/png`) and JPEG (`image/jpeg`); they must fully decode with Pillow at the boundary. Reference videos are routed by container signature (MP4 `video/mp4`, AVI `video/x-msvideo`) and decoded on the workers' NVDEC — tested codec is H.264 in both containers; other codecs/profiles depend on the GPU's decoder capabilities and are best-effort. Filename and MIME metadata are never used for routing. HEIF/AVIF still images (which share the ISO-BMFF `ftyp` signature with MP4) are detected and rejected with a 400 asking for PNG or JPEG, rather than being sent to the video decoder. Corrupt or undecodable content returns 400; a valid reference the deployment cannot fit returns 503. Reference decoding is bounded to 7200 frames by default (override with `TRTLLM_MAX_REFERENCE_DECODE_FRAMES`, `0` disables). - `extra_params`: model-specific overflow (see below) - `response_format`: `"b64_json"` or `"url"` - `format`: Generation content encoding. Video encoders: `"mp4"`, `"avi"`, `"auto"`. Tensor formats: `"safetensors"`, `"pt"` (carries video + audio + scalar metadata in one payload for LTX-2). diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index df6c01f8dedb..188237b6e4dc 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -356,16 +356,98 @@ def _get_cv2(): # with what actually decodes. +# ISO-BMFF brands that mark a still image or image sequence (HEIF/AVIF/AVC +# image). `ftyp` identifies the ISO-BMFF *family*, not video — HEIC photos +# (the iOS camera default) share it with MP4 — so these are split out and +# rejected at the boundary instead of being sent to a video decoder that +# cannot use them. +# +# Deliberately NOT an exhaustive image registry: an unfamiliar ISO-BMFF +# still-image brand simply stays a video candidate and gets a 400 from the +# worker's decoder, which is safe. Registered names per MP4RA +# (https://mp4ra.org/registered-types/brands). +_ISOBMFF_IMAGE_BRANDS = frozenset( + { + b"mif1", + b"mif2", + b"msf1", # MIAF image / image sequence + b"heic", + b"heix", + b"heim", + b"heis", # HEVC image + b"hevc", + b"hevx", + b"hevm", + b"hevs", # HEVC image sequence + b"avif", + b"avis", # AVIF image / sequence + b"avci", + b"avcs", # AVC image / sequence + } +) + +# Enough bytes for the whole `ftyp` box in practice: 16 header bytes plus a +# compatible-brands list, which realistically holds a handful of entries. +_FTYP_SCAN_BYTES = 64 + + +def _isobmff_brands(data) -> frozenset: + """Brands declared by a leading ``ftyp`` box: major + compatible. + + Layout: ``size`` [0:4], ``'ftyp'`` [4:8], ``major_brand`` [8:12], + ``minor_version`` [12:16], then four-byte compatible brands to the end of + the box. Bounds-checked against hostile input — a declared size that is + too small, truncated, or uses the 64-bit ``largesize`` escape degrades to + the major brand alone rather than reading brands at the wrong offset. + """ + head = bytes(data[:_FTYP_SCAN_BYTES]) + if len(head) < 12: + return frozenset() + brands = {head[8:12]} + + size = int.from_bytes(head[0:4], "big") + # size 0 -> box runs to EOF; size 1 -> 64-bit largesize shifts every + # following field, so do not guess at brand offsets. + if size == 1 or 0 < size < 16: + return frozenset(brands) + end = len(head) if size == 0 else min(size, len(head)) + for off in range(16, end - 3, 4): + brands.add(head[off : off + 4]) + return frozenset(brands) + + +def is_isobmff_image_bytes(data) -> bool: + """True when the payload is an ISO-BMFF still image (HEIF/AVIF/AVC image). + + Lets the boundary say "this format is unsupported, convert it" instead of + "this file is corrupt" — the payload is perfectly valid, we just do not + decode it. + """ + if bytes(data[4:8]) != b"ftyp": + return False + return bool(_isobmff_brands(data) & _ISOBMFF_IMAGE_BRANDS) + + def sniff_media_kind(data) -> Optional[str]: """Classify a reference payload by container signature. - Returns ``"image"`` (PNG/JPEG), ``"video"`` (ISO-BMFF/MP4 family or AVI), - or ``None`` for anything unrecognized. + Returns ``"image"`` (PNG/JPEG, or an ISO-BMFF still-image brand such as + HEIF/AVIF), ``"video"`` (other ISO-BMFF/MP4 family, or AVI), or ``None`` + for anything unrecognized. """ header = bytes(data[:12]) if header.startswith(b"\x89PNG\r\n\x1a\n") or header.startswith(b"\xff\xd8\xff"): return "image" if header[4:8] == b"ftyp": + brands = _isobmff_brands(data) + if not brands: + # `ftyp` present but truncated before the major brand: nothing to + # classify on, so reject here rather than sending it to a decoder. + return None + # AVIF requires `avif`/`avis` among the *compatible* brands, so the + # major brand alone is not sufficient to identify still images. + if brands & _ISOBMFF_IMAGE_BRANDS: + return "image" return "video" if header.startswith(b"RIFF") and header[8:12] == b"AVI ": return "video" diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index feee9a192692..69ac70217ba6 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -1697,6 +1697,9 @@ class VideoGenerationRequest(OpenAIBaseModel): "are decoded on the workers' NVDEC; tested codec is H.264 in " "both containers, other codecs/profiles depend on the GPU's " "decoder capabilities and are best-effort. Undecodable or " + "HEIF/AVIF still images share the ISO-BMFF signature with " + "MP4 and are rejected with a 400 asking for PNG/JPEG rather " + "than routed to the video decoder. Undecodable or " "corrupt content fails as a client error (400); a valid " "reference the deployment cannot fit fails as capacity " "(503). Reference decoding is bounded to 7200 frames by " diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index da3803485073..ace91c299cdc 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -3,7 +3,11 @@ import os from typing import Any, Dict, List, Optional -from tensorrt_llm.inputs.media_io import is_decodable_image_bytes, sniff_media_kind +from tensorrt_llm.inputs.media_io import ( + is_decodable_image_bytes, + is_isobmff_image_bytes, + sniff_media_kind, +) from tensorrt_llm.logger import logger from tensorrt_llm.serve.openai_protocol import ImageGenerationRequest, VideoGenerationRequest from tensorrt_llm.visual_gen import VisualGen, VisualGenParams @@ -178,6 +182,12 @@ def parse_visual_gen_params( # check, so a truncated PNG 400s here instead of 500ing at # the worker's load. if not is_decodable_image_bytes(payload): + if is_isobmff_image_bytes(payload): + raise ValueError( + "input_reference is a HEIF/AVIF image, which is not " + "a supported reference format; convert it to PNG or " + "JPEG." + ) raise ValueError( "input_reference has an image container signature but " "does not fully decode; the file may be truncated or " diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index 7a4e0d83d98d..538c1e850844 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -455,6 +455,65 @@ def test_sniff_media_kind(self): # RIFF alone is not AVI (e.g. WAV audio is RIFF too). assert sniff_media_kind(b"RIFF\x00\x00\x00\x00WAVEfmt ") is None + @staticmethod + def _ftyp(major: bytes, compatible: tuple = (), *, size: int = None) -> bytes: + """Build an ISO-BMFF `ftyp` box: size|'ftyp'|major|minor|compatible*.""" + body = major + b"\x00\x00\x00\x00" + b"".join(compatible) + declared = len(body) + 8 if size is None else size + return declared.to_bytes(4, "big") + b"ftyp" + body + + def test_sniff_isobmff_still_images_are_not_video(self): + """`ftyp` marks the ISO-BMFF family, not video: HEIF/AVIF photos share + it with MP4. Routing them to the video slot would hand a still image to + NVDEC; they must classify as image instead.""" + from tensorrt_llm.inputs.media_io import sniff_media_kind + + # HEIC as major brand (the iOS camera default). + assert sniff_media_kind(self._ftyp(b"heic", (b"mif1", b"heic"))) == "image" + # `mif1` major with `heic` compatible. + assert sniff_media_kind(self._ftyp(b"mif1", (b"heic",))) == "image" + # AVIF: the spec requires avif/avis among the COMPATIBLE brands, so a + # major-brand-only check would miss this one. + assert sniff_media_kind(self._ftyp(b"iso8", (b"avif", b"mif1"))) == "image" + assert sniff_media_kind(self._ftyp(b"avis", (b"avif",))) == "image" + # HEVC image sequence brands. + for brand in (b"heix", b"heim", b"heis", b"hevc", b"hevx", b"hevm", b"hevs"): + assert sniff_media_kind(self._ftyp(brand)) == "image", brand + # Ordinary video ISO-BMFF stays video. + assert sniff_media_kind(self._ftyp(b"mp42", (b"isom", b"mp42"))) == "video" + assert sniff_media_kind(self._ftyp(b"isom", (b"iso2", b"avc1"))) == "video" + assert sniff_media_kind(self._ftyp(b"qt ")) == "video" + + def test_sniff_ftyp_bounds_are_checked(self): + """Hostile/degenerate `ftyp` boxes must not read brands at the wrong + offset; they degrade to the major brand rather than misclassifying.""" + from tensorrt_llm.inputs.media_io import sniff_media_kind + + # Declared size larger than the payload (truncated file). + assert sniff_media_kind(self._ftyp(b"heic", (b"mif1",), size=4096)) == "image" + # Declared size 0 means "box runs to EOF". + assert sniff_media_kind(self._ftyp(b"iso8", (b"avif",), size=0)) == "image" + # Nonsense small size: fall back to the major brand only. + assert sniff_media_kind(self._ftyp(b"heic", (b"mif1",), size=8)) == "image" + assert sniff_media_kind(self._ftyp(b"mp42", (b"avif",), size=8)) == "video" + # size == 1 is the 64-bit largesize escape: every later field shifts, + # so compatible brands must be ignored, not read at bogus offsets. + assert sniff_media_kind(self._ftyp(b"mp42", (b"avif",), size=1)) == "video" + assert sniff_media_kind(self._ftyp(b"heic", (), size=1)) == "image" + # Header shorter than a brand. + assert sniff_media_kind(b"\x00\x00\x00\x18ftyp") is None + + def test_heif_reference_rejected_with_actionable_message(self): + """A HEIC upload is a valid file we simply do not support — the 400 + must say so, not claim the file is corrupt.""" + generator = _StubVisualGen() + heic = self._ftyp(b"heic", (b"mif1", b"heic")) + b"\x00" * 64 + request = VideoGenerationRequest( + prompt="x", input_reference=base64.b64encode(heic).decode() + ) + with pytest.raises(ValueError, match="HEIF/AVIF"): + parse_visual_gen_params(request, "vid-heic", generator, media_storage_path=None) + def test_truncated_image_bytes_are_not_decodable(self): # A truncated PNG still opens (the header parses) but cannot decode # its pixels; the probe must reject it so the boundary 400s instead From 3a235d35669c17cd62f39974f5052add983fe11a Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 29 Jul 2026 09:23:10 -0700 Subject: [PATCH 50/64] Parse the ISO-BMFF ftyp box instead of peeking at a fixed window The sniffer read 64 bytes and classified on whatever it found there, so an image brand later in a long compatible-brand list was missed and the payload went to the video decoder. It also took the major brand from [8:12] unconditionally, which is the high half of the largesize when the box uses the 64-bit size escape. Scan the declared box instead, and return None -- unrecognized, a 400 -- whenever the box cannot be read in full, rather than guessing "video" off a partial scan. Enforce the FileTypeBox minimums (a major brand plus the mandatory minor_version: 16 bytes, 24 with largesize) and reject a compatible-brand area that does not hold whole four-byte brands. The declared size is client-controlled, so the scan is capped: without a ceiling it costs O(payload), 1.6 s of interpreter time for a 64 MB buffer, on the serving event loop. Reject HEIF/AVIF on the signature alone rather than after Pillow fails to decode. Whether Pillow reads them depends on optional plugins, and the worker is a separate process that need not have the same ones, so a deployment with pillow-heif installed would have accepted a HEIC and written it to disk as an I2V reference. One test fixture declared a 24-byte box while supplying 16; it had been classified as video by rounding down to what parsed. Signed-off-by: Igor Shovkun --- tensorrt_llm/inputs/media_io.py | 74 ++++++++++++------- tensorrt_llm/serve/openai_protocol.py | 8 +- tensorrt_llm/serve/visual_gen_utils.py | 15 ++-- .../visual_gen/test_visual_gen_params.py | 2 +- .../visual_gen/test_visual_gen_utils.py | 73 +++++++++++++++--- 5 files changed, 124 insertions(+), 48 deletions(-) diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index 188237b6e4dc..4a578ab5f13b 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -386,33 +386,54 @@ def _get_cv2(): } ) -# Enough bytes for the whole `ftyp` box in practice: 16 header bytes plus a -# compatible-brands list, which realistically holds a handful of entries. -_FTYP_SCAN_BYTES = 64 +# Work bound, not a format rule. The declared box size is client-controlled, +# so without a ceiling the scan below costs O(payload): 1.6 s of interpreter +# time for a 64 MB buffer, on the serving event loop. This admits 1020 +# compatible brands against a registry holding a few hundred in total, so no +# encoder output comes near it; real boxes are tens of bytes. +_MAX_FTYP_BOX_BYTES = 4096 -def _isobmff_brands(data) -> frozenset: - """Brands declared by a leading ``ftyp`` box: major + compatible. +def _isobmff_brands(data) -> Optional[frozenset]: + """Brands declared by a leading ``ftyp`` box, or ``None`` if undeterminable. - Layout: ``size`` [0:4], ``'ftyp'`` [4:8], ``major_brand`` [8:12], - ``minor_version`` [12:16], then four-byte compatible brands to the end of - the box. Bounds-checked against hostile input — a declared size that is - too small, truncated, or uses the 64-bit ``largesize`` escape degrades to - the major brand alone rather than reading brands at the wrong offset. + Layout per ISO/IEC 14496-12 — ``size`` [0:4], ``'ftyp'`` [4:8], then:: + + normal major [8:12] minor [12:16] compatible from 16 + size == 1 largesize [8:16], major [16:20], minor [20:24], compat from 24 + + ``None`` means "do not classify" and callers must treat the payload as + unrecognized rather than assuming video: the box is malformed, extends + past the payload, or exceeds the scan bound, and a partial scan could + miss an image brand late in the compatible-brand list. """ - head = bytes(data[:_FTYP_SCAN_BYTES]) - if len(head) < 12: - return frozenset() - brands = {head[8:12]} - - size = int.from_bytes(head[0:4], "big") - # size 0 -> box runs to EOF; size 1 -> 64-bit largesize shifts every - # following field, so do not guess at brand offsets. - if size == 1 or 0 < size < 16: - return frozenset(brands) - end = len(head) if size == 0 else min(size, len(head)) - for off in range(16, end - 3, 4): - brands.add(head[off : off + 4]) + buf = bytes(data[:_MAX_FTYP_BOX_BYTES]) + if len(buf) < 16: + return None + + size = int.from_bytes(buf[0:4], "big") + if size == 1: + # Extended size: the 64-bit largesize occupies 8:16, shifting the + # brands. Reading 8:12 as the major brand here is what makes naive + # parsers misclassify. + box_end, brands_at = int.from_bytes(buf[8:16], "big"), 16 + elif size == 0: + box_end, brands_at = len(data), 8 # box runs to EOF + else: + box_end, brands_at = size, 8 + + if box_end > len(data) or box_end > _MAX_FTYP_BOX_BYTES: + return None # truncated, or past the work bound — refuse to guess + # `FileTypeBox` is a major brand plus a mandatory `minor_version`, + # followed by whole compatible brands. A short or misaligned box is + # malformed; reject it rather than rounding down to what parses. + if box_end < brands_at + 8 or (box_end - brands_at) % 4: + return None + + brands = {buf[brands_at : brands_at + 4]} + # Skip the 4-byte minor_version between the major and compatible brands. + for off in range(brands_at + 8, box_end - 3, 4): + brands.add(buf[off : off + 4]) return frozenset(brands) @@ -425,7 +446,8 @@ def is_isobmff_image_bytes(data) -> bool: """ if bytes(data[4:8]) != b"ftyp": return False - return bool(_isobmff_brands(data) & _ISOBMFF_IMAGE_BRANDS) + brands = _isobmff_brands(data) + return bool(brands and brands & _ISOBMFF_IMAGE_BRANDS) def sniff_media_kind(data) -> Optional[str]: @@ -441,8 +463,8 @@ def sniff_media_kind(data) -> Optional[str]: if header[4:8] == b"ftyp": brands = _isobmff_brands(data) if not brands: - # `ftyp` present but truncated before the major brand: nothing to - # classify on, so reject here rather than sending it to a decoder. + # Box unreadable (truncated / malformed / past the work bound): + # classify nothing rather than defaulting to video on a partial scan. return None # AVIF requires `avif`/`avis` among the *compatible* brands, so the # major brand alone is not sufficient to identify still images. diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index 69ac70217ba6..f39293a1ddcb 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -1696,10 +1696,10 @@ class VideoGenerationRequest(OpenAIBaseModel): "video/mp4, AVI video/x-msvideo) pass through encoded and " "are decoded on the workers' NVDEC; tested codec is H.264 in " "both containers, other codecs/profiles depend on the GPU's " - "decoder capabilities and are best-effort. Undecodable or " - "HEIF/AVIF still images share the ISO-BMFF signature with " - "MP4 and are rejected with a 400 asking for PNG/JPEG rather " - "than routed to the video decoder. Undecodable or " + "decoder capabilities and are best-effort. HEIF/AVIF still " + "images share the ISO-BMFF signature with MP4 and are " + "rejected with a 400 asking for PNG/JPEG rather than routed " + "to the video decoder. Undecodable or " "corrupt content fails as a client error (400); a valid " "reference the deployment cannot fit fails as capacity " "(503). Reference decoding is bounded to 7200 frames by " diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index ace91c299cdc..6d043db614d9 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -178,16 +178,19 @@ def parse_visual_gen_params( payload = _read_reference_payload(request.input_reference) kind = sniff_media_kind(payload) if kind == "image": + # Rejected on signature alone, not on a failed decode: + # whether Pillow reads HEIF/AVIF depends on optional plugins, + # and the worker process need not have the same ones. + if is_isobmff_image_bytes(payload): + raise ValueError( + "input_reference is a HEIF/AVIF image, which is not " + "a supported reference format; convert it to PNG or " + "JPEG." + ) # Signature routes; the full decode is still the acceptance # check, so a truncated PNG 400s here instead of 500ing at # the worker's load. if not is_decodable_image_bytes(payload): - if is_isobmff_image_bytes(payload): - raise ValueError( - "input_reference is a HEIF/AVIF image, which is not " - "a supported reference format; convert it to PNG or " - "JPEG." - ) raise ValueError( "input_reference has an image container signature but " "does not fully decode; the file may be truncated or " diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py index abcb784e4084..31d657e84160 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py @@ -780,7 +780,7 @@ def _validate(extras): # ``video`` carries encoded MP4/AVI bytes: a video signature passes, # empty / non-video bytes are client errors, and anything that is not # bytes (e.g. a decoded tensor) fails the type check. - _validate({"video": b"\x00\x00\x00\x18ftypmp42\x00\x00\x00\x00"}) + _validate({"video": b"\x00\x00\x00\x18ftypmp42\x00\x00\x00\x00mp42isom"}) with pytest.raises(ValueError, match="empty"): _validate({"video": b""}) with pytest.raises(ValueError, match="not a recognized video container"): diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index 538c1e850844..78b9a3239ceb 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -462,6 +462,17 @@ def _ftyp(major: bytes, compatible: tuple = (), *, size: int = None) -> bytes: declared = len(body) + 8 if size is None else size return declared.to_bytes(4, "big") + b"ftyp" + body + @staticmethod + def _ftyp_ext(major: bytes, compatible: tuple = (), *, largesize: int = None) -> bytes: + """Build a 64-bit `ftyp` box: 1|'ftyp'|largesize|major|minor|compat*. + + ``size == 1`` inserts an 8-byte largesize between the type and the + brands, so the major brand lives at [16:20], not [8:12]. + """ + body = major + b"\x00\x00\x00\x00" + b"".join(compatible) + declared = len(body) + 16 if largesize is None else largesize + return (1).to_bytes(4, "big") + b"ftyp" + declared.to_bytes(8, "big") + body + def test_sniff_isobmff_still_images_are_not_video(self): """`ftyp` marks the ISO-BMFF family, not video: HEIF/AVIF photos share it with MP4. Routing them to the video slot would hand a still image to @@ -484,22 +495,46 @@ def test_sniff_isobmff_still_images_are_not_video(self): assert sniff_media_kind(self._ftyp(b"isom", (b"iso2", b"avc1"))) == "video" assert sniff_media_kind(self._ftyp(b"qt ")) == "video" + def test_sniff_ftyp_reads_the_whole_declared_box(self): + """A brand list longer than a fixed peek window must still be seen — + an `avif` brand at the tail decides image-vs-video.""" + from tensorrt_llm.inputs.media_io import sniff_media_kind + + # 21 compatible brands = a 100-byte box; `avif` sits past byte 64. + padded = tuple([b"free"] * 20 + [b"avif"]) + box = self._ftyp(b"isom", padded) + assert len(box) > 64 and box.index(b"avif") > 64 + assert sniff_media_kind(box) == "image" + + def test_sniff_ftyp_extended_size_shifts_the_brands(self): + """`size == 1` puts a 64-bit largesize at [8:16]; a parser that reads + the major brand at [8:12] sees half of that integer instead.""" + from tensorrt_llm.inputs.media_io import sniff_media_kind + + assert sniff_media_kind(self._ftyp_ext(b"heic", (b"mif1",))) == "image" + assert sniff_media_kind(self._ftyp_ext(b"isom", (b"avif",))) == "image" + assert sniff_media_kind(self._ftyp_ext(b"mp42", (b"isom",))) == "video" + def test_sniff_ftyp_bounds_are_checked(self): - """Hostile/degenerate `ftyp` boxes must not read brands at the wrong - offset; they degrade to the major brand rather than misclassifying.""" + """An `ftyp` box we cannot read in full is unclassifiable: guessing + "video" off a partial scan is how a HEIC reaches NVDEC.""" from tensorrt_llm.inputs.media_io import sniff_media_kind # Declared size larger than the payload (truncated file). - assert sniff_media_kind(self._ftyp(b"heic", (b"mif1",), size=4096)) == "image" - # Declared size 0 means "box runs to EOF". + assert sniff_media_kind(self._ftyp(b"heic", (b"mif1",), size=4096)) is None + # Declared size 0 means "box runs to EOF" — readable, so classify. assert sniff_media_kind(self._ftyp(b"iso8", (b"avif",), size=0)) == "image" - # Nonsense small size: fall back to the major brand only. - assert sniff_media_kind(self._ftyp(b"heic", (b"mif1",), size=8)) == "image" - assert sniff_media_kind(self._ftyp(b"mp42", (b"avif",), size=8)) == "video" - # size == 1 is the 64-bit largesize escape: every later field shifts, - # so compatible brands must be ignored, not read at bogus offsets. - assert sniff_media_kind(self._ftyp(b"mp42", (b"avif",), size=1)) == "video" - assert sniff_media_kind(self._ftyp(b"heic", (), size=1)) == "image" + # Nonsense small size: no room for even a major brand. + assert sniff_media_kind(self._ftyp(b"heic", (b"mif1",), size=8)) is None + assert sniff_media_kind(self._ftyp(b"mp42", (b"avif",), size=8)) is None + # `minor_version` is mandatory, so a box is at least 16 bytes (24 with + # the largesize escape); stopping at the major brand is malformed. + assert sniff_media_kind(self._ftyp(b"heic", (b"mif1",), size=12)) is None + assert sniff_media_kind(self._ftyp_ext(b"heic", (b"mif1",), largesize=20)) is None + # Compatible-brand area must hold whole four-byte brands. + assert sniff_media_kind(self._ftyp(b"mp42", (b"isom",), size=18) + b"\x00\x00") is None + # Past the scan bound: bounded rather than scanned. + assert sniff_media_kind(self._ftyp(b"mp42", tuple([b"free"] * 2000))) is None # Header shorter than a brand. assert sniff_media_kind(b"\x00\x00\x00\x18ftyp") is None @@ -514,6 +549,22 @@ def test_heif_reference_rejected_with_actionable_message(self): with pytest.raises(ValueError, match="HEIF/AVIF"): parse_visual_gen_params(request, "vid-heic", generator, media_storage_path=None) + def test_heif_rejected_even_when_pillow_can_decode_it(self, monkeypatch): + """Acceptance must not hinge on optional Pillow plugins: a deployment + with pillow-heif installed would otherwise admit a HEIC that the + worker — a separate process, possibly without the plugin — cannot + read, and the documented contract says HEIF/AVIF is a 400.""" + import tensorrt_llm.serve.visual_gen_utils as serve_utils + + monkeypatch.setattr(serve_utils, "is_decodable_image_bytes", lambda _: True) + generator = _StubVisualGen() + heic = self._ftyp(b"heic", (b"mif1", b"heic")) + b"\x00" * 64 + request = VideoGenerationRequest( + prompt="x", input_reference=base64.b64encode(heic).decode() + ) + with pytest.raises(ValueError, match="HEIF/AVIF"): + parse_visual_gen_params(request, "vid-heic-plugin", generator, media_storage_path=None) + def test_truncated_image_bytes_are_not_decodable(self): # A truncated PNG still opens (the header parses) but cannot decode # its pixels; the probe must reject it so the boundary 400s instead From 457794384a38d16a9007617330b3ba2290e5e2bb Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 29 Jul 2026 12:12:12 -0700 Subject: [PATCH 51/64] Report VisualGen failures with built-in exception types MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MediaDecodeError and VisualGenCapacityError were a VisualGen-specific taxonomy for something the built-ins already express, and the capacity class was documented in the public VisualGenResult contract while living in tensorrt_llm._torch — so no caller could actually catch it without reaching into a private module. Drop both. The failure class is now carried by type alone: ValueError for unusable request content (400), MemoryError for a valid request that does not fit the deployment (503), RuntimeError for anything unclassified (500), with the detail in the message. torch.cuda.OutOfMemoryError is still named explicitly because it derives from RuntimeError rather than MemoryError. The wire stays as it was -- DiffusionResponse.error_type is a string, not a type -- so the executor and the rank-convergence protocol are unchanged apart from which exception they reconstruct. Classify the I2V image load too. Its OSError family (UnidentifiedImageError for a bad header, plain OSError partway through a truncated file, FileNotFoundError for a missing path) would otherwise report a bad upload as a server fault. Route it through the same convergence the V2V branch uses: every rank loads the reference independently, so a rank that failed while the others entered the transformer's collectives would hang the job. The 503 route had no test at all, so add one, plus its 400 counterpart -- both asserting the message survives, which is the point of carrying detail in a primitive rather than in a type. Signed-off-by: Igor Shovkun --- tensorrt_llm/_torch/visual_gen/executor.py | 4 +- .../_torch/visual_gen/media_decode.py | 47 +++++--------- .../models/cosmos3/pipeline_cosmos3.py | 63 +++++++++++++------ tensorrt_llm/serve/openai_video_routes.py | 3 +- tensorrt_llm/visual_gen/visual_gen.py | 43 ++++++------- .../visual_gen/test_cosmos3_pipeline.py | 37 +++++++++++ .../_torch/visual_gen/test_media_decode.py | 18 +++--- .../visual_gen/test_trtllm_serve_endpoints.py | 43 +++++++++++++ tests/unittest/visual_gen/test_output.py | 15 ++--- 9 files changed, 178 insertions(+), 95 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/executor.py b/tensorrt_llm/_torch/visual_gen/executor.py index 1d4580bdd6ee..673a7435561d 100644 --- a/tensorrt_llm/_torch/visual_gen/executor.py +++ b/tensorrt_llm/_torch/visual_gen/executor.py @@ -257,8 +257,8 @@ class DiffusionResponse: error_type: Failure class when ``error_msg`` is set: ``"client"`` (unusable request content → 400 / ``ValueError``), ``"capacity"`` (valid request does not fit the deployment → 503 / - ``RuntimeError``), or ``None`` for unclassified runtime failures - (500). + ``MemoryError``), or ``None`` for unclassified runtime failures + (500 / ``RuntimeError``). generation: Wall-clock time the executor measured around the engine's inference call (host ``time.perf_counter()``), in seconds. Default ``0.0`` so the dataclass round-trips through diff --git a/tensorrt_llm/_torch/visual_gen/media_decode.py b/tensorrt_llm/_torch/visual_gen/media_decode.py index a9f9e810424b..57bf9982d45d 100644 --- a/tensorrt_llm/_torch/visual_gen/media_decode.py +++ b/tensorrt_llm/_torch/visual_gen/media_decode.py @@ -31,33 +31,20 @@ import torch -class MediaDecodeError(ValueError): - """Client-class failure: the reference itself is unusable. - - Corrupt/undecodable content, unsupported codec, zero frames, or the - decode-work limit. Maps to HTTP 400 / a plain ``ValueError`` at the - public Python boundary. - """ - - -class VisualGenCapacityError(RuntimeError): - """Capacity-class failure: a valid request does not fit this deployment. - - CUDA/NVDEC allocation or session-init failure. Maps to HTTP 503; never - a client error — the input is not malformed. - """ - - def classify_worker_error(exc: BaseException) -> str | None: """Failure class for the response channel: "client", "capacity", or None. - Only the two dedicated classes are mapped — a bare ``ValueError`` from a - model bug must stay an unclassified runtime failure, not become a 400. + Keyed off built-in exception types rather than a VisualGen-specific + hierarchy: ``ValueError`` means the request's content was unusable + (400), ``MemoryError`` means a valid request did not fit (503), and + anything else is an unclassified runtime failure (500). Detail travels + in the message. ``torch.cuda.OutOfMemoryError`` is spelled out because + it derives from ``RuntimeError``, not ``MemoryError``. """ - if isinstance(exc, MediaDecodeError): - return "client" - if isinstance(exc, (VisualGenCapacityError, torch.cuda.OutOfMemoryError)): + if isinstance(exc, (MemoryError, torch.cuda.OutOfMemoryError)): return "capacity" + if isinstance(exc, ValueError): + return "client" return None @@ -99,9 +86,9 @@ def synchronize_media_prepare_status(exc: Exception | None) -> None: kind, message = payload[0] message = f"[rank {failing_rank}] {message}" if kind == "client": - raise MediaDecodeError(message) + raise ValueError(message) if kind == "capacity": - raise VisualGenCapacityError(message) + raise MemoryError(message) raise RuntimeError(message) @@ -231,7 +218,7 @@ def decode_video_reference_window( frame_cap = max_reference_decode_frames() if frame_cap is not None and keep == "first" and window > frame_cap: - raise MediaDecodeError( + raise ValueError( f"Conditioning window of {window} frames exceeds the reference " f"decode limit of {frame_cap} (TRTLLM_MAX_REFERENCE_DECODE_FRAMES)." ) @@ -253,7 +240,7 @@ def _read(buf: bytearray) -> int: # readable stream — a content problem, not a capacity one. demuxer = nvc.CreateDemuxer(_read) except nvc.PyNvVCException as exc: - raise MediaDecodeError( + raise ValueError( f"Video reference could not be demuxed (corrupt or not a " f"supported container): {exc}" ) from exc @@ -283,7 +270,7 @@ def _read(buf: bytearray) -> int: for packet in demuxer: for frame in decoder.Decode(packet): if frame_cap is not None and count >= frame_cap: - raise MediaDecodeError( + raise ValueError( f"Video reference exceeds the decode limit of " f"{frame_cap} frames " f"(TRTLLM_MAX_REFERENCE_DECODE_FRAMES); trim the " @@ -301,18 +288,18 @@ def _read(buf: bytearray) -> int: if keep == "first" and count >= window: break except torch.cuda.OutOfMemoryError as exc: - raise VisualGenCapacityError( + raise MemoryError( f"Out of device memory while decoding the video reference " f"({window} frames @ {target_w}x{target_h} retained): {exc}" ) from exc except nvc.PyNvVCException as exc: - raise MediaDecodeError( + raise ValueError( f"Video reference failed to decode (corrupt or unsupported " f"stream for this deployment's decoder): {exc}" ) from exc if count == 0: - raise MediaDecodeError( + raise ValueError( "Video reference contains no decodable frames; the payload " "may be corrupt or use an unsupported codec." ) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index 9050e77c8a61..91d2fd913ae3 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -27,7 +27,6 @@ from transformers import Qwen2Tokenizer from tensorrt_llm._torch.visual_gen.media_decode import ( - MediaDecodeError, decode_video_reference_window, synchronize_media_prepare_status, ) @@ -73,6 +72,24 @@ def _condition_pixel_frame_count( return max(condition_video_latent_indexes) * int(temporal_compression) + 1 +def _load_reference_image(path: str): + """Load an I2V reference, reporting unreadable content as a client error. + + The worker's load is the acceptance check — the serve boundary only + routes on the container signature — so PIL's ``OSError`` + (``UnidentifiedImageError`` for a bad header, plain ``OSError`` partway + through a truncated file) has to become a ``ValueError`` here, or a bad + upload would be reported as a server fault. + """ + try: + return load_image(path, format="pil") + except OSError as exc: + raise ValueError( + f"Image reference could not be decoded; it may be truncated, " + f"corrupt, or in an unsupported format: {exc}" + ) from exc + + @register_pipeline( "Cosmos3OmniMoTPipeline", hf_ids=[ @@ -609,7 +626,7 @@ def _condition_frames_to_video_tensor(self, frames: torch.Tensor) -> torch.Tenso decode (``decode_video_reference_window``) retains. """ if frames.shape[0] < 1: - raise MediaDecodeError("Cosmos3 condition video must contain at least one frame.") + raise ValueError("Cosmos3 condition video must contain at least one frame.") x = frames.to(torch.float32).div_(255.0).mul_(2.0).sub_(1.0) return x.permute(3, 0, 1, 2).unsqueeze(0).contiguous() @@ -674,7 +691,7 @@ def _prepare_latents_v2v( if out_of_range: # Mode-aware bound (num_frames may be a mode-deferred default, so # this cannot run at coordinator preflight); client error class. - raise MediaDecodeError( + raise ValueError( "Cosmos3 condition_video_latent_indexes contains indexes outside the latent video: " f"indexes={indexes}, latent_frames={T_lat}." ) @@ -921,21 +938,29 @@ def forward( velocity_mask = None if image is not None: - if isinstance(image, str): - image = load_image(image, format="pil") - - if isinstance(image, PIL.Image.Image): - image = image.convert("RGB") - image = self._resize_and_center_crop_image(image, height=height, width=width) - image = self.video_processor.preprocess( - image, - height=height, - width=width, - ) + prepare_error: Optional[Exception] = None + try: + if isinstance(image, str): + image = _load_reference_image(image) + + if isinstance(image, PIL.Image.Image): + image = image.convert("RGB") + image = self._resize_and_center_crop_image(image, height=height, width=width) + image = self.video_processor.preprocess( + image, + height=height, + width=width, + ) - latents, velocity_mask, image_latent = self._prepare_latents_i2v( - image, height=height, width=width, num_frames=num_frames, generator=generator - ) + latents, velocity_mask, image_latent = self._prepare_latents_i2v( + image, height=height, width=width, num_frames=num_frames, generator=generator + ) + except Exception as exc: + prepare_error = exc + # Same convergence as the V2V branch: every rank loads the image + # independently, so a rank that failed while others entered the + # transformer's collectives would hang the job. + synchronize_media_prepare_status(prepare_error) elif video is not None: prepare_error: Optional[Exception] = None try: @@ -950,7 +975,7 @@ def forward( num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 out_of_range = [i for i in condition_video_latent_indexes if i >= num_latent_frames] if out_of_range: - raise MediaDecodeError( + raise ValueError( f"Cosmos3 condition_video_latent_indexes {out_of_range} are out " f"of range for a {num_frames}-frame output " f"({num_latent_frames} latent frames)." @@ -968,7 +993,7 @@ def forward( device=self.device, ) else: - raise MediaDecodeError( + raise ValueError( "Cosmos3 V2V reference must be encoded MP4/AVI bytes " f"(the 'video' extra-param contract), got " f"{type(video).__name__}." diff --git a/tensorrt_llm/serve/openai_video_routes.py b/tensorrt_llm/serve/openai_video_routes.py index 49689ff1f946..3601c6f9624b 100644 --- a/tensorrt_llm/serve/openai_video_routes.py +++ b/tensorrt_llm/serve/openai_video_routes.py @@ -24,7 +24,6 @@ from fastapi.responses import FileResponse, JSONResponse, Response from pydantic import ValidationError -from tensorrt_llm._torch.visual_gen.media_decode import VisualGenCapacityError from tensorrt_llm.logger import logger from tensorrt_llm.media.encoding import resolve_video_format from tensorrt_llm.media.tensor_payload import is_tensor_format @@ -125,7 +124,7 @@ async def openai_video_generation_sync(self, raw_request: Request) -> Response: except ValueError as exc: logger.error(f"Video request error: {exc}") return self.create_error_response(str(exc), status_code=HTTPStatus.BAD_REQUEST) - except VisualGenCapacityError as exc: + except MemoryError as exc: # Valid request that does not fit this deployment (decode / # allocation capacity) — a server condition, not client error. logger.error(f"Video request capacity error: {exc}") diff --git a/tensorrt_llm/visual_gen/visual_gen.py b/tensorrt_llm/visual_gen/visual_gen.py index ecfc0a98f893..f8a604379a33 100644 --- a/tensorrt_llm/visual_gen/visual_gen.py +++ b/tensorrt_llm/visual_gen/visual_gen.py @@ -51,12 +51,12 @@ class VisualGenResult: A single instance backs both single-prompt and batch-prompt requests: - Single prompt: ``await handle`` resolves to a :class:`VisualGenOutput`. - Underlying-request failure raises by failure class: - :class:`ValueError` for client errors (unusable request content — an - undecodable media reference, out-of-range conditioning), - ``VisualGenCapacityError`` for capacity failures (a valid request that - does not fit this deployment), and :class:`RuntimeError` for anything - unclassified. + Underlying-request failure raises a built-in exception carrying the + detail in its message: :class:`ValueError` for client errors (unusable + request content — an undecodable media reference, out-of-range + conditioning), :class:`MemoryError` for capacity failures (a valid + request that does not fit this deployment), and :class:`RuntimeError` + for anything unclassified. - Batch prompt: ``await handle`` resolves to ``List[VisualGenOutput]``. Per-item or whole-batch failure never raises; failed items carry ``error != None`` (Option B semantics). @@ -97,7 +97,7 @@ async def aresult(self, timeout: Optional[float] = None): For single-prompt requests, returns a :class:`VisualGenOutput`. Underlying-request failure raises :class:`ValueError` (client), - ``VisualGenCapacityError`` (capacity), or :class:`RuntimeError` + :class:`MemoryError` (capacity), or :class:`RuntimeError` (unclassified) — see the class docstring. For batch-prompt requests, returns ``List[VisualGenOutput]``. Never @@ -171,14 +171,13 @@ def _build_resolved(self, response: "DiffusionResponse"): return split_visual_gen_output(response, self._batch_size) def _resolved_value(self): - # For single prompts, surface engine-side failure typed by class: - # worker-classified client errors (unusable reference content, - # conditioning bounds) raise ``ValueError`` — uniform with the - # synchronous parameter validation at ``generate_async`` entry — - # and capacity failures raise ``RuntimeError`` - # (``VisualGenCapacityError``), like any unclassified runtime - # failure. For batches, return the list as-is so callers iterate - # per-item ``error``. + # For single prompts, surface engine-side failure as a built-in + # exception carrying the detail in its message: ``ValueError`` for + # client errors (unusable reference content, conditioning bounds) — + # uniform with the synchronous parameter validation at + # ``generate_async`` entry — ``MemoryError`` for capacity, and + # ``RuntimeError`` for anything unclassified. For batches, return the + # list as-is so callers iterate per-item ``error``. if self._batch_size is None and isinstance(self._resolved, VisualGenOutput): if self._resolved.error is not None: message = f"Generation failed: {self._resolved.error}" @@ -186,9 +185,7 @@ def _resolved_value(self): if error_type == "client": raise ValueError(message) if error_type == "capacity": - from tensorrt_llm._torch.visual_gen.media_decode import VisualGenCapacityError - - raise VisualGenCapacityError(message) + raise MemoryError(message) raise RuntimeError(message) return self._resolved @@ -347,9 +344,13 @@ def generate( prompts, ``List[VisualGenOutput]`` of the same length. Raises: - RuntimeError: Single-prompt path on underlying-request failure. - The batch path never raises on per-item or whole-batch - failure; failed items carry ``error != None``. + ValueError: Single-prompt path, client-class failure (unusable + request content). + MemoryError: Single-prompt path, capacity failure (a valid + request that does not fit this deployment). + RuntimeError: Single-prompt path, any unclassified failure. The + batch path never raises on per-item or whole-batch failure; + failed items carry ``error != None``. NotImplementedError: ``params`` is a list (per-item parameters are not yet supported). """ diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index f03f28738481..3cb6a2d40ed5 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -51,6 +51,7 @@ COSMOS3_IMAGE_RESOLUTION_TEMPLATE, Cosmos3OmniMoTPipeline, _condition_pixel_frame_count, + _load_reference_image, _normalize_condition_video_latent_indexes, ) from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader @@ -566,6 +567,42 @@ def test_invalid_condition_video_keep_raises(self): _normalize_condition_video_keep("middle") +class TestReferenceImageLoad: + """The worker's image load is the acceptance check for an I2V reference. + + The serve boundary only routes on the container signature, so unreadable + content has to surface here as a client error (``ValueError`` → 400) and + not as a server fault. + """ + + def test_truncated_image_is_a_client_error(self, tmp_path): + # Incompressible content, so half the file is genuinely half the image. + noise = PIL.Image.frombytes("RGB", (64, 64), os.urandom(64 * 64 * 3)) + whole = tmp_path / "whole.png" + noise.save(whole, format="PNG") + data = whole.read_bytes() + path = tmp_path / "truncated.png" + path.write_bytes(data[: len(data) // 2]) + + with pytest.raises(ValueError, match="could not be decoded"): + _load_reference_image(str(path)) + + def test_unidentifiable_content_is_a_client_error(self, tmp_path): + path = tmp_path / "notreally.png" + path.write_bytes(b"not an image at all") + with pytest.raises(ValueError, match="could not be decoded"): + _load_reference_image(str(path)) + + def test_missing_file_is_a_client_error(self, tmp_path): + with pytest.raises(ValueError, match="could not be decoded"): + _load_reference_image(str(tmp_path / "nope.png")) + + def test_valid_image_loads(self, tmp_path): + path = tmp_path / "ok.png" + PIL.Image.new("RGB", (8, 8), (1, 2, 3)).save(path, format="PNG") + assert _load_reference_image(str(path)).size == (8, 8) + + _V2V_FIXTURE_MP4 = Path(__file__).parent / "test_data" / "cosmos3_v2v_ref_9f_bframes.mp4" diff --git a/tests/unittest/_torch/visual_gen/test_media_decode.py b/tests/unittest/_torch/visual_gen/test_media_decode.py index e5cda550d66c..ac1153ebe84d 100644 --- a/tests/unittest/_torch/visual_gen/test_media_decode.py +++ b/tests/unittest/_torch/visual_gen/test_media_decode.py @@ -24,7 +24,6 @@ from PIL import Image from tensorrt_llm._torch.visual_gen.media_decode import ( - MediaDecodeError, _lanczos_taps, decode_video_reference_window, max_reference_decode_frames, @@ -116,20 +115,17 @@ def _status_protocol_rank(rank: int, world_size: int, init_file: str, results_di import torch.distributed as dist - from tensorrt_llm._torch.visual_gen.media_decode import ( - MediaDecodeError, - synchronize_media_prepare_status, - ) + from tensorrt_llm._torch.visual_gen.media_decode import synchronize_media_prepare_status dist.init_process_group( "gloo", init_method=f"file://{init_file}", rank=rank, world_size=world_size ) try: - local_error = MediaDecodeError("rank-local decode failure") if rank == 1 else None + local_error = ValueError("rank-local decode failure") if rank == 1 else None try: synchronize_media_prepare_status(local_error) outcome = "no-error" - except MediaDecodeError as exc: + except ValueError as exc: outcome = f"client:{exc}" except Exception as exc: # pragma: no cover - diagnostic path outcome = f"unexpected:{type(exc).__name__}:{exc}" @@ -140,8 +136,8 @@ def _status_protocol_rank(rank: int, world_size: int, init_file: str, results_di class TestPrepareStatusProtocol: def test_local_failure_is_reraised_without_group(self): - err = MediaDecodeError("boom") - with pytest.raises(MediaDecodeError, match="boom"): + err = ValueError("boom") + with pytest.raises(ValueError, match="boom"): synchronize_media_prepare_status(err) def test_success_passes_through_without_group(self): @@ -243,12 +239,12 @@ def test_window_longer_than_clip_returns_all(self): def test_corrupt_bytes_with_valid_magic_is_client_error(self): payload = b"\x00\x00\x00\x18ftypmp42" + b"\x00" * 64 - with pytest.raises(MediaDecodeError): + with pytest.raises(ValueError): self._decode(payload) def test_frame_limit_trips_on_emitted_frames(self, monkeypatch): monkeypatch.setenv("TRTLLM_MAX_REFERENCE_DECODE_FRAMES", "4") - with pytest.raises(MediaDecodeError, match="decode limit"): + with pytest.raises(ValueError, match="decode limit"): self._decode(_MP4.read_bytes(), keep="last") def test_frame_limit_disabled(self, monkeypatch): diff --git a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py index a905b4ca4084..69d17ea523ef 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py @@ -134,6 +134,7 @@ def __init__( should_fail: bool = False, batch_aware: bool = True, validation_error: Optional[ValueError] = None, + generate_error: Optional[BaseException] = None, ): from types import SimpleNamespace @@ -143,6 +144,9 @@ def __init__( self._should_fail = should_fail self._batch_aware = batch_aware self._validation_error = validation_error + # Raised out of generate(): models an engine-side failure class, + # where validation_error models a coordinator preflight rejection. + self._generate_error = generate_error self._healthy = True self._req_counter = 0 # Captured arguments of the most recent generate / generate_async call, @@ -187,6 +191,8 @@ def generate(self, inputs=None, params=None) -> VisualGenOutput: self.last_params = params if self._validation_error is not None: raise self._validation_error + if self._generate_error is not None: + raise self._generate_error if self._should_fail: raise RuntimeError("Generation intentionally failed") n = getattr(params, "num_images_per_prompt", 1) if params else 1 @@ -913,6 +919,43 @@ def test_sync_video_null_output(self, tmp_path): assert resp.status_code == 500 os.environ.pop("TRTLLM_MEDIA_STORAGE_PATH", None) + def test_sync_video_capacity_failure_is_503(self, tmp_path, monkeypatch): + """A valid request the deployment cannot fit is retryable (503). + + The engine signals capacity with ``MemoryError`` — a built-in, so the + error contract carries no VisualGen-specific exception type — and the + route must not fold it into the generic 500. + """ + gen = MockVisualGen( + generate_error=MemoryError("Out of device memory while preparing generation") + ) + monkeypatch.setenv("TRTLLM_MEDIA_STORAGE_PATH", str(tmp_path)) + client = _create_server(gen) + resp = client.post( + "/v1/videos/generations", + json={"prompt": "big", "size": "64x64", "seconds": 1.0, "fps": 8}, + headers={"content-type": "application/json"}, + ) + assert resp.status_code == 503 + assert "Out of device memory" in resp.text + + def test_sync_video_client_failure_is_400(self, tmp_path, monkeypatch): + """Engine-side client errors stay 400, distinct from capacity. + + The detail rides in the message — that is what a primitive exception + type buys instead of a taxonomy — so it must reach the client. + """ + gen = MockVisualGen(generate_error=ValueError("reference has no decodable frames")) + monkeypatch.setenv("TRTLLM_MEDIA_STORAGE_PATH", str(tmp_path)) + client = _create_server(gen) + resp = client.post( + "/v1/videos/generations", + json={"prompt": "bad ref", "size": "64x64", "seconds": 1.0, "fps": 8}, + headers={"content-type": "application/json"}, + ) + assert resp.status_code == 400 + assert "no decodable frames" in resp.text + def test_sync_video_unsupported_content_type(self, video_client): resp = video_client.post( "/v1/videos/generations", diff --git a/tests/unittest/visual_gen/test_output.py b/tests/unittest/visual_gen/test_output.py index 002813f1f495..a4e0540d624e 100644 --- a/tests/unittest/visual_gen/test_output.py +++ b/tests/unittest/visual_gen/test_output.py @@ -796,25 +796,20 @@ def test_media_output_unimportable(): "error_type, expected", [ ("client", ValueError), - ("capacity", "capacity"), # resolved below; avoids a module-level import + ("capacity", MemoryError), (None, RuntimeError), ], ids=["client", "capacity", "unclassified"], ) def test_single_prompt_failure_raises_by_error_class(error_type, expected): - """The worker's failure class survives to the public API. + """The worker's failure class survives to the public API as a built-in. - ``VisualGenCapacityError`` subclasses ``RuntimeError``, so these assert the - *exact* type — otherwise the unclassified case would pass for a capacity - failure and the distinction the routes rely on (400 vs 503) would be - untested. + Asserts the *exact* type: the routes map these three to 400 / 503 / 500, + so a subclass relationship creeping back in would silently collapse two + of the cases into one. """ - from tensorrt_llm._torch.visual_gen.media_decode import VisualGenCapacityError from tensorrt_llm.visual_gen.visual_gen import VisualGenResult - if expected == "capacity": - expected = VisualGenCapacityError - resp = DiffusionResponse(request_id=20, error_msg="boom", error_type=error_type) fx = _FakeExecutor(resp) try: From 84eca2a2cb12b8707dd18e55cae7339b0a996560 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 29 Jul 2026 12:22:23 -0700 Subject: [PATCH 52/64] Stop decoding reference images at the serve boundary The boundary ran a full pixel decode (Image.open().load()) to decide whether to accept an image reference. That put unbounded, client-controlled CPU on an async request handler, and duplicated work the worker repeats anyway. It was also asymmetric: video was never decoded at the boundary -- routed by signature, accepted by NVDEC. Drop the probe. The container signature only routes; for both modalities the worker's decoder is what accepts, and it now reports a decode failure as a client error, so a corrupt upload is still a 400 rather than a 500. The HEIF/AVIF check stays -- it reads the ftyp brands, not pixels, and is what lets the boundary say "convert this to PNG or JPEG" instead of "this file is corrupt". Narrow the documented contract to match. An unrecognized container is a 400 at the boundary; how a decoder failure past it is reported is pipeline-specific, and Cosmos3's 400/503 classification is stated as Cosmos3's. Signed-off-by: Igor Shovkun --- examples/visual_gen/serve/README.md | 4 +- tensorrt_llm/inputs/media_io.py | 25 ++------ tensorrt_llm/serve/openai_protocol.py | 37 +++++------ tensorrt_llm/serve/visual_gen_utils.py | 15 +---- .../visual_gen/test_visual_gen_utils.py | 64 ++++--------------- 5 files changed, 40 insertions(+), 105 deletions(-) diff --git a/examples/visual_gen/serve/README.md b/examples/visual_gen/serve/README.md index 03cd358da68c..a1c6628fd06a 100644 --- a/examples/visual_gen/serve/README.md +++ b/examples/visual_gen/serve/README.md @@ -286,8 +286,8 @@ You can customize these by: - `frame_rate` (canonical) or `fps` (alias): frames per second - `num_frames`: when set, wins over the `seconds * frame_rate` derivation - `seed`, `num_inference_steps`, `guidance_scale`, `max_sequence_length`, `negative_prompt`: per-request denoise controls -- `input_reference`: Reference image (I2V/TI2V) or video (V2V), routed by container signature — filename and content type are ignored; accepted as base64-encoded string in JSON or as a file in multipart form-data. Images must fully decode at the boundary; video bytes pass through and decode on the workers' NVDEC (PyNvVideoCodec, a declared dependency). Unrecognized or undecodable content returns HTTP 400. - - **Supported formats**: reference images are tested with PNG (`image/png`) and JPEG (`image/jpeg`); they must fully decode with Pillow at the boundary. Reference videos are routed by container signature (MP4 `video/mp4`, AVI `video/x-msvideo`) and decoded on the workers' NVDEC — tested codec is H.264 in both containers; other codecs/profiles depend on the GPU's decoder capabilities and are best-effort. Filename and MIME metadata are never used for routing. HEIF/AVIF still images (which share the ISO-BMFF `ftyp` signature with MP4) are detected and rejected with a 400 asking for PNG or JPEG, rather than being sent to the video decoder. Corrupt or undecodable content returns 400; a valid reference the deployment cannot fit returns 503. Reference decoding is bounded to 7200 frames by default (override with `TRTLLM_MAX_REFERENCE_DECODE_FRAMES`, `0` disables). +- `input_reference`: Reference image (I2V/TI2V) or video (V2V), routed by container signature — filename and content type are ignored; accepted as base64-encoded string in JSON or as a file in multipart form-data. The signature only routes; the bytes pass through and the worker's decoder is what accepts them — Pillow for images, NVDEC for video (PyNvVideoCodec, a declared dependency). An unrecognized container is rejected with HTTP 400 at the boundary. + - **Supported formats**: reference images are tested with PNG (`image/png`) and JPEG (`image/jpeg`) and are decoded on the workers with Pillow. Reference videos are routed by container signature (MP4 `video/mp4`, AVI `video/x-msvideo`) and decoded on the workers' NVDEC — tested codec is H.264 in both containers; other codecs/profiles depend on the GPU's decoder capabilities and are best-effort. Filename and MIME metadata are never used for routing. HEIF/AVIF still images (which share the ISO-BMFF `ftyp` signature with MP4) are detected and rejected with a 400 asking for PNG or JPEG, rather than being sent to the video decoder. How a decoder failure past the boundary is reported is pipeline-specific: Cosmos3 classifies corrupt or undecodable references as client errors (400) and device-memory exhaustion as capacity (503). Reference decoding is bounded to 7200 frames by default (override with `TRTLLM_MAX_REFERENCE_DECODE_FRAMES`, `0` disables). - `extra_params`: model-specific overflow (see below) - `response_format`: `"b64_json"` or `"url"` - `format`: Generation content encoding. Video encoders: `"mp4"`, `"avi"`, `"auto"`. Tensor formats: `"safetensors"`, `"pt"` (carries video + audio + scalar metadata in one payload for LTX-2). diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index 4a578ab5f13b..76a748c3fe42 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -350,10 +350,10 @@ def _get_cv2(): # The serve boundary routes a metadata-free reference payload (JSON base64 # carries no filename or MIME type; multipart metadata is client-typed) to the # image or video slot by its container signature. Routing only, never -# acceptance: an image is accepted by a full PIL decode at the boundary, a -# video by the worker's decoder — so corrupt content behind a valid signature -# still fails cleanly as a client error, and the signature can never disagree -# with what actually decodes. +# acceptance: for both modalities the worker's decoder is what accepts, and +# its failure is reported as a client error. So the signature can never +# disagree with what actually decodes, and the boundary never spends a full +# decode — on an async request path — to answer a routing question. # ISO-BMFF brands that mark a still image or image sequence (HEIF/AVIF/AVC @@ -476,23 +476,6 @@ def sniff_media_kind(data) -> Optional[str]: return None -def is_decodable_image_bytes(data) -> bool: - """True when ``data`` holds still-image content PIL can fully decode. - - Strict on purpose: ``Image.open`` is lazy, so this also decodes the - pixels (``load``) — a truncated file passes a header-only probe and - would then 500 at the worker's load instead of 400ing at the boundary. - """ - try: - with Image.open(BytesIO(data)) as image: - image.load() - return True - except OSError: - # ``UnidentifiedImageError`` (bad header) subclasses ``OSError``; - # truncated files raise plain ``OSError`` from ``load``. - return False - - # Longest video, in frames, a client may request as *output* (the serve's # ``num_frames`` cap in ``openai_protocol``) and the most reference frames a # worker will decode from a video reference before raising. diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index f39293a1ddcb..16abbde0efb3 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -1687,24 +1687,25 @@ class VideoGenerationRequest(OpenAIBaseModel): input_reference: Optional[Union[str, UploadFile]] = Field( default=None, description=( - "Optional image or video reference that guides generation. " - "Content is routed by its container signature — filename and " - "MIME metadata are ignored (JSON requests carry bare base64 " - "with no such metadata). Images (PNG image/png, JPEG " - "image/jpeg) must fully decode with Pillow at the boundary " - "and condition image-to-video. Video containers (MP4 " - "video/mp4, AVI video/x-msvideo) pass through encoded and " - "are decoded on the workers' NVDEC; tested codec is H.264 in " - "both containers, other codecs/profiles depend on the GPU's " - "decoder capabilities and are best-effort. HEIF/AVIF still " - "images share the ISO-BMFF signature with MP4 and are " - "rejected with a 400 asking for PNG/JPEG rather than routed " - "to the video decoder. Undecodable or " - "corrupt content fails as a client error (400); a valid " - "reference the deployment cannot fit fails as capacity " - "(503). Reference decoding is bounded to 7200 frames by " - "default (TRTLLM_MAX_REFERENCE_DECODE_FRAMES). JSON requests " - "carry base64 bytes; multipart requests upload the file."), + "Optional image or video reference that guides generation. Content " + "is routed by its container signature — filename and MIME metadata " + "are ignored (JSON requests carry bare base64 with no such " + "metadata). Images (PNG image/png, JPEG image/jpeg) condition " + "image-to-video. Video containers (MP4 video/mp4, AVI " + "video/x-msvideo) pass through encoded and are decoded on the " + "workers' NVDEC; tested codec is H.264 in both containers, other " + "codecs/profiles depend on the GPU's decoder capabilities and are " + "best-effort. The signature only routes: the worker's decoder is " + "what accepts a reference, for both modalities. HEIF/AVIF still " + "images share the ISO-BMFF signature with MP4 and are rejected " + "with a 400 asking for PNG/JPEG rather than routed to the video " + "decoder. An unrecognized container is a 400 at the boundary; how " + "a decoder failure past it is reported is pipeline-specific — " + "Cosmos3 reports undecodable or corrupt references as client " + "errors (400) and device-memory exhaustion as capacity (503). " + "Reference decoding is bounded to 7200 frames by default " + "(TRTLLM_MAX_REFERENCE_DECODE_FRAMES). JSON requests carry base64 " + "bytes; multipart requests upload the file."), ) # Resolution diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index 6d043db614d9..31dee1f07d6d 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -3,11 +3,7 @@ import os from typing import Any, Dict, List, Optional -from tensorrt_llm.inputs.media_io import ( - is_decodable_image_bytes, - is_isobmff_image_bytes, - sniff_media_kind, -) +from tensorrt_llm.inputs.media_io import is_isobmff_image_bytes, sniff_media_kind from tensorrt_llm.logger import logger from tensorrt_llm.serve.openai_protocol import ImageGenerationRequest, VideoGenerationRequest from tensorrt_llm.visual_gen import VisualGen, VisualGenParams @@ -187,15 +183,6 @@ def parse_visual_gen_params( "a supported reference format; convert it to PNG or " "JPEG." ) - # Signature routes; the full decode is still the acceptance - # check, so a truncated PNG 400s here instead of 500ing at - # the worker's load. - if not is_decodable_image_bytes(payload): - raise ValueError( - "input_reference has an image container signature but " - "does not fully decode; the file may be truncated or " - "corrupt." - ) # I2V: the stored image file is the cross-model contract. # every I2V pipeline reads ``params.image`` as a path. if media_storage_path is None: diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index 78b9a3239ceb..d6bf720b12f2 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -424,20 +424,7 @@ def read(self, *args, **kwargs): class TestMediaBytesProbes: - """The in-memory probe/decode primitives the serve boundary runs on.""" - - def test_is_decodable_image_bytes(self): - from tensorrt_llm.inputs.media_io import is_decodable_image_bytes - - # PNG and JPEG are the two image formats in the documented support - # contract; both must probe as decodable. - for fmt in ("PNG", "JPEG"): - buf = BytesIO() - Image.new("RGB", (4, 4), (1, 2, 3)).save(buf, format=fmt) - assert is_decodable_image_bytes(buf.getvalue()), fmt - assert not is_decodable_image_bytes(b"definitely not an image") - # Video bytes are not an image (mp4 has no PIL-openable header). - assert not is_decodable_image_bytes(TestInputReferenceMaterialization._mp4_bytes()) + """The in-memory signature probes the serve boundary routes on.""" def test_sniff_media_kind(self): from tensorrt_llm.inputs.media_io import sniff_media_kind @@ -549,53 +536,30 @@ def test_heif_reference_rejected_with_actionable_message(self): with pytest.raises(ValueError, match="HEIF/AVIF"): parse_visual_gen_params(request, "vid-heic", generator, media_storage_path=None) - def test_heif_rejected_even_when_pillow_can_decode_it(self, monkeypatch): - """Acceptance must not hinge on optional Pillow plugins: a deployment - with pillow-heif installed would otherwise admit a HEIC that the - worker — a separate process, possibly without the plugin — cannot - read, and the documented contract says HEIF/AVIF is a 400.""" - import tensorrt_llm.serve.visual_gen_utils as serve_utils - - monkeypatch.setattr(serve_utils, "is_decodable_image_bytes", lambda _: True) - generator = _StubVisualGen() - heic = self._ftyp(b"heic", (b"mif1", b"heic")) + b"\x00" * 64 - request = VideoGenerationRequest( - prompt="x", input_reference=base64.b64encode(heic).decode() - ) - with pytest.raises(ValueError, match="HEIF/AVIF"): - parse_visual_gen_params(request, "vid-heic-plugin", generator, media_storage_path=None) + def test_truncated_image_reference_is_routed_not_decoded(self, tmp_path): + """The boundary routes on signature and never decodes. - def test_truncated_image_bytes_are_not_decodable(self): - # A truncated PNG still opens (the header parses) but cannot decode - # its pixels; the probe must reject it so the boundary 400s instead - # of the worker 500ing at load time. + A truncated PNG reaches the image slot; the worker's load is what + rejects it (as a client error). Decoding here would put unbounded, + client-controlled CPU on an async request path and duplicate work + the worker repeats anyway. + """ rng_pixels = np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8) buf = BytesIO() Image.fromarray(rng_pixels).save(buf, format="PNG") whole = buf.getvalue() truncated = whole[: len(whole) // 2] - Image.open(BytesIO(truncated)) # sanity: header-only open succeeds - - from tensorrt_llm.inputs.media_io import is_decodable_image_bytes - - assert is_decodable_image_bytes(whole) - assert not is_decodable_image_bytes(truncated) - - def test_truncated_image_reference_rejected_at_parse(self): - # End of the chain: a truncated image upload sniffs as an image but - # fails the strict decode -> client error at the boundary (never - # routed into the worker). - rng_pixels = np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8) - buf = BytesIO() - Image.fromarray(rng_pixels).save(buf, format="PNG") - truncated = buf.getvalue()[: len(buf.getvalue()) // 2] + with pytest.raises(OSError): + Image.open(BytesIO(truncated)).load() # sanity: it really is broken generator = _StubVisualGen() request = VideoGenerationRequest( prompt="x", input_reference=base64.b64encode(truncated).decode() ) - with pytest.raises(ValueError, match="does not fully decode"): - parse_visual_gen_params(request, "vid-12", generator, media_storage_path=None) + params = parse_visual_gen_params( + request, "vid-12", generator, media_storage_path=str(tmp_path) + ) + assert Path(params.image).read_bytes() == truncated # ============================================================================= From a91438e75c119f00daad7a6fa74c2c482e7b1abf Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 30 Jul 2026 09:44:33 -0700 Subject: [PATCH 53/64] Model the sampling policy in the V2V flow-shift stub forward() consults sampling.validate_request and is_distilled before the flow-shift block, so a stub carrying only set_flow_shift never reaches the code the test is about. Signed-off-by: Igor Shovkun --- .../visual_gen/test_cosmos3_pipeline.py | 22 +++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index 4bc2690e5f54..ee9dfa62142c 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -672,16 +672,30 @@ def test_v2v_flow_shift_override_request_path(self): class StopAfterTokenize(Exception): pass - def fake_set_flow_shift(scheduler, target, *, use_karras_sigmas=None): - calls.append((target, use_karras_sigmas)) - return scheduler + class FakeSampling: + """Minimal Cosmos3SamplingPolicy stand-in recording flow-shift calls. + + forward() consults the policy for request validation and the + distilled guard before it ever reaches the flow-shift block, so a + stub carrying only ``set_flow_shift`` never gets there. + """ + + is_distilled = False + checkpoint_flow_shift = 1.0 + + def validate_request(self, num_inference_steps, guidance_scale): + return None + + def set_flow_shift(self, scheduler, target, *, use_karras_sigmas=None): + calls.append((target, use_karras_sigmas)) + return scheduler def fake_tokenize_prompt(text, max_sequence_length, use_system_prompt, system_prompt=None): token_calls.append((text, max_sequence_length, use_system_prompt, system_prompt)) raise StopAfterTokenize pipeline.scheduler = SimpleNamespace(config=SimpleNamespace(flow_shift=1.0)) - pipeline.sampling = SimpleNamespace(set_flow_shift=fake_set_flow_shift) + pipeline.sampling = FakeSampling() pipeline._tokenize_prompt = fake_tokenize_prompt with pytest.raises(StopAfterTokenize): From 51799b4dd6019cdd7cc3366ec590126a2a8dfdb1 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 30 Jul 2026 09:46:21 -0700 Subject: [PATCH 54/64] Drop code the merge left dead infer() resolved height/width/steps/guidance into locals it never used: this branch's infer() passes req.params.* through and forward() does the same mode-aware resolution, for direct forward() callers too. is_t2i was also derived twice in forward() after output_type normalization. Signed-off-by: Igor Shovkun --- .../visual_gen/models/cosmos3/pipeline_cosmos3.py | 15 +-------------- .../_torch/visual_gen/models/cosmos3/sampling.py | 4 +--- 2 files changed, 2 insertions(+), 17 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index bc942bb5ef3e..4bd1b77557c4 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -17,7 +17,7 @@ import math import os import time -from typing import Any, Iterable, List, Optional, Union +from typing import Iterable, List, Optional, Union import PIL.Image import torch @@ -263,18 +263,6 @@ def _run_warmup(self, height: int, width: int, num_frames: int, steps: int) -> N def infer(self, req): extra_params = req.params.extra_params or {} output_type = extra_params.get("output_type", "video") - is_t2i = str(output_type).lower() == "image" - - # None = unset; resolve by mode exactly once. Non-None values pass through. - mode_params = COSMOS3_T2I_PARAMS if is_t2i else COSMOS3_720P_PARAMS - - def resolved(value, field_name): - return value if value is not None else mode_params[field_name] - - height = resolved(req.params.height, "height") - width = resolved(req.params.width, "width") - num_inference_steps = resolved(req.params.num_inference_steps, "num_inference_steps") - guidance_scale = resolved(req.params.guidance_scale, "guidance_scale") video = extra_params.get("video") # encoded MP4/AVI bytes (the extra-param contract) return self.forward( @@ -772,7 +760,6 @@ def forward( "frame at every step, and this pipeline does not re-anchor it per step." ) - is_t2i = str(output_type).lower() == "image" if image is not None and video is not None: raise ValueError( "Cosmos3 generation supports text-only, text + image, " diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/sampling.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/sampling.py index e0645e5f0666..59e36e203ae6 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/sampling.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/sampling.py @@ -254,9 +254,7 @@ def set_flow_shift( return scheduler current_shift = float(_config_get(scheduler.config, "flow_shift", 1.0) or 1.0) - target_shift = ( - self.checkpoint_flow_shift if target_shift is None else float(target_shift) - ) + target_shift = self.checkpoint_flow_shift if target_shift is None else float(target_shift) current_karras = bool(_config_get(scheduler.config, "use_karras_sigmas", False)) base_karras = bool(_config_get(self.unipc_base_config, "use_karras_sigmas", False)) target_karras = base_karras if use_karras_sigmas is None else bool(use_karras_sigmas) From ce201a32656cda35410ceb43cec6c2cb6d26d749 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 30 Jul 2026 10:10:09 -0700 Subject: [PATCH 55/64] Trim the documented failure contract and the public description Drop the per-exception Raises block from VisualGen.generate() and the restatements on the result handle: enumerating every failure class in a docstring is a maintenance burden, and the single-vs-batch distinction (raises vs never raises) is what callers actually need. Strip worker-side and pipeline-specific behaviour out of the input_reference description. It is a public request schema, so it should say what to send and what comes back, not how decoding is arranged internally or which env var bounds it. Signed-off-by: Igor Shovkun --- tensorrt_llm/serve/openai_protocol.py | 22 ++++++---------------- tensorrt_llm/visual_gen/visual_gen.py | 18 ++---------------- 2 files changed, 8 insertions(+), 32 deletions(-) diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index 16abbde0efb3..47d5afe2986f 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -1689,22 +1689,12 @@ class VideoGenerationRequest(OpenAIBaseModel): description=( "Optional image or video reference that guides generation. Content " "is routed by its container signature — filename and MIME metadata " - "are ignored (JSON requests carry bare base64 with no such " - "metadata). Images (PNG image/png, JPEG image/jpeg) condition " - "image-to-video. Video containers (MP4 video/mp4, AVI " - "video/x-msvideo) pass through encoded and are decoded on the " - "workers' NVDEC; tested codec is H.264 in both containers, other " - "codecs/profiles depend on the GPU's decoder capabilities and are " - "best-effort. The signature only routes: the worker's decoder is " - "what accepts a reference, for both modalities. HEIF/AVIF still " - "images share the ISO-BMFF signature with MP4 and are rejected " - "with a 400 asking for PNG/JPEG rather than routed to the video " - "decoder. An unrecognized container is a 400 at the boundary; how " - "a decoder failure past it is reported is pipeline-specific — " - "Cosmos3 reports undecodable or corrupt references as client " - "errors (400) and device-memory exhaustion as capacity (503). " - "Reference decoding is bounded to 7200 frames by default " - "(TRTLLM_MAX_REFERENCE_DECODE_FRAMES). JSON requests carry base64 " + "are ignored. Supported references are PNG and JPEG images (which " + "condition image-to-video) and MP4 and AVI video; H.264 is the " + "tested video codec, other codecs are best-effort. HEIF/AVIF still " + "images share a container signature with MP4 and are rejected with " + "a 400 asking for PNG or JPEG. Unrecognized, corrupt or " + "undecodable content returns 400. JSON requests carry base64 " "bytes; multipart requests upload the file."), ) diff --git a/tensorrt_llm/visual_gen/visual_gen.py b/tensorrt_llm/visual_gen/visual_gen.py index 9aacc2fdd217..d597bc88b238 100644 --- a/tensorrt_llm/visual_gen/visual_gen.py +++ b/tensorrt_llm/visual_gen/visual_gen.py @@ -51,12 +51,7 @@ class VisualGenResult: A single instance backs both single-prompt and batch-prompt requests: - Single prompt: ``await handle`` resolves to a :class:`VisualGenOutput`. - Underlying-request failure raises a built-in exception carrying the - detail in its message: :class:`ValueError` for client errors (unusable - request content — an undecodable media reference, out-of-range - conditioning), :class:`MemoryError` for capacity failures (a valid - request that does not fit this deployment), and :class:`RuntimeError` - for anything unclassified. + Underlying-request failure raises. - Batch prompt: ``await handle`` resolves to ``List[VisualGenOutput]``. Per-item or whole-batch failure never raises; failed items carry ``error != None`` (Option B semantics). @@ -96,9 +91,7 @@ async def aresult(self, timeout: Optional[float] = None): """Wait for the underlying request and return the resolved value. For single-prompt requests, returns a :class:`VisualGenOutput`. - Underlying-request failure raises :class:`ValueError` (client), - :class:`MemoryError` (capacity), or :class:`RuntimeError` - (unclassified) — see the class docstring. + Underlying-request failure raises. For batch-prompt requests, returns ``List[VisualGenOutput]``. Never raises; failed items carry ``error != None``. @@ -349,13 +342,6 @@ def generate( prompts, ``List[VisualGenOutput]`` of the same length. Raises: - ValueError: Single-prompt path, client-class failure (unusable - request content). - MemoryError: Single-prompt path, capacity failure (a valid - request that does not fit this deployment). - RuntimeError: Single-prompt path, any unclassified failure. The - batch path never raises on per-item or whole-batch failure; - failed items carry ``error != None``. NotImplementedError: ``params`` is a list (per-item parameters are not yet supported). """ From 6657db155e0ac0d67ff5c2b4aaee1640408ca4f2 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 30 Jul 2026 11:09:11 -0700 Subject: [PATCH 56/64] Give the decoder a frame range, not a conditioning policy decode_video_reference_window took (window, keep) and enforced a frame cap read from TRTLLM_MAX_REFERENCE_DECODE_FRAMES. Both were Cosmos3 concerns living in a decode helper: the window is derived from indexes into the *output* latent timeline, which the pipeline already bound-checks, and the cap duplicated that reasoning one layer down. Take [first_frame, last_frame] instead, Python-style, so -1 is the last frame and (-8, -1) the final eight. The decoder returns what it is asked for and holds no policy; Cosmos3 turns condition_video_keep into a range. Only a range anchored at the end still costs a decode to EOS, which is now visibly the caller's choice. Also drops the constant and the env knob, so the module no longer imports os. Signed-off-by: Igor Shovkun --- examples/visual_gen/serve/README.md | 2 +- .../_torch/visual_gen/media_decode.py | 118 ++++++++---------- .../models/cosmos3/pipeline_cosmos3.py | 13 +- tensorrt_llm/deep_ep | 1 + tensorrt_llm/deep_gemm | 1 + tensorrt_llm/flash_mla | 1 + tensorrt_llm/inputs/media_io.py | 6 - tensorrt_llm/serve/openai_protocol.py | 4 +- .../_torch/visual_gen/test_media_decode.py | 67 ++++++---- 9 files changed, 118 insertions(+), 95 deletions(-) create mode 120000 tensorrt_llm/deep_ep create mode 120000 tensorrt_llm/deep_gemm create mode 120000 tensorrt_llm/flash_mla diff --git a/examples/visual_gen/serve/README.md b/examples/visual_gen/serve/README.md index a1c6628fd06a..0f2283afdb9c 100644 --- a/examples/visual_gen/serve/README.md +++ b/examples/visual_gen/serve/README.md @@ -287,7 +287,7 @@ You can customize these by: - `num_frames`: when set, wins over the `seconds * frame_rate` derivation - `seed`, `num_inference_steps`, `guidance_scale`, `max_sequence_length`, `negative_prompt`: per-request denoise controls - `input_reference`: Reference image (I2V/TI2V) or video (V2V), routed by container signature — filename and content type are ignored; accepted as base64-encoded string in JSON or as a file in multipart form-data. The signature only routes; the bytes pass through and the worker's decoder is what accepts them — Pillow for images, NVDEC for video (PyNvVideoCodec, a declared dependency). An unrecognized container is rejected with HTTP 400 at the boundary. - - **Supported formats**: reference images are tested with PNG (`image/png`) and JPEG (`image/jpeg`) and are decoded on the workers with Pillow. Reference videos are routed by container signature (MP4 `video/mp4`, AVI `video/x-msvideo`) and decoded on the workers' NVDEC — tested codec is H.264 in both containers; other codecs/profiles depend on the GPU's decoder capabilities and are best-effort. Filename and MIME metadata are never used for routing. HEIF/AVIF still images (which share the ISO-BMFF `ftyp` signature with MP4) are detected and rejected with a 400 asking for PNG or JPEG, rather than being sent to the video decoder. How a decoder failure past the boundary is reported is pipeline-specific: Cosmos3 classifies corrupt or undecodable references as client errors (400) and device-memory exhaustion as capacity (503). Reference decoding is bounded to 7200 frames by default (override with `TRTLLM_MAX_REFERENCE_DECODE_FRAMES`, `0` disables). + - **Supported formats**: reference images are tested with PNG (`image/png`) and JPEG (`image/jpeg`) and are decoded on the workers with Pillow. Reference videos are routed by container signature (MP4 `video/mp4`, AVI `video/x-msvideo`) and decoded on the workers' NVDEC — tested codec is H.264 in both containers; other codecs/profiles depend on the GPU's decoder capabilities and are best-effort. Filename and MIME metadata are never used for routing. HEIF/AVIF still images (which share the ISO-BMFF `ftyp` signature with MP4) are detected and rejected with a 400 asking for PNG or JPEG, rather than being sent to the video decoder. How a decoder failure past the boundary is reported is pipeline-specific: Cosmos3 classifies corrupt or undecodable references as client errors (400) and device-memory exhaustion as capacity (503). The conditioning window a model asks for bounds how much of a reference is decoded; asking for it from the end of a clip costs a decode of the whole clip. - `extra_params`: model-specific overflow (see below) - `response_format`: `"b64_json"` or `"url"` - `format`: Generation content encoding. Video encoders: `"mp4"`, `"avi"`, `"auto"`. Tensor formats: `"safetensors"`, `"pt"` (carries video + audio + scalar metadata in one payload for LTX-2). diff --git a/tensorrt_llm/_torch/visual_gen/media_decode.py b/tensorrt_llm/_torch/visual_gen/media_decode.py index 57bf9982d45d..fc4b9034d2b8 100644 --- a/tensorrt_llm/_torch/visual_gen/media_decode.py +++ b/tensorrt_llm/_torch/visual_gen/media_decode.py @@ -26,7 +26,6 @@ import functools import math -import os import torch @@ -92,27 +91,6 @@ def synchronize_media_prepare_status(exc: Exception | None) -> None: raise RuntimeError(message) -# Decode-work default: 5 min @ 24 fps, far above the canonical ~8 s reference. -# Deliberately its own constant — the serve's output cap (``MAX_VIDEO_FRAMES``) -# is a different policy that happens to share the value today. -DEFAULT_MAX_REFERENCE_DECODE_FRAMES = 7200 - - -def max_reference_decode_frames() -> int | None: - """Decode-work limit for a video reference, or ``None`` when disabled. - - Bounds serial worker occupancy for forward-only ``keep="last"`` decoding - (encoded size cannot: hours of low-bitrate video are small on disk). The - default sits far above the canonical ~8 s reference; trusted deployments - may raise it or disable it with ``TRTLLM_MAX_REFERENCE_DECODE_FRAMES=0``. - """ - raw = os.environ.get("TRTLLM_MAX_REFERENCE_DECODE_FRAMES") - if raw is None: - return DEFAULT_MAX_REFERENCE_DECODE_FRAMES - limit = int(raw) - return None if limit <= 0 else limit - - @functools.lru_cache(maxsize=32) def _lanczos_taps( in_size: int, out_size: int, device_str: str, a: int = 3 @@ -192,22 +170,39 @@ def resize_center_crop_uint8(frames: torch.Tensor, target_h: int, target_w: int) def decode_video_reference_window( data: bytes, *, - window: int, - keep: str, + first_frame: int, + last_frame: int, target_h: int, target_w: int, device: torch.device, ) -> torch.Tensor: - """Decode encoded reference bytes into the conditioning window on device. - - Returns a uint8 ``[T, target_h, target_w, 3]`` tensor, ``T <= window``: - the first ``window`` frames (``keep="first"``, decode stops early) or the - last ``window`` (``keep="last"``, sequential decode to EOS through a - preallocated ring — the memory-buffer demuxer is a forward-only feeder, - seeking is not assumed). Shorter-than-window clips return what exists; - the pipeline right-pads. Frames are resized to the target resolution - before retention, so a high-resolution source never dominates memory. + """Decode frames ``[first_frame, last_frame]`` of a reference on device. + + Returns uint8 ``[T, target_h, target_w, 3]``. Indices are Python-style: + non-negative counts from the start, negative from the end, so ``-1`` is + the last frame and ``(-8, -1)`` the final eight. Both ends are inclusive. + + A negative index costs a decode to EOS — the memory-buffer demuxer is a + forward-only feeder, seeking is not assumed — so the caller pays for the + whole clip when asking from the end. Non-negative ranges stop as soon as + the range is filled. Frames are resized to the target resolution before + retention, so a high-resolution source never dominates memory. Clips + shorter than the range return what exists; the caller pads. + + This decodes what it is asked for and imposes no policy of its own: any + bound on range size belongs to the model that knows what it can use. """ + if (first_frame < 0) != (last_frame < 0): + raise ValueError( + f"first_frame and last_frame must both count from the start or " + f"both from the end, got ({first_frame}, {last_frame})." + ) + if first_frame > last_frame: + raise ValueError( + f"first_frame must not exceed last_frame, got ({first_frame}, {last_frame})." + ) + window = last_frame - first_frame + 1 + from_end = first_frame < 0 try: import PyNvVideoCodec as nvc except ImportError as exc: @@ -216,13 +211,6 @@ def decode_video_reference_window( "install the declared dependency (pip install PyNvVideoCodec)." ) from exc - frame_cap = max_reference_decode_frames() - if frame_cap is not None and keep == "first" and window > frame_cap: - raise ValueError( - f"Conditioning window of {window} frames exceeds the reference " - f"decode limit of {frame_cap} (TRTLLM_MAX_REFERENCE_DECODE_FRAMES)." - ) - position = 0 def _read(buf: bytearray) -> int: @@ -264,28 +252,30 @@ def _read(buf: bytearray) -> int: f"decoder session could not be created: {exc}" ) from exc - ring = torch.empty(window, target_h, target_w, 3, dtype=torch.uint8, device=device) - count = 0 + # Non-negative ranges retain exactly the requested slice, so the ring + # is filled once; negative ranges cannot know the length up front, so + # it wraps and holds the trailing `tail` frames until EOS. + tail = -first_frame if from_end else window + ring = torch.empty(tail, target_h, target_w, 3, dtype=torch.uint8, device=device) + count = 0 # frames decoded so far, i.e. the index of the next frame + kept = 0 # frames written into the ring try: for packet in demuxer: + done = False for frame in decoder.Decode(packet): - if frame_cap is not None and count >= frame_cap: - raise ValueError( - f"Video reference exceeds the decode limit of " - f"{frame_cap} frames " - f"(TRTLLM_MAX_REFERENCE_DECODE_FRAMES); trim the " - f"clip{' or use condition_video_keep=first' if keep == 'last' else ''}." + if not from_end and count > last_frame: + done = True + break + if from_end or count >= first_frame: + decoded = torch.from_dlpack(frame) + # Ownership copy off the NVDEC surface (recycled by + # the decoder) and resize-before-retain in one step. + ring[kept % tail].copy_( + resize_center_crop_uint8(decoded.unsqueeze(0), target_h, target_w)[0] ) - decoded = torch.from_dlpack(frame) - # Ownership copy off the NVDEC surface (recycled by the - # decoder) and resize-before-retain in one step. - ring[count % window].copy_( - resize_center_crop_uint8(decoded.unsqueeze(0), target_h, target_w)[0] - ) + kept += 1 count += 1 - if keep == "first" and count >= window: - break - if keep == "first" and count >= window: + if done: break except torch.cuda.OutOfMemoryError as exc: raise MemoryError( @@ -303,12 +293,14 @@ def _read(buf: bytearray) -> int: "Video reference contains no decodable frames; the payload " "may be corrupt or use an unsupported codec." ) - if count <= window: - return ring[:count] - start = count % window - if start == 0: - return ring - return torch.cat([ring[start:], ring[:start]]) + if kept < tail: + frames = ring[:kept] + else: + start = kept % tail + frames = ring if start == 0 else torch.cat([ring[start:], ring[:start]]) + # `frames` now holds the trailing `tail` frames in order; a negative + # last_frame other than -1 drops the ones after it. + return frames[:window] if from_end else frames finally: del decoder del demuxer diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index 4bd1b77557c4..a44fa8c5fc34 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -972,10 +972,19 @@ def forward( window = _condition_pixel_frame_count( condition_video_latent_indexes, self.vae_scale_factor_temporal ) + # The conditioning window is a Cosmos3 constraint: it is + # derived from indexes into the *output* latent timeline, + # already bound-checked above. The decoder just returns + # the frames asked for. "last" is a negative range, which + # costs a decode to EOS -- the caller's choice to make. + if _normalize_condition_video_keep(condition_video_keep) == "first": + first_frame, last_frame = 0, window - 1 + else: + first_frame, last_frame = -window, -1 frames_u8 = decode_video_reference_window( video, - window=window, - keep=_normalize_condition_video_keep(condition_video_keep), + first_frame=first_frame, + last_frame=last_frame, target_h=height, target_w=width, device=self.device, diff --git a/tensorrt_llm/deep_ep b/tensorrt_llm/deep_ep new file mode 120000 index 000000000000..bfec2c2cde26 --- /dev/null +++ b/tensorrt_llm/deep_ep @@ -0,0 +1 @@ +/home/scratch.ishovkun_gpu/wt-v2v-rebase/cpp/build/tensorrt_llm/deep_ep/python/deep_ep \ No newline at end of file diff --git a/tensorrt_llm/deep_gemm b/tensorrt_llm/deep_gemm new file mode 120000 index 000000000000..3e72470b5636 --- /dev/null +++ b/tensorrt_llm/deep_gemm @@ -0,0 +1 @@ +/home/scratch.ishovkun_gpu/wt-v2v-rebase/cpp/build/tensorrt_llm/deep_gemm/python/deep_gemm \ No newline at end of file diff --git a/tensorrt_llm/flash_mla b/tensorrt_llm/flash_mla new file mode 120000 index 000000000000..b92b2be6f6d7 --- /dev/null +++ b/tensorrt_llm/flash_mla @@ -0,0 +1 @@ +/home/scratch.ishovkun_gpu/wt-v2v-rebase/cpp/build/tensorrt_llm/flash_mla/python/flash_mla \ No newline at end of file diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index 76a748c3fe42..19c10697416f 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -476,12 +476,6 @@ def sniff_media_kind(data) -> Optional[str]: return None -# Longest video, in frames, a client may request as *output* (the serve's -# ``num_frames`` cap in ``openai_protocol``) and the most reference frames a -# worker will decode from a video reference before raising. -MAX_VIDEO_FRAMES = 7200 - - def _select_cv2_stream_buffered_backend() -> Optional[int]: """Return a VideoCapture backend that can read from a Python `BytesIO`. diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index 47d5afe2986f..50893a3fdec3 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -37,7 +37,7 @@ from typing_extensions import Annotated, Required, TypeAlias, TypedDict from tensorrt_llm.executor.request import LoRARequest -from tensorrt_llm.inputs.media_io import MAX_VIDEO_FRAMES, MediaModality +from tensorrt_llm.inputs.media_io import MediaModality from tensorrt_llm.llmapi import ConversationParams as LlmConversationParams from tensorrt_llm.llmapi import DisaggregatedParams as LlmDisaggregatedParams from tensorrt_llm.llmapi import (DisaggScheduleStyle, GuidedDecodingParams, @@ -1716,7 +1716,7 @@ class VideoGenerationRequest(OpenAIBaseModel): # The numbers are generous (a minute of video at 120 fps) so common # workloads pass; clients that need larger budgets can lift the cap # at deployment time. - num_frames: Optional[int] = Field(default=None, gt=0, le=MAX_VIDEO_FRAMES) + num_frames: Optional[int] = Field(default=None, gt=0, le=7200) seconds: Optional[float] = Field(default=None, gt=0, le=60.0) frame_rate: Optional[float] = Field(default=None, alias="fps", diff --git a/tests/unittest/_torch/visual_gen/test_media_decode.py b/tests/unittest/_torch/visual_gen/test_media_decode.py index ac1153ebe84d..ff510057689e 100644 --- a/tests/unittest/_torch/visual_gen/test_media_decode.py +++ b/tests/unittest/_torch/visual_gen/test_media_decode.py @@ -26,7 +26,6 @@ from tensorrt_llm._torch.visual_gen.media_decode import ( _lanczos_taps, decode_video_reference_window, - max_reference_decode_frames, resize_center_crop_uint8, synchronize_media_prepare_status, ) @@ -176,24 +175,18 @@ def test_two_rank_convergence_over_gloo(self, tmp_path): assert healthy == "client:[rank 1] rank-local decode failure" -class TestDecodeFrameLimit: - def test_default_and_env_override(self, monkeypatch): - monkeypatch.delenv("TRTLLM_MAX_REFERENCE_DECODE_FRAMES", raising=False) - assert max_reference_decode_frames() == 7200 - monkeypatch.setenv("TRTLLM_MAX_REFERENCE_DECODE_FRAMES", "12") - assert max_reference_decode_frames() == 12 - monkeypatch.setenv("TRTLLM_MAX_REFERENCE_DECODE_FRAMES", "0") - assert max_reference_decode_frames() is None - - @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") class TestDecodeVideoReferenceWindow: _DEVICE = torch.device("cuda:0") - def _decode(self, data: bytes, **kwargs): - defaults = dict(window=5, keep="first", target_h=64, target_w=64, device=self._DEVICE) + def _decode(self, data: bytes, *, window=5, keep="first", **kwargs): + """Express a (window, keep) request as the decoder's frame range.""" + span = (0, window - 1) if keep == "first" else (-window, -1) + defaults = dict(target_h=64, target_w=64, device=self._DEVICE) defaults.update(kwargs) - return decode_video_reference_window(data, **defaults) + return decode_video_reference_window( + data, first_frame=span[0], last_frame=span[1], **defaults + ) @pytest.mark.parametrize("fixture", [_MP4, _AVI], ids=["mp4", "avi"]) def test_keep_first_display_order(self, fixture): @@ -242,16 +235,48 @@ def test_corrupt_bytes_with_valid_magic_is_client_error(self): with pytest.raises(ValueError): self._decode(payload) - def test_frame_limit_trips_on_emitted_frames(self, monkeypatch): - monkeypatch.setenv("TRTLLM_MAX_REFERENCE_DECODE_FRAMES", "4") - with pytest.raises(ValueError, match="decode limit"): - self._decode(_MP4.read_bytes(), keep="last") - - def test_frame_limit_disabled(self, monkeypatch): - monkeypatch.setenv("TRTLLM_MAX_REFERENCE_DECODE_FRAMES", "0") + def test_decoder_imposes_no_frame_limit(self): + # The decoder returns what it is asked for; bounding the request is + # the caller's business (Cosmos3 derives it from the output length). window = self._decode(_MP4.read_bytes(), window=20, keep="last") assert window.shape[0] == 9 + def test_interior_range(self): + window = decode_video_reference_window( + _MP4.read_bytes(), + first_frame=3, + last_frame=5, + target_h=64, + target_w=64, + device=self._DEVICE, + ) + assert _frame_indices(window) == [3, 4, 5] + + def test_negative_range_excluding_final_frames(self): + window = decode_video_reference_window( + _MP4.read_bytes(), + first_frame=-4, + last_frame=-3, + target_h=64, + target_w=64, + device=self._DEVICE, + ) + assert _frame_indices(window) == [5, 6] + + @pytest.mark.parametrize( + "span", [(0, -1), (-1, 0), (5, 2)], ids=["mixed", "mixed-rev", "reversed"] + ) + def test_malformed_range_rejected(self, span): + with pytest.raises(ValueError): + decode_video_reference_window( + _MP4.read_bytes(), + first_frame=span[0], + last_frame=span[1], + target_h=64, + target_w=64, + device=self._DEVICE, + ) + def test_resize_perf_representative(self): # Representative evidence for the local-tap resample: a 1080p frame # to 720p-cover must be in the low-millisecond range. The bound is From 8a0e4c7a96425c0bb9e6d1629809c31ff73a8435 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 30 Jul 2026 13:16:46 -0700 Subject: [PATCH 57/64] Move the decode to tensorrt_llm/media media_decode.py held two unrelated things. The decoding half is generic media handling and belongs next to its counterpart, media/encoding.py; tensorrt_llm/media has the same code owner, so nothing about the review surface changes. The other half stays under visual_gen. classify_worker_error is not a media concern: executor.py calls it in the generic worker handler to fill DiffusionResponse.error_type, and it is what remained after we settled on primitive exceptions instead of a VisualGen error hierarchy. synchronize_media_prepare_status is the rank protocol and calls it, so the two share a home in visual_gen/utils.py rather than drifting apart. Tests keep their path so the CI lists are untouched. Signed-off-by: Igor Shovkun --- examples/visual_gen/serve/README.md | 2 +- tensorrt_llm/_torch/visual_gen/executor.py | 2 +- .../models/cosmos3/pipeline_cosmos3.py | 10 +-- tensorrt_llm/_torch/visual_gen/utils.py | 59 ++++++++++++++ .../media_decode.py => media/decoding.py} | 76 ++----------------- .../_torch/visual_gen/test_media_decode.py | 12 +-- 6 files changed, 81 insertions(+), 80 deletions(-) rename tensorrt_llm/{_torch/visual_gen/media_decode.py => media/decoding.py} (78%) diff --git a/examples/visual_gen/serve/README.md b/examples/visual_gen/serve/README.md index 0f2283afdb9c..09fbd87f7667 100644 --- a/examples/visual_gen/serve/README.md +++ b/examples/visual_gen/serve/README.md @@ -287,7 +287,7 @@ You can customize these by: - `num_frames`: when set, wins over the `seconds * frame_rate` derivation - `seed`, `num_inference_steps`, `guidance_scale`, `max_sequence_length`, `negative_prompt`: per-request denoise controls - `input_reference`: Reference image (I2V/TI2V) or video (V2V), routed by container signature — filename and content type are ignored; accepted as base64-encoded string in JSON or as a file in multipart form-data. The signature only routes; the bytes pass through and the worker's decoder is what accepts them — Pillow for images, NVDEC for video (PyNvVideoCodec, a declared dependency). An unrecognized container is rejected with HTTP 400 at the boundary. - - **Supported formats**: reference images are tested with PNG (`image/png`) and JPEG (`image/jpeg`) and are decoded on the workers with Pillow. Reference videos are routed by container signature (MP4 `video/mp4`, AVI `video/x-msvideo`) and decoded on the workers' NVDEC — tested codec is H.264 in both containers; other codecs/profiles depend on the GPU's decoder capabilities and are best-effort. Filename and MIME metadata are never used for routing. HEIF/AVIF still images (which share the ISO-BMFF `ftyp` signature with MP4) are detected and rejected with a 400 asking for PNG or JPEG, rather than being sent to the video decoder. How a decoder failure past the boundary is reported is pipeline-specific: Cosmos3 classifies corrupt or undecodable references as client errors (400) and device-memory exhaustion as capacity (503). The conditioning window a model asks for bounds how much of a reference is decoded; asking for it from the end of a clip costs a decode of the whole clip. + - **Supported formats**: reference images are tested with PNG (`image/png`) and JPEG (`image/jpeg`) and are decoded on the workers with Pillow. Reference videos are routed by container signature (MP4 `video/mp4`, AVI `video/x-msvideo`) and decoded on the workers' NVDEC — tested codec is H.264 in both containers; other codecs/profiles depend on the GPU's decoder capabilities and are best-effort. Filename and MIME metadata are never used for routing. HEIF/AVIF still images (which share the ISO-BMFF `ftyp` signature with MP4) are detected and rejected with a 400 asking for PNG or JPEG, rather than being sent to the video decoder. - `extra_params`: model-specific overflow (see below) - `response_format`: `"b64_json"` or `"url"` - `format`: Generation content encoding. Video encoders: `"mp4"`, `"avi"`, `"auto"`. Tensor formats: `"safetensors"`, `"pt"` (carries video + audio + scalar metadata in one payload for LTX-2). diff --git a/tensorrt_llm/_torch/visual_gen/executor.py b/tensorrt_llm/_torch/visual_gen/executor.py index 673a7435561d..2da03be7b32b 100644 --- a/tensorrt_llm/_torch/visual_gen/executor.py +++ b/tensorrt_llm/_torch/visual_gen/executor.py @@ -15,9 +15,9 @@ import torch.multiprocessing as mp import zmq -from tensorrt_llm._torch.visual_gen.media_decode import classify_worker_error from tensorrt_llm._torch.visual_gen.output import PipelineOutput from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader +from tensorrt_llm._torch.visual_gen.utils import classify_worker_error from tensorrt_llm.executor.ipc import ZeroMqQueue from tensorrt_llm.llmapi.utils import configure_cpu_affinity from tensorrt_llm.logger import logger diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index a44fa8c5fc34..bf5de96eca50 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -26,17 +26,17 @@ from diffusers.video_processor import VideoProcessor from transformers import Qwen2Tokenizer -from tensorrt_llm._torch.visual_gen.media_decode import ( - decode_video_reference_window, - synchronize_media_prepare_status, -) from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline from tensorrt_llm._torch.visual_gen.pipeline_registry import PipelineComponent, register_pipeline -from tensorrt_llm._torch.visual_gen.utils import postprocess_video_tensor +from tensorrt_llm._torch.visual_gen.utils import ( + postprocess_video_tensor, + synchronize_media_prepare_status, +) from tensorrt_llm._utils import nvtx_range from tensorrt_llm.inputs.utils import load_image from tensorrt_llm.logger import logger +from tensorrt_llm.media.decoding import decode_video_reference_window from .defaults import ( COSMOS3_720P_PARAMS, diff --git a/tensorrt_llm/_torch/visual_gen/utils.py b/tensorrt_llm/_torch/visual_gen/utils.py index 8140dc959122..04d9c9db00ec 100644 --- a/tensorrt_llm/_torch/visual_gen/utils.py +++ b/tensorrt_llm/_torch/visual_gen/utils.py @@ -43,6 +43,65 @@ def as_tuple(x): return x if isinstance(x, tuple) else (x, x) +def classify_worker_error(exc: BaseException) -> str | None: + """Failure class for the response channel: "client", "capacity", or None. + + Keyed off built-in exception types rather than a VisualGen-specific + hierarchy: ``ValueError`` means the request's content was unusable + (400), ``MemoryError`` means a valid request did not fit (503), and + anything else is an unclassified runtime failure (500). Detail travels + in the message. ``torch.cuda.OutOfMemoryError`` is spelled out because + it derives from ``RuntimeError``, not ``MemoryError``. + """ + if isinstance(exc, (MemoryError, torch.cuda.OutOfMemoryError)): + return "capacity" + if isinstance(exc, ValueError): + return "client" + return None + + +def synchronize_media_prepare_status(exc: Exception | None) -> None: + """All-rank convergence point between media prepare and model collectives. + + Every rank decodes/prepares its media independently; a rank that failed + while others proceed into the transformer's collectives would hang the + job. All ranks call this with their local outcome; if any failed, the + lowest failing rank's error class + message is broadcast, the failing + rank(s) re-raise their own exception, and every healthy rank raises a + reconstructed equivalent in lockstep. Runs on CPU tensors so + the hybrid (``cpu:gloo``) process group carries it even when the failure + was CUDA/NVDEC initialization. Converges *caught* failures only — a fatal + process or context death is beyond its reach. + """ + if not (dist.is_available() and dist.is_initialized()) or dist.get_world_size() == 1: + if exc is not None: + raise exc + return + + healthy_sentinel = 2**31 - 1 + rank = dist.get_rank() + flag = torch.tensor([rank if exc is not None else healthy_sentinel], dtype=torch.int64) + dist.all_reduce(flag, op=dist.ReduceOp.MIN) + failing_rank = int(flag.item()) + if failing_rank == healthy_sentinel: + return + + payload = [None] + if rank == failing_rank: + payload = [(classify_worker_error(exc), str(exc))] + dist.broadcast_object_list(payload, src=failing_rank) + + if exc is not None: + raise exc + kind, message = payload[0] + message = f"[rank {failing_rank}] {message}" + if kind == "client": + raise ValueError(message) + if kind == "capacity": + raise MemoryError(message) + raise RuntimeError(message) + + class SequenceSharder: """Block-shard / all-gather a tensor along its sequence dimension. diff --git a/tensorrt_llm/_torch/visual_gen/media_decode.py b/tensorrt_llm/media/decoding.py similarity index 78% rename from tensorrt_llm/_torch/visual_gen/media_decode.py rename to tensorrt_llm/media/decoding.py index fc4b9034d2b8..9fa72833bc0d 100644 --- a/tensorrt_llm/_torch/visual_gen/media_decode.py +++ b/tensorrt_llm/media/decoding.py @@ -12,13 +12,14 @@ # 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. -"""Worker-side video-reference decoding on NVDEC (PyNvVideoCodec). +"""Video-reference decoding on NVDEC (PyNvVideoCodec). -Encoded reference bytes arrive from the coordinator; each worker rank demuxes -them from memory, decodes on NVDEC, and retains only the conditioning window, -resized to the request's output resolution — so retained memory is bounded by -the request's own output shape and at any instant only one source-resolution -frame is alive. +Encoded bytes are demuxed from memory, decoded on NVDEC, and only the +requested window is retained, resized to the caller's target resolution — so +retained memory is bounded by that target and at any instant only one +source-resolution frame is alive. + +Counterpart to :mod:`tensorrt_llm.media.encoding`. PyNvVideoCodec is imported function-locally: ``import tensorrt_llm`` on a CPU-only host must never load the driver-linked extension. @@ -30,67 +31,6 @@ import torch -def classify_worker_error(exc: BaseException) -> str | None: - """Failure class for the response channel: "client", "capacity", or None. - - Keyed off built-in exception types rather than a VisualGen-specific - hierarchy: ``ValueError`` means the request's content was unusable - (400), ``MemoryError`` means a valid request did not fit (503), and - anything else is an unclassified runtime failure (500). Detail travels - in the message. ``torch.cuda.OutOfMemoryError`` is spelled out because - it derives from ``RuntimeError``, not ``MemoryError``. - """ - if isinstance(exc, (MemoryError, torch.cuda.OutOfMemoryError)): - return "capacity" - if isinstance(exc, ValueError): - return "client" - return None - - -def synchronize_media_prepare_status(exc: Exception | None) -> None: - """All-rank convergence point between media prepare and model collectives. - - Every rank decodes/prepares its media independently; a rank that failed - while others proceed into the transformer's collectives would hang the - job. All ranks call this with their local outcome; if any failed, the - lowest failing rank's error class + message is broadcast, the failing - rank(s) re-raise their own exception, and every healthy rank raises a - reconstructed equivalent in lockstep. Runs on CPU tensors so - the hybrid (``cpu:gloo``) process group carries it even when the failure - was CUDA/NVDEC initialization. Converges *caught* failures only — a fatal - process or context death is beyond its reach. - """ - import torch.distributed as dist - - if not (dist.is_available() and dist.is_initialized()) or dist.get_world_size() == 1: - if exc is not None: - raise exc - return - - healthy_sentinel = 2**31 - 1 - rank = dist.get_rank() - flag = torch.tensor([rank if exc is not None else healthy_sentinel], dtype=torch.int64) - dist.all_reduce(flag, op=dist.ReduceOp.MIN) - failing_rank = int(flag.item()) - if failing_rank == healthy_sentinel: - return - - payload = [None] - if rank == failing_rank: - payload = [(classify_worker_error(exc), str(exc))] - dist.broadcast_object_list(payload, src=failing_rank) - - if exc is not None: - raise exc - kind, message = payload[0] - message = f"[rank {failing_rank}] {message}" - if kind == "client": - raise ValueError(message) - if kind == "capacity": - raise MemoryError(message) - raise RuntimeError(message) - - @functools.lru_cache(maxsize=32) def _lanczos_taps( in_size: int, out_size: int, device_str: str, a: int = 3 @@ -135,7 +75,7 @@ def _resample_last_dim(x: torch.Tensor, weights: torch.Tensor, taps: torch.Tenso def resize_center_crop_uint8(frames: torch.Tensor, target_h: int, target_w: int) -> torch.Tensor: """Resize + center-crop uint8 ``[T, H, W, C]`` frames to the target size. - Applied to the worker-decoded reference frames before retention. + Applied to the decoded reference frames before retention. Semantics mirror the reference implementation's PIL path (cover-scale by ``max(target/source)``, ceil-rounded resize with Lanczos-3, center crop), implemented as separable local-tap resampling: per output pixel only the diff --git a/tests/unittest/_torch/visual_gen/test_media_decode.py b/tests/unittest/_torch/visual_gen/test_media_decode.py index ff510057689e..7fc9f3a84d6f 100644 --- a/tests/unittest/_torch/visual_gen/test_media_decode.py +++ b/tests/unittest/_torch/visual_gen/test_media_decode.py @@ -1,6 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Direct tests for :mod:`tensorrt_llm._torch.visual_gen.media_decode`. +"""Direct tests for :mod:`tensorrt_llm.media.decoding`, plus the rank +convergence protocol that guards it +(``tensorrt_llm._torch.visual_gen.utils.synchronize_media_prepare_status``). The decode tests run on the checked-in H.264 fixtures (see ``test_data/README.md`` for provenance). Each fixture frame encodes its own @@ -23,11 +25,11 @@ import torch from PIL import Image -from tensorrt_llm._torch.visual_gen.media_decode import ( +from tensorrt_llm._torch.visual_gen.utils import synchronize_media_prepare_status +from tensorrt_llm.media.decoding import ( _lanczos_taps, decode_video_reference_window, resize_center_crop_uint8, - synchronize_media_prepare_status, ) _TEST_DATA = Path(__file__).parent / "test_data" @@ -96,7 +98,7 @@ def test_import_tensorrt_llm_does_not_load_pynvvideocodec(self): decode module itself) must not load it — only an actual decode may.""" code = ( "import sys; import tensorrt_llm; " - "import tensorrt_llm._torch.visual_gen.media_decode; " + "import tensorrt_llm.media.decoding; " "assert 'PyNvVideoCodec' not in sys.modules, " "'driver-linked PyNvVideoCodec loaded at import time'" ) @@ -114,7 +116,7 @@ def _status_protocol_rank(rank: int, world_size: int, init_file: str, results_di import torch.distributed as dist - from tensorrt_llm._torch.visual_gen.media_decode import synchronize_media_prepare_status + from tensorrt_llm._torch.visual_gen.utils import synchronize_media_prepare_status dist.init_process_group( "gloo", init_method=f"file://{init_file}", rank=rank, world_size=world_size From 76fda37c3157b3480b4b4e1dfa7ee5676d0f0c72 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 30 Jul 2026 13:51:55 -0700 Subject: [PATCH 58/64] Drop build symlinks from the index 6657db155e committed tensorrt_llm/{deep_ep,deep_gemm,flash_mla} as symlinks holding absolute paths into this worktree's cpp/build, which resolve nowhere else. They slipped past .gitignore because the patterns carried a trailing slash, which git matches against directories only; a native build leaves symlinks there instead. Drop the slash so both forms are ignored. Signed-off-by: Igor Shovkun --- .gitignore | 6 +++--- tensorrt_llm/deep_ep | 1 - tensorrt_llm/deep_gemm | 1 - tensorrt_llm/flash_mla | 1 - 4 files changed, 3 insertions(+), 6 deletions(-) delete mode 120000 tensorrt_llm/deep_ep delete mode 120000 tensorrt_llm/deep_gemm delete mode 120000 tensorrt_llm/flash_mla diff --git a/.gitignore b/.gitignore index 512373ae8788..9417ba9a7ebc 100644 --- a/.gitignore +++ b/.gitignore @@ -47,14 +47,14 @@ tensorrt_llm/bindings.pyi tensorrt_llm/bindings/**/*.pyi tensorrt_llm/tensorrt_llm_transfer_agent_binding.*.so tensorrt_llm/tensorrt_llm_transfer_agent_binding.pyi -tensorrt_llm/deep_ep/ +tensorrt_llm/deep_ep tensorrt_llm/deep_ep_cpp_tllm.*.so tensorrt_llm/deep_ep_cpp_tllm.pyi -tensorrt_llm/deep_gemm/ +tensorrt_llm/deep_gemm tensorrt_llm/deep_gemm_cpp_tllm.*.so tensorrt_llm/deep_gemm_cpp_tllm.pyi tensorrt_llm/pg_utils_bindings.*.so -tensorrt_llm/flash_mla/ +tensorrt_llm/flash_mla tensorrt_llm/flash_mla_cpp_tllm.*.so tensorrt_llm/flash_mla_cpp_tllm.pyi tensorrt_llm/runtime/kv_cache_manager_v2/**/*.so diff --git a/tensorrt_llm/deep_ep b/tensorrt_llm/deep_ep deleted file mode 120000 index bfec2c2cde26..000000000000 --- a/tensorrt_llm/deep_ep +++ /dev/null @@ -1 +0,0 @@ -/home/scratch.ishovkun_gpu/wt-v2v-rebase/cpp/build/tensorrt_llm/deep_ep/python/deep_ep \ No newline at end of file diff --git a/tensorrt_llm/deep_gemm b/tensorrt_llm/deep_gemm deleted file mode 120000 index 3e72470b5636..000000000000 --- a/tensorrt_llm/deep_gemm +++ /dev/null @@ -1 +0,0 @@ -/home/scratch.ishovkun_gpu/wt-v2v-rebase/cpp/build/tensorrt_llm/deep_gemm/python/deep_gemm \ No newline at end of file diff --git a/tensorrt_llm/flash_mla b/tensorrt_llm/flash_mla deleted file mode 120000 index b92b2be6f6d7..000000000000 --- a/tensorrt_llm/flash_mla +++ /dev/null @@ -1 +0,0 @@ -/home/scratch.ishovkun_gpu/wt-v2v-rebase/cpp/build/tensorrt_llm/flash_mla/python/flash_mla \ No newline at end of file From 79b7d1aeba127aaf0eeddfb7581e589eebe99d99 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 30 Jul 2026 14:04:58 -0700 Subject: [PATCH 59/64] Restore the action state the merge dropped The upstream merge removed TransformerOutput.action, the action_gen flag, and narrowed the skipped action weight prefixes to a blanket "action_"; it also added start_frame_offset, which no caller passes. None of that belongs to V2V. transformer_cosmos3.py is now byte-identical to upstream/main, so this PR no longer touches the file. Signed-off-by: Igor Shovkun --- .../models/cosmos3/transformer_cosmos3.py | 34 +++++++++---------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py index 8bf5fb1367f4..1a85cf2b5784 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -93,6 +93,9 @@ class TransformerOutput: audio: Optional[torch.Tensor] = None """[B, audio_dim, T_audio] audio velocity prediction, or None.""" + action: Optional[torch.Tensor] = None + """[B, T_action, action_dim] action velocity prediction, or None.""" + def compute_mrope_position_ids_text( num_tokens: int, @@ -120,7 +123,6 @@ def compute_mrope_position_ids_vision( base_fps: float = 24.0, temporal_compression_factor: int = 4, enable_fps_modulation: bool = False, - start_frame_offset: int = 0, ) -> tuple[torch.Tensor, int | float]: """Generate 3D mRoPE position IDs for vision tokens. @@ -140,17 +142,15 @@ def compute_mrope_position_ids_vision( base_tps = base_fps / temporal_compression_factor frame_indices = torch.arange(grid_t, dtype=torch.float32) t_index = ( - ((frame_indices + start_frame_offset) / tps * base_tps + temporal_offset) + (frame_indices / tps * base_tps + temporal_offset) .view(-1, 1) .expand(-1, grid_h * grid_w) .flatten() ) else: - t_index = ( - torch.arange(grid_t, dtype=torch.long).view(-1, 1).expand(-1, grid_h * grid_w).flatten() - + int(temporal_offset) - + start_frame_offset - ) + t_index = torch.arange(grid_t, dtype=torch.long).view(-1, 1).expand( + -1, grid_h * grid_w + ).flatten() + int(temporal_offset) h_index = ( torch.arange(grid_h, dtype=torch.long).view(1, -1, 1).expand(grid_t, -1, grid_w).flatten() @@ -725,6 +725,7 @@ def __init__(self, model_config: DiffusionModelConfig): super().__init__(model_config) pretrained_config = apply_pretrained_config_compat_defaults(model_config.pretrained_config) self.audio_gen = getattr(pretrained_config, "sound_gen", False) + self.action_gen = getattr(pretrained_config, "action_gen", False) self.hidden_size = pretrained_config.hidden_size self.num_hidden_layers = pretrained_config.num_hidden_layers @@ -1026,14 +1027,13 @@ def forward( Returns: TransformerOutput with video (and image alias) always set. audio is set to the predicted audio velocity when audio_latents is - provided; otherwise None. + provided; otherwise None. action is always None for now. """ del kwargs # Kept for diffusers API compatibility. if timestep is None: raise ValueError("Cosmos3VFMTransformer.forward requires normalized timestep.") if raw_timestep is None: raise ValueError("Cosmos3VFMTransformer.forward requires raw_timestep.") - T, H, W = video_shape Hp, Wp, _, _ = self._pad_to_patch_size(H, W) max_real_len = text_mask.sum(dim=1).max().item() @@ -1096,7 +1096,7 @@ def forward( else: self.cached_kv = cached_kv_full - # --- Extra modality token injection (audio) --- + # --- Audio token injection ------------------------------------------------- T_vid_tokens = hidden_gen.shape[1] # T * Hp * Wp T_audio = 0 if audio_latents is not None and self.audio_gen: @@ -1111,6 +1111,7 @@ def forward( hidden_states.device, hidden_gen.dtype, ) + # [B, T_vid+T_audio, hidden_size] hidden_gen = torch.cat([hidden_gen, hidden_audio], dim=1) cos_v, sin_v = self.cached_freqs_gen freqs_gen_combined = ( @@ -1157,12 +1158,13 @@ def forward( # --- Decode video velocity ------------------------------------------------ video_vel = self.unpatchify(self.llm2vae(hidden_gen[:, :T_vid_tokens]), T, H, W) - # --- Decode extra-modality velocity (audio; follows video) --- - extra_start = T_vid_tokens + # --- Decode audio velocity (if requested) --------------------------------- audio_vel = None if T_audio > 0 and audio_latents is not None and self.audio_gen: + # hidden_gen[:, T_vid_tokens:] → [B, T_audio, hidden_size] + # → llm2audio → [B, T_audio, audio_dim] → unpack → [B, audio_dim, T_audio] audio_vel = self.unpack_audio_latents( - self.llm2audio(hidden_gen[:, extra_start : extra_start + T_audio]) + self.llm2audio(hidden_gen[:, T_vid_tokens : T_vid_tokens + T_audio]) ) return TransformerOutput(video=video_vel, image=video_vel, audio=audio_vel) @@ -1175,12 +1177,10 @@ def load_weights(self, weights: dict) -> None: Maps UND vs GEN blocks into this module's layout (causal self-attn vs cross-attn + MLPs). """ remapped = {} - # The Cosmos3 checkpoint ships action modules (action_gen=true), but this - # transformer no longer builds them — skip their weights explicitly so the - # load stays quiet instead of warning on each as an unknown key. skip_prefixes = ( "lm_head.", - "action_", + "action_modality_embed", + "action_proj_", ) for key, value in weights.items(): From bb74c9ce76fbe872b0914fde4597f7b367c7368b Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 30 Jul 2026 14:05:10 -0700 Subject: [PATCH 60/64] Fix two sampling regressions the merge introduced infer() computed mode-resolved height/width/steps/guidance and then passed req.params.* to forward() anyway -- the merge kept upstream's resolution block but reverted the call site. 51799b4dd6 then correctly observed the locals were unused and deleted them, removing the evidence rather than the bug. Restore the block and wire it into the call, which is what upstream/main does; the four TestInferModeResolution cases in test_cosmos3_distilled.py cover it. V2V also rebuilt only the video scheduler. Audio kept the checkpoint's flow shift and Karras setting and was merely handed timesteps, so the two streams denoised on different schedules in the same loop; the pre-merge code rebuilt both. Route every mode through _apply_flow_shift so neither stream can be rebuilt without the other, and test it -- the existing V2V audio smoke only inspects the video scheduler. Signed-off-by: Igor Shovkun --- .../models/cosmos3/pipeline_cosmos3.py | 46 ++++++++++++---- .../visual_gen/test_cosmos3_pipeline.py | 55 +++++++++++++++++++ 2 files changed, 91 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index bf5de96eca50..05663c6633aa 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -260,20 +260,49 @@ def _run_warmup(self, height: int, width: int, num_frames: int, steps: int) -> N enable_audio=False, ) + def _apply_flow_shift( + self, target_shift: Optional[float], *, use_karras_sigmas: Optional[bool] = None + ) -> None: + """Rebuild both stream schedulers for the requested sampling knobs. + + Video and audio denoise in lockstep in one loop, so a mode that + rebuilds only the video scheduler leaves audio on the checkpoint's + sigmas and the two streams step on different schedules. + """ + self.scheduler = self.sampling.set_flow_shift( + self.scheduler, target_shift, use_karras_sigmas=use_karras_sigmas + ) + if getattr(self, "audio_scheduler", None) is not None: + self.audio_scheduler = self.sampling.set_flow_shift( + self.audio_scheduler, target_shift, use_karras_sigmas=use_karras_sigmas + ) + def infer(self, req): extra_params = req.params.extra_params or {} output_type = extra_params.get("output_type", "video") + is_t2i = str(output_type).lower() == "image" + + # None = unset; resolve by mode exactly once. Non-None values pass through. + mode_params = COSMOS3_T2I_PARAMS if is_t2i else COSMOS3_720P_PARAMS + + def resolved(value, field_name): + return value if value is not None else mode_params[field_name] + + height = resolved(req.params.height, "height") + width = resolved(req.params.width, "width") + num_inference_steps = resolved(req.params.num_inference_steps, "num_inference_steps") + guidance_scale = resolved(req.params.guidance_scale, "guidance_scale") video = extra_params.get("video") # encoded MP4/AVI bytes (the extra-param contract) return self.forward( prompt=req.prompt, negative_prompt=req.params.negative_prompt, image=req.params.image, - height=req.params.height, - width=req.params.width, + height=height, + width=width, num_frames=req.params.num_frames, - num_inference_steps=req.params.num_inference_steps, - guidance_scale=req.params.guidance_scale, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, seed=req.params.seed, max_sequence_length=req.params.max_sequence_length, frame_rate=req.params.frame_rate, @@ -789,8 +818,7 @@ def forward( if guidance_scale is None: guidance_scale = COSMOS3_T2I_PARAMS["guidance_scale"] guidance_interval = COSMOS3_T2I_PARAMS["guidance_interval"] - self.scheduler = self.sampling.set_flow_shift( - self.scheduler, + self._apply_flow_shift( flow_shift if flow_shift is not None else COSMOS3_T2I_PARAMS["flow_shift"], ) else: @@ -802,16 +830,14 @@ def forward( guidance_scale = COSMOS3_720P_PARAMS["guidance_scale"] if is_v2v: # V2V wants a stronger shift and the uniform sigma schedule. - self.scheduler = self.sampling.set_flow_shift( - self.scheduler, + self._apply_flow_shift( flow_shift if flow_shift is not None else 10.0, use_karras_sigmas=False, ) else: # Restore the checkpoint sampling knobs in case a prior T2I or # V2V request rebuilt the scheduler with mode-specific values. - self.scheduler = self.sampling.set_flow_shift( - self.scheduler, + self._apply_flow_shift( flow_shift if flow_shift is not None else self.sampling.checkpoint_flow_shift, use_karras_sigmas=None, ) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index ee9dfa62142c..e1f95f593178 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -721,6 +721,61 @@ def fake_tokenize_prompt(text, max_sequence_length, use_system_prompt, system_pr assert token_calls[0][2] is True assert token_calls[0][3] == COSMOS3_DEFAULT_SYSTEM_PROMPT + def test_v2v_rebuilds_the_audio_scheduler_too(self): + """Video and audio denoise in lockstep in one loop, so a V2V request + must rebuild both. Rebuilding only the video scheduler leaves audio on + the checkpoint's flow shift / Karras sigmas and the streams step on + different schedules.""" + pipeline = Cosmos3OmniMoTPipeline.__new__(Cosmos3OmniMoTPipeline) + pipeline.transformer = SimpleNamespace(device=torch.device("cpu")) + pipeline.audio_gen = True + rebuilt = [] + + class StopAfterTokenize(Exception): + pass + + class FakeSampling: + is_distilled = False + checkpoint_flow_shift = 1.0 + + def validate_request(self, num_inference_steps, guidance_scale): + return None + + def set_flow_shift(self, scheduler, target, *, use_karras_sigmas=None): + rebuilt.append((scheduler.name, target, use_karras_sigmas)) + return scheduler + + def fake_tokenize_prompt(text, max_sequence_length, use_system_prompt, system_prompt=None): + raise StopAfterTokenize + + pipeline.scheduler = SimpleNamespace(name="video", config=SimpleNamespace(flow_shift=1.0)) + pipeline.audio_scheduler = SimpleNamespace( + name="audio", config=SimpleNamespace(flow_shift=1.0) + ) + pipeline.sampling = FakeSampling() + pipeline._tokenize_prompt = fake_tokenize_prompt + + with pytest.raises(StopAfterTokenize): + pipeline.forward( + prompt="continue", + video=_V2V_FIXTURE_MP4.read_bytes(), + height=16, + width=16, + num_frames=5, + num_inference_steps=1, + guidance_scale=1.0, + seed=1, + max_sequence_length=8, + frame_rate=8.0, + use_duration_template=False, + use_resolution_template=False, + use_system_prompt=None, + use_guardrails=False, + enable_audio=True, + ) + + assert rebuilt == [("video", 10.0, False), ("audio", 10.0, False)] + def test_image_and_video_rejected(self, cosmos3_pipeline): with pytest.raises(ValueError, match="not both image and video"): _run_forward( From 7094e101f6a3157d68104545be077b60b9bddda3 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 30 Jul 2026 15:09:45 -0700 Subject: [PATCH 61/64] Trim input_reference down to the request contract The description still explained how routing works (container signatures, the HEIF/AVIF-vs-MP4 signature collision) and restated status codes. None of that helps a caller build a request, and the field already names the supported formats, so "send PNG or JPEG" was saying it twice. Keep what you need to send one: which format conditions which mode, that HEIF/AVIF won't work, and the two transports. Signed-off-by: Igor Shovkun --- tensorrt_llm/serve/openai_protocol.py | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index 50893a3fdec3..7ec9706e99ee 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -1687,15 +1687,11 @@ class VideoGenerationRequest(OpenAIBaseModel): input_reference: Optional[Union[str, UploadFile]] = Field( default=None, description=( - "Optional image or video reference that guides generation. Content " - "is routed by its container signature — filename and MIME metadata " - "are ignored. Supported references are PNG and JPEG images (which " - "condition image-to-video) and MP4 and AVI video; H.264 is the " - "tested video codec, other codecs are best-effort. HEIF/AVIF still " - "images share a container signature with MP4 and are rejected with " - "a 400 asking for PNG or JPEG. Unrecognized, corrupt or " - "undecodable content returns 400. JSON requests carry base64 " - "bytes; multipart requests upload the file."), + "Optional image or video reference that guides generation. PNG or " + "JPEG images condition image-to-video; MP4 or AVI video conditions " + "video-to-video, with H.264 the tested codec and others " + "best-effort. HEIF/AVIF are not supported. JSON requests carry " + "base64 bytes; multipart requests upload the file."), ) # Resolution From 5445bee810362396f65792291a0c56b86256f831 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 30 Jul 2026 15:13:01 -0700 Subject: [PATCH 62/64] Trim the serve README's input_reference bullet Same cut as the request schema: drop how routing and decoding are arranged (container signatures, the ftyp collision, Pillow/NVDEC by name) and keep what a caller sends. The two bullets ran five lines each while every neighbour is a one-liner. Deps stay documented in the Cosmos3 README's Media I/O section, which is where someone asking "what do I need installed" would look. Signed-off-by: Igor Shovkun --- examples/visual_gen/serve/README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/visual_gen/serve/README.md b/examples/visual_gen/serve/README.md index 09fbd87f7667..fa9c55c8ca70 100644 --- a/examples/visual_gen/serve/README.md +++ b/examples/visual_gen/serve/README.md @@ -286,8 +286,8 @@ You can customize these by: - `frame_rate` (canonical) or `fps` (alias): frames per second - `num_frames`: when set, wins over the `seconds * frame_rate` derivation - `seed`, `num_inference_steps`, `guidance_scale`, `max_sequence_length`, `negative_prompt`: per-request denoise controls -- `input_reference`: Reference image (I2V/TI2V) or video (V2V), routed by container signature — filename and content type are ignored; accepted as base64-encoded string in JSON or as a file in multipart form-data. The signature only routes; the bytes pass through and the worker's decoder is what accepts them — Pillow for images, NVDEC for video (PyNvVideoCodec, a declared dependency). An unrecognized container is rejected with HTTP 400 at the boundary. - - **Supported formats**: reference images are tested with PNG (`image/png`) and JPEG (`image/jpeg`) and are decoded on the workers with Pillow. Reference videos are routed by container signature (MP4 `video/mp4`, AVI `video/x-msvideo`) and decoded on the workers' NVDEC — tested codec is H.264 in both containers; other codecs/profiles depend on the GPU's decoder capabilities and are best-effort. Filename and MIME metadata are never used for routing. HEIF/AVIF still images (which share the ISO-BMFF `ftyp` signature with MP4) are detected and rejected with a 400 asking for PNG or JPEG, rather than being sent to the video decoder. +- `input_reference`: Reference image (I2V/TI2V) or video (V2V), accepted as a base64-encoded string in JSON or as a file in multipart form-data + - **Supported formats**: PNG and JPEG images; MP4 and AVI video, with H.264 the tested codec and others best-effort. HEIF/AVIF are not supported. - `extra_params`: model-specific overflow (see below) - `response_format`: `"b64_json"` or `"url"` - `format`: Generation content encoding. Video encoders: `"mp4"`, `"avi"`, `"auto"`. Tensor formats: `"safetensors"`, `"pt"` (carries video + audio + scalar metadata in one payload for LTX-2). From 55fb3b502581e6340a513f173afc8c41e913a5be Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 30 Jul 2026 15:30:36 -0700 Subject: [PATCH 63/64] Resolve distilled sampling from the checkpoint, restore action_gen Two things the merge left behind. The pipeline stopped reading action_gen off the checkpoint config, so a checkpoint that ships action modules initialized with it off. 79b7d1aeba restored the transformer half of this and missed the pipeline; restore upstream's block verbatim. forward() also resolved an unset step count and guidance scale from the mode tables, which carry base values. validate_request only rejects values that *conflict* with a distilled checkpoint, so leaving both unset passed and then ran 35 steps at guidance 6.0 against weights with guidance baked in. The step count was masked -- set_timesteps ignores it for distilled and programs the fixed sigmas -- but the guidance was not. Requests through infer() were unaffected: they arrive carrying the checkpoint's defaults via default_generation_params. Resolve from the same policy before the mode tables, so both entry points agree. test_cosmos3_distilled.py covers this, and joins the b200 list -- it was in no test list, which is why the earlier infer() regression reached review. Signed-off-by: Igor Shovkun --- .../models/cosmos3/pipeline_cosmos3.py | 16 ++++ tensorrt_llm/media/__init__.py | 9 +- .../test_lists/test-db/l0_b200.yml | 1 + .../visual_gen/test_cosmos3_distilled.py | 93 +++++++++++++++++++ .../visual_gen/test_cosmos3_pipeline.py | 6 ++ 5 files changed, 121 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index 05663c6633aa..97ca3fed6482 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -118,6 +118,10 @@ def __init__(self, pipeline_config): logger.info("Initializing Cosmos3OmniMoTPipeline with audio generation.") self.audio_gen = True + if getattr(primary_pretrained_config, "action_gen", False): + logger.info("Initializing Cosmos3OmniMoTPipeline with action generation.") + self.action_gen = True + super().__init__(pipeline_config) def _init_transformer(self) -> None: @@ -782,6 +786,18 @@ def forward( self.sampling.validate_request(num_inference_steps, guidance_scale) + # A distilled checkpoint's step count and guidance are checkpoint facts, + # not mode defaults, so they resolve before the mode tables below. + # Requests through infer() already carry them via + # default_generation_params; a direct forward() call would otherwise + # fall through to the base tables and run CFG at 6.0 against weights + # with guidance baked in. + checkpoint_defaults = self.sampling.generation_default_overrides() + if num_inference_steps is None: + num_inference_steps = checkpoint_defaults.get("num_inference_steps") + if guidance_scale is None: + guidance_scale = checkpoint_defaults.get("guidance_scale") + if image is not None and self.sampling.is_distilled: raise ValueError( "Image-conditioned generation is not supported on distilled Cosmos3 " diff --git a/tensorrt_llm/media/__init__.py b/tensorrt_llm/media/__init__.py index 4ecd06ee3b58..e99f96800a2b 100644 --- a/tensorrt_llm/media/__init__.py +++ b/tensorrt_llm/media/__init__.py @@ -1,8 +1,9 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Media encoding utilities for TensorRT-LLM. +"""Media encoding and decoding utilities for TensorRT-LLM. -Free functions for encoding tensors to image / video files or in-memory bytes. -Internal-by-convention: not re-exported from ``tensorrt_llm`` so the public -API surface is reached via :class:`tensorrt_llm.visual_gen.VisualGenOutput`. +Free functions for encoding tensors to image / video files or in-memory bytes, +and for decoding encoded video back to frames. Internal-by-convention: not +re-exported from ``tensorrt_llm`` so the public API surface is reached via +:class:`tensorrt_llm.visual_gen.VisualGenOutput`. """ diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 8e39b9b84a21..c5709fe4a2f5 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -230,6 +230,7 @@ l0_b200: - unittest/_torch/visual_gen/test_wan_transformer.py - unittest/_torch/visual_gen/test_cosmos3_transformer.py - unittest/_torch/visual_gen/test_cosmos3_pipeline.py + - unittest/_torch/visual_gen/test_cosmos3_distilled.py - examples/visual_gen/test_visual_gen.py::test_wan_t2v_example - examples/visual_gen/test_visual_gen.py::test_flux1_example - examples/visual_gen/test_visual_gen.py::test_flux2_example diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py index 955eec6c29ab..5534c8fc900b 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py @@ -599,3 +599,96 @@ def test_model_index_class_name_dispatches(self, tmp_path): def test_hf_id_registered(self): entry = PIPELINE_REGISTRY["Cosmos3OmniMoTPipeline"] assert "nvidia/Cosmos3-Super-Text2Image-4Step" in entry.hf_ids + + +class TestDistilledForwardDefaults: + """A distilled checkpoint's step count and guidance are checkpoint facts. + + Requests through ``infer()`` carry them already, merged from + ``default_generation_params``. A direct ``forward()`` leaving them unset + passes ``validate_request`` (it only rejects *conflicting* values), so + without checkpoint-first resolution it would fall through to the base mode + tables and run CFG at 6.0 against weights with guidance baked in. + """ + + def _resolved_recipe(self, monkeypatch, **forward_kwargs): + """Run forward() far enough to see the recipe it reports, then bail.""" + import tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 as mod + + pipeline = _bare_pipeline(sampling=_distilled_policy()) + pipeline.transformer = SimpleNamespace(device=torch.device("cpu")) + pipeline.scheduler = SimpleNamespace(config=SimpleNamespace(flow_shift=1.0)) + + class StopAfterDims(Exception): + pass + + def stop(*args, **kwargs): + raise StopAfterDims + + pipeline._tokenize_prompt = stop + + lines = [] + monkeypatch.setattr(mod, "logger", SimpleNamespace(info=lines.append, warning=print)) + + with pytest.raises(StopAfterDims): + pipeline.forward( + prompt="a distilled render", + negative_prompt="", + num_frames=COSMOS3_720P_PARAMS["num_frames"], + seed=1, + max_sequence_length=8, + frame_rate=COSMOS3_720P_PARAMS["frame_rate"], + use_duration_template=False, + use_resolution_template=False, + use_system_prompt=False, + use_guardrails=False, + **forward_kwargs, + ) + dims = next(line for line in lines if "Cosmos3 generation dims" in line) + return dims + + def test_unset_resolves_to_checkpoint_not_base_table(self, monkeypatch): + dims = self._resolved_recipe(monkeypatch, num_inference_steps=None, guidance_scale=None) + assert f"num_inference_steps={len(DISTILLED_SIGMAS)}" in dims + assert f"guidance_scale={DISTILLED_GUIDANCE_SCALE:.2f}" in dims + # The base video table's values must not have been substituted. + assert f"num_inference_steps={COSMOS3_720P_PARAMS['num_inference_steps']}" not in dims + assert f"guidance_scale={COSMOS3_720P_PARAMS['guidance_scale']:.2f}" not in dims + + def test_base_checkpoint_still_uses_the_mode_table(self, monkeypatch): + """The checkpoint-first step must be a no-op for a base checkpoint: + ``generation_default_overrides()`` is empty, so the mode tables win.""" + import tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 as mod + + pipeline = _bare_pipeline() + pipeline.transformer = SimpleNamespace(device=torch.device("cpu")) + pipeline.scheduler = SimpleNamespace(config=SimpleNamespace(flow_shift=1.0)) + + class StopAfterDims(Exception): + pass + + def stop(*args, **kwargs): + raise StopAfterDims + + pipeline._tokenize_prompt = stop + lines = [] + monkeypatch.setattr(mod, "logger", SimpleNamespace(info=lines.append, warning=print)) + + with pytest.raises(StopAfterDims): + pipeline.forward( + prompt="a base render", + negative_prompt="", + num_frames=COSMOS3_720P_PARAMS["num_frames"], + num_inference_steps=None, + guidance_scale=None, + seed=1, + max_sequence_length=8, + frame_rate=COSMOS3_720P_PARAMS["frame_rate"], + use_duration_template=False, + use_resolution_template=False, + use_system_prompt=False, + use_guardrails=False, + ) + dims = next(line for line in lines if "Cosmos3 generation dims" in line) + assert f"num_inference_steps={COSMOS3_720P_PARAMS['num_inference_steps']}" in dims + assert f"guidance_scale={COSMOS3_720P_PARAMS['guidance_scale']:.2f}" in dims diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index e1f95f593178..8b2a56251149 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -686,6 +686,9 @@ class FakeSampling: def validate_request(self, num_inference_steps, guidance_scale): return None + def generation_default_overrides(self): + return {} + def set_flow_shift(self, scheduler, target, *, use_karras_sigmas=None): calls.append((target, use_karras_sigmas)) return scheduler @@ -741,6 +744,9 @@ class FakeSampling: def validate_request(self, num_inference_steps, guidance_scale): return None + def generation_default_overrides(self): + return {} + def set_flow_shift(self, scheduler, target, *, use_karras_sigmas=None): rebuilt.append((scheduler.name, target, use_karras_sigmas)) return scheduler From 43956e91bc53fd234ba83d1dcf9a5e30b7fd8748 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 30 Jul 2026 15:45:45 -0700 Subject: [PATCH 64/64] Stop duplicating the Cosmos3 README in the example The example's module docstring restated the README sitting next to it -- modes, checkpoints, guardrail setup, deployment configs and a worked command line per mode, right down to naming cosmos3-nano-1gpu.yaml eight times in each. Two copies of the same specifics drift; the README is the one people find. Keep a pointer and let --help cover the flags. Also revert the PipelineOutput docstring rewrap: it only moved a line break, and dropping it takes _torch/visual_gen/output.py out of this PR. Signed-off-by: Igor Shovkun --- examples/visual_gen/models/cosmos3/cosmos3.py | 89 ++----------------- tensorrt_llm/_torch/visual_gen/output.py | 5 +- 2 files changed, 7 insertions(+), 87 deletions(-) diff --git a/examples/visual_gen/models/cosmos3/cosmos3.py b/examples/visual_gen/models/cosmos3/cosmos3.py index f0bb6b2bc8aa..d3aed83a9a6b 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3.py +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -13,91 +13,12 @@ # 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"""Cosmos3 Text(+Image)-to-Video(+Audio) generation. +"""Cosmos3 Text(+Image/Video)-to-Video(+Audio) generation. -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 the - first (or last, per ``condition_video_keep``) frames of a reference video - via ``--video_path`` (a local MP4/AVI file; its encoded bytes pass through - and each worker decodes the conditioning window on NVDEC). -- **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 (pass the Hub ID or local path via ``--model``): - -- `nvidia/Cosmos3-Nano `_ -- `nvidia/Cosmos3-Super `_ - -Guardrails are enabled by default (required by the -`NVIDIA Open Model License Agreement -`_). -Install and authenticate as follows:: - - pip install cosmos_guardrail==0.3.0 && pip uninstall opencv-python - -Accept the terms for the guardrail checkpoint at -https://huggingface.co/nvidia/Cosmos-1.0-Guardrail and set a valid ``HF_TOKEN`` -(the checkpoint is downloaded automatically on first run). - -To run without guardrails (you are responsible for safe deployment):: - - export TRTLLM_DISABLE_COSMOS3_GUARDRAILS=1 - -Deployment configs (``examples/visual_gen/configs/``): - -- ``cosmos3-nano-1gpu.yaml`` — 1 GPU -- ``cosmos3-super-4gpu.yaml`` — 4 GPU, CFG + Ulysses + parallel VAE - -Example prompts live under ``prompts/`` (mirroring ``cosmos3-internal/inputs/omni``). - -Usage:: - - # T2V: text-to-video - python cosmos3.py --model nvidia/Cosmos3-Nano \ - --prompt_file prompts/t2v.json \ - --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml - - # I2V/TI2V: image-conditioned video (vision_path is read from the prompt file; - # local path, file://, http(s):// URL, or data: URI are all accepted) - python cosmos3.py --model nvidia/Cosmos3-Nano \ - --prompt_file prompts/i2v.json \ - --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml - - # I2V with an explicit conditioning image (overrides the prompt file) - python cosmos3.py --model nvidia/Cosmos3-Nano \ - --prompt_file prompts/i2v.json \ - --image_path https://example.com/frame.jpg \ - --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml - - # V2V: video-conditioned video (continues the first frames of --video_path) - python cosmos3.py --model nvidia/Cosmos3-Nano \ - --prompt_file prompts/v2v.json \ - --video_path /path/to/reference.mp4 \ - --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml - - # T2AV: text-to-video with synchronized audio - python cosmos3.py --model nvidia/Cosmos3-Nano \ - --prompt_file prompts/t2av.json \ - --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml - - # T2I: text-to-image - python cosmos3.py --model nvidia/Cosmos3-Nano \ - --prompt_file prompts/t2i.json \ - --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml \ - --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" \ - --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml +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. """ import argparse diff --git a/tensorrt_llm/_torch/visual_gen/output.py b/tensorrt_llm/_torch/visual_gen/output.py index 64963edc76be..b77956b7b6d8 100644 --- a/tensorrt_llm/_torch/visual_gen/output.py +++ b/tensorrt_llm/_torch/visual_gen/output.py @@ -37,9 +37,8 @@ class PipelineOutput: """Internal per-pipeline output. Each pipeline ``infer()`` populates the media tensor it produces plus - the metadata it owns (``frame_rate``, ``audio_sample_rate``) - and the three CUDA-event-measured timing phases that decompose - ``pipeline.infer()``. + the metadata it owns (``frame_rate``, ``audio_sample_rate``) and the + three CUDA-event-measured timing phases that decompose ``pipeline.infer()``. Attributes: image: Generated image as ``torch.Tensor`` shape ``(B, H, W, C)``,