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

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

Browse Source
This commit is contained in:
Mick
2026-06-02 00:53:39 +08:00
committed by GitHub
parent ed24e3aae8
commit 1988a2c9ea
19 changed files with 637 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/configs/sample/sampling_params.py
+8
View File
@@ -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
+2
View File
@@ -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
+53 -4
View File
@@ -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 = (
python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py
+6
View File
@@ -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
+25
View File
@@ -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
+141 -36
View File
@@ -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(
python/sglang/multimodal_gen/runtime/entrypoints/utils.py
+3
View File
@@ -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
+15 -2
View File
@@ -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))
python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py
+7 -1
View File
@@ -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
+24 -4
View File
@@ -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,
)
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3_guardrails.py
+12 -1
View File
@@ -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
+11 -4
View File
@@ -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():
python/sglang/multimodal_gen/test/server/gpu_cases.py
+26
View File
@@ -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
+8 -16
View File
@@ -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"
python/sglang/multimodal_gen/test/server/test_server_utils.py
+2 -1
View File
@@ -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
+4 -1
View File
@@ -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
+87
View File
@@ -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()