[diffusion] fix: make input/output file save paths configurable and disableable (#19580)

Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Ruihang Li
2026-03-02 23:02:33 +08:00
committed by GitHub
parent 53de53fb53
commit 5833ea684d
9 changed files with 315 additions and 155 deletions

View File

@@ -539,7 +539,7 @@ class PipelineConfig:
cls, kwargs: dict[str, Any], config_cli_prefix: str = ""
) -> "PipelineConfig":
"""
Load PipelineConfig from kwargs Dictionary.
Load PipelineConfig from kwargs Dictionary, as part of the ServerArg initialization process
kwargs: dictionary of kwargs
config_cli_prefix: prefix of CLI arguments for this PipelineConfig instance
"""
@@ -583,7 +583,11 @@ class PipelineConfig:
f"using {pipeline_config_cls.__name__} directly without model_index.json"
)
else:
model_info = get_model_info(model_path, backend=kwargs.get("backend"))
model_info = get_model_info(
model_path,
backend=kwargs.get("backend"),
model_id=kwargs.get("model_id"),
)
if model_info is None:
from sglang.multimodal_gen.registry import (
_PIPELINE_CONFIG_REGISTRY,
@@ -599,7 +603,11 @@ class PipelineConfig:
)
pipeline_config_cls = model_info.pipeline_config_cls
else:
model_info = get_model_info(model_path, backend=kwargs.get("backend"))
model_info = get_model_info(
model_path,
backend=kwargs.get("backend"),
model_id=kwargs.get("model_id"),
)
if model_info is None:
raise ValueError(
f"Could not get model info for '{model_path}'. "

View File

@@ -14,13 +14,16 @@ import unicodedata
import uuid
from dataclasses import dataclass
from enum import Enum, auto
from typing import Any
from typing import TYPE_CHECKING, Any
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.utils import StoreBoolean
from sglang.multimodal_gen.utils import StoreBoolean, expand_path_fields
logger = init_logger(__name__)
if TYPE_CHECKING:
from sglang.multimodal_gen.runtime.server_args import ServerArgs
def _json_safe(obj: Any):
"""
@@ -354,6 +357,8 @@ class SamplingParams:
"""
final adjustment, called after merged with user params
"""
expand_path_fields(self)
# TODO: SamplingParams should not rely on ServerArgs
pipeline_config = server_args.pipeline_config
if not isinstance(self.prompt, str):
@@ -374,11 +379,14 @@ class SamplingParams:
self.data_type = server_args.pipeline_config.task_type.data_type()
if self.output_path is None and server_args.output_path is not None:
self.output_path = server_args.output_path
logger.debug(
f"Overriding output_path with server configuration: {self.output_path}"
)
if self.output_path is None:
if server_args.output_path is not None:
self.output_path = server_args.output_path
logger.debug(
f"Overriding output_path with server configuration: {self.output_path}"
)
else:
self.save_output = False
# Process negative prompt
if self.negative_prompt is not None and not self.negative_prompt.isspace():
@@ -504,15 +512,18 @@ class SamplingParams:
from sglang.multimodal_gen.registry import get_model_info
backend = kwargs.pop("backend", None)
model_info = get_model_info(model_path, backend=backend)
model_id = kwargs.pop("model_id", None)
model_info = get_model_info(model_path, backend=backend, model_id=model_id)
sampling_params: SamplingParams = model_info.sampling_param_cls(**kwargs)
return sampling_params
@staticmethod
def from_user_sampling_params_args(model_path: str, server_args, *args, **kwargs):
def from_user_sampling_params_args(
model_path: str, server_args: "ServerArgs", *args, **kwargs
):
try:
sampling_params = SamplingParams.from_pretrained(
model_path, backend=server_args.backend
model_path, backend=server_args.backend, model_id=server_args.model_id
)
except (AttributeError, ValueError) as e:
# Handle safetensors files or other cases where model_index.json is not available
@@ -890,6 +901,8 @@ class SamplingParams:
return {attr: getattr(args, attr) for attr in attrs if hasattr(args, attr)}
def output_file_path(self):
if self.output_path is None:
return None
return os.path.join(self.output_path, self.output_file_name)
def _merge_with_user_params(self, user_params: "SamplingParams"):

View File

@@ -1,6 +1,7 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
import base64
import contextlib
import os
import time
from typing import List, Optional
@@ -23,6 +24,7 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
merge_image_input_list,
process_generation_batch,
save_image_to_path,
temp_dir_if_disabled,
)
from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
@@ -53,11 +55,13 @@ def _build_image_response_kwargs(
b64_list: list[str] | None = None,
cloud_url: str | None = None,
fallback_url: str | None = None,
is_persistent: bool = True,
) -> 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.
file_path is omitted when is_persistent=False to avoid exposing stale temp paths.
"""
ret = None
if resp_format == "b64_json":
@@ -67,7 +71,7 @@ def _build_image_response_kwargs(
ImageResponseData(
b64_json=b64,
revised_prompt=prompt,
file_path=os.path.abspath(path),
file_path=os.path.abspath(path) if is_persistent else None,
)
for b64, path in zip(b64_list, save_file_path_list)
]
@@ -84,7 +88,11 @@ def _build_image_response_kwargs(
ImageResponseData(
url=url,
revised_prompt=prompt,
file_path=os.path.abspath(save_file_path_list[0]),
file_path=(
os.path.abspath(save_file_path_list[0])
if is_persistent
else None
),
)
],
}
@@ -103,65 +111,72 @@ async def generations(
request: ImageGenerationsRequest,
):
request_id = generate_request_id()
server_args = get_global_server_args()
ext = choose_output_image_ext(request.output_format, request.background)
sampling = build_sampling_params(
request_id,
prompt=request.prompt,
size=request.size,
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,
guidance_scale=request.guidance_scale,
true_cfg_scale=request.true_cfg_scale,
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,
)
# Add diffusers_kwargs if provided
if request.diffusers_kwargs:
batch.extra["diffusers_kwargs"] = request.diffusers_kwargs
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()
with temp_dir_if_disabled(server_args.output_path) as output_dir:
sampling = build_sampling_params(
request_id,
prompt=request.prompt,
size=request.size,
num_outputs_per_prompt=max(1, min(int(request.n or 1), 10)),
output_file_name=f"{request_id}.{ext}",
output_path=output_dir,
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,
enable_teacache=request.enable_teacache,
output_compression=request.output_compression,
output_quality=request.output_quality,
)
batch = prepare_request(
server_args=server_args,
sampling_params=sampling,
)
# Add diffusers_kwargs if provided
if request.diffusers_kwargs:
batch.extra["diffusers_kwargs"] = request.diffusers_kwargs
# 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
)
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()
cloud_url = await cloud_storage.upload_and_cleanup(save_file_path)
# 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
)
await IMAGE_STORE.upsert(
request_id,
{
"id": request_id,
"created_at": int(time.time()),
"file_path": None if cloud_url else save_file_path,
"url": cloud_url,
},
)
cloud_url = await cloud_storage.upload_and_cleanup(save_file_path)
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,
# Provide a local fallback URL when cloud upload is unavailable, mirroring /v1/images/edits.
fallback_url=f"/v1/images/{request_id}/content",
)
is_persistent = server_args.output_path is not None
await IMAGE_STORE.upsert(
request_id,
{
"id": request_id,
"created_at": int(time.time()),
"file_path": None if cloud_url or not is_persistent else save_file_path,
"url": cloud_url,
},
)
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,
fallback_url=f"/v1/images/{request_id}/content" if is_persistent else None,
is_persistent=is_persistent,
)
return ImageResponse(**response_kwargs)
@@ -193,6 +208,7 @@ async def edits(
num_frames: int = Form(1),
):
request_id = generate_request_id()
server_args = get_global_server_args()
# Resolve images from either `image` or `image[]` (OpenAI SDK sends `image[]` when list is provided)
images = image or image_array
urls = url or url_array
@@ -202,83 +218,93 @@ async def edits(
status_code=422, detail="Field 'image' or 'url' is required"
)
# Save all input images; additional images beyond the first are saved for potential future use
uploads_dir = os.path.join("inputs", "uploads")
os.makedirs(uploads_dir, exist_ok=True)
image_list = merge_image_input_list(images, urls)
input_paths = []
try:
for idx, img in enumerate(image_list):
filename = img.filename if hasattr(img, "filename") else f"image_{idx}"
input_path = await save_image_to_path(
img, os.path.join(uploads_dir, f"{request_id}_{idx}_{filename}")
with contextlib.ExitStack() as stack:
uploads_dir = stack.enter_context(
temp_dir_if_disabled(server_args.input_save_path)
)
output_dir = stack.enter_context(temp_dir_if_disabled(server_args.output_path))
input_paths = []
try:
for idx, img in enumerate(image_list):
filename = img.filename if hasattr(img, "filename") else f"image_{idx}"
input_path = await save_image_to_path(
img,
os.path.join(uploads_dir, f"{request_id}_{idx}_{filename}"),
)
input_paths.append(input_path)
except Exception as e:
raise HTTPException(
status_code=400,
detail=f"Failed to process image source: {str(e)}",
)
input_paths.append(input_path)
except Exception as e:
raise HTTPException(
status_code=400, detail=f"Failed to process image source: {str(e)}"
ext = choose_output_image_ext(output_format, background)
sampling = build_sampling_params(
request_id,
prompt=prompt,
size=size,
num_outputs_per_prompt=max(1, min(int(n or 1), 10)),
output_file_name=f"{request_id}.{ext}",
output_path=output_dir,
image_path=input_paths,
seed=seed,
generator_device=generator_device,
negative_prompt=negative_prompt,
guidance_scale=guidance_scale,
true_cfg_scale=true_cfg_scale,
num_inference_steps=num_inference_steps,
enable_teacache=enable_teacache,
num_frames=num_frames,
output_compression=output_compression,
output_quality=output_quality,
)
batch = prepare_request(
server_args=server_args,
sampling_params=sampling,
)
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()
# 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
)
ext = choose_output_image_ext(output_format, background)
sampling = build_sampling_params(
request_id,
prompt=prompt,
size=size,
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,
negative_prompt=negative_prompt,
guidance_scale=guidance_scale,
true_cfg_scale=true_cfg_scale,
num_inference_steps=num_inference_steps,
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,
)
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()
cloud_url = await cloud_storage.upload_and_cleanup(save_file_path)
# 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
)
is_persistent = server_args.output_path is not None
is_input_persistent = server_args.input_save_path is not None
await IMAGE_STORE.upsert(
request_id,
{
"id": request_id,
"created_at": int(time.time()),
"file_path": None if cloud_url or not is_persistent else save_file_path,
"url": cloud_url,
"input_image_paths": input_paths if is_input_persistent else None,
"num_input_images": len(input_paths),
},
)
cloud_url = await cloud_storage.upload_and_cleanup(save_file_path)
await IMAGE_STORE.upsert(
request_id,
{
"id": request_id,
"created_at": int(time.time()),
"file_path": None if cloud_url else save_file_path,
"url": cloud_url,
"input_image_paths": input_paths,
"num_input_images": len(input_paths),
},
)
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",
)
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" if is_persistent else None,
is_persistent=is_persistent,
)
return ImageResponse(**response_kwargs)
@@ -298,7 +324,12 @@ async def download_image_content(
)
file_path = item.get("file_path")
if not file_path or not os.path.exists(file_path):
if not file_path:
raise HTTPException(
status_code=404,
detail="Image was not persisted on disk (output_path is disabled). Use b64_json response_format or configure cloud storage.",
)
if not os.path.exists(file_path):
raise HTTPException(status_code=404, detail="Image is still being generated")
ext = os.path.splitext(file_path)[1].lower()

View File

@@ -2,8 +2,11 @@
import base64
import os
import re
import shutil
import tempfile
import time
from typing import Any, List, Optional, Union
from contextlib import contextmanager
from typing import Any, Generator, List, Optional, Union
import httpx
from fastapi import UploadFile
@@ -47,6 +50,23 @@ DEFAULT_FPS = 24
DEFAULT_VIDEO_SECONDS = 4
@contextmanager
def temp_dir_if_disabled(
configured_path: str | None,
) -> Generator[str, None, None]:
"""Yield *configured_path* when it is set, otherwise create a temporary
directory that is automatically removed when the context exits."""
if configured_path is not None:
os.makedirs(configured_path, exist_ok=True)
yield configured_path
else:
tmp = tempfile.mkdtemp(prefix="sglang_")
try:
yield tmp
finally:
shutil.rmtree(tmp, ignore_errors=True)
def _parse_size(size: str) -> tuple[int, int] | tuple[None, None]:
try:
parts = size.lower().replace(" ", "").split("x")

View File

@@ -3,6 +3,8 @@
import asyncio
import json
import os
import shutil
import tempfile
import time
from typing import Any, Dict, Optional
@@ -99,14 +101,15 @@ def _video_job_from_sampling(
}
async def _save_first_input_image(image_sources, request_id: str) -> str | None:
async def _save_first_input_image(
image_sources, request_id: str, uploads_dir: 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"
@@ -114,7 +117,13 @@ async def _save_first_input_image(image_sources, request_id: str) -> str | None:
return await save_image_to_path(image, target_path)
async def _dispatch_job_async(job_id: str, batch: Req) -> None:
async def _dispatch_job_async(
job_id: str,
batch: Req,
*,
temp_dirs: list[str] | None = None,
output_persistent: bool = True,
) -> None:
from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client
try:
@@ -125,12 +134,15 @@ async def _dispatch_job_async(job_id: str, batch: Req) -> None:
cloud_url = await cloud_storage.upload_and_cleanup(save_file_path)
persistent_path = (
save_file_path if not cloud_url and output_persistent else None
)
update_fields = {
"status": "completed",
"progress": 100,
"completed_at": int(time.time()),
"url": cloud_url,
"file_path": save_file_path if not cloud_url else None,
"file_path": persistent_path,
}
update_fields = add_common_data_to_response(
update_fields, request_id=job_id, result=result
@@ -141,6 +153,9 @@ async def _dispatch_job_async(job_id: str, batch: Req) -> None:
await VIDEO_STORE.update_fields(
job_id, {"status": "failed", "error": {"message": str(e)}}
)
finally:
for td in temp_dirs or []:
shutil.rmtree(td, ignore_errors=True)
# TODO: support image to video generation
@@ -176,6 +191,22 @@ async def create_video(
server_args = get_global_server_args()
task_type = server_args.pipeline_config.task_type
# Resolve input upload directory (may be a temp dir when saving is disabled)
temp_dirs: list[str] = []
if server_args.input_save_path is not None:
uploads_dir = server_args.input_save_path
os.makedirs(uploads_dir, exist_ok=True)
else:
uploads_dir = tempfile.mkdtemp(prefix="sglang_input_")
temp_dirs.append(uploads_dir)
# Resolve output directory
effective_output_path = server_args.output_path
output_persistent = True
if "multipart/form-data" not in content_type:
# JSON body may carry a per-request output_path; checked after parsing below
pass
if "multipart/form-data" in content_type:
if not prompt:
raise HTTPException(status_code=400, detail="prompt is required")
@@ -187,7 +218,9 @@ async def create_video(
detail="input_reference or reference_url is required for image-to-video generation",
)
try:
input_path = await _save_first_input_image(image_sources, request_id)
input_path = await _save_first_input_image(
image_sources, request_id, uploads_dir
)
except Exception as e:
raise HTTPException(
status_code=400, detail=f"Failed to process image source: {str(e)}"
@@ -258,7 +291,7 @@ async def create_video(
if payload.get("reference_url"):
try:
input_path = await _save_first_input_image(
payload.get("reference_url"), request_id
payload.get("reference_url"), request_id, uploads_dir
)
except Exception as e:
raise HTTPException(
@@ -270,6 +303,17 @@ async def create_video(
except Exception as e:
raise HTTPException(status_code=400, detail=f"Invalid request body: {e}")
# Resolve per-request output_path override
effective_output_path = req.output_path or server_args.output_path
if effective_output_path is None:
output_tmp = tempfile.mkdtemp(prefix="sglang_output_")
temp_dirs.append(output_tmp)
effective_output_path = output_tmp
output_persistent = False
# Inject resolved output_path so _build_video_sampling_params picks it up
req.output_path = effective_output_path
logger.debug(f"Server received from create_video endpoint: req={req}")
try:
@@ -289,7 +333,14 @@ async def create_video(
if req.diffusers_kwargs:
batch.extra["diffusers_kwargs"] = req.diffusers_kwargs
# Enqueue the job asynchronously and return immediately
asyncio.create_task(_dispatch_job_async(request_id, batch))
asyncio.create_task(
_dispatch_job_async(
request_id,
batch,
temp_dirs=temp_dirs or None,
output_persistent=output_persistent,
)
)
return VideoResponse(**job)

View File

@@ -247,11 +247,9 @@ class Req:
base, ext = os.path.splitext(output_file_name)
output_file_name = f"{base}_{output_idx}{ext}"
return (
os.path.join(self.output_path, output_file_name)
if output_file_name
else None
)
if self.output_path is None or not output_file_name:
return None
return os.path.join(self.output_path, output_file_name)
def set_as_warmup(self):
self.is_warmup = True

View File

@@ -39,7 +39,11 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import (
configure_logger,
init_logger,
)
from sglang.multimodal_gen.utils import FlexibleArgumentParser, StoreBoolean
from sglang.multimodal_gen.utils import (
FlexibleArgumentParser,
StoreBoolean,
expand_path_fields,
)
logger = init_logger(__name__)
@@ -293,6 +297,7 @@ class ServerArgs:
scheduler_port: int = 5555
output_path: str | None = "outputs/"
input_save_path: str | None = "inputs/uploads"
# Prompt text file for batch processing
prompt_file_path: str | None = None
@@ -332,7 +337,8 @@ class ServerArgs:
return self.host is None or self.port is None
def _adjust_path(self):
self.model_path = os.path.expanduser(self.model_path)
expand_path_fields(self)
self._adjust_save_paths()
def _adjust_parameters(self):
"""set defaults and normalize values."""
@@ -354,6 +360,13 @@ class ServerArgs:
self._validate_parallelism()
self._validate_cfg_parallel()
def _adjust_save_paths(self):
"""Normalize empty-string save paths to None (disabled)."""
if self.output_path is not None and self.output_path.strip() == "":
self.output_path = None
if self.input_save_path is not None and self.input_save_path.strip() == "":
self.input_save_path = None
def _adjust_quant_config(self):
"""validate and adjust"""
@@ -838,7 +851,13 @@ class ServerArgs:
"--output-path",
type=str,
default=ServerArgs.output_path,
help="Directory path to save generated images/videos",
help='Directory path to save generated images/videos. Set to "" to disable persistent saving.',
)
parser.add_argument(
"--input-save-path",
type=str,
default=ServerArgs.input_save_path,
help='Directory path to save uploaded input images/videos. Set to "" to disable persistent saving.',
)
# LoRA

View File

@@ -34,6 +34,24 @@ logger = init_logger(__name__)
T = TypeVar("T")
def expand_path_fields(obj) -> None:
"""In-place expanduser on all dataclass fields whose name ends with '_path' or '_paths'."""
eu = os.path.expanduser
for f in fields(obj):
v = getattr(obj, f.name)
if f.name.endswith("_path") and isinstance(v, str):
setattr(obj, f.name, eu(v))
elif f.name.endswith("_path") and isinstance(v, list):
setattr(obj, f.name, [eu(x) if isinstance(x, str) else x for x in v])
elif f.name.endswith("_paths") and isinstance(v, dict):
setattr(
obj,
f.name,
{k: eu(p) if isinstance(p, str) else p for k, p in v.items()},
)
# TODO(will): used to convert server_args.precision to torch.dtype. Find a
# cleaner way to do this.
PRECISION_TO_TYPE = {