[diffusion] feat: support --served-model-name in sglang serve (#34228)

Co-authored-by: TobyMint <tobymint@users.noreply.github.com>
Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
TobyMint
2026-08-10 22:27:09 +08:00
committed by GitHub
co-authored by TobyMint Mick
parent f5f0c3ee7a
commit d07ac32d05
13 changed files with 198 additions and 25 deletions
@@ -151,7 +151,7 @@ def action_metadata(server_args: ServerArgs) -> dict[str, Any]:
defaults = Cosmos3SamplingParams()
return {
"object": "action.metadata",
"model": server_args.model_id or server_args.model_path,
"model": server_args.served_model_name,
"model_path": server_args.model_path,
"policy_family": "cosmos3",
"input": {
@@ -201,7 +201,7 @@ def action_metadata(server_args: ServerArgs) -> dict[str, Any]:
)
return {
"object": "action.metadata",
"model": server_args.model_id or server_args.model_path,
"model": server_args.served_model_name,
"model_path": server_args.model_path,
"policy_family": policy_family,
"input": {
@@ -622,7 +622,7 @@ def action_generation_response(
"id": output.get("request_id") or f"act_{uuid.uuid4().hex}",
"object": "action.generation",
"created": int(time.time()),
"model": server_args.model_id or server_args.model_path,
"model": server_args.served_model_name,
"data": [
{
"index": 0,
@@ -205,7 +205,7 @@ async def server_info_endpoint(request: Request):
return {
"model_path": server_args.model_path,
"served_model_name": server_args.model_id or server_args.model_path,
"served_model_name": server_args.served_model_name,
"tp_size": server_args.tp_size,
"dp_size": server_args.dp_size,
"version": __version__,
@@ -45,15 +45,17 @@ class DiffusionModelCard(ModelCard):
pipeline_class: Optional[str] = None
def _build_model_card(server_args: ServerArgs, model_id: str) -> DiffusionModelCard:
def _build_model_card(
server_args: ServerArgs, served_model_name: str
) -> DiffusionModelCard:
model_info = get_model_info(
server_args.model_path,
backend=server_args.backend,
model_id=server_args.model_id,
)
card_kwargs: dict[str, Any] = {
"id": model_id,
"root": model_id,
"id": served_model_name,
"root": served_model_name,
"num_gpus": server_args.num_gpus,
"task_type": server_args.pipeline_config.task_type.name,
"dit_precision": server_args.pipeline_config.dit_precision,
@@ -203,7 +205,7 @@ async def available_models():
if not server_args:
raise HTTPException(status_code=500, detail="Server args not initialized")
model_card = _build_model_card(server_args, server_args.model_path)
model_card = _build_model_card(server_args, server_args.served_model_name)
# Return dict directly to preserve extended fields (ModelList strips them)
return {"object": "list", "data": [model_card.model_dump()]}
@@ -216,7 +218,7 @@ async def retrieve_model(model: str):
if not server_args:
raise HTTPException(status_code=500, detail="Server args not initialized")
if model != server_args.model_path:
if model != server_args.served_model_name:
return orjson_response(
{
"error": {
@@ -230,4 +232,4 @@ async def retrieve_model(model: str):
)
# Return dict to preserve extended fields
return _build_model_card(server_args, model).model_dump()
return _build_model_card(server_args, server_args.served_model_name).model_dump()
@@ -425,14 +425,17 @@ def _build_video_sampling_params(request_id: str, request: VideoGenerationsReque
# extract metadata which http_server needs to know
def _video_job_from_sampling(
request_id: str, req: VideoGenerationsRequest, sampling: SamplingParams
request_id: str,
req: VideoGenerationsRequest,
sampling: SamplingParams,
served_model_name: str,
) -> Dict[str, Any]:
size_str = f"{sampling.width}x{sampling.height}"
seconds = int(round((sampling.num_frames or 0) / float(sampling.fps or 24)))
return {
"id": request_id,
"object": "video",
"model": req.model or "sora-2",
"model": req.model or served_model_name,
"status": "queued",
"progress": 0,
"created_at": int(time.time()),
@@ -841,7 +844,12 @@ async def create_video(
scheduler_batches = sampling_params.expand_video_request_outputs_for_queue(
batch
)
job = _video_job_from_sampling(request_id, req, sampling_params)
job = _video_job_from_sampling(
request_id,
req,
sampling_params,
server_args.served_model_name,
)
job.update(sampling_params.project_video_queued_job_fields(batch))
await VIDEO_STORE.upsert(request_id, job)
except Exception as e:
@@ -309,7 +309,7 @@ def launch_server(server_args: ServerArgs, launch_http_server: bool = True):
logger.info(
"Action generation endpoint ready: model=%s; per-request details are "
"debug-only (use --log-level debug).",
server_args.model_id or server_args.model_path,
server_args.served_model_name,
)
logger.info("Starting FastAPI server.")
if server_args.webui:
@@ -199,6 +199,9 @@ class ServerArgs(DisaggServerArgsMixin):
# explicit model ID override (e.g. "Qwen-Image")
model_id: str | None = None
# served model name exposed via /v1/models and generation responses
served_model_name: str | None = None
# Model backend (sglang native or diffusers)
backend: Backend = Backend.AUTO
@@ -492,6 +495,8 @@ class ServerArgs(DisaggServerArgsMixin):
auto_tuner.maybe_adjust_auto_fsdp_with_offload_enabled()
auto_tuner.maybe_replace_cpu_offloaded_components_with_layerwise()
self._adjust_path()
if self.served_model_name is None:
self.served_model_name = self.model_id or self.model_path
self._adjust_quant_config()
self._adjust_breakable_cuda_graph_support()
self._adjust_warmup()
@@ -1428,6 +1433,15 @@ class ServerArgs(DisaggServerArgsMixin):
"(e.g. 'Qwen-Image' for 'Qwen/Qwen-Image')."
),
)
parser.add_argument(
"--served-model-name",
type=str,
default=ServerArgs.served_model_name,
help=(
"Override the model name exposed by /v1/models and used in generation "
"responses. Defaults to --model-id if set, otherwise --model-path."
),
)
parser.add_argument(
"--pipeline",
"--pipeline-class-name",
@@ -66,6 +66,7 @@ def _cosmos3_server_args(config=None):
return types.SimpleNamespace(
model_id=None,
model_path="nvidia/Cosmos3-Nano",
served_model_name="cosmos3-production",
backend=None,
pipeline_class_name=None,
output_path=None,
@@ -453,6 +454,7 @@ class TestCosmos3ActionEndpoint(unittest.TestCase):
def test_metadata_describes_cosmos_action_contract(self):
metadata = action_metadata(_cosmos3_server_args())
self.assertEqual(metadata["model"], "cosmos3-production")
self.assertEqual(metadata["policy_family"], "cosmos3")
self.assertEqual(metadata["input"]["modalities"], ["image", "video"])
self.assertEqual(metadata["output"]["action_horizon"], 16)
@@ -26,6 +26,7 @@ def _server_args(config: Pi05PipelineConfig | None = None) -> SimpleNamespace:
return SimpleNamespace(
model_id=None,
model_path="lerobot/pi05_base",
served_model_name="pi05-production",
output_path=None,
comfyui_mode=False,
num_gpus=1,
@@ -165,6 +166,7 @@ def test_action_metadata_reports_policy_shape_and_capabilities():
metadata = action_metadata(_server_args(config))
assert metadata["object"] == "action.metadata"
assert metadata["model"] == "pi05-production"
assert metadata["policy_family"] == "pi05"
assert metadata["input"]["image_keys"] == ["front", "wrist"]
assert metadata["input"]["image_size"] == [256, 256]
@@ -197,6 +199,7 @@ def test_action_generation_response_uses_actual_output_parameters():
assert response["id"] == "action-response-1"
assert response["object"] == "action.generation"
assert response["model"] == "pi05-production"
assert response["data"][0]["action"]["shape"] == [2, 2]
assert response["data"][0]["action"]["values"] == output["actions"]
assert response["usage"]["action_horizon"] == 2
@@ -0,0 +1,58 @@
import asyncio
from types import SimpleNamespace
from unittest.mock import patch
from sglang.multimodal_gen.runtime.entrypoints.openai import common_api
from sglang.multimodal_gen.runtime.entrypoints.openai.common_api import (
DiffusionModelCard,
)
from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import (
VideoGenerationsRequest,
)
from sglang.multimodal_gen.runtime.entrypoints.openai.video_api import (
_video_job_from_sampling,
)
def test_model_list_and_retrieve_use_served_model_name():
server_args = SimpleNamespace(
model_path="/private/checkpoint",
served_model_name="public-model",
)
def build_model_card(_server_args, model_name):
return DiffusionModelCard(id=model_name, root=model_name)
with (
patch.object(common_api, "get_global_server_args", return_value=server_args),
patch.object(common_api, "_build_model_card", side_effect=build_model_card),
):
models = asyncio.run(common_api.available_models())
model = asyncio.run(common_api.retrieve_model("public-model"))
missing = asyncio.run(common_api.retrieve_model("/private/checkpoint"))
assert models["data"][0]["id"] == "public-model"
assert model["id"] == "public-model"
assert missing.status_code == 404
def test_video_job_uses_served_model_name_unless_requested():
sampling = SimpleNamespace(
width=512,
height=512,
num_frames=49,
fps=24,
output_file_path=lambda: "/tmp/output.mp4",
)
request = VideoGenerationsRequest(prompt="test")
job = _video_job_from_sampling("request-id", request, sampling, "public-model")
explicit_job = _video_job_from_sampling(
"request-id",
request.model_copy(update={"model": "requested-model"}),
sampling,
"public-model",
)
assert job["model"] == "public-model"
assert explicit_job["model"] == "requested-model"
@@ -304,6 +304,41 @@ class TestServerArgsPathExpansion(unittest.TestCase):
["text_encoder", "image_encoder", "vae"],
)
def test_served_model_name_cli_arg(self):
parser = FlexibleArgumentParser()
ServerArgs.add_cli_args(parser)
cases = [
(
[
"--model-path",
"/fake",
"--model-id",
"Qwen-Image",
"--served-model-name",
"my-served-name",
],
"my-served-name",
),
(
["--model-path", "/fake", "--model-id", "Qwen-Image"],
"Qwen-Image",
),
(["--model-path", "/fake"], "/fake"),
]
for argv, expected in cases:
with self.subTest(argv=argv):
with patch.object(sys, "argv", ["sglang"] + argv):
args, unknown_args = parser.parse_known_args(argv)
with patch.object(
PipelineConfig,
"from_kwargs",
return_value=QwenImagePipelineConfig(),
):
server_args = ServerArgs.from_cli_args(args, unknown_args)
self.assertEqual(server_args.served_model_name, expected)
def test_dit_layerwise_offload_cli_arg(self):
parser = FlexibleArgumentParser()
ServerArgs.add_cli_args(parser)