From 1988a2c9ea2135d3e8d00c1a034464e88d92739c Mon Sep 17 00:00:00 2001 From: Mick Date: Tue, 2 Jun 2026 00:53:39 +0800 Subject: [PATCH] [diffusion] feat: improve cosmos3 serve API support (#26926) --- .../cookbook/diffusion/Cosmos/Cosmos3.mdx | 197 ++++++++++++++++++ docs_new/docs.json | 6 + .../configs/sample/sampling_params.py | 8 + python/sglang/multimodal_gen/registry.py | 2 + .../runtime/entrypoints/openai/image_api.py | 57 ++++- .../runtime/entrypoints/openai/protocol.py | 6 + .../runtime/entrypoints/openai/utils.py | 25 +++ .../runtime/entrypoints/openai/video_api.py | 177 ++++++++++++---- .../runtime/entrypoints/utils.py | 3 + .../runtime/pipelines/cosmos3_pipeline.py | 17 +- .../runtime/pipelines_core/schedule_batch.py | 8 +- .../stages/model_specific_stages/cosmos3.py | 28 ++- .../cosmos3_guardrails.py | 13 +- .../multimodal_gen/runtime/server_args.py | 15 +- .../multimodal_gen/test/server/gpu_cases.py | 26 +++ .../test/server/test_server_common.py | 24 +-- .../test/server/test_server_utils.py | 3 +- .../sglang/multimodal_gen/test/test_utils.py | 5 +- .../multimodal_gen/test/unit/test_cosmos3.py | 87 ++++++++ 19 files changed, 637 insertions(+), 70 deletions(-) create mode 100644 docs_new/cookbook/diffusion/Cosmos/Cosmos3.mdx diff --git a/docs_new/cookbook/diffusion/Cosmos/Cosmos3.mdx b/docs_new/cookbook/diffusion/Cosmos/Cosmos3.mdx new file mode 100644 index 000000000..c3c5d3a89 --- /dev/null +++ b/docs_new/cookbook/diffusion/Cosmos/Cosmos3.mdx @@ -0,0 +1,197 @@ +--- +title: Cosmos3 +metatags: + description: "Serve NVIDIA Cosmos3 text-to-video, image-to-video, and text-to-image generation with SGLang Diffusion." +--- + +## 1. Model Introduction + +[NVIDIA Cosmos3](https://huggingface.co/collections/nvidia/cosmos3) is a world-generation model family for text-to-image, text-to-video, and image-to-video generation. SGLang Diffusion serves the public generator checkpoints with the native `Cosmos3OmniDiffusersPipeline`. + +| Model | Status | Notes | +| --- | --- | --- | +| `nvidia/Cosmos3-Nano` | Supported | T2I, T2V, I2V | +| `nvidia/Cosmos3-Super` | Supported | T2I, T2V, I2V; use multi-GPU for the 64B checkpoint | +| `nvidia/Cosmos3-Super-Text2Image` | Supported | T2I-specialized checkpoint | +| `nvidia/Cosmos3-Super-Image2Video` | Supported | I2V-specialized checkpoint | +| `nvidia/Cosmos3-Nano-Policy-DROID` | Not supported yet | Action/policy model; planned separately from visual generation | + +Cosmos3 video-with-sound, video-to-video conditioning, and action generation are not supported yet. Requests that set `generate_sound`, `action_mode`, or video-to-video conditioning fields return a clear error instead of being silently ignored. + +## 2. Installation + +Install SGLang with the diffusion dependencies: + +```bash Command +pip install -e "python[diffusion]" +``` + +Cosmos3 guardrails are enabled by default when the package is available: + +```bash Command +pip install "cosmos-guardrail==0.3.1" +``` + +`cosmos-guardrail` downloads gated NVIDIA guardrail weights, so pass a Hugging Face token if your environment needs one. If the package is not installed, SGLang skips Cosmos3 guardrails and logs a warning. To disable Cosmos3 guardrails for local experiments, set `SGLANG_DISABLE_COSMOS3_GUARDRAILS=1` before starting the server. + +## 3. Serve Cosmos3 + +Serve `Cosmos3-Nano` directly from the Hugging Face model ID: + +```bash Command +sglang serve \ + --model-type diffusion \ + --model-path nvidia/Cosmos3-Nano \ + --num-gpus 1 \ + --host 0.0.0.0 \ + --port 30010 \ + --output-path /tmp/sglang-cosmos3 +``` + +For `Cosmos3-Super`, split the model across multiple GPUs: + +```bash Command +sglang serve \ + --model-type diffusion \ + --model-path nvidia/Cosmos3-Super \ + --num-gpus 4 \ + --host 0.0.0.0 \ + --port 30010 \ + --output-path /tmp/sglang-cosmos3 +``` + +The server also accepts the specialized `nvidia/Cosmos3-Super-Text2Image` and `nvidia/Cosmos3-Super-Image2Video` checkpoint IDs. + +## 4. OpenAI-Compatible Requests + +### Text to image + +Cosmos3 text-to-image uses `/v1/images/generations`. The default Cosmos3 image response is `b64_json`, matching vLLM-Omni's examples. + +```bash Command +curl -sS -X POST http://127.0.0.1:30010/v1/images/generations \ + -H "Content-Type: application/json" \ + -d '{ + "prompt": "A warehouse robot folds a blue cloth on a clean workbench.", + "size": "1280x720", + "n": 1, + "num_inference_steps": 35, + "guidance_scale": 6.0, + "flow_shift": 10.0, + "seed": 0, + "extra_args": { + "use_resolution_template": false, + "guardrails": true + } + }' +``` + +### Text to video + +Use `/v1/videos` to create an asynchronous job, then poll the job and download the completed MP4. + +```bash Command +job_id=$(curl -sS -X POST http://127.0.0.1:30010/v1/videos \ + --form-string "prompt=A small warehouse robot moves a blue box across a clean floor." \ + --form-string "negative_prompt=blurry, distorted, low quality" \ + --form-string "size=1280x720" \ + --form-string "num_frames=81" \ + --form-string "fps=24" \ + --form-string "num_inference_steps=35" \ + --form-string "guidance_scale=4.0" \ + --form-string "flow_shift=10.0" \ + --form-string "seed=42" \ + --form-string 'extra_params={"guardrails":true,"use_resolution_template":false,"use_duration_template":false}' \ + | python -c 'import json, sys; print(json.load(sys.stdin)["id"])') + +while true; do + status=$(curl -sS "http://127.0.0.1:30010/v1/videos/${job_id}" \ + | python -c 'import json, sys; print(json.load(sys.stdin)["status"])') + [ "$status" = "completed" ] && break + [ "$status" = "failed" ] && exit 1 + sleep 1 +done + +curl -sS -L "http://127.0.0.1:30010/v1/videos/${job_id}/content" \ + -o cosmos3_t2v.mp4 +``` + +### Image to video + +This mirrors the official `nvidia/Cosmos3-Nano` Hugging Face image-to-video example: + +```python Python +import json +import time +from pathlib import Path + +import requests +from huggingface_hub import snapshot_download + +base_url = "http://127.0.0.1:30010" +model_dir = Path(snapshot_download("nvidia/Cosmos3-Nano")) +asset_dir = model_dir / "assets" + +prompt = json.dumps(json.loads((asset_dir / "example_i2v_prompt.json").read_text())) +negative_prompt = json.dumps( + json.loads((asset_dir / "negative_prompt.json").read_text()) +) + +data = { + "prompt": prompt, + "negative_prompt": negative_prompt, + "size": "1280x720", + "num_frames": "189", + "fps": "24", + "num_inference_steps": "35", + "guidance_scale": "6.0", + "max_sequence_length": "4096", + "flow_shift": "10.0", + "seed": "1111", + "extra_params": json.dumps( + { + "use_resolution_template": False, + "use_duration_template": False, + "guardrails": True, + } + ), +} + +with (asset_dir / "example_i2v_input.jpg").open("rb") as image: + response = requests.post( + f"{base_url}/v1/videos", + data=data, + files={"input_reference": ("example_i2v_input.jpg", image, "image/jpeg")}, + timeout=60, + ) +response.raise_for_status() +video_id = response.json()["id"] + +while True: + job = requests.get(f"{base_url}/v1/videos/{video_id}", timeout=30).json() + if job["status"] == "completed": + break + if job["status"] == "failed": + raise RuntimeError(job.get("error") or "Video generation failed") + time.sleep(1) + +response = requests.get(f"{base_url}/v1/videos/{video_id}/content", timeout=300) +response.raise_for_status() +Path("cosmos3_i2v.mp4").write_bytes(response.content) +``` + +## 5. Cosmos3 Parameters + +Cosmos3 supports the standard SGLang video and image fields such as `size`, `num_frames`, `fps`, `num_inference_steps`, `guidance_scale`, `negative_prompt`, and `seed`. + +Top-level Cosmos3 request fields: + +- `max_sequence_length`: maximum text token length used by the Cosmos3 tokenizer. +- `flow_shift`: per-request scheduler flow shift. If omitted, SGLang uses `--flow-shift`, then the checkpoint scheduler default. + +Put model-specific compatibility knobs in `extra_params` for video requests, or `extra_args` for image requests: + +- `use_duration_template`: whether to append SGLang's generated duration suffix to video prompts. +- `use_resolution_template`: accepted for vLLM-Omni request compatibility. +- `use_system_prompt`: whether to add the Cosmos3 system prompt to the chat template. +- `guardrails` or `use_guardrails`: per-request guardrail toggle when the server started with guardrails enabled. diff --git a/docs_new/docs.json b/docs_new/docs.json index 09a37e672..c55d4b426 100644 --- a/docs_new/docs.json +++ b/docs_new/docs.json @@ -1114,6 +1114,12 @@ "group": "Diffusion Models", "pages": [ "cookbook/diffusion/intro", + { + "group": "Cosmos", + "pages": [ + "cookbook/diffusion/Cosmos/Cosmos3" + ] + }, { "group": "FLUX", "pages": [ diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index 84fdf29c7..50f446480 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -219,6 +219,14 @@ class SamplingParams: return_file_paths_only: bool = True enable_sequence_shard: bool | None = None diffusers_kwargs: dict | None = None + max_sequence_length: int | None = None + flow_shift: float | None = None + + # cosmos-related + use_duration_template: bool | None = None + use_resolution_template: bool | None = None + use_system_prompt: bool | None = None + use_guardrails: bool | None = None # Prompt enhancement (ErnieImage) use_pe: bool | None = None diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index 9f310a8da..d557516fb 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -961,6 +961,8 @@ def _register_configs(): hf_model_paths=[ "nvidia/Cosmos3-Nano", "nvidia/Cosmos3-Super", + "nvidia/Cosmos3-Super-Text2Image", + "nvidia/Cosmos3-Super-Image2Video", ], model_detectors=[lambda hf_id: "cosmos3omnidiffuserspipeline" in hf_id.lower()], ) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py index 848d4f151..8ceb2856b 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py @@ -2,9 +2,10 @@ import base64 import contextlib +import json import os import time -from typing import List, Optional +from typing import Any, List, Optional from fastapi import ( APIRouter, @@ -30,6 +31,7 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.utils import ( add_common_data_to_response, build_sampling_params, choose_output_image_ext, + flatten_extra_params, merge_image_input_list, process_generation_batch, save_image_to_path, @@ -48,11 +50,30 @@ def _get_extra_field(request, field_name): """Get a field from model_extra, with fallback to nested extra_body dict.""" extra = request.model_extra or {} value = extra.get(field_name) - if value is None and isinstance(extra.get("extra_body"), dict): - value = extra["extra_body"].get(field_name) + if value is not None: + return value + if field_name == "use_guardrails" and extra.get("guardrails") is not None: + return extra["guardrails"] + + for container_name in ("extra_body", "extra_json", "extra_args", "extra_params"): + value = _parse_extra_container(extra.get(container_name)).get(field_name) + if value is not None: + return value + return value +def _parse_extra_container(value: Any) -> dict[str, Any]: + if isinstance(value, str): + try: + value = json.loads(value) + except Exception: + return {} + if isinstance(value, dict): + return flatten_extra_params(dict(value)) + return {} + + def _read_b64_for_paths(paths: list[str]) -> list[str]: """Read and base64-encode each file. Must be called before cloud upload deletes them.""" result = [] @@ -130,7 +151,12 @@ async def generations( ): request_id = generate_request_id() server_args = get_global_server_args() - ext = choose_output_image_ext(request.output_format, request.background) + is_cosmos3 = "cosmos3" in (server_args.model_path or "").lower() + ext = ( + "png" + if is_cosmos3 and request.output_format is None + else choose_output_image_ext(request.output_format, request.background) + ) with temp_dir_if_disabled(server_args.output_path) as output_dir: sampling = build_sampling_params( @@ -142,12 +168,29 @@ async def generations( num_outputs_per_prompt=max(1, min(int(request.n or 1), 10)), output_file_name=f"{request_id}.{ext}", output_path=output_dir, + num_frames=1, seed=request.seed, generator_device=request.generator_device, num_inference_steps=request.num_inference_steps, guidance_scale=request.guidance_scale, true_cfg_scale=request.true_cfg_scale, negative_prompt=request.negative_prompt, + max_sequence_length=( + request.max_sequence_length + if request.max_sequence_length is not None + else _get_extra_field(request, "max_sequence_length") + ), + flow_shift=( + request.flow_shift + if request.flow_shift is not None + else _get_extra_field(request, "flow_shift") + ), + use_duration_template=_get_extra_field(request, "use_duration_template"), + use_resolution_template=_get_extra_field( + request, "use_resolution_template" + ), + use_system_prompt=_get_extra_field(request, "use_system_prompt"), + use_guardrails=_get_extra_field(request, "use_guardrails"), enable_teacache=request.enable_teacache, output_compression=request.output_compression, output_quality=request.output_quality, @@ -173,6 +216,12 @@ async def generations( ) save_file_path = save_file_path_list[0] resp_format = (request.response_format or "b64_json").lower() + if ( + is_cosmos3 + and "response_format" not in request.model_fields_set + and request.response_format == "url" + ): + resp_format = "b64_json" # read b64 before cloud upload may delete the local file b64_list = ( diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py index ed627777d..3d6540546 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py @@ -50,6 +50,8 @@ class ImageGenerationsRequest(BaseModel): output_quality: Optional[str] = "default" output_compression: Optional[int] = None enable_teacache: Optional[bool] = False + max_sequence_length: Optional[int] = None + flow_shift: Optional[float] = None # Upscaling enable_upscaling: Optional[bool] = False upscaling_model_path: Optional[str] = None @@ -83,6 +85,8 @@ class VideoResponse(BaseModel): class VideoGenerationsRequest(BaseModel): + model_config = ConfigDict(extra="allow") + prompt: str input_reference: Optional[str] = None reference_url: Optional[str] = None @@ -105,6 +109,8 @@ class VideoGenerationsRequest(BaseModel): None # for CFG vs guidance distillation (e.g., QwenImage) ) negative_prompt: Optional[str] = None + max_sequence_length: Optional[int] = None + flow_shift: Optional[float] = None enable_teacache: Optional[bool] = False # Frame interpolation enable_frame_interpolation: Optional[bool] = False diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py index bdc2d2fae..ccca4d3e4 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py @@ -1,6 +1,7 @@ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo import asyncio import base64 +import json import os import re import shutil @@ -52,6 +53,30 @@ DEFAULT_FPS = 24 DEFAULT_VIDEO_SECONDS = 4 +def flatten_extra_params(payload: Any) -> dict[str, Any]: + """Promote vLLM-Omni-style extra_params into regular request fields.""" + if not isinstance(payload, dict): + return {} + + extra_params = payload.pop("extra_params", None) + if isinstance(extra_params, str): + try: + extra_params = json.loads(extra_params) + except Exception: + extra_params = None + if not isinstance(extra_params, dict): + if "guardrails" in payload: + payload.setdefault("use_guardrails", payload["guardrails"]) + return payload + + for key, value in extra_params.items(): + payload.setdefault(key, value) + if "guardrails" in extra_params: + payload.setdefault("use_guardrails", extra_params["guardrails"]) + + return payload + + @contextmanager def temp_dir_if_disabled( configured_path: str | None, 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 ccf07918d..4ee93723b 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py @@ -36,6 +36,7 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.utils import ( DEFAULT_VIDEO_SECONDS, add_common_data_to_response, build_sampling_params, + flatten_extra_params, merge_image_input_list, process_generation_batch, save_image_to_path, @@ -50,6 +51,19 @@ logger = init_logger(__name__) router = APIRouter(prefix="/v1/videos", tags=["videos"]) +def _extra_value(request: VideoGenerationsRequest, name: str) -> Any: + return (request.model_extra or {}).get(name) + + +def _parse_form_extra_value(value: Any) -> Any: + if not isinstance(value, str): + return value + try: + return json.loads(value) + except Exception: + return value + + def _build_video_sampling_params(request_id: str, request: VideoGenerationsRequest): """Resolve video-specific defaults (fps, seconds → num_frames) then delegate to the shared build_sampling_params.""" @@ -77,6 +91,12 @@ def _build_video_sampling_params(request_id: str, request: VideoGenerationsReque guidance_scale=request.guidance_scale, guidance_scale_2=request.guidance_scale_2, negative_prompt=request.negative_prompt, + max_sequence_length=request.max_sequence_length, + flow_shift=request.flow_shift, + use_duration_template=_extra_value(request, "use_duration_template"), + use_resolution_template=_extra_value(request, "use_resolution_template"), + use_system_prompt=_extra_value(request, "use_system_prompt"), + use_guardrails=_extra_value(request, "use_guardrails"), enable_teacache=request.enable_teacache, enable_frame_interpolation=request.enable_frame_interpolation, frame_interpolation_exp=request.frame_interpolation_exp, @@ -89,9 +109,34 @@ def _build_video_sampling_params(request_id: str, request: VideoGenerationsReque output_compression=request.output_compression, output_quality=request.output_quality, perf_dump_path=request.perf_dump_path, + diffusers_kwargs=request.diffusers_kwargs, ) +def _reject_unsupported_cosmos3_modes( + req: VideoGenerationsRequest, model_path: str | None +) -> None: + if "cosmos3" not in (model_path or "").lower(): + return + + extra = req.model_extra or {} + if extra.get("generate_sound"): + raise HTTPException( + status_code=400, + detail="Cosmos3 video-with-sound is not supported by SGLang yet; omit generate_sound for video-only generation.", + ) + if extra.get("action_mode"): + raise HTTPException( + status_code=400, + detail="Cosmos3 action generation is not supported by SGLang yet.", + ) + if extra.get("condition_frame_indexes_vision") or extra.get("condition_video_keep"): + raise HTTPException( + status_code=400, + detail="Cosmos3 video-to-video conditioning is not supported by SGLang yet.", + ) + + # extract metadata which http_server needs to know def _video_job_from_sampling( request_id: str, req: VideoGenerationsRequest, sampling: SamplingParams @@ -201,16 +246,20 @@ async def create_video( negative_prompt: Optional[str] = Form(None), guidance_scale: Optional[float] = Form(None), num_inference_steps: Optional[int] = Form(None), - enable_teacache: Optional[bool] = Form(False), - enable_frame_interpolation: Optional[bool] = Form(False), - frame_interpolation_exp: Optional[int] = Form(1), - frame_interpolation_scale: Optional[float] = Form(1.0), + max_sequence_length: Optional[int] = Form(None), + flow_shift: Optional[float] = Form(None), + enable_teacache: Optional[bool] = Form(None), + enable_frame_interpolation: Optional[bool] = Form(None), + frame_interpolation_exp: Optional[int] = Form(None), + frame_interpolation_scale: Optional[float] = Form(None), frame_interpolation_model_path: Optional[str] = Form(None), - enable_upscaling: Optional[bool] = Form(False), + enable_upscaling: Optional[bool] = Form(None), upscaling_model_path: Optional[str] = Form(None), - upscaling_scale: Optional[int] = Form(4), - output_quality: Optional[str] = Form("default"), + upscaling_scale: Optional[int] = Form(None), + output_quality: Optional[str] = Form(None), output_compression: Optional[int] = Form(None), + output_path: Optional[str] = Form(None), + extra_params: Optional[str] = Form(None), extra_body: Optional[str] = Form(None), ): content_type = request.headers.get("content-type", "").lower() @@ -261,42 +310,88 @@ async def create_video( extra_from_form: Dict[str, Any] = {} if extra_body: try: - extra_from_form = json.loads(extra_body) + extra_from_form = flatten_extra_params(json.loads(extra_body)) except Exception: extra_from_form = {} + if extra_params: + try: + extra_from_form.update( + flatten_extra_params({"extra_params": json.loads(extra_params)}) + ) + except Exception: + pass - fps_val = fps if fps is not None else extra_from_form.get("fps") - num_frames_val = ( - num_frames if num_frames is not None else extra_from_form.get("num_frames") - ) + def form_value(name: str, value: Any) -> Any: + return value if value is not None else extra_from_form.get(name) + + raw_form = await request.form() + for key in ( + "use_duration_template", + "use_resolution_template", + "use_system_prompt", + "use_guardrails", + "guardrails", + "generate_sound", + "sound_duration", + "action_mode", + "condition_frame_indexes_vision", + "condition_video_keep", + ): + if key in raw_form and key not in extra_from_form: + extra_from_form[key] = _parse_form_extra_value(raw_form[key]) + flatten_extra_params(extra_from_form) + + request_field_names = set(VideoGenerationsRequest.model_fields) + extra_request_fields = { + key: value + for key, value in extra_from_form.items() + if key not in request_field_names + } + fps_val = form_value("fps", fps) + num_frames_val = form_value("num_frames", num_frames) req = VideoGenerationsRequest( prompt=prompt, input_reference=input_path, - model=model, - n=n, - num_outputs_per_prompt=num_outputs_per_prompt, - seconds=seconds if seconds is not None else 4, - size=size, + model=form_value("model", model), + n=form_value("n", n), + num_outputs_per_prompt=form_value( + "num_outputs_per_prompt", num_outputs_per_prompt + ), + seconds=form_value("seconds", seconds) or 4, + size=form_value("size", size), fps=fps_val, num_frames=num_frames_val, - seed=seed, - generator_device=generator_device, - negative_prompt=negative_prompt, - num_inference_steps=num_inference_steps, - enable_teacache=enable_teacache, - enable_frame_interpolation=enable_frame_interpolation, - frame_interpolation_exp=frame_interpolation_exp, - frame_interpolation_scale=frame_interpolation_scale, - frame_interpolation_model_path=frame_interpolation_model_path, - enable_upscaling=enable_upscaling, - upscaling_model_path=upscaling_model_path, - upscaling_scale=upscaling_scale, - output_compression=output_compression, - output_quality=output_quality, - **( - {"guidance_scale": guidance_scale} if guidance_scale is not None else {} + seed=form_value("seed", seed), + generator_device=form_value("generator_device", generator_device), + negative_prompt=form_value("negative_prompt", negative_prompt), + num_inference_steps=form_value("num_inference_steps", num_inference_steps), + guidance_scale=form_value("guidance_scale", guidance_scale), + max_sequence_length=form_value("max_sequence_length", max_sequence_length), + flow_shift=form_value("flow_shift", flow_shift), + enable_teacache=form_value("enable_teacache", enable_teacache), + enable_frame_interpolation=form_value( + "enable_frame_interpolation", enable_frame_interpolation ), + frame_interpolation_exp=form_value( + "frame_interpolation_exp", frame_interpolation_exp + ), + frame_interpolation_scale=form_value( + "frame_interpolation_scale", frame_interpolation_scale + ), + frame_interpolation_model_path=form_value( + "frame_interpolation_model_path", frame_interpolation_model_path + ), + enable_upscaling=form_value("enable_upscaling", enable_upscaling), + upscaling_model_path=form_value( + "upscaling_model_path", upscaling_model_path + ), + upscaling_scale=form_value("upscaling_scale", upscaling_scale), + output_compression=form_value("output_compression", output_compression), + output_quality=form_value("output_quality", output_quality), + output_path=form_value("output_path", output_path), + diffusers_kwargs=form_value("diffusers_kwargs", None), + **extra_request_fields, ) else: try: @@ -307,13 +402,17 @@ async def create_video( # If client uses extra_body, merge it into the top-level payload payload: Dict[str, Any] = dict(body or {}) extra = payload.pop("extra_body", None) + if isinstance(extra, str): + extra = json.loads(extra) if isinstance(extra, dict): - # Shallow-merge: only keys like fps/num_frames are expected - payload.update(extra) + payload.update(flatten_extra_params(extra)) # openai may turn extra_body to extra_json extra_json = payload.pop("extra_json", None) + if isinstance(extra_json, str): + extra_json = json.loads(extra_json) if isinstance(extra_json, dict): - payload.update(extra_json) + payload.update(flatten_extra_params(extra_json)) + flatten_extra_params(payload) # Validate image input based on model task type has_image_input = payload.get("reference_url") or payload.get( "input_reference" @@ -355,6 +454,8 @@ async def create_video( logger.debug(f"Server received from create_video endpoint: req={req}") + _reject_unsupported_cosmos3_modes(req, server_args.model_path) + try: sampling_params = _build_video_sampling_params(request_id, req) except (ValueError, TypeError) as e: @@ -373,6 +474,10 @@ async def create_video( # Add diffusers_kwargs if provided if req.diffusers_kwargs: batch.extra["diffusers_kwargs"] = req.diffusers_kwargs + if "max_sequence_length" in req.diffusers_kwargs: + batch.max_sequence_length = req.diffusers_kwargs["max_sequence_length"] + if "flow_shift" in req.diffusers_kwargs: + batch.flow_shift = req.diffusers_kwargs["flow_shift"] # Enqueue the job asynchronously and return immediately asyncio.create_task( _dispatch_job_async( diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py index 96fa94386..509030551 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py @@ -420,6 +420,9 @@ def prepare_request( VSA_sparsity=server_args.attention_backend_config.VSA_sparsity, ) sampling_params.apply_request_extra(req) + if getattr(sampling_params, "max_sequence_length", None) is not None: + req.max_sequence_length = sampling_params.max_sequence_length + diffusers_kwargs = getattr(sampling_params, "diffusers_kwargs", None) if diffusers_kwargs and "max_sequence_length" in diffusers_kwargs: req.max_sequence_length = diffusers_kwargs["max_sequence_length"] diff --git a/python/sglang/multimodal_gen/runtime/pipelines/cosmos3_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/cosmos3_pipeline.py index db667dd02..2b47d2276 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/cosmos3_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/cosmos3_pipeline.py @@ -59,8 +59,21 @@ class Cosmos3Pipeline(ComposedPipelineBase): transformer = self.get_module("transformer") scheduler = self.get_module("scheduler") - # Guardrails on by default; opt out with SGLANG_DISABLE_COSMOS3_GUARDRAILS=1. - guardrails_on = os.environ.get("SGLANG_DISABLE_COSMOS3_GUARDRAILS", "0") != "1" + guardrails_disabled = ( + os.environ.get("SGLANG_DISABLE_COSMOS3_GUARDRAILS", "0") == "1" + ) + guardrails_on = False + if not guardrails_disabled: + from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3_guardrails import ( + is_cosmos_guardrail_available, + ) + + guardrails_on = is_cosmos_guardrail_available() + if not guardrails_on: + logger.warning( + "Cosmos3 guardrails disabled because cosmos-guardrail is not " + "installed. Install it with: pip install cosmos-guardrail==0.3.1" + ) self.add_stage(Cosmos3ImagePreprocessStage()) self.add_stage(Cosmos3TokenizationStage(tokenizer=text_tokenizer)) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py index 2e0cc0a70..0491dc24f 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py @@ -353,6 +353,12 @@ class Req: self.negative_prompt, key_hint="negative_prompt" ) + effective_flow_shift = ( + self.flow_shift + if self.flow_shift is not None + else getattr(server_args.pipeline_config, "flow_shift", None) + ) + debug_str = f"""Sampling params: width: {target_width} height: {target_height} @@ -366,7 +372,7 @@ class Req: guidance_scale: {self.guidance_scale} embedded_guidance_scale: {server_args.pipeline_config.embedded_cfg_scale} n_tokens: {self.n_tokens} - flow_shift: {server_args.pipeline_config.flow_shift} + flow_shift: {effective_flow_shift} image_path: {self.image_path} save_output: {self.save_output} output_file_path: {self.output_file_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 fed13f075..8a53b38f7 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 @@ -190,8 +190,16 @@ class Cosmos3TokenizationStage(PipelineStage): # Get parameters max_sequence_length = getattr(batch, "max_sequence_length", None) or 512 - use_duration_template = getattr(batch, "use_duration_template", True) - use_system_prompt = getattr(batch, "use_system_prompt", False) + use_duration_template = getattr(batch, "use_duration_template", None) + if use_duration_template is None: + use_duration_template = getattr( + server_args.pipeline_config, "use_duration_template", True + ) + use_system_prompt = getattr(batch, "use_system_prompt", None) + if use_system_prompt is None: + use_system_prompt = getattr( + server_args.pipeline_config, "use_system_prompt", False + ) fps = batch.fps or 24.0 num_frames = batch.num_frames is_image_gen = batch.data_type == DataType.IMAGE @@ -349,16 +357,28 @@ class Cosmos3TimestepPreparationStage(PipelineStage): def __init__(self, scheduler): super().__init__() self.scheduler = scheduler + self.default_flow_shift = getattr( + getattr(scheduler, "config", None), "flow_shift", None + ) def forward(self, batch: Req, server_args: ServerArgs) -> Req: """Prepare scheduler timesteps.""" device = get_local_torch_device() num_inference_steps = batch.num_inference_steps + flow_shift = getattr(batch, "flow_shift", None) + if flow_shift is None: + flow_shift = server_args.pipeline_config.flow_shift + if flow_shift is None: + flow_shift = self.default_flow_shift + if flow_shift is not None and hasattr(self.scheduler, "set_shift"): + self.scheduler.set_shift(float(flow_shift)) self.scheduler.set_timesteps(num_inference_steps, device=device) batch.timesteps = self.scheduler.timesteps - self.log_info(f"Prepared {len(batch.timesteps)} timesteps") + self.log_info( + f"Prepared {len(batch.timesteps)} timesteps (flow_shift={flow_shift})" + ) return batch @@ -865,7 +885,7 @@ class Cosmos3DecodingStage(PipelineStage): output = self.video_processor.postprocess_video(decoded, output_type="np") self.log_info(f"Postprocessed video shape: {output.shape}") - if self._guardrails: + if self._guardrails and batch.use_guardrails is not False: from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3_guardrails import ( check_video_safety, ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3_guardrails.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3_guardrails.py index 808aa1a1c..c3e05e53a 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3_guardrails.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3_guardrails.py @@ -4,11 +4,15 @@ Text and video safety checks via the ``cosmos_guardrail`` package. Install with: pip install cosmos-guardrail==0.3.1 -Enabled by default; opt out with ``SGLANG_DISABLE_COSMOS3_GUARDRAILS=1``. +Enabled by default when available; opt out with +``SGLANG_DISABLE_COSMOS3_GUARDRAILS=1``. """ from __future__ import annotations +import importlib.util +from functools import lru_cache + import numpy as np import torch @@ -25,6 +29,11 @@ logger = init_logger(__name__) _checker = None +@lru_cache(maxsize=1) +def is_cosmos_guardrail_available() -> bool: + return importlib.util.find_spec("cosmos_guardrail") is not None + + def _init_guardrails(offload_to_cpu: bool = False) -> None: global _checker if _checker is not None: @@ -91,6 +100,8 @@ class Cosmos3TextGuardrailStage(PipelineStage): _init_guardrails(offload_to_cpu) def forward(self, batch: Req, server_args: ServerArgs) -> Req: + if batch.use_guardrails is False: + return batch prompt = batch.prompt if prompt is None: return batch diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index 331c9e9d8..d758acaf3 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args.py @@ -1144,6 +1144,7 @@ class ServerArgs(DisaggArgsMixin): "--data-parallel-size", "--dp-size", "--dp", + dest="dp_size", type=int, default=ServerArgs.dp_size, help="The data parallelism size.", @@ -1794,10 +1795,16 @@ class ServerArgs(DisaggArgsMixin): # For '--arg=value', this gets 'arg'; for '--arg', this also gets 'arg'. arg_name = arg.split("=", 1)[0].replace("-", "_").lstrip("_") provided_arg_names.add(arg_name) - if "mode" in provided_arg_names: - provided_arg_names.add("performance_mode") - if "layerwise_offload_modules" in provided_arg_names: - provided_arg_names.add("layerwise_offload_components") + cli_aliases = { + "cfg_parallel_size": "cfg_parallel_degree", + "data_parallel_size": "dp_size", + "dp": "dp_size", + "layerwise_offload_modules": "layerwise_offload_components", + "mode": "performance_mode", + } + for alias_name, dest_name in cli_aliases.items(): + if alias_name in provided_arg_names: + provided_arg_names.add(dest_name) # Populate provided_args if the argument from the namespace was on the command line. for k, v in vars(args).items(): diff --git a/python/sglang/multimodal_gen/test/server/gpu_cases.py b/python/sglang/multimodal_gen/test/server/gpu_cases.py index 885582d95..f141f4efd 100644 --- a/python/sglang/multimodal_gen/test/server/gpu_cases.py +++ b/python/sglang/multimodal_gen/test/server/gpu_cases.py @@ -30,6 +30,7 @@ from sglang.multimodal_gen.test.server.testcase_configs import ( _with_default_num_gpus, ) from sglang.multimodal_gen.test.test_utils import ( + DEFAULT_COSMOS3_NANO_MODEL_NAME_FOR_TEST, DEFAULT_FLUX_1_DEV_MODEL_NAME_FOR_TEST, DEFAULT_FLUX_2_DEV_MODEL_NAME_FOR_TEST, DEFAULT_FLUX_2_KLEIN_4B_MODEL_NAME_FOR_TEST, @@ -161,6 +162,31 @@ ONE_GPU_CASES: list[DiffusionTestCase] = [ run_lora_dynamic_switch_check=True, run_multi_lora_api_check=True, ), + DiffusionTestCase( + "cosmos3_nano_t2i", + DiffusionServerArgs( + model_path=DEFAULT_COSMOS3_NANO_MODEL_NAME_FOR_TEST, + modality="image", + ), + DiffusionSamplingParams( + prompt="A red cube on a white table, product photo.", + output_size="832x480", + output_format="png", + extras={ + "num_inference_steps": 35, + "seed": 0, + "max_sequence_length": 128, + "flow_shift": 10.0, + "extra_args": { + "guardrails": False, + "use_resolution_template": False, + }, + }, + ), + run_perf_check=False, + run_consistency_check=True, + run_component_accuracy_check=False, + ), # === Text and Image to Image (TI2I) === DiffusionTestCase( "qwen_image_edit_ti2i", diff --git a/python/sglang/multimodal_gen/test/server/test_server_common.py b/python/sglang/multimodal_gen/test/server/test_server_common.py index 246f2841f..3390dc5c3 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_common.py +++ b/python/sglang/multimodal_gen/test/server/test_server_common.py @@ -35,6 +35,7 @@ from sglang.multimodal_gen.test.server.testcase_configs import ( DiffusionTestCase, PerformanceSummary, ScenarioConfig, + get_model_task_type_for_server_args, ) from sglang.multimodal_gen.test.test_utils import ( SGL_TEST_FILES_CI_DATA_REVISION, @@ -1085,19 +1086,10 @@ Pinned revision used by this check: {SGL_TEST_FILES_CI_DATA_REVISION} assert ( model["num_gpus"] == case.server_args.num_gpus ), f"num_gpus mismatch: expected {case.server_args.num_gpus}, got {model['num_gpus']}" - # Verify task_type is consistent with the modality specified in the test config. - # We can't access pipeline_config from test config, but we can validate against modality. - modality_to_valid_task_types = { - "image": {"T2I", "I2I", "TI2I"}, - "video": {"T2V", "I2V", "TI2V"}, - "3d": {"I2M"}, - } - valid_task_types = modality_to_valid_task_types.get( - case.server_args.modality, set() - ) - assert model["task_type"] in valid_task_types, ( - f"task_type '{model['task_type']}' not valid for modality " - f"'{case.server_args.modality}'. Expected one of: {valid_task_types}" + expected_task_type = get_model_task_type_for_server_args(case.server_args).name + assert model["task_type"] == expected_task_type, ( + f"task_type mismatch: expected {expected_task_type}, " + f"got {model['task_type']}" ) logger.info( "[Models API] GET /v1/models returned valid response with extended fields" @@ -1118,9 +1110,9 @@ Pinned revision used by this check: {SGL_TEST_FILES_CI_DATA_REVISION} # Verify extended fields on single model endpoint too assert "num_gpus" in single_model, "Single model missing 'num_gpus' field" assert "task_type" in single_model, "Single model missing 'task_type' field" - assert single_model["task_type"] in valid_task_types, ( - f"Single model task_type '{single_model['task_type']}' not valid for modality " - f"'{case.server_args.modality}'. Expected one of: {valid_task_types}" + assert single_model["task_type"] == expected_task_type, ( + f"Single model task_type mismatch: expected {expected_task_type}, " + f"got {single_model['task_type']}" ) logger.info( "[Models API] GET /v1/models/{model_path} returned valid response with extended fields" diff --git a/python/sglang/multimodal_gen/test/server/test_server_utils.py b/python/sglang/multimodal_gen/test/server/test_server_utils.py index e207e21d5..59be5da5a 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_utils.py +++ b/python/sglang/multimodal_gen/test/server/test_server_utils.py @@ -986,7 +986,7 @@ def get_generate_fn( pytest.skip(f"{case_id}: no text prompt configured") # Request parameters that affect output format - req_output_format = None # Not specified in current request + req_output_format = sampling_params.output_format req_background = None # Not specified in current request # Build extra_body for optional features @@ -998,6 +998,7 @@ def get_generate_fn( n=n, size=output_size, response_format="b64_json", + output_format=req_output_format, extra_body=extra_body if extra_body else None, ) result = response.parse() diff --git a/python/sglang/multimodal_gen/test/test_utils.py b/python/sglang/multimodal_gen/test/test_utils.py index ff1e467aa..679c2d7bd 100644 --- a/python/sglang/multimodal_gen/test/test_utils.py +++ b/python/sglang/multimodal_gen/test/test_utils.py @@ -33,7 +33,7 @@ if TYPE_CHECKING: logger = init_logger(__name__) -SGL_TEST_FILES_CI_DATA_REVISION = "a17a6cd676d16d0f6c93cc80d0144138ab87dca1" +SGL_TEST_FILES_CI_DATA_REVISION = "10f3826199ae524b3af5026a57c8f817d207b2e5" SGL_TEST_FILES_CONSISTENCY_GT_ROOT = ( "https://raw.githubusercontent.com/" f"sgl-project/ci-data/{SGL_TEST_FILES_CI_DATA_REVISION}/" @@ -108,6 +108,9 @@ def _load_clip_processor_with_roberta_processing_compat( DEFAULT_SMALL_MODEL_NAME_FOR_TEST = "Tongyi-MAI/Z-Image-Turbo" +# Cosmos3 generation models +DEFAULT_COSMOS3_NANO_MODEL_NAME_FOR_TEST = "nvidia/Cosmos3-Nano" + # Qwen image generation models DEFAULT_QWEN_IMAGE_MODEL_NAME_FOR_TEST = "Qwen/Qwen-Image" DEFAULT_QWEN_IMAGE_2512_MODEL_NAME_FOR_TEST = "Qwen/Qwen-Image-2512" diff --git a/python/sglang/multimodal_gen/test/unit/test_cosmos3.py b/python/sglang/multimodal_gen/test/unit/test_cosmos3.py index 31d3b4296..3d120e1f8 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cosmos3.py +++ b/python/sglang/multimodal_gen/test/unit/test_cosmos3.py @@ -1,7 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 """Unit tests for Cosmos3 config, weight mapping, and sampling params.""" +import importlib.util import unittest +from unittest import mock + +from fastapi import HTTPException from sglang.multimodal_gen.configs.models.dits.cosmos3video import ( _build_cosmos3_param_names_mapping, @@ -9,7 +13,21 @@ from sglang.multimodal_gen.configs.models.dits.cosmos3video import ( from sglang.multimodal_gen.configs.pipeline_configs.cosmos3 import Cosmos3Config from sglang.multimodal_gen.configs.sample.cosmos3 import Cosmos3SamplingParams from sglang.multimodal_gen.configs.sample.sampling_params import DataType +from sglang.multimodal_gen.registry import ( + _get_config_info, + get_non_diffusers_pipeline_name, +) +from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import ( + ImageGenerationsRequest, + VideoGenerationsRequest, +) +from sglang.multimodal_gen.runtime.entrypoints.openai.video_api import ( + _reject_unsupported_cosmos3_modes, +) from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3_guardrails import ( + is_cosmos_guardrail_available, +) def _apply(mapping_fn, key): @@ -209,5 +227,74 @@ class TestCosmos3SamplingParamsDataType(unittest.TestCase): self.assertEqual(params.data_type, DataType.VIDEO) +class TestCosmos3ModelResolution(unittest.TestCase): + """Verify Cosmos3 checkpoints resolve to the native SGLang pipeline.""" + + def test_hf_checkpoint_uses_registered_native_pipeline_config(self): + for model_path in ( + "nvidia/Cosmos3-Nano", + "nvidia/Cosmos3-Super", + "nvidia/Cosmos3-Super-Text2Image", + "nvidia/Cosmos3-Super-Image2Video", + ): + with self.subTest(model_path=model_path): + self.assertIsNone(get_non_diffusers_pipeline_name(model_path)) + config_info = _get_config_info(model_path) + self.assertIsNotNone(config_info) + self.assertIs(config_info.sampling_param_cls, Cosmos3SamplingParams) + self.assertIs(config_info.pipeline_config_cls, Cosmos3Config) + + +class TestCosmos3OpenAIProtocol(unittest.TestCase): + """Verify Cosmos3-only knobs stay out of the stable request schema.""" + + def test_cosmos3_private_fields_are_extra_fields(self): + for request_cls in (ImageGenerationsRequest, VideoGenerationsRequest): + with self.subTest(request_cls=request_cls.__name__): + self.assertIn("max_sequence_length", request_cls.model_fields) + self.assertIn("flow_shift", request_cls.model_fields) + self.assertNotIn("use_duration_template", request_cls.model_fields) + self.assertNotIn("use_resolution_template", request_cls.model_fields) + self.assertNotIn("use_system_prompt", request_cls.model_fields) + self.assertNotIn("use_guardrails", request_cls.model_fields) + + self.assertNotIn("generate_sound", VideoGenerationsRequest.model_fields) + self.assertNotIn("sound_duration", VideoGenerationsRequest.model_fields) + + def test_unsupported_cosmos3_modes_allow_falsy_extra_fields(self): + req = VideoGenerationsRequest( + prompt="test", + generate_sound=False, + action_mode="", + condition_frame_indexes_vision=[], + condition_video_keep={}, + ) + _reject_unsupported_cosmos3_modes(req, "nvidia/Cosmos3-Nano") + + req = VideoGenerationsRequest(prompt="test", generate_sound=True) + with self.assertRaises(HTTPException): + _reject_unsupported_cosmos3_modes(req, "nvidia/Cosmos3-Nano") + + +class TestCosmos3Guardrails(unittest.TestCase): + """Verify optional guardrail dependency handling.""" + + def setUp(self): + is_cosmos_guardrail_available.cache_clear() + + def tearDown(self): + is_cosmos_guardrail_available.cache_clear() + + def test_guardrail_availability_matches_package_spec(self): + self.assertEqual( + is_cosmos_guardrail_available(), + importlib.util.find_spec("cosmos_guardrail") is not None, + ) + + @mock.patch("importlib.util.find_spec", return_value=None) + def test_missing_guardrail_package_reports_unavailable(self, _): + self.assertFalse(is_cosmos_guardrail_available()) + + if __name__ == "__main__": unittest.main()