diff --git a/examples/visual_gen/models/cosmos3/README.md b/examples/visual_gen/models/cosmos3/README.md index 69be21fe4880..65d5595e7f99 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. Save action runs as `safetensors` or `pt` so the rollout and action tensor stay in one payload. ## 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.safetensors \ + --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.safetensors + +# 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.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 de9e9e5010ad..1144c87eb0f6 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3.py +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -23,6 +23,8 @@ from tensorrt_llm import VisualGen, VisualGenArgs _SCRIPT_DIR = Path(__file__).resolve().parent +_ACTION_MODES = ("policy", "forward_dynamics", "inverse_dynamics") +_TENSOR_OUTPUT_SUFFIXES = {".pt", ".safetensors"} def _resolve_path(path: str) -> str: @@ -80,6 +82,91 @@ 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 {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": + 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 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)." + ) + 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 _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: + 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 +226,68 @@ 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=list(_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="Action tokens to generate. Defaults to the domain preset or model default.", + ) + 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, .mp4/.avi video, or image path for inverse_dynamics", + ) + parser.add_argument( + "--action_resolution", + type=int, + default=None, + choices=[256, 480, 704, 720], + help=("Resolution bucket for action image sizing. Defaults to the domain preset or 480."), + ) + parser.add_argument( + "--action_fps", + type=float, + default=None, + help="Action-token temporal rate for mRoPE (Hz). Defaults to frame_rate.", + ) + parser.add_argument( + "--action_output_path", + type=str, + default=None, + help="Path to save predicted action JSON (default: _action.json)", + ) parser.add_argument( "--output_type", type=str, default="video", help="Output type (video, image)" ) @@ -156,6 +305,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 @@ -166,6 +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 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: @@ -186,6 +338,24 @@ def main(): params.extra_params["use_guardrails"] = not args.disable_guardrails params.extra_params["output_type"] = output_type + 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_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: + with open(args.action_json, encoding="utf-8") as f: + params.extra_params["action"] = json.load(f) + if negative_prompt is None: params.negative_prompt = None elif isinstance(negative_prompt, str): @@ -198,8 +368,22 @@ 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(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..fea4a71a5ec6 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py @@ -0,0 +1,470 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Action-token helpers for Cosmos3 UVA/action generation.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +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" +ACTION_MODE_INVERSE_DYNAMICS = "inverse_dynamics" +ACTION_MODES = { + ACTION_MODE_POLICY, + ACTION_MODE_FORWARD_DYNAMICS, + ACTION_MODE_INVERSE_DYNAMICS, +} + +EMBODIMENT_TO_DOMAIN_ID: dict[str, int] = { + "no_action": 0, + "av": 1, + "camera_pose": 2, + "hand_pose": 3, + "pusht": 4, + "libero": 5, + "umi": 6, + "bridge_orig_lerobot": 7, + "droid_lerobot": 8, + "robomind-franka": 8, + "galbot": 9, + "robomind-franka-dual": 12, + "robomind-ur": 13, + "agibotworld": 15, + "agibot_gear_gripper": 15, + "agibot_gear_gripper_ext": 15, + "fractal": 20, +} + +VIDEO_RES_SIZE_INFO: dict[str, dict[str, tuple[int, int]]] = { + "256": { + "1,1": (256, 256), + "4,3": (320, 256), + "3,4": (256, 320), + "16,9": (320, 192), + "9,16": (192, 320), + }, + "480": { + "1,1": (640, 640), + "4,3": (736, 544), + "3,4": (544, 736), + "16,9": (832, 480), + "9,16": (480, 832), + }, + "704": { + "1,1": (960, 960), + "4,3": (1088, 832), + "3,4": (832, 1088), + "16,9": (1280, 704), + "9,16": (704, 1280), + }, + "720": { + "1,1": (960, 960), + "4,3": (1104, 832), + "3,4": (832, 1104), + "16,9": (1280, 720), + "9,16": (720, 1280), + }, +} + + +COSMOS3_ACTION_RESOLUTIONS = tuple(int(key) for key in sorted(VIDEO_RES_SIZE_INFO, key=int)) + + +def normalize_action_resolution(resolution: Any) -> int: + 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 + 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 + + +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, Path)): + return PIL.Image.open(Path(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: 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, + 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 = 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}." + ) + + +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: 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: + 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]) + 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: + 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 f5747544946d..e0e8c06ba384 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, 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,22 @@ "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 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"], +} + +# Text-to-image (``output_type="image"``) defaults. Applied when the request +# field is ``None``. COSMOS3_T2I_PARAMS = { "height": 1024, "width": 1024, @@ -48,7 +74,254 @@ "guidance_interval": (400.0, 1000.0), } -COSMOS3_EXTRA_SPECS: Dict[str, ExtraParamSchema] = { +COSMOS3_ACTION_PARAMS = { + "action_chunk_size": 16, + "num_frames": 17, + "num_inference_steps": 30, + "guidance_scale": 1.0, + "flow_shift": 5.0, + "frame_rate": 24.0, +} + + +class Cosmos3DomainPreset(TypedDict, total=False): + """Recommended action-generation settings for a trained embodiment.""" + + raw_action_dim: int + action_chunk_size: int + num_frames: int + action_resolution: int + frame_rate: float + + +# Training-aligned defaults. Values mirror Cosmos3 omni action JSON examples where available. +COSMOS3_DOMAIN_PRESETS: dict[str, Cosmos3DomainPreset] = { + # WidowX bridge; 7-DOF arm + gripper in 10-D state. + "bridge_orig_lerobot": { + "raw_action_dim": 10, + "action_chunk_size": 16, + "num_frames": 17, + "action_resolution": 480, + "frame_rate": 5.0, + }, + # Autonomous-vehicle steering/throttle; longer action horizon. + "av": { + "raw_action_dim": 9, + "action_chunk_size": 60, + "num_frames": 61, + "action_resolution": 480, + "frame_rate": 10.0, + }, + # 6-DoF camera pose + shutter; matches AV-style horizon. + "camera_pose": { + "raw_action_dim": 9, + "action_chunk_size": 60, + "num_frames": 61, + "action_resolution": 480, + "frame_rate": 30.0, + }, + # Franka single-arm tabletop; same domain_id as robomind-franka. + "droid_lerobot": { + "raw_action_dim": 10, + "action_chunk_size": 16, + "num_frames": 17, + "action_resolution": 480, + "frame_rate": 15.0, + }, + # LIBERO sim single-arm; lower action resolution bucket. + "libero": { + "raw_action_dim": 10, + "action_chunk_size": 16, + "num_frames": 17, + "action_resolution": 256, + "frame_rate": 10.0, + }, + # MANO hand pose; high-DOF hand articulation. + "hand_pose": { + "raw_action_dim": 57, + "action_chunk_size": 16, + "num_frames": 17, + "action_resolution": 480, + "frame_rate": 24.0, + }, + # AgiBot humanoid; shared domain_id with agibot_gear_gripper*. + "agibotworld": { + "raw_action_dim": 29, + "action_chunk_size": 16, + "num_frames": 17, + "action_resolution": 480, + "frame_rate": 10.0, + }, + # Google Robot (RT-1 / fractal) single-arm. + "fractal": { + "raw_action_dim": 10, + "action_chunk_size": 16, + "num_frames": 17, + "action_resolution": 480, + "frame_rate": 5.0, + }, + # 2-D planar push task. + "pusht": { + "raw_action_dim": 2, + "action_chunk_size": 16, + "num_frames": 17, + "action_resolution": 256, + "frame_rate": 10.0, + }, + # UMI handheld gripper setup. + "umi": { + "raw_action_dim": 10, + "action_chunk_size": 16, + "num_frames": 17, + "action_resolution": 480, + "frame_rate": 10.0, + }, +} + +# Map alias domain_name keys to a canonical preset entry. +COSMOS3_DOMAIN_PRESET_ALIASES: dict[str, str] = { + "robomind-franka": "droid_lerobot", + "robomind-franka-dual": "droid_lerobot", + "robomind-ur": "droid_lerobot", + "agibot_gear_gripper": "agibotworld", + "agibot_gear_gripper_ext": "agibotworld", + "galbot": "agibotworld", +} + + +def canonical_domain_preset_key( + domain_name: 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) + 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: 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 + return COSMOS3_DOMAIN_PRESETS[key] + + +def resolve_domain_action_config( + *, + 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] = [] + + 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, + current: Any, + *, + fallback: Any = None, + ) -> Any: + recommended = preset.get(field) if preset else None + if current is not None: + if recommended is not None and current != recommended: + warnings.append( + f"Cosmos3 {field}={current} differs from recommended " + f"{recommended} for domain {preset_key!r}." + ) + return current + if recommended is not None: + return recommended + return fallback + + resolved_raw_action_dim = _resolve_field("raw_action_dim", raw_action_dim) + resolved_chunk = _resolve_field( + "action_chunk_size", + action_chunk_size, + fallback=COSMOS3_ACTION_PARAMS["action_chunk_size"], + ) + resolved_resolution = normalize_action_resolution( + _resolve_field( + "action_resolution", + action_resolution, + fallback=480, + ) + ) + resolved_frame_rate = _resolve_field( + "frame_rate", + frame_rate, + fallback=COSMOS3_ACTION_PARAMS["frame_rate"], + ) + resolved_num_frames = _resolve_field("num_frames", num_frames) + if resolved_num_frames is None: + resolved_num_frames = int(resolved_chunk) + 1 + resolved_action_fps = ( + float(action_fps) if action_fps is not None else float(resolved_frame_rate) + ) + 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, + "action_chunk_size": int(resolved_chunk), + "action_resolution": resolved_resolution, + "frame_rate": float(resolved_frame_rate), + "action_fps": resolved_action_fps, + "num_frames": int(resolved_num_frames), + "preset_key": preset_key, + "warnings": warnings, + } + + +COSMOS3_EXTRA_SPECS: dict[str, ExtraParamSchema] = { "use_duration_template": ExtraParamSchema( type="bool", default=True, @@ -79,4 +352,68 @@ default="video", description="Output modality: 'video' (T2V/I2V) or 'image' (text-to-image).", ), + "action_mode": ExtraParamSchema( + type="Literal['policy', 'forward_dynamics', 'inverse_dynamics']", + default=None, + description="Action generation mode: policy, forward_dynamics, or inverse_dynamics.", + ), + "domain_name": ExtraParamSchema( + type="str", + default=None, + description=( + "Embodiment domain name for action generation (e.g. bridge_orig_lerobot, av). " + "When set, omitted raw_action_dim/action_chunk_size/action_resolution/frame_rate " + "are filled from COSMOS3_DOMAIN_PRESETS; mismatches are logged as warnings." + ), + ), + "domain_id": ExtraParamSchema( + type="int", + default=None, + description="Embodiment domain id for action generation.", + ), + "raw_action_dim": ExtraParamSchema( + type="int", + default=None, + description=( + "Raw action DOF for policy/inverse_dynamics (e.g. 10 bridge, 9 av, 29 agibot). " + "Inferred from domain_name preset when omitted." + ), + ), + "action_chunk_size": ExtraParamSchema( + type="int", + default=None, + description=( + "Number of action tokens to generate (16 for most robots, 60 for av/camera_pose). " + "Inferred from domain_name preset when omitted." + ), + ), + "action": ExtraParamSchema( + type="list", + default=None, + description="Action trajectory [T, D] for forward_dynamics mode.", + ), + "action_resolution": ExtraParamSchema( + type="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." + ), + 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, + 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 2dac241f2115..a22fb6810aa9 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -17,7 +17,7 @@ import math import os import time -from typing import List, Optional, Union +from typing import Any, List, Optional, Union import PIL.Image import torch @@ -34,7 +34,27 @@ 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_INVERSE_DYNAMICS, + action_reference_image, + action_start_frame_offset, + build_vision_condition_mask, + normalize_action_mode, + normalize_action_video_input, + pil_to_rgb, + prepare_action_latents, + resize_and_pad_action_image, + resolve_action_size, + resolve_domain_id, +) +from .defaults import ( + COSMOS3_720P_PARAMS, + COSMOS3_ACTION_PARAMS, + 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 from .transformer_cosmos3 import Cosmos3VFMTransformer @@ -156,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 @@ -200,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 @@ -212,7 +244,7 @@ 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): @@ -235,53 +267,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 +295,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"), + action_fps=extra_params.get("action_fps"), + video=extra_params.get("video"), ) def _apply_metadata_templates( @@ -452,7 +460,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, @@ -608,6 +618,112 @@ def decode_audio(self, latent: torch.Tensor) -> torch.Tensor: """ return self.audio_tokenizer.decode(latent).float() # [B, audio_channels, N_samples] + # ========================================================================= + # Action generation + # ========================================================================= + + def _preprocess_action_image( + self, image: PIL.Image.Image, target_h: int, target_w: int + ) -> torch.Tensor: + image = resize_and_pad_action_image(image, target_h, target_w) + return self.video_processor.preprocess(image, height=target_h, width=target_w) + + def _preprocess_action_video( + self, frames: List[Any], target_h: int, target_w: int + ) -> torch.Tensor: + if not frames: + raise ValueError("Cosmos3 action video input must contain at least one frame.") + processed = [ + self._preprocess_action_image(pil_to_rgb(frame), target_h, target_w).squeeze(0) + for frame in frames + ] + return torch.stack(processed, dim=1).unsqueeze(0).contiguous() + + def _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]: + 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, + ) + # ========================================================================= # Forward (main generation entry point) # ========================================================================= @@ -620,19 +736,28 @@ def forward( seed: int, negative_prompt: Optional[str] = None, image: Optional[Union[PIL.Image.Image, torch.Tensor, str]] = None, - height: int = COSMOS3_720P_PARAMS["height"], - width: int = COSMOS3_720P_PARAMS["width"], - num_frames: int = COSMOS3_720P_PARAMS["num_frames"], - num_inference_steps: int = COSMOS3_720P_PARAMS["num_inference_steps"], - guidance_scale: float = COSMOS3_720P_PARAMS["guidance_scale"], - max_sequence_length: int = COSMOS3_720P_PARAMS["max_sequence_length"], - frame_rate: float = COSMOS3_720P_PARAMS["frame_rate"], + height: Optional[int] = None, + width: Optional[int] = None, + num_frames: Optional[int] = None, + num_inference_steps: Optional[int] = None, + guidance_scale: Optional[float] = None, + 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: Optional[int] = None, + action_fps: Optional[float] = None, + video: Any = None, ): pipeline_start = time.time() timer = CudaPhaseTimer() @@ -640,25 +765,122 @@ def forward( use_guardrails = use_guardrails and not TRTLLM_DISABLE_COSMOS3_GUARDRAILS + normalized_action_mode = normalize_action_mode(action_mode) + do_action = normalized_action_mode is not None + if do_action and not self.action_gen: + raise ValueError( + "Cosmos3 action generation was requested, but this checkpoint " + "does not enable action_gen." + ) + if do_action and enable_audio: + raise ValueError("Cosmos3 does not support joint action and audio generation.") + # Text-to-image mode: same checkpoint/forward path as T2V, but a single # latent frame, image-flavored prompt templates, flow_shift=3.0, a CFG # guidance interval, and an image (rather than video) output. is_t2i = str(output_type).lower() == "image" guidance_interval = None + resolved_action_fps: Optional[float] = None if is_t2i: if image is not None: raise ValueError( "Cosmos3 text-to-image (output_type='image') does not accept an image input." ) + if do_action: + raise ValueError("Cosmos3 action generation does not support output_type='image'.") + if enable_audio: + raise ValueError("Cosmos3 audio generation does not support output_type='image'.") num_frames = 1 - 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_cfg = resolve_domain_action_config( + domain_name=domain_name, + domain_id=domain_id, + raw_action_dim=raw_action_dim, + action_chunk_size=action_chunk_size, + action_resolution=action_resolution, + frame_rate=frame_rate, + action_fps=action_fps, + num_frames=num_frames, + ) + if self.rank == 0: + for warning in action_cfg["warnings"]: + logger.warning(warning) + if action_cfg["preset_key"] is not None: + logger.info( + f"Cosmos3 action domain preset {action_cfg['preset_key']!r}: " + f"raw_action_dim={action_cfg['raw_action_dim']}, " + f"action_chunk_size={action_cfg['action_chunk_size']}, " + f"action_resolution={action_cfg['action_resolution']}, " + f"frame_rate={action_cfg['frame_rate']:.1f}, " + f"action_fps={action_cfg['action_fps']:.1f}, " + f"num_frames={action_cfg['num_frames']}" + ) + + raw_action_dim = action_cfg["raw_action_dim"] + action_chunk_size = action_cfg["action_chunk_size"] + action_resolution = action_cfg["action_resolution"] + num_frames = action_cfg["num_frames"] + frame_rate = action_cfg["frame_rate"] + resolved_action_fps = action_cfg["action_fps"] + num_inference_steps = ( + num_inference_steps or COSMOS3_ACTION_PARAMS["num_inference_steps"] + ) + if guidance_scale is None: + guidance_scale = COSMOS3_ACTION_PARAMS["guidance_scale"] + self._set_flow_shift(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"] + if resolved_action_fps is None: + resolved_action_fps = frame_rate + + 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, + video=video, + ) + height, width = resolve_action_size(height, width, action_ref_image, action_resolution) + + if self.rank == 0: + logger.info( + 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( + f"Cosmos3 action dims: action_chunk_size={action_chunk_size}, " + f"action_resolution={action_resolution}, " + f"action_fps={resolved_action_fps:.1f}, " + f"input_aspect={action_ref_image.width / action_ref_image.height:.3f}" + ) + if isinstance(prompt, str): prompt = [prompt] batch_size = len(prompt) @@ -753,7 +975,73 @@ 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: + 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, + 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") @@ -782,6 +1070,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") @@ -801,6 +1091,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 +1111,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 +1132,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=resolved_action_fps, ) video_noise_pred = result.video audio_noise_pred = result.audio + action_noise_pred = result.action if velocity_mask is not None: video_noise_pred = video_noise_pred * velocity_mask + if action_noise_pred is not None: + if action_velocity_mask is not None: + action_noise_pred = action_noise_pred * action_velocity_mask + if ( + resolved_raw_action_dim is not None + and 0 < resolved_raw_action_dim < action_noise_pred.shape[-1] + ): + action_noise_pred = action_noise_pred.clone() + action_noise_pred[..., resolved_raw_action_dim:] = 0 + return video_noise_pred, {"action": action_noise_pred} + if audio_noise_pred is not None: return video_noise_pred, {"audio": audio_noise_pred} return video_noise_pred + 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 +1196,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.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. denoise_result = self.denoise( latents=latents, scheduler=self.scheduler, @@ -861,11 +1212,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 +1265,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..7740efc1a24b 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,78 @@ def compute_mrope_position_ids_vision( return mrope_ids, next_offset +def compute_mrope_position_ids_action( + grid_t: int, + temporal_offset: int | float, + action_fps: float | None, + base_fps: float = 24.0, + base_temporal_compression_factor: int = 4, + enable_fps_modulation: bool = True, + start_frame_offset: int = 1, +) -> tuple[torch.Tensor, int | float]: + """Generate mRoPE IDs for action tokens as a frame-rate (T, 1, 1) grid.""" + return compute_mrope_position_ids_vision( + grid_t=grid_t, + grid_h=1, + grid_w=1, + temporal_offset=temporal_offset, + fps=action_fps, + base_fps=base_fps, + temporal_compression_factor=1, + enable_fps_modulation=enable_fps_modulation, + start_frame_offset=start_frame_offset, + ) + + +class DomainAwareLinear(nn.Module): + """Linear projection with one weight/bias pair per action embodiment domain.""" + + def __init__( + self, + input_size: int, + output_size: int, + num_domains: int, + *, + dtype: torch.dtype = torch.bfloat16, + ) -> None: + super().__init__() + self.input_size = int(input_size) + self.output_size = int(output_size) + self.num_domains = int(num_domains) + self.dtype = dtype + self.fc = nn.Embedding(self.num_domains, self.output_size * self.input_size, dtype=dtype) + self.bias = nn.Embedding(self.num_domains, self.output_size, dtype=dtype) + + def post_load_weights(self) -> None: + self.fc.to(self.dtype) + self.bias.to(self.dtype) + + def forward(self, x: torch.Tensor, domain_id: torch.Tensor) -> torch.Tensor: + if domain_id.ndim == 0: + domain_id = domain_id.unsqueeze(0) + domain_id = domain_id.to(device=x.device, dtype=torch.long).reshape(-1) + if x.shape[0] != domain_id.shape[0]: + raise ValueError( + "Cosmos3 action domain_id batch size must match action batch: " + f"tokens={x.shape[0]}, domain_id={domain_id.shape[0]}." + ) + if torch.any((domain_id < 0) | (domain_id >= self.num_domains)): + raise ValueError( + f"Cosmos3 action domain_id must be in [0, {self.num_domains}), " + f"got {domain_id.tolist()}." + ) + + weight = self.fc(domain_id).view(domain_id.shape[0], self.input_size, self.output_size) + bias = self.bias(domain_id).view(domain_id.shape[0], self.output_size) + if x.ndim == 2: + return torch.bmm(x.unsqueeze(1), weight).squeeze(1) + bias + if x.ndim == 3: + return torch.bmm(x, weight) + bias.unsqueeze(1) + raise ValueError( + f"Cosmos3 DomainAwareLinear expected rank-2 or rank-3 input, got {tuple(x.shape)}." + ) + + class TimestepEmbedder(nn.Module): """ Embeds scalar timesteps into vector representations. @@ -731,6 +806,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 +1062,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 +1130,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 +1162,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 +1241,44 @@ def forward( else: self.cached_kv = cached_kv_full - # --- Audio token injection ------------------------------------------------- + # --- Extra modality token injection (mutually exclusive: action OR audio) --- T_vid_tokens = hidden_gen.shape[1] # T * Hp * Wp + T_action = 0 T_audio = 0 - if audio_latents is not None and self.audio_gen: + action_domain_ids_tensor = action_domain_ids + if action_latents is not None and self.action_gen: + # FUTURE(action+audio): concat order is video|action|audio; adjust slices below. + if action_domain_ids_tensor is None: + action_domain_ids_tensor = torch.zeros( + action_latents.shape[0], dtype=torch.long, device=action_latents.device + ) + T_action = action_latents.shape[1] + hidden_action = self.action_proj_in( + self.pack_action(action_latents), action_domain_ids_tensor + ) + hidden_action = hidden_action + self.action_modality_embed.to(hidden_action.dtype) + if action_noisy_mask is None: + hidden_action = hidden_action + time_embed.unsqueeze(1) + else: + hidden_action = hidden_action + time_embed.unsqueeze(1) * action_noisy_mask.to( + hidden_action.dtype + ) + hidden_gen = torch.cat([hidden_gen, hidden_action], dim=1) + effective_action_fps = action_fps if action_fps is not None else (fps or self.base_fps) + cos_a, sin_a = self._compute_action_rope_freqs( + T_action, + text_mask, + float(effective_action_fps), + action_start_frame_offset, + hidden_states.device, + hidden_gen.dtype, + ) + cos_v, sin_v = self.cached_freqs_gen + freqs_gen_combined = ( + torch.cat([cos_v, cos_a], dim=1), + torch.cat([sin_v, sin_a], dim=1), + ) + elif audio_latents is not None and self.audio_gen: T_audio = audio_latents.shape[2] hidden_audio = self.pack_audio_latents(audio_latents).to(hidden_gen.dtype) hidden_audio = self.audio2llm(hidden_audio) + self.audio_modality_embed @@ -1090,7 +1290,6 @@ def forward( hidden_states.device, hidden_gen.dtype, ) - # [B, T_vid+T_audio, hidden_size] hidden_gen = torch.cat([hidden_gen, hidden_audio], dim=1) cos_v, sin_v = self.cached_freqs_gen freqs_gen_combined = ( @@ -1137,16 +1336,27 @@ def forward( # --- Decode video velocity ------------------------------------------------ video_vel = self.unpatchify(self.llm2vae(hidden_gen[:, :T_vid_tokens]), T, H, W) - # --- Decode audio velocity (if requested) --------------------------------- + # --- Decode extra-modality velocity (action XOR audio; follows video) --- + extra_start = T_vid_tokens audio_vel = None if T_audio > 0 and audio_latents is not None and self.audio_gen: - # hidden_gen[:, T_vid_tokens:] → [B, T_audio, hidden_size] - # → llm2audio → [B, T_audio, audio_dim] → unpack → [B, audio_dim, T_audio] audio_vel = self.unpack_audio_latents( - self.llm2audio(hidden_gen[:, T_vid_tokens : T_vid_tokens + T_audio]) + self.llm2audio(hidden_gen[:, extra_start : extra_start + T_audio]) ) - return TransformerOutput(video=video_vel, image=video_vel, audio=audio_vel) + action_vel = None + if T_action > 0 and action_latents is not None and self.action_gen: + assert action_domain_ids_tensor is not None + action_vel = self.unpack_action( + self.action_proj_out( + hidden_gen[:, extra_start : extra_start + T_action], + action_domain_ids_tensor, + ) + ) + + return TransformerOutput( + video=video_vel, image=video_vel, audio=audio_vel, action=action_vel + ) def load_weights(self, weights: dict) -> None: """Load weights with key remapping from Cosmos3-Nano / Diffusers checkpoints. @@ -1158,8 +1368,7 @@ def load_weights(self, weights: dict) -> None: remapped = {} skip_prefixes = ( "lm_head.", - "action_modality_embed", - "action_proj_", + "action_pos_embed.", ) for key, value in weights.items(): @@ -1194,6 +1403,14 @@ def load_weights(self, weights: dict) -> None: remapped[k] = value continue + if k.startswith("action_modality_embed"): + remapped[k] = value + continue + + if k.startswith("action_proj_in.") or k.startswith("action_proj_out."): + remapped[k] = value + continue + if k.startswith("time_embedder.linear"): k = k.replace("time_embedder.linear_1.", "time_embedder.mlp.linear_1.") k = k.replace("time_embedder.linear_2.", "time_embedder.mlp.linear_2.") @@ -1326,6 +1543,11 @@ def post_load_weights(self) -> None: self.llm2audio.to(target_dtype) self.audio_modality_embed.data = self.audio_modality_embed.data.to(target_dtype) + if self.action_gen: + self.action_modality_embed.data = self.action_modality_embed.data.to(target_dtype) + self.action_proj_in.post_load_weights() + self.action_proj_out.post_load_weights() + for _, module in self.named_modules(): if isinstance(module, Linear) or isinstance(module, Qwen3VLTextRMSNorm): module.post_load_weights() diff --git a/tensorrt_llm/_torch/visual_gen/output.py b/tensorrt_llm/_torch/visual_gen/output.py index 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..8233042f346c 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline.py @@ -1042,8 +1042,11 @@ 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. + 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 +1177,7 @@ def denoise( ) if post_step_fn is not None: - 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 b61580d9f439..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``). 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)`` @@ -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. """ @@ -143,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 @@ -202,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 0660e661122f..099d40aa3a4c 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 @@ -121,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 @@ -138,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 @@ -152,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. @@ -166,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( @@ -209,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/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..100f42df710b --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py @@ -0,0 +1,280 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for Cosmos3 action sizing helpers (no checkpoint / GPU required). + +Run: + pytest tests/unittest/_torch/visual_gen/test_cosmos3_action.py -v +""" + +import numpy as np +import PIL.Image +import pytest +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, + get_domain_preset, + resolve_domain_action_config, +) + +pytestmark = pytest.mark.cosmos3 + + +class TestFindClosestTargetSize: + @pytest.mark.parametrize( + "input_h,input_w,action_resolution,expected", + [ + (480, 832, 480, (832, 480)), + (832, 480, 480, (480, 832)), + (512, 512, 480, (640, 640)), + (704, 1280, 704, (1280, 704)), + (256, 256, 256, (256, 256)), + (720, 1280, 720, (1280, 720)), + ], + ) + def test_picks_closest_aspect_bucket(self, input_h, input_w, action_resolution, expected): + assert find_closest_target_size(input_h, input_w, action_resolution) == expected + + def test_accepts_string_and_int_resolution_keys(self): + ref_h, ref_w = 480, 832 + assert find_closest_target_size(ref_h, ref_w, 480) == find_closest_target_size( + ref_h, ref_w, "480" + ) + + def test_unknown_resolution_raises(self): + with pytest.raises(ValueError, match="Unknown Cosmos3 action resolution"): + find_closest_target_size(480, 832, 1080) + + @pytest.mark.parametrize("action_resolution", sorted(VIDEO_RES_SIZE_INFO)) + def test_all_buckets_have_aspect_entries(self, action_resolution): + assert VIDEO_RES_SIZE_INFO[action_resolution] + + +class TestResolveActionSize: + @staticmethod + def _ref_image(width: int, height: int) -> PIL.Image.Image: + return PIL.Image.new("RGB", (width, height)) + + def test_explicit_height_and_width_are_unchanged(self): + ref = self._ref_image(832, 480) + assert resolve_action_size(400, 600, ref, 480) == (400, 600) + + def test_unset_height_and_width_use_action_resolution_bucket(self): + ref = self._ref_image(832, 480) + assert resolve_action_size(None, None, ref, 480) == (480, 832) + + def test_partial_height_fills_width_from_bucket(self): + ref = self._ref_image(832, 480) + assert resolve_action_size(400, None, ref, 480) == (400, 832) + + def test_partial_width_fills_height_from_bucket(self): + ref = self._ref_image(832, 480) + assert resolve_action_size(None, 600, ref, 480) == (480, 600) + + +class TestActionResolutionExtraParam: + def test_extra_param_spec_uses_action_resolution_key(self): + spec = COSMOS3_EXTRA_SPECS["action_resolution"] + assert spec.type == "Literal[256, 480, 704, 720]" + assert spec.default is None + + +class TestDomainActionPresets: + def test_bridge_preset_fills_missing_fields(self): + 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): + 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): + cfg = resolve_domain_action_config( + domain_name="bridge_orig_lerobot", + raw_action_dim=9, + ) + assert cfg["raw_action_dim"] == 9 + assert len(cfg["warnings"]) == 1 + assert "raw_action_dim=9" in cfg["warnings"][0] + + def test_action_fps_defaults_to_frame_rate(self): + 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): + 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): + preset = get_domain_preset("robomind-franka") + assert preset is not None + assert preset["raw_action_dim"] == 10 + + 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) + + +class TestActionReferenceImage: + def test_forward_dynamics_accepts_mp4_on_image_path(self, tmp_path, monkeypatch): + video_path = tmp_path / "clip.mp4" + video_path.write_bytes(b"fake") + expected = PIL.Image.new("RGB", (4, 2), "red") + + def _fake_read_video(path, pts_unit="sec"): + import torch + + tensor = torch.from_numpy(np.array(expected)).unsqueeze(0) + return tensor, None, {} + + monkeypatch.setattr("torchvision.io.read_video", _fake_read_video) + ref = action_reference_image( + action_mode="forward_dynamics", + image=str(video_path), + video=None, + ) + assert ref.size == expected.size + assert ref.getpixel((0, 0)) == (255, 0, 0) + + def test_policy_prefers_image_path_over_video(self, tmp_path): + image_path = tmp_path / "frame.png" + PIL.Image.new("RGB", (3, 3), "blue").save(image_path) + ref = action_reference_image( + action_mode="policy", + image=str(image_path), + video=str(tmp_path / "unused.mp4"), + ) + assert ref.getpixel((0, 0)) == (0, 0, 255) + + def test_policy_accepts_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) + 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 + + +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 c0329b71a060..422adaae8b6a 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,149 @@ 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, + ) + assert result.action_mode == "forward_dynamics" + assert result.domain_id == 7 + + def test_inverse_dynamics_smoke(self, cosmos3_pipeline): + _require_action_pipeline(cosmos3_pipeline) + image = _make_test_image().resize((self.ACTION_WIDTH, self.ACTION_HEIGHT)) + video = [image.copy() for _ in range(NUM_FRAMES)] + result = _run_forward( + cosmos3_pipeline, + image=None, + height=self.ACTION_HEIGHT, + width=self.ACTION_WIDTH, + num_frames=NUM_FRAMES, + guidance_scale=COSMOS3_ACTION_PARAMS["guidance_scale"], + action_mode="inverse_dynamics", + domain_name="bridge_orig_lerobot", + raw_action_dim=self.RAW_ACTION_DIM, + action_chunk_size=NUM_FRAMES, + video=video, + ) + _assert_valid_video( + result.video, + num_frames=NUM_FRAMES, + height=self.ACTION_HEIGHT, + width=self.ACTION_WIDTH, + ) + _assert_valid_action( + result.action, + raw_action_dim=self.RAW_ACTION_DIM, + chunk_size=NUM_FRAMES, + ) + assert 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) + 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..3b6cf23fcc04 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py @@ -397,6 +397,173 @@ 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") + 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): + 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_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 + 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.""" diff --git a/tests/unittest/_torch/visual_gen/test_tensor_payload.py b/tests/unittest/_torch/visual_gen/test_tensor_payload.py index 4ad482b886fd..3c9168ee92cb 100644 --- a/tests/unittest/_torch/visual_gen/test_tensor_payload.py +++ b/tests/unittest/_torch/visual_gen/test_tensor_payload.py @@ -52,6 +52,18 @@ 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, + ) + + class TestIsTensorFormat: def test_accepts_supported_tokens(self): for token in TENSOR_FORMATS: @@ -93,11 +105,77 @@ 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) + + 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: """A single path writes one logical output. Unbatched tensors and 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 e27f9c33e3b5..5ff8e6c7818f 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 @@ -323,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) # --------------------------------------------------------------------------- @@ -753,15 +773,19 @@ def test_encoding_not_top_level_reexport(): # --------------------------------------------------------------------------- -def test_pipeline_output_has_eight_fields(): - """PipelineOutput has the eight expected fields.""" +def test_pipeline_output_has_twelve_fields(): + """PipelineOutput has the twelve expected fields.""" field_names = {f.name for f in fields(PipelineOutput)} assert field_names == { "image", "video", "audio", + "action", "frame_rate", "audio_sample_rate", + "raw_action_dim", + "action_mode", + "domain_id", "pre_denoise", "denoise", "post_denoise",