diff --git a/docs/cookbook/diffusion/Cosmos/Cosmos3.mdx b/docs/cookbook/diffusion/Cosmos/Cosmos3.mdx index 51686bbdd..0dcb84c85 100644 --- a/docs/cookbook/diffusion/Cosmos/Cosmos3.mdx +++ b/docs/cookbook/diffusion/Cosmos/Cosmos3.mdx @@ -331,6 +331,37 @@ print(action["shape"], action["values"]) Use `GET /v1/actions/metadata` to inspect the action modes, default horizon, padded action dimension, and accepted observation modalities. Msgpack requests and the `/v1/actions/realtime` websocket use the same action envelope. +To batch policy observations inside one request, opt in with a bounded batch size: + +```bash Command +sglang serve \ + --model-path nvidia/Cosmos3-Nano-Policy-DROID \ + --num-gpus 1 \ + --batching-max-size 4 +``` + +Send one image per observation as a list or `[B, H, W, C]` uint8 array in `input.input_reference`, and either one prompt per image or one scalar prompt to broadcast across the batch. Batched prompts must currently tokenize to the same length because Cosmos3 GEN cross-attention does not mask padded text K/V. All items in one request share the domain, resolution, action horizon, and denoise settings. The standard action envelope returns one `data[i]` item per input, each with action shape `[H, D]`. For a compact msgpack response containing one `[B, H, D]` array, set `runtime.response_format="raw"` and read the top-level `actions` field. + +For JSON, `input_reference` can be a list of base64 image payloads. For msgpack, it can be a packed uint8 numpy array directly: + +```json JSON +{ + "input": { + "prompt": ["pick up the block", "close the drawer"], + "input_reference": [ + {"b64_json": ""}, + {"b64_json": ""} + ] + }, + "parameters": { + "action_mode": "policy", + "domain_name": "droid_lerobot" + } +} +``` + +The batch size cannot exceed `--batching-max-size`; this keeps one request from bypassing the server's configured memory limit. Batching applies to `action_mode="policy"` only. A request seed controls the random stream for the whole batch, so a batched result is deterministic for that request but is not expected to be bit-exact with separately seeded B=1 requests. + `inverse_dynamics` also uses `/v1/actions/generations`; set `action_mode="inverse_dynamics"` and pass an observation video URL or server-local path as `input.observation.video`. Select the embodiment head with `domain_name` or `domain_id`; set `raw_action_dim` explicitly when it cannot be inferred from the domain name. `forward_dynamics` is intentionally different: it consumes an action array and predicts video, so it remains on `/v1/videos`. Action-producing modes submitted to `/v1/videos` return HTTP 400 with the canonical action endpoint in the error message. diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/action/cosmos3.py b/python/sglang/multimodal_gen/runtime/entrypoints/action/cosmos3.py new file mode 100644 index 000000000..cbc223a3b --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/entrypoints/action/cosmos3.py @@ -0,0 +1,242 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Cosmos3 adapter for the generic action endpoint.""" + +from __future__ import annotations + +import dataclasses +from typing import Any + +import numpy as np +from PIL import Image + +from sglang.multimodal_gen.configs.sample.cosmos3 import Cosmos3SamplingParams +from sglang.multimodal_gen.runtime.server_args import ServerArgs + + +def cosmos3_action_metadata(server_args: ServerArgs) -> dict[str, Any]: + pipeline_config = server_args.pipeline_config + defaults = Cosmos3SamplingParams() + max_batch_size = max(1, int(getattr(server_args, "batching_max_size", 1))) + return { + "object": "action.metadata", + "model": server_args.served_model_name, + "model_path": server_args.model_path, + "policy_family": "cosmos3", + "input": { + "modalities": ["image", "video"], + "supported_resolutions": [ + list(resolution) for resolution in defaults.supported_resolutions + ], + "state_dim": None, + }, + "output": { + "action_type": "continuous", + "action_horizon": 16, + "action_dim": None, + "padded_action_dim": pipeline_config.dit_config.arch_config.action_dim, + "dtype": "float32", + }, + "runtime": { + "parallelism": { + "num_gpus": server_args.num_gpus, + "tp_size": server_args.tp_size, + "sp_degree": server_args.sp_degree, + "ulysses_degree": server_args.ulysses_degree, + "ring_degree": server_args.ring_degree, + } + }, + "defaults": { + "action_mode": "policy", + "action_horizon": 16, + "num_inference_steps": defaults.num_inference_steps, + "height": 480, + "width": 832, + "fps": 5, + }, + "capabilities": { + "action_modes": ["policy", "inverse_dynamics"], + "realtime_websocket": True, + "openpi_websocket": False, + "batch_inputs": max_batch_size > 1, + "max_batch_size": max_batch_size, + "batched_action_modes": ["policy"], + "multiple_candidates": False, + }, + } + + +def _images_from_observation(observation: dict[str, Any]) -> list[Any]: + image = None + for name in ("image", "image_path", "input_reference"): + if name in observation: + image = observation[name] + break + + if image is None: + images = observation.get("images") + if images is None or (isinstance(images, dict) and not images): + return [] + if not isinstance(images, dict) or len(images) != 1: + raise ValueError( + "Cosmos3 action input accepts one image field; use a list or " + "a [B, H, W, C] array in that field for batched observations" + ) + image = next(iter(images.values())) + + if isinstance(image, (list, tuple)): + images = list(image) + elif isinstance(image, np.ndarray) and image.ndim == 4: + images = list(image) + else: + images = [image] + + normalized_images: list[Any] = [] + for item in images: + if not isinstance(item, np.ndarray): + normalized_images.append(item) + continue + if item.dtype != np.uint8: + raise ValueError("Cosmos3 observation image arrays must use uint8 dtype") + if item.ndim not in (2, 3): + raise ValueError( + "Cosmos3 observation image arrays must have shape [H, W] " + f"or [H, W, C], got {tuple(item.shape)}" + ) + normalized_images.append(Image.fromarray(item)) + return normalized_images + + +def _action_prompt(prompt: Any, batch_size: int) -> str | list[str]: + if isinstance(prompt, str): + return prompt if batch_size == 1 else [prompt] * batch_size + if not isinstance(prompt, (list, tuple)) or not prompt: + raise ValueError("Cosmos3 action prompt must be a string or non-empty list") + if not all(isinstance(item, str) for item in prompt): + raise ValueError("Cosmos3 action prompt list must contain only strings") + prompts = list(prompt) + if len(prompts) == 1 and batch_size > 1: + prompts *= batch_size + if len(prompts) != batch_size: + raise ValueError( + "Cosmos3 batched action input requires one prompt per image, got " + f"{len(prompts)} prompt(s) and {batch_size} image(s)" + ) + return prompts[0] if batch_size == 1 else prompts + + +def build_cosmos3_action_sampling_params( + payload: dict[str, Any], + observation: dict[str, Any], + server_args: ServerArgs, + sampling_params_cls: type[Cosmos3SamplingParams], +) -> Cosmos3SamplingParams: + parameters = dict(payload.get("parameters") or {}) + options = {**observation, **parameters} + action_mode = str(options.get("action_mode", "policy")).strip().lower() + if action_mode == "forward_dynamics": + raise ValueError( + "Cosmos3 forward_dynamics produces video; use /v1/videos instead" + ) + if action_mode not in ("policy", "inverse_dynamics"): + raise ValueError( + "Cosmos3 action endpoint supports action_mode='policy' or " + "'inverse_dynamics'" + ) + + action_horizon = options.get("action_horizon") + num_frames = options.get("num_frames") + if action_horizon is None and num_frames is None: + action_horizon = 16 + if action_horizon is not None: + action_horizon = int(action_horizon) + if action_horizon <= 0: + raise ValueError("action_horizon must be a positive integer") + expected_num_frames = action_horizon + 1 + if num_frames is not None and int(num_frames) != expected_num_frames: + raise ValueError( + "Cosmos3 requires num_frames == action_horizon + 1, got " + f"num_frames={num_frames}, action_horizon={action_horizon}" + ) + num_frames = expected_num_frames + else: + num_frames = int(num_frames) + if num_frames <= 1: + raise ValueError("Cosmos3 action num_frames must be greater than 1") + if (num_frames - 1) % 4 != 0: + raise ValueError( + "Cosmos3 action_horizon must be divisible by 4 so num_frames " + "is compatible with the temporal VAE" + ) + + images = _images_from_observation(observation) + video_path = options.get("video_path") or observation.get("video") + if action_mode == "policy" and not images: + raise ValueError("Cosmos3 policy input requires an observation image") + if action_mode == "inverse_dynamics" and video_path is None: + raise ValueError("Cosmos3 inverse_dynamics input requires an observation video") + if images and video_path is not None: + raise ValueError("Cosmos3 action requests accept either an image or a video") + batch_size = len(images) if images else 1 + max_batch_size = max(1, int(getattr(server_args, "batching_max_size", 1))) + if batch_size > max_batch_size: + raise ValueError( + f"Cosmos3 action batch size {batch_size} exceeds " + f"--batching-max-size={max_batch_size}" + ) + image_path = None if not images else images[0] if batch_size == 1 else images + + domain_id = options.get("domain_id") + domain_name = options.get("domain_name") + raw_action_dim = options.get("raw_action_dim") + if domain_id is None and not domain_name: + raise ValueError("Cosmos3 action requests require domain_name or domain_id") + if domain_id is not None and not domain_name and raw_action_dim is None: + raise ValueError("raw_action_dim is required when only domain_id is provided") + + prompt = observation.get("prompt") + if prompt is None: + prompt = observation.get("task", "") + if action_mode != "policy" and not isinstance(prompt, str): + raise ValueError("Cosmos3 inverse_dynamics prompt must be a string") + prompt = _action_prompt(prompt, batch_size) + sampling_kwargs = { + "request_id": payload.get("request_id") or payload.get("id"), + "prompt": prompt, + "image_path": image_path, + "video_path": video_path, + "action_mode": action_mode, + "domain_id": domain_id, + "domain_name": domain_name, + "raw_action_dim": raw_action_dim, + "action_fps": options.get("action_fps"), + "action_view_point": options.get("action_view_point", "ego_view"), + "action_normalization": options.get("action_normalization", "quantile"), + "action_stats_path": server_args.pipeline_config.action_stats_path, + "num_frames": num_frames, + "fps": int(options.get("fps", 5)), + "height": int(options.get("height", 480)), + "width": int(options.get("width", 832)), + "num_inference_steps": int(options.get("num_inference_steps", 35)), + "guidance_scale": float(options.get("guidance_scale", 1.0)), + "seed": int(options.get("seed", 42)), + "flow_shift": options.get("flow_shift"), + "max_sequence_length": options.get("max_sequence_length"), + "condition_frame_indexes": options.get("condition_frame_indexes"), + "condition_video_keep": options.get("condition_video_keep", "first"), + "use_duration_template": False, + "use_system_prompt": False, + "use_guardrails": options.get("use_guardrails"), + "save_output": False, + "return_file_paths_only": False, + "return_frames": False, + } + supported_fields = {field.name for field in dataclasses.fields(sampling_params_cls)} + sampling_params = sampling_params_cls( + **{ + name: value + for name, value in sampling_kwargs.items() + if name in supported_fields and value is not None + } + ) + sampling_params._adjust(server_args) + return sampling_params diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/action/protocol.py b/python/sglang/multimodal_gen/runtime/entrypoints/action/protocol.py index d4f65b2be..da8ae3cdd 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/action/protocol.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/action/protocol.py @@ -17,6 +17,10 @@ from sglang.multimodal_gen.configs.pipeline_configs.cosmos3 import Cosmos3Config from sglang.multimodal_gen.configs.sample.action import ActionSamplingParams from sglang.multimodal_gen.configs.sample.cosmos3 import Cosmos3SamplingParams from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams +from sglang.multimodal_gen.runtime.entrypoints.action.cosmos3 import ( + build_cosmos3_action_sampling_params, + cosmos3_action_metadata, +) from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client from sglang.multimodal_gen.runtime.server_args import ServerArgs @@ -112,7 +116,12 @@ def _normalize_observation(observation: dict[str, Any]) -> dict[str, Any]: } for name in ("image", "image_path", "input_reference"): if name in normalized: - normalized[name] = _normalize_image_value(normalized[name]) + value = normalized[name] + normalized[name] = ( + [_normalize_image_value(item) for item in value] + if isinstance(value, (list, tuple)) + else _normalize_image_value(value) + ) state = normalized.get("state") if isinstance(state, dict): normalized["state"] = _decode_tensor_payload(state) @@ -148,51 +157,7 @@ def images_from_observation( def action_metadata(server_args: ServerArgs) -> dict[str, Any]: pipeline_config = server_args.pipeline_config if isinstance(pipeline_config, Cosmos3Config): - defaults = Cosmos3SamplingParams() - return { - "object": "action.metadata", - "model": server_args.served_model_name, - "model_path": server_args.model_path, - "policy_family": "cosmos3", - "input": { - "modalities": ["image", "video"], - "supported_resolutions": [ - list(resolution) for resolution in defaults.supported_resolutions - ], - "state_dim": None, - }, - "output": { - "action_type": "continuous", - "action_horizon": 16, - "action_dim": None, - "padded_action_dim": pipeline_config.dit_config.arch_config.action_dim, - "dtype": "float32", - }, - "runtime": { - "parallelism": { - "num_gpus": server_args.num_gpus, - "tp_size": server_args.tp_size, - "sp_degree": server_args.sp_degree, - "ulysses_degree": server_args.ulysses_degree, - "ring_degree": server_args.ring_degree, - } - }, - "defaults": { - "action_mode": "policy", - "action_horizon": 16, - "num_inference_steps": defaults.num_inference_steps, - "height": 480, - "width": 832, - "fps": 5, - }, - "capabilities": { - "action_modes": ["policy", "inverse_dynamics"], - "realtime_websocket": True, - "openpi_websocket": False, - "batch_inputs": False, - "multiple_candidates": False, - }, - } + return cosmos3_action_metadata(server_args) policy_family = getattr( pipeline_config, @@ -423,134 +388,6 @@ def _build_action_model_sampling_params( return sp -def _cosmos3_image_from_observation(observation: dict[str, Any]) -> Any: - image = None - for name in ("image", "image_path", "input_reference"): - if name in observation: - image = observation[name] - break - - if image is None: - images = observation.get("images") - if not images: - return None - if not isinstance(images, dict) or len(images) != 1: - raise ValueError( - "Cosmos3 policy input requires exactly one observation image" - ) - image = next(iter(images.values())) - if isinstance(image, np.ndarray): - if image.dtype != np.uint8: - raise ValueError("Cosmos3 observation image arrays must use uint8 dtype") - return Image.fromarray(image) - return image - - -def _build_cosmos3_action_sampling_params( - payload: dict[str, Any], - server_args: ServerArgs, - sampling_params_cls: type[Cosmos3SamplingParams], -) -> Cosmos3SamplingParams: - observation = _action_request_to_observation(payload) - parameters = dict(payload.get("parameters") or {}) - options = {**observation, **parameters} - action_mode = str(options.get("action_mode", "policy")).strip().lower() - if action_mode == "forward_dynamics": - raise ValueError( - "Cosmos3 forward_dynamics produces video; use /v1/videos instead" - ) - if action_mode not in ("policy", "inverse_dynamics"): - raise ValueError( - "Cosmos3 action endpoint supports action_mode='policy' or " - "'inverse_dynamics'" - ) - - action_horizon = options.get("action_horizon") - num_frames = options.get("num_frames") - if action_horizon is None and num_frames is None: - action_horizon = 16 - if action_horizon is not None: - action_horizon = int(action_horizon) - if action_horizon <= 0: - raise ValueError("action_horizon must be a positive integer") - expected_num_frames = action_horizon + 1 - if num_frames is not None and int(num_frames) != expected_num_frames: - raise ValueError( - "Cosmos3 requires num_frames == action_horizon + 1, got " - f"num_frames={num_frames}, action_horizon={action_horizon}" - ) - num_frames = expected_num_frames - else: - num_frames = int(num_frames) - if num_frames <= 1: - raise ValueError("Cosmos3 action num_frames must be greater than 1") - if (num_frames - 1) % 4 != 0: - raise ValueError( - "Cosmos3 action_horizon must be divisible by 4 so num_frames " - "is compatible with the temporal VAE" - ) - - image_path = _cosmos3_image_from_observation(observation) - video_path = options.get("video_path") or observation.get("video") - if action_mode == "policy" and image_path is None: - raise ValueError("Cosmos3 policy input requires an observation image") - if action_mode == "inverse_dynamics" and video_path is None: - raise ValueError("Cosmos3 inverse_dynamics input requires an observation video") - if image_path is not None and video_path is not None: - raise ValueError("Cosmos3 action requests accept either an image or a video") - - domain_id = options.get("domain_id") - domain_name = options.get("domain_name") - raw_action_dim = options.get("raw_action_dim") - if domain_id is None and not domain_name: - raise ValueError("Cosmos3 action requests require domain_name or domain_id") - if domain_id is not None and not domain_name and raw_action_dim is None: - raise ValueError("raw_action_dim is required when only domain_id is provided") - - prompt = observation.get("prompt") or observation.get("task") or "" - sampling_kwargs = { - "request_id": payload.get("request_id") or payload.get("id"), - "prompt": prompt, - "image_path": image_path, - "video_path": video_path, - "action_mode": action_mode, - "domain_id": domain_id, - "domain_name": domain_name, - "raw_action_dim": raw_action_dim, - "action_fps": options.get("action_fps"), - "action_view_point": options.get("action_view_point", "ego_view"), - "action_normalization": options.get("action_normalization", "quantile"), - "action_stats_path": server_args.pipeline_config.action_stats_path, - "num_frames": num_frames, - "fps": int(options.get("fps", 5)), - "height": int(options.get("height", 480)), - "width": int(options.get("width", 832)), - "num_inference_steps": int(options.get("num_inference_steps", 35)), - "guidance_scale": float(options.get("guidance_scale", 1.0)), - "seed": int(options.get("seed", 42)), - "flow_shift": options.get("flow_shift"), - "max_sequence_length": options.get("max_sequence_length"), - "condition_frame_indexes": options.get("condition_frame_indexes"), - "condition_video_keep": options.get("condition_video_keep", "first"), - "use_duration_template": False, - "use_system_prompt": False, - "use_guardrails": options.get("use_guardrails"), - "save_output": False, - "return_file_paths_only": False, - "return_frames": False, - } - supported_fields = _sampling_params_field_names(sampling_params_cls) - sp = sampling_params_cls( - **{ - name: value - for name, value in sampling_kwargs.items() - if name in supported_fields and value is not None - } - ) - sp._adjust(server_args) - return sp - - def build_action_sampling_params( payload: dict[str, Any], server_args: ServerArgs, @@ -563,8 +400,9 @@ def build_action_sampling_params( sampling_params_cls, ) if issubclass(sampling_params_cls, Cosmos3SamplingParams): - return _build_cosmos3_action_sampling_params( + return build_cosmos3_action_sampling_params( payload, + _action_request_to_observation(payload), server_args, sampling_params_cls, ) @@ -594,23 +432,44 @@ def action_generation_response( preserve_numpy: bool = False, ) -> dict[str, Any]: actions = output["actions"] - if isinstance(actions, np.ndarray): - action_shape = list(actions.shape) - action_values = actions if preserve_numpy else actions.tolist() - else: - horizon = len(actions) if isinstance(actions, list) else 0 - action_dim = len(actions[0]) if horizon and isinstance(actions[0], list) else 0 - action_shape = [horizon, action_dim] - action_values = actions - action = { - "type": "continuous", - "dtype": "float32", - "shape": action_shape, - "values": action_values, - } - for name in ("action_mode", "domain_id", "raw_action_dim"): - if output.get(name) is not None: - action[name] = output[name] + action_array = np.asarray(actions) + if any(size == 0 for size in action_array.shape): + raise ValueError( + "action output dimensions must be non-zero, got " + f"{tuple(action_array.shape)}" + ) + if action_array.ndim == 2: + action_array = action_array[None] + elif action_array.ndim != 3: + raise ValueError( + "action output must have shape [H, D] or [B, H, D], got " + f"{tuple(action_array.shape)}" + ) + + data = [] + for input_index, action_values in enumerate(action_array): + action_shape = list(action_values.shape) + if not preserve_numpy: + action_values = action_values.tolist() + action = { + "type": "continuous", + "dtype": "float32", + "shape": action_shape, + "values": action_values, + } + for name in ("action_mode", "domain_id", "raw_action_dim"): + if output.get(name) is not None: + action[name] = output[name] + data.append( + { + "index": input_index, + "input_index": input_index, + "candidate_index": 0, + "action": action, + } + ) + + action_shape = data[0]["action"]["shape"] pipeline_config = server_args.pipeline_config if isinstance(pipeline_config, Cosmos3Config): @@ -623,15 +482,9 @@ def action_generation_response( "object": "action.generation", "created": int(time.time()), "model": server_args.served_model_name, - "data": [ - { - "index": 0, - "input_index": 0, - "candidate_index": 0, - "action": action, - } - ], + "data": data, "usage": { + "batch_size": len(data), "action_horizon": action_shape[0] if action_shape else 0, "action_dim": action_shape[1] if len(action_shape) > 1 else 0, "denoise_steps": output.get("parameters", {}).get( diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py index e964b91a1..97a7d3588 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py @@ -899,8 +899,17 @@ def prepare_request( if diffusers_kwargs and "max_sequence_length" in diffusers_kwargs: req.max_sequence_length = diffusers_kwargs["max_sequence_length"] - if not isinstance(req.prompt, str): - raise TypeError(f"`prompt` must be a string, but got {type(req.prompt)}") + action_prompt = ( + req.data_type == DataType.ACTION + and isinstance(req.prompt, list) + and bool(req.prompt) + and all(isinstance(item, str) for item in req.prompt) + ) + if not isinstance(req.prompt, str) and not action_prompt: + raise TypeError( + "`prompt` must be a string, or a non-empty list of strings for " + f"batched action requests, but got {type(req.prompt)}" + ) req_width = getattr(req, "width", None) req_height = getattr(req, "height", None) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py index 80316b67c..ee496fe86 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py @@ -41,6 +41,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.base import ( from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3_action import ( ACTION_MODE_FORWARD_DYNAMICS, ACTION_MODE_INVERSE_DYNAMICS, + ACTION_MODE_POLICY, ACTION_MODES, EMBODIMENT_TO_DOMAIN_ID, build_action_prompt, @@ -141,7 +142,9 @@ def _pil_to_normalized_tensor(image: PIL.Image.Image) -> torch.Tensor: class Cosmos3ImagePreprocessStage(PipelineStage): """Load, aspect-resize, and center-crop the conditioning input. - For I2V: writes ``[1, 3, H, W]`` to ``batch.preprocessed_image``. + For I2V: writes ``[1, 3, H, W]`` to ``batch.preprocessed_image``. Batched + policy requests write ``[B, 3, H, W]``; regular visual generation remains + single-image conditioned. For V2V: writes ``[1, 3, T_in, H, W]`` to ``batch.preprocessed_video``. No-op for T2V / T2I. """ @@ -153,9 +156,14 @@ class Cosmos3ImagePreprocessStage(PipelineStage): def forward(self, batch: Req, server_args: ServerArgs) -> Req: image_path = batch.image_path - if isinstance(image_path, list): - image_path = image_path[0] if image_path else None video_path = batch.video_path + is_action_policy = ( + batch.data_type == DataType.ACTION + and getattr(batch.sampling_params, "action_mode", None) + == ACTION_MODE_POLICY + ) + if isinstance(image_path, list) and not is_action_policy: + image_path = image_path[0] if image_path else None if isinstance(video_path, list): video_path = video_path[0] if video_path else None @@ -168,10 +176,23 @@ class Cosmos3ImagePreprocessStage(PipelineStage): target_h, target_w = batch.height, batch.width if image_path is not None: - image = load_image(image_path) - image = _resize_crop_pil(image, target_w, target_h) - batch.preprocessed_image = _pil_to_normalized_tensor(image).unsqueeze(0) - self.log_info(f"Preprocessed conditioning image to {target_w}x{target_h}") + image_sources = ( + list(image_path) + if isinstance(image_path, (list, tuple)) + else [image_path] + ) + if not image_sources: + raise ValueError("Cosmos3 I2V image list is empty") + tensors: list[torch.Tensor] = [] + for src in image_sources: + image = load_image(src) + image = _resize_crop_pil(image, target_w, target_h) + tensors.append(_pil_to_normalized_tensor(image)) + batch.preprocessed_image = torch.stack(tensors, dim=0).contiguous() + self.log_info( + f"Preprocessed {len(tensors)} conditioning image(s) to " + f"{target_w}x{target_h}" + ) return batch if isinstance(video_path, str) and video_path: @@ -265,7 +286,7 @@ class Cosmos3TokenizationStage(PipelineStage): def _tokenize_prompt( self, - text: str, + text: str | list[str], max_sequence_length: int, device: torch.device, use_system_prompt: bool = False, @@ -273,58 +294,72 @@ class Cosmos3TokenizationStage(PipelineStage): ) -> tuple[torch.Tensor, torch.Tensor, int]: """Tokenize a prompt using Qwen2 chat template. - Returns (input_ids, attention_mask, seq_len) as [1, S] tensors. + Returns (input_ids, attention_mask, seq_len) as [B, S] tensors. """ - conversations = [] - if use_system_prompt: - conversations.append( - { - "role": "system", - "content": system_prompt or COSMOS3_VIDEO_SYSTEM_PROMPT, - } - ) - conversations.append({"role": "user", "content": text}) - - result = self.tokenizer.apply_chat_template( - conversations, - tokenize=True, - add_generation_prompt=True, - ) - - # Handle different return types from apply_chat_template - # Fast tokenizer returns BatchEncoding, slow tokenizer returns list[int] - if hasattr(result, "input_ids"): - # BatchEncoding from fast tokenizer - token_ids = list(result.input_ids) - elif isinstance(result, list): - # Already a list from slow tokenizer - token_ids = list(result) - else: - raise TypeError( - f"Unexpected return type from apply_chat_template: {type(result)}" - ) - - # Reserve room for the two special tokens (EOS + vision_start) so the - # final length cannot exceed ``max_sequence_length``. - token_ids = token_ids[: max_sequence_length - 2] - - # Add EOS and vision_start tokens - token_ids.append(self.tokenizer.eos_token_id) - vision_start_id = self.tokenizer.convert_tokens_to_ids("<|vision_start|>") - if vision_start_id is not None: - token_ids.append(vision_start_id) - - seq_len = len(token_ids) - - # Pad to max_sequence_length - pad_len = max_sequence_length - seq_len - attention_mask = [1] * seq_len + [0] * pad_len + texts = text if isinstance(text, (list, tuple)) else [text] + if not texts: + raise ValueError("Cosmos3 prompt batch must not be empty") + input_id_lists: list[list[int]] = [] + attention_mask_lists: list[list[int]] = [] + seq_lens: list[int] = [] pad_token_id = self.tokenizer.pad_token_id or 0 - token_ids = token_ids + [pad_token_id] * pad_len + vision_start_id = self.tokenizer.convert_tokens_to_ids("<|vision_start|>") + for text_item in texts: + conversations = [] + if use_system_prompt: + conversations.append( + { + "role": "system", + "content": system_prompt or COSMOS3_VIDEO_SYSTEM_PROMPT, + } + ) + conversations.append({"role": "user", "content": text_item}) - input_ids = torch.tensor([token_ids], dtype=torch.long, device=device) - attention_mask = torch.tensor([attention_mask], dtype=torch.long, device=device) - return input_ids, attention_mask, seq_len + result = self.tokenizer.apply_chat_template( + conversations, + tokenize=True, + add_generation_prompt=True, + ) + # Handle different return types from apply_chat_template + # Fast tokenizer returns BatchEncoding, slow tokenizer returns list[int] + if hasattr(result, "input_ids"): + # BatchEncoding from fast tokenizer + token_ids = list(result.input_ids) + elif isinstance(result, list): + # Already a list from slow tokenizer + token_ids = list(result) + else: + raise TypeError( + f"Unexpected return type from apply_chat_template: {type(result)}" + ) + + # Reserve room for the two special tokens (EOS + vision_start) so the + # final length cannot exceed ``max_sequence_length``. + token_ids = token_ids[: max_sequence_length - 2] + # Add EOS and vision_start tokens + token_ids.append(self.tokenizer.eos_token_id) + if vision_start_id is not None: + token_ids.append(vision_start_id) + + seq_len = len(token_ids) + pad_len = max_sequence_length - seq_len + attention_mask = [1] * seq_len + [0] * pad_len + token_ids = token_ids + [pad_token_id] * pad_len + input_id_lists.append(token_ids) + attention_mask_lists.append(attention_mask) + seq_lens.append(seq_len) + + if len(set(seq_lens)) != 1: + raise ValueError( + "Cosmos3 batched prompts must tokenize to the same length because " + "GEN cross-attention does not mask padded text K/V; split prompts " + f"into equal-length batches instead (lengths={seq_lens})" + ) + input_ids = torch.tensor(input_id_lists, dtype=torch.long, device=device) + attention_mask = torch.tensor( + attention_mask_lists, dtype=torch.long, device=device + ) + return input_ids, attention_mask, seq_lens[0] def forward(self, batch: Req, server_args: ServerArgs) -> Req: """Tokenize prompt and negative prompt.""" @@ -383,6 +418,9 @@ class Cosmos3TokenizationStage(PipelineStage): self.log_info(f"Prompt with duration: '{prompt}'") # Tokenize prompts + if isinstance(prompt, list) and not isinstance(negative_prompt, list): + negative_prompt = [negative_prompt] * len(prompt) + cond_ids, cond_mask, cond_seq_len = self._tokenize_prompt( prompt, max_sequence_length, device, use_system_prompt, system_prompt ) @@ -473,8 +511,13 @@ class Cosmos3LatentPreparationStage(PipelineStage): height_latent = batch.height // vae_scale_factor_spatial width_latent = batch.width // vae_scale_factor_spatial + if batch.preprocessed_image is not None: + batch_dim = int(batch.preprocessed_image.shape[0]) + else: + batch_dim = 1 + shape = ( - 1, + batch_dim, num_channels_latents, num_latent_frames, height_latent, @@ -523,7 +566,7 @@ class Cosmos3LatentPreparationStage(PipelineStage): condition_latents = torch.zeros_like(noise) condition_mask = torch.zeros( - 1, 1, num_latent_frames, 1, 1, device=device, dtype=dtype + batch_dim, 1, num_latent_frames, 1, 1, device=device, dtype=dtype ) for idx in cond_indexes: src = min(idx, cond_latent.shape[2] - 1) @@ -637,6 +680,11 @@ class Cosmos3LatentPreparationStage(PipelineStage): action_offset = 1 if action_chunk_size == num_frames - 1 else 0 domain_id = self._resolve_domain_id(batch) + batch_dim = ( + int(batch.raw_latent_shape[0]) + if getattr(batch, "raw_latent_shape", None) + else 1 + ) raw_action_dim = getattr(sp, "raw_action_dim", None) if raw_action_dim is None: embodiment = getattr(sp, "domain_name", None) @@ -678,7 +726,7 @@ class Cosmos3LatentPreparationStage(PipelineStage): if raw_action_dim is None: raise ValueError(f"action_mode={mode!r} requires --raw-action-dim.") clean_action = torch.zeros( - 1, action_chunk_size, action_dim, device=device, dtype=dtype + batch_dim, action_chunk_size, action_dim, device=device, dtype=dtype ) raw_action_dim = int(raw_action_dim) @@ -690,13 +738,13 @@ class Cosmos3LatentPreparationStage(PipelineStage): # condition_mask marks clean (given) action tokens. forward_dynamics # conditions on the whole action sequence; the others denoise it fully. condition_mask = torch.zeros( - 1, action_chunk_size, 1, device=device, dtype=dtype + batch_dim, action_chunk_size, 1, device=device, dtype=dtype ) if mode == ACTION_MODE_FORWARD_DYNAMICS: condition_mask[:] = 1.0 noise = torch.randn( - 1, + batch_dim, action_chunk_size, action_dim, generator=generator, @@ -709,7 +757,7 @@ class Cosmos3LatentPreparationStage(PipelineStage): batch.action_latents = action_latents batch.extra["action_domain_ids"] = torch.tensor( - [domain_id], dtype=torch.long, device=device + [domain_id] * batch_dim, dtype=torch.long, device=device ) batch.extra["action_velocity_mask"] = 1.0 - condition_mask batch.extra["action_condition_latents"] = clean_action @@ -1092,7 +1140,8 @@ class Cosmos3DenoisingStage(PipelineStage, RolloutDenoisingMixin): ) for i, t in progress_bar: - timestep = t.unsqueeze(0) if t.dim() == 0 else t + batch_dim = batch.latents.shape[0] if batch.latents is not None else 1 + timestep = t.unsqueeze(0).expand(batch_dim) if t.dim() == 0 else t # Outside the CFG window the effective scale collapses to 1.0, # which reduces CFG to the cond branch (cfg-parallel safe). effective_scale = ( @@ -1399,7 +1448,7 @@ class Cosmos3DenoisingStage(PipelineStage, RolloutDenoisingMixin): latents_batched = torch.cat([latents, latents], dim=0) text_ids_batched = torch.cat([uncond_text_ids, cond_text_ids], dim=0) text_mask_batched = torch.cat([uncond_text_mask, cond_text_mask], dim=0) - timestep_batched = timestep.expand(2) + timestep_batched = torch.cat([timestep, timestep], dim=0) mask_batched = ( torch.cat([noisy_frame_mask, noisy_frame_mask], dim=0) if noisy_frame_mask is not None @@ -1645,9 +1694,12 @@ class Cosmos3DecodingStage(PipelineStage): if batch.data_type == DataType.ACTION: if action_pred is None: raise RuntimeError("Cosmos3 action request produced no action tensor") + payload_actions = ( + action_pred[0] if action_pred.shape[0] == 1 else action_pred + ) payload = { "request_id": batch.request_id, - "actions": action_pred[0].numpy(), + "actions": payload_actions.numpy(), "action_mode": action_metadata["action_mode"], "domain_id": action_metadata["action_domain_id"], "raw_action_dim": action_metadata["action_raw_action_dim"], diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3_action.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3_action.py index a573673eb..e71440150 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3_action.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3_action.py @@ -123,31 +123,42 @@ def canonical_aspect_ratio(width: int, height: int) -> str: def build_action_prompt( - description: str, + description: str | list[str], view_point: str, num_frames: int, fps: float, height: int, width: int, -) -> str: +) -> str | list[str]: """Render the structured JSON action caption the action checkpoints expect.""" duration_seconds = num_frames / fps minutes, secs = divmod(round(duration_seconds), 60) - if description and description[-1] not in ".!?": - description = description + "." - prompt = { - "cinematography": { - "framing": VIEWPOINT_TEMPLATES.get( - view_point, VIEWPOINT_TEMPLATES["ego_view"] - ) - }, - "actions": [{"time": f"0:00-{minutes}:{secs:02d}", "description": description}], - "duration": f"{int(duration_seconds)}s", - "fps": float(fps), - "resolution": {"H": int(height), "W": int(width)}, - "aspect_ratio": canonical_aspect_ratio(int(width), int(height)), - } - return json.dumps(prompt) + if isinstance(description, (list, tuple)): + descriptions = [str(d) for d in description] + else: + descriptions = [description] + + prompts = [] + for desc in descriptions: + if desc and desc[-1] not in ".!?": + desc = desc + "." + prompt = { + "cinematography": { + "framing": VIEWPOINT_TEMPLATES.get( + view_point, VIEWPOINT_TEMPLATES["ego_view"] + ) + }, + "actions": [{"time": f"0:00-{minutes}:{secs:02d}", "description": desc}], + "duration": f"{int(duration_seconds)}s", + "fps": float(fps), + "resolution": {"H": int(height), "W": int(width)}, + "aspect_ratio": canonical_aspect_ratio(int(width), int(height)), + } + prompts.append(json.dumps(prompt)) + + if isinstance(description, (list, tuple)): + return prompts + return prompts[0] def load_action_stats( diff --git a/python/sglang/multimodal_gen/test/unit/test_cosmos3.py b/python/sglang/multimodal_gen/test/unit/test_cosmos3.py index f303a3e53..61f703beb 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cosmos3.py +++ b/python/sglang/multimodal_gen/test/unit/test_cosmos3.py @@ -8,6 +8,7 @@ import unittest from unittest import mock import torch +from PIL import Image from sglang.multimodal_gen.configs.models.dits.cosmos3video import ( _build_cosmos3_param_names_mapping, @@ -51,6 +52,7 @@ from sglang.multimodal_gen.runtime.models.dits.cosmos3video import ( ) from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3 import ( Cosmos3DecodingStage, + Cosmos3DenoisingStage, Cosmos3ImagePreprocessStage, Cosmos3LatentPreparationStage, Cosmos3TimestepPreparationStage, @@ -70,7 +72,7 @@ def _apply(mapping_fn, key): return mapping_fn(key) -def _cosmos3_server_args(config=None): +def _cosmos3_server_args(config=None, batching_max_size=1): return types.SimpleNamespace( model_id=None, model_path="nvidia/Cosmos3-Nano", @@ -84,6 +86,7 @@ def _cosmos3_server_args(config=None): sp_degree=1, ulysses_degree=1, ring_degree=1, + batching_max_size=batching_max_size, pipeline_config=config or Cosmos3Config(), ) @@ -497,6 +500,31 @@ class TestCosmos3SamplingParamsDataType(unittest.TestCase): class TestCosmos3ActionEndpoint(unittest.TestCase): + @staticmethod + def _policy_payload( + batch_size=2, + prompt="pick up the block", + *, + tensor_payload=False, + ): + images = torch.zeros(batch_size, 8, 8, 3, dtype=torch.uint8).numpy() + input_reference = ( + { + "dtype": "uint8", + "shape": list(images.shape), + "values": images.tolist(), + } + if tensor_payload + else images + ) + return { + "input": {"prompt": prompt, "input_reference": input_reference}, + "parameters": { + "action_mode": "policy", + "domain_name": "droid_lerobot", + }, + } + def test_policy_request_builds_action_sampling_params(self): image = torch.zeros(8, 8, 3, dtype=torch.uint8).numpy() payload = { @@ -535,6 +563,49 @@ class TestCosmos3ActionEndpoint(unittest.TestCase): self.assertEqual(params.seed, 7) self.assertEqual(params.image_path.size, (8, 8)) + def test_batched_policy_request_preserves_input_pairing(self): + payload = self._policy_payload( + prompt=["pick up the block", "close the drawer"], + tensor_payload=True, + ) + + params = build_action_sampling_params( + payload, _cosmos3_server_args(batching_max_size=2) + ) + + self.assertEqual(params.prompt, ["pick up the block", "close the drawer"]) + self.assertEqual(len(params.image_path), 2) + self.assertTrue( + all(isinstance(image, Image.Image) for image in params.image_path) + ) + + def test_batched_policy_request_broadcasts_scalar_prompt(self): + params = build_action_sampling_params( + self._policy_payload(), _cosmos3_server_args(batching_max_size=2) + ) + + self.assertEqual(params.prompt, ["pick up the block"] * 2) + + def test_batched_policy_request_rejects_cardinality_mismatch(self): + with self.assertRaisesRegex(ValueError, "one prompt per image"): + build_action_sampling_params( + self._policy_payload(prompt=["one", "two", "three"]), + _cosmos3_server_args(batching_max_size=3), + ) + + def test_batched_policy_request_honors_server_batch_limit(self): + with self.assertRaisesRegex(ValueError, "--batching-max-size=1"): + build_action_sampling_params(self._policy_payload(), _cosmos3_server_args()) + + def test_single_item_image_batch_keeps_single_input_contract(self): + params = build_action_sampling_params( + self._policy_payload(batch_size=1, prompt="pick"), + _cosmos3_server_args(), + ) + + self.assertEqual(params.prompt, "pick") + self.assertIsInstance(params.image_path, Image.Image) + def test_inverse_dynamics_maps_video_input(self): payload = { "input": { @@ -554,6 +625,21 @@ class TestCosmos3ActionEndpoint(unittest.TestCase): self.assertEqual(params.video_path, "observation.mp4") self.assertEqual(params.num_frames, 61) + def test_inverse_dynamics_rejects_prompt_batch(self): + payload = { + "input": { + "prompt": ["first", "second"], + "video": "observation.mp4", + }, + "parameters": { + "action_mode": "inverse_dynamics", + "domain_name": "droid_lerobot", + }, + } + + with self.assertRaisesRegex(ValueError, "prompt must be a string"): + build_action_sampling_params(payload, _cosmos3_server_args()) + def test_forward_dynamics_is_rejected_by_action_endpoint(self): payload = { "input": { @@ -567,7 +653,7 @@ class TestCosmos3ActionEndpoint(unittest.TestCase): build_action_sampling_params(payload, _cosmos3_server_args()) def test_metadata_describes_cosmos_action_contract(self): - metadata = action_metadata(_cosmos3_server_args()) + metadata = action_metadata(_cosmos3_server_args(batching_max_size=4)) self.assertEqual(metadata["model"], "cosmos3-production") self.assertEqual(metadata["policy_family"], "cosmos3") @@ -575,6 +661,15 @@ class TestCosmos3ActionEndpoint(unittest.TestCase): self.assertEqual(metadata["output"]["action_horizon"], 16) self.assertEqual(metadata["output"]["padded_action_dim"], 64) self.assertFalse(metadata["capabilities"]["openpi_websocket"]) + self.assertTrue(metadata["capabilities"]["batch_inputs"]) + self.assertEqual(metadata["capabilities"]["max_batch_size"], 4) + self.assertEqual(metadata["capabilities"]["batched_action_modes"], ["policy"]) + + def test_metadata_keeps_batching_opt_in(self): + metadata = action_metadata(_cosmos3_server_args()) + + self.assertFalse(metadata["capabilities"]["batch_inputs"]) + self.assertEqual(metadata["capabilities"]["max_batch_size"], 1) def test_action_response_includes_cosmos_metadata(self): output = { @@ -595,6 +690,33 @@ class TestCosmos3ActionEndpoint(unittest.TestCase): self.assertEqual(action["raw_action_dim"], 10) self.assertEqual(response["usage"]["denoise_steps"], 30) + def test_batched_action_response_emits_one_data_item_per_input(self): + output = { + "request_id": "cosmos-action-batch", + "actions": torch.arange(24, dtype=torch.float32).reshape(2, 4, 3).numpy(), + "action_mode": "policy", + "domain_id": 8, + "raw_action_dim": 3, + } + + response = action_generation_response(output, _cosmos3_server_args()) + + self.assertEqual(len(response["data"]), 2) + self.assertEqual(response["data"][0]["action"]["shape"], [4, 3]) + self.assertEqual(response["data"][1]["input_index"], 1) + self.assertEqual(response["usage"]["batch_size"], 2) + self.assertEqual(response["usage"]["action_horizon"], 4) + self.assertEqual(response["usage"]["action_dim"], 3) + + def test_action_response_rejects_empty_batch(self): + output = { + "request_id": "cosmos-action-empty", + "actions": torch.empty(0, 4, 3).numpy(), + } + + with self.assertRaisesRegex(ValueError, "dimensions must be non-zero"): + action_generation_response(output, _cosmos3_server_args()) + def test_action_decode_skips_vae(self): class FailIfDecoded: def decode(self, _latents): @@ -629,6 +751,35 @@ class TestCosmos3ActionEndpoint(unittest.TestCase): self.assertEqual(output.output[0]["domain_id"], 8) self.assertEqual(output.action_pred.shape, (1, 4, 3)) + def test_batched_action_decode_keeps_batch_dimension(self): + stage = Cosmos3DecodingStage.__new__(Cosmos3DecodingStage) + stage.vae = mock.Mock() + stage.sound_tokenizer = None + stage._guardrails = False + stage.log_info = lambda *_args, **_kwargs: None + batch = types.SimpleNamespace( + data_type=DataType.ACTION, + action_latents=torch.arange(48, dtype=torch.float32).reshape(2, 4, 6), + extra={ + "raw_action_dim": 3, + "action_domain_ids": torch.tensor([8, 8]), + }, + sampling_params=Cosmos3SamplingParams( + prompt=["one", "two"], + action_mode="policy", + domain_name="droid_lerobot", + ), + request_id="cosmos-action-batch", + num_inference_steps=30, + num_frames=5, + metrics=None, + ) + + output = stage.forward(batch, types.SimpleNamespace(vae_cpu_offload=False)) + + self.assertEqual(output.output[0]["actions"].shape, (2, 4, 3)) + self.assertEqual(output.action_pred.shape, (2, 4, 3)) + class TestCosmos3ModelResolution(unittest.TestCase): """Verify Cosmos3 checkpoints resolve to the native SGLang pipeline.""" @@ -972,10 +1123,13 @@ class TestCosmos3ActionLatentPrep(unittest.TestCase): cls.device = torch.device("cpu") cls.dtype = torch.float32 - def _run(self, num_frames=17, **sp_kwargs): + def _run(self, num_frames=17, batch_size=1, **sp_kwargs): sp = Cosmos3SamplingParams(prompt="t", num_frames=num_frames, **sp_kwargs) batch = types.SimpleNamespace( - sampling_params=sp, num_frames=num_frames, extra={} + sampling_params=sp, + num_frames=num_frames, + raw_latent_shape=(batch_size, 48, 1, 1, 1), + extra={}, ) gen = torch.Generator(device=self.device).manual_seed(0) self.stage._prepare_action_latents(batch, gen, self.device, self.dtype) @@ -1010,6 +1164,17 @@ class TestCosmos3ActionLatentPrep(unittest.TestCase): # padding dims beyond raw_action_dim start at zero. self.assertTrue(torch.all(batch.action_latents[:, :, 10:] == 0)) + def test_batched_policy_prepares_one_action_stream_per_observation(self): + batch = self._run( + batch_size=3, + action_mode="policy", + domain_name="droid_lerobot", + ) + + self.assertEqual(tuple(batch.action_latents.shape), (3, 16, 64)) + self.assertEqual(tuple(batch.extra["action_domain_ids"].shape), (3,)) + self.assertEqual(tuple(batch.extra["action_velocity_mask"].shape), (3, 16, 1)) + def test_inverse_dynamics_denoises_from_noise(self): batch = self._run( num_frames=61, @@ -1041,6 +1206,113 @@ class TestCosmos3ActionLatentPrep(unittest.TestCase): self._run(action_mode="teleport", domain_id=0) +class TestCosmos3BatchedActionStages(unittest.TestCase): + class _Tokenizer: + pad_token_id = 0 + eos_token_id = 1 + + @staticmethod + def convert_tokens_to_ids(_token): + return 2 + + @staticmethod + def apply_chat_template(conversations, **_kwargs): + text = conversations[-1]["content"] + return list(range(3, 3 + len(text.split()))) + + @staticmethod + def _preprocess_images(data_type, action_mode=None): + stage = Cosmos3ImagePreprocessStage.__new__(Cosmos3ImagePreprocessStage) + stage.log_info = lambda *_args, **_kwargs: None + seen = [] + batch = types.SimpleNamespace( + image_path=["first.png", "second.png"], + video_path=None, + height=16, + width=16, + data_type=data_type, + sampling_params=types.SimpleNamespace(action_mode=action_mode), + preprocessed_image=None, + ) + + def fake_load_image(path): + seen.append(path) + return Image.new("RGB", (16, 16)) + + with mock.patch( + "sglang.multimodal_gen.runtime.pipelines_core.stages." + "model_specific_stages.cosmos3.load_image", + side_effect=fake_load_image, + ): + stage.forward(batch, types.SimpleNamespace()) + return seen, batch.preprocessed_image + + def test_tokenization_rejects_unequal_prompt_lengths(self): + stage = Cosmos3TokenizationStage.__new__(Cosmos3TokenizationStage) + stage.tokenizer = self._Tokenizer() + + with self.assertRaisesRegex(ValueError, "same length"): + stage._tokenize_prompt( + ["pick block", "close the top drawer"], + max_sequence_length=16, + device=torch.device("cpu"), + ) + + def test_tokenization_preserves_equal_length_prompt_batch(self): + stage = Cosmos3TokenizationStage.__new__(Cosmos3TokenizationStage) + stage.tokenizer = self._Tokenizer() + + input_ids, attention_mask, seq_len = stage._tokenize_prompt( + ["pick block", "push cube"], + max_sequence_length=16, + device=torch.device("cpu"), + ) + + self.assertEqual(tuple(input_ids.shape), (2, 16)) + self.assertEqual(tuple(attention_mask.shape), (2, 16)) + self.assertEqual(seq_len, 4) + + def test_cfg_duplicates_batched_timesteps(self): + stage = Cosmos3DenoisingStage.__new__(Cosmos3DenoisingStage) + captured = {} + + def fake_run_transformer(**kwargs): + captured.update(kwargs) + return kwargs["latents"] + + stage._run_transformer = fake_run_transformer + latents = torch.zeros(3, 1, 1, 1, 1) + text_ids = torch.ones(3, 4, dtype=torch.long) + text_mask = torch.ones_like(text_ids) + + output = stage._predict_noise_cfg_batched( + latents=latents, + timestep=torch.ones(3), + cond_text_ids=text_ids, + cond_text_mask=text_mask, + uncond_text_ids=text_ids, + uncond_text_mask=text_mask, + video_shape=(1, 1, 1), + fps=20.0, + guidance_scale=2.0, + ) + + self.assertEqual(tuple(captured["timestep"].shape), (6,)) + self.assertEqual(tuple(output.shape), tuple(latents.shape)) + + def test_visual_i2v_keeps_single_conditioning_image(self): + seen, image = self._preprocess_images(DataType.VIDEO) + + self.assertEqual(seen, ["first.png"]) + self.assertEqual(tuple(image.shape), (1, 3, 16, 16)) + + def test_action_preprocess_stacks_all_images(self): + seen, image = self._preprocess_images(DataType.ACTION, "policy") + + self.assertEqual(seen, ["first.png", "second.png"]) + self.assertEqual(tuple(image.shape), (2, 3, 16, 16)) + + class TestCosmos3ModalitySamplingParams(unittest.TestCase): """Sound / V2V / action sampling-param fields and defaults."""