[diffusion] feat: improve cosmos3 serve API support (#26926)

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.

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": [

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

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()],
     )

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):
+    if value is not None:
-        value = extra["extra_body"].get(field_name)
+        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 = (

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

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,

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),
+    max_sequence_length: Optional[int] = Form(None),
-    enable_frame_interpolation: Optional[bool] = Form(False),
+    flow_shift: Optional[float] = Form(None),
-    frame_interpolation_exp: Optional[int] = Form(1),
+    enable_teacache: Optional[bool] = Form(None),
-    frame_interpolation_scale: Optional[float] = Form(1.0),
+    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),
+    upscaling_scale: Optional[int] = Form(None),
-    output_quality: Optional[str] = Form("default"),
+    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")
+        def form_value(name: str, value: Any) -> Any:
-        num_frames_val = (
+            return value if value is not None else extra_from_form.get(name)
-            num_frames if num_frames is not None else extra_from_form.get("num_frames")
+
-        )
+        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,
+            model=form_value("model", model),
-            n=n,
+            n=form_value("n", n),
-            num_outputs_per_prompt=num_outputs_per_prompt,
+            num_outputs_per_prompt=form_value(
-            seconds=seconds if seconds is not None else 4,
+                "num_outputs_per_prompt", num_outputs_per_prompt
-            size=size,
+            ),
+            seconds=form_value("seconds", seconds) or 4,
+            size=form_value("size", size),
             fps=fps_val,
             num_frames=num_frames_val,
-            seed=seed,
+            seed=form_value("seed", seed),
-            generator_device=generator_device,
+            generator_device=form_value("generator_device", generator_device),
-            negative_prompt=negative_prompt,
+            negative_prompt=form_value("negative_prompt", negative_prompt),
-            num_inference_steps=num_inference_steps,
+            num_inference_steps=form_value("num_inference_steps", num_inference_steps),
-            enable_teacache=enable_teacache,
+            guidance_scale=form_value("guidance_scale", guidance_scale),
-            enable_frame_interpolation=enable_frame_interpolation,
+            max_sequence_length=form_value("max_sequence_length", max_sequence_length),
-            frame_interpolation_exp=frame_interpolation_exp,
+            flow_shift=form_value("flow_shift", flow_shift),
-            frame_interpolation_scale=frame_interpolation_scale,
+            enable_teacache=form_value("enable_teacache", enable_teacache),
-            frame_interpolation_model_path=frame_interpolation_model_path,
+            enable_frame_interpolation=form_value(
-            enable_upscaling=enable_upscaling,
+                "enable_frame_interpolation", enable_frame_interpolation
-            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 {}
             ),
+            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(flatten_extra_params(extra))
-                payload.update(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(

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"]

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_disabled = (
-        guardrails_on = os.environ.get("SGLANG_DISABLE_COSMOS3_GUARDRAILS", "0") != "1"
+            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))

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()}

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_duration_template = getattr(batch, "use_duration_template", None)
-        use_system_prompt = getattr(batch, "use_system_prompt", False)
+        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,
             )

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

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:
+        cli_aliases = {
-            provided_arg_names.add("performance_mode")
+            "cfg_parallel_size": "cfg_parallel_degree",
-        if "layerwise_offload_modules" in provided_arg_names:
+            "data_parallel_size": "dp_size",
-            provided_arg_names.add("layerwise_offload_components")
+            "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():

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",

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.
+        expected_task_type = get_model_task_type_for_server_args(case.server_args).name
-        # We can't access pipeline_config from test config, but we can validate against modality.
+        assert model["task_type"] == expected_task_type, (
-        modality_to_valid_task_types = {
+            f"task_type mismatch: expected {expected_task_type}, "
-            "image": {"T2I", "I2I", "TI2I"},
+            f"got {model['task_type']}"
-            "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}"
         )
         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, (
+        assert single_model["task_type"] == expected_task_type, (
-            f"Single model task_type '{single_model['task_type']}' not valid for modality "
+            f"Single model task_type mismatch: expected {expected_task_type}, "
-            f"'{case.server_args.modality}'. Expected one of: {valid_task_types}"
+            f"got {single_model['task_type']}"
         )
         logger.info(
             "[Models API] GET /v1/models/{model_path} returned valid response with extended fields"

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()

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"

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()