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

[diffusion] feat: add cloud storage support for API (#14579)

Browse Source
This commit is contained in:
Ratish P
2026-01-16 07:59:38 +05:30
committed by GitHub
parent 0e86de7c0b
commit c81bad1bf7
7 changed files with 467 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -17,6 +17,7 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import (
ImageResponse,
ImageResponseData,
)
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,
@@ -136,39 +137,64 @@ async def generations(
)
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")
# 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,
{
"id": request_id,
"created_at": int(time.time()),
"file_path": save_file_path,
"file_path": None if cloud_url else save_file_path,
"url": cloud_url,
},
)
resp_format = (request.response_format or "b64_json").lower()
# 4. Return Response
if resp_format == "b64_json":
with open(save_file_path, "rb") as f:
b64 = base64.b64encode(f.read()).decode("utf-8")
response_kwargs = {
"data": [
ImageResponseData(
b64_json=b64,
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),
)
],
}
response_kwargs = add_common_data_to_response(
response_kwargs, request_id=request_id, result=result
)
return ImageResponse(**response_kwargs)
else:
# Return error, not supported
raise HTTPException(
status_code=400, detail="response_format=url is not supported"
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
)
return ImageResponse(**response_kwargs)
@router.post("/edits", response_model=ImageResponse)
async def edits(
@@ -248,38 +274,54 @@ async def edits(
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")
# 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,
{
"id": request_id,
"created_at": int(time.time()),
"file_path": save_file_path,
"file_path": None if cloud_url else save_file_path,
"url": cloud_url,
"input_image_paths": input_paths, # Store all input image paths
"num_input_images": len(input_paths),
},
)
# Default to b64_json to align with gpt-image-1 behavior in OpenAI examples
# 4. Return Response
if (response_format or "b64_json").lower() == "b64_json":
response_kwargs = {"data": []}
for save_file_path in save_file_path_list:
with open(save_file_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(save_file_path),
)
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:
url = f"/v1/images/{request_id}/content"
response_kwargs = {
"data": [
ImageResponseData(
url=url,
url=cloud_url if cloud_url else f"/v1/images/{request_id}/content",
revised_prompt=prompt,
file_path=os.path.abspath(save_file_path),
)
@@ -289,7 +331,6 @@ async def edits(
response_kwargs = add_common_data_to_response(
response_kwargs, request_id=request_id, result=result
)
return ImageResponse(**response_kwargs)
@@ -301,6 +342,12 @@ async def download_image_content(
if not item:
raise HTTPException(status_code=404, detail="Image not found")
if item.get("url"):
raise HTTPException(
status_code=400,
detail=f"Image has been uploaded to cloud storage. Please use the cloud URL: {item.get('url')}",
)
file_path = item.get("file_path")
if not file_path or not os.path.exists(file_path):
raise HTTPException(status_code=404, detail="Image is still being generated")

1
python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py
View File

@@ -55,6 +55,7 @@ class VideoResponse(BaseModel):
size: str = ""
seconds: str = "4"
quality: str = "standard"
url: Optional[str] = None
remixed_from_video_id: Optional[str] = None
completed_at: Optional[int] = None
expires_at: Optional[int] = None

107
python/sglang/multimodal_gen/runtime/entrypoints/openai/storage.py Normal file
View File

@@ -0,0 +1,107 @@
import asyncio
import os
from typing import Optional
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
class CloudStorage:
def __init__(self):
self.enabled = os.getenv("SGLANG_CLOUD_STORAGE_TYPE", "").lower() == "s3"
if not self.enabled:
return
try:
import boto3
except ImportError:
logger.error(
"boto3 is not installed. Please install it with `pip install boto3` to use cloud storage."
)
self.enabled = False
return
self.bucket_name = os.getenv("SGLANG_S3_BUCKET_NAME")
if not self.bucket_name:
self.enabled = False
return
endpoint_url = os.getenv("SGLANG_S3_ENDPOINT_URL") or None
region_name = os.getenv("SGLANG_S3_REGION_NAME") or None
self.client = boto3.client(
"s3",
aws_access_key_id=os.getenv("SGLANG_S3_ACCESS_KEY_ID"),
aws_secret_access_key=os.getenv("SGLANG_S3_SECRET_ACCESS_KEY"),
endpoint_url=endpoint_url,
region_name=region_name,
)
self.endpoint_url = endpoint_url
self.region_name = region_name
def is_enabled(self) -> bool:
return self.enabled
async def upload_file(self, local_path: str, destination_key: str) -> Optional[str]:
if not self.is_enabled():
return None
def _sync_upload():
"""Synchronous part of the upload to run in a thread."""
ext = os.path.splitext(local_path)[1].lower()
content_type = {
".png": "image/png",
".jpg": "image/jpeg",
".jpeg": "image/jpeg",
".webp": "image/webp",
".mp4": "video/mp4",
}.get(ext, "application/octet-stream")
# Use the client created once in __init__
self.client.upload_file(
local_path,
self.bucket_name,
destination_key,
ExtraArgs={"ContentType": content_type},
)
try:
# Offload the blocking I/O call to a thread executor
await asyncio.get_running_loop().run_in_executor(None, _sync_upload)
except Exception as e:
# If upload fails, log the error and return None for fallback
logger.error(f"Upload failed for {destination_key}: {e}")
return None
# Simplified URL generation with a default region
if self.endpoint_url:
url = (
f"{self.endpoint_url.rstrip('/')}/{self.bucket_name}/{destination_key}"
)
else:
region = self.region_name or "us-east-1"
url = f"https://{self.bucket_name}.s3.{region}.amazonaws.com/{destination_key}"
logger.info(f"Uploaded {local_path} to {url}")
return url
async def upload_and_cleanup(self, file_path: str) -> Optional[str]:
"""Helper to upload a file and delete the local copy if successful."""
if not self.is_enabled():
return None
key = os.path.basename(file_path)
url = await self.upload_file(file_path, key)
if url:
try:
# pass if removal fails
os.remove(file_path)
except OSError as e:
logger.warning(f"Failed to remove temporary file {file_path}: {e}")
return url
# Global instance
cloud_storage = CloudStorage()

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

@@ -27,6 +27,7 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import (
VideoListResponse,
VideoResponse,
)
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,
@@ -127,11 +128,19 @@ async def _dispatch_job_async(job_id: str, batch: Req) -> None:
from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client
try:
_, result = await process_generation_batch(async_scheduler_client, batch)
save_file_path_list, result = await process_generation_batch(
async_scheduler_client, batch
)
save_file_path = save_file_path_list[0]
cloud_url = await cloud_storage.upload_and_cleanup(save_file_path)
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,
}
update_fields = add_common_data_to_response(
update_fields, request_id=job_id, result=result
@@ -334,6 +343,12 @@ async def download_video_content(
if not job:
raise HTTPException(status_code=404, detail="Video not found")
if job.get("url"):
raise HTTPException(
status_code=400,
detail=f"Video has been uploaded to cloud storage. Please use the cloud URL: {job.get('url')}",
)
file_path = job.get("file_path")
if not file_path or not os.path.exists(file_path):
raise HTTPException(status_code=404, detail="Generation is still in-progress")