From 955569a2dc993174a27c6e8bfea70e2e212afad1 Mon Sep 17 00:00:00 2001 From: Mick Date: Mon, 10 Aug 2026 18:16:20 +0800 Subject: [PATCH] [diffusion] feat: expose cosmos3 policies through the Action API (#34243) --- docs/cookbook/diffusion/Cosmos/Cosmos3.mdx | 55 ++-- .../benchmarks/bench_pi05_openpi.py | 10 +- .../configs/pipeline_configs/base.py | 10 + .../configs/pipeline_configs/cosmos3.py | 6 + .../configs/pipeline_configs/pi05.py | 3 + .../multimodal_gen/configs/sample/__init__.py | 4 +- .../configs/sample/{vla.py => action.py} | 12 +- .../multimodal_gen/configs/sample/cosmos3.py | 31 +++ .../multimodal_gen/configs/sample/pi05.py | 4 +- python/sglang/multimodal_gen/registry.py | 1 + .../entrypoints/{vla => action}/__init__.py | 0 .../entrypoints/{vla => action}/api.py | 4 +- .../entrypoints/{vla => action}/openpi.py | 2 +- .../entrypoints/{vla => action}/protocol.py | 252 ++++++++++++++++-- .../entrypoints/{vla => action}/ws_utils.py | 2 +- .../runtime/entrypoints/http_server.py | 9 +- .../runtime/entrypoints/openai/video_api.py | 12 +- .../multimodal_gen/runtime/launch_server.py | 4 +- .../stages/model_specific_stages/cosmos3.py | 86 ++++-- .../multimodal_gen/test/unit/test_cosmos3.py | 185 +++++++++++++ .../test/unit/test_pi05_action_api.py | 8 +- 21 files changed, 608 insertions(+), 92 deletions(-) rename python/sglang/multimodal_gen/configs/sample/{vla.py => action.py} (96%) rename python/sglang/multimodal_gen/runtime/entrypoints/{vla => action}/__init__.py (100%) rename python/sglang/multimodal_gen/runtime/entrypoints/{vla => action}/api.py (95%) rename python/sglang/multimodal_gen/runtime/entrypoints/{vla => action}/openpi.py (90%) rename python/sglang/multimodal_gen/runtime/entrypoints/{vla => action}/protocol.py (60%) rename python/sglang/multimodal_gen/runtime/entrypoints/{vla => action}/ws_utils.py (96%) diff --git a/docs/cookbook/diffusion/Cosmos/Cosmos3.mdx b/docs/cookbook/diffusion/Cosmos/Cosmos3.mdx index d375b8f25..34ec9e720 100644 --- a/docs/cookbook/diffusion/Cosmos/Cosmos3.mdx +++ b/docs/cookbook/diffusion/Cosmos/Cosmos3.mdx @@ -207,27 +207,48 @@ sglang serve \ --num-gpus 1 ``` -The following request predicts a 16-step action chunk from one observation. The chunk length is `num_frames - 1`, and the completed job's `action` field contains the tensor data, shape, mode, and active action dimension. +`policy` and `inverse_dynamics` return actions, so their canonical API is the synchronous `/v1/actions/generations` endpoint. The following request predicts a 16-step action chunk from one observation image. `action_horizon=16` maps to the model's `num_frames=17` convention. -```bash Command -job_id=$(curl -sS -X POST http://127.0.0.1:30010/v1/videos \ - --form-string "prompt=Put the pot to the left of the purple item." \ - --form "input_reference=@observation.png;type=image/png" \ - --form-string "size=832x480" \ - --form-string "num_frames=17" \ - --form-string "fps=5" \ - --form-string "num_inference_steps=30" \ - --form-string "guidance_scale=1.0" \ - --form-string "action_mode=policy" \ - --form-string "domain_name=droid_lerobot" \ - | python -c 'import json, sys; print(json.load(sys.stdin)["id"])') +```python Python +import base64 +from pathlib import Path -# After the job reaches "completed": -curl -sS "http://127.0.0.1:30010/v1/videos/${job_id}" \ - | python -c 'import json, sys; print(json.dumps(json.load(sys.stdin)["action"], indent=2))' +import requests + +image_b64 = base64.b64encode(Path("observation.png").read_bytes()).decode() +response = requests.post( + "http://127.0.0.1:30010/v1/actions/generations", + json={ + "input": { + "task": "Put the pot to the left of the purple item.", + "observation": { + "image": {"b64_json": image_b64}, + }, + }, + "parameters": { + "action_mode": "policy", + "action_horizon": 16, + "domain_name": "droid_lerobot", + "height": 480, + "width": 832, + "fps": 5, + "num_inference_steps": 30, + "guidance_scale": 1.0, + "seed": 42, + }, + }, + timeout=300, +) +response.raise_for_status() +action = response.json()["data"][0]["action"] +print(action["shape"], action["values"]) ``` -The other action modes are `forward_dynamics` (condition on an observation and an `action` JSON array to generate video) and `inverse_dynamics` (condition on a full video to predict action). Select the embodiment head with `domain_name` or `domain_id`; set `raw_action_dim` explicitly when it cannot be inferred from the domain name. +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. + +`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. ## 5. Cosmos3 Parameters diff --git a/python/sglang/multimodal_gen/benchmarks/bench_pi05_openpi.py b/python/sglang/multimodal_gen/benchmarks/bench_pi05_openpi.py index b711e6d69..7632cc626 100644 --- a/python/sglang/multimodal_gen/benchmarks/bench_pi05_openpi.py +++ b/python/sglang/multimodal_gen/benchmarks/bench_pi05_openpi.py @@ -27,7 +27,7 @@ SCRIPT_DIR = Path(__file__).resolve().parent if sys.path and Path(sys.path[0]).resolve() == SCRIPT_DIR: sys.path.pop(0) -from sglang.multimodal_gen.runtime.entrypoints.vla.protocol import ( # noqa: E402 +from sglang.multimodal_gen.runtime.entrypoints.action.protocol import ( # noqa: E402 pack_msgpack, unpack_msgpack, ) @@ -681,10 +681,10 @@ def create_sglang_python_pipeline( def _make_sglang_python_req(server_args, payload: dict[str, Any]): - from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request - from sglang.multimodal_gen.runtime.entrypoints.vla.protocol import ( + from sglang.multimodal_gen.runtime.entrypoints.action.protocol import ( build_action_sampling_params, ) + from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request sampling_params = build_action_sampling_params(payload, server_args) req = prepare_request(server_args, sampling_params) @@ -729,7 +729,9 @@ def run_sglang_python( model_path, pipeline_config_path=pipeline_config_path, ) - from sglang.multimodal_gen.runtime.entrypoints.vla.protocol import action_metadata + from sglang.multimodal_gen.runtime.entrypoints.action.protocol import ( + action_metadata, + ) metadata = action_metadata(server_args) metadata["precision"] = sglang_precision_metadata(pipeline) diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index 887d87ed1..9ad6c558e 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py @@ -272,6 +272,16 @@ class PipelineConfig: del server_args + def supports_action_endpoint(self) -> bool: + """Whether this pipeline exposes the generic action generation API.""" + + return self.task_type.is_action_gen() + + def supports_openpi_endpoint(self) -> bool: + """Whether this pipeline implements the OpenPI policy websocket.""" + + return False + # Wan2.2 TI2V parameters boundary_ratio: float | None = None diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py b/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py index 947d6388e..0a4034989 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py @@ -84,3 +84,9 @@ class Cosmos3Config(PipelineConfig): (num_frames - 1) // vae_scale_factor_temporal ) * vae_scale_factor_temporal + 1 return num_frames + + def supports_action_endpoint(self) -> bool: + # The public Cosmos3 family shares one pipeline/config across visual-only + # and action-capable checkpoints. The loaded transformer validates that + # an action head is actually present when an action request is submitted. + return True diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/pi05.py b/python/sglang/multimodal_gen/configs/pipeline_configs/pi05.py index 0a72daede..71ff1cf9f 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/pi05.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/pi05.py @@ -90,6 +90,9 @@ class Pi05PipelineConfig(PipelineConfig): def supports_native_grouped_requests(self): return True + def supports_openpi_endpoint(self) -> bool: + return True + def estimate_request_cost(self, batch) -> float: return float( self.action_horizon * self.action_dim * self.default_num_inference_steps diff --git a/python/sglang/multimodal_gen/configs/sample/__init__.py b/python/sglang/multimodal_gen/configs/sample/__init__.py index d62cd057c..0251fc720 100644 --- a/python/sglang/multimodal_gen/configs/sample/__init__.py +++ b/python/sglang/multimodal_gen/configs/sample/__init__.py @@ -1,5 +1,6 @@ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo +from sglang.multimodal_gen.configs.sample.action import ActionSamplingParams from sglang.multimodal_gen.configs.sample.diffusers_generic import ( DiffusersGenericSamplingParams, ) @@ -9,11 +10,10 @@ from sglang.multimodal_gen.configs.sample.lingbot_video_moe import ( ) from sglang.multimodal_gen.configs.sample.pi05 import Pi05SamplingParams from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams -from sglang.multimodal_gen.configs.sample.vla import VLASamplingParams __all__ = [ "SamplingParams", - "VLASamplingParams", + "ActionSamplingParams", "DiffusersGenericSamplingParams", "Ideogram4SamplingParams", "Pi05SamplingParams", diff --git a/python/sglang/multimodal_gen/configs/sample/vla.py b/python/sglang/multimodal_gen/configs/sample/action.py similarity index 96% rename from python/sglang/multimodal_gen/configs/sample/vla.py rename to python/sglang/multimodal_gen/configs/sample/action.py index 571f9428d..1062a0379 100644 --- a/python/sglang/multimodal_gen/configs/sample/vla.py +++ b/python/sglang/multimodal_gen/configs/sample/action.py @@ -17,8 +17,8 @@ if TYPE_CHECKING: @dataclass -class VLASamplingParams: - """Sampling parameters for VLA/action-generation policies.""" +class ActionSamplingParams: + """Sampling parameters for policies that generate continuous actions.""" data_type: DataType = DataType.ACTION request_id: str | None = field(default=None, metadata={"batch_sig_exclude": True}) @@ -120,7 +120,7 @@ class VLASamplingParams: def _validate_with_pipeline_config(self, pipeline_config): if not pipeline_config.task_type.is_action_gen(): raise ValueError( - f"VLASamplingParams requires an ACTION pipeline, got {pipeline_config.task_type.name}" + f"ActionSamplingParams requires an ACTION pipeline, got {pipeline_config.task_type.name}" ) def _adjust(self, server_args: "ServerArgs"): @@ -140,7 +140,7 @@ class VLASamplingParams: def _set_output_file_name(self): if self.output_file_name is None: - self.output_file_name = "vla_action" + self.output_file_name = "action" self.output_file_name = _sanitize_filename(self.output_file_name) self._set_output_file_ext() @@ -151,7 +151,7 @@ class VLASamplingParams: def _merge_with_user_params( self, - user_params: "VLASamplingParams", + user_params: "ActionSamplingParams", explicit_fields: set[str] | None = None, ): if user_params is None: @@ -193,7 +193,7 @@ class VLASamplingParams: "--prompt", type=str, nargs="+", - help="Language instruction(s) for the VLA policy.", + help="Language instruction(s) for the action policy.", ) add_argument( "--num-inference-steps", diff --git a/python/sglang/multimodal_gen/configs/sample/cosmos3.py b/python/sglang/multimodal_gen/configs/sample/cosmos3.py index b1de5d02e..e55f46d32 100644 --- a/python/sglang/multimodal_gen/configs/sample/cosmos3.py +++ b/python/sglang/multimodal_gen/configs/sample/cosmos3.py @@ -71,7 +71,38 @@ class Cosmos3SamplingParams(SamplingParams): action_stats_path: str | None = None action_normalization: str = "quantile" + def _adjust(self, server_args) -> None: + action_output = False + if self.action_mode is not None: + self.action_mode = str(self.action_mode).strip().lower() + if self.action_mode not in ( + "policy", + "forward_dynamics", + "inverse_dynamics", + ): + raise ValueError( + f"Unsupported action_mode={self.action_mode!r}; expected " + "'policy', 'forward_dynamics', or 'inverse_dynamics'." + ) + action_output = self.action_mode != "forward_dynamics" + + super()._adjust(server_args) + + # Policy and inverse dynamics produce actions. Forward dynamics consumes + # actions to produce video and therefore remains a visual request. + if action_output: + self.data_type = DataType.ACTION + self.save_output = False + self.return_file_paths_only = False + self.return_frames = False + self.output_file_name = None + self.output_compression = 0 + def _set_output_file_name(self) -> None: + # Action outputs never need a visual filename. This also avoids hashing + # in-memory observation images while base visual adjustment is running. + if self.action_mode in ("policy", "inverse_dynamics"): + return # The pipeline config's ``task_type=TI2V`` drives ``data_type`` to # VIDEO, but a single-frame request is a T2I and must pick the IMAGE # extension. Flip before the base derives the file name. diff --git a/python/sglang/multimodal_gen/configs/sample/pi05.py b/python/sglang/multimodal_gen/configs/sample/pi05.py index 273b04d32..eb0d64706 100644 --- a/python/sglang/multimodal_gen/configs/sample/pi05.py +++ b/python/sglang/multimodal_gen/configs/sample/pi05.py @@ -3,11 +3,11 @@ from dataclasses import dataclass, field from typing import Any -from sglang.multimodal_gen.configs.sample.vla import VLASamplingParams +from sglang.multimodal_gen.configs.sample.action import ActionSamplingParams @dataclass -class Pi05SamplingParams(VLASamplingParams): +class Pi05SamplingParams(ActionSamplingParams): """Sampling parameters for Pi0.5 flow-matching action inference.""" num_inference_steps: int = 10 diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index cb89dd589..27e574f4d 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -1077,6 +1077,7 @@ def _register_configs(): pipeline_config_cls=Cosmos3Config, hf_model_paths=[ "nvidia/Cosmos3-Nano", + "nvidia/Cosmos3-Nano-Policy-DROID", "nvidia/Cosmos3-Super", "nvidia/Cosmos3-Super-Text2Image", "nvidia/Cosmos3-Super-Image2Video", diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/vla/__init__.py b/python/sglang/multimodal_gen/runtime/entrypoints/action/__init__.py similarity index 100% rename from python/sglang/multimodal_gen/runtime/entrypoints/vla/__init__.py rename to python/sglang/multimodal_gen/runtime/entrypoints/action/__init__.py diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/vla/api.py b/python/sglang/multimodal_gen/runtime/entrypoints/action/api.py similarity index 95% rename from python/sglang/multimodal_gen/runtime/entrypoints/vla/api.py rename to python/sglang/multimodal_gen/runtime/entrypoints/action/api.py index cc901ea83..dc2a7ea5f 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/vla/api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/action/api.py @@ -4,7 +4,7 @@ from __future__ import annotations from fastapi import APIRouter, HTTPException, Request, Response, WebSocket -from sglang.multimodal_gen.runtime.entrypoints.vla.protocol import ( +from sglang.multimodal_gen.runtime.entrypoints.action.protocol import ( action_generation_response, action_metadata, action_raw_response, @@ -12,7 +12,7 @@ from sglang.multimodal_gen.runtime.entrypoints.vla.protocol import ( pack_msgpack, unpack_msgpack, ) -from sglang.multimodal_gen.runtime.entrypoints.vla.ws_utils import ( +from sglang.multimodal_gen.runtime.entrypoints.action.ws_utils import ( run_action_msgpack_ws, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/vla/openpi.py b/python/sglang/multimodal_gen/runtime/entrypoints/action/openpi.py similarity index 90% rename from python/sglang/multimodal_gen/runtime/entrypoints/vla/openpi.py rename to python/sglang/multimodal_gen/runtime/entrypoints/action/openpi.py index b12c3cbe4..35cfa314e 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/vla/openpi.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/action/openpi.py @@ -6,7 +6,7 @@ from typing import Any from fastapi import APIRouter, WebSocket -from sglang.multimodal_gen.runtime.entrypoints.vla.ws_utils import ( +from sglang.multimodal_gen.runtime.entrypoints.action.ws_utils import ( run_action_msgpack_ws, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/vla/protocol.py b/python/sglang/multimodal_gen/runtime/entrypoints/action/protocol.py similarity index 60% rename from python/sglang/multimodal_gen/runtime/entrypoints/vla/protocol.py rename to python/sglang/multimodal_gen/runtime/entrypoints/action/protocol.py index c8ef4ddba..2e4ced790 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/vla/protocol.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/action/protocol.py @@ -13,7 +13,10 @@ from typing import Any import numpy as np from PIL import Image -from sglang.multimodal_gen.configs.sample.vla import VLASamplingParams +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.utils import prepare_request from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client from sglang.multimodal_gen.runtime.server_args import ServerArgs @@ -107,6 +110,9 @@ def _normalize_observation(observation: dict[str, Any]) -> dict[str, Any]: normalized["images"] = { name: _normalize_image_value(value) for name, value in images.items() } + for name in ("image", "image_path", "input_reference"): + if name in normalized: + normalized[name] = _normalize_image_value(normalized[name]) state = normalized.get("state") if isinstance(state, dict): normalized["state"] = _decode_tensor_payload(state) @@ -141,6 +147,53 @@ 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.model_id or server_args.model_path, + "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, + }, + } + policy_family = getattr( pipeline_config, "policy_family", @@ -226,6 +279,9 @@ def _action_request_to_observation(payload: dict[str, Any]) -> dict[str, Any]: observation["state"] = input_payload["state"] if "noise" in input_payload: observation["noise"] = input_payload["noise"] + for name in ("image", "image_path", "input_reference", "video", "video_path"): + if name in input_payload: + observation[name] = input_payload[name] return _normalize_observation(observation) @@ -235,14 +291,14 @@ def _resolve_action_sampling_params_cls_cached( backend: str | None, model_id: str | None, pipeline_class_name: str | None, -) -> type[VLASamplingParams]: +) -> type[SamplingParams] | type[ActionSamplingParams]: if pipeline_class_name: from sglang.multimodal_gen.registry import get_pipeline_config_classes config_classes = get_pipeline_config_classes(pipeline_class_name) if config_classes is not None: _, sampling_params_cls = config_classes - if issubclass(sampling_params_cls, VLASamplingParams): + if issubclass(sampling_params_cls, (SamplingParams, ActionSamplingParams)): return sampling_params_cls from sglang.multimodal_gen.registry import get_model_info @@ -253,16 +309,17 @@ def _resolve_action_sampling_params_cls_cached( model_id=model_id, ) sampling_params_cls = model_info.sampling_param_cls - if not issubclass(sampling_params_cls, VLASamplingParams): + if not issubclass(sampling_params_cls, (SamplingParams, ActionSamplingParams)): raise ValueError( - f"Action endpoint requires VLASamplingParams, got {sampling_params_cls.__name__}" + "Action endpoint requires SamplingParams or ActionSamplingParams, got " + f"{sampling_params_cls.__name__}" ) return sampling_params_cls def _resolve_action_sampling_params_cls( server_args: ServerArgs, -) -> type[VLASamplingParams]: +) -> type[SamplingParams] | type[ActionSamplingParams]: return _resolve_action_sampling_params_cls_cached( server_args.model_path, getattr(server_args, "backend", None), @@ -273,15 +330,16 @@ def _resolve_action_sampling_params_cls( @lru_cache(maxsize=32) def _sampling_params_field_names( - sampling_params_cls: type[VLASamplingParams], + sampling_params_cls: type[SamplingParams] | type[ActionSamplingParams], ) -> frozenset[str]: return frozenset(field.name for field in dataclasses.fields(sampling_params_cls)) -def build_action_sampling_params( +def _build_action_model_sampling_params( payload: dict[str, Any], server_args: ServerArgs, -) -> VLASamplingParams: + sampling_params_cls: type[ActionSamplingParams], +) -> ActionSamplingParams: pipeline_config = server_args.pipeline_config observation = _action_request_to_observation(payload) parameters = dict(payload.get("parameters") or {}) @@ -318,7 +376,6 @@ def build_action_sampling_params( if output_format not in ("list", "numpy"): raise ValueError("output_format must be 'list' or 'numpy'") - sampling_params_cls = _resolve_action_sampling_params_cls(server_args) sampling_kwargs = { "request_id": payload.get("request_id") or payload.get("id"), "prompt": prompt, @@ -366,6 +423,156 @@ def build_action_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, +) -> SamplingParams | ActionSamplingParams: + sampling_params_cls = _resolve_action_sampling_params_cls(server_args) + if issubclass(sampling_params_cls, ActionSamplingParams): + return _build_action_model_sampling_params( + payload, + server_args, + sampling_params_cls, + ) + if issubclass(sampling_params_cls, Cosmos3SamplingParams): + return _build_cosmos3_action_sampling_params( + payload, + server_args, + sampling_params_cls, + ) + raise ValueError( + f"Action endpoint is not implemented for {sampling_params_cls.__name__}" + ) + + async def infer_action( payload: dict[str, Any], server_args: ServerArgs, @@ -395,6 +602,22 @@ def action_generation_response( 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] + + pipeline_config = server_args.pipeline_config + if isinstance(pipeline_config, Cosmos3Config): + default_num_inference_steps = Cosmos3SamplingParams().num_inference_steps + else: + default_num_inference_steps = pipeline_config.default_num_inference_steps + response = { "id": output.get("request_id") or f"act_{uuid.uuid4().hex}", "object": "action.generation", @@ -405,12 +628,7 @@ def action_generation_response( "index": 0, "input_index": 0, "candidate_index": 0, - "action": { - "type": "continuous", - "dtype": "float32", - "shape": action_shape, - "values": action_values, - }, + "action": action, } ], "usage": { @@ -418,7 +636,7 @@ def action_generation_response( "action_dim": action_shape[1] if len(action_shape) > 1 else 0, "denoise_steps": output.get("parameters", {}).get( "num_inference_steps", - server_args.pipeline_config.default_num_inference_steps, + default_num_inference_steps, ), "prefix_cache_hit": bool(output.get("cache", {}).get("hit", False)), }, diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/vla/ws_utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/action/ws_utils.py similarity index 96% rename from python/sglang/multimodal_gen/runtime/entrypoints/vla/ws_utils.py rename to python/sglang/multimodal_gen/runtime/entrypoints/action/ws_utils.py index d6916c0da..1251f13ce 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/vla/ws_utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/action/ws_utils.py @@ -9,7 +9,7 @@ from typing import Any from fastapi import WebSocket, WebSocketDisconnect -from sglang.multimodal_gen.runtime.entrypoints.vla.protocol import ( +from sglang.multimodal_gen.runtime.entrypoints.action.protocol import ( action_metadata, infer_action, pack_msgpack, diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py b/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py index e29123a2b..2b16736d4 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py @@ -14,6 +14,8 @@ from fastapi import APIRouter, FastAPI, Request, Response from fastapi.middleware.cors import CORSMiddleware from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams +from sglang.multimodal_gen.runtime.entrypoints.action import api as action_api +from sglang.multimodal_gen.runtime.entrypoints.action import openpi from sglang.multimodal_gen.runtime.entrypoints.openai import image_api, video_api from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import ( VertexGenerateReqInput, @@ -30,8 +32,6 @@ from sglang.multimodal_gen.runtime.entrypoints.utils import ( prepare_request, save_outputs, ) -from sglang.multimodal_gen.runtime.entrypoints.vla import api as vla_api -from sglang.multimodal_gen.runtime.entrypoints.vla import openpi from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client from sglang.multimodal_gen.runtime.server_args import ServerArgs, get_global_server_args from sglang.multimodal_gen.runtime.server_warmup import ( @@ -423,8 +423,9 @@ def create_app(server_args: ServerArgs): app.include_router(image_api.router) app.include_router(video_api.router) app.include_router(realtime_video_api.router) - if server_args.pipeline_config.task_type.is_action_gen(): - app.include_router(vla_api.router) + if server_args.pipeline_config.supports_action_endpoint(): + app.include_router(action_api.router) + if server_args.pipeline_config.supports_openpi_endpoint(): app.include_router(openpi.router) app.include_router(mesh_api.router) app.include_router(weights_api.router) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py index d96df3ae6..676b3ad1a 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py @@ -23,6 +23,7 @@ from fastapi import ( from fastapi.responses import FileResponse from sglang.multimodal_gen.configs.sample.sampling_params import ( + DataType, SamplingParams, generate_request_id, ) @@ -410,7 +411,16 @@ def _build_video_sampling_params(request_id: str, request: VideoGenerationsReque sampling_params_cls = _video_sampling_params_cls(server_args) kwargs = sampling_params_cls.lower_video_request_kwargs(request, kwargs) - return build_sampling_params(request_id, **kwargs) + sampling_params = build_sampling_params(request_id, **kwargs) + if ( + isinstance(sampling_params, SamplingParams) + and sampling_params.data_type == DataType.ACTION + ): + raise ValueError( + "Action-producing policy and inverse-dynamics requests use " + "/v1/actions/generations; /v1/videos is reserved for visual outputs" + ) + return sampling_params # extract metadata which http_server needs to know diff --git a/python/sglang/multimodal_gen/runtime/launch_server.py b/python/sglang/multimodal_gen/runtime/launch_server.py index 85fcbeddb..2820c3df0 100644 --- a/python/sglang/multimodal_gen/runtime/launch_server.py +++ b/python/sglang/multimodal_gen/runtime/launch_server.py @@ -305,9 +305,9 @@ def launch_server(server_args: ServerArgs, launch_http_server: bool = True): return processes if launch_http_server: - if server_args.pipeline_config.task_type.is_action_gen(): + if server_args.pipeline_config.supports_action_endpoint(): logger.info( - "VLA pipeline ready: model=%s; per-request details are " + "Action generation endpoint ready: model=%s; per-request details are " "debug-only (use --log-level debug).", server_args.model_id or server_args.model_path, ) 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 169a79459..ac24aaff9 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 @@ -56,7 +56,7 @@ from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.profiler import SGLDiffusionProfiler -from sglang.multimodal_gen.runtime.utils.vision import load_video +from sglang.multimodal_gen.runtime.utils.vision import load_image, load_video from sglang.srt.utils.common import get_compiler_backend logger = init_logger(__name__) @@ -129,8 +129,8 @@ class Cosmos3ImagePreprocessStage(PipelineStage): target_h, target_w = batch.height, batch.width - if isinstance(image_path, str) and image_path: - image = PIL.Image.open(image_path).convert("RGB") + 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}") @@ -440,9 +440,9 @@ class Cosmos3LatentPreparationStage(PipelineStage): noise = torch.randn(shape, generator=generator, device=device, dtype=dtype) - is_video_gen = batch.data_type == DataType.VIDEO - has_image_cond = batch.preprocessed_image is not None and is_video_gen - has_video_cond = batch.preprocessed_video is not None and is_video_gen + uses_visual_latents = batch.data_type in (DataType.VIDEO, DataType.ACTION) + has_image_cond = batch.preprocessed_image is not None and uses_visual_latents + has_video_cond = batch.preprocessed_video is not None and uses_visual_latents if has_image_cond or has_video_cond: vae_dtype = next(self.vae.parameters()).dtype @@ -1387,6 +1387,56 @@ class Cosmos3DecodingStage(PipelineStage): OutputBatch, ) + action_pred = None + if getattr(batch, "action_latents", None) is not None: + raw_action_dim = batch.extra.get("raw_action_dim") + action_pred = batch.action_latents.float().cpu() + if raw_action_dim is not None: + action_pred = action_pred[:, :, :raw_action_dim] + stats_path = getattr(batch.sampling_params, "action_stats_path", None) + if stats_path is not None: + method = getattr( + batch.sampling_params, "action_normalization", "quantile" + ) + action_pred = denormalize_action( + action_pred, method, load_action_stats(stats_path) + ) + self.log_info(f"Action predictions shape: {tuple(action_pred.shape)}") + + action_domain_ids = batch.extra.get("action_domain_ids") + action_domain_id = ( + int(action_domain_ids[0].item()) if action_domain_ids is not None else None + ) + action_metadata = { + "action_mode": getattr(batch.sampling_params, "action_mode", None), + "action_domain_id": action_domain_id, + "action_raw_action_dim": ( + batch.extra.get("raw_action_dim") + if getattr(batch, "extra", None) + else None + ), + } + if batch.data_type == DataType.ACTION: + if action_pred is None: + raise RuntimeError("Cosmos3 action request produced no action tensor") + payload = { + "request_id": batch.request_id, + "actions": action_pred[0].numpy(), + "action_mode": action_metadata["action_mode"], + "domain_id": action_metadata["action_domain_id"], + "raw_action_dim": action_metadata["action_raw_action_dim"], + "parameters": { + "num_inference_steps": batch.num_inference_steps, + "num_frames": batch.num_frames, + }, + } + return OutputBatch( + output=[payload], + action_pred=action_pred, + metrics=batch.metrics if hasattr(batch, "metrics") else None, + **action_metadata, + ) + is_image_gen = batch.data_type == DataType.IMAGE self.log_info( "Decoding latents to image..." @@ -1438,33 +1488,11 @@ class Cosmos3DecodingStage(PipelineStage): f"Decoded audio tensor shape: {tuple(audio.shape)} @ {audio_sample_rate} Hz" ) - action_pred = None - if getattr(batch, "action_latents", None) is not None: - raw_action_dim = batch.extra.get("raw_action_dim") - action_pred = batch.action_latents.float().cpu() - if raw_action_dim is not None: - action_pred = action_pred[:, :, :raw_action_dim] - stats_path = getattr(batch.sampling_params, "action_stats_path", None) - if stats_path is not None: - method = getattr( - batch.sampling_params, "action_normalization", "quantile" - ) - action_pred = denormalize_action( - action_pred, method, load_action_stats(stats_path) - ) - self.log_info(f"Action predictions shape: {tuple(action_pred.shape)}") - return OutputBatch( output=output, audio=audio, audio_sample_rate=audio_sample_rate, action_pred=action_pred, - action_mode=getattr(batch.sampling_params, "action_mode", None), - action_domain_id=getattr(batch.sampling_params, "domain_id", None), - action_raw_action_dim=( - batch.extra.get("raw_action_dim") - if getattr(batch, "extra", None) - else None - ), metrics=batch.metrics if hasattr(batch, "metrics") else None, + **action_metadata, ) diff --git a/python/sglang/multimodal_gen/test/unit/test_cosmos3.py b/python/sglang/multimodal_gen/test/unit/test_cosmos3.py index 42c9f7a82..68cd4239e 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cosmos3.py +++ b/python/sglang/multimodal_gen/test/unit/test_cosmos3.py @@ -18,6 +18,11 @@ from sglang.multimodal_gen.registry import ( _get_config_info, get_non_diffusers_pipeline_name, ) +from sglang.multimodal_gen.runtime.entrypoints.action.protocol import ( + action_generation_response, + action_metadata, + build_action_sampling_params, +) from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import ( ImageGenerationsRequest, VideoGenerationsRequest, @@ -39,6 +44,7 @@ from sglang.multimodal_gen.runtime.models.dits.cosmos3video import ( compute_mrope_position_ids_vision, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3 import ( + Cosmos3DecodingStage, Cosmos3ImagePreprocessStage, Cosmos3LatentPreparationStage, Cosmos3TimestepPreparationStage, @@ -56,6 +62,23 @@ def _apply(mapping_fn, key): return mapping_fn(key) +def _cosmos3_server_args(config=None): + return types.SimpleNamespace( + model_id=None, + model_path="nvidia/Cosmos3-Nano", + backend=None, + pipeline_class_name=None, + output_path=None, + comfyui_mode=False, + num_gpus=1, + tp_size=1, + sp_degree=1, + ulysses_degree=1, + ring_degree=1, + pipeline_config=config or Cosmos3Config(), + ) + + class TestCosmos3ParamNamesMapping(unittest.TestCase): """Verify diffusers → sglang weight key translations.""" @@ -328,6 +351,167 @@ class TestCosmos3SamplingParamsDataType(unittest.TestCase): params._set_output_file_name() self.assertEqual(params.data_type, DataType.VIDEO) + def test_policy_adjusts_to_action_output(self): + params = Cosmos3SamplingParams( + prompt="test", + action_mode="policy", + num_frames=17, + image_path="observation.png", + ) + + params._adjust(_cosmos3_server_args()) + + self.assertEqual(params.data_type, DataType.ACTION) + self.assertFalse(params.save_output) + self.assertFalse(params.return_file_paths_only) + self.assertIsNone(params.output_file_name) + self.assertEqual(params.num_frames, 17) + + def test_forward_dynamics_remains_video_output(self): + params = Cosmos3SamplingParams( + prompt="test", + action_mode="forward_dynamics", + num_frames=17, + image_path="observation.png", + ) + + params._adjust(_cosmos3_server_args()) + + self.assertEqual(params.data_type, DataType.VIDEO) + + +class TestCosmos3ActionEndpoint(unittest.TestCase): + def test_policy_request_builds_action_sampling_params(self): + image = torch.zeros(8, 8, 3, dtype=torch.uint8).numpy() + payload = { + "request_id": "cosmos-action-1", + "input": { + "task": "pick up the block", + "observation": { + "image": { + "dtype": "uint8", + "shape": [8, 8, 3], + "values": image.tolist(), + } + }, + }, + "parameters": { + "action_mode": "policy", + "action_horizon": 16, + "domain_name": "droid_lerobot", + "num_inference_steps": 30, + "height": 480, + "width": 832, + "fps": 5, + "seed": 7, + }, + } + + params = build_action_sampling_params(payload, _cosmos3_server_args()) + + self.assertIsInstance(params, Cosmos3SamplingParams) + self.assertEqual(params.data_type, DataType.ACTION) + self.assertEqual(params.prompt, "pick up the block") + self.assertEqual(params.action_mode, "policy") + self.assertEqual(params.domain_name, "droid_lerobot") + self.assertEqual(params.num_frames, 17) + self.assertEqual(params.num_inference_steps, 30) + self.assertEqual(params.seed, 7) + self.assertEqual(params.image_path.size, (8, 8)) + + def test_inverse_dynamics_maps_video_input(self): + payload = { + "input": { + "task": "infer the robot motion", + "observation": {"video": "observation.mp4"}, + }, + "parameters": { + "action_mode": "inverse_dynamics", + "num_frames": 61, + "domain_name": "av", + }, + } + + params = build_action_sampling_params(payload, _cosmos3_server_args()) + + self.assertEqual(params.data_type, DataType.ACTION) + self.assertEqual(params.video_path, "observation.mp4") + self.assertEqual(params.num_frames, 61) + + def test_forward_dynamics_is_rejected_by_action_endpoint(self): + payload = { + "input": { + "task": "predict the next frames", + "observation": {"image": "observation.png"}, + }, + "parameters": {"action_mode": "forward_dynamics"}, + } + + with self.assertRaisesRegex(ValueError, "/v1/videos"): + build_action_sampling_params(payload, _cosmos3_server_args()) + + def test_metadata_describes_cosmos_action_contract(self): + metadata = action_metadata(_cosmos3_server_args()) + + self.assertEqual(metadata["policy_family"], "cosmos3") + self.assertEqual(metadata["input"]["modalities"], ["image", "video"]) + self.assertEqual(metadata["output"]["action_horizon"], 16) + self.assertEqual(metadata["output"]["padded_action_dim"], 64) + self.assertFalse(metadata["capabilities"]["openpi_websocket"]) + + def test_action_response_includes_cosmos_metadata(self): + output = { + "request_id": "cosmos-action-2", + "actions": torch.zeros(16, 10).numpy(), + "action_mode": "policy", + "domain_id": 8, + "raw_action_dim": 10, + "parameters": {"num_inference_steps": 30}, + } + + response = action_generation_response(output, _cosmos3_server_args()) + action = response["data"][0]["action"] + + self.assertEqual(action["shape"], [16, 10]) + self.assertEqual(action["action_mode"], "policy") + self.assertEqual(action["domain_id"], 8) + self.assertEqual(action["raw_action_dim"], 10) + self.assertEqual(response["usage"]["denoise_steps"], 30) + + def test_action_decode_skips_vae(self): + class FailIfDecoded: + def decode(self, _latents): + raise AssertionError("VAE decode must not run for action output") + + stage = Cosmos3DecodingStage.__new__(Cosmos3DecodingStage) + stage.vae = FailIfDecoded() + 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(24, dtype=torch.float32).reshape(1, 4, 6), + extra={ + "raw_action_dim": 3, + "action_domain_ids": torch.tensor([8]), + }, + sampling_params=Cosmos3SamplingParams( + prompt="test", + action_mode="policy", + domain_name="droid_lerobot", + ), + request_id="cosmos-action-3", + 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, (4, 3)) + self.assertEqual(output.output[0]["domain_id"], 8) + self.assertEqual(output.action_pred.shape, (1, 4, 3)) + class TestCosmos3ModelResolution(unittest.TestCase): """Verify Cosmos3 checkpoints resolve to the native SGLang pipeline.""" @@ -335,6 +519,7 @@ class TestCosmos3ModelResolution(unittest.TestCase): def test_hf_checkpoint_uses_registered_native_pipeline_config(self): for model_path in ( "nvidia/Cosmos3-Nano", + "nvidia/Cosmos3-Nano-Policy-DROID", "nvidia/Cosmos3-Super", "nvidia/Cosmos3-Super-Text2Image", "nvidia/Cosmos3-Super-Image2Video", diff --git a/python/sglang/multimodal_gen/test/unit/test_pi05_action_api.py b/python/sglang/multimodal_gen/test/unit/test_pi05_action_api.py index 720e12086..0c26b0e47 100644 --- a/python/sglang/multimodal_gen/test/unit/test_pi05_action_api.py +++ b/python/sglang/multimodal_gen/test/unit/test_pi05_action_api.py @@ -6,13 +6,13 @@ from types import SimpleNamespace import numpy as np from sglang.multimodal_gen.configs.pipeline_configs.pi05 import Pi05PipelineConfig +from sglang.multimodal_gen.configs.sample.action import ActionSamplingParams from sglang.multimodal_gen.configs.sample.pi05 import Pi05SamplingParams from sglang.multimodal_gen.configs.sample.sampling_params import ( DataType, SamplingParams, ) -from sglang.multimodal_gen.configs.sample.vla import VLASamplingParams -from sglang.multimodal_gen.runtime.entrypoints.vla.protocol import ( +from sglang.multimodal_gen.runtime.entrypoints.action.protocol import ( action_generation_response, action_metadata, action_raw_response, @@ -38,11 +38,11 @@ def _server_args(config: Pi05PipelineConfig | None = None) -> SimpleNamespace: ) -def test_pi05_uses_vla_sampling_params_not_visual_sampling_params(): +def test_pi05_uses_action_sampling_params_not_visual_sampling_params(): params = Pi05SamplingParams() field_names = {field.name for field in dataclasses.fields(params)} - assert isinstance(params, VLASamplingParams) + assert isinstance(params, ActionSamplingParams) assert not isinstance(params, SamplingParams) assert "action_horizon" in field_names assert "action_dim" in field_names