minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] feat: expose cosmos3 policies through the Action API (#34243)

Browse Source
This commit is contained in:
Mick
2026-08-10 18:16:20 +08:00
committed by GitHub
parent c971d7ac9c
commit 955569a2dc
21 changed files with 608 additions and 92 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/benchmarks/bench_pi05_openpi.py
+6 -4
View File
@@ -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)
python/sglang/multimodal_gen/configs/pipeline_configs/base.py
+10
View File
@@ -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
python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py
+6
View File
@@ -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
python/sglang/multimodal_gen/configs/pipeline_configs/pi05.py
+3
View File
@@ -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
python/sglang/multimodal_gen/configs/sample/__init__.py
+2 -2
View File
@@ -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",
python/sglang/multimodal_gen/configs/sample/vla.py → python/sglang/multimodal_gen/configs/sample/action.py
+6 -6
View File
@@ -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",
python/sglang/multimodal_gen/configs/sample/cosmos3.py
+31
View File
@@ -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.
python/sglang/multimodal_gen/configs/sample/pi05.py
+2 -2
View File
@@ -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
python/sglang/multimodal_gen/registry.py
+1
View File
@@ -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",
python/sglang/multimodal_gen/runtime/entrypoints/vla/__init__.py → python/sglang/multimodal_gen/runtime/entrypoints/action/__init__.py
View File
python/sglang/multimodal_gen/runtime/entrypoints/vla/api.py → python/sglang/multimodal_gen/runtime/entrypoints/action/api.py
+2 -2
View File
@@ -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
python/sglang/multimodal_gen/runtime/entrypoints/vla/openpi.py → python/sglang/multimodal_gen/runtime/entrypoints/action/openpi.py
+1 -1
View File
@@ -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
python/sglang/multimodal_gen/runtime/entrypoints/vla/protocol.py → python/sglang/multimodal_gen/runtime/entrypoints/action/protocol.py
+235 -17
View File
@@ -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)),
},
python/sglang/multimodal_gen/runtime/entrypoints/vla/ws_utils.py → python/sglang/multimodal_gen/runtime/entrypoints/action/ws_utils.py
+1 -1
View File
@@ -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,
python/sglang/multimodal_gen/runtime/entrypoints/http_server.py
+5 -4
View File
@@ -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)
python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py
+11 -1
View File
@@ -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
python/sglang/multimodal_gen/runtime/launch_server.py
+2 -2
View File
@@ -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,
)
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py
+57 -29
View File
@@ -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,
)
python/sglang/multimodal_gen/test/unit/test_cosmos3.py
+185
View File
@@ -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",
python/sglang/multimodal_gen/test/unit/test_pi05_action_api.py
+4 -4
View File
@@ -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