diff --git a/examples/visual_gen/models/cosmos3/README.md b/examples/visual_gen/models/cosmos3/README.md index 880f4b536d8e..56aa7480c903 100644 --- a/examples/visual_gen/models/cosmos3/README.md +++ b/examples/visual_gen/models/cosmos3/README.md @@ -8,6 +8,7 @@ Cosmos3 supports the following generation modes from a single checkpoint: - **V2V** — video-conditioned video (`prompts/v2v.json`). Condition on a reference video via `--video_path` (a local MP4/AVI file). Only the first (or last, per `condition_video_keep`) `max(condition_video_latent_indexes) * 4 + 1` input frames condition the output (5 by default); the encoded bytes pass through and each worker decodes just that window on NVDEC (see [Media I/O dependencies](#media-io-dependencies)). Validated for Nano / Super only. - **Transfer** — control-video conditioning (`edge`/`blur`/`depth`/`seg`/`wsm` hints via `--extra_params`). The control constrains structure frame by frame; the prompt supplies appearance. `edge` and `blur` are auto-computed from `--video_path`; any hint also accepts a precomputed control clip (`{"edge": "control.mp4"}` — the example reads it and sends encoded bytes, the same contract as the `video` reference). Multiple hints compose (each adds a full control-token copy of the video sequence); long videos run chunked (93 frames/chunk, stitched on overlap frames) — but only past the first chunk, so raise `num_frames` above the pipeline default to generate one: it bounds how many frames are decoded from the inputs, and so how long the output can be. A single-hint request picks up that hint's tuned sampling preset — guidance scale, control guidance and flow shift — for any of those the request leaves unset; requests with several hints fall back to the generic video defaults. The active hint names are also appended to the prompt as a one-sentence control-adherence directive; pass `"emphasize_control_in_prompt": false` to suppress it for clean baselines or ablations. - **T2AV** — text-to-video with synchronized audio (`prompts/t2av.json` with `enable_audio: true`, or pass `--enable_audio`). Combine with a `vision_path` for image-conditioned audio-video (TI2AV). +- **Action** — policy / forward dynamics / inverse dynamics generation (pass `--action_mode`); `inverse_dynamics` reads its observation clip from `--video_path` (MP4/AVI, decoded on worker NVDEC like V2V). Action and audio generation are mutually exclusive. A predicted trajectory has no representation in a video container, so action runs are saved as `safetensors` or `pt`, keeping the rollout and the action tensor in one payload — over `trtllm-serve` the default `format=auto` selects that payload automatically, and an explicit `mp4`/`avi` is rejected. ## Checkpoints @@ -169,4 +170,36 @@ python cosmos3.py --model nvidia/Cosmos3-Edge \ 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) +# action_trajectory.json is a [T, D] list of lists; D is the embodiment's action +# width (9 for av) and a mismatch is rejected. +python cosmos3.py --model nvidia/Cosmos3-Nano \ + --prompt_file prompts/action_forward_dynamics.json \ + --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml \ + --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/observation_clip.mp4 \ + --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 9664940c71eb..18b1cba051c8 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3.py +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -31,6 +31,8 @@ from tensorrt_llm._torch.visual_gen.models.cosmos3.transfer import TRANSFER_HINT_KEYS _SCRIPT_DIR = Path(__file__).resolve().parent +_ACTION_MODES = ("policy", "forward_dynamics", "inverse_dynamics") +_TENSOR_OUTPUT_SUFFIXES = {".pt", ".safetensors"} DEFAULT_PROMPT_FILE = "prompts/t2v.json" DEFAULT_NEGATIVE_PROMPT_FILE = "cosmos3_negative_prompt.json" @@ -199,6 +201,95 @@ 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 (an .mp4 or .avi file).") + if args.raw_action_dim is None and args.domain_name is None and args.domain_id is None: + raise SystemExit(f"{mode} requires --raw_action_dim, --domain_name, or --domain_id.") + + +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, args: argparse.Namespace) -> None: + """Write the trajectory plus the request that produced it. + + The mode and embodiment are this script's own inputs, so they are read + from *args* rather than echoed back through the output schema. + """ + if output.action is None: + return + + 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": args.action_mode, + "domain_name": args.domain_name, + "domain_id": args.domain_id, + "raw_action_dim": action.shape[-1], + "shape": shape, + "dtype": str(action.dtype).replace("torch.", ""), + "data": action_data, + } + 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( @@ -274,11 +365,77 @@ def main(): ), ) 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="Reference video for V2V: a local MP4/AVI file (decoded on worker NVDEC)", + help=( + "Reference video (MP4/AVI, decoded on worker NVDEC): V2V conditioning, " + "or the observation clip for action 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( + "--view_point", + type=str, + default=None, + choices=["ego_view", "third_person_view", "wrist_view", "concat_view"], + help="Camera perspective for the action caption (default: ego_view).", + ) + parser.add_argument( + "--action_output_path", + type=str, + 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)" @@ -310,6 +467,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 @@ -336,6 +494,25 @@ 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.view_point is not None: + params.extra_params["view_point"] = args.view_point + if args.action_json is not None: + with open(args.action_json, encoding="utf-8") as f: + params.extra_params["action"] = json.load(f) if args.video_path is not None: params.extra_params["video"] = Path(args.video_path).read_bytes() if args.extra_params: @@ -352,8 +529,18 @@ 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, args) + if output.action is not None: + print(f"Saved action: {action_path}") + print(f"Action shape: {tuple(output.action.shape)}") + 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..fcc0a92feeeb --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py @@ -0,0 +1,584 @@ +# 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 + +import json +import math +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 + +from tensorrt_llm.logger import logger + +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, +} + +# Canonical unpadded action width per embodiment. Widths compose the Cosmos3 +# unified action representation from shared geometric blocks: a 9-D pose (3-D +# translation + 6-D rotation), a 1-D grasp state, and a 15-D fingertip state. +# One arm is 9 + 1 = 10; a dual-arm setup is 20; the AgiBot humanoid is +# 9 + 2 x (9 + 1) = 29; two-hand egocentric motion is 9 + 2 x (9 + 15) = 57. +# +# This is a property of the embodiment, not a tunable, so it is keyed by the +# real domain name rather than by the sampling presets in ``defaults.py`` (where +# several of these names share one preset). ``libero`` is absent on purpose: +# its width depends on the dataset's rotation space (7/10/13), so callers must +# pass ``raw_action_dim`` explicitly. +EMBODIMENT_TO_RAW_ACTION_DIM: dict[str, int] = { + "av": 9, + "camera_pose": 9, + "hand_pose": 57, + "pusht": 2, + "umi": 10, + "bridge_orig_lerobot": 10, + "droid_lerobot": 10, + "robomind-franka": 10, + "robomind-franka-dual": 20, + "robomind-ur": 10, + "galbot": 30, + "agibotworld": 29, + "agibot_gear_gripper": 29, + "agibot_gear_gripper_ext": 29, + "fractal": 10, +} + + +def resolve_raw_action_dim( + domain_name: Any = None, + domain_id: Any = None, +) -> int | None: + """Look up the canonical action width, or None when it cannot be determined. + + Resolves by name first. A bare ``domain_id`` is only usable when every + embodiment sharing that id agrees on a width (true for all current ids). + """ + if domain_name is not None and str(domain_name).strip(): + return EMBODIMENT_TO_RAW_ACTION_DIM.get(str(domain_name).strip().lower()) + + if domain_id is None: + return None + + widths = { + EMBODIMENT_TO_RAW_ACTION_DIM[name] + for name, mapped_id in EMBODIMENT_TO_DOMAIN_ID.items() + if mapped_id == int(domain_id) and name in EMBODIMENT_TO_RAW_ACTION_DIM + } + return widths.pop() if len(widths) == 1 else None + + +# Camera perspective -> framing sentence. The action model was trained on these +# exact sentences, so they are reproduced verbatim rather than paraphrased. +ACTION_VIEWPOINT_TEMPLATES: dict[str, str] = { + "ego_view": "This video is captured from a first-person perspective looking at the scene.", + "third_person_view": ( + "This video is captured from a third-person perspective looking towards the agent " + "from the front." + ), + "wrist_view": "This video is captured from a wrist-mounted camera.", + "concat_view": "This video contains concatenated views from multiple camera perspectives.", +} + +DEFAULT_ACTION_VIEW_POINT = "ego_view" + +# Canonical ``W,H`` labels; every action canvas is one of these bucket shapes. +ACTION_ASPECT_RATIO_LABELS = ("1,1", "4,3", "3,4", "16,9", "9,16") + + +def action_aspect_ratio_label(height: int, width: int) -> str: + """Closest canonical aspect label, e.g. 832x480 -> ``"16,9"``. + + Bucket sizes are only approximately their label (832/480 is 1.733, not + 1.778), so the label is matched by nearest ratio instead of reducing H/W. + """ + ratio = width / height if height > 0 else 1.0 + return min( + ACTION_ASPECT_RATIO_LABELS, + key=lambda label: abs(int(label.split(",")[0]) / int(label.split(",")[1]) - ratio), + ) + + +def build_action_json_prompt( + description: str, + *, + view_point: str | None, + num_frames: int, + frame_rate: float, + height: int, + width: int, +) -> str: + """Build the structured action caption the action model was trained on. + + Replaces the flat duration/resolution templates used by the video paths: the + JSON already carries duration, fps, resolution and aspect ratio. Key order is + part of the trained format and is preserved. + """ + duration_seconds = num_frames / frame_rate if frame_rate > 0 else 0.0 + if not math.isfinite(duration_seconds) or duration_seconds < 0: + duration_seconds = 0.0 + minutes, seconds = divmod(round(duration_seconds), 60) + + text = description.strip() + if text and not text.endswith((".", "!", "?")): + text = f"{text}." + + framing = ACTION_VIEWPOINT_TEMPLATES.get(view_point) if view_point is not None else None + if view_point is not None and framing is None: + logger.warning( + f"Unrecognized Cosmos3 action view_point={view_point!r}; expected one of " + f"{sorted(ACTION_VIEWPOINT_TEMPLATES)}. Dropping the cinematography.framing field." + ) + + prompt: dict[str, Any] = {} + if framing: + prompt["cinematography"] = {"framing": framing} + prompt["actions"] = [{"time": f"0:00-{minutes}:{seconds:02d}", "description": text}] + prompt["duration"] = f"{int(duration_seconds)}s" + prompt["fps"] = float(frame_rate) + prompt["resolution"] = {"H": int(height), "W": int(width)} + prompt["aspect_ratio"] = action_aspect_ratio_label(height, width) + return json.dumps(prompt) + + +VIDEO_RES_SIZE_INFO: dict[str, dict[str, tuple[int, int]]] = { + "256": { + "1,1": (256, 256), + "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}.") + # domain_id wins so unlisted embodiments stay reachable, but a caller + # that passes both and disagrees would otherwise silently get a + # trajectory in a different robot's dialect. + if domain_name is not None and str(domain_name).strip(): + key = str(domain_name).strip().lower() + named_id = EMBODIMENT_TO_DOMAIN_ID.get(key) + if named_id is not None and named_id != resolved: + raise ValueError( + f"Cosmos3 domain_id={resolved} contradicts domain_name={domain_name!r}, " + f"which maps to domain_id={named_id}. Pass only one, or make them agree." + ) + return resolved + + if domain_name is None or str(domain_name).strip() == "": + 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)}.") + if tensor.shape[0] == 0: + raise ValueError( + f"Cosmos3 action trajectory must have at least one timestep, got shape " + f"{tuple(tensor.shape)}." + ) + return tensor + + +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)): + # load_image, not PIL.Image.open: the bundled action prompts carry + # https:// frame references, and it also handles file:// and data: URIs. + from tensorrt_llm.inputs.utils import load_image + + return load_image(str(value), format="pil").convert("RGB") + if isinstance(value, PIL.Image.Image): + return value.convert("RGB") + raise TypeError( + f"Cosmos3 action preprocessing expected PIL image or image path, got {type(value)!r}." + ) + + +def resolve_action_size( + height: int | None, + width: int | None, + source_h: int, + source_w: int, + action_resolution: int, +) -> tuple[int, int]: + """Fill unset action H/W from the action resolution bucket; honor explicit values. + + The bucket is the canvas whose shape is closest to the source's, so the + reference only ever needs a modest pad to reach it. + """ + if height is not None and width is not None: + return height, width + target_w, target_h = find_closest_target_size(source_h, source_w, action_resolution) + return ( + height if height is not None else target_h, + width if width is not None else target_w, + ) + + +def action_reference_size( + *, + action_mode: str, + image: Any, + video: Any, +) -> tuple[int, int]: + """Source ``(height, width)`` of the reference, for choosing the canvas. + + Video references are encoded bytes, so their size comes from the container + header rather than a decode; images are measured directly. + """ + source = video if action_mode == ACTION_MODE_INVERSE_DYNAMICS else (image or video) + if source is None: + raise ValueError(f"Cosmos3 action_mode={action_mode!r} requires an image or video input.") + if isinstance(source, bytes): + from tensorrt_llm.media.decoding import video_stream_info + + info = video_stream_info(source) + if info is None: + raise ValueError( + f"Cosmos3 action_mode={action_mode!r} video reference could not be demuxed " + "(corrupt or not a supported container)." + ) + return info.height, info.width + reference = pil_to_rgb(source) + return reference.height, reference.width + + +def action_reference_frame_step(source_frame_rate: float | None, target_frame_rate: float) -> int: + """Source frames to advance per reference frame retained. + + An embodiment's frame rate is a property of what it was trained on, not of + the clip a caller happens to send: bridge learned one command per frame at + 5 Hz, so 200ms of gripper motion between frames. A 30 fps clip read + consecutively shows the model a sixth of that motion while the caption and + the mRoPE positions still claim 5 Hz, so the reference is thinned to match + -- every sixth frame here. + + A clip slower than the embodiment returns 1: selection can drop frames, + never invent them, and closing that gap needs interpolation rather than a + step. The caller is expected to say so rather than let it pass silently. + """ + if not source_frame_rate or not target_frame_rate: + return 1 + if source_frame_rate <= 0 or target_frame_rate <= 0: + return 1 + return max(1, round(source_frame_rate / target_frame_rate)) + + +def resize_and_pad_action_image( + image: PIL.Image.Image, target_h: int, target_w: int +) -> PIL.Image.Image: + 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 1661adfe364c..cf18b82bb498 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -18,8 +18,15 @@ """ from collections.abc import Mapping -from typing import Any, Dict, Iterable - +from typing import Any, Dict, Iterable, TypedDict + +from tensorrt_llm._torch.visual_gen.models.cosmos3.action import ( + COSMOS3_ACTION_RESOLUTIONS, + DEFAULT_ACTION_VIEW_POINT, + EMBODIMENT_TO_DOMAIN_ID, + normalize_action_resolution, + resolve_raw_action_dim, +) from tensorrt_llm._torch.visual_gen.pipeline import ExtraParamSchema from tensorrt_llm.inputs.media_io import sniff_media_kind @@ -310,11 +317,275 @@ def _validate_non_negative_frames(value: Any) -> None: ("nemotron_dense", "image"): COSMOS3_EDGE_T2I_PARAMS, } +# Action's table carries the sampling recipe only; the canvas, clip length and +# frame rate resolve from the embodiment preset in the pipeline, and anything +# missing here falls back to the family's video table inside +# _resolve_generation_params. Same recipe for both families. + + # Families without an entry get no envelope advisory. COSMOS3_ENVELOPES: Dict = { "nemotron_dense": COSMOS3_EDGE_ENVELOPE, } +COSMOS3_ACTION_PARAMS = { + "action_chunk_size": 16, + "num_inference_steps": 30, + "guidance_scale": 1.0, + "frame_rate": 24.0, +} + +COSMOS3_GENERATION_DEFAULTS[("qwen3", "action")] = COSMOS3_ACTION_PARAMS +COSMOS3_GENERATION_DEFAULTS[("nemotron_dense", "action")] = COSMOS3_ACTION_PARAMS + + +class Cosmos3DomainPreset(TypedDict, total=False): + """Recommended action sampling settings for a trained embodiment. + + Sampling settings only — the embodiment's action width lives in + ``action.EMBODIMENT_TO_RAW_ACTION_DIM``, keyed by the unaliased domain name. + """ + + action_chunk_size: int + action_resolution: int + frame_rate: float + + +# Training-aligned defaults, mirroring the Cosmos3 omni ``action_*.json`` inputs +# (bridge, av, droid, libero, ...) where those exist. +COSMOS3_DOMAIN_PRESETS: dict[str, Cosmos3DomainPreset] = { + # WidowX bridge. + "bridge_orig_lerobot": { + "action_chunk_size": 16, + "action_resolution": 480, + "frame_rate": 5.0, + }, + # Autonomous-vehicle steering/throttle; longer action horizon. + "av": { + "action_chunk_size": 60, + "action_resolution": 480, + "frame_rate": 10.0, + }, + # 6-DoF camera pose + shutter; matches AV-style horizon. + "camera_pose": { + "action_chunk_size": 60, + "action_resolution": 480, + "frame_rate": 30.0, + }, + # Franka single-arm tabletop; same domain_id as robomind-franka. + "droid_lerobot": { + "action_chunk_size": 16, + "action_resolution": 480, + "frame_rate": 15.0, + }, + # LIBERO sim single-arm; lower action resolution bucket. + "libero": { + "action_chunk_size": 16, + "action_resolution": 256, + "frame_rate": 10.0, + }, + # MANO hand pose. + "hand_pose": { + "action_chunk_size": 16, + "action_resolution": 480, + "frame_rate": 24.0, + }, + # AgiBot humanoid; shared domain_id with agibot_gear_gripper*. + "agibotworld": { + "action_chunk_size": 16, + "action_resolution": 480, + "frame_rate": 10.0, + }, + # Google Robot (RT-1 / fractal) single-arm. + "fractal": { + "action_chunk_size": 16, + "action_resolution": 480, + "frame_rate": 5.0, + }, + # 2-D planar push task. + "pusht": { + "action_chunk_size": 16, + "action_resolution": 256, + "frame_rate": 10.0, + }, + # UMI handheld gripper setup. + "umi": { + "action_chunk_size": 16, + "action_resolution": 480, + "frame_rate": 10.0, + }, +} + +# Map alias domain_name keys to a canonical preset entry. These share *sampling* +# settings only; each alias keeps its own action width (e.g. robomind-franka-dual +# is 20-D and galbot is 30-D, unlike the presets they borrow here). +COSMOS3_DOMAIN_PRESET_ALIASES: dict[str, str] = { + "robomind-franka": "droid_lerobot", + "robomind-franka-dual": "droid_lerobot", + "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, +) -> dict[str, Any]: + """Merge user action params with domain presets and generic fallbacks. + + A recognized ``domain_name`` (or a uniquely mapped ``domain_id``) fills + whichever of ``action_chunk_size``, ``action_resolution`` and ``frame_rate`` + the caller left unset; an explicit value wins but is reported in + ``warnings`` when it differs from the preset. ``num_frames`` is derived as + ``action_chunk_size + 1`` and is never a preset field. + """ + preset_key = canonical_domain_preset_key(domain_name, domain_id) + preset = COSMOS3_DOMAIN_PRESETS.get(preset_key) if preset_key else None + warnings: list[str] = [] + + 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 + + # The action width is canonical per embodiment, so it comes from the + # embodiment table rather than the (alias-shared) sampling preset. + canonical_raw_action_dim = resolve_raw_action_dim(domain_name=domain_name, domain_id=domain_id) + if raw_action_dim is not None: + if canonical_raw_action_dim is not None and int(raw_action_dim) != canonical_raw_action_dim: + warnings.append( + f"Cosmos3 raw_action_dim={raw_action_dim} differs from the canonical width " + f"{canonical_raw_action_dim} for domain_name={domain_name!r}." + ) + resolved_raw_action_dim = raw_action_dim + else: + resolved_raw_action_dim = canonical_raw_action_dim + if domain_requested and canonical_raw_action_dim is None: + warnings.append( + "Cosmos3 has no canonical action width for " + f"domain_name={domain_name!r}, domain_id={domain_id!r}; " + "pass raw_action_dim explicitly for policy/inverse_dynamics." + ) + resolved_chunk = _resolve_field( + "action_chunk_size", + action_chunk_size, + 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"], + ) + # Always derived: an action clip is the chunk plus its initial frame. Both + # references fix this, and diffusers rejects a caller-supplied num_frames + # for action runs outright, so a preset must not pin it independently of + # an overridden action_chunk_size. + resolved_num_frames = int(resolved_chunk) + 1 + resolved_action_fps = ( + float(action_fps) if action_fps is not None else float(resolved_frame_rate) + ) + 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_V2V_DEFAULT_FLOW_SHIFT = 10.0 COSMOS3_EXTRA_SPECS: Dict[str, ExtraParamSchema] = { @@ -387,6 +658,77 @@ def _validate_non_negative_frames(value: Any) -> None: ), validator=_validate_video_reference, ), + "action_mode": ExtraParamSchema( + type="Literal['policy', 'forward_dynamics', 'inverse_dynamics']", + default=None, + description=( + "Action generation mode: policy, forward_dynamics, or inverse_dynamics. " + "The predicted trajectory is not representable in a video container, so " + "an action request is served as a tensor payload." + ), + requires_tensor_output=True, + ), + "domain_name": ExtraParamSchema( + type="str", + 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). " + "Resolved from the embodiment when omitted; required for domains with no " + "canonical width (libero)." + ), + ), + "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." + ), + # No range: the buckets are not an interval, and validation stops at the + # literal check anyway. + ), + "view_point": ExtraParamSchema( + type="Literal['ego_view', 'third_person_view', 'wrist_view', 'concat_view']", + default=DEFAULT_ACTION_VIEW_POINT, + description=( + "Camera perspective for action generation. Fills the trained action caption's " + "cinematography.framing field." + ), + ), + "action_fps": ExtraParamSchema( + type="float", + default=None, + description=( + "Action-token temporal rate for mRoPE (Hz). Defaults to frame_rate when omitted." + ), + ), # Transfer "edge": ExtraParamSchema( type="bool_or_bytes_or_dict", 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 165135a7ae8b..494a415ad99a 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -48,6 +48,21 @@ def tqdm(iterable, **kwargs): from tensorrt_llm.logger import logger from tensorrt_llm.media.decoding import decode_video_reference_window, video_stream_info +from .action import ( + ACTION_MODE_INVERSE_DYNAMICS, + DEFAULT_ACTION_VIEW_POINT, + action_reference_frame_step, + action_reference_size, + action_start_frame_offset, + build_action_json_prompt, + build_vision_condition_mask, + normalize_action_mode, + pil_to_rgb, + prepare_action_latents, + resize_and_pad_action_image, + resolve_action_size, + resolve_domain_id, +) from .defaults import ( COSMOS3_720P_PARAMS, COSMOS3_ENVELOPES, @@ -56,6 +71,7 @@ def tqdm(iterable, **kwargs): COSMOS3_V2V_DEFAULT_FLOW_SHIFT, _normalize_condition_video_keep, _normalize_condition_video_latent_indexes, + resolve_domain_action_config, ) from .guardrails import check_video_safety, download_guardrail_checkpoint from .negative_prompt import COSMOS3_VIDEO_NEGATIVE_PROMPT @@ -262,11 +278,9 @@ def __init__(self, pipeline_config): self.audio_gen = True if getattr(primary_pretrained_config, "action_gen", False): - logger.info( - "Checkpoint declares action weights; action generation is not supported " - "by this pipeline (weights are skipped)." - ) + logger.info("Initializing Cosmos3OmniMoTPipeline with action generation.") self.has_action_weights = True + self.action_gen = True super().__init__(pipeline_config) @@ -426,6 +440,11 @@ def load_standard_components( self.audio_scheduler = type(self.scheduler).from_config(self.scheduler.config) self._base_audio_scheduler = self.audio_scheduler + if self.action_gen: + # Action uses its own scheduler for the same reason as audio. + self.action_scheduler = type(self.scheduler).from_config(self.scheduler.config) + self._base_action_scheduler = self.action_scheduler + # 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 @@ -539,6 +558,8 @@ def _scheduler_for( """ if stream == "audio": base = getattr(self, "_base_audio_scheduler", None) or self.audio_scheduler + elif stream == "action": + base = getattr(self, "_base_action_scheduler", None) or self.action_scheduler else: base = getattr(self, "_base_scheduler", None) or self.scheduler # Only configurations this checkpoint can resolve to on its own are @@ -637,55 +658,70 @@ def as_given(field_name): return value if field_name in specified else None video = extra_params.get("video") # encoded MP4/AVI bytes (the extra-param contract) - # Source-derived sizes go in as non-None, so they win over the mode - # table: _resolve_generation_params only fills what is still None. - height, width = as_given("height"), as_given("width") - frame_rate = req.params.frame_rate - wants_source_size = "height" not in specified and "width" not in specified - # Timing resolves as a unit. A caller who pinned `num_frames` -- directly, - # or via `seconds`, which the serving layer already converted at the - # default rate -- has fixed the output duration; adopting a different - # frame rate underneath that would silently stretch or compress it. - wants_source_fps = "frame_rate" not in specified and "num_frames" not in specified - if wants_source_size or wants_source_fps: - source = video - if source is None and transfer_config is not None: - # Transfer can run on precomputed controls alone; then the - # control clip is what defines the structure to match. - source = next( - (hint.control for hint in transfer_config.ordered_hints if hint.control), - None, - ) - # One header read serves both: no GPU, no frame decoded. - info = video_stream_info(source) if source is not None else None - if info is not None: - if wants_source_size: - # Follow the reference's aspect, so a portrait or square - # source is not center-cropped into the default landscape - # bucket. Skipped when either dimension was named: that - # states an intent, and overriding half of it would be - # worse than leaving it alone. - width, height = find_closest_target_size(info.height, info.width, 720) - if wants_source_fps and info.frame_rate is not None: - # Emitting an 8 fps source at the 24 fps default would play - # it back at 3x speed and misreport its duration to the - # text conditioning. - frame_rate = info.frame_rate - if self.rank == 0: - logger.info( - f"Cosmos3 following the source: {width}x{height} (WxH) @ {frame_rate} fps" + is_action = extra_params.get("action_mode") is not None + if is_action: + # Action resolves its whole recipe in forward() -- the canvas from + # the resolution bucket, the frame rate from the embodiment, the + # sampling recipe from the action table -- so a caller value passes + # through and an untouched field arrives as None. The embodiment's + # rate is a trained property, so an action request does not follow + # the source clip the way a video request below does. + height = as_given("height") + width = as_given("width") + num_inference_steps = as_given("num_inference_steps") + guidance_scale = as_given("guidance_scale") + frame_rate = as_given("frame_rate") + else: + # Source-derived sizes go in as non-None, so they win over the mode + # table: _resolve_generation_params only fills what is still None. + height, width = as_given("height"), as_given("width") + frame_rate = req.params.frame_rate + wants_source_size = "height" not in specified and "width" not in specified + # Timing resolves as a unit. A caller who pinned `num_frames` -- directly, + # or via `seconds`, which the serving layer already converted at the + # default rate -- has fixed the output duration; adopting a different + # frame rate underneath that would silently stretch or compress it. + wants_source_fps = "frame_rate" not in specified and "num_frames" not in specified + if wants_source_size or wants_source_fps: + source = video + if source is None and transfer_config is not None: + # Transfer can run on precomputed controls alone; then the + # control clip is what defines the structure to match. + source = next( + (hint.control for hint in transfer_config.ordered_hints if hint.control), + None, ) - resolved = self._resolve_generation_params( - "image" if is_t2i else "video", - height=height, - width=width, - num_inference_steps=as_given("num_inference_steps"), - guidance_scale=as_given("guidance_scale"), - ) - height = resolved["height"] - width = resolved["width"] - num_inference_steps = resolved["num_inference_steps"] - guidance_scale = resolved["guidance_scale"] + # One header read serves both: no GPU, no frame decoded. + info = video_stream_info(source) if source is not None else None + if info is not None: + if wants_source_size: + # Follow the reference's aspect, so a portrait or square + # source is not center-cropped into the default landscape + # bucket. Skipped when either dimension was named: that + # states an intent, and overriding half of it would be + # worse than leaving it alone. + width, height = find_closest_target_size(info.height, info.width, 720) + if wants_source_fps and info.frame_rate is not None: + # Emitting an 8 fps source at the 24 fps default would play + # it back at 3x speed and misreport its duration to the + # text conditioning. + frame_rate = info.frame_rate + if self.rank == 0: + logger.info( + f"Cosmos3 following the source: " + f"{width}x{height} (WxH) @ {frame_rate} fps" + ) + resolved = self._resolve_generation_params( + "image" if is_t2i else "video", + height=height, + width=width, + num_inference_steps=as_given("num_inference_steps"), + guidance_scale=as_given("guidance_scale"), + ) + height = resolved["height"] + width = resolved["width"] + num_inference_steps = resolved["num_inference_steps"] + guidance_scale = resolved["guidance_scale"] return self.forward( prompt=req.prompt, @@ -716,6 +752,15 @@ def as_given(field_name): condition_video_latent_indexes=extra_params.get("condition_video_latent_indexes"), condition_video_keep=extra_params.get("condition_video_keep"), flow_shift=extra_params.get("flow_shift"), + 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"), + action_fps=extra_params.get("action_fps"), + view_point=extra_params.get("view_point", DEFAULT_ACTION_VIEW_POINT), transfer_config=transfer_config, ) @@ -987,11 +1032,11 @@ def _conditioning_anchor_post_step(self, image_latent: Optional[torch.Tensor]): if not self.sampling.is_distilled or image_latent is None: return None - def post_step_fn(latents: torch.Tensor) -> torch.Tensor: + def post_step_fn(latents: torch.Tensor, extra_stream_latents): # In-place: writes one latent frame, no full-tensor copies. _assert_anchor_matches(image_latent, latents) latents[:, :, 0:1] = image_latent - return latents + return latents, extra_stream_latents return post_step_fn @@ -1157,6 +1202,98 @@ def _prepare_latents_v2v( velocity_mask = 1.0 - condition_mask return latents, velocity_mask, condition_latents + # ========================================================================= + # 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_first_frame( + self, image: Any, video: Any, target_h: int, target_w: int + ) -> torch.Tensor: + """Conditioning frame for policy / forward_dynamics as ``[1, 3, H, W]``. + + Either source is accepted: an image goes through PIL, video bytes take + frame 0 off NVDEC. Both land on the padded canvas, so the two entry + points produce the same conditioning for the same picture. + """ + if image is not None: + return self._preprocess_action_image(pil_to_rgb(image), target_h, target_w) + if not isinstance(video, bytes): + raise ValueError( + "Cosmos3 action conditioning requires an image or encoded MP4/AVI " + f"bytes, got {type(video).__name__}." + ) + frames_u8 = decode_video_reference_window( + video, + first_frame=0, + last_frame=0, + target_h=target_h, + target_w=target_w, + device=self.device, + resize="fit", + ) + return self._condition_frames_to_video_tensor(frames_u8).squeeze(2) + + def _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(self.transformer.action_dim), + generator=generator, + device=self.device, + dtype=self.dtype, + action_input=action_input, + ) + # ========================================================================= # Transfer # ========================================================================= @@ -1215,10 +1352,19 @@ def forward( use_guardrails: bool = COSMOS3_EXTRA_SPECS["use_guardrails"].default, enable_audio: bool = COSMOS3_EXTRA_SPECS["enable_audio"].default, output_type: str = COSMOS3_EXTRA_SPECS["output_type"].default, - video: bytes | None = None, # encoded MP4/AVI reference (V2V) + video: bytes | None = None, # encoded MP4/AVI reference (V2V and action) condition_video_latent_indexes: Iterable[int] | None = None, condition_video_keep: str | None = None, flow_shift: Optional[float] = None, + 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, + view_point: Optional[str] = DEFAULT_ACTION_VIEW_POINT, transfer_config: Optional[Cosmos3TransferConfig] = None, ): """Run one generation. ``infer()`` is the resolved entry point. @@ -1240,6 +1386,16 @@ def forward( use_guardrails = use_guardrails and not TRTLLM_DISABLE_COSMOS3_GUARDRAILS + normalized_action_mode = normalize_action_mode(action_mode) + do_action = normalized_action_mode is not None + if do_action and not self.action_gen: + raise ValueError( + "Cosmos3 action generation was requested, but this checkpoint " + "does not enable action_gen." + ) + if do_action and enable_audio: + raise ValueError("Cosmos3 does not support joint action and audio generation.") + # Text-to-image mode: same checkpoint/forward path as T2V, but a single # latent frame, image-flavored prompt templates, flow_shift=3.0, a CFG # guidance interval, and an image (rather than video) output. @@ -1249,34 +1405,51 @@ def forward( is_t2i = output_type == "image" mode_params = self._mode_params(output_type) - resolved = self._resolve_generation_params( - output_type, - height=height, - width=width, - num_frames=num_frames, - num_inference_steps=num_inference_steps, - guidance_scale=guidance_scale, - max_sequence_length=max_sequence_length, - frame_rate=frame_rate, - ) - height = resolved["height"] - width = resolved["width"] - num_frames = resolved["num_frames"] - num_inference_steps = resolved["num_inference_steps"] - guidance_scale = resolved["guidance_scale"] - max_sequence_length = resolved["max_sequence_length"] - frame_rate = resolved["frame_rate"] + if do_action: + # The embodiment resolves the canvas, clip length and frame rate + # below; only the sampling recipe and the text budget come from + # tables (distilled overrides still win inside the resolver). + resolved = self._resolve_generation_params( + "action", + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + max_sequence_length=max_sequence_length, + ) + num_inference_steps = resolved["num_inference_steps"] + guidance_scale = resolved["guidance_scale"] + max_sequence_length = resolved["max_sequence_length"] + else: + resolved = self._resolve_generation_params( + output_type, + height=height, + width=width, + num_frames=num_frames, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + max_sequence_length=max_sequence_length, + frame_rate=frame_rate, + ) + height = resolved["height"] + width = resolved["width"] + num_frames = resolved["num_frames"] + num_inference_steps = resolved["num_inference_steps"] + guidance_scale = resolved["guidance_scale"] + max_sequence_length = resolved["max_sequence_length"] + frame_rate = resolved["frame_rate"] self.sampling.validate_request(num_inference_steps, guidance_scale) - self._log_envelope_advisory( - is_t2i=is_t2i, - height=height, - width=width, - num_frames=num_frames, - frame_rate=frame_rate, - max_sequence_length=max_sequence_length, - ) + # Skipped for action: its dims resolve from the embodiment below, and + # the envelope is a video/image model-card claim. + if not do_action: + self._log_envelope_advisory( + is_t2i=is_t2i, + height=height, + width=width, + num_frames=num_frames, + frame_rate=frame_rate, + max_sequence_length=max_sequence_length, + ) if image is not None and video is not None: raise ValueError( @@ -1287,7 +1460,11 @@ def forward( raise ValueError( "Cosmos3 video-to-video generation is supported only for video outputs." ) - is_v2v = video is not None and not is_t2i + # Action reads its reference through the same `video` bytes, but it is + # not V2V: the reference is an observation, not a clip to continue. Left + # in, an action request's prompt would depend on whether the caller + # passed the same frame as an image or as a one-frame clip. + is_v2v = video is not None and not is_t2i and not do_action if use_system_prompt is None: # V2V always wants it; otherwise the checkpoint declares the default. # Transfer opts out for reference parity (vllm-omni @@ -1315,17 +1492,55 @@ def forward( ) 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'.") num_frames = 1 # T2I force-disables audio instead of rejecting it, so an image # request never trips the audio-weight presence check below. enable_audio = False guidance_interval = mode_params["guidance_interval"] + if do_action: + # num_frames is derived from the action chunk, never taken from the + # request: both references fix it at chunk_size + 1 (diffusers + # rejects a caller-supplied num_frames for action runs outright). + 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, + ) + 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"] + enable_audio = False + # Flow shift is a mode table fact unless the request overrides it, and # V2V additionally wants the uniform sigma grid. Both streams take the # same knobs so video and audio never step on different schedules. @@ -1342,11 +1557,17 @@ def forward( if transfer_config is None: # Transfer applies the hint's own shift inside _forward_transfer, so # rebuilding the schedulers here would only be undone a few lines later. + # Action never arrives with a transfer_config (the two are mutually + # exclusive), so its stream is always rebuilt here. self.scheduler = self._scheduler_for(target_shift, target_karras) if getattr(self, "audio_scheduler", None) is not None: self.audio_scheduler = self._scheduler_for( target_shift, target_karras, stream="audio" ) + if getattr(self, "action_scheduler", None) is not None: + self.action_scheduler = self._scheduler_for( + target_shift, target_karras, stream="action" + ) if self.rank == 0: logger.info( @@ -1365,6 +1586,56 @@ def forward( "or use an audio-capable Cosmos3 checkpoint." ) + if resolved_action_fps is None: + resolved_action_fps = frame_rate + + action_source_h = action_source_w = None + if do_action: + if isinstance(image, torch.Tensor) or isinstance(video, torch.Tensor): + raise ValueError( + "Cosmos3 action generation does not support tensor image/video inputs; " + "pass a PIL image or image path, or encoded MP4/AVI video bytes." + ) + # Header probe for video bytes, direct measure for an image: the + # canvas is the bucket closest to the source's shape, so the size + # has to be known before anything is decoded at it. + # + # This is the first per-rank read of the reference, so it converges + # like the decode below: a missing file or an unreachable URL on one + # rank must not leave the others walking into the collectives. + probe_error: Optional[Exception] = None + try: + if image is not None and not isinstance(image, PIL.Image.Image): + # Resolve once: the bundled prompts point at https frames, + # and the probe and the conditioning frame both read it. + image = pil_to_rgb(image) + action_source_h, action_source_w = action_reference_size( + action_mode=normalized_action_mode, + image=image, + video=video, + ) + except Exception as exc: + probe_error = exc + synchronize_media_prepare_status(probe_error) + height, width = resolve_action_size( + height, width, action_source_h, action_source_w, action_resolution + ) + + 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"source={action_source_w}x{action_source_h} " + f"(aspect {action_source_w / action_source_h:.3f})" + ) + if isinstance(prompt, str): prompt = [prompt] batch_size = len(prompt) @@ -1427,47 +1698,65 @@ def forward( if negative_prompt is None: negative_prompt = default_negative_prompt(output_type) - # Positive prompt: forward duration/resolution templates. T2I has no - # duration concept (single image) and uses the image-flavored - # resolution template. - use_duration_template = use_duration_template and not is_t2i - dur_tmpl = COSMOS3_DURATION_TEMPLATE if use_duration_template else None - if use_resolution_template: - res_tmpl = ( - COSMOS3_IMAGE_RESOLUTION_TEMPLATE if is_t2i else COSMOS3_DEFAULT_RESOLUTION_TEMPLATE - ) + if do_action: + # Action checkpoints were trained on a structured JSON caption that + # already carries duration/fps/resolution/aspect_ratio, so the flat + # templates are skipped here and the negative prompt stays verbatim. + prompt = [ + build_action_json_prompt( + p, + view_point=view_point, + num_frames=num_frames, + frame_rate=frame_rate, + height=height, + width=width, + ) + for p in prompt + ] else: - res_tmpl = None - - # Negative prompt: mirror positive metadata (cosmos-framework CLI default - # when ``negative_prompt_keep_metadata`` promotes mode to ``same``). - # Always the plain-text templates, never the JSON field injection the - # positive branch uses: the reference appends these sentences to the - # negative prompt whether or not it is a JSON object, so a JSON negative - # prompt ends up as the serialized object followed by the sentences. - negative_prompt = self._apply_metadata_templates( - negative_prompt, - height=height, - width=width, - num_frames=num_frames, - frame_rate=frame_rate, - duration_template=dur_tmpl, - resolution_template=res_tmpl, - force_duration_template=False, - ) - - prompt = [ - self._format_prompt_with_metadata( - p, + # Positive prompt: forward duration/resolution templates. T2I has no + # duration concept (single image) and uses the image-flavored + # resolution template. + use_duration_template = use_duration_template and not is_t2i + dur_tmpl = COSMOS3_DURATION_TEMPLATE if use_duration_template else None + if use_resolution_template: + res_tmpl = ( + COSMOS3_IMAGE_RESOLUTION_TEMPLATE + if is_t2i + else COSMOS3_DEFAULT_RESOLUTION_TEMPLATE + ) + else: + res_tmpl = None + + # Negative prompt: mirror positive metadata (cosmos-framework CLI default + # when ``negative_prompt_keep_metadata`` promotes mode to ``same``). + # Always the plain-text templates, never the JSON field injection the + # positive branch uses: the reference appends these sentences to the + # negative prompt whether or not it is a JSON object, so a JSON negative + # prompt ends up as the serialized object followed by the sentences. + negative_prompt = self._apply_metadata_templates( + negative_prompt, height=height, width=width, num_frames=num_frames, frame_rate=frame_rate, duration_template=dur_tmpl, resolution_template=res_tmpl, + force_duration_template=False, ) - for p in prompt - ] + + prompt = [ + self._format_prompt_with_metadata( + p, + height=height, + width=width, + num_frames=num_frames, + frame_rate=frame_rate, + duration_template=dur_tmpl, + resolution_template=res_tmpl, + ) + for p in prompt + ] logger.info(f"Prompt with metadata: '{prompt}'") prompt = prompt[0] @@ -1486,8 +1775,129 @@ def forward( condition_latents = None image_latent = None velocity_mask = 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 + + 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, + ) + num_domains = getattr(self.transformer, "num_embodiment_domains", None) + if num_domains is not None and not 0 <= action_domain_id < num_domains: + raise ValueError( + f"Cosmos3 action domain_id must be in [0, {num_domains}), " + f"got {action_domain_id}." + ) + action_frame_offset = action_start_frame_offset( + normalized_action_mode, action_chunk_size, num_frames + ) - if image is not None: + if normalized_action_mode == ACTION_MODE_INVERSE_DYNAMICS: + if not isinstance(video, bytes): + raise ValueError( + "Cosmos3 inverse_dynamics requires encoded MP4/AVI bytes " + f"(the 'video' extra-param contract), got {type(video).__name__}." + ) + prepare_error: Optional[Exception] = None + try: + source_info = video_stream_info(video) + source_frame_rate = source_info.frame_rate if source_info else None + frame_step = action_reference_frame_step(source_frame_rate, frame_rate) + if self.rank == 0: + if frame_step > 1: + logger.info( + f"Cosmos3 action reference: {source_frame_rate} fps source " + f"thinned to {frame_rate} fps, keeping every {frame_step} " + f"frames of {(num_frames - 1) * frame_step + 1}" + ) + elif source_frame_rate is not None and source_frame_rate < frame_rate: + logger.warning( + f"Cosmos3 action reference is {source_frame_rate} fps but " + f"{normalized_action_mode} expects {frame_rate} fps: frames are " + "further apart than the model was trained on and cannot be " + "thinned to match. Re-encode the reference at the higher rate, " + "or pass frame_rate explicitly to accept this spacing." + ) + # "fit" rather than the V2V default: an action reference is + # padded to the canvas, never cropped to it, because the + # gripper and target sit at the frame edge. + frames_u8 = decode_video_reference_window( + video, + first_frame=0, + last_frame=(num_frames - 1) * frame_step, + target_h=height, + target_w=width, + device=self.device, + resize="fit", + frame_step=frame_step, + ) + if frames_u8.shape[0] < num_frames: + raise ValueError( + f"Cosmos3 inverse_dynamics requires {num_frames} frames at " + f"{frame_rate} fps; a {source_frame_rate} fps reference supplies " + f"{frames_u8.shape[0]} once thinned by {frame_step} " + f"({(num_frames - 1) * frame_step + 1} source frames needed)." + ) + video_tensor = self._condition_frames_to_video_tensor(frames_u8) + del frames_u8 + latents, velocity_mask, condition_latents = self._prepare_latents_action_video( + video_tensor, + normalized_action_mode, + num_frames, + generator, + ) + del video_tensor + except Exception as exc: + prepare_error = exc + # Every rank decodes independently; converge before the + # transformer's collectives so a failure cannot hang the job. + synchronize_media_prepare_status(prepare_error) + else: + prepare_error = None + try: + image_tensor = self._preprocess_action_first_frame(image, video, height, width) + if image_tensor.ndim == 4: + video_tensor = ( + image_tensor.unsqueeze(2) + .expand(-1, -1, num_frames, -1, -1) + .contiguous() + ) + else: + video_tensor = image_tensor + latents, velocity_mask, condition_latents = self._prepare_latents_action_video( + video_tensor, + normalized_action_mode, + num_frames, + generator, + ) + except Exception as exc: + prepare_error = exc + synchronize_media_prepare_status(prepare_error) + 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: prepare_error: Optional[Exception] = None try: if isinstance(image, str): @@ -1591,6 +2001,10 @@ def forward( # 3. Set up scheduler self.sampling.set_timesteps(self.scheduler, num_inference_steps, device=self.device) + if do_action: + self.sampling.set_timesteps( + self.action_scheduler, 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") @@ -1612,6 +2026,12 @@ def forward( ) # 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, @@ -1626,6 +2046,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, @@ -1637,25 +2067,56 @@ 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): - # V2V only: re-impose the clean condition latents after every - # scheduler step. I2V deliberately keeps its pre-existing behavior - # (velocity mask during the loop, one write-back after it) so this - # PR does not alter I2V denoising; per-step anchoring for - # stochastic distilled schedulers belongs to the distilled work. - return velocity_mask * step_latents + (1.0 - velocity_mask) * condition_latents + def post_step_fn(step_latents, step_extra_stream_latents): + # V2V (and action's clean vision frames): re-impose the condition + # latents after every scheduler step. I2V deliberately keeps its + # pre-existing behavior -- velocity mask during the loop, one + # write-back after it -- so this does not alter I2V denoising. + if velocity_mask is not None and condition_latents is not None: + step_latents = ( + velocity_mask * step_latents + (1.0 - velocity_mask) * condition_latents + ) + 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 @@ -1670,8 +2131,11 @@ def post_step_fn(step_latents): # 6. Denoise timer.mark_denoise_start() extra_streams = None - if do_audio: + 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. should_pin_condition_latents = condition_latents is not None and velocity_mask is not None denoise_result = self.denoise( latents=latents, @@ -1683,12 +2147,12 @@ def post_step_fn(step_latents): extra_cfg_tensors=extra_cfg_tensors, extra_streams=extra_streams, guidance_interval=guidance_interval, - # V2V pins the conditioning latents; distilled I2V re-anchors the - # conditioning frame. A request carries an image or a video, never - # both, so at most one of these applies. + # V2V and action pin the conditioning latents; distilled I2V + # re-anchors the conditioning frame. A request carries an image or + # a video, never both, so at most one of these applies. post_step_fn=( post_step_fn - if should_pin_condition_latents + if (do_action or should_pin_condition_latents) else self._conditioning_anchor_post_step(image_latent) ), scheduler_step_kwargs=self.sampling.scheduler_step_kwargs(generator), @@ -1697,6 +2161,8 @@ def post_step_fn(step_latents): if extra_streams is not None: latents, extra_latents = denoise_result audio_latents = extra_latents.get("audio") + if do_action: + action_latents = extra_latents.get("action") else: latents = denoise_result audio_latents = None @@ -1745,6 +2211,12 @@ def post_step_fn(step_latents): audio_sample_rate=self.audio_tokenizer.model_config["sampling_rate"] if waveform is not None else None, + # Sliced to the embodiment's real width, so the trailing dim is + # raw_action_dim; the mode and embodiment are the caller's own + # request and are not echoed back. + action=action_latents[:, :, :resolved_raw_action_dim].float().cpu() + if do_action and action_latents is not None + else None, ) ) 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 ebe83e605626..4eb2ebd0dbd1 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -271,6 +271,8 @@ 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, + base_temporal_compression_factor: int | None = None, ) -> tuple[torch.Tensor, int | float]: """Generate 3D mRoPE position IDs for vision tokens. @@ -282,23 +284,33 @@ def compute_mrope_position_ids_vision( to reflect real time so that videos at different frame rates get comparable temporal embeddings. + ``base_temporal_compression_factor`` sets the temporal grid the scaled + positions land on, and defaults to ``temporal_compression_factor``. Action + tokens run at frame rate (``temporal_compression_factor=1``) but must share + the vision latent-frame grid, so they pass the vision VAE factor here. + Returns: (position_ids [3, grid_t * grid_h * grid_w], next_temporal_offset) """ + if base_temporal_compression_factor is None: + base_temporal_compression_factor = temporal_compression_factor + if enable_fps_modulation and fps is not None: tps = fps / temporal_compression_factor - base_tps = base_fps / temporal_compression_factor + base_tps = base_fps / base_temporal_compression_factor frame_indices = torch.arange(grid_t, dtype=torch.float32) t_index = ( - (frame_indices / 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() @@ -318,6 +330,94 @@ 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. + + Action tokens are uncompressed in time, so they advance one source frame per + token while vision latent frames advance ``base_temporal_compression_factor`` + source frames. Scaling against the vision base rate keeps both streams on + one shared timeline. + """ + return compute_mrope_position_ids_vision( + grid_t=grid_t, + grid_h=1, + 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, + base_temporal_compression_factor=base_temporal_compression_factor, + ) + + +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 validate_domain_ids(self, domain_id: torch.Tensor) -> None: + """Range-check the ids. Reads a device tensor, so call once per request. + + Out-of-range ids index ``nn.Embedding`` out of bounds, which on GPU is a + device-side assert with no useful message; this turns it into a real + error. Kept out of forward() because the ``if`` on a device predicate is + a blocking sync, and forward() runs twice on every denoise step. + """ + if torch.any((domain_id < 0) | (domain_id >= self.num_domains)): + raise ValueError( + f"Cosmos3 action domain_id must be in [0, {self.num_domains}), " + f"got {domain_id.tolist()}." + ) + + def forward(self, x: torch.Tensor, domain_id: torch.Tensor) -> torch.Tensor: + if domain_id.ndim == 0: + domain_id = domain_id.unsqueeze(0) + 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]}." + ) + + 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. @@ -941,8 +1041,9 @@ def __init__(self, model_config: DiffusionModelConfig): pretrained_config = apply_pretrained_config_compat_defaults(model_config.pretrained_config) self.recipe = resolve_arch_recipe(pretrained_config) self.audio_gen = getattr(pretrained_config, "sound_gen", False) - # Config fact only: the transformer never constructs action modules. - self.has_action_weights = getattr(pretrained_config, "action_gen", False) + self.action_gen = getattr(pretrained_config, "action_gen", False) + # Config-fact alias kept for callers that predate action support. + self.has_action_weights = self.action_gen self.hidden_size = pretrained_config.hidden_size self.num_hidden_layers = pretrained_config.num_hidden_layers @@ -970,6 +1071,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" @@ -1031,6 +1155,8 @@ def __init__(self, model_config: DiffusionModelConfig): self.cached_kv = None self.cached_freqs_gen = None + self.cached_freqs_gen_combined = None + self.domain_ids_validated = False self.__post_init__() @@ -1228,9 +1354,97 @@ 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 register_cuda_graph_extra_key_fns(self, runner) -> None: + """Make the position-determining scalars part of the graph key. + + The base key is tensor shapes only, but the rotary tables are built + from Python scalars that leave every shape unchanged: the frame rate, + the action clock, and the offset of the first action step. Two requests + differing only in these produce different positions at identical + shapes, so without them one captured graph would be replayed for both. + Each returns ``None`` when absent, which drops that part of the key -- + a video-only request keys exactly as it did before. + """ + super().register_cuda_graph_extra_key_fns(runner) + + def _float_key(name): + def fn(*args, **kwargs): + value = kwargs.get(name) + return None if value is None else float(value) + + return fn + + def _int_key(name): + def fn(*args, **kwargs): + value = kwargs.get(name) + return None if value is None else int(value) + + return fn + + runner.register_extra_key_fn("fps", _float_key("fps")) + runner.register_extra_key_fn("action_fps", _float_key("action_fps")) + runner.register_extra_key_fn( + "action_start_frame_offset", _int_key("action_start_frame_offset") + ) + def reset_cache(self): self.cached_kv = None self.cached_freqs_gen = None + self.cached_freqs_gen_combined = None + self.domain_ids_validated = False def forward( self, @@ -1243,6 +1457,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, control_latents: list[torch.Tensor] | tuple[torch.Tensor, ...] | torch.Tensor | None = None, transfer_share_vision_temporal_positions: bool = True, **kwargs, @@ -1280,13 +1499,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() @@ -1358,8 +1587,9 @@ def forward( else: self.cached_kv = cached_kv_full - # --- Audio token injection ------------------------------------------------- + # --- Extra modality token injection (mutually exclusive: action, audio or control) --- T_vid_tokens = hidden_gen.shape[1] # T * Hp * Wp + T_action = 0 T_audio = 0 T_control = 0 hidden_controls: list[torch.Tensor] = [] @@ -1369,26 +1599,88 @@ def forward( raise ValueError( "Cosmos3 transfer control latents cannot be combined with sound latents" ) + if has_control and action_latents is not None: + raise ValueError( + "Cosmos3 transfer control latents cannot be combined with action latents" + ) - 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 control|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 + ) + if not self.domain_ids_validated: + # Once per request, alongside the other first-step host work. + self.action_proj_in.validate_domain_ids( + action_domain_ids_tensor.to(dtype=torch.long).reshape(-1) + ) + self.domain_ids_validated = True + T_action = action_latents.shape[1] + # Checked, not cast: bmm in DomainAwareLinear needs the latents to + # match the projection weights, and silently converting a whole + # stream every step would hide the misconfiguration that produced + # the mismatch. + if action_latents.dtype != self.action_proj_in.dtype: + raise ValueError( + "Cosmos3 action latents must match the action projection dtype: " + f"latents={action_latents.dtype}, projection={self.action_proj_in.dtype}." + ) + hidden_action = self.action_proj_in( + self.pack_action(action_latents), action_domain_ids_tensor + ) + 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) + # The rotary table is request-invariant: chunk size, prompt lengths, + # fps and the frame offset are all fixed once the request starts. + # Recomputing it per step costs a device-to-host sync per batch + # element (the prompt-length readback), an H2D copy of the position + # ids, and two concatenations -- all for the same numbers. + if self.cached_freqs_gen_combined is None: + effective_action_fps = ( + action_fps if action_fps is not None else (fps or self.base_fps) + ) + cos_a, sin_a = self._compute_action_rope_freqs( + T_action, + text_mask, + float(effective_action_fps), + action_start_frame_offset, + hidden_states.device, + hidden_gen.dtype, + ) + cos_v, sin_v = self.cached_freqs_gen + self.cached_freqs_gen_combined = ( + torch.cat([cos_v, cos_a], dim=1), + torch.cat([sin_v, sin_a], dim=1), + ) + freqs_gen_combined = self.cached_freqs_gen_combined + elif audio_latents is not None and self.audio_gen: T_audio = audio_latents.shape[2] hidden_audio = self.pack_audio_latents(audio_latents).to(hidden_gen.dtype) hidden_audio = self.audio2llm(hidden_audio) + self.audio_modality_embed hidden_audio = hidden_audio + time_embed.unsqueeze(1) - cos_a, sin_a = self._compute_audio_rope_freqs( - T_audio, - text_mask, - float(self.audio_latent_fps), - hidden_states.device, - hidden_gen.dtype, - ) - # [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 = ( - torch.cat([cos_v, cos_a], dim=1), - torch.cat([sin_v, sin_a], dim=1), - ) + if self.cached_freqs_gen_combined is None: + cos_a, sin_a = self._compute_audio_rope_freqs( + T_audio, + text_mask, + float(self.audio_latent_fps), + hidden_states.device, + hidden_gen.dtype, + ) + cos_v, sin_v = self.cached_freqs_gen + self.cached_freqs_gen_combined = ( + torch.cat([cos_v, cos_a], dim=1), + torch.cat([sin_v, sin_a], dim=1), + ) + freqs_gen_combined = self.cached_freqs_gen_combined elif has_control: for idx, control in enumerate(control_lantent_list): if control.shape != hidden_states.shape: @@ -1447,18 +1739,31 @@ def forward( self.llm2vae(hidden_gen[:, T_control : T_control + T_vid_tokens]), T, H, W ) - # --- Decode audio velocity (if requested) --------------------------------- - # Control latents and audio are mutually exclusive (guarded above), so the - # audio span always starts right after the video tokens. + # --- Decode extra-modality velocity (action XOR audio; follows video) --- + # Sequence layout is control|video|action-or-audio. Controls are prepended + # and are mutually exclusive with both action and audio, so T_control is + # zero whenever this span is non-empty; carrying it keeps the offset right + # if that exclusion is ever relaxed. + extra_start = T_control + 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. @@ -1470,8 +1775,7 @@ def load_weights(self, weights: dict) -> None: remapped = {} skip_prefixes = ( "lm_head.", - "action_modality_embed", - "action_proj_", + "action_pos_embed.", ) skipped_keys = [] @@ -1509,6 +1813,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.") @@ -1691,6 +2003,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/models/wan/pipeline_wan.py b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py index 2a2407752f15..f686b61753e4 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py @@ -627,11 +627,14 @@ def forward_fn( encoder_hidden_states=encoder_hidden_states, ) - # Pin reference image to latent after each scheduler step (Wan 2.2 5B I2V only) - def _pin_i2v_first_frame(x): - return ((1 - i2v_first_frame_mask) * i2v_condition + i2v_first_frame_mask * x).to( + # Pin reference image to latent after each scheduler step (Wan 2.2 5B I2V only). + # post_step_fn also carries the denoise loop's side-stream latents (Cosmos3 + # denoises action/audio alongside video); Wan has none, so they pass through. + def _pin_i2v_first_frame(x, extra_stream_latents): + pinned = ((1 - i2v_first_frame_mask) * i2v_condition + i2v_first_frame_mask * x).to( self.dtype ) + return pinned, extra_stream_latents post_step_fn = _pin_i2v_first_frame if (self.is_wan22_5b and is_i2v) else None diff --git a/tensorrt_llm/_torch/visual_gen/output.py b/tensorrt_llm/_torch/visual_gen/output.py index b77956b7b6d8..76cf825cb4d9 100644 --- a/tensorrt_llm/_torch/visual_gen/output.py +++ b/tensorrt_llm/_torch/visual_gen/output.py @@ -51,6 +51,12 @@ class PipelineOutput: ``(B, channels, T_audio)``, dtype ``float32``. Populated by LTX-2. The leading batch dim is always present, even for single-prompt requests (size 1). + action: Predicted or refined action trajectory as ``torch.Tensor`` + shape ``(B, T_action, D_raw)``, dtype ``float32``. Populated by + Cosmos3 action generation (``policy`` / ``forward_dynamics`` / + ``inverse_dynamics``), already sliced to the embodiment's real + degrees of freedom, so ``D_raw`` states them; no VAE decode. + ``None`` when action generation was not requested. frame_rate: Video frame rate in fps. Populated by video pipelines (Wan T2V/I2V emit ``16.0``; LTX-2 emits ``params.frame_rate``). ``None`` for image-only pipelines. @@ -72,6 +78,7 @@ 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 pre_denoise: float = 0.0 @@ -201,6 +208,7 @@ 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, metrics=metrics, @@ -245,6 +253,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,6 +271,7 @@ 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, metrics=metrics, diff --git a/tensorrt_llm/_torch/visual_gen/pipeline.py b/tensorrt_llm/_torch/visual_gen/pipeline.py index 78627069396a..2cf24173699c 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline.py @@ -58,6 +58,15 @@ class ExtraParamSchema(StrictBaseModel): "values. Must be a module-level function (specs are pickled to the " "coordinator in the READY handshake).", ) + requires_tensor_output: bool = Field( + default=False, + description="Setting this parameter makes the request produce a result " + "the media encoders cannot represent (a non-image/video modality), so " + "the response must be a tensor payload. Serve resolves 'auto' to a " + "tensor format and rejects an explicit encoder format. Declared here " + "rather than hard-coded in the routes so the serving layer needs no " + "per-model knowledge.", + ) if TYPE_CHECKING: @@ -1186,8 +1195,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 after each scheduler step, - invoked as ``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. scheduler_step_kwargs: Extra keyword arguments forwarded to every scheduler's ``step()`` call. @@ -1297,7 +1309,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/decoding.py b/tensorrt_llm/media/decoding.py index 71708229cc38..87c1b9e68288 100644 --- a/tensorrt_llm/media/decoding.py +++ b/tensorrt_llm/media/decoding.py @@ -108,6 +108,51 @@ def resize_center_crop_uint8(frames: torch.Tensor, target_h: int, target_w: int) return x.round_().clamp_(0, 255).to(torch.uint8).permute(0, 2, 3, 1).contiguous() +def resize_fit_pad_uint8(frames: torch.Tensor, target_h: int, target_w: int) -> torch.Tensor: + """Resize to fit inside the target, then pad bottom/right to fill it. + + The counterpart to :func:`resize_center_crop_uint8`, for references whose + periphery carries signal — a robot gripper works at the frame edge, so + cropping it away costs the model the thing it is meant to act on. + + Semantics mirror the action reference's ``reflection_pad_to_target``: + contain-scale by ``min(target/source, 1.0)`` (never enlarge — a small clip + keeps its own pixels and gets a wider border), round-rather-than-ceil + resize, then pad bottom/right by reflection, switching to edge replication + when a pad run reaches the resized extent (reflection has no source pixels + left to mirror). The resampling filter stays this module's Lanczos-3 rather + than the reference's bicubic: the geometry is what preserves content, and a + second filter would buy sub-pixel differences for a second code path. + """ + t, h, w, c = frames.shape + if (h, w) == (target_h, target_w): + return frames + ratio = min(target_w / w, target_h / h, 1.0) + resize_w = min(int(ratio * w + 0.5), target_w) + resize_h = min(int(ratio * h + 0.5), target_h) + + # Two passes with a uint8-quantized intermediate, as in the cover path. + x = frames.permute(0, 3, 1, 2).to(torch.float32) # [T, C, H, W] + if resize_w != w: + weights, taps = _lanczos_taps(w, resize_w, str(frames.device)) + x = _resample_last_dim(x, weights, taps) + x = x.round_().clamp_(0, 255) + if resize_h != h: + weights, taps = _lanczos_taps(h, resize_h, str(frames.device)) + x = _resample_last_dim(x.transpose(-1, -2), weights, taps).transpose(-1, -2) + x = x.round_().clamp_(0, 255) + + pad_w = target_w - resize_w + pad_h = target_h - resize_h + if pad_w or pad_h: + mode = "replicate" if (pad_w >= resize_w or pad_h >= resize_h) else "reflect" + x = torch.nn.functional.pad(x, (0, pad_w, 0, pad_h), mode=mode) + return x.to(torch.uint8).permute(0, 2, 3, 1).contiguous() + + +_RESIZE_MODES = {"cover": resize_center_crop_uint8, "fit": resize_fit_pad_uint8} + + class VideoStreamInfo(NamedTuple): """What a container header reports about its video stream.""" @@ -169,6 +214,8 @@ def decode_video_reference_window( target_h: int, target_w: int, device: torch.device, + resize: str = "cover", + frame_step: int = 1, ) -> torch.Tensor: """Decode frames ``[first_frame, last_frame]`` of a reference on device. @@ -176,6 +223,22 @@ def decode_video_reference_window( non-negative counts from the start, negative from the end, so ``-1`` is the last frame and ``(-8, -1)`` the final eight. Both ends are inclusive. + ``frame_step`` retains every n-th frame of the range, so ``(0, 96)`` with + ``frame_step=6`` yields frames 0, 6, ... 96 — seventeen frames, not + ninety-seven. A caller whose model expects a frame spacing the source was + not shot at uses this to pick the right frames; the ratio of the two rates + is the caller's to compute, and no rate is named here. Skipped frames are + still decoded (inter-frame compression leaves no choice) but are neither + resized nor retained, so the cost is decode time, not memory. Only + non-negative ranges may step: the negative form wraps a ring whose length + is not known until EOS, and combining the two is not supported. + + ``resize`` selects how each frame reaches ``target_h x target_w``: + ``"cover"`` scales to fill and center-crops (the default, and what video + continuation wants); ``"fit"`` scales to fit and pads, for references whose + frame edges carry signal. See :func:`resize_center_crop_uint8` and + :func:`resize_fit_pad_uint8`. + A negative index costs a decode to EOS — the memory-buffer demuxer is a forward-only feeder, seeking is not assumed — so the caller pays for the whole clip when asking from the end. Non-negative ranges stop as soon as @@ -195,6 +258,18 @@ def decode_video_reference_window( raise ValueError( f"first_frame must not exceed last_frame, got ({first_frame}, {last_frame})." ) + if frame_step < 1: + raise ValueError(f"frame_step must be at least 1, got {frame_step}.") + if frame_step > 1 and first_frame < 0: + raise ValueError( + f"frame_step > 1 is only supported for non-negative ranges, got " + f"({first_frame}, {last_frame}) with frame_step={frame_step}." + ) + resize_frames = _RESIZE_MODES.get(resize) + if resize_frames is None: + raise ValueError( + f"Unknown resize mode {resize!r}; expected one of {sorted(_RESIZE_MODES)}." + ) window = last_frame - first_frame + 1 from_end = first_frame < 0 try: @@ -249,7 +324,7 @@ def _read(buf: bytearray) -> int: # Non-negative ranges retain exactly the requested slice, so the ring # is filled once; negative ranges cannot know the length up front, so # it wraps and holds the trailing `tail` frames until EOS. - tail = -first_frame if from_end else window + tail = -first_frame if from_end else (window + frame_step - 1) // frame_step ring = torch.empty(tail, target_h, target_w, 3, dtype=torch.uint8, device=device) count = 0 # frames decoded so far, i.e. the index of the next frame kept = 0 # frames written into the ring @@ -260,12 +335,14 @@ def _read(buf: bytearray) -> int: if not from_end and count > last_frame: done = True break - if from_end or count >= first_frame: + if from_end or ( + count >= first_frame and (count - first_frame) % frame_step == 0 + ): decoded = torch.from_dlpack(frame) # Ownership copy off the NVDEC surface (recycled by # the decoder) and resize-before-retain in one step. ring[kept % tail].copy_( - resize_center_crop_uint8(decoded.unsqueeze(0), target_h, target_w)[0] + resize_frames(decoded.unsqueeze(0), target_h, target_w)[0] ) kept += 1 count += 1 @@ -274,7 +351,7 @@ def _read(buf: bytearray) -> int: except torch.cuda.OutOfMemoryError as exc: raise MemoryError( f"Out of device memory while decoding the video reference " - f"({window} frames @ {target_w}x{target_h} retained): {exc}" + f"({tail} frames @ {target_w}x{target_h} retained): {exc}" ) from exc except nvc.PyNvVCException as exc: raise ValueError( diff --git a/tensorrt_llm/media/tensor_payload.py b/tensorrt_llm/media/tensor_payload.py index b61580d9f439..70fb0cdadb3b 100644 --- a/tensorrt_llm/media/tensor_payload.py +++ b/tensorrt_llm/media/tensor_payload.py @@ -5,13 +5,12 @@ Two payload formats are supported: - ``"safetensors"``: writes a single file with named tensors - (``image``/``video``/``audio``). 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``). Scalar metadata + (``frame_rate``, ``audio_sample_rate``) is stored two ways: as a 0-d + tensor under the same key (so ``safetensors.torch.load(bytes)`` returns + it alongside the media tensors — consumers call ``.item()`` to unbox) + and as a stringified value in the file header (preserved for callers + using ``safe_open(...).metadata()``). No pickle on load. - ``"pt"``: writes a single file via :func:`torch.save` with the same tensor keys plus scalar metadata as native Python values. Clients should load with ``torch.load(buf, weights_only=True)`` @@ -49,12 +48,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 +64,15 @@ def _modalities(output: "VisualGenOutput") -> Tuple[Tuple[str, Optional[torch.Te ("image", output.image), ("video", output.video), ("audio", output.audio), + ("action", output.action), ) def infer_batch_size(output: "VisualGenOutput") -> int: """Return the leading batch dimension across the populated media tensors. - Image is batched only at rank 4, video at rank 5, audio at rank 3. - An unbatched media tensor reports a batch size of 1 so list-path + Image is batched only at rank 4, video at rank 5, audio and action + at rank 3. An unbatched media tensor reports a batch size of 1 so list-path callers can still ask for ``[0]`` and get a single-item payload. Raises :class:`ValueError` when *output* carries no media tensor. """ diff --git a/tensorrt_llm/serve/openai_video_routes.py b/tensorrt_llm/serve/openai_video_routes.py index d8dea240196d..cf69905ad505 100644 --- a/tensorrt_llm/serve/openai_video_routes.py +++ b/tensorrt_llm/serve/openai_video_routes.py @@ -56,6 +56,42 @@ def _video_content_type(suffix: str) -> str: _KNOWN_VIDEO_OUTPUT_SUFFIXES = (".mp4", ".avi", ".safetensors", ".pt") +def _resolve_tensor_only_format(fmt, extra_params, extra_param_specs): + """Resolve ``format`` for a request whose result an encoder cannot carry. + + A pipeline marks such parameters with ``requires_tensor_output`` on their + :class:`ExtraParamSchema` (Cosmos3 does so for ``action_mode``: a predicted + trajectory has no representation in a video container). The rule keeps this + route model-agnostic -- it reads the declaration, never the parameter's + meaning: + + * ``auto`` resolves to ``safetensors``, so the default request returns + everything it generated instead of silently dropping a modality; + * an explicit tensor format passes through; + * an explicit encoder format is a contradiction the caller stated -- two + incompatible things in one request -- so it is rejected rather than + guessed at. + """ + if not extra_params or not extra_param_specs: + return fmt + triggered = sorted( + key + for key, spec in extra_param_specs.items() + if getattr(spec, "requires_tensor_output", False) and extra_params.get(key) is not None + ) + if not triggered: + return fmt + if is_tensor_format(fmt): + return fmt + if fmt == "auto": + return _DEFAULT_TENSOR_FORMAT + raise ValueError( + f"format={fmt!r} cannot carry the result of {', '.join(triggered)}: a " + f"video container holds only video. Use format='safetensors' or 'pt', " + f"or omit format so 'auto' selects a payload that carries everything." + ) + + def _preflight_encoder_format(fmt): """Pre-flight an encoder format string before any GPU work. @@ -74,6 +110,9 @@ def _preflight_encoder_format(fmt): raise ValueError(str(exc)) from exc +_DEFAULT_TENSOR_FORMAT = "safetensors" + + def _path_json_video_response( video_id: str, path: Union[str, Path], headers: Optional[dict[str, str]] = None ) -> JSONResponse: @@ -125,7 +164,10 @@ async def openai_video_generation_sync(self, raw_request: Request) -> Response: self.generator, media_storage_path=str(self.media_storage_path), ) - resolved_encoder_fmt = _preflight_encoder_format(request.format) + request_format = _resolve_tensor_only_format( + request.format, request.extra_params, self.generator.extra_param_specs + ) + resolved_encoder_fmt = _preflight_encoder_format(request_format) logger.info( f"Generating video: {video_id} with params: {params} and prompt: {request.prompt}" ) @@ -153,8 +195,8 @@ async def openai_video_generation_sync(self, raw_request: Request) -> Response: status_code=HTTPStatus.INTERNAL_SERVER_ERROR, ) - if is_tensor_format(request.format): - ext = f".{request.format}" + if is_tensor_format(request_format): + ext = f".{request_format}" media_type = "application/octet-stream" # Match the encoder-format path: persist one file per batch # item, ship the first as the route's primary download @@ -164,7 +206,7 @@ async def openai_video_generation_sync(self, raw_request: Request) -> Response: tensor_paths = [ self.media_storage_path / f"{video_id}_{i}{ext}" for i in range(batch_size) ] - saved_paths = output.save(tensor_paths, format=request.format) + saved_paths = output.save(tensor_paths, format=request_format) target = saved_paths[0] latency = time.perf_counter() - sync_video_start logger.info( @@ -361,7 +403,10 @@ async def openai_video_generation_async( declared_defaults=self.generator.executor.default_generation_params, extra_param_specs=self.generator.executor.extra_param_specs, ) - _preflight_encoder_format(request.format) + request_format = _resolve_tensor_only_format( + request.format, request.extra_params, self.generator.extra_param_specs + ) + _preflight_encoder_format(request_format) logger.info( f"Generating video: {video_id} with params: {params} and prompt: {request.prompt}" ) @@ -388,6 +433,7 @@ async def openai_video_generation_async( video_id=video_id, request=request, params=params, + request_format=request_format, ) ) self.video_gen_tasks[video_id] = task @@ -413,8 +459,15 @@ async def _generate_video_background( video_id: str, request: VideoGenerationRequest, params: VisualGenParams, + request_format: str, ): - """Background task to generate video and save to storage.""" + """Background task to generate video and save to storage. + + ``request_format`` is the format already resolved by the route (see + :func:`_resolve_tensor_only_format`), not ``request.format``: the + resolution happens before the job is queued so a rejected request never + becomes a background task. + """ try: background_start = time.perf_counter() job = await VIDEO_STORE.get(video_id) @@ -441,18 +494,18 @@ async def _generate_video_background( job.status = "postprocessing" await VIDEO_STORE.upsert(video_id, job) - if is_tensor_format(request.format): + if is_tensor_format(request_format): # One tensor file per batch item, mirroring the encoder # path; the async job records all paths on # ``output_paths`` so subsequent GETs can find each item. batch_size = output.video.shape[0] if output.video.dim() == 5 else 1 tensor_paths = [ - self.media_storage_path / f"{video_id}_{i}.{request.format}" + self.media_storage_path / f"{video_id}_{i}.{request_format}" for i in range(batch_size) ] - saved_paths = output.save(tensor_paths, format=request.format) + saved_paths = output.save(tensor_paths, format=request_format) else: - resolved_fmt, _ = resolve_video_format(request.format) + resolved_fmt, _ = resolve_video_format(request_format) batch_size = output.video.shape[0] if output.video.dim() == 5 else 1 paths_in = [self.media_storage_path / f"{video_id}_{i}" for i in range(batch_size)] _save_kwargs = dict( diff --git a/tensorrt_llm/visual_gen/output.py b/tensorrt_llm/visual_gen/output.py index 0660e661122f..88ba6bc5082a 100644 --- a/tensorrt_llm/visual_gen/output.py +++ b/tensorrt_llm/visual_gen/output.py @@ -96,6 +96,7 @@ 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 error: Optional[str] = None @@ -121,8 +122,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 +139,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 +153,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 +167,13 @@ 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'); no video or image container can carry " + "an action trajectory." + ) + if self.image is not None: if is_batch: saved_list = save_images( @@ -209,7 +217,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 a3145ebbe81f..cd9d86ae19b8 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 @@ -121,6 +122,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, *, @@ -175,6 +188,18 @@ 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}" + ) + # Terminal on purpose: membership in the literal set already + # decides the value, so the type, validator and range checks + # below cannot add anything. A literal spec that also declares + # one of those has a redundant declaration, not a skipped check. + continue # Type check expected_types = _TYPE_MAP.get(spec.type) 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 9321ea3d4242..37d6b3bc5697 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -269,6 +269,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 - unittest/_torch/visual_gen/test_cosmos3_distilled.py 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 afc39db47aa4..b9029e7946f8 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 @@ -157,6 +157,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 @@ -462,6 +471,34 @@ 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 / _NUM_TRAIN_TIMESTEPS, + raw_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, @@ -628,6 +665,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 @@ -789,6 +857,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..fa594ba78de6 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_action.py @@ -0,0 +1,796 @@ +# 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 json + +import numpy as np +import PIL.Image +import pytest +import torch + +from tensorrt_llm._torch.visual_gen.models.cosmos3.action import ( + ACTION_ASPECT_RATIO_LABELS, + ACTION_VIEWPOINT_TEMPLATES, + DEFAULT_ACTION_VIEW_POINT, + EMBODIMENT_TO_DOMAIN_ID, + EMBODIMENT_TO_RAW_ACTION_DIM, + VIDEO_RES_SIZE_INFO, + action_aspect_ratio_label, + action_reference_frame_step, + action_reference_size, + build_action_json_prompt, + find_closest_target_size, + normalize_action_resolution, + prepare_action_latents, + resize_and_pad_action_image, + resolve_action_size, + resolve_domain_id, + resolve_raw_action_dim, +) +from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( + COSMOS3_DOMAIN_PRESET_ALIASES, + COSMOS3_DOMAIN_PRESETS, + COSMOS3_EXTRA_SPECS, + get_domain_preset, + resolve_domain_action_config, +) +from tensorrt_llm._torch.visual_gen.models.cosmos3.transformer_cosmos3 import ( + compute_mrope_position_ids_action, + compute_mrope_position_ids_vision, +) + +pytestmark = pytest.mark.cosmos3 + + +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): + # find_closest_target_size picks from whatever entries exist, so a bucket + # that lost "9,16" would silently land portrait sources on another canvas. + assert set(VIDEO_RES_SIZE_INFO[action_resolution]) == set(ACTION_ASPECT_RATIO_LABELS) + + +class TestResolveActionSize: + SOURCE_H, SOURCE_W = 480, 832 + + def test_explicit_height_and_width_are_unchanged(self): + assert resolve_action_size(400, 600, self.SOURCE_H, self.SOURCE_W, 480) == (400, 600) + + def test_unset_height_and_width_use_action_resolution_bucket(self): + assert resolve_action_size(None, None, self.SOURCE_H, self.SOURCE_W, 480) == (480, 832) + + def test_partial_height_fills_width_from_bucket(self): + assert resolve_action_size(400, None, self.SOURCE_H, self.SOURCE_W, 480) == (400, 832) + + def test_partial_width_fills_height_from_bucket(self): + assert resolve_action_size(None, 600, self.SOURCE_H, self.SOURCE_W, 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 == get_domain_preset("droid_lerobot") + + def test_presets_carry_sampling_settings_only(self): + """Width is per-embodiment; aliases share presets, so it must not live there.""" + for preset in COSMOS3_DOMAIN_PRESETS.values(): + assert "raw_action_dim" not in preset + + @pytest.mark.parametrize( + "domain_name,expected", + [ + ("bridge_orig_lerobot", 10), + ("droid_lerobot", 10), + ("robomind-franka", 10), + ("robomind-ur", 10), + ("robomind-franka-dual", 20), # dual arm, not the droid preset's 10 + ("galbot", 30), # humanoid stack, not agibotworld's 29 + ("agibotworld", 29), + ("agibot_gear_gripper", 29), + ("agibot_gear_gripper_ext", 29), + ("av", 9), + ("camera_pose", 9), + ("hand_pose", 57), + ("pusht", 2), + ("umi", 10), + ("fractal", 10), + ], + ) + def test_canonical_action_width_per_embodiment(self, domain_name, expected): + assert resolve_raw_action_dim(domain_name=domain_name) == expected + assert resolve_domain_action_config(domain_name=domain_name)["raw_action_dim"] == expected + + def test_aliased_domains_keep_their_own_width(self): + """Sharing a sampling preset must not import that preset's action width.""" + for alias, canonical in COSMOS3_DOMAIN_PRESET_ALIASES.items(): + alias_dim = resolve_raw_action_dim(domain_name=alias) + if alias_dim is None: + continue + assert alias_dim == EMBODIMENT_TO_RAW_ACTION_DIM[alias], ( + f"{alias} must keep its own width, not {canonical}'s" + ) + + def test_libero_has_no_canonical_width(self): + """LIBERO's width depends on the dataset's rotation space (7/10/13).""" + assert "libero" not in EMBODIMENT_TO_RAW_ACTION_DIM + cfg = resolve_domain_action_config(domain_name="libero") + assert cfg["raw_action_dim"] is None + assert cfg["action_resolution"] == 256 # sampling preset still applies + assert any("canonical action width" in w for w in cfg["warnings"]) + + def test_explicit_raw_action_dim_overrides_with_warning(self): + cfg = resolve_domain_action_config(domain_name="libero", raw_action_dim=7) + assert cfg["raw_action_dim"] == 7 + assert cfg["warnings"] == [] + + def test_domain_id_resolves_width_when_unambiguous(self): + assert resolve_raw_action_dim(domain_id=12) == 20 # robomind-franka-dual + assert resolve_raw_action_dim(domain_id=8) == 10 # droid / robomind-franka agree + assert resolve_raw_action_dim(domain_id=15) == 29 # all agibot variants agree + assert resolve_raw_action_dim(domain_id=5) is None # libero + assert resolve_raw_action_dim(domain_id=0) is None # no_action + + def test_every_width_entry_has_a_domain_id(self): + assert set(EMBODIMENT_TO_RAW_ACTION_DIM) <= set(EMBODIMENT_TO_DOMAIN_ID) + + def test_unknown_domain_warns_and_uses_generic_defaults(self): + cfg = resolve_domain_action_config(domain_name="typo_domain") + 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 TestActionReferenceSize: + """Canvas selection needs the source size, not a decoded frame.""" + + def test_policy_measures_image(self, tmp_path): + image_path = tmp_path / "frame.png" + PIL.Image.new("RGB", (640, 480), "blue").save(image_path) + assert action_reference_size(action_mode="policy", image=str(image_path), video=None) == ( + 480, + 640, + ) + + def test_policy_prefers_image_over_video(self, tmp_path): + image_path = tmp_path / "frame.png" + PIL.Image.new("RGB", (320, 240), "blue").save(image_path) + # Bytes would raise if consulted: the image must win. + assert action_reference_size( + action_mode="policy", image=str(image_path), video=b"not-a-video" + ) == (240, 320) + + def test_accepts_pil_image_directly(self): + assert action_reference_size( + action_mode="forward_dynamics", + image=PIL.Image.new("RGB", (256, 128)), + video=None, + ) == (128, 256) + + def test_missing_source_raises(self): + with pytest.raises(ValueError, match="requires an image or video"): + action_reference_size(action_mode="policy", image=None, video=None) + + def test_inverse_dynamics_ignores_image(self, tmp_path): + image_path = tmp_path / "frame.png" + PIL.Image.new("RGB", (640, 480), "blue").save(image_path) + # inverse_dynamics conditions on the clip, so an image is not a source. + with pytest.raises(ValueError, match="requires an image or video"): + action_reference_size(action_mode="inverse_dynamics", image=str(image_path), video=None) + + def test_https_reference_goes_through_the_repo_loader(self, monkeypatch): + """Bundled action prompts point at https:// frames, so a bare + PIL.Image.open(path) would fail on every one of them.""" + import tensorrt_llm.inputs.utils as inputs_utils + + requested = [] + + def fake_load_image(source, format="pt", device="cpu"): + requested.append((source, format)) + return PIL.Image.new("RGB", (640, 480), "blue") + + monkeypatch.setattr(inputs_utils, "load_image", fake_load_image) + assert action_reference_size( + action_mode="policy", + image="https://example.invalid/frame.png", + video=None, + ) == (480, 640) + assert requested == [("https://example.invalid/frame.png", "pil")] + + def test_video_bytes_probe_the_container_header(self, monkeypatch): + """Bytes are measured from the header, never by decoding a frame.""" + import tensorrt_llm.media.decoding as decoding + + monkeypatch.setattr( + decoding, + "video_stream_info", + lambda data: decoding.VideoStreamInfo(480, 640, 30.0), + ) + assert action_reference_size( + action_mode="inverse_dynamics", image=None, video=b"\x00mp4" + ) == (480, 640) + + def test_unreadable_video_bytes_are_rejected(self, monkeypatch): + """An unreadable container fails here, not silently at decode.""" + import tensorrt_llm.media.decoding as decoding + + monkeypatch.setattr(decoding, "video_stream_info", lambda data: None) + with pytest.raises(ValueError, match="could not be demuxed"): + action_reference_size(action_mode="inverse_dynamics", image=None, video=b"\x00bad") + + +class TestActionReferenceFrameStep: + """The reference is thinned to the embodiment's rate, never invented.""" + + @pytest.mark.parametrize( + "source_frame_rate, target_frame_rate, expected", + [ + (30.0, 5.0, 6), # bridge: every sixth frame of a 30 fps clip + (5.0, 5.0, 1), # already at the trained rate + (24.0, 5.0, 5), # 4.8 rounds to 5 + (10.0, 30.0, 1), # slower than trained: cannot be thinned + (None, 5.0, 1), # header unreadable + (0.0, 5.0, 1), # header reported nothing usable + ], + ) + def test_step_from_rates(self, source_frame_rate, target_frame_rate, expected): + assert action_reference_frame_step(source_frame_rate, target_frame_rate) == expected + + +class TestActionJsonPrompt: + """The trained action caption: structured JSON, not the flat video templates.""" + + BRIDGE = dict(num_frames=17, frame_rate=5.0, height=480, width=832) + + def test_matches_trained_shape(self): + payload = json.loads( + build_action_json_prompt( + "Pick up the pear and place it in the bag", + view_point="ego_view", + **self.BRIDGE, + ) + ) + assert payload == { + "cinematography": { + "framing": ( + "This video is captured from a first-person perspective looking at the scene." + ) + }, + "actions": [ + { + "time": "0:00-0:03", + "description": "Pick up the pear and place it in the bag.", + } + ], + "duration": "3s", + "fps": 5.0, + "resolution": {"H": 480, "W": 832}, + "aspect_ratio": "16,9", + } + + def test_key_order_is_preserved(self): + """Field order is part of the trained caption format.""" + text = build_action_json_prompt("Do a thing", view_point="ego_view", **self.BRIDGE) + assert list(json.loads(text).keys()) == [ + "cinematography", + "actions", + "duration", + "fps", + "resolution", + "aspect_ratio", + ] + + @pytest.mark.parametrize("view_point", sorted(ACTION_VIEWPOINT_TEMPLATES)) + def test_every_viewpoint_emits_its_trained_sentence(self, view_point): + payload = json.loads( + build_action_json_prompt("Do a thing", view_point=view_point, **self.BRIDGE) + ) + assert payload["cinematography"]["framing"] == ACTION_VIEWPOINT_TEMPLATES[view_point] + + def test_default_view_point_is_known(self): + assert DEFAULT_ACTION_VIEW_POINT in ACTION_VIEWPOINT_TEMPLATES + + @pytest.mark.parametrize("view_point", [None, "sideways_view"]) + def test_unknown_or_missing_view_point_drops_framing(self, view_point): + payload = json.loads( + build_action_json_prompt("Do a thing", view_point=view_point, **self.BRIDGE) + ) + assert "cinematography" not in payload + assert list(payload.keys())[0] == "actions" + + @pytest.mark.parametrize( + "description,expected", + [ + ("Pick up the pear", "Pick up the pear."), + ("Pick up the pear.", "Pick up the pear."), + ("Is it a pear?", "Is it a pear?"), + ("Grab it!", "Grab it!"), + (" padded ", "padded."), + ("", ""), + ], + ) + def test_description_is_terminated_once(self, description, expected): + payload = json.loads(build_action_json_prompt(description, view_point=None, **self.BRIDGE)) + assert payload["actions"][0]["description"] == expected + + @pytest.mark.parametrize( + "num_frames,frame_rate,duration,time_range", + [ + (17, 5.0, "3s", "0:00-0:03"), # bridge: 3.4s truncates, rounds to 3 + (17, 24.0, "0s", "0:00-0:01"), # 0.708s truncates to 0, rounds to 1 + (61, 10.0, "6s", "0:00-0:06"), # av preset + (241, 2.0, "120s", "0:00-2:00"), # crosses the minute boundary + ], + ) + def test_duration_truncates_while_time_range_rounds( + self, num_frames, frame_rate, duration, time_range + ): + payload = json.loads( + build_action_json_prompt( + "Do a thing", + view_point=None, + num_frames=num_frames, + frame_rate=frame_rate, + height=480, + width=832, + ) + ) + assert payload["duration"] == duration + assert payload["actions"][0]["time"] == time_range + assert payload["fps"] == float(frame_rate) + + @pytest.mark.parametrize("action_resolution", sorted(VIDEO_RES_SIZE_INFO)) + def test_aspect_label_matches_the_bucket_it_came_from(self, action_resolution): + """Every canvas is a bucket entry, so its label must round-trip.""" + for label, (width, height) in VIDEO_RES_SIZE_INFO[action_resolution].items(): + assert action_aspect_ratio_label(height, width) == label + + def test_aspect_label_is_not_a_reduced_fraction(self): + """832x480 reduces to 26,15 but the trained label is 16,9.""" + assert action_aspect_ratio_label(480, 832) == "16,9" + + def test_resolution_is_reported_as_the_padded_canvas(self): + payload = json.loads(build_action_json_prompt("Do a thing", view_point=None, **self.BRIDGE)) + assert payload["resolution"] == {"H": 480, "W": 832} + + def test_zero_frame_rate_does_not_raise(self): + payload = json.loads( + build_action_json_prompt( + "Do a thing", + view_point=None, + num_frames=17, + frame_rate=0.0, + height=480, + width=832, + ) + ) + assert payload["duration"] == "0s" + assert payload["actions"][0]["time"] == "0:00-0:00" + + def test_view_point_spec_defaults_to_ego_view(self): + spec = COSMOS3_EXTRA_SPECS["view_point"] + assert spec.default == DEFAULT_ACTION_VIEW_POINT + for view_point in ACTION_VIEWPOINT_TEMPLATES: + assert repr(view_point) in spec.type + + +def _reference_scaled_positions( + *, + grid_t: int, + temporal_offset: float, + fps: float, + base_fps: float, + temporal_compression_factor: int, + base_temporal_compression_factor: int, + start_frame_offset: int, +) -> list[float]: + """Transcription of cosmos-framework ``get_3d_mrope_ids_vae_tokens``. + + Reference: ``cosmos_framework/data/generator/sequence_packing/mrope.py`` + (mirrored by diffusers ``pipeline_cosmos3_omni.get_3d_mrope_ids_vae_tokens``). + """ + tps = fps / temporal_compression_factor + base_tps = base_fps / base_temporal_compression_factor + return [(i + start_frame_offset) / tps * base_tps + temporal_offset for i in range(grid_t)] + + +class TestActionMropePositionIds: + """Action tokens run at frame rate but must share the vision latent timeline.""" + + VISION_TCF = 4 + + @pytest.mark.parametrize( + "grid_t,temporal_offset,action_fps,base_fps,start_frame_offset", + [ + (4, 0.0, 24.0, 24.0, 1), + (4, 15032.0, 24.0, 24.0, 1), + (16, 0.0, 5.0, 24.0, 1), + (60, 0.0, 10.0, 24.0, 1), + (4, 0.0, 24.0, 24.0, 0), + ], + ) + def test_matches_reference_formula( + self, grid_t, temporal_offset, action_fps, base_fps, start_frame_offset + ): + ids, _ = compute_mrope_position_ids_action( + grid_t, + temporal_offset=temporal_offset, + action_fps=action_fps, + base_fps=base_fps, + base_temporal_compression_factor=self.VISION_TCF, + enable_fps_modulation=True, + start_frame_offset=start_frame_offset, + ) + expected = _reference_scaled_positions( + grid_t=grid_t, + temporal_offset=temporal_offset, + fps=action_fps, + base_fps=base_fps, + temporal_compression_factor=1, + base_temporal_compression_factor=self.VISION_TCF, + start_frame_offset=start_frame_offset, + ) + torch.testing.assert_close( + ids[0], torch.tensor(expected, dtype=ids.dtype), rtol=0, atol=1e-5 + ) + + def test_action_step_advances_one_source_frame(self): + """Consecutive action tokens are 1/vision_tcf of a latent frame apart.""" + ids, _ = compute_mrope_position_ids_action( + 8, + temporal_offset=0.0, + action_fps=24.0, + base_fps=24.0, + base_temporal_compression_factor=self.VISION_TCF, + enable_fps_modulation=True, + start_frame_offset=1, + ) + deltas = ids[0, 1:] - ids[0, :-1] + torch.testing.assert_close( + deltas, torch.full_like(deltas, 1.0 / self.VISION_TCF), rtol=0, atol=1e-5 + ) + + @pytest.mark.parametrize( + "action_chunk_size,num_frames,fps", + [ + (16, 17, 24.0), # generic COSMOS3_ACTION_PARAMS default + (16, 17, 5.0), # bridge_orig_lerobot preset + (60, 61, 10.0), # av preset + ], + ) + def test_last_action_token_lands_on_last_vision_latent_frame( + self, action_chunk_size, num_frames, fps + ): + """The 4x-scaling regression: action must not outrun the video it conditions. + + Vision and action are packed into one temporal axis, so the paired + (num_frames, action_chunk_size) config must place the final action token + exactly on the final vision latent frame. + """ + latent_t = (num_frames - 1) // self.VISION_TCF + 1 + vision_ids, _ = compute_mrope_position_ids_vision( + latent_t, + 1, + 1, + temporal_offset=0.0, + fps=fps, + base_fps=24.0, + temporal_compression_factor=self.VISION_TCF, + enable_fps_modulation=True, + ) + action_ids, _ = compute_mrope_position_ids_action( + action_chunk_size, + temporal_offset=0.0, + action_fps=fps, + base_fps=24.0, + base_temporal_compression_factor=self.VISION_TCF, + enable_fps_modulation=True, + start_frame_offset=1, + ) + assert action_ids[0, -1].item() == pytest.approx(vision_ids[0, -1].item(), abs=1e-5) + assert action_ids[0, 0].item() > vision_ids[0, 0].item() + + def test_spatial_rows_are_zero(self): + ids, _ = compute_mrope_position_ids_action( + 5, + temporal_offset=0.0, + action_fps=24.0, + base_fps=24.0, + base_temporal_compression_factor=self.VISION_TCF, + enable_fps_modulation=True, + ) + assert ids.shape == (3, 5) + assert torch.all(ids[1] == 0) + assert torch.all(ids[2] == 0) + + def test_fps_modulation_disabled_gives_integer_frame_indices(self): + ids, _ = compute_mrope_position_ids_action( + 4, + temporal_offset=7.0, + action_fps=24.0, + base_fps=24.0, + base_temporal_compression_factor=self.VISION_TCF, + enable_fps_modulation=False, + start_frame_offset=1, + ) + assert ids[0].tolist() == [8, 9, 10, 11] + + def test_lower_action_fps_stretches_positions(self): + kwargs = dict( + temporal_offset=0.0, + base_fps=24.0, + base_temporal_compression_factor=self.VISION_TCF, + enable_fps_modulation=True, + start_frame_offset=1, + ) + fast, _ = compute_mrope_position_ids_action(4, action_fps=24.0, **kwargs) + slow, _ = compute_mrope_position_ids_action(4, action_fps=12.0, **kwargs) + torch.testing.assert_close(slow[0], fast[0] * 2.0, rtol=0, atol=1e-5) + + +class TestVisionMropeBaseCompressionDefault: + """``base_temporal_compression_factor=None`` must not disturb vision or audio.""" + + def test_vision_positions_unchanged_by_default(self): + kwargs = dict( + temporal_offset=3.0, + fps=30.0, + base_fps=24.0, + temporal_compression_factor=4, + enable_fps_modulation=True, + ) + implicit, next_implicit = compute_mrope_position_ids_vision(5, 2, 2, **kwargs) + explicit, next_explicit = compute_mrope_position_ids_vision( + 5, 2, 2, base_temporal_compression_factor=4, **kwargs + ) + torch.testing.assert_close(implicit, explicit, rtol=0, atol=0) + assert next_implicit == next_explicit + + def test_audio_style_call_unchanged(self): + """Audio packs with tcf=1 and no base override (sound base tcf is also 1).""" + ids, _ = compute_mrope_position_ids_vision( + 3, + 1, + 1, + temporal_offset=0.0, + fps=25.0, + base_fps=24.0, + temporal_compression_factor=1, + enable_fps_modulation=True, + ) + expected = _reference_scaled_positions( + grid_t=3, + temporal_offset=0.0, + fps=25.0, + base_fps=24.0, + temporal_compression_factor=1, + base_temporal_compression_factor=1, + start_frame_offset=0, + ) + torch.testing.assert_close( + ids[0], torch.tensor(expected, dtype=ids.dtype), rtol=0, atol=1e-5 + ) + + +class TestPrepareActionLatents: + """The CPU half of the action contract: which tokens start clean, what the + mask says, and how a caller's trajectory is fitted to the chunk.""" + + ACTION_DIM = 8 + CHUNK = 4 + + def _prepare(self, mode, **kwargs): + kwargs.setdefault("action_chunk_size", self.CHUNK) + kwargs.setdefault("action_dim", self.ACTION_DIM) + return prepare_action_latents( + mode=mode, + generator=torch.Generator(device="cpu").manual_seed(0), + device=torch.device("cpu"), + dtype=torch.float32, + **kwargs, + ) + + def test_forward_dynamics_pads_a_short_trajectory_by_holding_the_last_step(self): + latents, _, clean, raw_dim = self._prepare( + "forward_dynamics", raw_action_dim=None, action_input=[[1.0, 2.0], [3.0, 4.0]] + ) + assert raw_dim == 2 + assert latents.shape == (1, self.CHUNK, self.ACTION_DIM) + # Steps 2 and 3 repeat the supplied final step rather than going to zero. + for step in range(2, self.CHUNK): + torch.testing.assert_close(clean[0, step, :2], torch.tensor([3.0, 4.0])) + + def test_forward_dynamics_truncates_a_long_trajectory(self): + _, _, clean, _ = self._prepare( + "forward_dynamics", + raw_action_dim=None, + action_input=[[float(i), float(i)] for i in range(self.CHUNK + 3)], + ) + torch.testing.assert_close(clean[0, -1, :2], torch.tensor([3.0, 3.0])) + + def test_forward_dynamics_conditions_every_step(self): + """All action tokens are given, so none carries velocity.""" + latents, mask, clean, _ = self._prepare( + "forward_dynamics", raw_action_dim=None, action_input=[[1.0, 2.0]] * self.CHUNK + ) + assert torch.all(mask == 0.0) + torch.testing.assert_close(latents, clean) + + @pytest.mark.parametrize("mode", ["policy", "inverse_dynamics"]) + def test_predicted_modes_start_from_noise_everywhere(self, mode): + latents, mask, clean, _ = self._prepare(mode, raw_action_dim=2) + assert torch.all(mask == 1.0) + assert torch.all(clean == 0.0) + assert torch.any(latents[..., :2] != 0.0) + + @pytest.mark.parametrize("mode", ["policy", "forward_dynamics", "inverse_dynamics"]) + def test_columns_above_raw_action_dim_are_zero(self, mode): + """The head is action_dim wide but only raw_action_dim is meaningful; + padding must not carry noise into the model or out of it.""" + kwargs = ( + {"raw_action_dim": None, "action_input": [[1.0, 2.0]] * self.CHUNK} + if mode == "forward_dynamics" + else {"raw_action_dim": 2} + ) + latents, _, clean, raw_dim = self._prepare(mode, **kwargs) + assert raw_dim == 2 + assert torch.all(latents[..., raw_dim:] == 0.0) + assert torch.all(clean[..., raw_dim:] == 0.0) + + def test_empty_trajectory_raises(self): + """Without this the empty slice reaches the mask broadcast and dies with + an opaque torch shape error instead of a client-facing one.""" + with pytest.raises(ValueError, match="at least one timestep"): + self._prepare("forward_dynamics", raw_action_dim=None, action_input=torch.zeros(0, 2)) + + @pytest.mark.parametrize("raw_action_dim", [0, -1, ACTION_DIM + 1]) + def test_out_of_range_raw_action_dim_raises(self, raw_action_dim): + with pytest.raises(ValueError, match=r"raw_action_dim must be in \[1, \d+\]"): + self._prepare("policy", raw_action_dim=raw_action_dim) + + def test_forward_dynamics_raw_dim_mismatch_raises(self): + with pytest.raises(ValueError, match="raw_action_dim must match"): + prepare_action_latents( + 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]], + ) + + +class TestResizeAndPadActionImage: + """Action pads to the canvas where V2V crops to it: a gripper works at the + frame edge, so cover-scale would cut away what the policy acts on.""" + + def test_contain_scale_then_pad_to_canvas(self): + image = PIL.Image.new("RGB", (800, 400), "blue") + out = resize_and_pad_action_image(image, target_h=480, target_w=832) + assert (out.height, out.width) == (480, 832) + + def test_small_source_is_never_enlarged(self): + """min(..., 1.0): a small clip keeps its own pixels and a wider border + rather than being upscaled into blur.""" + image = PIL.Image.new("RGB", (100, 50), "blue") + out = resize_and_pad_action_image(image, target_h=480, target_w=832) + assert (out.height, out.width) == (480, 832) + assert np.asarray(out)[:50, :100].any() + + def test_exact_size_is_returned_unchanged(self): + image = PIL.Image.new("RGB", (832, 480), "blue") + assert resize_and_pad_action_image(image, 480, 832).size == (832, 480) + + def test_aspect_ratio_is_preserved(self): + """Contain-scale keeps the source's own aspect; the leftover strip is + padding, not stretch.""" + image = PIL.Image.new("RGB", (400, 400), "blue") + out = np.asarray(resize_and_pad_action_image(image, target_h=480, target_w=832)) + assert out.shape[:2] == (480, 832) + # A square source contained in a 16:9 canvas fills the height, so the + # scaled content is square and the pad lands on the right. + assert out[:480, :480].any() + + +class TestResolveDomainId: + """domain_id wins, but a caller that contradicts itself is a real mistake: + the wrong embodiment yields a fluent trajectory in another robot's dialect.""" + + def test_agreeing_pair_is_accepted(self): + assert resolve_domain_id(domain_id=7, domain_name="bridge_orig_lerobot") == 7 + + def test_contradicting_pair_raises(self): + with pytest.raises(ValueError, match="contradicts domain_name"): + resolve_domain_id(domain_id=20, domain_name="bridge_orig_lerobot") + + def test_unlisted_name_leaves_domain_id_authoritative(self): + assert resolve_domain_id(domain_id=31, domain_name="some-new-robot") == 31 + + def test_name_alone_still_resolves(self): + assert resolve_domain_id(domain_name="fractal") == 20 + + def test_negative_domain_id_raises(self): + with pytest.raises(ValueError, match="must be non-negative"): + resolve_domain_id(domain_id=-1) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py index 02280c2f4419..368af6943542 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py @@ -414,6 +414,54 @@ def test_explicit_values_pass_through(self): assert got["num_inference_steps"] == 20 assert got["width"] == COSMOS3_T2I_PARAMS["width"] + def test_action_keeps_every_mode_field_unset(self): + """Action resolves its own canvas, steps, guidance and frame rate from + the embodiment preset. Filling them from the video table here would run + every action request at 720p, 35 steps and guidance 6 (CFG on).""" + req = _fake_request("video", extra_params={"action_mode": "policy"}) + got = self._captured_forward_kwargs(_bare_pipeline(), req) + for field in ("height", "width", "num_inference_steps", "guidance_scale", "frame_rate"): + assert got[field] is None, field + + def test_video_keeps_its_materialised_frame_rate(self): + """Only action drops it: the serve layer derives num_frames from + seconds x frame_rate, so the video default has to stay materialised.""" + got = self._captured_forward_kwargs(_bare_pipeline(), _fake_request("video")) + assert got["frame_rate"] == COSMOS3_720P_PARAMS["frame_rate"] + + def test_action_keeps_an_explicit_frame_rate(self): + """Dropping the materialised default is what lets the embodiment preset + win; dropping a caller's own value would make it unsettable.""" + req = _fake_request("video", frame_rate=30.0, extra_params={"action_mode": "policy"}) + got = self._captured_forward_kwargs(_bare_pipeline(), req) + assert got["frame_rate"] == 30.0 + + def test_action_honors_an_explicit_frame_rate_equal_to_the_video_default(self): + """The collision case: 24.0 is both a legal caller choice and the + materialized video default. Value equality reads it as unset and hands + back the embodiment preset; provenance keeps the caller's 24.""" + req = _fake_request("video", frame_rate=24.0, extra_params={"action_mode": "policy"}) + got = self._captured_forward_kwargs(_bare_pipeline(), req) + assert got["frame_rate"] == 24.0 + + def test_action_drops_a_frame_rate_the_caller_never_set(self): + req = _fake_request("video", extra_params={"action_mode": "policy"}) + got = self._captured_forward_kwargs(_bare_pipeline(), req) + assert got["frame_rate"] is None + + def test_action_honors_explicit_values_equal_to_their_defaults(self): + """Same hazard for the other mode-dependent fields.""" + req = _fake_request( + "video", + height=COSMOS3_720P_PARAMS["height"], + guidance_scale=COSMOS3_720P_PARAMS["guidance_scale"], + extra_params={"action_mode": "policy"}, + ) + got = self._captured_forward_kwargs(_bare_pipeline(), req) + assert got["height"] == COSMOS3_720P_PARAMS["height"] + assert got["guidance_scale"] == COSMOS3_720P_PARAMS["guidance_scale"] + assert got["width"] is None + def test_distilled_merged_defaults_pass_through(self): req = _fake_request("image", num_inference_steps=4, guidance_scale=1.0) got = self._captured_forward_kwargs(_bare_pipeline(sampling=_distilled_policy()), req) @@ -766,9 +814,10 @@ def test_anchor_writes_only_frame_zero_in_place(self): latents = torch.arange(48, dtype=torch.float32).reshape(1, 4, 3, 2, 2) untouched = latents[:, :, 1:].clone() - returned = post_step_fn(latents) + returned, extra = post_step_fn(latents, None) assert returned is latents, "must write in place, not copy" + assert extra is None, "extra-stream latents pass through untouched" assert torch.all(latents[:, :, 0:1] == self.CLEAN) assert torch.equal(latents[:, :, 1:], untouched) @@ -782,7 +831,7 @@ def test_anchor_rejects_dtype_mismatch(self): latents = torch.zeros(1, 4, 3, 2, 2, dtype=torch.float32) with pytest.raises(RuntimeError, match="must match the denoised latents"): - post_step_fn(latents) + post_step_fn(latents, None) def test_anchor_accepts_matching_dtype(self): pipeline = _bare_pipeline(sampling=_distilled_policy()) @@ -791,7 +840,7 @@ def test_anchor_accepts_matching_dtype(self): ) latents = torch.zeros(1, 4, 3, 2, 2, dtype=torch.bfloat16) - post_step_fn(latents) + post_step_fn(latents, None) assert torch.all(latents[:, :, 0:1] == self.CLEAN) def _run_denoise(self, with_anchor: bool): @@ -922,7 +971,7 @@ def test_i2v_request_wires_anchor_and_seeded_steps(self): post_step_fn = captured["post_step_fn"] assert post_step_fn is not None latents = torch.zeros(1, 4, self.T_LAT, self.H_LAT, self.W_LAT) - post_step_fn(latents) + post_step_fn(latents, None) assert torch.all(latents[:, :, 0:1] == self.CLEAN) assert torch.all(latents[:, :, 1:] == 0.0) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py b/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py index ffe3777774d4..9feed9e2dad2 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py @@ -164,6 +164,20 @@ def _synthetic_state_dict(cfg: SimpleNamespace) -> dict: "time_embedder.linear_2.weight": torch.randn(h, h), "time_embedder.linear_2.bias": torch.randn(h), } + if getattr(cfg, "action_gen", False): + # DomainAwareLinear stores per-domain weights as nn.Embedding rows of + # flattened [out * in] matrices; the modality embed is a root parameter. + n_dom = cfg.num_embodiment_domains + a = cfg.action_dim + sd.update( + { + "action_modality_embed": torch.randn(h), + "action_proj_in.fc.weight": torch.randn(n_dom, h * a), + "action_proj_in.bias.weight": torch.randn(n_dom, h), + "action_proj_out.fc.weight": torch.randn(n_dom, a * h), + "action_proj_out.bias.weight": torch.randn(n_dom, a), + } + ) if getattr(cfg, "sound_gen", False): sd.update( { @@ -611,6 +625,8 @@ def test_full_checkpoint_loads(self): "missing_key", [ "layers.0.self_attn.k_norm_und_for_gen.weight", + "action_proj_in.fc.weight", + "action_modality_embed", "layers.0.mlp.up_proj.weight", "layers.1.mlp_moe_gen.down_proj.weight", "layers.0.input_layernorm_moe_gen.weight", @@ -638,15 +654,6 @@ def test_intentional_skips_are_logged_with_names(self, monkeypatch): monkeypatch.setattr(tf_module.logger, "info", infos.append) cfg = _reduced_edge_config() sd = _edge_state_dict(cfg) - sd.update( - { - "action_modality_embed": torch.randn(cfg.hidden_size), - "action_proj_in.fc.weight": torch.randn(4, 8), - "action_proj_in.bias.weight": torch.randn(4, 8), - "action_proj_out.fc.weight": torch.randn(4, 8), - "action_proj_out.bias.weight": torch.randn(4, 8), - } - ) model = self._model() model.load_weights(sd) @@ -659,9 +666,6 @@ def test_intentional_skips_are_logged_with_names(self, monkeypatch): # text also mentions lm_head/norm, so assert the parsed set exactly. skipped_families = {name.strip() for name in skip_logs[0].rsplit(": ", 1)[1].split(",")} assert skipped_families == { - "action_modality_embed", - "action_proj_in", - "action_proj_out", "lm_head", "norm", } @@ -679,14 +683,13 @@ def test_model_prefixed_skip_keys_are_intentional(self, monkeypatch): cfg = _reduced_edge_config() sd = _edge_state_dict(cfg) sd["model.lm_head.weight"] = torch.randn(cfg.vocab_size, cfg.hidden_size) - sd["model.action_modality_embed"] = torch.randn(cfg.hidden_size) self._model().load_weights(sd) assert not any("unknown checkpoint key" in m for m in warnings) skip_logs = [m for m in infos if "intentionally unused" in m] assert len(skip_logs) == 1 skipped_families = {name.strip() for name in skip_logs[0].rsplit(": ", 1)[1].split(",")} - assert {"lm_head", "action_modality_embed"} <= skipped_families + assert "lm_head" in skipped_families def test_unconsumed_mapped_tensor_warns(self, monkeypatch): """A checkpoint tensor that remaps to a module the recipe didn't diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index 5c4fd0df80c8..f49af5e7473f 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -39,6 +39,7 @@ import tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 as pipe_mod from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( + COSMOS3_ACTION_PARAMS, COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP, COSMOS3_DEFAULT_CONDITION_VIDEO_LATENT_INDEXES, COSMOS3_EXTRA_SPECS, @@ -147,6 +148,7 @@ def _run_forward( height=HEIGHT, width=WIDTH, guidance_scale=GUIDANCE_SCALE, + frame_rate=FRAME_RATE, **extra, ): return pipeline.forward( @@ -158,7 +160,7 @@ def _run_forward( num_inference_steps=NUM_STEPS, guidance_scale=guidance_scale, seed=SEED, - frame_rate=FRAME_RATE, + frame_rate=frame_rate, use_guardrails=False, **extra, ) @@ -231,6 +233,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 _scheduler_use_karras_sigmas(scheduler) -> bool | None: value = getattr(scheduler.config, "use_karras_sigmas", None) return None if value is None else bool(value) @@ -981,6 +1001,153 @@ def fake_tokenize_prompt(text, max_sequence_length, use_system_prompt, system_pr assert rebuilt == [("video", 10.0, False), ("audio", 10.0, False)] + def test_action_video_is_not_classified_as_v2v(self): + """An action reference arrives as the same `video` bytes V2V uses, but + it is an observation, not a clip to continue. Treating it as V2V forces + the system prompt, so the same frame would tokenize differently + depending on whether it was passed as an image or a one-frame clip.""" + pipeline = Cosmos3OmniMoTPipeline.__new__(Cosmos3OmniMoTPipeline) + pipeline.transformer = SimpleNamespace( + device=torch.device("cpu"), num_embodiment_domains=32 + ) + pipeline.audio_gen = False + pipeline.action_gen = True + pipeline.default_use_system_prompt = False + pipeline.family = QWEN3_RECIPE.name + token_calls = [] + + class StopAfterTokenize(Exception): + pass + + class FakeSampling: + is_distilled = False + checkpoint_flow_shift = 1.0 + + def validate_request(self, num_inference_steps, guidance_scale): + return None + + def generation_default_overrides(self): + return {} + + def set_flow_shift(self, scheduler, target, *, use_karras_sigmas=None): + return scheduler + + def fake_tokenize_prompt(text, max_sequence_length, use_system_prompt, system_prompt=None): + token_calls.append(use_system_prompt) + raise StopAfterTokenize + + pipeline.scheduler = SimpleNamespace(config=SimpleNamespace(flow_shift=1.0)) + pipeline.sampling = FakeSampling() + pipeline._tokenize_prompt = fake_tokenize_prompt + + with pytest.raises(StopAfterTokenize): + pipeline.forward( + prompt="pick up the block", + video=_V2V_FIXTURE_MP4.read_bytes(), + num_frames=NUM_FRAMES, + num_inference_steps=1, + guidance_scale=1.0, + seed=1, + max_sequence_length=8, + use_system_prompt=None, + use_guardrails=False, + action_mode="inverse_dynamics", + domain_name="bridge_orig_lerobot", + raw_action_dim=10, + action_chunk_size=NUM_FRAMES - 1, + ) + + assert token_calls[0] is False + + def test_action_restores_the_checkpoint_flow_shift(self): + """set_flow_shift is a no-op when both knobs are None, and the pipeline + instance outlives the request. An action request that followed a V2V one + would otherwise keep V2V's uniform sigmas instead of the checkpoint's.""" + pipeline = Cosmos3OmniMoTPipeline.__new__(Cosmos3OmniMoTPipeline) + pipeline.transformer = SimpleNamespace( + device=torch.device("cpu"), num_embodiment_domains=32 + ) + pipeline.audio_gen = False + pipeline.action_gen = True + pipeline.default_use_system_prompt = False + pipeline.family = QWEN3_RECIPE.name + calls = [] + + class StopAfterTokenize(Exception): + pass + + class FakeSampling: + is_distilled = False + checkpoint_flow_shift = 7.0 + + def validate_request(self, num_inference_steps, guidance_scale): + return None + + def generation_default_overrides(self): + return {} + + def set_flow_shift(self, scheduler, target, *, use_karras_sigmas=None): + calls.append((scheduler.name, target, use_karras_sigmas)) + return scheduler + + def fake_tokenize_prompt(text, max_sequence_length, use_system_prompt, system_prompt=None): + raise StopAfterTokenize + + pipeline.scheduler = SimpleNamespace(name="video", config=SimpleNamespace(flow_shift=1.0)) + pipeline.action_scheduler = SimpleNamespace( + name="action", config=SimpleNamespace(flow_shift=1.0) + ) + pipeline.sampling = FakeSampling() + pipeline._tokenize_prompt = fake_tokenize_prompt + + with pytest.raises(StopAfterTokenize): + pipeline.forward( + prompt="pick up the block", + video=_V2V_FIXTURE_MP4.read_bytes(), + num_frames=NUM_FRAMES, + num_inference_steps=1, + guidance_scale=1.0, + seed=1, + max_sequence_length=8, + use_guardrails=False, + action_mode="inverse_dynamics", + domain_name="bridge_orig_lerobot", + raw_action_dim=10, + action_chunk_size=NUM_FRAMES - 1, + ) + + assert calls == [("video", 7.0, None), ("action", 7.0, None)] + + def test_scheduler_for_serves_every_stream(self): + """Video, audio and action denoise in lockstep, so one resolved + (shift, karras) configuration must yield a scheduler per stream -- + separate instances (schedulers mutate state on every .step()) built + from that stream's own base.""" + pipeline = Cosmos3OmniMoTPipeline.__new__(Cosmos3OmniMoTPipeline) + pipeline.family = QWEN3_RECIPE.name + rebuilt = [] + + class FakeSampling: + checkpoint_flow_shift = 1.0 + + def set_flow_shift(self, scheduler, target, *, use_karras_sigmas=None): + rebuilt.append((scheduler.name, target, use_karras_sigmas)) + return scheduler + + pipeline.sampling = FakeSampling() + pipeline.scheduler = SimpleNamespace(name="video") + pipeline.audio_scheduler = SimpleNamespace(name="audio") + pipeline.action_scheduler = SimpleNamespace(name="action") + + for stream in ("video", "audio", "action"): + pipeline._scheduler_for(10.0, False, stream=stream) + + assert rebuilt == [ + ("video", 10.0, False), + ("audio", 10.0, False), + ("action", 10.0, False), + ] + def test_image_and_video_rejected(self, cosmos3_pipeline): with pytest.raises(ValueError, match="not both image and video"): _run_forward( @@ -1165,6 +1332,208 @@ def test_v2v_audio_smoke(self, cosmos3_pipeline): ) +@pytest.mark.integration +@pytest.mark.cosmos3_action +@pytest.mark.high_cuda_memory +class TestCosmos3Action: + ACTION_HEIGHT = 480 + ACTION_WIDTH = 832 + ACTION_CHUNK = COSMOS3_ACTION_PARAMS["action_chunk_size"] + # Derived, not configured: both references fix the clip at chunk + 1. + ACTION_FRAMES = ACTION_CHUNK + 1 + RAW_ACTION_DIM = 10 + + def test_policy_smoke(self, cosmos3_pipeline): + _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, + ) + + def test_forward_dynamics_smoke(self, cosmos3_pipeline): + _require_action_pipeline(cosmos3_pipeline) + image = _make_test_image().resize((self.ACTION_WIDTH, self.ACTION_HEIGHT)) + action_traj = [[0.1] * self.RAW_ACTION_DIM for _ in range(self.ACTION_CHUNK)] + result = _run_forward( + cosmos3_pipeline, + image=image, + height=self.ACTION_HEIGHT, + width=self.ACTION_WIDTH, + num_frames=self.ACTION_FRAMES, + guidance_scale=COSMOS3_ACTION_PARAMS["guidance_scale"], + action_mode="forward_dynamics", + domain_name="bridge_orig_lerobot", + action=action_traj, + action_chunk_size=self.ACTION_CHUNK, + ) + _assert_valid_video( + result.video, + num_frames=self.ACTION_FRAMES, + height=self.ACTION_HEIGHT, + width=self.ACTION_WIDTH, + ) + _assert_valid_action( + result.action, + raw_action_dim=self.RAW_ACTION_DIM, + chunk_size=self.ACTION_CHUNK, + ) + + def test_inverse_dynamics_smoke(self, cosmos3_pipeline): + _require_action_pipeline(cosmos3_pipeline) + 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, + # The clip is chunk + 1 frames, and the fixture holds NUM_FRAMES. + action_chunk_size=NUM_FRAMES - 1, + video=_V2V_FIXTURE_MP4.read_bytes(), + ) + _assert_valid_video( + result.video, + 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 - 1, + ) + + def test_inverse_dynamics_rejects_short_video(self, cosmos3_pipeline): + _require_action_pipeline(cosmos3_pipeline) + with pytest.raises(ValueError, match=r"requires \d+ frames at"): + _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=_V2V_FIXTURE_MP4.read_bytes(), + ) + + def test_inverse_dynamics_thins_to_the_requested_rate(self, cosmos3_pipeline): + """A 24 fps reference asked for at 5 fps keeps every 5th frame, so the + window widens accordingly and the 9-frame fixture comes up short.""" + _require_action_pipeline(cosmos3_pipeline) + with pytest.raises(ValueError, match=r"thinned by 5 \(41 source frames needed\)"): + _run_forward( + cosmos3_pipeline, + image=None, + height=self.ACTION_HEIGHT, + width=self.ACTION_WIDTH, + num_frames=NUM_FRAMES, + guidance_scale=COSMOS3_ACTION_PARAMS["guidance_scale"], + frame_rate=5.0, + action_mode="inverse_dynamics", + domain_name="bridge_orig_lerobot", + raw_action_dim=self.RAW_ACTION_DIM, + action_chunk_size=NUM_FRAMES - 1, + video=_V2V_FIXTURE_MP4.read_bytes(), + ) + + def test_out_of_range_domain_id_rejected_before_decode(self, cosmos3_pipeline): + _require_action_pipeline(cosmos3_pipeline) + with pytest.raises(ValueError, match=r"domain_id must be in \[0, \d+\)"): + _run_forward( + cosmos3_pipeline, + image=_make_test_image(), + height=self.ACTION_HEIGHT, + width=self.ACTION_WIDTH, + num_frames=self.ACTION_FRAMES, + guidance_scale=COSMOS3_ACTION_PARAMS["guidance_scale"], + action_mode="policy", + domain_id=10_000, + raw_action_dim=self.RAW_ACTION_DIM, + action_chunk_size=self.ACTION_CHUNK, + ) + + def test_first_frame_failure_is_synchronized(self, cosmos3_pipeline, monkeypatch): + """Every rank decodes its own reference. A failure on one rank has to + reach the others before the transformer's collectives, or the job hangs.""" + _require_action_pipeline(cosmos3_pipeline) + from tensorrt_llm._torch.visual_gen.models.cosmos3 import pipeline_cosmos3 + + seen = [] + real = pipeline_cosmos3.synchronize_media_prepare_status + + def spy(error): + seen.append(error) + return real(error) + + monkeypatch.setattr(pipeline_cosmos3, "synchronize_media_prepare_status", spy) + + with pytest.raises(Exception): + _run_forward( + cosmos3_pipeline, + image="/nonexistent/action_reference_frame.png", + height=self.ACTION_HEIGHT, + width=self.ACTION_WIDTH, + num_frames=self.ACTION_FRAMES, + guidance_scale=COSMOS3_ACTION_PARAMS["guidance_scale"], + action_mode="policy", + domain_name="bridge_orig_lerobot", + raw_action_dim=self.RAW_ACTION_DIM, + action_chunk_size=self.ACTION_CHUNK, + ) + + assert seen and isinstance(seen[0], Exception) + + def test_action_and_audio_rejected(self, cosmos3_pipeline): + _require_action_pipeline(cosmos3_pipeline) + with pytest.raises(ValueError, match="joint action and audio"): + _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 6041300cdf6e..5a53683ab9f1 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py @@ -401,6 +401,289 @@ 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 / _NUM_TRAIN_TIMESTEPS, + raw_timestep=ts, + text_ids=text_ids, + text_mask=text_mask, + video_shape=video_shape, + fps=24.0, + action_latents=action_latents, + action_domain_ids=domain_ids, + ) + _assert_finite_output(out.video, hs.shape) + assert out.action is not None + _assert_finite_output(out.action, torch.Size([1, self.T_ACTION, model.action_dim])) + + @pytest.mark.high_cuda_memory + def test_domain_ids_validated_once_per_request(self, action_model_config): + """The range check reads a device tensor, so it is a blocking sync. It + belongs on the first step of a request, not on every denoise step.""" + cfg = action_model_config.pretrained_config + model = _build_random_weight_model(action_model_config) + hs, ts, text_ids, text_mask, video_shape = _cosmos3_inputs( + DEVICE, channels=cfg.latent_channel + ) + action_latents = torch.randn(1, self.T_ACTION, model.action_dim, device=DEVICE, dtype=DTYPE) + domain_ids = torch.tensor([7], dtype=torch.long, device=DEVICE) + + calls = [] + real_validate = model.action_proj_in.validate_domain_ids + model.action_proj_in.validate_domain_ids = lambda ids: ( + calls.append(ids), + real_validate(ids), + )[1] + + def run_step(): + with torch.inference_mode(): + model( + hidden_states=hs, + timestep=ts / _NUM_TRAIN_TIMESTEPS, + raw_timestep=ts, + text_ids=text_ids, + text_mask=text_mask, + video_shape=video_shape, + fps=24.0, + action_latents=action_latents, + action_domain_ids=domain_ids, + ) + + run_step() + run_step() + assert len(calls) == 1 + + model.reset_cache() + run_step() + assert len(calls) == 2 + + def test_graph_key_separates_requests_that_differ_only_in_scalars(self, action_model_config): + """TRT-LLM captures a family of graphs and dispatches by key. fps, the + action clock and the start offset change the rotary positions without + changing any tensor shape, so they must discriminate keys or two such + requests would replay the same graph.""" + from tensorrt_llm._torch.visual_gen.cuda_graph_runner import ( + CUDAGraphRunner, + CUDAGraphRunnerConfig, + ) + + model = Cosmos3VFMTransformer(model_config=action_model_config) + runner = CUDAGraphRunner(CUDAGraphRunnerConfig(use_cuda_graph=True)) + model.register_cuda_graph_extra_key_fns(runner) + + base = dict(fps=24.0, action_fps=5.0, action_start_frame_offset=1) + key = runner.get_graph_key(**base) + for field, other in ( + ("fps", 16.0), + ("action_fps", 10.0), + ("action_start_frame_offset", 0), + ): + assert runner.get_graph_key(**{**base, field: other}) != key, field + + # A video-only request keys exactly as before: absent scalars drop out. + assert runner.get_graph_key(fps=None, action_fps=None) == runner.get_graph_key() + + @pytest.mark.high_cuda_memory + def test_action_rope_table_built_once_per_request(self, action_model_config): + """Chunk size, prompt lengths, fps and the frame offset are fixed for a + request, so the rotary table is too. Rebuilding it per step costs a + device-to-host sync per batch element plus an H2D copy of the position + ids -- for identical numbers.""" + cfg = action_model_config.pretrained_config + model = _build_random_weight_model(action_model_config) + hs, ts, text_ids, text_mask, video_shape = _cosmos3_inputs( + DEVICE, channels=cfg.latent_channel + ) + action_latents = torch.randn(1, self.T_ACTION, model.action_dim, device=DEVICE, dtype=DTYPE) + domain_ids = torch.tensor([7], dtype=torch.long, device=DEVICE) + + calls = [] + real = model._compute_action_rope_freqs + model._compute_action_rope_freqs = lambda *a, **k: (calls.append(1), real(*a, **k))[1] + + def run_step(): + with torch.inference_mode(): + model( + hidden_states=hs, + timestep=ts / _NUM_TRAIN_TIMESTEPS, + raw_timestep=ts, + text_ids=text_ids, + text_mask=text_mask, + video_shape=video_shape, + fps=24.0, + action_latents=action_latents, + action_domain_ids=domain_ids, + ) + + run_step() + run_step() + run_step() + assert len(calls) == 1 + + model.reset_cache() + run_step() + assert len(calls) == 2 + + @pytest.mark.high_cuda_memory + def test_forward_with_action_domain_id_out_of_range_raises(self, action_model_config): + cfg = action_model_config.pretrained_config + 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=r"domain_id must be in \[0, \d+\)"), + ): + model( + hidden_states=hs, + timestep=ts / _NUM_TRAIN_TIMESTEPS, + raw_timestep=ts, + text_ids=text_ids, + text_mask=text_mask, + video_shape=video_shape, + fps=24.0, + action_latents=action_latents, + action_domain_ids=domain_ids, + ) + + @pytest.mark.high_cuda_memory + def test_forward_without_action_latents_returns_none(self, action_model_config): + cfg = action_model_config.pretrained_config + model = _build_random_weight_model(action_model_config) + hs, ts, text_ids, text_mask, video_shape = _cosmos3_inputs( + DEVICE, channels=cfg.latent_channel + ) + with torch.inference_mode(): + out = model( + hidden_states=hs, + timestep=ts / _NUM_TRAIN_TIMESTEPS, + raw_timestep=ts, + text_ids=text_ids, + text_mask=text_mask, + video_shape=video_shape, + ) + _assert_finite_output(out.video, hs.shape) + assert out.action is None + + @pytest.mark.high_cuda_memory + def test_forward_with_action_noisy_mask(self, action_model_config): + cfg = action_model_config.pretrained_config + model = _build_random_weight_model(action_model_config) + hs, ts, text_ids, text_mask, video_shape = _cosmos3_inputs( + DEVICE, channels=cfg.latent_channel, t=2 + ) + action_latents = torch.randn(1, self.T_ACTION, model.action_dim, device=DEVICE, dtype=DTYPE) + noisy_mask = torch.ones(1, self.T_ACTION, 1, device=DEVICE, dtype=DTYPE) + noisy_mask[:, 0, :] = 0.0 + domain_ids = torch.tensor([7], dtype=torch.long, device=DEVICE) + with torch.inference_mode(): + out = model( + hidden_states=hs, + timestep=ts / _NUM_TRAIN_TIMESTEPS, + raw_timestep=ts, + text_ids=text_ids, + text_mask=text_mask, + video_shape=video_shape, + fps=24.0, + action_latents=action_latents, + action_domain_ids=domain_ids, + action_noisy_mask=noisy_mask, + ) + _assert_finite_output(out.video, hs.shape) + _assert_finite_output(out.action, torch.Size([1, self.T_ACTION, model.action_dim])) + + @pytest.mark.high_cuda_memory + def test_forward_with_action_multiframe(self, action_model_config): + cfg = action_model_config.pretrained_config + model = _build_random_weight_model(action_model_config) + hs, ts, text_ids, text_mask, video_shape = _cosmos3_inputs( + DEVICE, channels=cfg.latent_channel, t=3 + ) + action_latents = torch.randn(1, self.T_ACTION, model.action_dim, device=DEVICE, dtype=DTYPE) + domain_ids = torch.tensor([7], dtype=torch.long, device=DEVICE) + with torch.inference_mode(): + out = model( + hidden_states=hs, + timestep=ts / _NUM_TRAIN_TIMESTEPS, + raw_timestep=ts, + text_ids=text_ids, + text_mask=text_mask, + video_shape=video_shape, + fps=24.0, + action_latents=action_latents, + action_domain_ids=domain_ids, + ) + _assert_finite_output(out.video, hs.shape) + _assert_finite_output(out.action, torch.Size([1, self.T_ACTION, model.action_dim])) + + @pytest.mark.integration class TestCosmos3TransformerCheckpoint: """Load Cosmos3-Nano transformer weights and run a single forward step.""" diff --git a/tests/unittest/_torch/visual_gen/test_media_decode.py b/tests/unittest/_torch/visual_gen/test_media_decode.py index 7fc9f3a84d6f..0c88f47ff86b 100644 --- a/tests/unittest/_torch/visual_gen/test_media_decode.py +++ b/tests/unittest/_torch/visual_gen/test_media_decode.py @@ -30,6 +30,7 @@ _lanczos_taps, decode_video_reference_window, resize_center_crop_uint8, + resize_fit_pad_uint8, ) _TEST_DATA = Path(__file__).parent / "test_data" @@ -42,6 +43,49 @@ def _frame_indices(frames: torch.Tensor) -> list[int]: return [round((f[:, :, 0].float().mean().item() - 20) / 25) for f in frames] +class TestResizeFitPad: + """The action counterpart to cover-scale + center-crop: contain-scale, then + pad. A gripper works at the frame edge, so cropping costs the policy the + evidence it acts on.""" + + def test_contains_the_whole_source(self): + """Cover-scale would crop the wide source; fit keeps all of it.""" + frames = torch.full((2, 100, 400, 3), 200, dtype=torch.uint8) + out = resize_fit_pad_uint8(frames, target_h=200, target_w=400) + assert out.shape == (2, 200, 400, 3) + # 400x100 contained in 400x200 scales by 1.0 and pads the bottom half. + assert (out[:, :100] == 200).all() + + def test_never_enlarges_a_small_source(self): + """min(..., 1.0): a small clip keeps its own pixels and a wider border + rather than being upscaled.""" + frames = torch.full((1, 10, 20, 3), 255, dtype=torch.uint8) + out = resize_fit_pad_uint8(frames, target_h=64, target_w=64) + assert out.shape == (1, 64, 64, 3) + assert (out[0, :10, :20] == 255).all() + + def test_native_resolution_is_identity(self): + frames = torch.zeros(2, 32, 32, 3, dtype=torch.uint8) + assert resize_fit_pad_uint8(frames, 32, 32) is frames + + def test_aspect_ratio_is_preserved(self): + """The source is contained, not stretched: a square stays square.""" + frames = torch.zeros(1, 64, 64, 3, dtype=torch.uint8) + frames[:, :, :, 0] = 255 + out = resize_fit_pad_uint8(frames, target_h=64, target_w=128) + assert out.shape == (1, 64, 128, 3) + # A 1:1 source in a 2:1 canvas fills the height, so content is 64 wide. + assert (out[0, :, :64, 0] == 255).all() + + def test_pad_falls_back_to_replicate_when_reflection_has_no_source(self): + """A pad run wider than the resized extent has nothing left to mirror; + reflect would raise, so the helper switches to edge replication.""" + frames = torch.full((1, 4, 4, 3), 128, dtype=torch.uint8) + out = resize_fit_pad_uint8(frames, target_h=64, target_w=64) + assert out.shape == (1, 64, 64, 3) + assert (out[0, 4:, :4] == 128).all() + + class TestResizeCenterCrop: """CPU-runnable checks of the shared Lanczos resize/crop front.""" @@ -203,6 +247,48 @@ def test_keep_last_ring_reorder(self, fixture): window = self._decode(fixture.read_bytes(), keep="last") assert _frame_indices(window) == [4, 5, 6, 7, 8] + def test_frame_step_thins_the_window(self): + window = decode_video_reference_window( + _MP4.read_bytes(), + first_frame=0, + last_frame=8, + target_h=64, + target_w=64, + device=self._DEVICE, + frame_step=2, + ) + assert window.shape[0] == 5 + assert _frame_indices(window) == [0, 2, 4, 6, 8] + + def test_frame_step_keeps_a_partial_tail(self): + # The range end is not a multiple of the step: 0, 3, 6 and stop. + window = decode_video_reference_window( + _MP4.read_bytes(), + first_frame=0, + last_frame=7, + target_h=64, + target_w=64, + device=self._DEVICE, + frame_step=3, + ) + assert _frame_indices(window) == [0, 3, 6] + + def test_frame_step_one_matches_the_default(self): + span = dict(first_frame=0, last_frame=4, target_h=64, target_w=64, device=self._DEVICE) + assert torch.equal( + decode_video_reference_window(_MP4.read_bytes(), **span), + decode_video_reference_window(_MP4.read_bytes(), frame_step=1, **span), + ) + + def test_frame_step_rejects_trailing_windows(self): + # The trailing form wraps a ring whose length is unknown until EOS. + with pytest.raises(ValueError, match="non-negative ranges"): + self._decode(_MP4.read_bytes(), keep="last", frame_step=2) + + def test_frame_step_must_be_positive(self): + with pytest.raises(ValueError, match="at least 1"): + self._decode(_MP4.read_bytes(), frame_step=0) + def test_rgb_channel_layout(self): # Frame i carries a green horizontal bar at rows [7i, 7i+7) and a # blue vertical bar at cols [7i, 7i+7): asserts the NVDEC output is diff --git a/tests/unittest/_torch/visual_gen/test_tensor_payload.py b/tests/unittest/_torch/visual_gen/test_tensor_payload.py index 6338f58989d0..744efe98762b 100644 --- a/tests/unittest/_torch/visual_gen/test_tensor_payload.py +++ b/tests/unittest/_torch/visual_gen/test_tensor_payload.py @@ -54,6 +54,12 @@ 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) + + class TestIsTensorFormat: def test_accepts_supported_tokens(self): for token in TENSOR_FORMATS: @@ -95,11 +101,65 @@ 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_carries_no_request_metadata(self, fmt): + """The trajectory's own shape states its DOF, and the mode and + embodiment are the caller's request. The payload stays model-agnostic: + media tensors plus rates, nothing Cosmos3-shaped.""" + output = _make_action_output(batch=1, t=4, action_dim=7) + loaded = self._load(serialize_visual_gen_output(output, fmt, batch_index=0), fmt) + assert loaded["action"].shape == (4, 7) + assert set(loaded) == {"action"} + + @pytest.mark.parametrize("fmt", ["safetensors", "pt"]) 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 f238019b2333..1a35d3161abc 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py @@ -17,6 +17,7 @@ import base64 import json import os +import threading import time from io import BytesIO from pathlib import Path @@ -360,6 +361,28 @@ def result(self, timeout=None): # --------------------------------------------------------------------------- +class _ThreadSettlingTestClient(TestClient): + """TestClient that outwaits starlette's FileResponse reader thread. + + pytest-threadleak samples running threads at the end of a test's call + phase. ``FileResponse`` reads the download on an anyio pool thread + ("AnyIO worker thread") that is told to stop when the request's portal + closes but exits a few milliseconds later -- inside the sampling window, + so any file-shipping test in this module can fail as a phantom leak + (pipeline #52355 did). Joining here, still in the call phase, waits out + those milliseconds deterministically; stop is already queued, so the + deadline is a guard, not an expected wait. + """ + + def request(self, *args, **kwargs): + response = super().request(*args, **kwargs) + deadline = time.monotonic() + 10.0 + for thread in threading.enumerate(): + if thread.name == "AnyIO worker thread": + thread.join(max(0.0, deadline - time.monotonic())) + return response + + def _create_server( generator: MockVisualGen, model_name: str = "test-model", @@ -385,7 +408,7 @@ def _create_server( server_role=ServerRole.VISUAL_GEN, metadata_server_cfg=None, ) - client = TestClient(server.app) + client = _ThreadSettlingTestClient(server.app) # Expose the mock so tests can assert captured generate() arguments. client.mock_gen = generator return client @@ -459,6 +482,28 @@ async def async_video_client(tmp_path): os.environ.pop("TRTLLM_MEDIA_STORAGE_PATH", None) +@pytest.fixture() +def action_video_client(tmp_path): + """Video client whose pipeline declares a tensor-only extra param. + + Stands in for Cosmos3 action: the route must learn "this result needs a + tensor payload" from the spec, never from the parameter's name. + """ + from tensorrt_llm._torch.visual_gen.pipeline import ExtraParamSchema + + gen = MockVisualGen(video_output=_make_dummy_video_tensor()) + specs = { + "action_mode": ExtraParamSchema(type="str", default=None, requires_tensor_output=True), + } + gen.executor.extra_param_specs = specs + type(gen).extra_param_specs = property(lambda self: specs) + os.environ["TRTLLM_MEDIA_STORAGE_PATH"] = str(tmp_path) + client = _create_server(gen) + yield client + os.environ.pop("TRTLLM_MEDIA_STORAGE_PATH", None) + del type(gen).extra_param_specs + + @pytest.fixture() def video_audio_client(tmp_path): """TestClient backed by a MockVisualGen that produces videos with audio.""" @@ -2921,3 +2966,86 @@ async def test_async_file_still_returns_file_response(self, async_video_client): # AVI FileResponse carries ``video/x-msvideo``; the path # branch would have set ``application/json``. assert content.headers["content-type"] == "video/x-msvideo" + + +class TestTensorOnlyFormatResolution: + """A request whose result an encoder cannot carry must not be served as video.""" + + @staticmethod + def _post(client, **body): + return client.post( + "/v1/videos/generations", + json={"prompt": "pick up the block", "size": "64x64", "seconds": 1.0, "fps": 8, **body}, + headers={"content-type": "application/json"}, + ) + + def test_auto_resolves_to_tensor_payload(self, action_video_client): + resp = self._post(action_video_client, extra_params={"action_mode": "policy"}) + assert resp.status_code == 200 + # 'auto' would otherwise have produced an encoded video, silently + # dropping the modality the request was made for. + assert resp.headers["content-type"] == "application/octet-stream" + assert resp.headers["content-disposition"].endswith('.safetensors"') + + def test_explicit_tensor_format_passes_through(self, action_video_client): + resp = self._post(action_video_client, format="pt", extra_params={"action_mode": "policy"}) + assert resp.status_code == 200 + assert resp.headers["content-disposition"].endswith('.pt"') + + @pytest.mark.parametrize("fmt", ["mp4", "avi"]) + def test_explicit_encoder_format_is_rejected(self, action_video_client, fmt): + """The caller stated two incompatible things; guessing either way is wrong.""" + resp = self._post(action_video_client, format=fmt, extra_params={"action_mode": "policy"}) + assert resp.status_code == 400 + body = resp.json() + _assert_llm_envelope(body, code=400, message_contains="action_mode") + # The message must name the way out, not just the refusal. + assert "safetensors" in body["message"] + + def test_untriggered_request_keeps_encoder_default(self, action_video_client): + """No tensor-only param set -> ordinary video request, unchanged.""" + resp = self._post(action_video_client) + assert resp.status_code == 200 + assert resp.headers["content-type"].startswith("video/") + + +class TestTensorOnlyFormatRule: + """Unit coverage of the resolution rule itself, without a server.""" + + @staticmethod + def _specs(**flags): + from tensorrt_llm._torch.visual_gen.pipeline import ExtraParamSchema + + return { + name: ExtraParamSchema(type="str", default=None, requires_tensor_output=flag) + for name, flag in flags.items() + } + + def test_no_specs_is_a_noop(self): + from tensorrt_llm.serve.openai_video_routes import _resolve_tensor_only_format + + assert _resolve_tensor_only_format("auto", {"action_mode": "policy"}, None) == "auto" + assert _resolve_tensor_only_format("auto", None, self._specs(a=True)) == "auto" + + def test_only_a_declared_param_triggers(self): + from tensorrt_llm.serve.openai_video_routes import _resolve_tensor_only_format + + specs = self._specs(action_mode=True, stg_scale=False) + # a non-declaring param must not force a tensor payload + assert _resolve_tensor_only_format("auto", {"stg_scale": 2.0}, specs) == "auto" + assert ( + _resolve_tensor_only_format("auto", {"action_mode": "policy"}, specs) == "safetensors" + ) + + def test_null_value_does_not_trigger(self): + from tensorrt_llm.serve.openai_video_routes import _resolve_tensor_only_format + + specs = self._specs(action_mode=True) + assert _resolve_tensor_only_format("auto", {"action_mode": None}, specs) == "auto" + + def test_encoder_format_raises_naming_the_parameter(self): + from tensorrt_llm.serve.openai_video_routes import _resolve_tensor_only_format + + specs = self._specs(action_mode=True) + with pytest.raises(ValueError, match="action_mode"): + _resolve_tensor_only_format("mp4", {"action_mode": "policy"}, specs) 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 978c9f8743c4..7dbf95922e8b 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py @@ -1093,6 +1093,45 @@ 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_video_reference_must_be_bytes(self): + """A server-local path must not reach the worker: the ``video`` contract + is encoded bytes, so a string (or anything else) fails preflight.""" + from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import COSMOS3_EXTRA_SPECS + + req = self._make_request(extra_params={"video": "/server/local/path.mp4"}) + with pytest.raises(ValueError, match="expected type 'bytes'"): + from tensorrt_llm.visual_gen.params import validate_visual_gen_params + + validate_visual_gen_params( + 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 73dbb32ce9dd..bfbaafdc4380 100644 --- a/tests/unittest/visual_gen/test_output.py +++ b/tests/unittest/visual_gen/test_output.py @@ -35,6 +35,7 @@ def test_visual_gen_output_is_dataclass(): "image", "video", "audio", + "action", "frame_rate", "audio_sample_rate", "error", @@ -63,6 +64,7 @@ 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.frame_rate is None assert out.audio_sample_rate is None assert out.error is None @@ -326,6 +328,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) # --------------------------------------------------------------------------- @@ -756,13 +770,14 @@ 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_nine_fields(): + """PipelineOutput has the nine expected fields.""" field_names = {f.name for f in fields(PipelineOutput)} assert field_names == { "image", "video", "audio", + "action", "frame_rate", "audio_sample_rate", "pre_denoise",