chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[diffusion] refactor: unify SamplingParams construction and improve DiffGenerator return types (#18928)

Browse Source
This commit is contained in:
Mick
2026-02-18 14:56:58 +08:00
committed by GitHub
parent 9d138685c1
commit 420a611275
12 changed files with 478 additions and 537 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
python/sglang/multimodal_gen/configs/pipeline_configs/base.py
View File

@@ -352,6 +352,9 @@ class PipelineConfig:
# Pad to next multiple of SP degree if needed
if time_dim > 0 and time_dim % sp_world_size != 0:
logger.debug(
"Padding latents to next multiple of SP degree, performance is sub-optimal"
)
pad_len = sp_world_size - (time_dim % sp_world_size)
pad = torch.zeros(
(*latents.shape[:2], pad_len, *latents.shape[3:]),

18
python/sglang/multimodal_gen/runtime/entrypoints/cli/generate.py
View File

@@ -21,7 +21,6 @@ from sglang.multimodal_gen.runtime.entrypoints.cli.utils import (
from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.runtime.utils.perf_logger import (
MemorySnapshot,
PerformanceLogger,
RequestMetrics,
)
@@ -65,11 +64,13 @@ def maybe_dump_performance(args: argparse.Namespace, server_args, prompt: str, r
return
if isinstance(results, list):
result = results[0] if results else {}
result = results[0] if results else None
else:
result = results
timings_dict = result.get("timings")
timings_dict = getattr(result, "timings", None) or (
result.get("timings") if isinstance(result, dict) else None
)
if not (args.perf_dump_path and timings_dict):
return
@@ -78,17 +79,6 @@ def maybe_dump_performance(args: argparse.Namespace, server_args, prompt: str, r
timings.steps = timings_dict.get("steps", [])
timings.total_duration_ms = timings_dict.get("total_duration_ms", 0)
# restore memory snapshots from serialized dict
memory_snapshots_dict = timings_dict.get("memory_snapshots", {})
for checkpoint_name, snapshot_dict in memory_snapshots_dict.items():
snapshot = MemorySnapshot(
allocated_mb=snapshot_dict.get("allocated_mb", 0.0),
reserved_mb=snapshot_dict.get("reserved_mb", 0.0),
peak_allocated_mb=snapshot_dict.get("peak_allocated_mb", 0.0),
peak_reserved_mb=snapshot_dict.get("peak_reserved_mb", 0.0),
)
timings.memory_snapshots[checkpoint_name] = snapshot
PerformanceLogger.dump_benchmark_report(
file_path=args.perf_dump_path,
timings=timings,

226
python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py
View File

@@ -13,18 +13,15 @@ import os
import time
from typing import Any, List, Union
import numpy as np
from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
from sglang.multimodal_gen.runtime.entrypoints.utils import (
GenerationResult,
ListLorasReq,
MergeLoraWeightsReq,
SetLoraReq,
ShutdownReq,
UnmergeLoraWeightsReq,
format_lora_message,
)
from sglang.multimodal_gen.runtime.entrypoints.utils import (
prepare_request,
save_outputs,
)
@@ -157,65 +154,33 @@ class DiffGenerator:
def generate(
self,
sampling_params_kwargs: dict | None = None,
) -> dict[str, Any] | list[np.ndarray] | list[dict[str, Any]] | None:
"""
Generate a image/video based on the given prompt.
) -> GenerationResult | list[GenerationResult] | None:
"""Generate image(s)/video(s) based on the given prompt(s).
Args:
Returns:
Either the output dictionary, list of frames, or list of results for batch processing
Returns a single GenerationResult for a single prompt, a list for
multiple prompts, or None when every request failed.
"""
# 1. prepare requests
prompt = sampling_params_kwargs.get("prompt", None)
prompts: list[str] = []
# Handle batch processing from text file
if self.server_args.prompt_file_path is not None:
prompt_txt_path = self.server_args.prompt_file_path
if not os.path.exists(prompt_txt_path):
raise FileNotFoundError(
f"Prompt text file not found: {prompt_txt_path}"
)
# Read prompts from file
with open(prompt_txt_path, encoding="utf-8") as f:
prompts.extend(line.strip() for line in f if line.strip())
if not prompts:
raise ValueError(f"No prompts found in file: {prompt_txt_path}")
logger.info("Found %d prompts in %s", len(prompts), prompt_txt_path)
else:
if prompt is None:
prompt = " "
if isinstance(prompt, str):
prompts.append(prompt)
elif isinstance(prompt, list):
prompts.extend(prompt)
prompts = self._resolve_prompts(sampling_params_kwargs.get("prompt"))
sampling_params = SamplingParams.from_user_sampling_params_args(
self.server_args.model_path,
server_args=self.server_args,
**sampling_params_kwargs,
)
# Extract diffusers_kwargs if passed
diffusers_kwargs = sampling_params_kwargs.pop("diffusers_kwargs", None)
requests: list[Req] = []
for output_idx, p in enumerate(prompts):
for p in prompts:
sampling_params.prompt = p
req = prepare_request(
server_args=self.server_args,
sampling_params=sampling_params,
)
# Add diffusers_kwargs to request's extra dict
if diffusers_kwargs:
req.extra["diffusers_kwargs"] = diffusers_kwargs
requests.append(req)
results = []
results: list[GenerationResult] = []
total_start_time = time.perf_counter()
# 2. send requests to scheduler, one at a time
# 2. send requests to scheduler one at a time
# TODO: send batch when supported
for request_idx, req in enumerate(requests):
try:
@@ -235,102 +200,115 @@ class DiffGenerator:
request_idx + 1,
)
continue
audio_sample_rate = output_batch.audio_sample_rate
common = dict(
prompt=req.prompt,
size=(req.height, req.width, req.num_frames),
generation_time=timer.duration,
peak_memory_mb=output_batch.peak_memory_mb,
timings=(
output_batch.timings.to_dict()
if output_batch.timings
else {}
),
trajectory_latents=output_batch.trajectory_latents,
trajectory_timesteps=output_batch.trajectory_timesteps,
trajectory_decoded=output_batch.trajectory_decoded,
)
if req.save_output and req.return_file_paths_only:
for output_idx, output_path in enumerate(
output_batch.output_file_paths
):
result_item: dict[str, Any] = {
"samples": None,
"frames": None,
"audio": None,
"prompts": req.prompt,
"size": (req.height, req.width, req.num_frames),
"generation_time": timer.duration,
"peak_memory_mb": output_batch.peak_memory_mb,
"timings": (
output_batch.timings.to_dict()
if output_batch.timings
else {}
),
"trajectory": output_batch.trajectory_latents,
"trajectory_timesteps": output_batch.trajectory_timesteps,
"trajectory_decoded": output_batch.trajectory_decoded,
"prompt_index": output_idx,
"output_file_path": output_path,
}
results.append(result_item)
for idx, path in enumerate(output_batch.output_file_paths):
results.append(
GenerationResult(
**common,
prompt_index=idx,
output_file_path=path,
)
)
continue
samples_out: list[Any] = []
audios_out: list[Any] = []
frames_out: list[Any] = []
num_outputs = len(output_batch.output)
save_outputs(
output_batch.output,
req.data_type,
req.fps,
req.save_output,
lambda idx: req.output_file_path(len(output_batch.output), idx),
lambda idx: req.output_file_path(num_outputs, idx),
audio=output_batch.audio,
audio_sample_rate=audio_sample_rate,
audio_sample_rate=output_batch.audio_sample_rate,
samples_out=samples_out,
audios_out=audios_out,
frames_out=frames_out,
output_compression=req.output_compression,
)
for output_idx in range(len(samples_out)):
result_item: dict[str, Any] = {
"samples": samples_out[output_idx],
"frames": frames_out[output_idx],
"audio": audios_out[output_idx],
"prompts": req.prompt,
"size": (req.height, req.width, req.num_frames),
"generation_time": timer.duration,
"peak_memory_mb": output_batch.peak_memory_mb,
"timings": (
output_batch.timings.to_dict()
if output_batch.timings
else {}
),
"trajectory": output_batch.trajectory_latents,
"trajectory_timesteps": output_batch.trajectory_timesteps,
"trajectory_decoded": output_batch.trajectory_decoded,
"prompt_index": output_idx,
}
results.append(result_item)
except Exception:
for idx in range(len(samples_out)):
results.append(
GenerationResult(
**common,
samples=samples_out[idx],
frames=frames_out[idx],
audio=audios_out[idx],
prompt_index=idx,
output_file_path=req.output_file_path(num_outputs, idx),
)
)
except Exception as e:
logger.error(
"Generation failed for prompt %d/%d: %s",
request_idx + 1,
len(requests),
e,
exc_info=True,
)
continue
total_gen_time = time.perf_counter() - total_start_time
log_batch_completion(logger, len(results), total_gen_time)
self._log_summary(results)
if results:
if self.server_args.warmup:
total_duration_ms = results[0]["timings"]["total_duration_ms"]
logger.info(
f"Warmed-up request processed in {GREEN}%.2f{RESET} seconds (with warmup excluded)",
total_duration_ms / 1000.0,
)
peak_memories = [r.get("peak_memory_mb", 0) for r in results]
if peak_memories:
max_peak_memory = max(peak_memories)
avg_peak_memory = sum(peak_memories) / len(peak_memories)
logger.info(
f"Memory usage - Max peak: {max_peak_memory:.2f} MB, "
f"Avg peak: {avg_peak_memory:.2f} MB"
)
if len(results) == 0:
if not results:
return None
else:
if requests[0].return_frames:
results = [r["frames"] for r in results]
if len(results) == 1:
return results[0]
return results
return results[0] if len(results) == 1 else results
def _resolve_prompts(self, prompt: str | list[str] | None) -> list[str]:
"""Collect prompts from the argument or from a prompt file."""
if self.server_args.prompt_file_path is not None:
path = self.server_args.prompt_file_path
if not os.path.exists(path):
raise FileNotFoundError(f"Prompt text file not found: {path}")
with open(path, encoding="utf-8") as f:
prompts = [line.strip() for line in f if line.strip()]
if not prompts:
raise ValueError(f"No prompts found in file: {path}")
logger.info("Found %d prompts in %s", len(prompts), path)
return prompts
if prompt is None:
return [" "]
if isinstance(prompt, str):
return [prompt]
return list(prompt)
def _log_summary(self, results: list[GenerationResult]) -> None:
if not results:
return
if self.server_args.warmup:
total_duration_ms = results[0].timings.get("total_duration_ms", 0)
logger.info(
f"Warmed-up request processed in {GREEN}%.2f{RESET} seconds (with warmup excluded)",
total_duration_ms / 1000.0,
)
peak_memories = [r.peak_memory_mb for r in results if r.peak_memory_mb]
if peak_memories:
logger.info(
f"Memory usage - Max peak: {max(peak_memories):.2f} MB, "
f"Avg peak: {sum(peak_memories) / len(peak_memories):.2f} MB"
)
def _send_to_scheduler_and_wait_for_response(self, batch: list[Req]) -> OutputBatch:
"""
@@ -415,20 +393,15 @@ class DiffGenerator:
"Failed to merge LoRA weights",
)
def list_loras(self) -> OutputBatch:
"""
List loaded LoRA adapters and current application status per module.
"""
def list_loras(self) -> dict:
"""List loaded LoRA adapters and current application status per module."""
output = self._send_lora_request(
req=ListLorasReq(),
success_msg="Successfully listed LoRA adapters",
failure_msg="Failed to list LoRA adapters",
)
if output.error is None:
return output.output or {}
else:
raise RuntimeError(f"Failed to list LoRA adapters: {output.error}")
# _send_lora_request already raises on error, so output.error is always None here
return output.output or {}
def _ensure_lora_state(
self,
@@ -505,6 +478,9 @@ class DiffGenerator:
sync_scheduler_client.close()
self.owns_scheduler_client = False
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.shutdown()

56
python/sglang/multimodal_gen/runtime/entrypoints/http_server.py
View File

@@ -5,6 +5,7 @@ import base64
import os
import uuid
from contextlib import asynccontextmanager
from typing import TYPE_CHECKING
import torch
from fastapi import APIRouter, FastAPI, Request
@@ -15,12 +16,19 @@ from sglang.multimodal_gen.runtime.entrypoints.openai import image_api, video_ap
from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import (
VertexGenerateReqInput,
)
from sglang.multimodal_gen.runtime.entrypoints.openai.utils import build_sampling_params
from sglang.multimodal_gen.runtime.entrypoints.utils import (
prepare_request,
save_outputs,
)
from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client
from sglang.multimodal_gen.runtime.server_args import ServerArgs, get_global_server_args
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
if TYPE_CHECKING:
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
logger = init_logger(__name__)
DEFAULT_SEED = 1024
VERTEX_ROUTE = os.environ.get("AIP_PREDICT_ROUTE", "/vertex_generate")
@@ -43,7 +51,7 @@ async def lifespan(app: FastAPI):
yield
# On shutdown
print("FastAPI app is shutting down...")
logger.info("FastAPI app is shutting down...")
broker_task.cancel()
async_scheduler_client.close()
@@ -110,7 +118,10 @@ def encode_video_to_base64(file_path: str):
return base64.b64encode(f.read()).decode("utf-8")
async def forward_to_scheduler(req_obj, sp):
async def forward_to_scheduler(
req_obj: "Req",
sp: SamplingParams,
):
"""Forwards request to scheduler and processes the result."""
try:
response = await async_scheduler_client.forward(req_obj)
@@ -137,7 +148,7 @@ async def forward_to_scheduler(req_obj, sp):
data = response if isinstance(response, dict) else vars(response)
if output_file_path:
print(f"Processing output file: {output_file_path}")
logger.info("Processing output file: %s", output_file_path)
b64_video = encode_video_to_base64(output_file_path)
if b64_video:
@@ -148,7 +159,7 @@ async def forward_to_scheduler(req_obj, sp):
return make_serializable(data)
except Exception as e:
print(f"Error during generation: {e}")
logger.error("Error during generation: %s", e, exc_info=True)
return {"error": str(e)}
@@ -168,32 +179,17 @@ async def vertex_generate(vertex_req: VertexGenerateReqInput):
for inst in vertex_req.instances:
rid = f"vertex_{uuid.uuid4()}"
prompt = inst.get("prompt") or inst.get("text")
image_input = inst.get("image") or inst.get("image_url")
seed_val = params.get("seed", DEFAULT_SEED)
# Create a dictionary of provided parameters
# This filters out None values so the dataclass defaults kick in
user_params = {
"num_frames": params.get("num_frames"),
"fps": params.get("fps"),
"width": params.get("width"),
"height": params.get("height"),
"guidance_scale": params.get("guidance_scale"),
"save_output": params.get("save_output"),
}
# Remove None values to allow SamplingParams defaults to take over
valid_params = {k: v for k, v in user_params.items() if v is not None}
sp = SamplingParams.from_user_sampling_params_args(
model_path=server_args.model_path,
request_id=rid,
prompt=prompt,
image_path=image_input,
seed=seed_val,
server_args=server_args,
**valid_params, # Unpack the filtered dictionary
sp = build_sampling_params(
rid,
prompt=inst.get("prompt") or inst.get("text"),
image_path=inst.get("image") or inst.get("image_url"),
seed=params.get("seed", DEFAULT_SEED),
num_frames=params.get("num_frames"),
fps=params.get("fps"),
width=params.get("width"),
height=params.get("height"),
guidance_scale=params.get("guidance_scale"),
save_output=params.get("save_output"),
)
backend_req = prepare_request(server_args, sampling_params=sp)

2
python/sglang/multimodal_gen/runtime/entrypoints/openai/common_api.py
View File

@@ -6,7 +6,7 @@ from fastapi.responses import ORJSONResponse
from pydantic import BaseModel, Field
from sglang.multimodal_gen.registry import get_model_info
from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
from sglang.multimodal_gen.runtime.entrypoints.utils import (
ListLorasReq,
MergeLoraWeightsReq,
SetLoraReq,

275
python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py
View File

@@ -8,10 +8,7 @@ from typing import List, Optional
from fastapi import APIRouter, File, Form, HTTPException, Path, Query, UploadFile
from fastapi.responses import FileResponse
from sglang.multimodal_gen.configs.sample.sampling_params import (
SamplingParams,
generate_request_id,
)
from sglang.multimodal_gen.configs.sample.sampling_params import generate_request_id
from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import (
ImageGenerationsRequest,
ImageResponse,
@@ -20,14 +17,15 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import (
from sglang.multimodal_gen.runtime.entrypoints.openai.storage import cloud_storage
from sglang.multimodal_gen.runtime.entrypoints.openai.stores import IMAGE_STORE
from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
_parse_size,
add_common_data_to_response,
adjust_output_quality,
build_sampling_params,
choose_output_image_ext,
merge_image_input_list,
process_generation_batch,
save_image_to_path,
)
from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client
from sglang.multimodal_gen.runtime.server_args import get_global_server_args
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
@@ -36,92 +34,82 @@ router = APIRouter(prefix="/v1/images", tags=["images"])
logger = init_logger(__name__)
def _choose_ext(output_format: Optional[str], background: Optional[str]) -> str:
# Normalize and choose extension
fmt = (output_format or "").lower()
if fmt in {"png", "webp", "jpeg", "jpg"}:
return "jpg" if fmt == "jpeg" else fmt
# If transparency requested, prefer png
if (background or "auto").lower() == "transparent":
return "png"
# Default
return "jpg"
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 = []
for path in paths:
with open(path, "rb") as f:
result.append(base64.b64encode(f.read()).decode("utf-8"))
return result
def _build_sampling_params_from_request(
request_id: str,
def _build_image_response_kwargs(
save_file_path_list: list[str],
resp_format: str,
prompt: str,
n: int,
size: Optional[str],
output_format: Optional[str],
background: Optional[str],
image_path: Optional[list[str]] = None,
seed: Optional[int] = None,
generator_device: Optional[str] = None,
num_inference_steps: Optional[int] = None,
guidance_scale: Optional[float] = None,
true_cfg_scale: Optional[float] = None,
negative_prompt: Optional[str] = None,
enable_teacache: Optional[bool] = None,
num_frames: int = 1,
output_compression: Optional[int] = None,
) -> SamplingParams:
if size is None:
width, height = None, None
request_id: str,
result: OutputBatch,
*,
b64_list: list[str] | None = None,
cloud_url: str | None = None,
fallback_url: str | None = None,
) -> dict:
"""Build ImageResponse data list.
For b64_json: uses pre-read b64_list (call _read_b64_for_paths first).
For url: uses cloud_url or fallback_url.
"""
ret = None
if resp_format == "b64_json":
if not b64_list:
raise ValueError("b64_list required for b64_json response_format")
data = [
ImageResponseData(
b64_json=b64,
revised_prompt=prompt,
file_path=os.path.abspath(path),
)
for b64, path in zip(b64_list, save_file_path_list)
]
ret = {"data": data}
elif resp_format == "url":
url = cloud_url or fallback_url
if not url:
raise HTTPException(
status_code=400,
detail="response_format='url' requires cloud storage to be configured.",
)
ret = {
"data": [
ImageResponseData(
url=url,
revised_prompt=prompt,
file_path=os.path.abspath(save_file_path_list[0]),
)
],
}
else:
width, height = _parse_size(size)
ext = _choose_ext(output_format, background)
raise HTTPException(
status_code=400, detail=f"response_format={resp_format} is not supported"
)
server_args = get_global_server_args()
sampling_params = SamplingParams.from_user_sampling_params_args(
model_path=server_args.model_path,
request_id=request_id,
prompt=prompt,
image_path=image_path,
num_frames=num_frames,
width=width,
height=height,
num_outputs_per_prompt=max(1, min(int(n or 1), 10)),
save_output=True,
server_args=server_args,
output_file_name=f"{request_id}.{ext}",
seed=seed,
generator_device=generator_device,
num_inference_steps=num_inference_steps,
enable_teacache=enable_teacache,
**({"guidance_scale": guidance_scale} if guidance_scale is not None else {}),
**({"negative_prompt": negative_prompt} if negative_prompt is not None else {}),
**({"true_cfg_scale": true_cfg_scale} if true_cfg_scale is not None else {}),
**(
{"output_compression": output_compression}
if output_compression is not None
else {}
),
)
ret = add_common_data_to_response(ret, request_id=request_id, result=result)
if num_inference_steps is not None:
sampling_params.num_inference_steps = num_inference_steps
if guidance_scale is not None:
sampling_params.guidance_scale = guidance_scale
if seed is not None:
sampling_params.seed = seed
return sampling_params
return ret
@router.post("/generations", response_model=ImageResponse)
async def generations(
request: ImageGenerationsRequest,
):
request_id = generate_request_id()
sampling = _build_sampling_params_from_request(
request_id=request_id,
ext = choose_output_image_ext(request.output_format, request.background)
sampling = build_sampling_params(
request_id,
prompt=request.prompt,
n=request.n or 1,
size=request.size,
output_format=request.output_format,
background=request.background,
num_outputs_per_prompt=max(1, min(int(request.n or 1), 10)),
output_file_name=f"{request_id}.{ext}",
seed=request.seed,
generator_device=request.generator_device,
num_inference_steps=request.num_inference_steps,
@@ -130,37 +118,29 @@ async def generations(
negative_prompt=request.negative_prompt,
enable_teacache=request.enable_teacache,
output_compression=request.output_compression,
output_quality=request.output_quality,
)
batch = prepare_request(
server_args=get_global_server_args(),
sampling_params=sampling,
)
if batch.output_compression is None:
batch.output_compression = adjust_output_quality(
request.output_quality, batch.data_type
)
# Add diffusers_kwargs if provided
if request.diffusers_kwargs:
batch.extra["diffusers_kwargs"] = request.diffusers_kwargs
# Run synchronously for images and save to disk
save_file_path_list, result = await process_generation_batch(
async_scheduler_client, batch
)
save_file_path = save_file_path_list[0]
resp_format = (request.response_format or "b64_json").lower()
b64_data = None
# 1. Read content first if needed (while file exists)
if resp_format == "b64_json":
with open(save_file_path, "rb") as f:
b64_data = base64.b64encode(f.read()).decode("utf-8")
# read b64 before cloud upload may delete the local file
b64_list = (
_read_b64_for_paths(save_file_path_list) if resp_format == "b64_json" else None
)
# 2. Upload and Delete local file
cloud_url = await cloud_storage.upload_and_cleanup(save_file_path)
# 3. Update Database
await IMAGE_STORE.upsert(
request_id,
{
@@ -171,40 +151,16 @@ async def generations(
},
)
# 4. Return Response
if resp_format == "b64_json":
response_kwargs = {
"data": [
ImageResponseData(
b64_json=b64_data,
revised_prompt=request.prompt,
)
]
}
elif resp_format == "url":
if not cloud_url:
raise HTTPException(
status_code=400,
detail="response_format='url' requires cloud storage to be configured.",
)
response_kwargs = {
"data": [
ImageResponseData(
url=cloud_url,
revised_prompt=request.prompt,
file_path=os.path.abspath(save_file_path),
)
],
}
else:
# Return error, not supported
raise HTTPException(
status_code=400, detail=f"response_format={resp_format} is not supported"
)
response_kwargs = add_common_data_to_response(
response_kwargs, request_id=request_id, result=result
response_kwargs = _build_image_response_kwargs(
save_file_path_list,
resp_format,
request.prompt,
request_id,
result,
b64_list=b64_list,
cloud_url=cloud_url,
)
return ImageResponse(**response_kwargs)
@@ -245,8 +201,9 @@ async def edits(
)
# Save all input images; additional images beyond the first are saved for potential future use
uploads_dir = os.path.join("outputs", "uploads")
uploads_dir = os.path.join("inputs", "uploads")
os.makedirs(uploads_dir, exist_ok=True)
image_list = merge_image_input_list(images, urls)
input_paths = []
@@ -262,13 +219,13 @@ async def edits(
status_code=400, detail=f"Failed to process image source: {str(e)}"
)
sampling = _build_sampling_params_from_request(
request_id=request_id,
ext = choose_output_image_ext(output_format, background)
sampling = build_sampling_params(
request_id,
prompt=prompt,
n=n or 1,
size=size,
output_format=output_format,
background=background,
num_outputs_per_prompt=max(1, min(int(n or 1), 10)),
output_file_name=f"{request_id}.{ext}",
image_path=input_paths,
seed=seed,
generator_device=generator_device,
@@ -279,32 +236,25 @@ async def edits(
enable_teacache=enable_teacache,
num_frames=num_frames,
output_compression=output_compression,
output_quality=output_quality,
)
batch = prepare_request(
server_args=get_global_server_args(),
sampling_params=sampling,
)
if batch.output_compression is None:
batch.output_compression = adjust_output_quality(
output_quality, batch.data_type
)
save_file_path_list, result = await process_generation_batch(
async_scheduler_client, batch
)
save_file_path = save_file_path_list[0]
resp_format = (response_format or "b64_json").lower()
b64_data = None
# 1. Read content first if needed (while file exists)
if resp_format == "b64_json":
with open(save_file_path, "rb") as f:
b64_data = base64.b64encode(f.read()).decode("utf-8")
# read b64 before cloud upload may delete the local file
b64_list = (
_read_b64_for_paths(save_file_path_list) if resp_format == "b64_json" else None
)
# 2. Upload and Delete local file
cloud_url = await cloud_storage.upload_and_cleanup(save_file_path)
# 3. Update Database
await IMAGE_STORE.upsert(
request_id,
{
@@ -312,43 +262,22 @@ async def edits(
"created_at": int(time.time()),
"file_path": None if cloud_url else save_file_path,
"url": cloud_url,
"input_image_paths": input_paths, # Store all input image paths
"input_image_paths": input_paths,
"num_input_images": len(input_paths),
},
)
# 4. Return Response
if (response_format or "b64_json").lower() == "b64_json":
response_kwargs = {"data": []}
for path in save_file_path_list:
if path == save_file_path and b64_data is not None:
b64 = b64_data
else:
with open(path, "rb") as f:
b64 = base64.b64encode(f.read()).decode("utf-8")
response_kwargs["data"].append(
ImageResponseData(
b64_json=b64,
revised_prompt=prompt,
file_path=os.path.abspath(path),
)
)
if result.peak_memory_mb and result.peak_memory_mb > 0:
response_kwargs["peak_memory_mb"] = result.peak_memory_mb
else:
response_kwargs = {
"data": [
ImageResponseData(
url=cloud_url if cloud_url else f"/v1/images/{request_id}/content",
revised_prompt=prompt,
file_path=os.path.abspath(save_file_path),
)
],
}
response_kwargs = add_common_data_to_response(
response_kwargs, request_id=request_id, result=result
response_kwargs = _build_image_response_kwargs(
save_file_path_list,
resp_format,
prompt,
request_id,
result,
b64_list=b64_list,
cloud_url=cloud_url,
fallback_url=f"/v1/images/{request_id}/content",
)
return ImageResponse(**response_kwargs)

202
python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py
View File

@@ -1,6 +1,5 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
import base64
import dataclasses
import os
import re
import time
@@ -9,74 +8,43 @@ from typing import Any, List, Optional, Union
import httpx
from fastapi import UploadFile
from sglang.multimodal_gen.configs.sample.sampling_params import DataType
from sglang.multimodal_gen.runtime.entrypoints.utils import save_outputs
from sglang.multimodal_gen.configs.sample.sampling_params import (
DataType,
SamplingParams,
)
from sglang.multimodal_gen.runtime.entrypoints.utils import (
ListLorasReq,
MergeLoraWeightsReq,
SetLoraReq,
ShutdownReq,
UnmergeLoraWeightsReq,
format_lora_message,
save_outputs,
)
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
from sglang.multimodal_gen.runtime.scheduler_client import AsyncSchedulerClient
from sglang.multimodal_gen.runtime.server_args import get_global_server_args
from sglang.multimodal_gen.runtime.utils.logging_utils import (
init_logger,
log_batch_completion,
log_generation_timer,
)
# re-export LoRA protocol types for backward compatibility
__all__ = [
"SetLoraReq",
"MergeLoraWeightsReq",
"UnmergeLoraWeightsReq",
"ListLorasReq",
"ShutdownReq",
"format_lora_message",
]
logger = init_logger(__name__)
OUTPUT_QUALITY_MAPPER = {"maximum": 100, "high": 90, "medium": 55, "low": 35}
@dataclasses.dataclass
class SetLoraReq:
lora_nickname: Union[str, List[str]]
lora_path: Optional[Union[str, List[Optional[str]]]] = None
target: Union[str, List[str]] = "all"
strength: Union[float, List[float]] = 1.0 # LoRA strength for merge, default 1.0
@dataclasses.dataclass
class MergeLoraWeightsReq:
target: str = "all" # "all", "transformer", "transformer_2", "critic"
strength: float = 1.0 # LoRA strength for merge, default 1.0
@dataclasses.dataclass
class UnmergeLoraWeightsReq:
target: str = "all" # "all", "transformer", "transformer_2", "critic"
@dataclasses.dataclass
class ListLorasReq:
# Empty payload; used only as a type marker for listing LoRA status
pass
@dataclasses.dataclass
class ShutdownReq:
pass
def format_lora_message(
lora_nickname: Union[str, List[str]],
target: Union[str, List[str]],
strength: Union[float, List[float]],
) -> tuple[str, str, str]:
"""Format success message for single or multiple LoRAs"""
if isinstance(lora_nickname, list):
nickname_str = ", ".join(lora_nickname)
target_str = ", ".join(target) if isinstance(target, list) else target
strength_str = (
", ".join(f"{s:.2f}" for s in strength)
if isinstance(strength, list)
else f"{strength:.2f}"
)
else:
nickname_str = lora_nickname
target_str = target if isinstance(target, str) else ", ".join(target)
strength_str = (
f"{strength:.2f}"
if isinstance(strength, (int, float))
else ", ".join(f"{s:.2f}" for s in strength)
)
return nickname_str, target_str, strength_str
DEFAULT_FPS = 24
DEFAULT_VIDEO_SECONDS = 4
def _parse_size(size: str) -> tuple[int, int] | tuple[None, None]:
@@ -90,6 +58,62 @@ def _parse_size(size: str) -> tuple[int, int] | tuple[None, None]:
return None, None
def choose_output_image_ext(
output_format: Optional[str], background: Optional[str]
) -> str:
fmt = (output_format or "").lower()
if fmt in {"png", "webp", "jpeg", "jpg"}:
return "jpg" if fmt == "jpeg" else fmt
if (background or "auto").lower() == "transparent":
return "png"
return "jpg"
def build_sampling_params(request_id: str, **kwargs) -> SamplingParams:
"""Build SamplingParams from request parameters.
Handles size parsing, output_quality resolution, and None filtering before
delegating to SamplingParams.from_user_sampling_params_args. Callers pass
only the parameters they have; None values are stripped automatically so
that SamplingParams defaults apply.
"""
server_args = get_global_server_args()
# pop HTTP-layer params that aren't SamplingParams fields
output_quality = kwargs.pop("output_quality", None)
has_explicit_compression = kwargs.get("output_compression") is not None
# parse "WxH" size string if provided
size = kwargs.pop("size", None)
if size:
w, h = _parse_size(size)
if w is not None:
kwargs.setdefault("width", w)
kwargs.setdefault("height", h)
# filter out None values to let SamplingParams defaults apply
kwargs = {k: v for k, v in kwargs.items() if v is not None}
kwargs.setdefault("save_output", True)
sampling_params = SamplingParams.from_user_sampling_params_args(
model_path=server_args.model_path,
server_args=server_args,
request_id=request_id,
**kwargs,
)
# resolve output_quality → output_compression with the correct data_type.
# SamplingParams.__post_init__ may have resolved with the wrong data_type
# (default VIDEO) before _adjust() set the correct one.
if not has_explicit_compression and output_quality is not None:
resolved = adjust_output_quality(output_quality, sampling_params.data_type)
if resolved is not None:
sampling_params.output_compression = resolved
return sampling_params
async def save_image_to_path(image: Union[UploadFile, str], target_path: str) -> str:
input_path = await _maybe_url_image(image, target_path)
if input_path is None:
@@ -171,24 +195,23 @@ async def _save_url_image_to_path(image_url: str, target_path: str) -> str:
async def _save_base64_image_to_path(base64_data: str, target_path: str) -> str:
"""Decode base64 image data and save to target path."""
_B64_FMT_HINT = (
"Failed to decode base64 image. "
"Expected format: `data:[<media-type>];base64,<data>`"
)
# split `data:[<media-type>][;base64],<data>` to media-type base64 data
pattern = r"data:(.*?)(;base64)?,(.*)"
match = re.match(pattern, base64_data)
if not match:
raise ValueError(
f"Failed to decoding base64 image, please make sure the url format `data:[<media-type>][;base64],<data>` "
)
raise ValueError(_B64_FMT_HINT)
media_type = match.group(1)
is_base64 = match.group(2)
if not is_base64:
raise ValueError(
f"Failed to decoding base64 image, please make sure the url format `data:[<media-type>][;base64],<data>` "
)
raise ValueError(f"{_B64_FMT_HINT} (missing ;base64 marker)")
data = match.group(3)
if not data:
raise ValueError(
f"Failed to decoding base64 image, please make sure the url format `data:[<media-type>][;base64],<data>` "
)
raise ValueError(f"{_B64_FMT_HINT} (empty data payload)")
# get ext from url
if media_type.startswith("image/"):
ext = media_type.split("/")[-1].lower()
@@ -212,7 +235,7 @@ async def _save_base64_image_to_path(base64_data: str, target_path: str) -> str:
async def process_generation_batch(
scheduler_client: AsyncSchedulerClient,
batch,
) -> tuple[str, OutputBatch]:
) -> tuple[list[str], OutputBatch]:
total_start_time = time.perf_counter()
with log_generation_timer(logger, batch.prompt):
result = await scheduler_client.forward([batch])
@@ -222,40 +245,21 @@ async def process_generation_batch(
raise RuntimeError(
f"Model generation returned no output. Error from scheduler: {error_msg}"
)
save_file_path_list = []
# If output_file_paths is provided, use it instead of output.
if result.output_file_paths:
save_file_path_list = result.output_file_paths
else:
audio_sample_rate = result.audio_sample_rate
if batch.data_type == DataType.VIDEO:
save_file_path_list = save_outputs(
result.output,
batch.data_type,
batch.fps,
batch.save_output,
lambda _idx: str(
os.path.join(batch.output_path, batch.output_file_name)
),
audio=result.audio,
audio_sample_rate=audio_sample_rate,
output_compression=batch.output_compression,
)
else:
save_file_path_list = save_outputs(
result.output,
batch.data_type,
batch.fps,
batch.save_output,
lambda idx: str(
os.path.join(
batch.output_path,
f"sample_{idx}_" + batch.output_file_name,
)
),
audio_sample_rate=audio_sample_rate,
output_compression=batch.output_compression,
)
num_outputs = len(result.output)
save_file_path_list = save_outputs(
result.output,
batch.data_type,
batch.fps,
batch.save_output,
lambda idx: str(batch.output_file_path(num_outputs, idx)),
audio=result.audio,
audio_sample_rate=result.audio_sample_rate,
output_compression=batch.output_compression,
)
total_time = time.perf_counter() - total_start_time
log_batch_completion(logger, 1, total_time)

146
python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py
View File

@@ -30,9 +30,10 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import (
from sglang.multimodal_gen.runtime.entrypoints.openai.storage import cloud_storage
from sglang.multimodal_gen.runtime.entrypoints.openai.stores import VIDEO_STORE
from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
_parse_size,
DEFAULT_FPS,
DEFAULT_VIDEO_SECONDS,
add_common_data_to_response,
adjust_output_quality,
build_sampling_params,
merge_image_input_list,
process_generation_batch,
save_image_to_path,
@@ -46,66 +47,33 @@ logger = init_logger(__name__)
router = APIRouter(prefix="/v1/videos", tags=["videos"])
# NOTE(mick): the sampling params needs to be further adjusted
# FIXME: duplicated with the one in `image_api.py`
def _build_sampling_params_from_request(
request_id: str, request: VideoGenerationsRequest
) -> SamplingParams:
if request.size is None:
width, height = None, None
else:
width, height = _parse_size(request.size)
seconds = request.seconds if request.seconds is not None else 4
fps_default = 24
fps = request.fps if request.fps is not None else fps_default
derived_num_frames = fps * seconds
num_frames = (
request.num_frames if request.num_frames is not None else derived_num_frames
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."""
seconds = request.seconds if request.seconds is not None else DEFAULT_VIDEO_SECONDS
fps = request.fps if request.fps is not None else DEFAULT_FPS
num_frames = request.num_frames if request.num_frames is not None else fps * seconds
return build_sampling_params(
request_id,
prompt=request.prompt,
size=request.size,
num_frames=num_frames,
fps=fps,
image_path=request.input_reference,
output_file_name=request_id,
seed=request.seed,
generator_device=request.generator_device,
num_inference_steps=request.num_inference_steps,
guidance_scale=request.guidance_scale,
guidance_scale_2=request.guidance_scale_2,
negative_prompt=request.negative_prompt,
enable_teacache=request.enable_teacache,
output_path=request.output_path,
output_compression=request.output_compression,
output_quality=request.output_quality,
)
server_args = get_global_server_args()
sampling_kwargs = {
"request_id": request_id,
"prompt": request.prompt,
"num_frames": num_frames,
"fps": fps,
"width": width,
"height": height,
"image_path": request.input_reference,
"save_output": True,
"output_file_name": request_id,
"seed": request.seed,
"generator_device": request.generator_device,
}
if request.num_inference_steps is not None:
sampling_kwargs["num_inference_steps"] = request.num_inference_steps
if request.guidance_scale is not None:
sampling_kwargs["guidance_scale"] = request.guidance_scale
if request.guidance_scale_2 is not None:
sampling_kwargs["guidance_scale_2"] = request.guidance_scale_2
if request.negative_prompt is not None:
sampling_kwargs["negative_prompt"] = request.negative_prompt
if request.enable_teacache is not None:
sampling_kwargs["enable_teacache"] = request.enable_teacache
if request.output_path is not None:
sampling_kwargs["output_path"] = request.output_path
if request.output_compression is not None:
sampling_kwargs["output_compression"] = request.output_compression
sampling_params = SamplingParams.from_user_sampling_params_args(
model_path=server_args.model_path,
server_args=server_args,
**sampling_kwargs,
)
if request.num_inference_steps is not None:
sampling_params.num_inference_steps = request.num_inference_steps
if request.guidance_scale is not None:
sampling_params.guidance_scale = request.guidance_scale
if request.seed is not None:
sampling_params.seed = request.seed
return sampling_params
# extract metadata which http_server needs to know
def _video_job_from_sampling(
@@ -127,6 +95,21 @@ def _video_job_from_sampling(
}
async def _save_first_input_image(image_sources, request_id: str) -> str | None:
"""Save the first input image from a list of sources and return its path."""
image_list = merge_image_input_list(image_sources)
if not image_list:
return None
image = image_list[0]
uploads_dir = os.path.join("inputs", "uploads")
os.makedirs(uploads_dir, exist_ok=True)
filename = image.filename if hasattr(image, "filename") else "url_image"
target_path = os.path.join(uploads_dir, f"{request_id}_{filename}")
return await save_image_to_path(image, target_path)
async def _dispatch_job_async(job_id: str, batch: Req) -> None:
from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client
@@ -190,26 +173,18 @@ async def create_video(
if not prompt:
raise HTTPException(status_code=400, detail="prompt is required")
# Validate image input based on model task type
image_list = merge_image_input_list(input_reference, reference_url)
if task_type.requires_image_input() and not image_list:
image_sources = merge_image_input_list(input_reference, reference_url)
if task_type.requires_image_input() and not image_sources:
raise HTTPException(
status_code=400,
detail="input_reference or reference_url is required for image-to-video generation",
)
input_path = None
if image_list:
# Save first input image for image-to-video generation
image = image_list[0]
uploads_dir = os.path.join("outputs", "uploads")
os.makedirs(uploads_dir, exist_ok=True)
filename = image.filename if hasattr(image, "filename") else "url_image"
input_path = os.path.join(uploads_dir, f"{request_id}_{filename}")
try:
input_path = await save_image_to_path(image, input_path)
except Exception as e:
raise HTTPException(
status_code=400, detail=f"Failed to process image source: {str(e)}"
)
try:
input_path = await _save_first_input_image(image_sources, request_id)
except Exception as e:
raise HTTPException(
status_code=400, detail=f"Failed to process image source: {str(e)}"
)
# Parse extra_body JSON (if provided in multipart form) to get fps/num_frames overrides
extra_from_form: Dict[str, Any] = {}
@@ -268,17 +243,12 @@ async def create_video(
status_code=400,
detail="input_reference or reference_url is required for image-to-video generation",
)
# for not multipart/form-data type
# for non-multipart/form-data type
if payload.get("reference_url"):
image_list = merge_image_input_list(payload.get("reference_url"))
# Save first input image
image = image_list[0]
uploads_dir = os.path.join("outputs", "uploads")
os.makedirs(uploads_dir, exist_ok=True)
filename = image.filename if hasattr(image, "filename") else "url_image"
input_path = os.path.join(uploads_dir, f"{request_id}_{filename}")
try:
input_path = await save_image_to_path(image, input_path)
input_path = await _save_first_input_image(
payload.get("reference_url"), request_id
)
except Exception as e:
raise HTTPException(
status_code=400,
@@ -291,7 +261,7 @@ async def create_video(
logger.debug(f"Server received from create_video endpoint: req={req}")
sampling_params = _build_sampling_params_from_request(request_id, req)
sampling_params = _build_video_sampling_params(request_id, req)
job = _video_job_from_sampling(request_id, req, sampling_params)
await VIDEO_STORE.upsert(request_id, job)
@@ -300,10 +270,6 @@ async def create_video(
server_args=server_args,
sampling_params=sampling_params,
)
if batch.output_compression is None:
batch.output_compression = adjust_output_quality(
req.output_quality, batch.data_type
)
# Add diffusers_kwargs if provided
if req.diffusers_kwargs:
batch.extra["diffusers_kwargs"] = req.diffusers_kwargs

76
python/sglang/multimodal_gen/runtime/entrypoints/utils.py
View File

@@ -12,7 +12,8 @@ import os
import shutil
import subprocess
import tempfile
from typing import Any, Callable, Optional, Sequence
from dataclasses import dataclass, field
from typing import Any, Callable, List, Optional, Sequence, Union
import imageio
import numpy as np
@@ -39,6 +40,79 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import CYAN, RESET, init_
logger = init_logger(__name__)
@dataclass
class SetLoraReq:
lora_nickname: Union[str, List[str]]
lora_path: Optional[Union[str, List[Optional[str]]]] = None
target: Union[str, List[str]] = "all"
strength: Union[float, List[float]] = 1.0
@dataclass
class MergeLoraWeightsReq:
target: str = "all"
strength: float = 1.0
@dataclass
class UnmergeLoraWeightsReq:
target: str = "all"
@dataclass
class ListLorasReq:
pass
@dataclass
class ShutdownReq:
pass
def format_lora_message(
lora_nickname: Union[str, List[str]],
target: Union[str, List[str]],
strength: Union[float, List[float]],
) -> tuple[str, str, str]:
"""Format success message for single or multiple LoRAs."""
if isinstance(lora_nickname, list):
nickname_str = ", ".join(lora_nickname)
target_str = ", ".join(target) if isinstance(target, list) else target
strength_str = (
", ".join(f"{s:.2f}" for s in strength)
if isinstance(strength, list)
else f"{strength:.2f}"
)
else:
nickname_str = lora_nickname
target_str = target if isinstance(target, str) else ", ".join(target)
strength_str = (
f"{strength:.2f}"
if isinstance(strength, (int, float))
else ", ".join(f"{s:.2f}" for s in strength)
)
return nickname_str, target_str, strength_str
@dataclass
class GenerationResult:
"""Result of a single generation request from DiffGenerator."""
samples: Any = None
frames: Any = None
audio: Any = None
prompt: str | None = None
size: tuple | None = None # (height, width, num_frames)
generation_time: float = 0.0
peak_memory_mb: float = 0.0
timings: dict = field(default_factory=dict)
trajectory_latents: Any = None
trajectory_timesteps: Any = None
trajectory_decoded: Any = None
prompt_index: int = 0
output_file_path: str | None = None
def _normalize_audio_to_numpy(audio: Any) -> np.ndarray | None:
"""Convert audio (torch / numpy) into a float32 numpy array in [-1, 1], best-effort."""
if audio is None:

6
python/sglang/multimodal_gen/runtime/managers/scheduler.py
View File

@@ -12,13 +12,15 @@ import zmq
from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType
from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
_parse_size,
save_image_to_path,
)
from sglang.multimodal_gen.runtime.entrypoints.utils import (
ListLorasReq,
MergeLoraWeightsReq,
SetLoraReq,
ShutdownReq,
UnmergeLoraWeightsReq,
_parse_size,
save_image_to_path,
)
from sglang.multimodal_gen.runtime.managers.gpu_worker import GPUWorker
from sglang.multimodal_gen.runtime.pipelines_core import Req

2
python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py
View File

@@ -159,7 +159,7 @@ class InputValidationStage(PipelineStage):
scale = max(ow / iw, oh / ih)
img = img.resize((round(iw * scale), round(ih * scale)), Image.LANCZOS)
logger.debug("resized img height: %s, img width: %s", img.height, img.width)
logger.debug("resized condition image to: %sx%s", img.height, img.width)
# center-crop
x1 = (img.width - ow) // 2

3
python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py
View File

@@ -131,7 +131,8 @@ class TimestepPreparationStage(PipelineStage):
# Update batch with prepared timesteps
batch.timesteps = timesteps
self.log_debug("timesteps: %s", timesteps)
if not batch.is_warmup:
self.log_debug("timesteps: %s", timesteps)
return batch
def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: