From d2a42c3d650d0a004490a174f3024e951ad5d0e2 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Thu, 11 Jun 2026 08:11:55 -0700 Subject: [PATCH 01/35] cosmos3 action init Signed-off-by: Shreyas Misra --- examples/visual_gen/models/cosmos3/README.md | 35 +- examples/visual_gen/models/cosmos3/cosmos3.py | 180 ++++++ .../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 | 547 ++++++++++++++++-- .../models/cosmos3/transformer_cosmos3.py | 247 +++++++- 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 +++++ 16 files changed, 1716 insertions(+), 78 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/README.md b/examples/visual_gen/models/cosmos3/README.md index 69be21fe4880..6f4d8e54a1a2 100644 --- a/examples/visual_gen/models/cosmos3/README.md +++ b/examples/visual_gen/models/cosmos3/README.md @@ -1,11 +1,12 @@ -# Cosmos3 Text(+Image)-to-Video(+Audio) generation +# Cosmos3 Text(+Image)-to-Video(+Audio) and action 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. - **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 @@ -74,4 +75,34 @@ python cosmos3.py --model nvidia/Cosmos3-Nano \ 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 ``` diff --git a/examples/visual_gen/models/cosmos3/cosmos3.py b/examples/visual_gen/models/cosmos3/cosmos3.py index de9e9e5010ad..da242b1cb761 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3.py +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -21,9 +21,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 +83,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 +229,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 +302,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 +314,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 +335,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 +367,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..407cd53dbe4e 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,27 @@ 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, + 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 +884,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 +904,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 +1068,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 +1176,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 +1196,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 +1217,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 +1281,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 +1297,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 +1350,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..90d282b89980 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,23 @@ 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. 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." + ) 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 +1238,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 +1287,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 +1333,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 +1365,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 +1400,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 +1540,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 13cbe77bdf12..e38e82939866 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline.py @@ -1,3 +1,4 @@ +import inspect import itertools import os import time @@ -1042,8 +1043,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: @@ -1174,7 +1181,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..37dc235c16b5 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -36,7 +36,10 @@ 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.defaults import ( + COSMOS3_ACTION_PARAMS, + COSMOS3_T2I_PARAMS, +) from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import ( COSMOS3_DEFAULT_RESOLUTION_TEMPLATE, COSMOS3_DURATION_TEMPLATE, @@ -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 c369c86a3fdb6d969b93a37751ca4ab645cd4775 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Thu, 11 Jun 2026 09:26:53 -0700 Subject: [PATCH 02/35] 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 da242b1cb761..f9b81d751e5b 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3.py +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -21,7 +21,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 @@ -116,25 +121,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: @@ -270,14 +296,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 407cd53dbe4e..b4ff368f47d5 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) @@ -875,7 +754,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() @@ -917,15 +796,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: @@ -945,33 +847,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): @@ -1091,7 +984,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, @@ -1157,6 +1051,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") @@ -1283,7 +1179,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 90d282b89980..7740efc1a24b 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: @@ -1542,6 +1545,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 8816958dd63ddc0a80c095e01d8310dfdea14abb Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Thu, 11 Jun 2026 09:32:46 -0700 Subject: [PATCH 03/35] 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 f9b81d751e5b..8ec40344fb51 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3.py +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -305,6 +305,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, @@ -373,6 +379,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 b4ff368f47d5..0bc024240d46 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"), ) @@ -755,6 +756,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() @@ -778,6 +780,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( @@ -803,6 +806,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: @@ -815,6 +819,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']}" ) @@ -823,6 +828,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"] ) @@ -844,6 +850,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: @@ -864,6 +872,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}" ) @@ -1117,7 +1126,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 d7b3b4ba9537eea3d6e0fef25ca5ba99bced0d95 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Thu, 2 Jul 2026 11:21:51 -0700 Subject: [PATCH 04/35] add action to media API, example script fixes Signed-off-by: Shreyas Misra --- examples/visual_gen/models/cosmos3/cosmos3.py | 3 +- tensorrt_llm/media/tensor_payload.py | 11 ++-- .../_torch/visual_gen/test_tensor_payload.py | 56 +++++++++++++++++++ 3 files changed, 65 insertions(+), 5 deletions(-) diff --git a/examples/visual_gen/models/cosmos3/cosmos3.py b/examples/visual_gen/models/cosmos3/cosmos3.py index 8ec40344fb51..59a51c686054 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3.py +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -367,7 +367,6 @@ 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 @@ -379,6 +378,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_resolution is not None: + params.extra_params["action_resolution"] = args.action_resolution if args.action_fps is not None: params.extra_params["action_fps"] = args.action_fps if args.action_json is not None: diff --git a/tensorrt_llm/media/tensor_payload.py b/tensorrt_llm/media/tensor_payload.py index b61580d9f439..e78f02c6ff6a 100644 --- a/tensorrt_llm/media/tensor_payload.py +++ b/tensorrt_llm/media/tensor_payload.py @@ -5,7 +5,7 @@ Two payload formats are supported: - ``"safetensors"``: writes a single file with named tensors - (``image``/``video``/``audio``). Scalar metadata (``frame_rate``, + (``image``/``video``/``audio``/``action``). Scalar metadata (``frame_rate``, ``audio_sample_rate``) is stored two ways: as a 0-d tensor under the same key (so ``safetensors.torch.load(bytes)`` returns it alongside the media tensors — consumers call ``.item()`` to @@ -49,12 +49,14 @@ def is_tensor_format(fmt: Optional[str]) -> bool: # canonical ``(H, W, C)`` shape is unbatched at rank 3 and batched at # rank 4; video is unbatched at rank 4 ``(T, H, W, C)`` and batched at # rank 5; audio is unbatched at rank 2 ``(channels, T_audio)`` and -# batched at rank 3. The serializer uses these to decide whether a +# batched at rank 3; action is unbatched at rank 2 ``(T, action_dim)`` +# and batched at rank 3. The serializer uses these to decide whether a # media tensor has a true batch axis to slice along. _BATCHED_RANKS: Dict[str, int] = { "image": 4, "video": 5, "audio": 3, + "action": 3, } @@ -63,14 +65,15 @@ def _modalities(output: "VisualGenOutput") -> Tuple[Tuple[str, Optional[torch.Te ("image", output.image), ("video", output.video), ("audio", output.audio), + ("action", output.action), ) def infer_batch_size(output: "VisualGenOutput") -> int: """Return the leading batch dimension across the populated media tensors. - Image is batched only at rank 4, video at rank 5, audio at rank 3. - An unbatched media tensor reports a batch size of 1 so list-path + Image is batched only at rank 4, video at rank 5, audio and action + at rank 3. An unbatched media tensor reports a batch size of 1 so list-path callers can still ask for ``[0]`` and get a single-item payload. Raises :class:`ValueError` when *output* carries no media tensor. """ diff --git a/tests/unittest/_torch/visual_gen/test_tensor_payload.py b/tests/unittest/_torch/visual_gen/test_tensor_payload.py index 4ad482b886fd..d7b1913d5344 100644 --- a/tests/unittest/_torch/visual_gen/test_tensor_payload.py +++ b/tests/unittest/_torch/visual_gen/test_tensor_payload.py @@ -52,6 +52,17 @@ def _make_video_output(batch: int = 1, t: int = 2, h: int = 4, w: int = 4) -> Vi ) +def _make_action_output(batch: int = 1, t: int = 4, action_dim: int = 7) -> VisualGenOutput: + """Action uses ``(B, T, action_dim)``: batched at rank 3, unbatched at rank 2.""" + action = torch.arange(batch * t * action_dim, dtype=torch.float32).reshape(batch, t, action_dim) + return VisualGenOutput( + request_id=3, + action=action, + action_mode="policy", + raw_action_dim=action_dim, + ) + + class TestIsTensorFormat: def test_accepts_supported_tokens(self): for token in TENSOR_FORMATS: @@ -93,11 +104,56 @@ def test_inconsistent_batches_raise(self): with pytest.raises(ValueError, match="Inconsistent batch sizes"): infer_batch_size(output) + def test_action_rank_3_is_batched(self): + output = _make_action_output(batch=2) + assert infer_batch_size(output) == 2 + + def test_action_rank_2_is_unbatched(self): + """An unbatched action tensor has shape ``(T, action_dim)``; the + timestep axis must not be confused with a batch dimension.""" + output = VisualGenOutput(request_id=1, action=torch.zeros(4, 7)) + assert infer_batch_size(output) == 1 + + def test_action_only_output_infers_batch(self): + """Action counts as a media tensor, so an action-only output + (no image/video/audio) reports a batch size instead of raising.""" + assert infer_batch_size(_make_action_output(batch=3)) == 3 + def test_no_media_raises(self): with pytest.raises(ValueError, match="carries no media"): infer_batch_size(VisualGenOutput(request_id=1)) +@pytest.mark.parametrize("fmt", ["safetensors", "pt"]) +class TestActionRoundTrip: + """The ``action`` tensor must round-trip through both payload formats + alongside the other media modalities (regression: it was silently + dropped by :func:`serialize_visual_gen_output`).""" + + def _load(self, data: bytes, fmt: str) -> dict: + return _safetensors_load(data) if fmt == "safetensors" else _pt_load(data) + + def test_action_serialized_full(self, fmt): + output = _make_action_output(batch=2) + loaded = self._load(serialize_visual_gen_output(output, fmt), fmt) + assert "action" in loaded + assert torch.equal(loaded["action"], output.action) + + def test_action_sliced_drops_batch_axis(self, fmt): + output = _make_action_output(batch=3) + data = serialize_visual_gen_output(output, fmt, batch_index=1) + loaded = self._load(data, fmt) + assert loaded["action"].shape == (4, 7) + assert torch.equal(loaded["action"], output.action[1]) + + def test_unbatched_action_passthrough(self, fmt): + output = VisualGenOutput(request_id=1, action=torch.randn(4, 7)) + # batch_index set, but rank-2 action has no batch axis to slice. + loaded = self._load(serialize_visual_gen_output(output, fmt, batch_index=0), fmt) + assert loaded["action"].shape == (4, 7) + assert torch.equal(loaded["action"], output.action) + + @pytest.mark.parametrize("fmt", ["safetensors", "pt"]) class TestSingleSavePath: """A single path writes one logical output. Unbatched tensors and From e39caa89a27d42a3f52b842ff6a50101a3679523 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Tue, 14 Jul 2026 10:03:52 -0700 Subject: [PATCH 05/35] fix visualgen output test Signed-off-by: Shreyas Misra --- tests/unittest/visual_gen/test_output.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/unittest/visual_gen/test_output.py b/tests/unittest/visual_gen/test_output.py index e27f9c33e3b5..400e854791d8 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", } @@ -60,6 +64,10 @@ def test_minimal_construction_defaults(): assert out.image is not None assert out.video is None assert out.audio 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.frame_rate is None assert out.audio_sample_rate is None assert out.error is None From 944e5a856eb819464d0230ef586d483918647550 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Wed, 15 Jul 2026 11:09:01 -0700 Subject: [PATCH 06/35] unwaive test Signed-off-by: Shreyas Misra --- tests/integration/test_lists/waives.txt | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index e1d3b85239b6..047a63ea17fc 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -182,8 +182,6 @@ examples/test_ray.py::test_llm_inference_distributed_ray[pp2] SKIP (https://nvbu examples/test_ray.py::test_llm_inference_distributed_ray[tp2pp2] SKIP (https://nvbugs/6427411) examples/test_ray.py::test_ray_disaggregated_serving[tp2] SKIP (https://nvbugs/5612502) examples/test_whisper.py::test_llm_whisper_general[large-v3-disable_gemm_plugin-disable_attention_plugin-disable_weight_only-float16-nb:1-use_python_runtime] SKIP (https://nvbugs/5244570) -examples/visual_gen/test_visual_gen.py::test_cosmos3_nano_t2i_lpips_against_golden SKIP (https://nvbugs/6418815) -examples/visual_gen/test_visual_gen.py::test_cosmos3_nano_t2v_lpips_against_golden SKIP (https://nvbugs/6437341) examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_multi_gpu[attn2d_2x2] SKIP (https://nvbugs/6272644) examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_multi_gpu[cfg2_ulysses2] SKIP (https://nvbugs/6272644) examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_multi_gpu[cfg2_ulysses2_attn2d_2x1] SKIP (https://nvbugs/6272644) From 8b8aa44c3d792382b45aa778ecfe66df035c854a Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Wed, 15 Jul 2026 18:52:22 -0700 Subject: [PATCH 07/35] fix test Signed-off-by: Shreyas Misra --- tests/unittest/visual_gen/test_output.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/tests/unittest/visual_gen/test_output.py b/tests/unittest/visual_gen/test_output.py index 400e854791d8..c382107b291c 100644 --- a/tests/unittest/visual_gen/test_output.py +++ b/tests/unittest/visual_gen/test_output.py @@ -765,14 +765,18 @@ 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 == { + "request_id", "image", "video", "audio", + "action", "frame_rate", "audio_sample_rate", - "pre_denoise", - "denoise", - "post_denoise", + "raw_action_dim", + "action_mode", + "domain_id", + "error", + "metrics", } From f3b25da56bb5af4f50841de550d4ec7adbeec148 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Wed, 15 Jul 2026 18:52:32 -0700 Subject: [PATCH 08/35] Revert "unwaive test" This reverts commit c115ea698ff7e82e986de949d4017b86f175b2a0. Signed-off-by: Shreyas Misra --- tests/integration/test_lists/waives.txt | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 047a63ea17fc..e1d3b85239b6 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -182,6 +182,8 @@ examples/test_ray.py::test_llm_inference_distributed_ray[pp2] SKIP (https://nvbu examples/test_ray.py::test_llm_inference_distributed_ray[tp2pp2] SKIP (https://nvbugs/6427411) examples/test_ray.py::test_ray_disaggregated_serving[tp2] SKIP (https://nvbugs/5612502) examples/test_whisper.py::test_llm_whisper_general[large-v3-disable_gemm_plugin-disable_attention_plugin-disable_weight_only-float16-nb:1-use_python_runtime] SKIP (https://nvbugs/5244570) +examples/visual_gen/test_visual_gen.py::test_cosmos3_nano_t2i_lpips_against_golden SKIP (https://nvbugs/6418815) +examples/visual_gen/test_visual_gen.py::test_cosmos3_nano_t2v_lpips_against_golden SKIP (https://nvbugs/6437341) examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_multi_gpu[attn2d_2x2] SKIP (https://nvbugs/6272644) examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_multi_gpu[cfg2_ulysses2] SKIP (https://nvbugs/6272644) examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_multi_gpu[cfg2_ulysses2_attn2d_2x1] SKIP (https://nvbugs/6272644) From 6bf9ba9822558f6f02b24e8f80b487752bce694a Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Fri, 17 Jul 2026 13:39:04 -0700 Subject: [PATCH 09/35] address Cosmos3 action review comments Signed-off-by: Shreyas Misra --- examples/visual_gen/models/cosmos3/README.md | 8 +- examples/visual_gen/models/cosmos3/cosmos3.py | 89 ++++++------------- .../visual_gen/models/cosmos3/action.py | 50 ++++++++--- .../visual_gen/models/cosmos3/defaults.py | 70 +++++++++------ .../models/cosmos3/pipeline_cosmos3.py | 10 +++ tensorrt_llm/_torch/visual_gen/pipeline.py | 16 +--- tensorrt_llm/media/tensor_payload.py | 24 +++-- tensorrt_llm/serve/visual_gen_utils.py | 6 ++ tensorrt_llm/visual_gen/output.py | 18 ++-- tensorrt_llm/visual_gen/params.py | 22 +++++ .../test_lists/test-db/l0_b200.yml | 1 + .../_torch/visual_gen/test_cosmos3_action.py | 80 +++++++++++------ .../visual_gen/test_cosmos3_pipeline.py | 23 +++++ .../visual_gen/test_cosmos3_transformer.py | 29 ++++++ .../_torch/visual_gen/test_tensor_payload.py | 22 +++++ .../visual_gen/test_trtllm_serve_endpoints.py | 15 ++++ .../visual_gen/test_visual_gen_params.py | 37 ++++++++ tests/unittest/visual_gen/test_output.py | 22 +++-- 18 files changed, 384 insertions(+), 158 deletions(-) diff --git a/examples/visual_gen/models/cosmos3/README.md b/examples/visual_gen/models/cosmos3/README.md index 6f4d8e54a1a2..65d5595e7f99 100644 --- a/examples/visual_gen/models/cosmos3/README.md +++ b/examples/visual_gen/models/cosmos3/README.md @@ -6,7 +6,7 @@ Cosmos3 supports the following generation modes from a single checkpoint: - **T2I** — text-to-image (`prompts/t2i.json`); emits a still frame (use `--output_type image` / a non-video `--output_path`). - **I2V / TI2V** — image-conditioned video (`prompts/i2v.json`). Condition on a reference frame via the prompt file's `vision_path` or `--image_path`. The image may be a local path, a `file://` / `http(s)://` URL, or a `data:` URI. - **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. +- **Action** — policy / forward dynamics / inverse dynamics generation (pass `--action_mode`). Action and audio generation are mutually exclusive. Save action runs as `safetensors` or `pt` so the rollout and action tensor stay in one payload. ## Checkpoints @@ -83,7 +83,7 @@ python cosmos3.py --model nvidia/Cosmos3-Nano \ --action_mode policy \ --domain_name bridge_orig_lerobot \ --raw_action_dim 10 \ - --output_path policy_rollout.mp4 \ + --output_path policy_rollout.safetensors \ --action_output_path policy_action.json # Action — forward dynamics (first frame + action trajectory -> rollout video) @@ -93,7 +93,7 @@ python cosmos3.py --model nvidia/Cosmos3-Nano \ --action_mode forward_dynamics \ --domain_name av \ --action_json action_trajectory.json \ - --output_path forward_dynamics.mp4 + --output_path forward_dynamics.safetensors # Action — inverse dynamics (video -> predicted action) python cosmos3.py --model nvidia/Cosmos3-Nano \ @@ -103,6 +103,6 @@ python cosmos3.py --model nvidia/Cosmos3-Nano \ --action_mode inverse_dynamics \ --domain_name bridge_orig_lerobot \ --raw_action_dim 10 \ - --output_path inverse_video.mp4 \ + --output_path inverse_video.safetensors \ --action_output_path inverse_action.json ``` diff --git a/examples/visual_gen/models/cosmos3/cosmos3.py b/examples/visual_gen/models/cosmos3/cosmos3.py index 59a51c686054..1144c87eb0f6 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3.py +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -21,16 +21,10 @@ 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"}) +_ACTION_MODES = ("policy", "forward_dynamics", "inverse_dynamics") +_TENSOR_OUTPUT_SUFFIXES = {".pt", ".safetensors"} def _resolve_path(path: str) -> str: @@ -98,10 +92,11 @@ def _validate_action_args( 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: + if mode not in _ACTION_MODES: raise SystemExit( - f"Invalid --action_mode {args.action_mode!r}; expected one of {sorted(ACTION_MODES)}." + f"Invalid --action_mode {args.action_mode!r}; expected one of {list(_ACTION_MODES)}." ) + args.action_mode = mode if args.enable_audio: raise SystemExit("Cosmos3 does not support joint action and audio generation.") if args.output_type != "video": @@ -121,53 +116,29 @@ def _validate_action_args( 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." - ) + if args.raw_action_dim is None and args.domain_name is None and args.domain_id is None: + raise SystemExit(f"{mode} requires --raw_action_dim, --domain_name, or --domain_id.") 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." - ) + if args.raw_action_dim is None and args.domain_name is None and args.domain_id is None: + raise SystemExit(f"{mode} requires --raw_action_dim, --domain_name, or --domain_id.") -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 _resolved_output_path(path: str, action_mode: Optional[str]) -> str: + if action_mode is None: + return path + output_path = Path(path) + if output_path.suffix.lower() in _TENSOR_OUTPUT_SUFFIXES: + return str(output_path) + return str(output_path.with_suffix(".safetensors")) + + +def _default_action_output_path(output_path: str) -> str: + stem = Path(output_path) + return str(stem.with_suffix(".action.json")) def _save_action_output(output, path: str) -> None: @@ -259,7 +230,7 @@ def main(): "--action_mode", type=str, default=None, - choices=sorted(ACTION_MODES), + choices=list(_ACTION_MODES), help="Action mode: policy, forward_dynamics, or inverse_dynamics", ) parser.add_argument( @@ -284,7 +255,7 @@ def main(): "--action_chunk_size", type=int, default=None, - help=f"Action tokens to generate (default {COSMOS3_ACTION_PARAMS['action_chunk_size']})", + help="Action tokens to generate. Defaults to the domain preset or model default.", ) parser.add_argument( "--action_json", @@ -345,9 +316,8 @@ def main(): params = visual_gen.default_params if image_path is not None: params.image = image_path - - if args.action_mode is not None: - _apply_action_generation_params(params, args) + if args.action_mode is not None and args.video_path is not None: + params.image = args.video_path negative_prompt_path = _resolve_path(args.negative_prompt) if args.negative_prompt is not None: @@ -385,8 +355,6 @@ def main(): 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 @@ -400,11 +368,12 @@ def main(): params=params, ) - output.save(args.output_path) - print(f"Saved: {args.output_path}") + output_path = _resolved_output_path(args.output_path, args.action_mode) + output.save(output_path) + print(f"Saved: {output_path}") if args.action_mode is not None: - action_path = args.action_output_path or _default_action_output_path(args.output_path) + action_path = args.action_output_path or _default_action_output_path(output_path) _save_action_output(output, action_path) if output.action is not None: print(f"Saved action: {action_path}") diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py index b66dda5b6ac2..fea4a71a5ec6 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py @@ -6,7 +6,7 @@ from __future__ import annotations from pathlib import Path -from typing import Any, List, Optional +from typing import Any import numpy as np import PIL.Image @@ -249,8 +249,8 @@ def find_closest_target_size(h: int, w: int, resolution: str | int) -> tuple[int def pil_to_rgb(value: Any) -> PIL.Image.Image: - if isinstance(value, str): - return PIL.Image.open(value).convert("RGB") + if isinstance(value, (str, Path)): + return PIL.Image.open(Path(value)).convert("RGB") if isinstance(value, PIL.Image.Image): return value.convert("RGB") raise TypeError( @@ -258,7 +258,7 @@ def pil_to_rgb(value: Any) -> PIL.Image.Image: ) -def decode_action_video_file(path: Path, max_frames: Optional[int] = None) -> List[PIL.Image.Image]: +def decode_action_video_file(path: Path, max_frames: int | None = None) -> list[PIL.Image.Image]: import torchvision.io as io frames, _, _ = io.read_video(str(path), pts_unit="sec") @@ -269,7 +269,7 @@ def decode_action_video_file(path: Path, max_frames: Optional[int] = None) -> Li 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]: +def normalize_action_video_path(path: Path, max_frames: int | None = None) -> list[Any]: if not path.exists(): raise ValueError(f"Cosmos3 action video path does not exist: {path}") if path.is_dir(): @@ -293,7 +293,7 @@ def normalize_action_video_path(path: Path, max_frames: Optional[int] = None) -> ) -def normalize_action_video_input(video: Any, max_frames: Optional[int] = None) -> List[Any]: +def normalize_action_video_input(video: Any, max_frames: int | None = 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, @@ -313,8 +313,8 @@ def normalize_action_video_input(video: Any, max_frames: Optional[int] = None) - def resolve_action_size( - height: Optional[int], - width: Optional[int], + height: int | None, + width: int | None, ref_image: PIL.Image.Image, action_resolution: int, ) -> tuple[int, int]: @@ -349,10 +349,10 @@ def action_reference_image( 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 isinstance(source, (str, Path)): path = Path(source) if path.is_file() and path.suffix.lower() in ACTION_IMAGE_EXTENSIONS: - return PIL.Image.open(source).convert("RGB") + return PIL.Image.open(path).convert("RGB") frames = normalize_action_video_input(source, max_frames=1) if not frames: raise ValueError( @@ -392,13 +392,36 @@ def prepare_action_latents( *, mode: str, action_chunk_size: int, - raw_action_dim: Optional[int], + raw_action_dim: int | None, action_dim: int, generator: torch.Generator, device: torch.device, dtype: torch.dtype, action_input: Any = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: + """Prepare action latents and conditioning masks for a denoise request. + + Args: + mode: Cosmos3 action mode. ``forward_dynamics`` conditions all action + tokens on ``action_input``; ``policy`` and ``inverse_dynamics`` start + from action noise and predict the raw action dimensions. + action_chunk_size: Number of action tokens in the generated trajectory. + raw_action_dim: Number of semantic action dimensions before padding. + Required for ``policy`` and ``inverse_dynamics``. For + ``forward_dynamics``, omitted values are inferred from + ``action_input.shape[-1]`` and explicit values must match it. + action_dim: Model action embedding width after zero-padding. + generator: Random generator used for action noise. + device: Target device for the returned tensors. + dtype: Target dtype for the returned tensors. + action_input: Forward-dynamics action trajectory with shape ``[T, D]``. + + Returns: + Tuple of ``(action_latents, action_velocity_mask, clean_action, + raw_action_dim)``. The tensors have shape ``[1, action_chunk_size, + action_dim]`` except the mask, which has shape ``[1, action_chunk_size, + 1]``. + """ if mode == ACTION_MODE_FORWARD_DYNAMICS: action = load_action_tensor(action_input) if action.shape[0] < action_chunk_size: @@ -408,6 +431,11 @@ def prepare_action_latents( action = action[:action_chunk_size] if raw_action_dim is None: raw_action_dim = int(action.shape[-1]) + elif int(raw_action_dim) != int(action.shape[-1]): + raise ValueError( + "Cosmos3 forward_dynamics raw_action_dim must match action input width; " + f"got raw_action_dim={raw_action_dim}, action width={action.shape[-1]}." + ) clean_action = pad_action_to_dim(action, action_dim) else: if raw_action_dim is None: diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index b01e98334276..e0e8c06ba384 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -26,7 +26,7 @@ configs (bridge, av, droid, libero, etc.). """ -from typing import Any, Dict, List, Optional, TypedDict +from typing import Any, TypedDict from tensorrt_llm._torch.visual_gen.models.cosmos3.action import ( COSMOS3_ACTION_RESOLUTIONS, @@ -95,7 +95,7 @@ class Cosmos3DomainPreset(TypedDict, total=False): # Training-aligned defaults. Values mirror Cosmos3 omni action JSON examples where available. -COSMOS3_DOMAIN_PRESETS: Dict[str, Cosmos3DomainPreset] = { +COSMOS3_DOMAIN_PRESETS: dict[str, Cosmos3DomainPreset] = { # WidowX bridge; 7-DOF arm + gripper in 10-D state. "bridge_orig_lerobot": { "raw_action_dim": 10, @@ -179,7 +179,7 @@ class Cosmos3DomainPreset(TypedDict, total=False): } # Map alias domain_name keys to a canonical preset entry. -COSMOS3_DOMAIN_PRESET_ALIASES: Dict[str, str] = { +COSMOS3_DOMAIN_PRESET_ALIASES: dict[str, str] = { "robomind-franka": "droid_lerobot", "robomind-franka-dual": "droid_lerobot", "robomind-ur": "droid_lerobot", @@ -190,9 +190,9 @@ class Cosmos3DomainPreset(TypedDict, total=False): def canonical_domain_preset_key( - domain_name: Any = None, - domain_id: Any = None, -) -> Optional[str]: + domain_name: str | None = None, + domain_id: str | int | None = None, +) -> str | None: 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) @@ -207,7 +207,7 @@ def canonical_domain_preset_key( if resolved_id == 0: return None - candidates: List[str] = [] + candidates: list[str] = [] for name, mapped_id in EMBODIMENT_TO_DOMAIN_ID.items(): if mapped_id != resolved_id: continue @@ -221,9 +221,9 @@ def canonical_domain_preset_key( def get_domain_preset( - domain_name: Any = None, - domain_id: Any = None, -) -> Optional[Cosmos3DomainPreset]: + domain_name: str | None = None, + domain_id: str | int | None = None, +) -> Cosmos3DomainPreset | None: key = canonical_domain_preset_key(domain_name, domain_id) if key is None: return None @@ -232,19 +232,29 @@ def get_domain_preset( 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]: + domain_name: str | None = None, + domain_id: str | int | None = None, + raw_action_dim: int | None = None, + action_chunk_size: int | None = None, + action_resolution: int | None = None, + frame_rate: float | None = None, + action_fps: float | None = None, + num_frames: int | None = 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] = [] + warnings: list[str] = [] + + domain_requested = (domain_name is not None and str(domain_name).strip() != "") or ( + domain_id is not None and str(domain_id).strip() not in {"", "0"} + ) + if domain_requested and preset is None: + warnings.append( + "Cosmos3 action domain preset was not found for " + f"domain_name={domain_name!r}, domain_id={domain_id!r}; " + "using generic action defaults for omitted fields." + ) def _resolve_field( field: str, @@ -288,6 +298,16 @@ def _resolve_field( resolved_action_fps = ( float(action_fps) if action_fps is not None else float(resolved_frame_rate) ) + if resolved_raw_action_dim is not None and int(resolved_raw_action_dim) <= 0: + raise ValueError(f"Cosmos3 raw_action_dim must be positive, got {resolved_raw_action_dim}.") + if int(resolved_chunk) <= 0: + raise ValueError(f"Cosmos3 action_chunk_size must be positive, got {resolved_chunk}.") + if float(resolved_frame_rate) <= 0.0: + raise ValueError(f"Cosmos3 frame_rate must be positive, got {resolved_frame_rate}.") + if resolved_action_fps <= 0.0: + raise ValueError(f"Cosmos3 action_fps must be positive, got {resolved_action_fps}.") + if int(resolved_num_frames) <= 0: + raise ValueError(f"Cosmos3 num_frames must be positive, got {resolved_num_frames}.") return { "raw_action_dim": resolved_raw_action_dim, @@ -301,7 +321,7 @@ def _resolve_field( } -COSMOS3_EXTRA_SPECS: Dict[str, ExtraParamSchema] = { +COSMOS3_EXTRA_SPECS: dict[str, ExtraParamSchema] = { "use_duration_template": ExtraParamSchema( type="bool", default=True, @@ -333,7 +353,7 @@ def _resolve_field( description="Output modality: 'video' (T2V/I2V) or 'image' (text-to-image).", ), "action_mode": ExtraParamSchema( - type="str", + type="Literal['policy', 'forward_dynamics', 'inverse_dynamics']", default=None, description="Action generation mode: policy, forward_dynamics, or inverse_dynamics.", ), @@ -361,7 +381,7 @@ def _resolve_field( ), "action_chunk_size": ExtraParamSchema( type="int", - default=COSMOS3_ACTION_PARAMS["action_chunk_size"], + default=None, description=( "Number of action tokens to generate (16 for most robots, 60 for av/camera_pose). " "Inferred from domain_name preset when omitted." @@ -373,8 +393,8 @@ def _resolve_field( description="Action trajectory [T, D] for forward_dynamics mode.", ), "action_resolution": ExtraParamSchema( - type="int", - default=480, + type="Literal[256, 480, 704, 720]", + default=None, description=( "Resolution bucket for action image sizing. Must be one of " f"{list(COSMOS3_ACTION_RESOLUTIONS)}. Inferred from domain_name preset when omitted." 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 0bc024240d46..a22fb6810aa9 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -855,6 +855,11 @@ def forward( action_ref_image = None if do_action: + if isinstance(image, torch.Tensor) or isinstance(video, torch.Tensor): + raise ValueError( + "Cosmos3 action generation does not support tensor image/video inputs; " + "pass a PIL image, image path, frame directory, video path, or frame list." + ) action_ref_image = action_reference_image( action_mode=normalized_action_mode, image=image, @@ -995,6 +1000,11 @@ 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) + if len(video) < num_frames: + raise ValueError( + "Cosmos3 inverse_dynamics requires at least " + f"{num_frames} frames, got {len(video)}." + ) video_tensor = self._preprocess_action_video(video, height, width) latents, velocity_mask, condition_latents = self._prepare_latents_action_video( video_tensor, diff --git a/tensorrt_llm/_torch/visual_gen/pipeline.py b/tensorrt_llm/_torch/visual_gen/pipeline.py index e38e82939866..8233042f346c 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline.py @@ -1,4 +1,3 @@ -import inspect import itertools import os import time @@ -1044,12 +1043,9 @@ def denoise( 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 + 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. @@ -1181,11 +1177,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, extra_stream_latents = post_step_fn(latents, extra_stream_latents) # Logging if self.rank == 0: diff --git a/tensorrt_llm/media/tensor_payload.py b/tensorrt_llm/media/tensor_payload.py index e78f02c6ff6a..f94ea52dce66 100644 --- a/tensorrt_llm/media/tensor_payload.py +++ b/tensorrt_llm/media/tensor_payload.py @@ -5,13 +5,13 @@ Two payload formats are supported: - ``"safetensors"``: writes a single file with named tensors - (``image``/``video``/``audio``/``action``). Scalar metadata (``frame_rate``, - ``audio_sample_rate``) is stored two ways: as a 0-d tensor under - the same key (so ``safetensors.torch.load(bytes)`` returns it - alongside the media tensors — consumers call ``.item()`` to - unbox) and as a stringified value in the file header (preserved - for callers using ``safe_open(...).metadata()``). No pickle on - load. + (``image``/``video``/``audio``/``action``). Numeric scalar metadata + (``frame_rate``, ``audio_sample_rate``, ``raw_action_dim``, ``domain_id``) + is stored two ways: as a 0-d tensor under the same key (so + ``safetensors.torch.load(bytes)`` returns it alongside the media tensors — + consumers call ``.item()`` to unbox) and as a stringified value in the file + header (preserved for callers using ``safe_open(...).metadata()``). String + metadata such as ``action_mode`` is header-only. No pickle on load. - ``"pt"``: writes a single file via :func:`torch.save` with the same tensor keys plus scalar metadata as native Python values. Clients should load with ``torch.load(buf, weights_only=True)`` @@ -146,6 +146,12 @@ def _collect_tensors_and_metadata( metadata["frame_rate"] = float(frame_rate) if audio_sample_rate is not None: metadata["audio_sample_rate"] = int(audio_sample_rate) + if output.raw_action_dim is not None: + metadata["raw_action_dim"] = int(output.raw_action_dim) + if output.domain_id is not None: + metadata["domain_id"] = int(output.domain_id) + if output.action_mode is not None: + metadata["action_mode"] = str(output.action_mode) return tensors, metadata @@ -205,7 +211,9 @@ def serialize_visual_gen_output( # ``loaded["frame_rate"].item()`` directly) and as a string in the # file header (preserved for callers that already use # ``safe_open(...).metadata()``). The two views always agree. - scalar_tensors = {k: torch.as_tensor(v) for k, v in metadata.items()} + scalar_tensors = { + k: torch.as_tensor(v) for k, v in metadata.items() if isinstance(v, (int, float)) + } return safetensors_save( {**tensors, **scalar_tensors}, metadata={k: str(v) for k, v in metadata.items()}, diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index 3094bf66cb04..3a699d4bdfe4 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -164,6 +164,12 @@ def parse_visual_gen_params( with open(ref_path, "wb") as f: shutil.copyfileobj(request.input_reference.file, f) params.image = ref_path + if request.extra_params and request.extra_params.get("video") is not None: + raise ValueError( + "extra_params['video'] is not accepted over trtllm-serve because it would " + "reference server-local paths. Use input_reference for uploaded conditioning " + "media, or call the offline VisualGen API for frame-sequence inputs." + ) _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/output.py b/tensorrt_llm/visual_gen/output.py index 9d1f6e99336a..099d40aa3a4c 100644 --- a/tensorrt_llm/visual_gen/output.py +++ b/tensorrt_llm/visual_gen/output.py @@ -125,8 +125,8 @@ def save( format: Explicit format. Image encoders: ``"png"``, ``"jpg"``, ``"webp"``. Video encoders: ``"mp4"``, ``"avi"``. Tensor payloads: ``"safetensors"``, ``"pt"`` — these carry every - populated modality (image/video/audio) plus scalar - metadata (frame_rate, audio_sample_rate) in one file. + populated modality (image/video/audio/action) plus scalar + metadata in one file. frame_rate: Override the frame rate for video output. Defaults to ``self.frame_rate`` when not provided. audio_sample_rate: Override the audio sample rate. Defaults to @@ -142,8 +142,8 @@ def save( ValueError: When video output lacks a frame rate, when the output carries no media tensor at all, or when the list length does not match the batch size. - NotImplementedError: When the output is audio-only and a - non-tensor format is requested. + NotImplementedError: When the output is audio-only or contains + action data and a non-tensor format is requested. """ from tensorrt_llm.media.encoding import save_image, save_images, save_video, save_videos from tensorrt_llm.media.tensor_payload import is_tensor_format @@ -156,7 +156,7 @@ def save( is_batch = isinstance(path, list) # Tensor formats carry every populated modality in one payload, - # so the dispatch table for image/video/audio below does not + # so the dispatch table for image/video/audio/action below does not # apply. When ``format`` is omitted, infer it from the path # suffix so callers using the documented extension convention # (``out.safetensors``/``out.pt``) reach the tensor path. @@ -170,6 +170,12 @@ def save( audio_sample_rate=audio_sample_rate, ) + if self.action is not None: + raise NotImplementedError( + "Saving action outputs requires a tensor payload format " + "('safetensors' or 'pt') so the action tensor and metadata are preserved." + ) + if self.image is not None: if is_batch: saved_list = save_images( @@ -213,7 +219,7 @@ def save( raise ValueError( f"Cannot save output: request {self.request_id} carries no media " - "(image/video/audio are all None)." + "(image/video/audio/action are all None)." ) def _save_tensor_payload( diff --git a/tensorrt_llm/visual_gen/params.py b/tensorrt_llm/visual_gen/params.py index 9e71a25d0505..ca8993c8a820 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 ast from typing import Any, Dict, List, Optional, Union from pydantic import Field @@ -93,6 +94,7 @@ class VisualGenParams(StrictBaseModel): "bool": (bool,), "str": (str,), "list": (list,), + "path_or_list": (str, list), } # Generation config fields that pipelines declare defaults for. If a user @@ -112,6 +114,18 @@ class VisualGenParams(StrictBaseModel): ) +def _literal_choices(type_expr: str) -> tuple[Any, ...] | None: + if not type_expr.startswith("Literal[") or not type_expr.endswith("]"): + return None + + literal_body = type_expr[len("Literal[") : -1] + try: + choices = ast.literal_eval(f"({literal_body},)") + except (SyntaxError, ValueError): + return None + return choices if isinstance(choices, tuple) else (choices,) + + def validate_visual_gen_params( params: VisualGenParams, *, @@ -166,6 +180,14 @@ def validate_visual_gen_params( # Skip None values (param left at its None default) if value is None: continue + literal_choices = _literal_choices(spec.type) + if literal_choices is not None: + if value not in literal_choices: + messages.append( + f"extra_params['{key}'] expected one of {list(literal_choices)}, " + f"got {value!r}" + ) + continue # Type check expected_types = _TYPE_MAP.get(spec.type) if expected_types and not isinstance(value, expected_types): diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 71e081414933..342be8ac0f41 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -251,6 +251,7 @@ l0_b200: - unittest/_torch/visual_gen/test_wan22_i2v_teacache.py - unittest/_torch/visual_gen/test_wan22_t2v_teacache.py - unittest/_torch/visual_gen/test_wan_transformer.py + - unittest/_torch/visual_gen/test_cosmos3_action.py - unittest/_torch/visual_gen/test_cosmos3_transformer.py - unittest/_torch/visual_gen/test_cosmos3_pipeline.py - examples/visual_gen/test_visual_gen.py::test_wan_t2v_example diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py index 9a64ebdf4048..100f42df710b 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py @@ -10,15 +10,22 @@ import numpy as np import PIL.Image import pytest +import torch from tensorrt_llm._torch.visual_gen.models.cosmos3.action import ( VIDEO_RES_SIZE_INFO, action_reference_image, find_closest_target_size, + normalize_action_resolution, normalize_action_video_input, + prepare_action_latents, resolve_action_size, ) -from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import COSMOS3_EXTRA_SPECS +from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( + COSMOS3_EXTRA_SPECS, + get_domain_preset, + resolve_domain_action_config, +) pytestmark = pytest.mark.cosmos3 @@ -78,16 +85,12 @@ def test_partial_width_fills_height_from_bucket(self): 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 + assert spec.type == "Literal[256, 480, 704, 720]" + assert spec.default is None 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 @@ -97,20 +100,12 @@ def test_bridge_preset_fills_missing_fields(self): 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, @@ -120,33 +115,32 @@ def test_mismatch_emits_warning(self): 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 + def test_unknown_domain_warns_and_uses_generic_defaults(self): + cfg = resolve_domain_action_config(domain_name="typo_domain") + assert cfg["action_chunk_size"] == 16 + assert cfg["action_resolution"] == 480 + assert cfg["warnings"] + assert "preset was not found" in cfg["warnings"][0] + + def test_non_positive_action_timing_raises(self): + with pytest.raises(ValueError, match="action_fps must be positive"): + resolve_domain_action_config(domain_name="av", action_fps=0.0) + def test_unknown_resolution_raises(self): with pytest.raises(ValueError, match="Unknown Cosmos3 action_resolution"): normalize_action_resolution(1080) @@ -182,8 +176,25 @@ def test_policy_prefers_image_path_over_video(self, tmp_path): ) assert ref.getpixel((0, 0)) == (0, 0, 255) + def test_policy_accepts_path_image(self, tmp_path): + image_path = tmp_path / "frame.png" + PIL.Image.new("RGB", (3, 3), "green").save(image_path) + ref = action_reference_image( + action_mode="policy", + image=image_path, + video=None, + ) + assert ref.getpixel((0, 0)) == (0, 128, 0) + class TestNormalizeActionVideoInput: + def test_none_returns_empty_list(self): + assert normalize_action_video_input(None) == [] + + def test_empty_list_raises(self): + with pytest.raises(ValueError, match="at least one frame"): + normalize_action_video_input([]) + 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) @@ -252,3 +263,18 @@ def _fake_read_video(path, pts_unit="sec"): max_frames=2, ) assert len(frames) == 2 + + +class TestPrepareActionLatents: + def test_forward_dynamics_raw_dim_mismatch_raises(self): + with pytest.raises(ValueError, match="raw_action_dim must match"): + prepare_action_latents( + mode="forward_dynamics", + action_chunk_size=2, + raw_action_dim=3, + action_dim=8, + generator=torch.Generator(device="cpu").manual_seed(0), + device=torch.device("cpu"), + dtype=torch.float32, + action_input=[[0.0, 1.0], [2.0, 3.0]], + ) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index 37dc235c16b5..422adaae8b6a 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -520,6 +520,8 @@ def test_forward_dynamics_smoke(self, cosmos3_pipeline): raw_action_dim=self.RAW_ACTION_DIM, chunk_size=self.ACTION_CHUNK, ) + assert result.action_mode == "forward_dynamics" + assert result.domain_id == 7 def test_inverse_dynamics_smoke(self, cosmos3_pipeline): _require_action_pipeline(cosmos3_pipeline) @@ -549,6 +551,27 @@ def test_inverse_dynamics_smoke(self, cosmos3_pipeline): raw_action_dim=self.RAW_ACTION_DIM, chunk_size=NUM_FRAMES, ) + assert result.action_mode == "inverse_dynamics" + assert result.domain_id == 7 + + def test_inverse_dynamics_rejects_short_video(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 - 1)] + with pytest.raises(ValueError, match="requires at least"): + _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, + ) def test_action_and_audio_rejected(self, cosmos3_pipeline): _require_action_pipeline(cosmos3_pipeline) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py index 7f02d97ad117..3b6cf23fcc04 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py @@ -442,6 +442,14 @@ def test_video_only_model_has_no_action_heads(self, cosmos3_model_config_noactio model = Cosmos3VFMTransformer(model_config=cosmos3_model_config_noaction) assert model.action_gen is False assert not hasattr(model, "action_proj_in") + assert not hasattr(model, "action_proj_out") + assert not hasattr(model, "action_modality_embed") + + def test_pack_action_rejects_wrong_last_dim(self, action_model_config): + model = Cosmos3VFMTransformer(model_config=action_model_config) + action_latents = torch.randn(1, self.T_ACTION, model.action_dim - 1) + with pytest.raises(ValueError, match="action latent dimension mismatch"): + model.pack_action(action_latents) @pytest.mark.high_cuda_memory def test_forward_with_action(self, action_model_config): @@ -467,6 +475,27 @@ def test_forward_with_action(self, action_model_config): 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_with_action_domain_id_out_of_range_raises(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([self.NUM_DOMAINS], dtype=torch.long, device=DEVICE) + with torch.inference_mode(), pytest.raises(ValueError, match="domain_id"): + 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, + ) + @pytest.mark.high_cuda_memory def test_forward_without_action_latents_returns_none(self, action_model_config): cfg = action_model_config.pretrained_config diff --git a/tests/unittest/_torch/visual_gen/test_tensor_payload.py b/tests/unittest/_torch/visual_gen/test_tensor_payload.py index d7b1913d5344..3c9168ee92cb 100644 --- a/tests/unittest/_torch/visual_gen/test_tensor_payload.py +++ b/tests/unittest/_torch/visual_gen/test_tensor_payload.py @@ -60,6 +60,7 @@ def _make_action_output(batch: int = 1, t: int = 4, action_dim: int = 7) -> Visu action=action, action_mode="policy", raw_action_dim=action_dim, + domain_id=7, ) @@ -153,6 +154,27 @@ def test_unbatched_action_passthrough(self, fmt): assert loaded["action"].shape == (4, 7) assert torch.equal(loaded["action"], output.action) + def test_action_metadata_serialized(self, fmt): + output = _make_action_output(batch=1) + data = serialize_visual_gen_output(output, fmt, batch_index=0) + loaded = self._load(data, fmt) + assert loaded["raw_action_dim"] == 7 + assert loaded["domain_id"] == 7 + if fmt == "pt": + assert loaded["action_mode"] == "policy" + else: + assert "action_mode" not in loaded + import tempfile + + from safetensors import safe_open + + with tempfile.NamedTemporaryFile(suffix=".safetensors") as tf: + tf.write(data) + tf.flush() + with safe_open(tf.name, framework="pt") as f: + meta = f.metadata() or {} + assert meta.get("action_mode") == "policy" + @pytest.mark.parametrize("fmt", ["safetensors", "pt"]) class TestSingleSavePath: 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..6f2ccaf876ef 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py @@ -798,6 +798,21 @@ def test_sync_video_generation_with_params(self, video_client): assert params.frame_rate == 8 assert params.num_frames == int(2.0 * 8) + def test_sync_video_rejects_extra_params_video_path(self, video_client): + resp = video_client.post( + "/v1/videos/generations", + json={ + "prompt": "reject raw path", + "size": "64x64", + "seconds": 1.0, + "fps": 8, + "extra_params": {"video": "/server/local/path.mp4"}, + }, + headers={"content-type": "application/json"}, + ) + assert resp.status_code == 400 + _assert_llm_envelope(resp.json(), code=400, message_contains="server-local paths") + def test_sync_video_generation_multipart(self, video_client, tmp_path): """Multipart sync request with a real ``input_reference`` file.""" ref_path = tmp_path / "ref.png" 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..3262584de652 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py @@ -963,6 +963,43 @@ def test_none_extra_param_value_skipped(self): req = self._make_request(extra_params={"boundary_ratio": None}) self._merge_and_validate(executor, req) + def test_literal_extra_param_rejects_bad_value(self): + from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import COSMOS3_EXTRA_SPECS + + req = self._make_request(extra_params={"action_mode": "not_a_mode"}) + with pytest.raises(ValueError, match="expected one of"): + from tensorrt_llm.visual_gen.params import validate_visual_gen_params + + validate_visual_gen_params( + req.params, + declared_defaults=None, + extra_param_specs=COSMOS3_EXTRA_SPECS, + ) + + def test_literal_extra_param_accepts_numeric_choice(self): + from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import COSMOS3_EXTRA_SPECS + from tensorrt_llm.visual_gen.params import validate_visual_gen_params + + req = self._make_request(extra_params={"action_resolution": 480}) + validate_visual_gen_params( + req.params, + declared_defaults=None, + extra_param_specs=COSMOS3_EXTRA_SPECS, + ) + + def test_path_or_list_extra_param_type_checked(self): + from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import COSMOS3_EXTRA_SPECS + + req = self._make_request(extra_params={"video": 123}) + with pytest.raises(ValueError, match="expected type 'path_or_list'"): + from tensorrt_llm.visual_gen.params import validate_visual_gen_params + + validate_visual_gen_params( + req.params, + declared_defaults=None, + extra_param_specs=COSMOS3_EXTRA_SPECS, + ) + # ============================================================================= # Parameter validation — message content per category diff --git a/tests/unittest/visual_gen/test_output.py b/tests/unittest/visual_gen/test_output.py index c382107b291c..5ff8e6c7818f 100644 --- a/tests/unittest/visual_gen/test_output.py +++ b/tests/unittest/visual_gen/test_output.py @@ -331,6 +331,18 @@ def test_save_no_media_raises(tmp_path): out.save(tmp_path / "x.png") +def test_save_action_with_non_tensor_format_raises(tmp_path): + """Action-bearing outputs require safetensors/pt to preserve action data.""" + out = VisualGenOutput( + request_id=6, + video=torch.zeros(4, 8, 8, 3, dtype=torch.uint8), + action=torch.zeros(4, 7), + frame_rate=16.0, + ) + with pytest.raises(NotImplementedError, match="tensor payload"): + out.save(tmp_path / "x.mp4") + + # --------------------------------------------------------------------------- # VisualGenOutput.save batch routing (list of paths) # --------------------------------------------------------------------------- @@ -761,11 +773,10 @@ def test_encoding_not_top_level_reexport(): # --------------------------------------------------------------------------- -def test_pipeline_output_has_eight_fields(): - """PipelineOutput has the eight expected fields.""" +def test_pipeline_output_has_twelve_fields(): + """PipelineOutput has the twelve expected fields.""" field_names = {f.name for f in fields(PipelineOutput)} assert field_names == { - "request_id", "image", "video", "audio", @@ -775,8 +786,9 @@ def test_pipeline_output_has_eight_fields(): "raw_action_dim", "action_mode", "domain_id", - "error", - "metrics", + "pre_denoise", + "denoise", + "post_denoise", } From bf13fa2223ed95872750c2ee86adaff74e10e7ff Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 3 Aug 2026 10:22:41 -0700 Subject: [PATCH 10/35] [None][fix] Cosmos3 action: mRoPE scaling, action widths, trained caption Adopted-PR review findings, each verified against cosmos-framework and the diffusers Cosmos3 port. mRoPE: action token positions were 4x too large. compute_mrope_position_ids_action accepted base_temporal_compression_factor and dropped it, so the scaled positions used the action tcf (1) for both rates instead of the vision tcf (4) for the base rate. Action tokens therefore advanced one latent frame per step instead of a quarter, desynchronising them from the video they condition. Threaded the factor through compute_mrope_position_ids_vision (defaulting to temporal_compression_factor, so vision and audio are bit-identical). Action widths: raw_action_dim is fixed per embodiment, but it lived in the sampling presets, which several embodiments share via COSMOS3_DOMAIN_PRESET_ALIASES. robomind-franka-dual resolved to 10 instead of 20 and galbot to 29 instead of 30. Moved the canonical widths to action.EMBODIMENT_TO_RAW_ACTION_DIM, keyed by the unaliased domain name; presets now carry sampling settings only. libero stays absent, as in both references, because its width is dataset-dependent. Caption: action checkpoints are trained on a structured JSON caption (cinematography/actions/duration/fps/resolution/aspect_ratio) rather than the flat duration/resolution templates. Added build_action_json_prompt with the four trained viewpoint sentences and a view_point extra param. Aspect ratio snaps to the canonical bucket label; reducing H/W gave "15,26" where the trained label is "16,9". Tests: the five new TestCosmos3Action forward tests and the multi-GPU _forward_with_action helper omitted raw_timestep, which forward() rejects - this is the L0 failure in pipelines 48022 and 48728. Also tightened the domain-id range test, whose match= was satisfied by the raw_timestep error. Signed-off-by: Igor Shovkun --- examples/visual_gen/models/cosmos3/cosmos3.py | 9 + .../visual_gen/models/cosmos3/action.py | 130 ++++++ .../visual_gen/models/cosmos3/defaults.py | 68 ++- .../models/cosmos3/pipeline_cosmos3.py | 78 ++-- .../models/cosmos3/transformer_cosmos3.py | 20 +- .../test_cosmos3_transformer_parallel.py | 3 +- .../_torch/visual_gen/test_cosmos3_action.py | 414 +++++++++++++++++- .../visual_gen/test_cosmos3_transformer.py | 20 +- 8 files changed, 684 insertions(+), 58 deletions(-) diff --git a/examples/visual_gen/models/cosmos3/cosmos3.py b/examples/visual_gen/models/cosmos3/cosmos3.py index 1144c87eb0f6..4f502cec8caa 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3.py +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -282,6 +282,13 @@ def main(): default=None, help="Action-token temporal rate for mRoPE (Hz). Defaults to frame_rate.", ) + parser.add_argument( + "--view_point", + type=str, + default=None, + choices=["ego_view", "third_person_view", "wrist_view", "concat_view"], + help="Camera perspective for the action caption (default: ego_view).", + ) parser.add_argument( "--action_output_path", type=str, @@ -352,6 +359,8 @@ def main(): params.extra_params["action_resolution"] = args.action_resolution if args.action_fps is not None: params.extra_params["action_fps"] = args.action_fps + if args.view_point is not None: + params.extra_params["view_point"] = args.view_point 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/action.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py index fea4a71a5ec6..cadce5784e74 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py @@ -5,6 +5,8 @@ from __future__ import annotations +import json +import math from pathlib import Path from typing import Any @@ -13,6 +15,8 @@ import torch from diffusers.utils.torch_utils import randn_tensor +from tensorrt_llm.logger import logger + ACTION_MODE_POLICY = "policy" ACTION_MODE_FORWARD_DYNAMICS = "forward_dynamics" ACTION_MODE_INVERSE_DYNAMICS = "inverse_dynamics" @@ -42,6 +46,132 @@ "fractal": 20, } +# Canonical unpadded action width per embodiment. Widths compose the Cosmos3 +# unified action representation from shared geometric blocks: a 9-D pose (3-D +# translation + 6-D rotation), a 1-D grasp state, and a 15-D fingertip state. +# One arm is 9 + 1 = 10; a dual-arm setup is 20; the AgiBot humanoid is +# 9 + 2 x (9 + 1) = 29; two-hand egocentric motion is 9 + 2 x (9 + 15) = 57. +# +# This is a property of the embodiment, not a tunable, so it is keyed by the +# real domain name rather than by the sampling presets in ``defaults.py`` (where +# several of these names share one preset). ``libero`` is absent on purpose: +# its width depends on the dataset's rotation space (7/10/13), so callers must +# pass ``raw_action_dim`` explicitly. +EMBODIMENT_TO_RAW_ACTION_DIM: dict[str, int] = { + "av": 9, + "camera_pose": 9, + "hand_pose": 57, + "pusht": 2, + "umi": 10, + "bridge_orig_lerobot": 10, + "droid_lerobot": 10, + "robomind-franka": 10, + "robomind-franka-dual": 20, + "robomind-ur": 10, + "galbot": 30, + "agibotworld": 29, + "agibot_gear_gripper": 29, + "agibot_gear_gripper_ext": 29, + "fractal": 10, +} + + +def resolve_raw_action_dim( + domain_name: Any = None, + domain_id: Any = None, +) -> int | None: + """Look up the canonical action width, or None when it cannot be determined. + + Resolves by name first. A bare ``domain_id`` is only usable when every + embodiment sharing that id agrees on a width (true for all current ids). + """ + if domain_name is not None and str(domain_name).strip(): + return EMBODIMENT_TO_RAW_ACTION_DIM.get(str(domain_name).strip().lower()) + + if domain_id is None: + return None + + widths = { + EMBODIMENT_TO_RAW_ACTION_DIM[name] + for name, mapped_id in EMBODIMENT_TO_DOMAIN_ID.items() + if mapped_id == int(domain_id) and name in EMBODIMENT_TO_RAW_ACTION_DIM + } + return widths.pop() if len(widths) == 1 else None + + +# Camera perspective -> framing sentence. The action model was trained on these +# exact sentences, so they are reproduced verbatim rather than paraphrased. +ACTION_VIEWPOINT_TEMPLATES: dict[str, str] = { + "ego_view": "This video is captured from a first-person perspective looking at the scene.", + "third_person_view": ( + "This video is captured from a third-person perspective looking towards the agent " + "from the front." + ), + "wrist_view": "This video is captured from a wrist-mounted camera.", + "concat_view": "This video contains concatenated views from multiple camera perspectives.", +} + +DEFAULT_ACTION_VIEW_POINT = "ego_view" + +# Canonical ``W,H`` labels; every action canvas is one of these bucket shapes. +ACTION_ASPECT_RATIO_LABELS = ("1,1", "4,3", "3,4", "16,9", "9,16") + + +def action_aspect_ratio_label(height: int, width: int) -> str: + """Closest canonical aspect label, e.g. 832x480 -> ``"16,9"``. + + Bucket sizes are only approximately their label (832/480 is 1.733, not + 1.778), so the label is matched by nearest ratio instead of reducing H/W. + """ + ratio = width / height if height > 0 else 1.0 + return min( + ACTION_ASPECT_RATIO_LABELS, + key=lambda label: abs(int(label.split(",")[0]) / int(label.split(",")[1]) - ratio), + ) + + +def build_action_json_prompt( + description: str, + *, + view_point: str | None, + num_frames: int, + frame_rate: float, + height: int, + width: int, +) -> str: + """Build the structured action caption the action model was trained on. + + Replaces the flat duration/resolution templates used by the video paths: the + JSON already carries duration, fps, resolution and aspect ratio. Key order is + part of the trained format and is preserved. + """ + duration_seconds = num_frames / frame_rate if frame_rate > 0 else 0.0 + if not math.isfinite(duration_seconds) or duration_seconds < 0: + duration_seconds = 0.0 + minutes, seconds = divmod(round(duration_seconds), 60) + + text = description.strip() + if text and not text.endswith((".", "!", "?")): + text = f"{text}." + + framing = ACTION_VIEWPOINT_TEMPLATES.get(view_point) if view_point is not None else None + if view_point is not None and framing is None: + logger.warning( + f"Unrecognized Cosmos3 action view_point={view_point!r}; expected one of " + f"{sorted(ACTION_VIEWPOINT_TEMPLATES)}. Dropping the cinematography.framing field." + ) + + prompt: dict[str, Any] = {} + if framing: + prompt["cinematography"] = {"framing": framing} + prompt["actions"] = [{"time": f"0:00-{minutes}:{seconds:02d}", "description": text}] + prompt["duration"] = f"{int(duration_seconds)}s" + prompt["fps"] = float(frame_rate) + prompt["resolution"] = {"H": int(height), "W": int(width)} + prompt["aspect_ratio"] = action_aspect_ratio_label(height, width) + return json.dumps(prompt) + + VIDEO_RES_SIZE_INFO: dict[str, dict[str, tuple[int, int]]] = { "256": { "1,1": (256, 256), diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index e0e8c06ba384..201ed0321e50 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -18,20 +18,27 @@ 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``, +``COSMOS3_DOMAIN_PRESETS`` lists training-aligned sampling defaults per +embodiment. When ``domain_name`` (or a uniquely mapped ``domain_id``) is set, +the pipeline fills omitted ``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.). + +``raw_action_dim`` is deliberately *not* a preset field: it is fixed by the +embodiment, while several embodiments share one preset via +``COSMOS3_DOMAIN_PRESET_ALIASES``. It resolves from +``action.EMBODIMENT_TO_RAW_ACTION_DIM`` instead. """ from typing import Any, TypedDict from tensorrt_llm._torch.visual_gen.models.cosmos3.action import ( COSMOS3_ACTION_RESOLUTIONS, + DEFAULT_ACTION_VIEW_POINT, EMBODIMENT_TO_DOMAIN_ID, normalize_action_resolution, + resolve_raw_action_dim, ) from tensorrt_llm._torch.visual_gen.pipeline import ExtraParamSchema @@ -85,9 +92,12 @@ class Cosmos3DomainPreset(TypedDict, total=False): - """Recommended action-generation settings for a trained embodiment.""" + """Recommended action sampling settings for a trained embodiment. + + Sampling settings only — the embodiment's action width lives in + ``action.EMBODIMENT_TO_RAW_ACTION_DIM``, keyed by the unaliased domain name. + """ - raw_action_dim: int action_chunk_size: int num_frames: int action_resolution: int @@ -96,9 +106,8 @@ class Cosmos3DomainPreset(TypedDict, total=False): # 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. + # WidowX bridge. "bridge_orig_lerobot": { - "raw_action_dim": 10, "action_chunk_size": 16, "num_frames": 17, "action_resolution": 480, @@ -106,7 +115,6 @@ class Cosmos3DomainPreset(TypedDict, total=False): }, # Autonomous-vehicle steering/throttle; longer action horizon. "av": { - "raw_action_dim": 9, "action_chunk_size": 60, "num_frames": 61, "action_resolution": 480, @@ -114,7 +122,6 @@ class Cosmos3DomainPreset(TypedDict, total=False): }, # 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, @@ -122,7 +129,6 @@ class Cosmos3DomainPreset(TypedDict, total=False): }, # 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, @@ -130,15 +136,13 @@ class Cosmos3DomainPreset(TypedDict, total=False): }, # 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. + # MANO hand pose. "hand_pose": { - "raw_action_dim": 57, "action_chunk_size": 16, "num_frames": 17, "action_resolution": 480, @@ -146,7 +150,6 @@ class Cosmos3DomainPreset(TypedDict, total=False): }, # AgiBot humanoid; shared domain_id with agibot_gear_gripper*. "agibotworld": { - "raw_action_dim": 29, "action_chunk_size": 16, "num_frames": 17, "action_resolution": 480, @@ -154,7 +157,6 @@ class Cosmos3DomainPreset(TypedDict, total=False): }, # Google Robot (RT-1 / fractal) single-arm. "fractal": { - "raw_action_dim": 10, "action_chunk_size": 16, "num_frames": 17, "action_resolution": 480, @@ -162,7 +164,6 @@ class Cosmos3DomainPreset(TypedDict, total=False): }, # 2-D planar push task. "pusht": { - "raw_action_dim": 2, "action_chunk_size": 16, "num_frames": 17, "action_resolution": 256, @@ -170,7 +171,6 @@ class Cosmos3DomainPreset(TypedDict, total=False): }, # UMI handheld gripper setup. "umi": { - "raw_action_dim": 10, "action_chunk_size": 16, "num_frames": 17, "action_resolution": 480, @@ -178,7 +178,9 @@ class Cosmos3DomainPreset(TypedDict, total=False): }, } -# Map alias domain_name keys to a canonical preset entry. +# Map alias domain_name keys to a canonical preset entry. These share *sampling* +# settings only; each alias keeps its own action width (e.g. robomind-franka-dual +# is 20-D and galbot is 30-D, unlike the presets they borrow here). COSMOS3_DOMAIN_PRESET_ALIASES: dict[str, str] = { "robomind-franka": "droid_lerobot", "robomind-franka-dual": "droid_lerobot", @@ -274,7 +276,24 @@ def _resolve_field( return recommended return fallback - resolved_raw_action_dim = _resolve_field("raw_action_dim", raw_action_dim) + # The action width is canonical per embodiment, so it comes from the + # embodiment table rather than the (alias-shared) sampling preset. + canonical_raw_action_dim = resolve_raw_action_dim(domain_name=domain_name, domain_id=domain_id) + if raw_action_dim is not None: + if canonical_raw_action_dim is not None and int(raw_action_dim) != canonical_raw_action_dim: + warnings.append( + f"Cosmos3 raw_action_dim={raw_action_dim} differs from the canonical width " + f"{canonical_raw_action_dim} for domain_name={domain_name!r}." + ) + resolved_raw_action_dim = raw_action_dim + else: + resolved_raw_action_dim = canonical_raw_action_dim + if domain_requested and canonical_raw_action_dim is None: + warnings.append( + "Cosmos3 has no canonical action width for " + f"domain_name={domain_name!r}, domain_id={domain_id!r}; " + "pass raw_action_dim explicitly for policy/inverse_dynamics." + ) resolved_chunk = _resolve_field( "action_chunk_size", action_chunk_size, @@ -376,7 +395,8 @@ def _resolve_field( 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." + "Resolved from the embodiment when omitted; required for domains with no " + "canonical width (libero)." ), ), "action_chunk_size": ExtraParamSchema( @@ -401,6 +421,14 @@ def _resolve_field( ), range=(min(COSMOS3_ACTION_RESOLUTIONS), max(COSMOS3_ACTION_RESOLUTIONS)), ), + "view_point": ExtraParamSchema( + type="Literal['ego_view', 'third_person_view', 'wrist_view', 'concat_view']", + default=DEFAULT_ACTION_VIEW_POINT, + description=( + "Camera perspective for action generation. Fills the trained action caption's " + "cinematography.framing field." + ), + ), "action_fps": ExtraParamSchema( type="float", 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 a22fb6810aa9..201296934142 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,10 @@ from .action import ( ACTION_MODE_INVERSE_DYNAMICS, + DEFAULT_ACTION_VIEW_POINT, action_reference_image, action_start_frame_offset, + build_action_json_prompt, build_vision_condition_mask, normalize_action_mode, normalize_action_video_input, @@ -304,6 +306,7 @@ def infer(self, req): action_resolution=extra_params.get("action_resolution") or extra_params.get("image_size"), action_fps=extra_params.get("action_fps"), + view_point=extra_params.get("view_point", DEFAULT_ACTION_VIEW_POINT), video=extra_params.get("video"), ) @@ -757,6 +760,7 @@ def forward( action: Any = None, action_resolution: Optional[int] = None, action_fps: Optional[float] = None, + view_point: Optional[str] = DEFAULT_ACTION_VIEW_POINT, video: Any = None, ): pipeline_start = time.time() @@ -923,43 +927,61 @@ def forward( if negative_prompt is None: negative_prompt = COSMOS3_DEFAULT_NEGATIVE_PROMPT - # Positive prompt: forward duration/resolution templates. T2I has no - # duration concept (single image) and uses the image-flavored - # resolution template. - use_duration_template = use_duration_template and not is_t2i - dur_tmpl = COSMOS3_DURATION_TEMPLATE if use_duration_template else None - if use_resolution_template: - res_tmpl = ( - COSMOS3_IMAGE_RESOLUTION_TEMPLATE if is_t2i else COSMOS3_DEFAULT_RESOLUTION_TEMPLATE - ) + if do_action: + # Action checkpoints were trained on a structured JSON caption that + # already carries duration/fps/resolution/aspect_ratio, so the flat + # templates are skipped here and the negative prompt stays verbatim. + prompt = [ + build_action_json_prompt( + p, + view_point=view_point, + num_frames=num_frames, + frame_rate=frame_rate, + height=height, + width=width, + ) + for p in prompt + ] else: - res_tmpl = None - - # Negative prompt: mirror positive metadata (cosmos-framework CLI default - # when ``negative_prompt_keep_metadata`` promotes mode to ``same``). - negative_prompt = self._format_prompt_with_metadata( - negative_prompt, - height=height, - width=width, - num_frames=num_frames, - frame_rate=frame_rate, - duration_template=dur_tmpl, - resolution_template=res_tmpl, - force_duration_template=False, - ) + # Positive prompt: forward duration/resolution templates. T2I has no + # duration concept (single image) and uses the image-flavored + # resolution template. + use_duration_template = use_duration_template and not is_t2i + dur_tmpl = COSMOS3_DURATION_TEMPLATE if use_duration_template else None + if use_resolution_template: + res_tmpl = ( + COSMOS3_IMAGE_RESOLUTION_TEMPLATE + if is_t2i + else COSMOS3_DEFAULT_RESOLUTION_TEMPLATE + ) + else: + res_tmpl = None - prompt = [ - self._format_prompt_with_metadata( - p, + # Negative prompt: mirror positive metadata (cosmos-framework CLI default + # when ``negative_prompt_keep_metadata`` promotes mode to ``same``). + negative_prompt = self._format_prompt_with_metadata( + negative_prompt, height=height, width=width, num_frames=num_frames, frame_rate=frame_rate, duration_template=dur_tmpl, resolution_template=res_tmpl, + force_duration_template=False, ) - for p in prompt - ] + + prompt = [ + self._format_prompt_with_metadata( + p, + height=height, + width=width, + num_frames=num_frames, + frame_rate=frame_rate, + duration_template=dur_tmpl, + resolution_template=res_tmpl, + ) + for p in prompt + ] logger.info(f"Prompt with metadata: '{prompt}'") prompt = prompt[0] 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..6261decf0311 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -103,6 +103,7 @@ def compute_mrope_position_ids_vision( temporal_compression_factor: int = 4, enable_fps_modulation: bool = False, start_frame_offset: int = 0, + base_temporal_compression_factor: int | None = None, ) -> tuple[torch.Tensor, int | float]: """Generate 3D mRoPE position IDs for vision tokens. @@ -114,12 +115,20 @@ def compute_mrope_position_ids_vision( to reflect real time so that videos at different frame rates get comparable temporal embeddings. + ``base_temporal_compression_factor`` sets the temporal grid the scaled + positions land on, and defaults to ``temporal_compression_factor``. Action + tokens run at frame rate (``temporal_compression_factor=1``) but must share + the vision latent-frame grid, so they pass the vision VAE factor here. + Returns: (position_ids [3, grid_t * grid_h * grid_w], next_temporal_offset) """ + if base_temporal_compression_factor is None: + base_temporal_compression_factor = temporal_compression_factor + if enable_fps_modulation and fps is not None: tps = fps / temporal_compression_factor - base_tps = base_fps / temporal_compression_factor + base_tps = base_fps / base_temporal_compression_factor frame_indices = torch.arange(grid_t, dtype=torch.float32) t_index = ( ((frame_indices + start_frame_offset) / tps * base_tps + temporal_offset) @@ -161,7 +170,13 @@ def compute_mrope_position_ids_action( 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.""" + """Generate mRoPE IDs for action tokens as a frame-rate (T, 1, 1) grid. + + Action tokens are uncompressed in time, so they advance one source frame per + token while vision latent frames advance ``base_temporal_compression_factor`` + source frames. Scaling against the vision base rate keeps both streams on + one shared timeline. + """ return compute_mrope_position_ids_vision( grid_t=grid_t, grid_h=1, @@ -172,6 +187,7 @@ def compute_mrope_position_ids_action( temporal_compression_factor=1, enable_fps_modulation=enable_fps_modulation, start_frame_offset=start_frame_offset, + base_temporal_compression_factor=base_temporal_compression_factor, ) 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..5f4c15195857 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 @@ -451,7 +451,8 @@ def _forward_with_action( 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, diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py index 100f42df710b..7e3fcbe3f4b3 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py @@ -7,25 +7,40 @@ pytest tests/unittest/_torch/visual_gen/test_cosmos3_action.py -v """ +import json + import numpy as np import PIL.Image import pytest import torch from tensorrt_llm._torch.visual_gen.models.cosmos3.action import ( + ACTION_VIEWPOINT_TEMPLATES, + DEFAULT_ACTION_VIEW_POINT, + EMBODIMENT_TO_DOMAIN_ID, + EMBODIMENT_TO_RAW_ACTION_DIM, VIDEO_RES_SIZE_INFO, + action_aspect_ratio_label, action_reference_image, + build_action_json_prompt, find_closest_target_size, normalize_action_resolution, normalize_action_video_input, prepare_action_latents, resolve_action_size, + resolve_raw_action_dim, ) from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( + COSMOS3_DOMAIN_PRESET_ALIASES, + COSMOS3_DOMAIN_PRESETS, COSMOS3_EXTRA_SPECS, get_domain_preset, resolve_domain_action_config, ) +from tensorrt_llm._torch.visual_gen.models.cosmos3.transformer_cosmos3 import ( + compute_mrope_position_ids_action, + compute_mrope_position_ids_vision, +) pytestmark = pytest.mark.cosmos3 @@ -127,7 +142,69 @@ def test_explicit_action_fps_overrides_default(self): def test_alias_maps_to_canonical_preset(self): preset = get_domain_preset("robomind-franka") assert preset is not None - assert preset["raw_action_dim"] == 10 + assert preset == get_domain_preset("droid_lerobot") + + def test_presets_carry_sampling_settings_only(self): + """Width is per-embodiment; aliases share presets, so it must not live there.""" + for preset in COSMOS3_DOMAIN_PRESETS.values(): + assert "raw_action_dim" not in preset + + @pytest.mark.parametrize( + "domain_name,expected", + [ + ("bridge_orig_lerobot", 10), + ("droid_lerobot", 10), + ("robomind-franka", 10), + ("robomind-ur", 10), + ("robomind-franka-dual", 20), # dual arm, not the droid preset's 10 + ("galbot", 30), # humanoid stack, not agibotworld's 29 + ("agibotworld", 29), + ("agibot_gear_gripper", 29), + ("agibot_gear_gripper_ext", 29), + ("av", 9), + ("camera_pose", 9), + ("hand_pose", 57), + ("pusht", 2), + ("umi", 10), + ("fractal", 10), + ], + ) + def test_canonical_action_width_per_embodiment(self, domain_name, expected): + assert resolve_raw_action_dim(domain_name=domain_name) == expected + assert resolve_domain_action_config(domain_name=domain_name)["raw_action_dim"] == expected + + def test_aliased_domains_keep_their_own_width(self): + """Sharing a sampling preset must not import that preset's action width.""" + for alias, canonical in COSMOS3_DOMAIN_PRESET_ALIASES.items(): + alias_dim = resolve_raw_action_dim(domain_name=alias) + if alias_dim is None: + continue + assert alias_dim == EMBODIMENT_TO_RAW_ACTION_DIM[alias], ( + f"{alias} must keep its own width, not {canonical}'s" + ) + + def test_libero_has_no_canonical_width(self): + """LIBERO's width depends on the dataset's rotation space (7/10/13).""" + assert "libero" not in EMBODIMENT_TO_RAW_ACTION_DIM + cfg = resolve_domain_action_config(domain_name="libero") + assert cfg["raw_action_dim"] is None + assert cfg["action_resolution"] == 256 # sampling preset still applies + assert any("canonical action width" in w for w in cfg["warnings"]) + + def test_explicit_raw_action_dim_overrides_with_warning(self): + cfg = resolve_domain_action_config(domain_name="libero", raw_action_dim=7) + assert cfg["raw_action_dim"] == 7 + assert cfg["warnings"] == [] + + def test_domain_id_resolves_width_when_unambiguous(self): + assert resolve_raw_action_dim(domain_id=12) == 20 # robomind-franka-dual + assert resolve_raw_action_dim(domain_id=8) == 10 # droid / robomind-franka agree + assert resolve_raw_action_dim(domain_id=15) == 29 # all agibot variants agree + assert resolve_raw_action_dim(domain_id=5) is None # libero + assert resolve_raw_action_dim(domain_id=0) is None # no_action + + def test_every_width_entry_has_a_domain_id(self): + assert set(EMBODIMENT_TO_RAW_ACTION_DIM) <= set(EMBODIMENT_TO_DOMAIN_ID) def test_unknown_domain_warns_and_uses_generic_defaults(self): cfg = resolve_domain_action_config(domain_name="typo_domain") @@ -265,6 +342,341 @@ def _fake_read_video(path, pts_unit="sec"): assert len(frames) == 2 +class TestActionJsonPrompt: + """The trained action caption: structured JSON, not the flat video templates.""" + + BRIDGE = dict(num_frames=17, frame_rate=5.0, height=480, width=832) + + def test_matches_trained_shape(self): + payload = json.loads( + build_action_json_prompt( + "Pick up the pear and place it in the bag", + view_point="ego_view", + **self.BRIDGE, + ) + ) + assert payload == { + "cinematography": { + "framing": ( + "This video is captured from a first-person perspective looking at the scene." + ) + }, + "actions": [ + { + "time": "0:00-0:03", + "description": "Pick up the pear and place it in the bag.", + } + ], + "duration": "3s", + "fps": 5.0, + "resolution": {"H": 480, "W": 832}, + "aspect_ratio": "16,9", + } + + def test_key_order_is_preserved(self): + """Field order is part of the trained caption format.""" + text = build_action_json_prompt("Do a thing", view_point="ego_view", **self.BRIDGE) + assert list(json.loads(text).keys()) == [ + "cinematography", + "actions", + "duration", + "fps", + "resolution", + "aspect_ratio", + ] + + @pytest.mark.parametrize("view_point", sorted(ACTION_VIEWPOINT_TEMPLATES)) + def test_every_viewpoint_emits_its_trained_sentence(self, view_point): + payload = json.loads( + build_action_json_prompt("Do a thing", view_point=view_point, **self.BRIDGE) + ) + assert payload["cinematography"]["framing"] == ACTION_VIEWPOINT_TEMPLATES[view_point] + + def test_default_view_point_is_known(self): + assert DEFAULT_ACTION_VIEW_POINT in ACTION_VIEWPOINT_TEMPLATES + + @pytest.mark.parametrize("view_point", [None, "sideways_view"]) + def test_unknown_or_missing_view_point_drops_framing(self, view_point): + payload = json.loads( + build_action_json_prompt("Do a thing", view_point=view_point, **self.BRIDGE) + ) + assert "cinematography" not in payload + assert list(payload.keys())[0] == "actions" + + @pytest.mark.parametrize( + "description,expected", + [ + ("Pick up the pear", "Pick up the pear."), + ("Pick up the pear.", "Pick up the pear."), + ("Is it a pear?", "Is it a pear?"), + ("Grab it!", "Grab it!"), + (" padded ", "padded."), + ("", ""), + ], + ) + def test_description_is_terminated_once(self, description, expected): + payload = json.loads(build_action_json_prompt(description, view_point=None, **self.BRIDGE)) + assert payload["actions"][0]["description"] == expected + + @pytest.mark.parametrize( + "num_frames,frame_rate,duration,time_range", + [ + (17, 5.0, "3s", "0:00-0:03"), # bridge: 3.4s truncates, rounds to 3 + (17, 24.0, "0s", "0:00-0:01"), # 0.708s truncates to 0, rounds to 1 + (61, 10.0, "6s", "0:00-0:06"), # av preset + (241, 2.0, "120s", "0:00-2:00"), # crosses the minute boundary + ], + ) + def test_duration_truncates_while_time_range_rounds( + self, num_frames, frame_rate, duration, time_range + ): + payload = json.loads( + build_action_json_prompt( + "Do a thing", + view_point=None, + num_frames=num_frames, + frame_rate=frame_rate, + height=480, + width=832, + ) + ) + assert payload["duration"] == duration + assert payload["actions"][0]["time"] == time_range + assert payload["fps"] == float(frame_rate) + + @pytest.mark.parametrize("action_resolution", sorted(VIDEO_RES_SIZE_INFO)) + def test_aspect_label_matches_the_bucket_it_came_from(self, action_resolution): + """Every canvas is a bucket entry, so its label must round-trip.""" + for label, (width, height) in VIDEO_RES_SIZE_INFO[action_resolution].items(): + assert action_aspect_ratio_label(height, width) == label + + def test_aspect_label_is_not_a_reduced_fraction(self): + """832x480 reduces to 26,15 but the trained label is 16,9.""" + assert action_aspect_ratio_label(480, 832) == "16,9" + + def test_resolution_is_reported_as_the_padded_canvas(self): + payload = json.loads(build_action_json_prompt("Do a thing", view_point=None, **self.BRIDGE)) + assert payload["resolution"] == {"H": 480, "W": 832} + + def test_zero_frame_rate_does_not_raise(self): + payload = json.loads( + build_action_json_prompt( + "Do a thing", + view_point=None, + num_frames=17, + frame_rate=0.0, + height=480, + width=832, + ) + ) + assert payload["duration"] == "0s" + assert payload["actions"][0]["time"] == "0:00-0:00" + + def test_view_point_spec_defaults_to_ego_view(self): + spec = COSMOS3_EXTRA_SPECS["view_point"] + assert spec.default == DEFAULT_ACTION_VIEW_POINT + for view_point in ACTION_VIEWPOINT_TEMPLATES: + assert repr(view_point) in spec.type + + +def _reference_scaled_positions( + *, + grid_t: int, + temporal_offset: float, + fps: float, + base_fps: float, + temporal_compression_factor: int, + base_temporal_compression_factor: int, + start_frame_offset: int, +) -> list[float]: + """Transcription of cosmos-framework ``get_3d_mrope_ids_vae_tokens``. + + Reference: ``cosmos_framework/data/generator/sequence_packing/mrope.py`` + (mirrored by diffusers ``pipeline_cosmos3_omni.get_3d_mrope_ids_vae_tokens``). + """ + tps = fps / temporal_compression_factor + base_tps = base_fps / base_temporal_compression_factor + return [(i + start_frame_offset) / tps * base_tps + temporal_offset for i in range(grid_t)] + + +class TestActionMropePositionIds: + """Action tokens run at frame rate but must share the vision latent timeline.""" + + VISION_TCF = 4 + + @pytest.mark.parametrize( + "grid_t,temporal_offset,action_fps,base_fps,start_frame_offset", + [ + (4, 0.0, 24.0, 24.0, 1), + (4, 15032.0, 24.0, 24.0, 1), + (16, 0.0, 5.0, 24.0, 1), + (60, 0.0, 10.0, 24.0, 1), + (4, 0.0, 24.0, 24.0, 0), + ], + ) + def test_matches_reference_formula( + self, grid_t, temporal_offset, action_fps, base_fps, start_frame_offset + ): + ids, _ = compute_mrope_position_ids_action( + grid_t, + temporal_offset=temporal_offset, + action_fps=action_fps, + base_fps=base_fps, + base_temporal_compression_factor=self.VISION_TCF, + enable_fps_modulation=True, + start_frame_offset=start_frame_offset, + ) + expected = _reference_scaled_positions( + grid_t=grid_t, + temporal_offset=temporal_offset, + fps=action_fps, + base_fps=base_fps, + temporal_compression_factor=1, + base_temporal_compression_factor=self.VISION_TCF, + start_frame_offset=start_frame_offset, + ) + torch.testing.assert_close( + ids[0], torch.tensor(expected, dtype=ids.dtype), rtol=0, atol=1e-5 + ) + + def test_action_step_advances_one_source_frame(self): + """Consecutive action tokens are 1/vision_tcf of a latent frame apart.""" + ids, _ = compute_mrope_position_ids_action( + 8, + temporal_offset=0.0, + action_fps=24.0, + base_fps=24.0, + base_temporal_compression_factor=self.VISION_TCF, + enable_fps_modulation=True, + start_frame_offset=1, + ) + deltas = ids[0, 1:] - ids[0, :-1] + torch.testing.assert_close( + deltas, torch.full_like(deltas, 1.0 / self.VISION_TCF), rtol=0, atol=1e-5 + ) + + @pytest.mark.parametrize( + "action_chunk_size,num_frames,fps", + [ + (16, 17, 24.0), # generic COSMOS3_ACTION_PARAMS default + (16, 17, 5.0), # bridge_orig_lerobot preset + (60, 61, 10.0), # av preset + ], + ) + def test_last_action_token_lands_on_last_vision_latent_frame( + self, action_chunk_size, num_frames, fps + ): + """The 4x-scaling regression: action must not outrun the video it conditions. + + Vision and action are packed into one temporal axis, so the paired + (num_frames, action_chunk_size) config must place the final action token + exactly on the final vision latent frame. + """ + latent_t = (num_frames - 1) // self.VISION_TCF + 1 + vision_ids, _ = compute_mrope_position_ids_vision( + latent_t, + 1, + 1, + temporal_offset=0.0, + fps=fps, + base_fps=24.0, + temporal_compression_factor=self.VISION_TCF, + enable_fps_modulation=True, + ) + action_ids, _ = compute_mrope_position_ids_action( + action_chunk_size, + temporal_offset=0.0, + action_fps=fps, + base_fps=24.0, + base_temporal_compression_factor=self.VISION_TCF, + enable_fps_modulation=True, + start_frame_offset=1, + ) + assert action_ids[0, -1].item() == pytest.approx(vision_ids[0, -1].item(), abs=1e-5) + assert action_ids[0, 0].item() > vision_ids[0, 0].item() + + def test_spatial_rows_are_zero(self): + ids, _ = compute_mrope_position_ids_action( + 5, + temporal_offset=0.0, + action_fps=24.0, + base_fps=24.0, + base_temporal_compression_factor=self.VISION_TCF, + enable_fps_modulation=True, + ) + assert ids.shape == (3, 5) + assert torch.all(ids[1] == 0) + assert torch.all(ids[2] == 0) + + def test_fps_modulation_disabled_gives_integer_frame_indices(self): + ids, _ = compute_mrope_position_ids_action( + 4, + temporal_offset=7.0, + action_fps=24.0, + base_fps=24.0, + base_temporal_compression_factor=self.VISION_TCF, + enable_fps_modulation=False, + start_frame_offset=1, + ) + assert ids[0].tolist() == [8, 9, 10, 11] + + def test_lower_action_fps_stretches_positions(self): + kwargs = dict( + temporal_offset=0.0, + base_fps=24.0, + base_temporal_compression_factor=self.VISION_TCF, + enable_fps_modulation=True, + start_frame_offset=1, + ) + fast, _ = compute_mrope_position_ids_action(4, action_fps=24.0, **kwargs) + slow, _ = compute_mrope_position_ids_action(4, action_fps=12.0, **kwargs) + torch.testing.assert_close(slow[0], fast[0] * 2.0, rtol=0, atol=1e-5) + + +class TestVisionMropeBaseCompressionDefault: + """``base_temporal_compression_factor=None`` must not disturb vision or audio.""" + + def test_vision_positions_unchanged_by_default(self): + kwargs = dict( + temporal_offset=3.0, + fps=30.0, + base_fps=24.0, + temporal_compression_factor=4, + enable_fps_modulation=True, + ) + implicit, next_implicit = compute_mrope_position_ids_vision(5, 2, 2, **kwargs) + explicit, next_explicit = compute_mrope_position_ids_vision( + 5, 2, 2, base_temporal_compression_factor=4, **kwargs + ) + torch.testing.assert_close(implicit, explicit, rtol=0, atol=0) + assert next_implicit == next_explicit + + def test_audio_style_call_unchanged(self): + """Audio packs with tcf=1 and no base override (sound base tcf is also 1).""" + ids, _ = compute_mrope_position_ids_vision( + 3, + 1, + 1, + temporal_offset=0.0, + fps=25.0, + base_fps=24.0, + temporal_compression_factor=1, + enable_fps_modulation=True, + ) + expected = _reference_scaled_positions( + grid_t=3, + temporal_offset=0.0, + fps=25.0, + base_fps=24.0, + temporal_compression_factor=1, + base_temporal_compression_factor=1, + start_frame_offset=0, + ) + torch.testing.assert_close( + ids[0], torch.tensor(expected, dtype=ids.dtype), rtol=0, atol=1e-5 + ) + + class TestPrepareActionLatents: def test_forward_dynamics_raw_dim_mismatch_raises(self): with pytest.raises(ValueError, match="raw_action_dim must match"): diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py index 3b6cf23fcc04..2abddf6be45e 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py @@ -463,7 +463,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, @@ -484,10 +485,14 @@ def test_forward_with_action_domain_id_out_of_range_raises(self, action_model_co ) action_latents = torch.randn(1, self.T_ACTION, model.action_dim, device=DEVICE, dtype=DTYPE) domain_ids = torch.tensor([self.NUM_DOMAINS], dtype=torch.long, device=DEVICE) - with torch.inference_mode(), pytest.raises(ValueError, match="domain_id"): + with ( + torch.inference_mode(), + pytest.raises(ValueError, match=r"domain_id must be in \[0, \d+\)"), + ): 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, @@ -506,7 +511,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, @@ -528,7 +534,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, @@ -552,7 +559,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 d9fd796006fcf4684ad6f1f81a775812fcf003e3 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Tue, 4 Aug 2026 11:10:57 -0700 Subject: [PATCH 11/35] [None][feat] Decode Cosmos3 action references on NVDEC with fit+pad Action's inverse_dynamics decoded its reference with torchvision.io.read_video, an API removed in torchvision 0.24+, so the mode did not run at all on a current install. It now takes the same encoded-bytes contract as V2V and decodes worker-side on NVDEC. The two modes need opposite framing, so decode_video_reference_window gains a resize selector. V2V continues to cover-scale and center-crop (default, unchanged): losing a border strip costs a scene continuation nothing, and no invented pixels enter the frame. Action gets "fit": contain-scale and pad, because a gripper and its target sit at the frame edge, and cropping them away removes exactly what the policy is supposed to act on. resize_fit_pad_uint8 mirrors the action reference's reflection_pad_to_target - contain-scale by min(target/source, 1.0) so a small clip keeps its own pixels rather than being enlarged, round-not-ceil resize, pad bottom/right by reflection, switching to edge replication once a pad run reaches the resized extent and reflection has no pixels left to mirror. It reuses this module's Lanczos-3 taps rather than the reference's bicubic: the geometry is what preserves content, and a second filter would buy sub-pixel differences at the cost of a second code path. The action decode is wrapped in the same try/except plus synchronize_media_prepare_status convergence as V2V and I2V, so a per-rank NVDEC failure surfaces on every rank instead of hanging the collectives. Signed-off-by: Igor Shovkun --- .../models/cosmos3/pipeline_cosmos3.py | 69 +++++++++++++------ tensorrt_llm/media/decoding.py | 59 +++++++++++++++- 2 files changed, 106 insertions(+), 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 b35b9b5e3ced..162849c1e150 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -47,7 +47,6 @@ build_action_json_prompt, build_vision_condition_mask, normalize_action_mode, - normalize_action_video_input, pil_to_rgb, prepare_action_latents, resize_and_pad_action_image, @@ -228,7 +227,7 @@ def load_standard_components( if self.action_gen: # Action uses its own scheduler for the same reason as audio. - self.action_scheduler = UniPCMultistepScheduler.from_config(self.scheduler.config) + self.action_scheduler = type(self.scheduler).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" @@ -342,8 +341,13 @@ def infer(self, req): 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") + # Action sizes its canvas from the resolution bucket and the reference + # frame's aspect (resolve_action_size), so leaving height/width unset + # here is what lets forward() do that; filling them from the video + # table would pin every action request to 720p. + is_action = extra_params.get("action_mode") is not None + height = req.params.height if is_action else resolved(req.params.height, "height") + width = req.params.width if is_action else 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) @@ -352,11 +356,11 @@ def resolved(value, field_name): 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, @@ -1299,21 +1303,44 @@ 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) - if len(video) < num_frames: + if not isinstance(video, bytes): raise ValueError( - "Cosmos3 inverse_dynamics requires at least " - f"{num_frames} frames, got {len(video)}." + "Cosmos3 inverse_dynamics requires encoded MP4/AVI bytes " + f"(the 'video' extra-param contract), got {type(video).__name__}." ) - 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 + prepare_error: Optional[Exception] = None + try: + # "fit" rather than the V2V default: an action reference is + # padded to the canvas, never cropped to it, because the + # gripper and target sit at the frame edge. + frames_u8 = decode_video_reference_window( + video, + first_frame=0, + last_frame=num_frames - 1, + target_h=height, + target_w=width, + device=self.device, + resize="fit", + ) + if frames_u8.shape[0] < num_frames: + raise ValueError( + "Cosmos3 inverse_dynamics requires at least " + f"{num_frames} frames, got {frames_u8.shape[0]}." + ) + video_tensor = self._condition_frames_to_video_tensor(frames_u8) + del frames_u8 + latents, velocity_mask, condition_latents = self._prepare_latents_action_video( + video_tensor, + normalized_action_mode, + num_frames, + generator, + ) + del video_tensor + except Exception as exc: + prepare_error = exc + # Every rank decodes independently; converge before the + # transformer's collectives so a failure cannot hang the job. + synchronize_media_prepare_status(prepare_error) else: image_tensor = self._preprocess_action_image(action_ref_image, height, width) if image_tensor.ndim == 4: diff --git a/tensorrt_llm/media/decoding.py b/tensorrt_llm/media/decoding.py index 9fa72833bc0d..df8287def3ae 100644 --- a/tensorrt_llm/media/decoding.py +++ b/tensorrt_llm/media/decoding.py @@ -107,6 +107,51 @@ def resize_center_crop_uint8(frames: torch.Tensor, target_h: int, target_w: int) return x.round_().clamp_(0, 255).to(torch.uint8).permute(0, 2, 3, 1).contiguous() +def resize_fit_pad_uint8(frames: torch.Tensor, target_h: int, target_w: int) -> torch.Tensor: + """Resize to fit inside the target, then pad bottom/right to fill it. + + The counterpart to :func:`resize_center_crop_uint8`, for references whose + periphery carries signal — a robot gripper works at the frame edge, so + cropping it away costs the model the thing it is meant to act on. + + Semantics mirror the action reference's ``reflection_pad_to_target``: + contain-scale by ``min(target/source, 1.0)`` (never enlarge — a small clip + keeps its own pixels and gets a wider border), round-rather-than-ceil + resize, then pad bottom/right by reflection, switching to edge replication + when a pad run reaches the resized extent (reflection has no source pixels + left to mirror). The resampling filter stays this module's Lanczos-3 rather + than the reference's bicubic: the geometry is what preserves content, and a + second filter would buy sub-pixel differences for a second code path. + """ + t, h, w, c = frames.shape + if (h, w) == (target_h, target_w): + return frames + ratio = min(target_w / w, target_h / h, 1.0) + resize_w = min(int(ratio * w + 0.5), target_w) + resize_h = min(int(ratio * h + 0.5), target_h) + + # Two passes with a uint8-quantized intermediate, as in the cover path. + 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) + x = x.round_().clamp_(0, 255) + + pad_w = target_w - resize_w + pad_h = target_h - resize_h + if pad_w or pad_h: + mode = "replicate" if (pad_w >= resize_w or pad_h >= resize_h) else "reflect" + x = torch.nn.functional.pad(x, (0, pad_w, 0, pad_h), mode=mode) + return x.to(torch.uint8).permute(0, 2, 3, 1).contiguous() + + +_RESIZE_MODES = {"cover": resize_center_crop_uint8, "fit": resize_fit_pad_uint8} + + def decode_video_reference_window( data: bytes, *, @@ -115,6 +160,7 @@ def decode_video_reference_window( target_h: int, target_w: int, device: torch.device, + resize: str = "cover", ) -> torch.Tensor: """Decode frames ``[first_frame, last_frame]`` of a reference on device. @@ -122,6 +168,12 @@ def decode_video_reference_window( 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. + ``resize`` selects how each frame reaches ``target_h x target_w``: + ``"cover"`` scales to fill and center-crops (the default, and what video + continuation wants); ``"fit"`` scales to fit and pads, for references whose + frame edges carry signal. See :func:`resize_center_crop_uint8` and + :func:`resize_fit_pad_uint8`. + 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 @@ -141,6 +193,11 @@ def decode_video_reference_window( raise ValueError( f"first_frame must not exceed last_frame, got ({first_frame}, {last_frame})." ) + resize_frames = _RESIZE_MODES.get(resize) + if resize_frames is None: + raise ValueError( + f"Unknown resize mode {resize!r}; expected one of {sorted(_RESIZE_MODES)}." + ) window = last_frame - first_frame + 1 from_end = first_frame < 0 try: @@ -211,7 +268,7 @@ def _read(buf: bytearray) -> int: # 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] + resize_frames(decoded.unsqueeze(0), target_h, target_w)[0] ) kept += 1 count += 1 From ed951130da4c9099eee9a5ffc108665c6b00ca86 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Tue, 4 Aug 2026 11:42:37 -0700 Subject: [PATCH 12/35] [None][fix] Retire the Cosmos3 action path's own video reader Follows the NVDEC swap: action.py carried a private mp4/frame-directory reader (decode_action_video_file, normalize_action_video_path, normalize_action_video_input) and the pipeline a _preprocess_action_video built on it. All of that is superseded by decode_video_reference_window. action_reference_image returned a decoded PIL frame that callers used only for its dimensions - the canvas is the resolution bucket closest to the source's aspect - so it becomes action_reference_size, and video bytes report their size from the container header via probe_video_dimensions rather than by decoding. resolve_action_size takes those dimensions instead of an image. policy and forward_dynamics accept either source: an image goes through PIL, video bytes take frame 0 off NVDEC, and both land on the padded canvas, so the two entry points produce the same conditioning for the same picture. num_frames is now derived as action_chunk_size + 1 wherever it is needed and is gone from the presets and from resolve_domain_action_config's signature. It had been resolvable from a preset, so overriding action_chunk_size left a frame count that disagreed with it - the inverse_dynamics smoke test hit exactly that, asking for a 9-frame chunk and being handed the preset's 17. Tests: the mp4-decoding cases are replaced by TestActionReferenceSize, which covers image measurement, image-over-video precedence, and the header probe for bytes. The inverse_dynamics pipeline tests now feed the checked-in 9-frame V2V fixture as bytes; 178 pass across the three cosmos3 files. Signed-off-by: Igor Shovkun --- .../visual_gen/models/cosmos3/action.py | 107 +++--------- .../visual_gen/models/cosmos3/defaults.py | 29 ++- .../models/cosmos3/pipeline_cosmos3.py | 54 ++++-- tensorrt_llm/media/decoding.py | 34 ++++ .../_torch/visual_gen/test_cosmos3_action.py | 165 +++++------------- .../visual_gen/test_cosmos3_pipeline.py | 18 +- 6 files changed, 155 insertions(+), 252 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py index cadce5784e74..1d32e90cac4a 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py @@ -388,110 +388,47 @@ def pil_to_rgb(value: Any) -> PIL.Image.Image: ) -def decode_action_video_file(path: Path, max_frames: int | None = 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: int | None = 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: int | None = 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: int | None, width: int | None, - ref_image: PIL.Image.Image, + source_h: int, + source_w: int, action_resolution: int, ) -> tuple[int, int]: - """Fill unset action H/W from the action resolution bucket; honor explicit values.""" + """Fill unset action H/W from the action resolution bucket; honor explicit values. + + The bucket is the canvas whose shape is closest to the source's, so the + reference only ever needs a modest pad to reach it. + """ 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 - ) + target_w, target_h = find_closest_target_size(source_h, source_w, 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( +def action_reference_size( *, 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 +) -> tuple[int, int]: + """Source ``(height, width)`` of the reference, for choosing the canvas. + + Video references are encoded bytes, so their size comes from the container + header rather than a decode; images are measured directly. + """ + source = video if action_mode == ACTION_MODE_INVERSE_DYNAMICS else (image or 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 = Path(source) - if path.is_file() and path.suffix.lower() in ACTION_IMAGE_EXTENSIONS: - return PIL.Image.open(path).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}." - ) + if isinstance(source, bytes): + from tensorrt_llm.media.decoding import probe_video_dimensions + + return probe_video_dimensions(source) + reference = pil_to_rgb(source) + return reference.height, reference.width def resize_and_pad_action_image( diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index 4d2f84418703..51b7ca8ee047 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -20,10 +20,11 @@ ----------------- ``COSMOS3_DOMAIN_PRESETS`` lists training-aligned sampling defaults per embodiment. When ``domain_name`` (or a uniquely mapped ``domain_id``) is set, -the pipeline fills omitted ``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.). +the pipeline fills omitted ``action_chunk_size``, ``action_resolution``, and +``frame_rate`` from the preset and logs a warning if explicit values differ. +``num_frames`` is always ``action_chunk_size + 1``, never a preset field. See +Cosmos3 omni ``action_*.json`` inputs for reference configs (bridge, av, droid, +libero, etc.). ``raw_action_dim`` is deliberately *not* a preset field: it is fixed by the embodiment, while several embodiments share one preset via @@ -159,7 +160,6 @@ class Cosmos3DomainPreset(TypedDict, total=False): """ action_chunk_size: int - num_frames: int action_resolution: int frame_rate: float @@ -169,70 +169,60 @@ class Cosmos3DomainPreset(TypedDict, total=False): # WidowX bridge. "bridge_orig_lerobot": { "action_chunk_size": 16, - "num_frames": 17, "action_resolution": 480, "frame_rate": 5.0, }, # Autonomous-vehicle steering/throttle; longer action horizon. "av": { "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": { "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": { "action_chunk_size": 16, - "num_frames": 17, "action_resolution": 480, "frame_rate": 15.0, }, # LIBERO sim single-arm; lower action resolution bucket. "libero": { "action_chunk_size": 16, - "num_frames": 17, "action_resolution": 256, "frame_rate": 10.0, }, # MANO hand pose. "hand_pose": { "action_chunk_size": 16, - "num_frames": 17, "action_resolution": 480, "frame_rate": 24.0, }, # AgiBot humanoid; shared domain_id with agibot_gear_gripper*. "agibotworld": { "action_chunk_size": 16, - "num_frames": 17, "action_resolution": 480, "frame_rate": 10.0, }, # Google Robot (RT-1 / fractal) single-arm. "fractal": { "action_chunk_size": 16, - "num_frames": 17, "action_resolution": 480, "frame_rate": 5.0, }, # 2-D planar push task. "pusht": { "action_chunk_size": 16, - "num_frames": 17, "action_resolution": 256, "frame_rate": 10.0, }, # UMI handheld gripper setup. "umi": { "action_chunk_size": 16, - "num_frames": 17, "action_resolution": 480, "frame_rate": 10.0, }, @@ -301,7 +291,6 @@ def resolve_domain_action_config( action_resolution: int | None = None, frame_rate: float | None = None, action_fps: float | None = None, - num_frames: int | None = None, ) -> dict[str, Any]: """Merge user action params with domain presets and generic fallbacks.""" preset_key = canonical_domain_preset_key(domain_name, domain_id) @@ -371,9 +360,11 @@ def _resolve_field( 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 + # Always derived: an action clip is the chunk plus its initial frame. Both + # references fix this, and diffusers rejects a caller-supplied num_frames + # for action runs outright, so a preset must not pin it independently of + # an overridden action_chunk_size. + resolved_num_frames = int(resolved_chunk) + 1 resolved_action_fps = ( float(action_fps) if action_fps is not None else float(resolved_frame_rate) ) 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 162849c1e150..9647d8e64473 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -42,7 +42,7 @@ from .action import ( ACTION_MODE_INVERSE_DYNAMICS, DEFAULT_ACTION_VIEW_POINT, - action_reference_image, + action_reference_size, action_start_frame_offset, build_action_json_prompt, build_vision_condition_mask, @@ -827,16 +827,32 @@ def _preprocess_action_image( 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 + def _preprocess_action_first_frame( + self, image: Any, video: 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() + """Conditioning frame for policy / forward_dynamics as ``[1, 3, H, W]``. + + Either source is accepted: an image goes through PIL, video bytes take + frame 0 off NVDEC. Both land on the padded canvas, so the two entry + points produce the same conditioning for the same picture. + """ + if image is not None: + return self._preprocess_action_image(pil_to_rgb(image), target_h, target_w) + if not isinstance(video, bytes): + raise ValueError( + "Cosmos3 action conditioning requires an image or encoded MP4/AVI " + f"bytes, got {type(video).__name__}." + ) + frames_u8 = decode_video_reference_window( + video, + first_frame=0, + last_frame=0, + target_h=target_h, + target_w=target_w, + device=self.device, + resize="fit", + ) + return self._condition_frames_to_video_tensor(frames_u8).squeeze(2) def _encode_video_tensor(self, video_tensor: torch.Tensor) -> torch.Tensor: """VAE-encode a preprocessed pixel video [1, 3, T, H, W].""" @@ -1138,19 +1154,24 @@ def forward( if resolved_action_fps is None: resolved_action_fps = frame_rate - action_ref_image = None + action_source_h = action_source_w = None if do_action: if isinstance(image, torch.Tensor) or isinstance(video, torch.Tensor): raise ValueError( "Cosmos3 action generation does not support tensor image/video inputs; " - "pass a PIL image, image path, frame directory, video path, or frame list." + "pass a PIL image or image path, or encoded MP4/AVI video bytes." ) - action_ref_image = action_reference_image( + # Header probe for video bytes, direct measure for an image: the + # canvas is the bucket closest to the source's shape, so the size + # has to be known before anything is decoded at it. + action_source_h, action_source_w = action_reference_size( action_mode=normalized_action_mode, image=image, video=video, ) - height, width = resolve_action_size(height, width, action_ref_image, action_resolution) + height, width = resolve_action_size( + height, width, action_source_h, action_source_w, action_resolution + ) if self.rank == 0: logger.info( @@ -1163,7 +1184,8 @@ def forward( 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}" + f"source={action_source_w}x{action_source_h} " + f"(aspect {action_source_w / action_source_h:.3f})" ) if isinstance(prompt, str): @@ -1342,7 +1364,7 @@ def forward( # transformer's collectives so a failure cannot hang the job. synchronize_media_prepare_status(prepare_error) else: - image_tensor = self._preprocess_action_image(action_ref_image, height, width) + image_tensor = self._preprocess_action_first_frame(image, video, height, width) if image_tensor.ndim == 4: video_tensor = ( image_tensor.unsqueeze(2).expand(-1, -1, num_frames, -1, -1).contiguous() diff --git a/tensorrt_llm/media/decoding.py b/tensorrt_llm/media/decoding.py index df8287def3ae..506ce0ab56c9 100644 --- a/tensorrt_llm/media/decoding.py +++ b/tensorrt_llm/media/decoding.py @@ -152,6 +152,40 @@ def resize_fit_pad_uint8(frames: torch.Tensor, target_h: int, target_w: int) -> _RESIZE_MODES = {"cover": resize_center_crop_uint8, "fit": resize_fit_pad_uint8} +def probe_video_dimensions(data: bytes) -> tuple[int, int]: + """Return the reference's ``(height, width)`` without decoding a frame. + + The container header carries the source resolution, so a caller that must + choose its target size from the source aspect ratio - Cosmos3 action picks + the canvas whose shape is closest to the reference - can do so before + committing to a decode. + """ + 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 + + 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) + + try: + demuxer = nvc.CreateDemuxer(_read) + return int(demuxer.Height()), int(demuxer.Width()) + except nvc.PyNvVCException as exc: + raise ValueError( + f"Video reference could not be demuxed (corrupt or not a supported container): {exc}" + ) from exc + + def decode_video_reference_window( data: bytes, *, diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py index 7e3fcbe3f4b3..897a8de0a045 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py @@ -9,7 +9,6 @@ import json -import numpy as np import PIL.Image import pytest import torch @@ -21,11 +20,10 @@ EMBODIMENT_TO_RAW_ACTION_DIM, VIDEO_RES_SIZE_INFO, action_aspect_ratio_label, - action_reference_image, + action_reference_size, build_action_json_prompt, find_closest_target_size, normalize_action_resolution, - normalize_action_video_input, prepare_action_latents, resolve_action_size, resolve_raw_action_dim, @@ -76,25 +74,19 @@ def test_all_buckets_have_aspect_entries(self, action_resolution): class TestResolveActionSize: - @staticmethod - def _ref_image(width: int, height: int) -> PIL.Image.Image: - return PIL.Image.new("RGB", (width, height)) + SOURCE_H, SOURCE_W = 480, 832 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) + assert resolve_action_size(400, 600, self.SOURCE_H, self.SOURCE_W, 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) + assert resolve_action_size(None, None, self.SOURCE_H, self.SOURCE_W, 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) + assert resolve_action_size(400, None, self.SOURCE_H, self.SOURCE_W, 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) + assert resolve_action_size(None, 600, self.SOURCE_H, self.SOURCE_W, 480) == (480, 600) class TestActionResolutionExtraParam: @@ -222,124 +214,51 @@ def test_unknown_resolution_raises(self): 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") +class TestActionReferenceSize: + """Canvas selection needs the source size, not a decoded frame.""" - 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): + def test_policy_measures_image(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"), + PIL.Image.new("RGB", (640, 480), "blue").save(image_path) + assert action_reference_size(action_mode="policy", image=str(image_path), video=None) == ( + 480, + 640, ) - assert ref.getpixel((0, 0)) == (0, 0, 255) - def test_policy_accepts_path_image(self, tmp_path): + def test_policy_prefers_image_over_video(self, tmp_path): image_path = tmp_path / "frame.png" - PIL.Image.new("RGB", (3, 3), "green").save(image_path) - ref = action_reference_image( - action_mode="policy", - image=image_path, + PIL.Image.new("RGB", (320, 240), "blue").save(image_path) + # Bytes would raise if consulted: the image must win. + assert action_reference_size( + action_mode="policy", image=str(image_path), video=b"not-a-video" + ) == (240, 320) + + def test_accepts_pil_image_directly(self): + assert action_reference_size( + action_mode="forward_dynamics", + image=PIL.Image.new("RGB", (256, 128)), video=None, - ) - assert ref.getpixel((0, 0)) == (0, 128, 0) - + ) == (128, 256) -class TestNormalizeActionVideoInput: - def test_none_returns_empty_list(self): - assert normalize_action_video_input(None) == [] + def test_missing_source_raises(self): + with pytest.raises(ValueError, match="requires an image or video"): + action_reference_size(action_mode="policy", image=None, video=None) - def test_empty_list_raises(self): - with pytest.raises(ValueError, match="at least one frame"): - normalize_action_video_input([]) - - def test_image_path_returns_singleton_list(self, tmp_path): + def test_inverse_dynamics_ignores_image(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 + PIL.Image.new("RGB", (640, 480), "blue").save(image_path) + # inverse_dynamics conditions on the clip, so an image is not a source. + with pytest.raises(ValueError, match="requires an image or video"): + action_reference_size(action_mode="inverse_dynamics", image=str(image_path), video=None) + + def test_video_bytes_probe_the_container_header(self, monkeypatch): + """Bytes are measured from the header, never by decoding a frame.""" + import tensorrt_llm.media.decoding as decoding + + monkeypatch.setattr(decoding, "probe_video_dimensions", lambda data: (480, 640)) + assert action_reference_size( + action_mode="inverse_dynamics", image=None, video=b"\x00mp4" + ) == (480, 640) class TestActionJsonPrompt: diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index d48b558c37d3..7b95a02258b3 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -242,6 +242,8 @@ def _assert_valid_action(action: torch.Tensor, *, raw_action_dim: int, chunk_siz 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) @@ -877,8 +879,9 @@ def test_v2v_audio_smoke(self, cosmos3_pipeline): class TestCosmos3Action: ACTION_HEIGHT = 480 ACTION_WIDTH = 832 - ACTION_FRAMES = COSMOS3_ACTION_PARAMS["num_frames"] ACTION_CHUNK = COSMOS3_ACTION_PARAMS["action_chunk_size"] + # Derived, not configured: both references fix the clip at chunk + 1. + ACTION_FRAMES = ACTION_CHUNK + 1 RAW_ACTION_DIM = 10 def test_policy_smoke(self, cosmos3_pipeline): @@ -942,8 +945,6 @@ def test_forward_dynamics_smoke(self, cosmos3_pipeline): 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, @@ -954,8 +955,9 @@ def test_inverse_dynamics_smoke(self, cosmos3_pipeline): action_mode="inverse_dynamics", domain_name="bridge_orig_lerobot", raw_action_dim=self.RAW_ACTION_DIM, - action_chunk_size=NUM_FRAMES, - video=video, + # The clip is chunk + 1 frames, and the fixture holds NUM_FRAMES. + action_chunk_size=NUM_FRAMES - 1, + video=_V2V_FIXTURE_MP4.read_bytes(), ) _assert_valid_video( result.video, @@ -966,15 +968,13 @@ def test_inverse_dynamics_smoke(self, cosmos3_pipeline): _assert_valid_action( result.action, raw_action_dim=self.RAW_ACTION_DIM, - chunk_size=NUM_FRAMES, + chunk_size=NUM_FRAMES - 1, ) assert result.action_mode == "inverse_dynamics" assert result.domain_id == 7 def test_inverse_dynamics_rejects_short_video(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 - 1)] with pytest.raises(ValueError, match="requires at least"): _run_forward( cosmos3_pipeline, @@ -987,7 +987,7 @@ def test_inverse_dynamics_rejects_short_video(self, cosmos3_pipeline): domain_name="bridge_orig_lerobot", raw_action_dim=self.RAW_ACTION_DIM, action_chunk_size=NUM_FRAMES, - video=video, + video=_V2V_FIXTURE_MP4.read_bytes(), ) def test_action_and_audio_rejected(self, cosmos3_pipeline): From 1f9c0d5a9746a2545a5f65eefd095b9b6e80ebce Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Tue, 4 Aug 2026 16:42:48 -0700 Subject: [PATCH 13/35] [None][fix] Reconcile Cosmos3 action with the merged main behaviours Five failures surfaced by running the whole visual-gen suite after the merge; all five are this branch's, and each is a place where the action work and main had independently changed the same behaviour. post_step_fn takes (latents, extra_streams) and returns both, so the distilled conditioning anchor's hook does too. Two tests in test_cosmos3_distilled called it with one argument. T2I + audio: main force-disables audio for image requests so they never reach the audio-weight presence check, and asserts that. The action PR had added a raise a few lines earlier, which pre-empted it. Main owns T2I, so the raise goes; action's own T2I rejection stays. The serve-side rejection of extra_params['video'] tested a workaround that no longer exists. It guarded against a client passing a server-local path, which main now prevents with a stronger mechanism: `video` is declared bytes, so a path fails preflight type validation - verified directly, a str is rejected with "expected type 'bytes'". The obsolete endpoint test is dropped and test_visual_gen_params' path_or_list case becomes the bytes case, keeping the property under test rather than the removed implementation of it. Suite status: 1444 passed. The remaining failures need assets this host lacks - seven Wan VAE checkpoint comparisons and twelve serve e2e errors from `trtllm-serve` not being on PATH - and are untouched by this branch. Signed-off-by: Igor Shovkun --- .../visual_gen/models/cosmos3/pipeline_cosmos3.py | 2 -- .../_torch/visual_gen/test_cosmos3_distilled.py | 5 +++-- .../visual_gen/test_trtllm_serve_endpoints.py | 15 --------------- .../_torch/visual_gen/test_visual_gen_params.py | 8 +++++--- 4 files changed, 8 insertions(+), 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 9647d8e64473..8457494cb663 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -1054,8 +1054,6 @@ def forward( ) 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 # T2I force-disables audio instead of rejecting it, so an image # request never trips the audio-weight presence check below. diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py index e949c0183be5..f389e271d1dd 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py @@ -633,9 +633,10 @@ def test_anchor_writes_only_frame_zero_in_place(self): latents = torch.arange(48, dtype=torch.float32).reshape(1, 4, 3, 2, 2) untouched = latents[:, :, 1:].clone() - returned = post_step_fn(latents) + returned, extra = post_step_fn(latents, None) assert returned is latents, "must write in place, not copy" + assert extra is None, "extra-stream latents pass through untouched" assert torch.all(latents[:, :, 0:1] == self.CLEAN) assert torch.equal(latents[:, :, 1:], untouched) @@ -767,7 +768,7 @@ def test_i2v_request_wires_anchor_and_seeded_steps(self): post_step_fn = captured["post_step_fn"] assert post_step_fn is not None latents = torch.zeros(1, 4, self.T_LAT, self.H_LAT, self.W_LAT) - post_step_fn(latents) + post_step_fn(latents, None) assert torch.all(latents[:, :, 0:1] == self.CLEAN) assert torch.all(latents[:, :, 1:] == 0.0) 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 8ddeecf96cea..69d17ea523ef 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py @@ -808,21 +808,6 @@ def test_sync_video_generation_with_params(self, video_client): assert params.frame_rate == 8 assert params.num_frames == int(2.0 * 8) - def test_sync_video_rejects_extra_params_video_path(self, video_client): - resp = video_client.post( - "/v1/videos/generations", - json={ - "prompt": "reject raw path", - "size": "64x64", - "seconds": 1.0, - "fps": 8, - "extra_params": {"video": "/server/local/path.mp4"}, - }, - headers={"content-type": "application/json"}, - ) - assert resp.status_code == 400 - _assert_llm_envelope(resp.json(), code=400, message_contains="server-local paths") - def test_sync_video_generation_multipart(self, video_client, tmp_path): """Multipart sync request with a real ``input_reference`` file.""" ref_path = tmp_path / "ref.png" 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 8f6a00ca469a..4220e5dd052d 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py @@ -1087,11 +1087,13 @@ def test_literal_extra_param_accepts_numeric_choice(self): extra_param_specs=COSMOS3_EXTRA_SPECS, ) - def test_path_or_list_extra_param_type_checked(self): + def test_video_reference_must_be_bytes(self): + """A server-local path must not reach the worker: the ``video`` contract + is encoded bytes, so a string (or anything else) fails preflight.""" from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import COSMOS3_EXTRA_SPECS - req = self._make_request(extra_params={"video": 123}) - with pytest.raises(ValueError, match="expected type 'path_or_list'"): + req = self._make_request(extra_params={"video": "/server/local/path.mp4"}) + with pytest.raises(ValueError, match="expected type 'bytes'"): from tensorrt_llm.visual_gen.params import validate_visual_gen_params validate_visual_gen_params( From bf1857ff08622426c134af628f19bc955ebd6fae Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Tue, 4 Aug 2026 22:54:03 -0700 Subject: [PATCH 14/35] [None][feat] Serve action requests as a tensor payload, not a silent drop An action request over /v1/videos returned 200 with an mp4 and threw the trajectory away. The action tensor reaches the coordinator intact - to_handle walks every field, so it rides back like video does - but the encoder branch of the route only ever reads output.video, so the thing the request was made for was discarded at the last step. format now resolves against what the request will actually produce. 'auto' selects a tensor payload, which carries every populated modality plus its scalar metadata; an explicit safetensors/pt passes through; an explicit mp4/avi is rejected with 400, because the caller has stated two incompatible things and neither guess is right - encoding it drops the trajectory, ignoring the format disregards what they asked for. The message names the parameter that forced the choice and the formats that work. The rule is declared, not hard-coded: ExtraParamSchema gains requires_tensor_output, Cosmos3 sets it on action_mode, and the route reads the declaration without knowing what an action is. Specs already travel to the coordinator in the READY handshake, so this needs no new plumbing, and any pipeline with a non-encodable modality gets the same behaviour by declaring it. The async route resolves the format before queueing the job and passes it to the background task, so a rejected request never becomes one. Closes the "silently dropping action" review thread on visual_gen/output.py. Signed-off-by: Igor Shovkun --- .../visual_gen/models/cosmos3/defaults.py | 7 +- tensorrt_llm/_torch/visual_gen/pipeline.py | 9 ++ tensorrt_llm/serve/openai_video_routes.py | 75 +++++++++++-- .../visual_gen/test_trtllm_serve_endpoints.py | 105 ++++++++++++++++++ 4 files changed, 184 insertions(+), 12 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index 51b7ca8ee047..2cf6c6700336 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -464,7 +464,12 @@ def _resolve_field( "action_mode": ExtraParamSchema( type="Literal['policy', 'forward_dynamics', 'inverse_dynamics']", default=None, - description="Action generation mode: policy, forward_dynamics, or inverse_dynamics.", + description=( + "Action generation mode: policy, forward_dynamics, or inverse_dynamics. " + "The predicted trajectory is not representable in a video container, so " + "an action request is served as a tensor payload." + ), + requires_tensor_output=True, ), "domain_name": ExtraParamSchema( type="str", diff --git a/tensorrt_llm/_torch/visual_gen/pipeline.py b/tensorrt_llm/_torch/visual_gen/pipeline.py index 2f24be8755e2..f7ac02c7b43c 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline.py @@ -58,6 +58,15 @@ class ExtraParamSchema(StrictBaseModel): "values. Must be a module-level function (specs are pickled to the " "coordinator in the READY handshake).", ) + requires_tensor_output: bool = Field( + default=False, + description="Setting this parameter makes the request produce a result " + "the media encoders cannot represent (a non-image/video modality), so " + "the response must be a tensor payload. Serve resolves 'auto' to a " + "tensor format and rejects an explicit encoder format. Declared here " + "rather than hard-coded in the routes so the serving layer needs no " + "per-model knowledge.", + ) if TYPE_CHECKING: diff --git a/tensorrt_llm/serve/openai_video_routes.py b/tensorrt_llm/serve/openai_video_routes.py index 3601c6f9624b..9524da5f74e4 100644 --- a/tensorrt_llm/serve/openai_video_routes.py +++ b/tensorrt_llm/serve/openai_video_routes.py @@ -47,6 +47,42 @@ def _video_content_type(suffix: str) -> str: _KNOWN_VIDEO_OUTPUT_SUFFIXES = (".mp4", ".avi", ".safetensors", ".pt") +def _resolve_tensor_only_format(fmt, extra_params, extra_param_specs): + """Resolve ``format`` for a request whose result an encoder cannot carry. + + A pipeline marks such parameters with ``requires_tensor_output`` on their + :class:`ExtraParamSchema` (Cosmos3 does so for ``action_mode``: a predicted + trajectory has no representation in a video container). The rule keeps this + route model-agnostic -- it reads the declaration, never the parameter's + meaning: + + * ``auto`` resolves to ``safetensors``, so the default request returns + everything it generated instead of silently dropping a modality; + * an explicit tensor format passes through; + * an explicit encoder format is a contradiction the caller stated -- two + incompatible things in one request -- so it is rejected rather than + guessed at. + """ + if not extra_params or not extra_param_specs: + return fmt + triggered = sorted( + key + for key, spec in extra_param_specs.items() + if getattr(spec, "requires_tensor_output", False) and extra_params.get(key) is not None + ) + if not triggered: + return fmt + if is_tensor_format(fmt): + return fmt + if fmt == "auto": + return _DEFAULT_TENSOR_FORMAT + raise ValueError( + f"format={fmt!r} cannot carry the result of {', '.join(triggered)}: a " + f"video container holds only video. Use format='safetensors' or 'pt', " + f"or omit format so 'auto' selects a payload that carries everything." + ) + + def _preflight_encoder_format(fmt): """Pre-flight an encoder format string before any GPU work. @@ -65,6 +101,9 @@ def _preflight_encoder_format(fmt): raise ValueError(str(exc)) from exc +_DEFAULT_TENSOR_FORMAT = "safetensors" + + def _b64_json_video_response(video_id: str, fmt: str, path: Path) -> JSONResponse: """Build the OpenAI-style ``{id, format, b64_json}`` envelope. @@ -113,7 +152,10 @@ async def openai_video_generation_sync(self, raw_request: Request) -> Response: self.generator, media_storage_path=str(self.media_storage_path), ) - resolved_encoder_fmt = _preflight_encoder_format(request.format) + request_format = _resolve_tensor_only_format( + request.format, request.extra_params, self.generator.extra_param_specs + ) + resolved_encoder_fmt = _preflight_encoder_format(request_format) logger.info( f"Generating video: {video_id} with params: {params} and prompt: {request.prompt}" ) @@ -141,8 +183,8 @@ async def openai_video_generation_sync(self, raw_request: Request) -> Response: status_code=HTTPStatus.INTERNAL_SERVER_ERROR, ) - if is_tensor_format(request.format): - ext = f".{request.format}" + if is_tensor_format(request_format): + ext = f".{request_format}" media_type = "application/octet-stream" # Match the encoder-format path: persist one file per batch # item, ship the first as the route's primary download @@ -152,7 +194,7 @@ async def openai_video_generation_sync(self, raw_request: Request) -> Response: tensor_paths = [ self.media_storage_path / f"{video_id}_{i}{ext}" for i in range(batch_size) ] - saved_paths = output.save(tensor_paths, format=request.format) + saved_paths = output.save(tensor_paths, format=request_format) target = saved_paths[0] latency = time.perf_counter() - sync_video_start logger.info( @@ -160,7 +202,7 @@ async def openai_video_generation_sync(self, raw_request: Request) -> Response: f"generation={getattr(output.metrics, 'generation', 0.0):.3f}s" ) if request.response_format == "b64_json": - return _b64_json_video_response(video_id, request.format, target) + return _b64_json_video_response(video_id, request_format, target) return FileResponse(str(target), media_type=media_type, filename=target.name) # Encoder formats: one file per item; ship the first item as @@ -335,7 +377,10 @@ async def openai_video_generation_async( declared_defaults=self.generator.executor.default_generation_params, extra_param_specs=self.generator.executor.extra_param_specs, ) - _preflight_encoder_format(request.format) + request_format = _resolve_tensor_only_format( + request.format, request.extra_params, self.generator.extra_param_specs + ) + _preflight_encoder_format(request_format) logger.info( f"Generating video: {video_id} with params: {params} and prompt: {request.prompt}" ) @@ -361,6 +406,7 @@ async def openai_video_generation_async( video_id=video_id, request=request, params=params, + request_format=request_format, ) ) self.video_gen_tasks[video_id] = task @@ -386,8 +432,15 @@ async def _generate_video_background( video_id: str, request: VideoGenerationRequest, params: VisualGenParams, + request_format: str, ): - """Background task to generate video and save to storage.""" + """Background task to generate video and save to storage. + + ``request_format`` is the format already resolved by the route (see + :func:`_resolve_tensor_only_format`), not ``request.format``: the + resolution happens before the job is queued so a rejected request never + becomes a background task. + """ try: background_start = time.perf_counter() future = self.generator.generate_async(inputs=request.prompt, params=params) @@ -403,18 +456,18 @@ async def _generate_video_background( await VIDEO_STORE.upsert(video_id, job) return - if is_tensor_format(request.format): + if is_tensor_format(request_format): # One tensor file per batch item, mirroring the encoder # path; the async job records all paths on # ``output_paths`` so subsequent GETs can find each item. batch_size = output.video.shape[0] if output.video.dim() == 5 else 1 tensor_paths = [ - self.media_storage_path / f"{video_id}_{i}.{request.format}" + self.media_storage_path / f"{video_id}_{i}.{request_format}" for i in range(batch_size) ] - saved_paths = output.save(tensor_paths, format=request.format) + saved_paths = output.save(tensor_paths, format=request_format) else: - resolved_fmt, _ = resolve_video_format(request.format) + resolved_fmt, _ = resolve_video_format(request_format) batch_size = output.video.shape[0] if output.video.dim() == 5 else 1 paths_in = [self.media_storage_path / f"{video_id}_{i}" for i in range(batch_size)] saved_paths = output.save( 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 69d17ea523ef..8cb87240c8f6 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py @@ -355,6 +355,28 @@ def video_client(tmp_path): os.environ.pop("TRTLLM_MEDIA_STORAGE_PATH", None) +@pytest.fixture() +def action_video_client(tmp_path): + """Video client whose pipeline declares a tensor-only extra param. + + Stands in for Cosmos3 action: the route must learn "this result needs a + tensor payload" from the spec, never from the parameter's name. + """ + from tensorrt_llm._torch.visual_gen.pipeline import ExtraParamSchema + + gen = MockVisualGen(video_output=_make_dummy_video_tensor()) + specs = { + "action_mode": ExtraParamSchema(type="str", default=None, requires_tensor_output=True), + } + gen.executor.extra_param_specs = specs + type(gen).extra_param_specs = property(lambda self: specs) + os.environ["TRTLLM_MEDIA_STORAGE_PATH"] = str(tmp_path) + client = _create_server(gen) + yield client + os.environ.pop("TRTLLM_MEDIA_STORAGE_PATH", None) + del type(gen).extra_param_specs + + @pytest.fixture() def video_audio_client(tmp_path): """TestClient backed by a MockVisualGen that produces videos with audio.""" @@ -2202,3 +2224,86 @@ def test_async_url_still_returns_file_response(self, video_client): # AVI FileResponse carries ``video/x-msvideo``; the b64_json # branch would have set ``application/json``. assert content.headers["content-type"] == "video/x-msvideo" + + +class TestTensorOnlyFormatResolution: + """A request whose result an encoder cannot carry must not be served as video.""" + + @staticmethod + def _post(client, **body): + return client.post( + "/v1/videos/generations", + json={"prompt": "pick up the block", "size": "64x64", "seconds": 1.0, "fps": 8, **body}, + headers={"content-type": "application/json"}, + ) + + def test_auto_resolves_to_tensor_payload(self, action_video_client): + resp = self._post(action_video_client, extra_params={"action_mode": "policy"}) + assert resp.status_code == 200 + # 'auto' would otherwise have produced an encoded video, silently + # dropping the modality the request was made for. + assert resp.headers["content-type"] == "application/octet-stream" + assert resp.headers["content-disposition"].endswith('.safetensors"') + + def test_explicit_tensor_format_passes_through(self, action_video_client): + resp = self._post(action_video_client, format="pt", extra_params={"action_mode": "policy"}) + assert resp.status_code == 200 + assert resp.headers["content-disposition"].endswith('.pt"') + + @pytest.mark.parametrize("fmt", ["mp4", "avi"]) + def test_explicit_encoder_format_is_rejected(self, action_video_client, fmt): + """The caller stated two incompatible things; guessing either way is wrong.""" + resp = self._post(action_video_client, format=fmt, extra_params={"action_mode": "policy"}) + assert resp.status_code == 400 + body = resp.json() + _assert_llm_envelope(body, code=400, message_contains="action_mode") + # The message must name the way out, not just the refusal. + assert "safetensors" in body["message"] + + def test_untriggered_request_keeps_encoder_default(self, action_video_client): + """No tensor-only param set -> ordinary video request, unchanged.""" + resp = self._post(action_video_client) + assert resp.status_code == 200 + assert resp.headers["content-type"].startswith("video/") + + +class TestTensorOnlyFormatRule: + """Unit coverage of the resolution rule itself, without a server.""" + + @staticmethod + def _specs(**flags): + from tensorrt_llm._torch.visual_gen.pipeline import ExtraParamSchema + + return { + name: ExtraParamSchema(type="str", default=None, requires_tensor_output=flag) + for name, flag in flags.items() + } + + def test_no_specs_is_a_noop(self): + from tensorrt_llm.serve.openai_video_routes import _resolve_tensor_only_format + + assert _resolve_tensor_only_format("auto", {"action_mode": "policy"}, None) == "auto" + assert _resolve_tensor_only_format("auto", None, self._specs(a=True)) == "auto" + + def test_only_a_declared_param_triggers(self): + from tensorrt_llm.serve.openai_video_routes import _resolve_tensor_only_format + + specs = self._specs(action_mode=True, stg_scale=False) + # a non-declaring param must not force a tensor payload + assert _resolve_tensor_only_format("auto", {"stg_scale": 2.0}, specs) == "auto" + assert ( + _resolve_tensor_only_format("auto", {"action_mode": "policy"}, specs) == "safetensors" + ) + + def test_null_value_does_not_trigger(self): + from tensorrt_llm.serve.openai_video_routes import _resolve_tensor_only_format + + specs = self._specs(action_mode=True) + assert _resolve_tensor_only_format("auto", {"action_mode": None}, specs) == "auto" + + def test_encoder_format_raises_naming_the_parameter(self): + from tensorrt_llm.serve.openai_video_routes import _resolve_tensor_only_format + + specs = self._specs(action_mode=True) + with pytest.raises(ValueError, match="action_mode"): + _resolve_tensor_only_format("mp4", {"action_mode": "policy"}, specs) From 79ee31eaf7bfbda78019f04068695a21daa7402a Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Tue, 4 Aug 2026 22:57:22 -0700 Subject: [PATCH 15/35] [None][doc] Cosmos3 action README: encoded clip input, tensor-payload output --video_path for inverse_dynamics took a frame directory before the NVDEC swap; it is an MP4/AVI clip now, decoded worker-side like V2V. Also states why action runs are tensor payloads and what trtllm-serve does with format. Signed-off-by: Igor Shovkun --- examples/visual_gen/models/cosmos3/README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/visual_gen/models/cosmos3/README.md b/examples/visual_gen/models/cosmos3/README.md index 83241537f3d7..9086ea1ac152 100644 --- a/examples/visual_gen/models/cosmos3/README.md +++ b/examples/visual_gen/models/cosmos3/README.md @@ -7,7 +7,7 @@ Cosmos3 supports the following generation modes from a single checkpoint: - **I2V / TI2V** — image-conditioned video (`prompts/i2v.json`). Condition on a reference frame via the prompt file's `vision_path` or `--image_path`. The image may be a local path, a `file://` / `http(s)://` URL, or a `data:` URI. - **V2V** — video-conditioned video (`prompts/v2v.json`). Condition on a reference video via `--video_path` (a local MP4/AVI file). Only the first (or last, per `condition_video_keep`) `max(condition_video_latent_indexes) * 4 + 1` input frames condition the output (5 by default); the encoded bytes pass through and each worker decodes just that window on NVDEC (see [Media I/O dependencies](#media-io-dependencies)). - **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. Save action runs as `safetensors` or `pt` so the rollout and action tensor stay in one payload. +- **Action** — policy / forward dynamics / inverse dynamics generation (pass `--action_mode`); `inverse_dynamics` reads its observation clip from `--video_path` (MP4/AVI, decoded on worker NVDEC like V2V). Action and audio generation are mutually exclusive. A predicted trajectory has no representation in a video container, so action runs are saved as `safetensors` or `pt`, keeping the rollout and the action tensor in one payload — over `trtllm-serve` the default `format=auto` selects that payload automatically, and an explicit `mp4`/`avi` is rejected. ## Checkpoints @@ -132,7 +132,7 @@ python cosmos3.py --model nvidia/Cosmos3-Nano \ # 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/ \ + --video_path /path/to/observation_clip.mp4 \ --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml \ --action_mode inverse_dynamics \ --domain_name bridge_orig_lerobot \ From 1b3daea1f6156644836b0786058df2a8d5c3f53b Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 5 Aug 2026 09:49:17 -0700 Subject: [PATCH 16/35] [None][fix] Cosmos3 action: review follow-ups across serve, media and sync Eight items from review, each with a test: - Wan 2.2 5B I2V pinned its first frame through a one-argument post-step callback while the shared loop now passes two. Restore the two-argument signature. - Action requests inherited the video sampling recipe through infer(): unset steps, guidance and frame rate were materialized as 35 / 6.0 / 24 fps, so the action branch's 30 / 1.0 and the embodiment's frame rate (bridge 5, av 10) could never win, and CFG doubled the transformer work. Pass them through unset for action, as height and width already were. frame_rate stays materialized in the pipeline defaults because the serve layer derives num_frames from seconds x frame_rate. - The example assigned the reference clip to both params.image and the video extra param, which the pipeline rejects, so the documented inverse_dynamics invocation could not run. - pil_to_rgb opened every string as a local path, while both bundled action prompts carry https frame URLs. Route through the repo's URL-aware loader, and resolve the reference once per request instead of once per read. - Both per-rank reads of the action reference (the canvas probe and the policy / forward_dynamics decode) now converge like inverse_dynamics and V2V already did; a rank-local failure otherwise leaves healthy ranks in the transformer's collectives. - _apply_flow_shift rebuilds the action scheduler alongside video and audio. Unreachable on Nano and Edge, whose UniPC config takes the karras branch before flow_shift is read, but the distilled checkpoints do not share that. - A domain_id that contradicts domain_name is rejected instead of silently applying one robot's timing to another robot's weights. - DomainAwareLinear's range check used a device predicate as a Python condition, forcing two blocking device-to-host syncs on every denoise step. It moves to a once-per-request check, plus a free host-side check on the scalar before the reference is decoded. Signed-off-by: Igor Shovkun --- examples/visual_gen/models/cosmos3/cosmos3.py | 2 - .../visual_gen/models/cosmos3/action.py | 17 ++- .../models/cosmos3/pipeline_cosmos3.py | 111 +++++++++++++----- .../models/cosmos3/transformer_cosmos3.py | 27 ++++- .../visual_gen/models/wan/pipeline_wan.py | 5 +- .../_torch/visual_gen/test_cosmos3_action.py | 42 +++++++ .../visual_gen/test_cosmos3_distilled.py | 15 +++ .../visual_gen/test_cosmos3_pipeline.py | 71 +++++++++++ .../visual_gen/test_cosmos3_transformer.py | 41 +++++++ 9 files changed, 289 insertions(+), 42 deletions(-) diff --git a/examples/visual_gen/models/cosmos3/cosmos3.py b/examples/visual_gen/models/cosmos3/cosmos3.py index c4917b8a1cc4..22fc33111cfb 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3.py +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -340,8 +340,6 @@ def main(): params = visual_gen.default_params if image_path is not None: params.image = image_path - if args.action_mode is not None and args.video_path is not None: - params.image = args.video_path negative_prompt_path = _resolve_path(args.negative_prompt) if args.negative_prompt is not None: diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py index 1d32e90cac4a..05463cde9d04 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py @@ -247,6 +247,17 @@ def resolve_domain_id( resolved = int(domain_id) if resolved < 0: raise ValueError(f"Cosmos3 domain_id must be non-negative, got {resolved}.") + # domain_id wins so unlisted embodiments stay reachable, but a caller + # that passes both and disagrees would otherwise silently get a + # trajectory in a different robot's dialect. + if domain_name is not None and str(domain_name).strip(): + key = str(domain_name).strip().lower() + named_id = EMBODIMENT_TO_DOMAIN_ID.get(key) + if named_id is not None and named_id != resolved: + raise ValueError( + f"Cosmos3 domain_id={resolved} contradicts domain_name={domain_name!r}, " + f"which maps to domain_id={named_id}. Pass only one, or make them agree." + ) return resolved if domain_name is None or str(domain_name).strip() == "": @@ -380,7 +391,11 @@ def find_closest_target_size(h: int, w: int, resolution: str | int) -> tuple[int def pil_to_rgb(value: Any) -> PIL.Image.Image: if isinstance(value, (str, Path)): - return PIL.Image.open(Path(value)).convert("RGB") + # load_image, not PIL.Image.open: the bundled action prompts carry + # https:// frame references, and it also handles file:// and data: URIs. + from tensorrt_llm.inputs.utils import load_image + + return load_image(str(value), format="pil").convert("RGB") if isinstance(value, PIL.Image.Image): return value.convert("RGB") raise TypeError( 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 8457494cb663..9c9813e75334 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -316,19 +316,25 @@ def _run_warmup(self, height: int, width: int, num_frames: int, steps: int) -> N 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. + """Rebuild every stream scheduler 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. + Video, audio and action denoise in lockstep in one loop, so a mode that + rebuilds only the video scheduler leaves the side streams on the + checkpoint's sigmas and the 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 - ) + for stream in ("audio_scheduler", "action_scheduler"): + scheduler = getattr(self, stream, None) + if scheduler is not None: + setattr( + self, + stream, + self.sampling.set_flow_shift( + scheduler, target_shift, use_karras_sigmas=use_karras_sigmas + ), + ) def infer(self, req): extra_params = req.params.extra_params or {} @@ -341,15 +347,30 @@ def infer(self, req): def resolved(value, field_name): return value if value is not None else mode_params[field_name] - # Action sizes its canvas from the resolution bucket and the reference - # frame's aspect (resolve_action_size), so leaving height/width unset - # here is what lets forward() do that; filling them from the video - # table would pin every action request to 720p. + # Action resolves every one of these itself, from the embodiment preset + # and COSMOS3_ACTION_PARAMS: the canvas from the resolution bucket, the + # frame rate from the robot, and steps/guidance from the action recipe. + # Filling them from the video table here would leave forward()'s + # fallbacks dead and silently run action at 720p, 35 steps, guidance 6 + # (CFG on, double the work) and 24 fps. is_action = extra_params.get("action_mode") is not None height = req.params.height if is_action else resolved(req.params.height, "height") width = req.params.width if is_action else 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") + num_inference_steps = ( + req.params.num_inference_steps + if is_action + else resolved(req.params.num_inference_steps, "num_inference_steps") + ) + guidance_scale = ( + req.params.guidance_scale + if is_action + else resolved(req.params.guidance_scale, "guidance_scale") + ) + # frame_rate keeps a materialised video default (the serve layer derives + # num_frames from seconds x frame_rate), so action has to drop it here + # instead: an incoming 24.0 is indistinguishable from a caller who chose + # 24, and the embodiment preset (bridge 5, av 10) would never win. + frame_rate = None if is_action else req.params.frame_rate video = extra_params.get("video") # encoded MP4/AVI bytes (the extra-param contract) return self.forward( @@ -363,7 +384,7 @@ def resolved(value, field_name): guidance_scale=guidance_scale, seed=req.params.seed, max_sequence_length=req.params.max_sequence_length, - frame_rate=req.params.frame_rate, + frame_rate=frame_rate, use_duration_template=extra_params.get( "use_duration_template", COSMOS3_EXTRA_SPECS["use_duration_template"].default, @@ -1162,11 +1183,24 @@ def forward( # Header probe for video bytes, direct measure for an image: the # canvas is the bucket closest to the source's shape, so the size # has to be known before anything is decoded at it. - action_source_h, action_source_w = action_reference_size( - action_mode=normalized_action_mode, - image=image, - video=video, - ) + # + # This is the first per-rank read of the reference, so it converges + # like the decode below: a missing file or an unreachable URL on one + # rank must not leave the others walking into the collectives. + probe_error: Optional[Exception] = None + try: + if image is not None and not isinstance(image, PIL.Image.Image): + # Resolve once: the bundled prompts point at https frames, + # and the probe and the conditioning frame both read it. + image = pil_to_rgb(image) + action_source_h, action_source_w = action_reference_size( + action_mode=normalized_action_mode, + image=image, + video=video, + ) + except Exception as exc: + probe_error = exc + synchronize_media_prepare_status(probe_error) height, width = resolve_action_size( height, width, action_source_h, action_source_w, action_resolution ) @@ -1318,6 +1352,12 @@ def forward( domain_name=domain_name, require_explicit=True, ) + num_domains = getattr(self.transformer, "num_embodiment_domains", None) + if num_domains is not None and not 0 <= action_domain_id < num_domains: + raise ValueError( + f"Cosmos3 action domain_id must be in [0, {num_domains}), " + f"got {action_domain_id}." + ) action_frame_offset = action_start_frame_offset( normalized_action_mode, action_chunk_size, num_frames ) @@ -1362,19 +1402,26 @@ def forward( # transformer's collectives so a failure cannot hang the job. synchronize_media_prepare_status(prepare_error) else: - image_tensor = self._preprocess_action_first_frame(image, video, height, width) - if image_tensor.ndim == 4: - video_tensor = ( - image_tensor.unsqueeze(2).expand(-1, -1, num_frames, -1, -1).contiguous() + prepare_error = None + try: + image_tensor = self._preprocess_action_first_frame(image, video, 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, ) - else: - video_tensor = image_tensor - latents, velocity_mask, condition_latents = self._prepare_latents_action_video( - video_tensor, - normalized_action_mode, - num_frames, - generator, - ) + except Exception as exc: + prepare_error = exc + synchronize_media_prepare_status(prepare_error) image_latent = 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 c9d050515332..248006ec149e 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -235,6 +235,20 @@ def post_load_weights(self) -> None: self.fc.to(self.dtype) self.bias.to(self.dtype) + def validate_domain_ids(self, domain_id: torch.Tensor) -> None: + """Range-check the ids. Reads a device tensor, so call once per request. + + Out-of-range ids index ``nn.Embedding`` out of bounds, which on GPU is a + device-side assert with no useful message; this turns it into a real + error. Kept out of forward() because the ``if`` on a device predicate is + a blocking sync, and forward() runs twice on every denoise step. + """ + 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()}." + ) + def forward(self, x: torch.Tensor, domain_id: torch.Tensor) -> torch.Tensor: if domain_id.ndim == 0: domain_id = domain_id.unsqueeze(0) @@ -244,11 +258,6 @@ def forward(self, x: torch.Tensor, domain_id: torch.Tensor) -> torch.Tensor: "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) @@ -927,6 +936,7 @@ def __init__(self, model_config: DiffusionModelConfig): self.cached_kv = None self.cached_freqs_gen = None + self.domain_ids_validated = False self.__post_init__() @@ -1155,6 +1165,7 @@ def unpack_action(tokens: torch.Tensor) -> torch.Tensor: def reset_cache(self): self.cached_kv = None self.cached_freqs_gen = None + self.domain_ids_validated = False def forward( self, @@ -1289,6 +1300,12 @@ def forward( action_domain_ids_tensor = torch.zeros( action_latents.shape[0], dtype=torch.long, device=action_latents.device ) + if not self.domain_ids_validated: + # Once per request, alongside the other first-step host work. + self.action_proj_in.validate_domain_ids( + action_domain_ids_tensor.to(dtype=torch.long).reshape(-1) + ) + self.domain_ids_validated = True T_action = action_latents.shape[1] hidden_action = self.action_proj_in( self.pack_action(action_latents), action_domain_ids_tensor diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py index 8230f10823e5..88f90b719d84 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py @@ -626,10 +626,11 @@ def forward_fn( ) # Pin reference image to latent after each scheduler step (Wan 2.2 5B I2V only) - def _pin_i2v_first_frame(x): - return ((1 - i2v_first_frame_mask) * i2v_condition + i2v_first_frame_mask * x).to( + def _pin_i2v_first_frame(x, extra_stream_latents): + pinned = ((1 - i2v_first_frame_mask) * i2v_condition + i2v_first_frame_mask * x).to( self.dtype ) + return pinned, extra_stream_latents post_step_fn = _pin_i2v_first_frame if (self.is_wan22_5b and is_i2v) else None diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py index 897a8de0a045..a42992705d61 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py @@ -26,6 +26,7 @@ normalize_action_resolution, prepare_action_latents, resolve_action_size, + resolve_domain_id, resolve_raw_action_dim, ) from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( @@ -251,6 +252,25 @@ def test_inverse_dynamics_ignores_image(self, tmp_path): with pytest.raises(ValueError, match="requires an image or video"): action_reference_size(action_mode="inverse_dynamics", image=str(image_path), video=None) + def test_https_reference_goes_through_the_repo_loader(self, monkeypatch): + """Bundled action prompts point at https:// frames, so a bare + PIL.Image.open(path) would fail on every one of them.""" + import tensorrt_llm.inputs.utils as inputs_utils + + requested = [] + + def fake_load_image(source, format="pt", device="cpu"): + requested.append((source, format)) + return PIL.Image.new("RGB", (640, 480), "blue") + + monkeypatch.setattr(inputs_utils, "load_image", fake_load_image) + assert action_reference_size( + action_mode="policy", + image="https://example.invalid/frame.png", + video=None, + ) == (480, 640) + assert requested == [("https://example.invalid/frame.png", "pil")] + def test_video_bytes_probe_the_container_header(self, monkeypatch): """Bytes are measured from the header, never by decoding a frame.""" import tensorrt_llm.media.decoding as decoding @@ -609,3 +629,25 @@ def test_forward_dynamics_raw_dim_mismatch_raises(self): dtype=torch.float32, action_input=[[0.0, 1.0], [2.0, 3.0]], ) + + +class TestResolveDomainId: + """domain_id wins, but a caller that contradicts itself is a real mistake: + the wrong embodiment yields a fluent trajectory in another robot's dialect.""" + + def test_agreeing_pair_is_accepted(self): + assert resolve_domain_id(domain_id=7, domain_name="bridge_orig_lerobot") == 7 + + def test_contradicting_pair_raises(self): + with pytest.raises(ValueError, match="contradicts domain_name"): + resolve_domain_id(domain_id=20, domain_name="bridge_orig_lerobot") + + def test_unlisted_name_leaves_domain_id_authoritative(self): + assert resolve_domain_id(domain_id=31, domain_name="some-new-robot") == 31 + + def test_name_alone_still_resolves(self): + assert resolve_domain_id(domain_name="fractal") == 20 + + def test_negative_domain_id_raises(self): + with pytest.raises(ValueError, match="must be non-negative"): + resolve_domain_id(domain_id=-1) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py index f389e271d1dd..7f5927ad6581 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py @@ -389,6 +389,21 @@ def test_explicit_values_pass_through(self): assert got["num_inference_steps"] == 20 assert got["width"] == COSMOS3_T2I_PARAMS["width"] + def test_action_keeps_every_mode_field_unset(self): + """Action resolves its own canvas, steps, guidance and frame rate from + the embodiment preset. Filling them from the video table here would run + every action request at 720p, 35 steps and guidance 6 (CFG on).""" + req = _fake_request("video", extra_params={"action_mode": "policy"}) + got = self._captured_forward_kwargs(_bare_pipeline(), req) + for field in ("height", "width", "num_inference_steps", "guidance_scale", "frame_rate"): + assert got[field] is None, field + + def test_video_keeps_its_materialised_frame_rate(self): + """Only action drops it: the serve layer derives num_frames from + seconds x frame_rate, so the video default has to stay materialised.""" + got = self._captured_forward_kwargs(_bare_pipeline(), _fake_request("video")) + assert got["frame_rate"] == COSMOS3_720P_PARAMS["frame_rate"] + def test_distilled_merged_defaults_pass_through(self): req = _fake_request("image", num_inference_steps=4, guidance_scale=1.0) got = self._captured_forward_kwargs(_bare_pipeline(sampling=_distilled_policy()), req) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index 7b95a02258b3..b1cac503e967 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -801,6 +801,30 @@ def fake_tokenize_prompt(text, max_sequence_length, use_system_prompt, system_pr assert rebuilt == [("video", 10.0, False), ("audio", 10.0, False)] + def test_apply_flow_shift_rebuilds_every_stream_scheduler(self): + """Action denoises in the same loop as video, on its own scheduler + instance, so a request that shifts the schedule must move all of them.""" + pipeline = Cosmos3OmniMoTPipeline.__new__(Cosmos3OmniMoTPipeline) + rebuilt = [] + + class FakeSampling: + def set_flow_shift(self, scheduler, target, *, use_karras_sigmas=None): + rebuilt.append((scheduler.name, target, use_karras_sigmas)) + return scheduler + + pipeline.sampling = FakeSampling() + pipeline.scheduler = SimpleNamespace(name="video") + pipeline.audio_scheduler = SimpleNamespace(name="audio") + pipeline.action_scheduler = SimpleNamespace(name="action") + + pipeline._apply_flow_shift(10.0, use_karras_sigmas=False) + + assert rebuilt == [ + ("video", 10.0, False), + ("audio", 10.0, False), + ("action", 10.0, False), + ] + def test_image_and_video_rejected(self, cosmos3_pipeline): with pytest.raises(ValueError, match="not both image and video"): _run_forward( @@ -990,6 +1014,53 @@ def test_inverse_dynamics_rejects_short_video(self, cosmos3_pipeline): video=_V2V_FIXTURE_MP4.read_bytes(), ) + def test_out_of_range_domain_id_rejected_before_decode(self, cosmos3_pipeline): + _require_action_pipeline(cosmos3_pipeline) + with pytest.raises(ValueError, match=r"domain_id must be in \[0, \d+\)"): + _run_forward( + cosmos3_pipeline, + image=_make_test_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_id=10_000, + raw_action_dim=self.RAW_ACTION_DIM, + action_chunk_size=self.ACTION_CHUNK, + ) + + def test_first_frame_failure_is_synchronized(self, cosmos3_pipeline, monkeypatch): + """Every rank decodes its own reference. A failure on one rank has to + reach the others before the transformer's collectives, or the job hangs.""" + _require_action_pipeline(cosmos3_pipeline) + from tensorrt_llm._torch.visual_gen.models.cosmos3 import pipeline_cosmos3 + + seen = [] + real = pipeline_cosmos3.synchronize_media_prepare_status + + def spy(error): + seen.append(error) + return real(error) + + monkeypatch.setattr(pipeline_cosmos3, "synchronize_media_prepare_status", spy) + + with pytest.raises(Exception): + _run_forward( + cosmos3_pipeline, + image="/nonexistent/action_reference_frame.png", + 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 seen and isinstance(seen[0], Exception) + def test_action_and_audio_rejected(self, cosmos3_pipeline): _require_action_pipeline(cosmos3_pipeline) with pytest.raises(ValueError, match="joint action and audio"): diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py index e04d7a525222..af7eb8922f06 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py @@ -480,6 +480,47 @@ def test_forward_with_action(self, action_model_config): 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_domain_ids_validated_once_per_request(self, action_model_config): + """The range check reads a device tensor, so it is a blocking sync. It + belongs on the first step of a request, not on every denoise step.""" + 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) + + calls = [] + real_validate = model.action_proj_in.validate_domain_ids + model.action_proj_in.validate_domain_ids = lambda ids: ( + calls.append(ids), + real_validate(ids), + )[1] + + def run_step(): + with torch.inference_mode(): + 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, + ) + + run_step() + run_step() + assert len(calls) == 1 + + model.reset_cache() + run_step() + assert len(calls) == 2 + @pytest.mark.high_cuda_memory def test_forward_with_action_domain_id_out_of_range_raises(self, action_model_config): cfg = action_model_config.pretrained_config From b573c00315f1d205fe01865769bea5a3e2d6e1d7 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 5 Aug 2026 10:32:37 -0700 Subject: [PATCH 17/35] [None][fix] Cosmos3 action: keep an explicit frame_rate, drop the V2V prompt Two findings from a second review pass. infer() nulled frame_rate for every action request, not only the one the executor had materialized, so a caller could not override the embodiment preset at all. The explicit/default distinction is not recoverable here -- both VisualGen.default_params and parse_visual_gen_params construct VisualGenParams(**defaults), so pydantic marks every field as set before the pipeline sees the request -- but the value is: drop frame_rate only while it still equals what default_generation_params supplied. An explicit value that happens to equal the video default remains indistinguishable. is_v2v was true whenever video bytes were present, which is how an action reference arrives too. inverse_dynamics therefore always forced the system prompt, and a video-backed policy request tokenized differently from the same frame passed as an image. cosmos-framework attaches a system prompt for image editing and transfer only, never for action, so the checkpoint default is the right answer. Re-scored inverse_dynamics against the framework golden under the corrected prompt: mse 0.009148, corr +0.9865, against golden_mse_max 0.05. Signed-off-by: Igor Shovkun --- .../models/cosmos3/pipeline_cosmos3.py | 19 +++++-- .../visual_gen/test_cosmos3_distilled.py | 7 +++ .../visual_gen/test_cosmos3_pipeline.py | 57 +++++++++++++++++++ 3 files changed, 77 insertions(+), 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 9c9813e75334..604120b51cb1 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -366,11 +366,14 @@ def resolved(value, field_name): if is_action else resolved(req.params.guidance_scale, "guidance_scale") ) - # frame_rate keeps a materialised video default (the serve layer derives - # num_frames from seconds x frame_rate), so action has to drop it here - # instead: an incoming 24.0 is indistinguishable from a caller who chose - # 24, and the embodiment preset (bridge 5, av 10) would never win. - frame_rate = None if is_action else req.params.frame_rate + # frame_rate cannot stay None in the pipeline defaults the way the four + # above do: the serve layer derives num_frames from seconds x frame_rate + # before the pipeline sees the request. So for action, "unset" survives + # only as "still equal to what the executor materialized" — any other + # value is the caller's own and outranks the embodiment preset. + frame_rate = req.params.frame_rate + if is_action and frame_rate == self.default_generation_params.get("frame_rate"): + frame_rate = None video = extra_params.get("video") # encoded MP4/AVI bytes (the extra-param contract) return self.forward( @@ -1060,7 +1063,11 @@ def forward( raise ValueError( "Cosmos3 video-to-video generation is supported only for video outputs." ) - is_v2v = video is not None and not is_t2i + # Action reads its reference through the same `video` bytes, but it is + # not V2V: the reference is an observation, not a clip to continue. Left + # in, an action request's prompt would depend on whether the caller + # passed the same frame as an image or as a one-frame clip. + is_v2v = video is not None and not is_t2i and not do_action if use_system_prompt is None: # V2V always wants it; otherwise the checkpoint declares the default. use_system_prompt = is_v2v or self.default_use_system_prompt diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py index 7f5927ad6581..950400bee46e 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py @@ -404,6 +404,13 @@ def test_video_keeps_its_materialised_frame_rate(self): got = self._captured_forward_kwargs(_bare_pipeline(), _fake_request("video")) assert got["frame_rate"] == COSMOS3_720P_PARAMS["frame_rate"] + def test_action_keeps_an_explicit_frame_rate(self): + """Dropping the materialised default is what lets the embodiment preset + win; dropping a caller's own value would make it unsettable.""" + req = _fake_request("video", frame_rate=30.0, extra_params={"action_mode": "policy"}) + got = self._captured_forward_kwargs(_bare_pipeline(), req) + assert got["frame_rate"] == 30.0 + def test_distilled_merged_defaults_pass_through(self): req = _fake_request("image", num_inference_steps=4, guidance_scale=1.0) got = self._captured_forward_kwargs(_bare_pipeline(sampling=_distilled_policy()), req) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index b1cac503e967..2e4b56df4008 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -801,6 +801,63 @@ def fake_tokenize_prompt(text, max_sequence_length, use_system_prompt, system_pr assert rebuilt == [("video", 10.0, False), ("audio", 10.0, False)] + def test_action_video_is_not_classified_as_v2v(self): + """An action reference arrives as the same `video` bytes V2V uses, but + it is an observation, not a clip to continue. Treating it as V2V forces + the system prompt, so the same frame would tokenize differently + depending on whether it was passed as an image or a one-frame clip.""" + pipeline = Cosmos3OmniMoTPipeline.__new__(Cosmos3OmniMoTPipeline) + pipeline.transformer = SimpleNamespace( + device=torch.device("cpu"), num_embodiment_domains=32 + ) + pipeline.audio_gen = False + pipeline.action_gen = True + pipeline.default_use_system_prompt = False + token_calls = [] + + 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 generation_default_overrides(self): + return {} + + def set_flow_shift(self, scheduler, target, *, use_karras_sigmas=None): + return scheduler + + def fake_tokenize_prompt(text, max_sequence_length, use_system_prompt, system_prompt=None): + token_calls.append(use_system_prompt) + raise StopAfterTokenize + + pipeline.scheduler = SimpleNamespace(config=SimpleNamespace(flow_shift=1.0)) + pipeline.sampling = FakeSampling() + pipeline._tokenize_prompt = fake_tokenize_prompt + + with pytest.raises(StopAfterTokenize): + pipeline.forward( + prompt="pick up the block", + video=_V2V_FIXTURE_MP4.read_bytes(), + num_frames=NUM_FRAMES, + num_inference_steps=1, + guidance_scale=1.0, + seed=1, + max_sequence_length=8, + use_system_prompt=None, + use_guardrails=False, + action_mode="inverse_dynamics", + domain_name="bridge_orig_lerobot", + raw_action_dim=10, + action_chunk_size=NUM_FRAMES - 1, + ) + + assert token_calls[0] is False + def test_apply_flow_shift_rebuilds_every_stream_scheduler(self): """Action denoises in the same loop as video, on its own scheduler instance, so a request that shifts the schedule must move all of them.""" From 29a2de978ffcbf06c4863b88bec808d19d0b46dd Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 5 Aug 2026 11:16:49 -0700 Subject: [PATCH 18/35] [None][fix] Cosmos3: drop a shadowed VAE encode and duplicated prose _encode_video_tensor was defined twice in Cosmos3OmniMoTPipeline, byte for byte. The action-section copy shadowed the one the I2V and V2V paths read, so a future fix to the first would have silently done nothing. It is not action code either -- it VAE-encodes any preprocessed pixel video. Keep the original. _prepare_action_latents defaulted action_dim to 64 if the transformer lacked it, but the transformer applies that same default when it builds the action heads and this path only runs when action_gen is true, so the fallback covered a state that cannot occur. defaults.py's module docstring described COSMOS3_DOMAIN_PRESETS, the preset merge and raw_action_dim's absence from the presets. The first and third were already stated at their definitions; the second describes resolve_domain_action_config, whose own docstring was one line. Moved there. Signed-off-by: Igor Shovkun --- .../visual_gen/models/cosmos3/defaults.py | 27 ++++++---------- .../models/cosmos3/pipeline_cosmos3.py | 32 +------------------ 2 files changed, 11 insertions(+), 48 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index 2cf6c6700336..394c1e0a6bcd 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -15,21 +15,6 @@ """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 sampling defaults per -embodiment. When ``domain_name`` (or a uniquely mapped ``domain_id``) is set, -the pipeline fills omitted ``action_chunk_size``, ``action_resolution``, and -``frame_rate`` from the preset and logs a warning if explicit values differ. -``num_frames`` is always ``action_chunk_size + 1``, never a preset field. See -Cosmos3 omni ``action_*.json`` inputs for reference configs (bridge, av, droid, -libero, etc.). - -``raw_action_dim`` is deliberately *not* a preset field: it is fixed by the -embodiment, while several embodiments share one preset via -``COSMOS3_DOMAIN_PRESET_ALIASES``. It resolves from -``action.EMBODIMENT_TO_RAW_ACTION_DIM`` instead. """ from typing import Any, Dict, Iterable, TypedDict @@ -164,7 +149,8 @@ class Cosmos3DomainPreset(TypedDict, total=False): frame_rate: float -# Training-aligned defaults. Values mirror Cosmos3 omni action JSON examples where available. +# Training-aligned defaults, mirroring the Cosmos3 omni ``action_*.json`` inputs +# (bridge, av, droid, libero, ...) where those exist. COSMOS3_DOMAIN_PRESETS: dict[str, Cosmos3DomainPreset] = { # WidowX bridge. "bridge_orig_lerobot": { @@ -292,7 +278,14 @@ def resolve_domain_action_config( frame_rate: float | None = None, action_fps: float | None = None, ) -> dict[str, Any]: - """Merge user action params with domain presets and generic fallbacks.""" + """Merge user action params with domain presets and generic fallbacks. + + A recognized ``domain_name`` (or a uniquely mapped ``domain_id``) fills + whichever of ``action_chunk_size``, ``action_resolution`` and ``frame_rate`` + the caller left unset; an explicit value wins but is reported in + ``warnings`` when it differs from the preset. ``num_frames`` is derived as + ``action_chunk_size + 1`` and is never a preset field. + """ 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] = [] 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 604120b51cb1..8eaabd6f5e64 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -878,36 +878,6 @@ def _preprocess_action_first_frame( ) return self._condition_frames_to_video_tensor(frames_u8).squeeze(2) - 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, @@ -956,7 +926,7 @@ def _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)), + action_dim=int(self.transformer.action_dim), generator=generator, device=self.device, dtype=self.dtype, From d9bd437225d0199109368b5406e1e4e250f99ed9 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 5 Aug 2026 15:47:46 -0700 Subject: [PATCH 19/35] [None][feat] Return only the action tensor, not the request that asked for it VisualGenOutput and PipelineOutput carried raw_action_dim, action_mode and domain_id alongside the action tensor. Nothing in the codebase read them back: they flowed pipeline -> output -> safetensors header and stopped. Each is recoverable without the worker sending it: - raw_action_dim is action.shape[-1]. One line produces both, since the tensor is sliced to that width on the way out. - action_mode is the caller's own string, normalized. - domain_id is an index into DomainAwareLinear's weight table, resolved from the domain_name the caller sent. It is also the lossier of the pair -- id 8 is droid_lerobot or robomind-franka, id 15 is any of three agibot variants -- so a consumer holding it cannot recover what it asked for, while the request's domain_name says exactly. Dropping them leaves the shared output schema one field wider than main, for the tensor itself, and takes tensor_payload back to media tensors plus rates with no model-specific keys (a test now pins that). The offline example writes its own sidecar from argparse instead, which also lets it record domain_name -- something the output never carried. Also comment why Wan's post_step_fn takes side-stream latents it does not use: the shared denoise loop passes them for Cosmos3, and the Wan diff otherwise shows an unexplained parameter. Signed-off-by: Igor Shovkun --- examples/visual_gen/models/cosmos3/cosmos3.py | 21 ++++++----- .../models/cosmos3/pipeline_cosmos3.py | 6 ++-- .../visual_gen/models/wan/pipeline_wan.py | 4 ++- tensorrt_llm/_torch/visual_gen/output.py | 27 +++----------- tensorrt_llm/media/tensor_payload.py | 23 ++++-------- tensorrt_llm/visual_gen/output.py | 6 ++-- .../visual_gen/test_cosmos3_pipeline.py | 6 ---- .../_torch/visual_gen/test_tensor_payload.py | 36 +++++-------------- tests/unittest/visual_gen/test_output.py | 13 ++----- 9 files changed, 43 insertions(+), 99 deletions(-) diff --git a/examples/visual_gen/models/cosmos3/cosmos3.py b/examples/visual_gen/models/cosmos3/cosmos3.py index 22fc33111cfb..f94afe3ed13d 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3.py +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -148,7 +148,12 @@ def _default_action_output_path(output_path: str) -> str: return str(stem.with_suffix(".action.json")) -def _save_action_output(output, path: str) -> None: +def _save_action_output(output, path: str, args: argparse.Namespace) -> None: + """Write the trajectory plus the request that produced it. + + The mode and embodiment are this script's own inputs, so they are read + from *args* rather than echoed back through the output schema. + """ if output.action is None: return @@ -161,9 +166,10 @@ def _save_action_output(output, path: str) -> None: shape = list(action.shape) payload = { - "action_mode": output.action_mode, - "domain_id": output.domain_id, - "raw_action_dim": output.raw_action_dim, + "action_mode": args.action_mode, + "domain_name": args.domain_name, + "domain_id": args.domain_id, + "raw_action_dim": action.shape[-1], "shape": shape, "dtype": str(action.dtype).replace("torch.", ""), "data": action_data, @@ -401,13 +407,10 @@ def main(): if args.action_mode is not None: action_path = args.action_output_path or _default_action_output_path(output_path) - _save_action_output(output, action_path) + _save_action_output(output, action_path, args) 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}" - ) + print(f"Action shape: {tuple(output.action.shape)}") else: print("Warning: action_mode was set but the output carried no action 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 8eaabd6f5e64..4c83b299c532 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -1724,11 +1724,11 @@ 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, + # Sliced to the embodiment's real width, so the trailing dim is + # raw_action_dim; the mode and embodiment are the caller's own + # request and are not echoed back. 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/wan/pipeline_wan.py b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py index 88f90b719d84..204d7e5a20cb 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py @@ -625,7 +625,9 @@ def forward_fn( encoder_hidden_states=encoder_hidden_states, ) - # Pin reference image to latent after each scheduler step (Wan 2.2 5B I2V only) + # Pin reference image to latent after each scheduler step (Wan 2.2 5B I2V only). + # post_step_fn also carries the denoise loop's side-stream latents (Cosmos3 + # denoises action/audio alongside video); Wan has none, so they pass through. def _pin_i2v_first_frame(x, extra_stream_latents): pinned = ((1 - i2v_first_frame_mask) * i2v_condition + i2v_first_frame_mask * x).to( self.dtype diff --git a/tensorrt_llm/_torch/visual_gen/output.py b/tensorrt_llm/_torch/visual_gen/output.py index ece7fd203f3f..76cf825cb4d9 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``, action fields) - 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)``, @@ -55,22 +54,15 @@ class PipelineOutput: 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. + ``inverse_dynamics``), already sliced to the embodiment's real + degrees of freedom, so ``D_raw`` states them; 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. @@ -89,9 +81,6 @@ class PipelineOutput: 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 @@ -222,9 +211,6 @@ def to_visual_gen_output(resp: "DiffusionResponse") -> "VisualGenOutput": 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, ) @@ -288,9 +274,6 @@ def split_visual_gen_output(resp: "DiffusionResponse", batch_size: int) -> List[ 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/media/tensor_payload.py b/tensorrt_llm/media/tensor_payload.py index f94ea52dce66..70fb0cdadb3b 100644 --- a/tensorrt_llm/media/tensor_payload.py +++ b/tensorrt_llm/media/tensor_payload.py @@ -5,13 +5,12 @@ Two payload formats are supported: - ``"safetensors"``: writes a single file with named tensors - (``image``/``video``/``audio``/``action``). Numeric scalar metadata - (``frame_rate``, ``audio_sample_rate``, ``raw_action_dim``, ``domain_id``) - is stored two ways: as a 0-d tensor under the same key (so - ``safetensors.torch.load(bytes)`` returns it alongside the media tensors — - consumers call ``.item()`` to unbox) and as a stringified value in the file - header (preserved for callers using ``safe_open(...).metadata()``). String - metadata such as ``action_mode`` is header-only. No pickle on load. + (``image``/``video``/``audio``/``action``). Scalar metadata + (``frame_rate``, ``audio_sample_rate``) is stored two ways: as a 0-d + tensor under the same key (so ``safetensors.torch.load(bytes)`` returns + it alongside the media tensors — consumers call ``.item()`` to unbox) + and as a stringified value in the file header (preserved for callers + using ``safe_open(...).metadata()``). No pickle on load. - ``"pt"``: writes a single file via :func:`torch.save` with the same tensor keys plus scalar metadata as native Python values. Clients should load with ``torch.load(buf, weights_only=True)`` @@ -146,12 +145,6 @@ def _collect_tensors_and_metadata( metadata["frame_rate"] = float(frame_rate) if audio_sample_rate is not None: metadata["audio_sample_rate"] = int(audio_sample_rate) - if output.raw_action_dim is not None: - metadata["raw_action_dim"] = int(output.raw_action_dim) - if output.domain_id is not None: - metadata["domain_id"] = int(output.domain_id) - if output.action_mode is not None: - metadata["action_mode"] = str(output.action_mode) return tensors, metadata @@ -211,9 +204,7 @@ def serialize_visual_gen_output( # ``loaded["frame_rate"].item()`` directly) and as a string in the # file header (preserved for callers that already use # ``safe_open(...).metadata()``). The two views always agree. - scalar_tensors = { - k: torch.as_tensor(v) for k, v in metadata.items() if isinstance(v, (int, float)) - } + scalar_tensors = {k: torch.as_tensor(v) for k, v in metadata.items()} return safetensors_save( {**tensors, **scalar_tensors}, metadata={k: str(v) for k, v in metadata.items()}, diff --git a/tensorrt_llm/visual_gen/output.py b/tensorrt_llm/visual_gen/output.py index 099d40aa3a4c..88ba6bc5082a 100644 --- a/tensorrt_llm/visual_gen/output.py +++ b/tensorrt_llm/visual_gen/output.py @@ -99,9 +99,6 @@ class VisualGenOutput: 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 @@ -173,7 +170,8 @@ def save( if self.action is not None: raise NotImplementedError( "Saving action outputs requires a tensor payload format " - "('safetensors' or 'pt') so the action tensor and metadata are preserved." + "('safetensors' or 'pt'); no video or image container can carry " + "an action trajectory." ) if self.image is not None: diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index 2e4b56df4008..c8fa4f442a1f 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -991,8 +991,6 @@ def test_policy_smoke(self, cosmos3_pipeline): 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) @@ -1021,8 +1019,6 @@ def test_forward_dynamics_smoke(self, cosmos3_pipeline): raw_action_dim=self.RAW_ACTION_DIM, chunk_size=self.ACTION_CHUNK, ) - assert result.action_mode == "forward_dynamics" - assert result.domain_id == 7 def test_inverse_dynamics_smoke(self, cosmos3_pipeline): _require_action_pipeline(cosmos3_pipeline) @@ -1051,8 +1047,6 @@ def test_inverse_dynamics_smoke(self, cosmos3_pipeline): raw_action_dim=self.RAW_ACTION_DIM, chunk_size=NUM_FRAMES - 1, ) - assert result.action_mode == "inverse_dynamics" - assert result.domain_id == 7 def test_inverse_dynamics_rejects_short_video(self, cosmos3_pipeline): _require_action_pipeline(cosmos3_pipeline) diff --git a/tests/unittest/_torch/visual_gen/test_tensor_payload.py b/tests/unittest/_torch/visual_gen/test_tensor_payload.py index 3c9168ee92cb..d76660263104 100644 --- a/tests/unittest/_torch/visual_gen/test_tensor_payload.py +++ b/tests/unittest/_torch/visual_gen/test_tensor_payload.py @@ -55,13 +55,7 @@ def _make_video_output(batch: int = 1, t: int = 2, h: int = 4, w: int = 4) -> Vi def _make_action_output(batch: int = 1, t: int = 4, action_dim: int = 7) -> VisualGenOutput: """Action uses ``(B, T, action_dim)``: batched at rank 3, unbatched at rank 2.""" action = torch.arange(batch * t * action_dim, dtype=torch.float32).reshape(batch, t, action_dim) - return VisualGenOutput( - request_id=3, - action=action, - action_mode="policy", - raw_action_dim=action_dim, - domain_id=7, - ) + return VisualGenOutput(request_id=3, action=action) class TestIsTensorFormat: @@ -154,26 +148,14 @@ def test_unbatched_action_passthrough(self, fmt): assert loaded["action"].shape == (4, 7) assert torch.equal(loaded["action"], output.action) - def test_action_metadata_serialized(self, fmt): - output = _make_action_output(batch=1) - data = serialize_visual_gen_output(output, fmt, batch_index=0) - loaded = self._load(data, fmt) - assert loaded["raw_action_dim"] == 7 - assert loaded["domain_id"] == 7 - if fmt == "pt": - assert loaded["action_mode"] == "policy" - else: - assert "action_mode" not in loaded - import tempfile - - from safetensors import safe_open - - with tempfile.NamedTemporaryFile(suffix=".safetensors") as tf: - tf.write(data) - tf.flush() - with safe_open(tf.name, framework="pt") as f: - meta = f.metadata() or {} - assert meta.get("action_mode") == "policy" + def test_action_carries_no_request_metadata(self, fmt): + """The trajectory's own shape states its DOF, and the mode and + embodiment are the caller's request. The payload stays model-agnostic: + media tensors plus rates, nothing Cosmos3-shaped.""" + output = _make_action_output(batch=1, t=4, action_dim=7) + loaded = self._load(serialize_visual_gen_output(output, fmt, batch_index=0), fmt) + assert loaded["action"].shape == (4, 7) + assert set(loaded) == {"action"} @pytest.mark.parametrize("fmt", ["safetensors", "pt"]) diff --git a/tests/unittest/visual_gen/test_output.py b/tests/unittest/visual_gen/test_output.py index 9cd5d53c2ee2..0f622f222679 100644 --- a/tests/unittest/visual_gen/test_output.py +++ b/tests/unittest/visual_gen/test_output.py @@ -35,9 +35,6 @@ def test_visual_gen_output_is_dataclass(): "action", "frame_rate", "audio_sample_rate", - "raw_action_dim", - "action_mode", - "domain_id", "error", "metrics", } @@ -65,9 +62,6 @@ def test_minimal_construction_defaults(): assert out.video is None assert out.audio 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.frame_rate is None assert out.audio_sample_rate is None assert out.error is None @@ -773,8 +767,8 @@ def test_encoding_not_top_level_reexport(): # --------------------------------------------------------------------------- -def test_pipeline_output_has_twelve_fields(): - """PipelineOutput has the twelve expected fields.""" +def test_pipeline_output_has_nine_fields(): + """PipelineOutput has the nine expected fields.""" field_names = {f.name for f in fields(PipelineOutput)} assert field_names == { "image", @@ -783,9 +777,6 @@ def test_pipeline_output_has_twelve_fields(): "action", "frame_rate", "audio_sample_rate", - "raw_action_dim", - "action_mode", - "domain_id", "pre_denoise", "denoise", "post_denoise", From a98e031fcb8b4ebab2ab781a2ff1262a627f3695 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 5 Aug 2026 16:06:29 -0700 Subject: [PATCH 20/35] [None][doc] Cosmos3 README: keep the original title The action modes are listed in the mode bullets; the heading does not need to enumerate them. Signed-off-by: Igor Shovkun --- examples/visual_gen/models/cosmos3/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/visual_gen/models/cosmos3/README.md b/examples/visual_gen/models/cosmos3/README.md index 9086ea1ac152..17dd4ca3229c 100644 --- a/examples/visual_gen/models/cosmos3/README.md +++ b/examples/visual_gen/models/cosmos3/README.md @@ -1,4 +1,4 @@ -# Cosmos3 Text(+Image)-to-Video(+Audio) and action generation +# Cosmos3 Text(+Image)-to-Video(+Audio) generation Cosmos3 supports the following generation modes from a single checkpoint: From 90ce4d17c27e2e71d0c9525095cded71231e9b33 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 5 Aug 2026 21:27:51 -0700 Subject: [PATCH 21/35] [None][fix] Cosmos3 action: restore the checkpoint sigmas, reject empty trajectories set_flow_shift returns the scheduler untouched when both knobs are None, and the action branch passed an unset flow_shift straight through. The pipeline instance outlives a request, so an action request following a T2I or V2V one denoised on that mode's schedule -- V2V forces use_karras_sigmas=False, which swaps the checkpoint's Karras sigmas for uniform ones and is not inert on any checkpoint. The video branch already restores checkpoint_flow_shift when unset; action now does the same. load_action_tensor accepted a zero-row trajectory. The empty slice survived padding and only failed at the mask broadcast, as an opaque torch shape error rather than a client-facing one -- a 500 where a 400 belongs. Also from review: the inverse_dynamics error still offered a frame directory or an image, both retired when the clip moved to encoded bytes; --action_output_path documented a default the code never produced; and the forward-dynamics README example did not say a trajectory's width must match the embodiment (9 for av). Tests: the bucket table is now asserted against the full aspect-label set rather than merely non-empty, and prepare_action_latents gains happy-path coverage -- padding by holding the last step, truncation, the zeroing above raw_action_dim, and the per-mode mask semantics that post_step_fn re-anchors against. Signed-off-by: Igor Shovkun --- examples/visual_gen/models/cosmos3/README.md | 2 + examples/visual_gen/models/cosmos3/cosmos3.py | 6 +- .../visual_gen/models/cosmos3/action.py | 5 ++ .../models/cosmos3/pipeline_cosmos3.py | 9 ++- .../_torch/visual_gen/test_cosmos3_action.py | 80 ++++++++++++++++++- .../visual_gen/test_cosmos3_pipeline.py | 55 +++++++++++++ 6 files changed, 151 insertions(+), 6 deletions(-) diff --git a/examples/visual_gen/models/cosmos3/README.md b/examples/visual_gen/models/cosmos3/README.md index 17dd4ca3229c..3453b7a245eb 100644 --- a/examples/visual_gen/models/cosmos3/README.md +++ b/examples/visual_gen/models/cosmos3/README.md @@ -121,6 +121,8 @@ python cosmos3.py --model nvidia/Cosmos3-Nano \ --action_output_path policy_action.json # Action — forward dynamics (first frame + action trajectory -> rollout video) +# action_trajectory.json is a [T, D] list of lists; D is the embodiment's action +# width (9 for av) and a mismatch is rejected. python cosmos3.py --model nvidia/Cosmos3-Nano \ --prompt_file prompts/action_forward_dynamics.json \ --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml \ diff --git a/examples/visual_gen/models/cosmos3/cosmos3.py b/examples/visual_gen/models/cosmos3/cosmos3.py index f94afe3ed13d..0c0c4e473e88 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3.py +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -127,9 +127,7 @@ def _validate_action_args( raise SystemExit(f"{mode} requires --raw_action_dim, --domain_name, or --domain_id.") elif mode == "inverse_dynamics": if args.video_path is None: - raise SystemExit( - f"{mode} requires --video_path (frame directory, .mp4/.avi, or image)." - ) + raise SystemExit(f"{mode} requires --video_path (an .mp4 or .avi file).") if args.raw_action_dim is None and args.domain_name is None and args.domain_id is None: raise SystemExit(f"{mode} requires --raw_action_dim, --domain_name, or --domain_id.") @@ -316,7 +314,7 @@ def main(): "--action_output_path", type=str, default=None, - help="Path to save predicted action JSON (default: _action.json)", + help="Path to save predicted action JSON (default: .action.json)", ) parser.add_argument( "--output_type", type=str, default="video", help="Output type (video, image)" diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py index 05463cde9d04..224c7474dbf0 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py @@ -363,6 +363,11 @@ def load_action_tensor(action: Any = None) -> torch.Tensor: tensor = tensor.squeeze(0) if tensor.ndim != 2: raise ValueError(f"Cosmos3 action must have shape [T, D], got {tuple(tensor.shape)}.") + if tensor.shape[0] == 0: + raise ValueError( + f"Cosmos3 action trajectory must have at least one timestep, got shape " + f"{tuple(tensor.shape)}." + ) return 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 4c83b299c532..2438c11db4c9 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -1103,7 +1103,14 @@ def forward( ) if guidance_scale is None: guidance_scale = COSMOS3_ACTION_PARAMS["guidance_scale"] - self._apply_flow_shift(flow_shift) + # Restore the checkpoint knobs when unset, like the video branch + # below: the pipeline instance outlives the request, and a prior + # T2I or V2V request left its own shift and sigma schedule behind + # (V2V forces uniform sigmas, which is not inert on any checkpoint). + self._apply_flow_shift( + flow_shift if flow_shift is not None else self.sampling.checkpoint_flow_shift, + use_karras_sigmas=None, + ) enable_audio = False else: height = height or COSMOS3_720P_PARAMS["height"] diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py index a42992705d61..1f62a0aa55f8 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py @@ -14,6 +14,7 @@ import torch from tensorrt_llm._torch.visual_gen.models.cosmos3.action import ( + ACTION_ASPECT_RATIO_LABELS, ACTION_VIEWPOINT_TEMPLATES, DEFAULT_ACTION_VIEW_POINT, EMBODIMENT_TO_DOMAIN_ID, @@ -71,7 +72,9 @@ def test_unknown_resolution_raises(self): @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] + # find_closest_target_size picks from whatever entries exist, so a bucket + # that lost "9,16" would silently land portrait sources on another canvas. + assert set(VIDEO_RES_SIZE_INFO[action_resolution]) == set(ACTION_ASPECT_RATIO_LABELS) class TestResolveActionSize: @@ -617,6 +620,81 @@ def test_audio_style_call_unchanged(self): class TestPrepareActionLatents: + """The CPU half of the action contract: which tokens start clean, what the + mask says, and how a caller's trajectory is fitted to the chunk.""" + + ACTION_DIM = 8 + CHUNK = 4 + + def _prepare(self, mode, **kwargs): + kwargs.setdefault("action_chunk_size", self.CHUNK) + kwargs.setdefault("action_dim", self.ACTION_DIM) + return prepare_action_latents( + mode=mode, + generator=torch.Generator(device="cpu").manual_seed(0), + device=torch.device("cpu"), + dtype=torch.float32, + **kwargs, + ) + + def test_forward_dynamics_pads_a_short_trajectory_by_holding_the_last_step(self): + latents, mask, clean, raw_dim = self._prepare( + "forward_dynamics", raw_action_dim=None, action_input=[[1.0, 2.0], [3.0, 4.0]] + ) + assert raw_dim == 2 + assert latents.shape == (1, self.CHUNK, self.ACTION_DIM) + # Steps 2 and 3 repeat the supplied final step rather than going to zero. + for step in range(2, self.CHUNK): + torch.testing.assert_close(clean[0, step, :2], torch.tensor([3.0, 4.0])) + + def test_forward_dynamics_truncates_a_long_trajectory(self): + _, _, clean, _ = self._prepare( + "forward_dynamics", + raw_action_dim=None, + action_input=[[float(i), float(i)] for i in range(self.CHUNK + 3)], + ) + torch.testing.assert_close(clean[0, -1, :2], torch.tensor([3.0, 3.0])) + + def test_forward_dynamics_conditions_every_step(self): + """All action tokens are given, so none carries velocity.""" + latents, mask, clean, _ = self._prepare( + "forward_dynamics", raw_action_dim=None, action_input=[[1.0, 2.0]] * self.CHUNK + ) + assert torch.all(mask == 0.0) + torch.testing.assert_close(latents, clean) + + @pytest.mark.parametrize("mode", ["policy", "inverse_dynamics"]) + def test_predicted_modes_start_from_noise_everywhere(self, mode): + latents, mask, clean, _ = self._prepare(mode, raw_action_dim=2) + assert torch.all(mask == 1.0) + assert torch.all(clean == 0.0) + assert torch.any(latents[..., :2] != 0.0) + + @pytest.mark.parametrize("mode", ["policy", "forward_dynamics", "inverse_dynamics"]) + def test_columns_above_raw_action_dim_are_zero(self, mode): + """The head is action_dim wide but only raw_action_dim is meaningful; + padding must not carry noise into the model or out of it.""" + kwargs = ( + {"raw_action_dim": None, "action_input": [[1.0, 2.0]] * self.CHUNK} + if mode == "forward_dynamics" + else {"raw_action_dim": 2} + ) + latents, _, clean, raw_dim = self._prepare(mode, **kwargs) + assert raw_dim == 2 + assert torch.all(latents[..., raw_dim:] == 0.0) + assert torch.all(clean[..., raw_dim:] == 0.0) + + def test_empty_trajectory_raises(self): + """Without this the empty slice reaches the mask broadcast and dies with + an opaque torch shape error instead of a client-facing one.""" + with pytest.raises(ValueError, match="at least one timestep"): + self._prepare("forward_dynamics", raw_action_dim=None, action_input=torch.zeros(0, 2)) + + @pytest.mark.parametrize("raw_action_dim", [0, -1, ACTION_DIM + 1]) + def test_out_of_range_raw_action_dim_raises(self, raw_action_dim): + with pytest.raises(ValueError, match=r"raw_action_dim must be in \[1, \d+\]"): + self._prepare("policy", raw_action_dim=raw_action_dim) + def test_forward_dynamics_raw_dim_mismatch_raises(self): with pytest.raises(ValueError, match="raw_action_dim must match"): prepare_action_latents( diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index c8fa4f442a1f..8f7f08157b15 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -858,6 +858,61 @@ def fake_tokenize_prompt(text, max_sequence_length, use_system_prompt, system_pr assert token_calls[0] is False + def test_action_restores_the_checkpoint_flow_shift(self): + """set_flow_shift is a no-op when both knobs are None, and the pipeline + instance outlives the request. An action request that followed a V2V one + would otherwise keep V2V's uniform sigmas instead of the checkpoint's.""" + pipeline = Cosmos3OmniMoTPipeline.__new__(Cosmos3OmniMoTPipeline) + pipeline.transformer = SimpleNamespace( + device=torch.device("cpu"), num_embodiment_domains=32 + ) + pipeline.audio_gen = False + pipeline.action_gen = True + pipeline.default_use_system_prompt = False + calls = [] + + class StopAfterTokenize(Exception): + pass + + class FakeSampling: + is_distilled = False + checkpoint_flow_shift = 7.0 + + 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 + + def fake_tokenize_prompt(text, max_sequence_length, use_system_prompt, system_prompt=None): + raise StopAfterTokenize + + pipeline.scheduler = SimpleNamespace(config=SimpleNamespace(flow_shift=1.0)) + pipeline.sampling = FakeSampling() + pipeline._tokenize_prompt = fake_tokenize_prompt + + with pytest.raises(StopAfterTokenize): + pipeline.forward( + prompt="pick up the block", + video=_V2V_FIXTURE_MP4.read_bytes(), + num_frames=NUM_FRAMES, + num_inference_steps=1, + guidance_scale=1.0, + seed=1, + max_sequence_length=8, + use_guardrails=False, + action_mode="inverse_dynamics", + domain_name="bridge_orig_lerobot", + raw_action_dim=10, + action_chunk_size=NUM_FRAMES - 1, + ) + + assert calls == [(7.0, None)] + def test_apply_flow_shift_rebuilds_every_stream_scheduler(self): """Action denoises in the same loop as video, on its own scheduler instance, so a request that shifts the schedule must move all of them.""" From 1c252f5b462fe0731c534eb3b817637b87a55f3f Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 5 Aug 2026 21:37:06 -0700 Subject: [PATCH 22/35] [None][chore] Drop an unused mask binding in the action latent test The padding test never reads the velocity mask; the per-mode mask semantics are asserted in their own cases. Signed-off-by: Igor Shovkun --- tests/unittest/_torch/visual_gen/test_cosmos3_action.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py index 1f62a0aa55f8..f398ecdf0e81 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py @@ -638,7 +638,7 @@ def _prepare(self, mode, **kwargs): ) def test_forward_dynamics_pads_a_short_trajectory_by_holding_the_last_step(self): - latents, mask, clean, raw_dim = self._prepare( + latents, _, clean, raw_dim = self._prepare( "forward_dynamics", raw_action_dim=None, action_input=[[1.0, 2.0], [3.0, 4.0]] ) assert raw_dim == 2 From 5ef4ef1a64213cc57b865288e94ab891ba01008e Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 6 Aug 2026 09:58:43 -0700 Subject: [PATCH 23/35] [None][chore] Drop a dead range on a literal-typed action param Validation stops at the literal-membership check, so action_resolution's range=(256, 720) never ran. It was harmless -- the bucket set is a strict subset of the interval -- but it read as active validation. Remove it and say in params.py why the literal branch is terminal. Signed-off-by: Igor Shovkun --- tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py | 3 ++- tensorrt_llm/visual_gen/params.py | 4 ++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index 394c1e0a6bcd..1914c3f92fca 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -507,7 +507,8 @@ def _resolve_field( "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)), + # No range: the buckets are not an interval, and validation stops at the + # literal check anyway. ), "view_point": ExtraParamSchema( type="Literal['ego_view', 'third_person_view', 'wrist_view', 'concat_view']", diff --git a/tensorrt_llm/visual_gen/params.py b/tensorrt_llm/visual_gen/params.py index 9b27163db9fc..989ff018e1ab 100644 --- a/tensorrt_llm/visual_gen/params.py +++ b/tensorrt_llm/visual_gen/params.py @@ -186,6 +186,10 @@ def validate_visual_gen_params( f"extra_params['{key}'] expected one of {list(literal_choices)}, " f"got {value!r}" ) + # Terminal on purpose: membership in the literal set already + # decides the value, so the type, validator and range checks + # below cannot add anything. A literal spec that also declares + # one of those has a redundant declaration, not a skipped check. continue # Type check expected_types = _TYPE_MAP.get(spec.type) From e5a67b84058f58addfff70f8f5dc1726f0415262 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 6 Aug 2026 10:58:46 -0700 Subject: [PATCH 24/35] [None][fix] Check the action latent dtype, cover the action resize helpers DomainAwareLinear builds its projection weights at the transformer's declared dtype while the action latents carry the pipeline's, and the action branch -- unlike audio -- never reconciled them, so a mismatch surfaced as a bmm dtype error from inside the projection. Raise a named error instead of casting: a per-step conversion of a whole stream would hide the misconfiguration that produced the mismatch rather than report it. Both action resize helpers were untested. They implement the deliberate divergence from V2V -- contain-scale and pad, where video covers and crops, because a gripper works at the frame edge and cropping removes what the policy acts on. Cover containment, the never-enlarge floor, identity at native size, aspect preservation, and the replicate-instead-of-reflect fallback for a pad run wider than the resized extent. Signed-off-by: Igor Shovkun --- .../models/cosmos3/transformer_cosmos3.py | 9 ++++ .../_torch/visual_gen/test_cosmos3_action.py | 34 ++++++++++++++ .../_torch/visual_gen/test_media_decode.py | 44 +++++++++++++++++++ 3 files changed, 87 insertions(+) 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 248006ec149e..76aea8503c39 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -1307,6 +1307,15 @@ def forward( ) self.domain_ids_validated = True T_action = action_latents.shape[1] + # Checked, not cast: bmm in DomainAwareLinear needs the latents to + # match the projection weights, and silently converting a whole + # stream every step would hide the misconfiguration that produced + # the mismatch. + if action_latents.dtype != self.action_proj_in.dtype: + raise ValueError( + "Cosmos3 action latents must match the action projection dtype: " + f"latents={action_latents.dtype}, projection={self.action_proj_in.dtype}." + ) hidden_action = self.action_proj_in( self.pack_action(action_latents), action_domain_ids_tensor ) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py index f398ecdf0e81..1377db1dddf5 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py @@ -9,6 +9,7 @@ import json +import numpy as np import PIL.Image import pytest import torch @@ -26,6 +27,7 @@ find_closest_target_size, normalize_action_resolution, prepare_action_latents, + resize_and_pad_action_image, resolve_action_size, resolve_domain_id, resolve_raw_action_dim, @@ -709,6 +711,38 @@ def test_forward_dynamics_raw_dim_mismatch_raises(self): ) +class TestResizeAndPadActionImage: + """Action pads to the canvas where V2V crops to it: a gripper works at the + frame edge, so cover-scale would cut away what the policy acts on.""" + + def test_contain_scale_then_pad_to_canvas(self): + image = PIL.Image.new("RGB", (800, 400), "blue") + out = resize_and_pad_action_image(image, target_h=480, target_w=832) + assert (out.height, out.width) == (480, 832) + + def test_small_source_is_never_enlarged(self): + """min(..., 1.0): a small clip keeps its own pixels and a wider border + rather than being upscaled into blur.""" + image = PIL.Image.new("RGB", (100, 50), "blue") + out = resize_and_pad_action_image(image, target_h=480, target_w=832) + assert (out.height, out.width) == (480, 832) + assert np.asarray(out)[:50, :100].any() + + def test_exact_size_is_returned_unchanged(self): + image = PIL.Image.new("RGB", (832, 480), "blue") + assert resize_and_pad_action_image(image, 480, 832).size == (832, 480) + + def test_aspect_ratio_is_preserved(self): + """Contain-scale keeps the source's own aspect; the leftover strip is + padding, not stretch.""" + image = PIL.Image.new("RGB", (400, 400), "blue") + out = np.asarray(resize_and_pad_action_image(image, target_h=480, target_w=832)) + assert out.shape[:2] == (480, 832) + # A square source contained in a 16:9 canvas fills the height, so the + # scaled content is square and the pad lands on the right. + assert out[:480, :480].any() + + class TestResolveDomainId: """domain_id wins, but a caller that contradicts itself is a real mistake: the wrong embodiment yields a fluent trajectory in another robot's dialect.""" diff --git a/tests/unittest/_torch/visual_gen/test_media_decode.py b/tests/unittest/_torch/visual_gen/test_media_decode.py index 7fc9f3a84d6f..d92ba8a6863f 100644 --- a/tests/unittest/_torch/visual_gen/test_media_decode.py +++ b/tests/unittest/_torch/visual_gen/test_media_decode.py @@ -30,6 +30,7 @@ _lanczos_taps, decode_video_reference_window, resize_center_crop_uint8, + resize_fit_pad_uint8, ) _TEST_DATA = Path(__file__).parent / "test_data" @@ -42,6 +43,49 @@ def _frame_indices(frames: torch.Tensor) -> list[int]: return [round((f[:, :, 0].float().mean().item() - 20) / 25) for f in frames] +class TestResizeFitPad: + """The action counterpart to cover-scale + center-crop: contain-scale, then + pad. A gripper works at the frame edge, so cropping costs the policy the + evidence it acts on.""" + + def test_contains_the_whole_source(self): + """Cover-scale would crop the wide source; fit keeps all of it.""" + frames = torch.full((2, 100, 400, 3), 200, dtype=torch.uint8) + out = resize_fit_pad_uint8(frames, target_h=200, target_w=400) + assert out.shape == (2, 200, 400, 3) + # 400x100 contained in 400x200 scales by 1.0 and pads the bottom half. + assert (out[:, :100] == 200).all() + + def test_never_enlarges_a_small_source(self): + """min(..., 1.0): a small clip keeps its own pixels and a wider border + rather than being upscaled.""" + frames = torch.full((1, 10, 20, 3), 255, dtype=torch.uint8) + out = resize_fit_pad_uint8(frames, target_h=64, target_w=64) + assert out.shape == (1, 64, 64, 3) + assert (out[0, :10, :20] == 255).all() + + def test_native_resolution_is_identity(self): + frames = torch.zeros(2, 32, 32, 3, dtype=torch.uint8) + assert resize_fit_pad_uint8(frames, 32, 32) is frames + + def test_aspect_ratio_is_preserved(self): + """The source is contained, not stretched: a square stays square.""" + frames = torch.zeros(1, 64, 64, 3, dtype=torch.uint8) + frames[:, :, :, 0] = 255 + out = resize_fit_pad_uint8(frames, target_h=64, target_w=128) + assert out.shape == (1, 64, 128, 3) + # A 1:1 source in a 2:1 canvas fills the height, so content is 64 wide. + assert (out[0, :, :64, 0] == 255).all() + + def test_pad_falls_back_to_replicate_when_reflection_has_no_source(self): + """A pad run wider than the resized extent has nothing left to mirror; + reflect would raise, so the helper switches to edge replication.""" + frames = torch.full((1, 4, 4, 3), 128, dtype=torch.uint8) + out = resize_fit_pad_uint8(frames, target_h=64, target_w=64) + assert out.shape == (1, 64, 64, 3) + assert (out[0, 4:, :4] == 128).all() + + class TestResizeCenterCrop: """CPU-runnable checks of the shared Lanczos resize/crop front.""" From 77b9be28bc78bc4508377184f662f1bdbf576101 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 6 Aug 2026 15:10:13 -0700 Subject: [PATCH 25/35] [None][perf] Build the Cosmos3 side-stream rope tables once per request The action and audio rotary tables were rebuilt on every denoise step, though every input that determines them -- chunk size, prompt lengths, frame rate, the action clock and the start offset -- is fixed once a request begins. Each rebuild read the prompt lengths back to the host (a device-to-host sync per batch element, so two per step under CFG), constructed the position ids on the CPU, copied them back, evaluated the rotary embedding and concatenated twice. For a 30-step request that is ~30 rebuilds of identical numbers. Cache the combined video+side-stream table on the same per-request lifecycle the video table already uses, cleared by reset_cache(). The scalars that select those positions -- fps, action_fps and action_start_frame_offset -- are Python values that leave every tensor shape unchanged, so the shape-derived CUDA graph key could not tell two such requests apart and would replay one's graph for the other. Register them as extra key functions, following the img_shapes precedent in the Qwen-Image transformer. Absent scalars drop out of the key, so a video-only request keys as before. Prompt length still is not in the key: text_mask keeps a constant padded shape while its real length shifts every table through the text offset. Reading it in a key function would reintroduce the per-dispatch sync removed above, so it needs the host-side length the pipeline already holds. That is a shared-signature change affecting T2V/I2V/V2V equally, and cached_kv has the same property. Signed-off-by: Igor Shovkun --- .../models/cosmos3/transformer_cosmos3.py | 98 ++++++++++++++----- .../visual_gen/test_cosmos3_transformer.py | 67 +++++++++++++ 2 files changed, 139 insertions(+), 26 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 76aea8503c39..432b9b18af86 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -936,6 +936,7 @@ def __init__(self, model_config: DiffusionModelConfig): self.cached_kv = None self.cached_freqs_gen = None + self.cached_freqs_gen_combined = None self.domain_ids_validated = False self.__post_init__() @@ -1162,9 +1163,43 @@ def pack_action(self, action_latents: torch.Tensor) -> torch.Tensor: def unpack_action(tokens: torch.Tensor) -> torch.Tensor: return tokens + def register_cuda_graph_extra_key_fns(self, runner) -> None: + """Make the position-determining scalars part of the graph key. + + The base key is tensor shapes only, but the rotary tables are built + from Python scalars that leave every shape unchanged: the frame rate, + the action clock, and the offset of the first action step. Two requests + differing only in these produce different positions at identical + shapes, so without them one captured graph would be replayed for both. + Each returns ``None`` when absent, which drops that part of the key -- + a video-only request keys exactly as it did before. + """ + super().register_cuda_graph_extra_key_fns(runner) + + def _float_key(name): + def fn(*args, **kwargs): + value = kwargs.get(name) + return None if value is None else float(value) + + return fn + + def _int_key(name): + def fn(*args, **kwargs): + value = kwargs.get(name) + return None if value is None else int(value) + + return fn + + runner.register_extra_key_fn("fps", _float_key("fps")) + runner.register_extra_key_fn("action_fps", _float_key("action_fps")) + runner.register_extra_key_fn( + "action_start_frame_offset", _int_key("action_start_frame_offset") + ) + def reset_cache(self): self.cached_kv = None self.cached_freqs_gen = None + self.cached_freqs_gen_combined = None self.domain_ids_validated = False def forward( @@ -1327,38 +1362,49 @@ def forward( 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), - ) + # The rotary table is request-invariant: chunk size, prompt lengths, + # fps and the frame offset are all fixed once the request starts. + # Recomputing it per step costs a device-to-host sync per batch + # element (the prompt-length readback), an H2D copy of the position + # ids, and two concatenations -- all for the same numbers. + if self.cached_freqs_gen_combined is None: + 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 + self.cached_freqs_gen_combined = ( + torch.cat([cos_v, cos_a], dim=1), + torch.cat([sin_v, sin_a], dim=1), + ) + freqs_gen_combined = self.cached_freqs_gen_combined 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 hidden_audio = hidden_audio + time_embed.unsqueeze(1) - cos_a, sin_a = self._compute_audio_rope_freqs( - T_audio, - text_mask, - float(self.audio_latent_fps), - hidden_states.device, - hidden_gen.dtype, - ) hidden_gen = torch.cat([hidden_gen, hidden_audio], dim=1) - 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), - ) + if self.cached_freqs_gen_combined is None: + cos_a, sin_a = self._compute_audio_rope_freqs( + T_audio, + text_mask, + float(self.audio_latent_fps), + hidden_states.device, + hidden_gen.dtype, + ) + cos_v, sin_v = self.cached_freqs_gen + self.cached_freqs_gen_combined = ( + torch.cat([cos_v, cos_a], dim=1), + torch.cat([sin_v, sin_a], dim=1), + ) + freqs_gen_combined = self.cached_freqs_gen_combined else: freqs_gen_combined = self.cached_freqs_gen # -------------------------------------------------------------------------- diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py index af7eb8922f06..1d726841dbce 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py @@ -521,6 +521,73 @@ def run_step(): run_step() assert len(calls) == 2 + def test_graph_key_separates_requests_that_differ_only_in_scalars(self, action_model_config): + """TRT-LLM captures a family of graphs and dispatches by key. fps, the + action clock and the start offset change the rotary positions without + changing any tensor shape, so they must discriminate keys or two such + requests would replay the same graph.""" + from tensorrt_llm._torch.visual_gen.cuda_graph_runner import ( + CUDAGraphRunner, + CUDAGraphRunnerConfig, + ) + + model = Cosmos3VFMTransformer(model_config=action_model_config) + runner = CUDAGraphRunner(CUDAGraphRunnerConfig(use_cuda_graph=True)) + model.register_cuda_graph_extra_key_fns(runner) + + base = dict(fps=24.0, action_fps=5.0, action_start_frame_offset=1) + key = runner.get_graph_key(**base) + for field, other in ( + ("fps", 16.0), + ("action_fps", 10.0), + ("action_start_frame_offset", 0), + ): + assert runner.get_graph_key(**{**base, field: other}) != key, field + + # A video-only request keys exactly as before: absent scalars drop out. + assert runner.get_graph_key(fps=None, action_fps=None) == runner.get_graph_key() + + @pytest.mark.high_cuda_memory + def test_action_rope_table_built_once_per_request(self, action_model_config): + """Chunk size, prompt lengths, fps and the frame offset are fixed for a + request, so the rotary table is too. Rebuilding it per step costs a + device-to-host sync per batch element plus an H2D copy of the position + ids -- for identical numbers.""" + 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) + + calls = [] + real = model._compute_action_rope_freqs + model._compute_action_rope_freqs = lambda *a, **k: (calls.append(1), real(*a, **k))[1] + + def run_step(): + with torch.inference_mode(): + 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, + ) + + run_step() + run_step() + run_step() + assert len(calls) == 1 + + model.reset_cache() + run_step() + assert len(calls) == 2 + @pytest.mark.high_cuda_memory def test_forward_with_action_domain_id_out_of_range_raises(self, action_model_config): cfg = action_model_config.pretrained_config From 8fa221f58580c863bc787f6e307ad6f8a7092730 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 6 Aug 2026 22:53:04 -0700 Subject: [PATCH 26/35] [None][fix] Resolve Cosmos3 action defaults by provenance, not by value An action request carries three resolution tiers -- caller, embodiment preset, video default -- where every other model has two. Telling them apart needs to know whether a value was chosen or defaulted, and comparing against the materialized default cannot: 24.0 is both the video default and a legal caller choice, so an explicit 24 fps on a bridge request was read as unset and silently replaced by the preset's 5, against resolve_domain_action_config()'s explicit-wins contract. default_params now discards the fields it filled from default_generation_params out of model_fields_set, leaving the values in place while reserving the set bit for real caller choices; assignment re-marks a field, so the serve layer and offline callers both record what they touched, and the bit survives the pickle hop to the worker. This is the mechanism cosmos3_edge already uses, adopted in the same form so the two do not diverge. infer() reads that instead of comparing values, for all five mode-dependent fields rather than frame_rate alone -- the others carry the same hazard latent and are saved only by their defaults being None today. Distilled checkpoints are unaffected: their steps and guidance come from generation_default_overrides() and are re-applied in forward() ahead of the action recipe's own fallbacks. Signed-off-by: Igor Shovkun --- .../models/cosmos3/pipeline_cosmos3.py | 29 +++++++--------- tensorrt_llm/visual_gen/visual_gen.py | 8 ++++- .../visual_gen/test_cosmos3_distilled.py | 33 ++++++++++++++++++- 3 files changed, 50 insertions(+), 20 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 2438c11db4c9..8f0f934d0fe1 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -347,33 +347,26 @@ def infer(self, req): def resolved(value, field_name): return value if value is not None else mode_params[field_name] - # Action resolves every one of these itself, from the embodiment preset - # and COSMOS3_ACTION_PARAMS: the canvas from the resolution bucket, the - # frame rate from the robot, and steps/guidance from the action recipe. - # Filling them from the video table here would leave forward()'s - # fallbacks dead and silently run action at 720p, 35 steps, guidance 6 - # (CFG on, double the work) and 24 fps. + specified = req.params.model_fields_set + + def as_given(field_name): + value = getattr(req.params, field_name) + return value if field_name in specified else None + is_action = extra_params.get("action_mode") is not None - height = req.params.height if is_action else resolved(req.params.height, "height") - width = req.params.width if is_action else resolved(req.params.width, "width") + height = as_given("height") if is_action else resolved(req.params.height, "height") + width = as_given("width") if is_action else resolved(req.params.width, "width") num_inference_steps = ( - req.params.num_inference_steps + as_given("num_inference_steps") if is_action else resolved(req.params.num_inference_steps, "num_inference_steps") ) guidance_scale = ( - req.params.guidance_scale + as_given("guidance_scale") if is_action else resolved(req.params.guidance_scale, "guidance_scale") ) - # frame_rate cannot stay None in the pipeline defaults the way the four - # above do: the serve layer derives num_frames from seconds x frame_rate - # before the pipeline sees the request. So for action, "unset" survives - # only as "still equal to what the executor materialized" — any other - # value is the caller's own and outranks the embodiment preset. - frame_rate = req.params.frame_rate - if is_action and frame_rate == self.default_generation_params.get("frame_rate"): - frame_rate = None + frame_rate = as_given("frame_rate") if is_action else req.params.frame_rate video = extra_params.get("video") # encoded MP4/AVI bytes (the extra-param contract) return self.forward( diff --git a/tensorrt_llm/visual_gen/visual_gen.py b/tensorrt_llm/visual_gen/visual_gen.py index d597bc88b238..95c57a9737e8 100644 --- a/tensorrt_llm/visual_gen/visual_gen.py +++ b/tensorrt_llm/visual_gen/visual_gen.py @@ -321,7 +321,13 @@ def default_params(self) -> "VisualGenParams": if extra: kwargs["extra_params"] = extra - return VisualGenParams(**kwargs) + params = VisualGenParams(**kwargs) + # These came from the pipeline, not the caller. Un-marking them keeps + # request-dependent defaults resolvable after a round trip through + # this object; assigning any of them re-marks it automatically. + for field_name in self.executor.default_generation_params: + params.model_fields_set.discard(field_name) + return params @set_api_status("prototype") def generate( diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py index 950400bee46e..cb0480a153ea 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py @@ -93,7 +93,11 @@ def _bare_pipeline(**attrs) -> Cosmos3OmniMoTPipeline: def _fake_request(output_type: str = "video", **param_overrides) -> SimpleNamespace: - """A DiffusionRequest look-alike with executor-merged (None = unset) params.""" + """A DiffusionRequest look-alike with executor-merged params. + + ``model_fields_set`` mirrors the real object: pipeline defaults are cleared + by ``VisualGen.default_params``, so only what this caller overrode is marked. + """ params = SimpleNamespace( height=None, width=None, @@ -109,6 +113,7 @@ def _fake_request(output_type: str = "video", **param_overrides) -> SimpleNamesp ) for key, value in param_overrides.items(): setattr(params, key, value) + params.model_fields_set = set(param_overrides) return SimpleNamespace(prompt="x", params=params) @@ -411,6 +416,32 @@ def test_action_keeps_an_explicit_frame_rate(self): got = self._captured_forward_kwargs(_bare_pipeline(), req) assert got["frame_rate"] == 30.0 + def test_action_honors_an_explicit_frame_rate_equal_to_the_video_default(self): + """The collision case: 24.0 is both a legal caller choice and the + materialized video default. Value equality reads it as unset and hands + back the embodiment preset; provenance keeps the caller's 24.""" + req = _fake_request("video", frame_rate=24.0, extra_params={"action_mode": "policy"}) + got = self._captured_forward_kwargs(_bare_pipeline(), req) + assert got["frame_rate"] == 24.0 + + def test_action_drops_a_frame_rate_the_caller_never_set(self): + req = _fake_request("video", extra_params={"action_mode": "policy"}) + got = self._captured_forward_kwargs(_bare_pipeline(), req) + assert got["frame_rate"] is None + + def test_action_honors_explicit_values_equal_to_their_defaults(self): + """Same hazard for the other mode-dependent fields.""" + req = _fake_request( + "video", + height=COSMOS3_720P_PARAMS["height"], + guidance_scale=COSMOS3_720P_PARAMS["guidance_scale"], + extra_params={"action_mode": "policy"}, + ) + got = self._captured_forward_kwargs(_bare_pipeline(), req) + assert got["height"] == COSMOS3_720P_PARAMS["height"] + assert got["guidance_scale"] == COSMOS3_720P_PARAMS["guidance_scale"] + assert got["width"] is None + def test_distilled_merged_defaults_pass_through(self): req = _fake_request("image", num_inference_steps=4, guidance_scale=1.0) got = self._captured_forward_kwargs(_bare_pipeline(sampling=_distilled_policy()), req) From 4a3cb48761de5a61d6beb9fc4a1bfcb15e057d60 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Sat, 8 Aug 2026 10:18:15 -0700 Subject: [PATCH 27/35] [None][test] Exclude anyio's pool thread from the thread-leak check test_auto_resolves_to_tensor_payload failed pipeline #52355 with "Test leaked []"; the rerun on identical code passed. It is a race, not a leak: starlette's FileResponse reads the download through anyio.to_thread, anyio stops the pool thread when the TestClient portal closes, but the thread's exit lands a few milliseconds after the request returns -- and pytest-threadleak samples live threads at the end of the call phase, inside that window. Any file-shipping test in the serve-endpoint module can lose the race; which one does varies by machine load. A teardown fixture cannot fix this (the plugin's check runs before teardown -- verified: the class still failed with one in place), and no in-test point deterministically outlives the thread, since its exit is asynchronous to the response. The exclude list is the designed mechanism and already carries this exact category: asyncio_\d+ (the encoder path's run_in_executor threads, which is why mp4 tests never tripped) and ThreadPoolExecutor-\d+_\d+. This edits the unittest-wide config, not just visual_gen. A genuinely leaked portal is still caught through its remaining threads, which carry other names. Signed-off-by: Igor Shovkun --- tests/unittest/pytest.ini | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unittest/pytest.ini b/tests/unittest/pytest.ini index 5eb84b512ab5..be3cea806d97 100644 --- a/tests/unittest/pytest.ini +++ b/tests/unittest/pytest.ini @@ -8,7 +8,7 @@ threadleak = True # session-reuse-* are the reuse layer's own threads (tests/test_common/session_reuse.py) # and session-prefetch-* are the prefetcher's (tests/test_common/session_prefetcher.py); # they legitimately outlive the test they start under. -threadleak_exclude = asyncio_\d+|rpc_client_loop|rpc_client_worker_\d+|rpc_server_worker_\d+|InductorSubproc|subproc_worker_timer|ThreadPoolExecutor-\d+_\d+|Thread-\d+ \(_manager_spawn\)|session-reuse-\w+|session-prefetch-\w+ +threadleak_exclude = AnyIO worker thread|asyncio_\d+|rpc_client_loop|rpc_client_worker_\d+|rpc_server_worker_\d+|InductorSubproc|subproc_worker_timer|ThreadPoolExecutor-\d+_\d+|Thread-\d+ \(_manager_spawn\)|session-reuse-\w+|session-prefetch-\w+ addopts = --durations=0 -W ignore::DeprecationWarning -p test_common.session_reuse_hooks -p test_common.s3_output_hooks pythonpath = auto_deploy/_utils_test From 68be7d786e005325f6109aa067464d2f4fa14737 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Sat, 8 Aug 2026 10:30:48 -0700 Subject: [PATCH 28/35] Revert "[None][test] Exclude anyio's pool thread from the thread-leak check" This reverts commit 4a3cb48761de5a61d6beb9fc4a1bfcb15e057d60. Signed-off-by: Igor Shovkun --- tests/unittest/pytest.ini | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unittest/pytest.ini b/tests/unittest/pytest.ini index be3cea806d97..5eb84b512ab5 100644 --- a/tests/unittest/pytest.ini +++ b/tests/unittest/pytest.ini @@ -8,7 +8,7 @@ threadleak = True # session-reuse-* are the reuse layer's own threads (tests/test_common/session_reuse.py) # and session-prefetch-* are the prefetcher's (tests/test_common/session_prefetcher.py); # they legitimately outlive the test they start under. -threadleak_exclude = AnyIO worker thread|asyncio_\d+|rpc_client_loop|rpc_client_worker_\d+|rpc_server_worker_\d+|InductorSubproc|subproc_worker_timer|ThreadPoolExecutor-\d+_\d+|Thread-\d+ \(_manager_spawn\)|session-reuse-\w+|session-prefetch-\w+ +threadleak_exclude = asyncio_\d+|rpc_client_loop|rpc_client_worker_\d+|rpc_server_worker_\d+|InductorSubproc|subproc_worker_timer|ThreadPoolExecutor-\d+_\d+|Thread-\d+ \(_manager_spawn\)|session-reuse-\w+|session-prefetch-\w+ addopts = --durations=0 -W ignore::DeprecationWarning -p test_common.session_reuse_hooks -p test_common.s3_output_hooks pythonpath = auto_deploy/_utils_test From 6860fbcb1a85ceef470ef287e04f0c22ef6f2c0a Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Sat, 8 Aug 2026 10:36:32 -0700 Subject: [PATCH 29/35] [None][test] Outwait FileResponse's reader thread in the serve-endpoint tests test_auto_resolves_to_tensor_payload failed pipeline #52355 as a phantom thread leak. FileResponse reads the download on an anyio pool thread that is told to stop when the request's portal closes but exits a few milliseconds later -- inside the window where pytest-threadleak samples running threads at the end of the call phase. Any file-shipping test in this module can lose that race; which one does varies with machine load. Wrap the module's TestClient so request() joins the worker before returning. The join runs in the call phase, ahead of the plugin's check, and stop is already queued by then, so it waits milliseconds; the deadline is a guard. Scoped to this module on purpose: excluding the thread name suite-wide in pytest.ini would be a policy change belonging to the test-config owners. Signed-off-by: Igor Shovkun --- .../visual_gen/test_trtllm_serve_endpoints.py | 26 ++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) 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 8cb87240c8f6..de7e675e5c6a 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py @@ -16,6 +16,8 @@ import asyncio import base64 import os +import threading +import time from io import BytesIO from pathlib import Path from typing import Optional @@ -307,6 +309,28 @@ def result(self, timeout=None): # --------------------------------------------------------------------------- +class _ThreadSettlingTestClient(TestClient): + """TestClient that outwaits starlette's FileResponse reader thread. + + pytest-threadleak samples running threads at the end of a test's call + phase. ``FileResponse`` reads the download on an anyio pool thread + ("AnyIO worker thread") that is told to stop when the request's portal + closes but exits a few milliseconds later -- inside the sampling window, + so any file-shipping test in this module can fail as a phantom leak + (pipeline #52355 did). Joining here, still in the call phase, waits out + those milliseconds deterministically; stop is already queued, so the + deadline is a guard, not an expected wait. + """ + + def request(self, *args, **kwargs): + response = super().request(*args, **kwargs) + deadline = time.monotonic() + 10.0 + for thread in threading.enumerate(): + if thread.name == "AnyIO worker thread": + thread.join(max(0.0, deadline - time.monotonic())) + return response + + def _create_server(generator: MockVisualGen, model_name: str = "test-model") -> TestClient: """Instantiate an OpenAIServer for VISUAL_GEN with a mocked generator. @@ -324,7 +348,7 @@ def _create_server(generator: MockVisualGen, model_name: str = "test-model") -> server_role=ServerRole.VISUAL_GEN, metadata_server_cfg=None, ) - client = TestClient(server.app) + client = _ThreadSettlingTestClient(server.app) # Expose the mock so tests can assert captured generate() arguments. client.mock_gen = generator return client From c45108e8d1053628ec8d643b645c1df62681185d Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Tue, 11 Aug 2026 14:22:54 -0700 Subject: [PATCH 30/35] [None][fix] Restore transformer action wiring lost in the Edge merge Edge's transformer __init__ recorded action_gen as a config fact only ("the transformer never constructs action modules") -- true on main, where action is stubbed, but the auto-merge stitched that line onto this branch's head-construction block, which reads self.action_gen thirty lines later: 47 tests failed on the missing attribute. action_gen is real again; has_action_weights stays as the alias Edge's callers use. Edge's new conditioning-anchor tests call the post-step callback with one argument, predating this branch's two-argument contract (the shared denoise loop threads extra_stream_latents through for the action/audio side streams). Production was consistent on both sides; only the two new test callsites move to (latents, None). Signed-off-by: Igor Shovkun --- .../_torch/visual_gen/models/cosmos3/transformer_cosmos3.py | 5 +++-- tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py | 4 ++-- 2 files changed, 5 insertions(+), 4 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 e813d68a6747..3608167d76a7 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -1034,8 +1034,9 @@ def __init__(self, model_config: DiffusionModelConfig): pretrained_config = apply_pretrained_config_compat_defaults(model_config.pretrained_config) self.recipe = resolve_arch_recipe(pretrained_config) self.audio_gen = getattr(pretrained_config, "sound_gen", False) - # Config fact only: the transformer never constructs action modules. - self.has_action_weights = getattr(pretrained_config, "action_gen", False) + self.action_gen = getattr(pretrained_config, "action_gen", False) + # Config-fact alias kept for callers that predate action support. + self.has_action_weights = self.action_gen self.hidden_size = pretrained_config.hidden_size self.num_hidden_layers = pretrained_config.num_hidden_layers diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py index 3c6770d4075d..368af6943542 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py @@ -831,7 +831,7 @@ def test_anchor_rejects_dtype_mismatch(self): latents = torch.zeros(1, 4, 3, 2, 2, dtype=torch.float32) with pytest.raises(RuntimeError, match="must match the denoised latents"): - post_step_fn(latents) + post_step_fn(latents, None) def test_anchor_accepts_matching_dtype(self): pipeline = _bare_pipeline(sampling=_distilled_policy()) @@ -840,7 +840,7 @@ def test_anchor_accepts_matching_dtype(self): ) latents = torch.zeros(1, 4, 3, 2, 2, dtype=torch.bfloat16) - post_step_fn(latents) + post_step_fn(latents, None) assert torch.all(latents[:, :, 0:1] == self.CLEAN) def _run_denoise(self, with_anchor: bool): From 4c6073e667f7cc81483370ecf09f3e9ddd844280 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Tue, 11 Aug 2026 20:30:35 -0700 Subject: [PATCH 31/35] [None][test] Edge synthetic checkpoints carry the action weights they declare Pipeline #53158: six TestStrictLoading tests failed because the Edge fixtures encode the pre-action contract. Their configs mirror the real checkpoint's transformer/config.json verbatim -- including action_gen: true -- but _synthetic_state_dict emitted no action tensors, and one test asserted the action families land among the intentional skips. On main that was consistent: the transformer never constructed action modules and skipped their weights. This PR makes declared action weights constructed and consumed, so a declared- but-absent set is now a strict-loading error by design. The real checkpoints already satisfy the new contract: Cosmos3-Edge, like Nano, declares action_gen and ships all five action tensors -- main skipped weights that exist; loading them is the point of this PR. The fixture now emits the five tensors for action_gen configs (DomainAwareLinear stores per-domain weights as nn.Embedding rows of flattened [out * in] matrices). The intentional-skip families shrink to lm_head and norm; the "model."-prefix probe keeps its intent through lm_head, since a prefixed action tensor now remaps to a constructed parameter; and the missing-weight matrix gains action_proj_in.fc.weight and action_modality_embed, pinning that a declared-action checkpoint without action weights fails by name. Signed-off-by: Igor Shovkun --- .../_torch/visual_gen/test_cosmos3_edge.py | 31 ++++++++++--------- 1 file changed, 17 insertions(+), 14 deletions(-) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py b/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py index ffe3777774d4..9feed9e2dad2 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py @@ -164,6 +164,20 @@ def _synthetic_state_dict(cfg: SimpleNamespace) -> dict: "time_embedder.linear_2.weight": torch.randn(h, h), "time_embedder.linear_2.bias": torch.randn(h), } + if getattr(cfg, "action_gen", False): + # DomainAwareLinear stores per-domain weights as nn.Embedding rows of + # flattened [out * in] matrices; the modality embed is a root parameter. + n_dom = cfg.num_embodiment_domains + a = cfg.action_dim + sd.update( + { + "action_modality_embed": torch.randn(h), + "action_proj_in.fc.weight": torch.randn(n_dom, h * a), + "action_proj_in.bias.weight": torch.randn(n_dom, h), + "action_proj_out.fc.weight": torch.randn(n_dom, a * h), + "action_proj_out.bias.weight": torch.randn(n_dom, a), + } + ) if getattr(cfg, "sound_gen", False): sd.update( { @@ -611,6 +625,8 @@ def test_full_checkpoint_loads(self): "missing_key", [ "layers.0.self_attn.k_norm_und_for_gen.weight", + "action_proj_in.fc.weight", + "action_modality_embed", "layers.0.mlp.up_proj.weight", "layers.1.mlp_moe_gen.down_proj.weight", "layers.0.input_layernorm_moe_gen.weight", @@ -638,15 +654,6 @@ def test_intentional_skips_are_logged_with_names(self, monkeypatch): monkeypatch.setattr(tf_module.logger, "info", infos.append) cfg = _reduced_edge_config() sd = _edge_state_dict(cfg) - sd.update( - { - "action_modality_embed": torch.randn(cfg.hidden_size), - "action_proj_in.fc.weight": torch.randn(4, 8), - "action_proj_in.bias.weight": torch.randn(4, 8), - "action_proj_out.fc.weight": torch.randn(4, 8), - "action_proj_out.bias.weight": torch.randn(4, 8), - } - ) model = self._model() model.load_weights(sd) @@ -659,9 +666,6 @@ def test_intentional_skips_are_logged_with_names(self, monkeypatch): # text also mentions lm_head/norm, so assert the parsed set exactly. skipped_families = {name.strip() for name in skip_logs[0].rsplit(": ", 1)[1].split(",")} assert skipped_families == { - "action_modality_embed", - "action_proj_in", - "action_proj_out", "lm_head", "norm", } @@ -679,14 +683,13 @@ def test_model_prefixed_skip_keys_are_intentional(self, monkeypatch): cfg = _reduced_edge_config() sd = _edge_state_dict(cfg) sd["model.lm_head.weight"] = torch.randn(cfg.vocab_size, cfg.hidden_size) - sd["model.action_modality_embed"] = torch.randn(cfg.hidden_size) self._model().load_weights(sd) assert not any("unknown checkpoint key" in m for m in warnings) skip_logs = [m for m in infos if "intentionally unused" in m] assert len(skip_logs) == 1 skipped_families = {name.strip() for name in skip_logs[0].rsplit(": ", 1)[1].split(",")} - assert {"lm_head", "action_modality_embed"} <= skipped_families + assert "lm_head" in skipped_families def test_unconsumed_mapped_tensor_warns(self, monkeypatch): """A checkpoint tensor that remaps to a module the recipe didn't From e6a76a7c5e699018ca562991976ccba6eec79c08 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Tue, 18 Aug 2026 15:50:08 -0700 Subject: [PATCH 32/35] [None][fix] Thin the Cosmos3 action reference to the embodiment's frame rate An embodiment's frame rate describes what the model was trained on, not the clip a caller sends. Bridge learned one command per frame at 5 Hz; the reference decode read frames consecutively at whatever rate the clip was shot. A 30 fps clip therefore handed the model 17 frames spanning 0.567s while the trained JSON caption said "duration": "3s", "fps": 5.0 and the mRoPE temporal positions were laid out for 3.4s -- a 6x disagreement between the pixels and everything describing them. Read the source rate from the container header (one CPU-side demux, no frame decoded) and keep every n-th frame so the retained window really is at the embodiment's rate. decode_video_reference_window gains frame_step. It stays mechanism: the module's vocabulary is frame indices, so no rate is named there and the ratio is the caller's to compute. Skipped frames are still decoded -- inter-frame compression leaves no choice -- but are neither resized nor retained, so the ring stays at the requested length instead of growing by the step. frame_step is rejected on trailing (negative) windows, whose ring length is unknown until EOS; only the action path steps, and it reads from the start. A clip slower than the embodiment cannot be thinned, since selection drops frames and never invents them. That warns rather than failing. A clip too short to supply the widened range fails with the numbers, because the caption is built before the decode: falling back to a consecutive read there would leave the prompt asserting a rate the frames do not have, which is the bug this fixes. Signed-off-by: Igor Shovkun --- .../visual_gen/models/cosmos3/action.py | 21 +++++++++ .../models/cosmos3/pipeline_cosmos3.py | 28 ++++++++++-- .../models/cosmos3/transformer_cosmos3.py | 2 +- tensorrt_llm/media/decoding.py | 26 +++++++++-- .../_torch/visual_gen/test_cosmos3_action.py | 19 ++++++++ .../_torch/visual_gen/test_media_decode.py | 44 +++++++++++++++++++ 6 files changed, 133 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py index f303538f2674..fcc0a92feeeb 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py @@ -457,6 +457,27 @@ def action_reference_size( return reference.height, reference.width +def action_reference_frame_step(source_frame_rate: float | None, target_frame_rate: float) -> int: + """Source frames to advance per reference frame retained. + + An embodiment's frame rate is a property of what it was trained on, not of + the clip a caller happens to send: bridge learned one command per frame at + 5 Hz, so 200ms of gripper motion between frames. A 30 fps clip read + consecutively shows the model a sixth of that motion while the caption and + the mRoPE positions still claim 5 Hz, so the reference is thinned to match + -- every sixth frame here. + + A clip slower than the embodiment returns 1: selection can drop frames, + never invent them, and closing that gap needs interpolation rather than a + step. The caller is expected to say so rather than let it pass silently. + """ + if not source_frame_rate or not target_frame_rate: + return 1 + if source_frame_rate <= 0 or target_frame_rate <= 0: + return 1 + return max(1, round(source_frame_rate / target_frame_rate)) + + def resize_and_pad_action_image( image: PIL.Image.Image, target_h: int, target_w: int ) -> PIL.Image.Image: 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 bfda2edc9808..06da6a8b7187 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -51,6 +51,7 @@ def tqdm(iterable, **kwargs): from .action import ( ACTION_MODE_INVERSE_DYNAMICS, DEFAULT_ACTION_VIEW_POINT, + action_reference_frame_step, action_reference_size, action_start_frame_offset, build_action_json_prompt, @@ -1809,22 +1810,43 @@ def forward( ) prepare_error: Optional[Exception] = None try: + source_info = video_stream_info(video) + source_frame_rate = source_info.frame_rate if source_info else None + frame_step = action_reference_frame_step(source_frame_rate, frame_rate) + if self.rank == 0: + if frame_step > 1: + logger.info( + f"Cosmos3 action reference: {source_frame_rate} fps source " + f"thinned to {frame_rate} fps, keeping every {frame_step} " + f"frames of {(num_frames - 1) * frame_step + 1}" + ) + elif source_frame_rate is not None and source_frame_rate < frame_rate: + logger.warning( + f"Cosmos3 action reference is {source_frame_rate} fps but " + f"{normalized_action_mode} expects {frame_rate} fps: frames are " + "further apart than the model was trained on and cannot be " + "thinned to match. Re-encode the reference at the higher rate, " + "or pass frame_rate explicitly to accept this spacing." + ) # "fit" rather than the V2V default: an action reference is # padded to the canvas, never cropped to it, because the # gripper and target sit at the frame edge. frames_u8 = decode_video_reference_window( video, first_frame=0, - last_frame=num_frames - 1, + last_frame=(num_frames - 1) * frame_step, target_h=height, target_w=width, device=self.device, resize="fit", + frame_step=frame_step, ) if frames_u8.shape[0] < num_frames: raise ValueError( - "Cosmos3 inverse_dynamics requires at least " - f"{num_frames} frames, got {frames_u8.shape[0]}." + f"Cosmos3 inverse_dynamics requires {num_frames} frames at " + f"{frame_rate} fps; a {source_frame_rate} fps reference supplies " + f"{frames_u8.shape[0]} once thinned by {frame_step} " + f"({(num_frames - 1) * frame_step + 1} source frames needed)." ) video_tensor = self._condition_frames_to_video_tensor(frames_u8) del frames_u8 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 45da389a0c6d..4eb2ebd0dbd1 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -1587,7 +1587,7 @@ def forward( else: self.cached_kv = cached_kv_full - # --- Extra modality token injection (mutually exclusive: action OR audio) --- + # --- Extra modality token injection (mutually exclusive: action, audio or control) --- T_vid_tokens = hidden_gen.shape[1] # T * Hp * Wp T_action = 0 T_audio = 0 diff --git a/tensorrt_llm/media/decoding.py b/tensorrt_llm/media/decoding.py index 539b7153e7f5..87c1b9e68288 100644 --- a/tensorrt_llm/media/decoding.py +++ b/tensorrt_llm/media/decoding.py @@ -215,6 +215,7 @@ def decode_video_reference_window( target_w: int, device: torch.device, resize: str = "cover", + frame_step: int = 1, ) -> torch.Tensor: """Decode frames ``[first_frame, last_frame]`` of a reference on device. @@ -222,6 +223,16 @@ def decode_video_reference_window( 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. + ``frame_step`` retains every n-th frame of the range, so ``(0, 96)`` with + ``frame_step=6`` yields frames 0, 6, ... 96 — seventeen frames, not + ninety-seven. A caller whose model expects a frame spacing the source was + not shot at uses this to pick the right frames; the ratio of the two rates + is the caller's to compute, and no rate is named here. Skipped frames are + still decoded (inter-frame compression leaves no choice) but are neither + resized nor retained, so the cost is decode time, not memory. Only + non-negative ranges may step: the negative form wraps a ring whose length + is not known until EOS, and combining the two is not supported. + ``resize`` selects how each frame reaches ``target_h x target_w``: ``"cover"`` scales to fill and center-crops (the default, and what video continuation wants); ``"fit"`` scales to fit and pads, for references whose @@ -247,6 +258,13 @@ def decode_video_reference_window( raise ValueError( f"first_frame must not exceed last_frame, got ({first_frame}, {last_frame})." ) + if frame_step < 1: + raise ValueError(f"frame_step must be at least 1, got {frame_step}.") + if frame_step > 1 and first_frame < 0: + raise ValueError( + f"frame_step > 1 is only supported for non-negative ranges, got " + f"({first_frame}, {last_frame}) with frame_step={frame_step}." + ) resize_frames = _RESIZE_MODES.get(resize) if resize_frames is None: raise ValueError( @@ -306,7 +324,7 @@ def _read(buf: bytearray) -> int: # 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 + tail = -first_frame if from_end else (window + frame_step - 1) // frame_step 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 @@ -317,7 +335,9 @@ def _read(buf: bytearray) -> int: if not from_end and count > last_frame: done = True break - if from_end or count >= first_frame: + if from_end or ( + count >= first_frame and (count - first_frame) % frame_step == 0 + ): decoded = torch.from_dlpack(frame) # Ownership copy off the NVDEC surface (recycled by # the decoder) and resize-before-retain in one step. @@ -331,7 +351,7 @@ def _read(buf: bytearray) -> int: except torch.cuda.OutOfMemoryError as exc: raise MemoryError( f"Out of device memory while decoding the video reference " - f"({window} frames @ {target_w}x{target_h} retained): {exc}" + f"({tail} frames @ {target_w}x{target_h} retained): {exc}" ) from exc except nvc.PyNvVCException as exc: raise ValueError( diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py index 61859e44ddbc..fa594ba78de6 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_action.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py @@ -22,6 +22,7 @@ EMBODIMENT_TO_RAW_ACTION_DIM, VIDEO_RES_SIZE_INFO, action_aspect_ratio_label, + action_reference_frame_step, action_reference_size, build_action_json_prompt, find_closest_target_size, @@ -298,6 +299,24 @@ def test_unreadable_video_bytes_are_rejected(self, monkeypatch): action_reference_size(action_mode="inverse_dynamics", image=None, video=b"\x00bad") +class TestActionReferenceFrameStep: + """The reference is thinned to the embodiment's rate, never invented.""" + + @pytest.mark.parametrize( + "source_frame_rate, target_frame_rate, expected", + [ + (30.0, 5.0, 6), # bridge: every sixth frame of a 30 fps clip + (5.0, 5.0, 1), # already at the trained rate + (24.0, 5.0, 5), # 4.8 rounds to 5 + (10.0, 30.0, 1), # slower than trained: cannot be thinned + (None, 5.0, 1), # header unreadable + (0.0, 5.0, 1), # header reported nothing usable + ], + ) + def test_step_from_rates(self, source_frame_rate, target_frame_rate, expected): + assert action_reference_frame_step(source_frame_rate, target_frame_rate) == expected + + class TestActionJsonPrompt: """The trained action caption: structured JSON, not the flat video templates.""" diff --git a/tests/unittest/_torch/visual_gen/test_media_decode.py b/tests/unittest/_torch/visual_gen/test_media_decode.py index d92ba8a6863f..90bd16fba0b8 100644 --- a/tests/unittest/_torch/visual_gen/test_media_decode.py +++ b/tests/unittest/_torch/visual_gen/test_media_decode.py @@ -247,6 +247,50 @@ 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_frame_step_thins_the_window(self): + window = decode_video_reference_window( + _MP4.read_bytes(), + first_frame=0, + last_frame=8, + target_h=64, + target_w=64, + device=self._DEVICE, + frame_step=2, + ) + assert window.shape[0] == 5 + assert _frame_indices(window) == [0, 2, 4, 6, 8] + + def test_frame_step_keeps_a_partial_tail(self): + # The range end is not a multiple of the step: 0, 3, 6 and stop. + window = decode_video_reference_window( + _MP4.read_bytes(), + first_frame=0, + last_frame=7, + target_h=64, + target_w=64, + device=self._DEVICE, + frame_step=3, + ) + assert _frame_indices(window) == [0, 3, 6] + + def test_frame_step_one_matches_the_default(self): + span = dict( + first_frame=0, last_frame=4, target_h=64, target_w=64, device=self._DEVICE + ) + assert torch.equal( + decode_video_reference_window(_MP4.read_bytes(), **span), + decode_video_reference_window(_MP4.read_bytes(), frame_step=1, **span), + ) + + def test_frame_step_rejects_trailing_windows(self): + # The trailing form wraps a ring whose length is unknown until EOS. + with pytest.raises(ValueError, match="non-negative ranges"): + self._decode(_MP4.read_bytes(), keep="last", frame_step=2) + + def test_frame_step_must_be_positive(self): + with pytest.raises(ValueError, match="at least 1"): + self._decode(_MP4.read_bytes(), frame_step=0) + 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 From fbbaa4831dc928dedb132b35ceb65ca6c5186991 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Fri, 21 Aug 2026 22:36:57 -0700 Subject: [PATCH 33/35] [None][test] Align the short-reference assertion with the thinning error; cover the widened window The frame-thinning fix reworded the inverse-dynamics short-clip error to carry the rates and the widened frame requirement; the pipeline suite still matched the old "requires at least" text. Match the new shape instead, and add a case that requests 5 fps from the 24 fps fixture so the widened window (step 5, 41 source frames from a 9-frame clip) is asserted at pipeline level rather than only in the unit helpers. _run_forward gains a frame_rate parameter (defaulting to the suite-wide pin) so a test can vary the requested rate without duplicating the keyword. Signed-off-by: Igor Shovkun --- .../visual_gen/test_cosmos3_pipeline.py | 25 +++++++++++++++++-- 1 file changed, 23 insertions(+), 2 deletions(-) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index b3c7befb86dc..e9c1babe7650 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -145,6 +145,7 @@ def _run_forward( height=HEIGHT, width=WIDTH, guidance_scale=GUIDANCE_SCALE, + frame_rate=FRAME_RATE, **extra, ): return pipeline.forward( @@ -156,7 +157,7 @@ def _run_forward( num_inference_steps=NUM_STEPS, guidance_scale=guidance_scale, seed=SEED, - frame_rate=FRAME_RATE, + frame_rate=frame_rate, use_guardrails=False, **extra, ) @@ -1413,7 +1414,7 @@ def test_inverse_dynamics_smoke(self, cosmos3_pipeline): def test_inverse_dynamics_rejects_short_video(self, cosmos3_pipeline): _require_action_pipeline(cosmos3_pipeline) - with pytest.raises(ValueError, match="requires at least"): + with pytest.raises(ValueError, match=r"requires \d+ frames at"): _run_forward( cosmos3_pipeline, image=None, @@ -1428,6 +1429,26 @@ def test_inverse_dynamics_rejects_short_video(self, cosmos3_pipeline): video=_V2V_FIXTURE_MP4.read_bytes(), ) + def test_inverse_dynamics_thins_to_the_requested_rate(self, cosmos3_pipeline): + """A 24 fps reference asked for at 5 fps keeps every 5th frame, so the + window widens accordingly and the 9-frame fixture comes up short.""" + _require_action_pipeline(cosmos3_pipeline) + with pytest.raises(ValueError, match=r"thinned by 5 \(41 source frames needed\)"): + _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"], + frame_rate=5.0, + action_mode="inverse_dynamics", + domain_name="bridge_orig_lerobot", + raw_action_dim=self.RAW_ACTION_DIM, + action_chunk_size=NUM_FRAMES - 1, + video=_V2V_FIXTURE_MP4.read_bytes(), + ) + def test_out_of_range_domain_id_rejected_before_decode(self, cosmos3_pipeline): _require_action_pipeline(cosmos3_pipeline) with pytest.raises(ValueError, match=r"domain_id must be in \[0, \d+\)"): From 1f90daac5de1b8290745f46dcdcc9ede499c06b0 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Fri, 21 Aug 2026 23:04:31 -0700 Subject: [PATCH 34/35] [None][chore] Collapse a dict literal ruff-format rejects Signed-off-by: Igor Shovkun --- tests/unittest/_torch/visual_gen/test_media_decode.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/unittest/_torch/visual_gen/test_media_decode.py b/tests/unittest/_torch/visual_gen/test_media_decode.py index 90bd16fba0b8..0c88f47ff86b 100644 --- a/tests/unittest/_torch/visual_gen/test_media_decode.py +++ b/tests/unittest/_torch/visual_gen/test_media_decode.py @@ -274,9 +274,7 @@ def test_frame_step_keeps_a_partial_tail(self): assert _frame_indices(window) == [0, 3, 6] def test_frame_step_one_matches_the_default(self): - span = dict( - first_frame=0, last_frame=4, target_h=64, target_w=64, device=self._DEVICE - ) + span = dict(first_frame=0, last_frame=4, target_h=64, target_w=64, device=self._DEVICE) assert torch.equal( decode_video_reference_window(_MP4.read_bytes(), **span), decode_video_reference_window(_MP4.read_bytes(), frame_step=1, **span), From 4b2b6d8ca78badb27fd79cd10e8a174692910714 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 26 Aug 2026 22:24:24 -0700 Subject: [PATCH 35/35] [None][test] Waive test_wan_t2v_example (corrupt NVFP4 checkpoint on CI storage) The Wan2.2-T2V-A14B-Diffusers-NVFP4 transformer shard 00002-of-00002 on CI's shared model storage fails safetensors header deserialization ("header too large"), so the example dies loading weights before any pipeline code runs. Deterministic across every PR whose selection includes visual-gen examples (L0_Test-x86_64-Single-GPU 7181, 7182, 7186); unrelated to this PR. Waived until the checkpoint is re-synced. Signed-off-by: Igor Shovkun --- tests/integration/test_lists/waives.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index bceeb940cd60..32b08da9a1ef 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -129,6 +129,7 @@ examples/visual_gen/test_visual_gen_wan.py::test_wan_feature_accuracy_against_go examples/visual_gen/test_visual_gen_wan.py::test_wan_feature_accuracy_against_golden[wan22-cuda-graph] SKIP (https://nvbugs/6572800) examples/visual_gen/test_visual_gen_wan.py::test_wan_feature_accuracy_against_golden[wan22-fp8-blockwise] SKIP (https://nvbugs/6572800) examples/visual_gen/test_visual_gen_wan.py::test_wan_feature_accuracy_against_golden[wan22-nvfp4] SKIP (https://nvbugs/6572800) +examples/visual_gen/test_visual_gen_wan.py::test_wan_t2v_example SKIP (corrupt Wan2.2-T2V-A14B-Diffusers-NVFP4 safetensors shard on CI model storage; fails all visual-gen PRs since L0_Test-x86_64-Single-GPU 7181) full:A100/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp SKIP (https://nvbugs/6275856) full:A100/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[forced_chunked_prefill] SKIP (https://nvbugs/6597570) full:A100/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[full_budget] SKIP (https://nvbugs/6597570)