From 9a0fd2ff0c830a30e8d7e67c493e928fa5fdbdae Mon Sep 17 00:00:00 2001 From: Mick Date: Mon, 20 Apr 2026 23:29:02 +0800 Subject: [PATCH] [diffusion] optimize: default to in-memory loading for URL/base64 image inputs (#23118) --- .../runtime/entrypoints/openai/image_api.py | 1 + .../runtime/entrypoints/openai/utils.py | 138 +++++++++++++----- .../runtime/entrypoints/openai/video_api.py | 20 ++- .../runtime/models/vision_utils.py | 18 +-- .../runtime/pipelines/diffusers_pipeline.py | 12 +- 5 files changed, 132 insertions(+), 57 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py index 3e02d1346..8e6697157 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py @@ -252,6 +252,7 @@ async def edits( input_path = await save_image_to_path( img, os.path.join(uploads_dir, f"{request_id}_{idx}_{filename}"), + prefer_remote_source=server_args.input_save_path is None, ) input_paths.append(input_path) except Exception as e: diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py index 29a944f59..d8a48c7d9 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py @@ -1,4 +1,5 @@ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo +import asyncio import base64 import os import re @@ -137,8 +138,15 @@ def build_sampling_params(request_id: str, **kwargs) -> SamplingParams: 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) +async def save_image_to_path( + image: Union[UploadFile, str], + target_path: str, + *, + prefer_remote_source: bool = False, +) -> str: + input_path = await _maybe_url_image( + image, target_path, prefer_remote_source=prefer_remote_source + ) if input_path is None: input_path = await _save_upload_to_path(image, target_path) return input_path @@ -153,16 +161,27 @@ async def _save_upload_to_path(upload: UploadFile, target_path: str) -> str: return target_path -async def _maybe_url_image(img_url: str, target_path: str) -> str | None: +async def _maybe_url_image( + img_url: str, + target_path: str, + *, + prefer_remote_source: bool = False, +) -> str | None: if not isinstance(img_url, str): return None if img_url.lower().startswith(("http://", "https://")): - # Download image from URL + # Only bypass persistence when the caller explicitly disables input saves. + # Otherwise keep the prefetch outside the measured server stages. + if prefer_remote_source: + return img_url + # download image from URL and persist on disk input_path = await _save_url_image_to_path(img_url, target_path) return input_path elif img_url.startswith("data:image"): - # encode image base64 url + if prefer_remote_source: + return img_url + # encode image base64 url and persist on disk input_path = await _save_base64_image_to_path(img_url, target_path) return input_path else: @@ -172,47 +191,92 @@ async def _maybe_url_image(img_url: str, target_path: str) -> str | None: async def _save_url_image_to_path(image_url: str, target_path: str) -> str: """Download image from URL and save to target path.""" + def _is_retryable_download_error(error: Exception) -> bool: + if isinstance(error, httpx.HTTPStatusError): + status_code = error.response.status_code + # Retry on rate limit and transient server-side failures. + return status_code == 429 or 500 <= status_code < 600 + # Retry on transient network/protocol issues. + return isinstance( + error, + ( + httpx.TimeoutException, + httpx.NetworkError, + httpx.RemoteProtocolError, + ), + ) + os.makedirs(os.path.dirname(target_path), exist_ok=True) + max_attempts = 3 + backoff_seconds = 0.2 + last_error: Exception | None = None + try: async with httpx.AsyncClient(follow_redirects=True) as client: - response = await client.get(image_url, timeout=10.0) - response.raise_for_status() + for attempt in range(1, max_attempts + 1): + try: + response = await client.get(image_url, timeout=10.0) + response.raise_for_status() - # Determine file extension from content type or URL after downloading - if not os.path.splitext(target_path)[1]: - content_type = response.headers.get("content-type", "").lower() + # Determine file extension from content type or URL after downloading + if not os.path.splitext(target_path)[1]: + content_type = response.headers.get("content-type", "").lower() - url_path = image_url.split("?")[0] - _, url_ext = os.path.splitext(url_path) - url_ext = url_ext.lower() + url_path = image_url.split("?")[0] + _, url_ext = os.path.splitext(url_path) + url_ext = url_ext.lower() - if url_ext in {".jpg", ".jpeg", ".png", ".webp", ".gif", ".bmp"}: - ext = ".jpg" if url_ext == ".jpeg" else url_ext - elif content_type.startswith("image/"): - if "jpeg" in content_type or "jpg" in content_type: - ext = ".jpg" - elif "png" in content_type: - ext = ".png" - elif "webp" in content_type: - ext = ".webp" - else: - ext = ".jpg" # Default to jpg - elif content_type == "application/octet-stream": - # for octet-stream, if we couldn't get it from URL, default to jpg - ext = ".jpg" - else: - raise ValueError( - f"URL does not point to an image. Content-Type: {content_type}" + if url_ext in { + ".jpg", + ".jpeg", + ".png", + ".webp", + ".gif", + ".bmp", + }: + ext = ".jpg" if url_ext == ".jpeg" else url_ext + elif content_type.startswith("image/"): + if "jpeg" in content_type or "jpg" in content_type: + ext = ".jpg" + elif "png" in content_type: + ext = ".png" + elif "webp" in content_type: + ext = ".webp" + else: + ext = ".jpg" # Default to jpg + elif content_type == "application/octet-stream": + # for octet-stream, if we couldn't get it from URL, default to jpg + ext = ".jpg" + else: + raise ValueError( + f"URL does not point to an image. Content-Type: {content_type}" + ) + target_path = f"{target_path}{ext}" + + with open(target_path, "wb") as f: + f.write(response.content) + + return target_path + except Exception as e: + last_error = e + if attempt == max_attempts or not _is_retryable_download_error(e): + raise + wait_s = backoff_seconds * (2 ** (attempt - 1)) + logger.warning( + "Retrying image download (%s/%s) for %s after %.1fs due to: %s", + attempt, + max_attempts, + image_url, + wait_s, + e, ) - target_path = f"{target_path}{ext}" - - with open(target_path, "wb") as f: - f.write(response.content) - - return target_path + await asyncio.sleep(wait_s) except Exception as e: - raise Exception(f"Failed to download image from URL: {str(e)}") + final_error = last_error or e + raise Exception( + f"Failed to download image from URL {image_url}: {str(final_error)}" + ) async def _save_base64_image_to_path(base64_data: str, target_path: str) -> str: diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py index abccf31bb..9798eff1b 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py @@ -108,7 +108,11 @@ def _video_job_from_sampling( async def _save_first_input_image( - image_sources, request_id: str, uploads_dir: str + image_sources, + request_id: str, + uploads_dir: str, + *, + prefer_remote_source: bool = False, ) -> str | None: """Save the first input image from a list of sources and return its path.""" image_list = merge_image_input_list(image_sources) @@ -120,7 +124,9 @@ async def _save_first_input_image( 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) + return await save_image_to_path( + image, target_path, prefer_remote_source=prefer_remote_source + ) async def _dispatch_job_async( @@ -228,7 +234,10 @@ async def create_video( ) try: input_path = await _save_first_input_image( - image_sources, request_id, uploads_dir + image_sources, + request_id, + uploads_dir, + prefer_remote_source=server_args.input_save_path is None, ) except Exception as e: raise HTTPException( @@ -303,7 +312,10 @@ async def create_video( if payload.get("reference_url"): try: input_path = await _save_first_input_image( - payload.get("reference_url"), request_id, uploads_dir + payload.get("reference_url"), + request_id, + uploads_dir, + prefer_remote_source=server_args.input_save_path is None, ) except Exception as e: raise HTTPException( diff --git a/python/sglang/multimodal_gen/runtime/models/vision_utils.py b/python/sglang/multimodal_gen/runtime/models/vision_utils.py index 2086170d5..d3610187e 100644 --- a/python/sglang/multimodal_gen/runtime/models/vision_utils.py +++ b/python/sglang/multimodal_gen/runtime/models/vision_utils.py @@ -5,6 +5,7 @@ import os import tempfile from collections.abc import Callable +from io import BytesIO from urllib.parse import unquote, urlparse import imageio @@ -15,6 +16,8 @@ import requests import torch from packaging import version +from sglang.srt.utils.common import get_image_bytes as srt_get_image_bytes + if version.parse(version.parse(PIL.__version__).base_version) >= version.parse("9.1.0"): PIL_INTERPOLATION = { "linear": PIL.Image.Resampling.BILINEAR, @@ -89,7 +92,7 @@ def normalize(images: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor: # adapted from diffusers.utils import load_image def load_image( - image: str | PIL.Image.Image, + image: str | bytes | PIL.Image.Image, convert_method: Callable[[PIL.Image.Image], PIL.Image.Image] | None = None, ) -> PIL.Image.Image: """ @@ -102,20 +105,17 @@ def load_image( A conversion method to apply to the image after loading it. When set to `None` the image will be converted "RGB". """ - if isinstance(image, str): - if image.startswith("http://") or image.startswith("https://"): - image = PIL.Image.open(requests.get(image, stream=True).raw) - elif os.path.isfile(image): + if isinstance(image, (str, bytes)): + if isinstance(image, str) and os.path.isfile(image): image = PIL.Image.open(image) else: - raise ValueError( - f"Incorrect path or URL. URLs must start with `http://` or `https://`, and {image} is not a valid path." - ) + # in-memory loading path + image = PIL.Image.open(BytesIO(srt_get_image_bytes(image))) elif isinstance(image, PIL.Image.Image): image = image else: raise ValueError( - "Incorrect format used for the image. Should be a URL linking to an image, a local path, or a PIL image." + "Incorrect format used for the image. Should be bytes, a URL, a local path, base64/data URL, or a PIL image." ) image = PIL.ImageOps.exif_transpose(image) diff --git a/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py index b0c5ea8ee..502239212 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py @@ -10,11 +10,9 @@ import argparse import inspect import re import warnings -from io import BytesIO from typing import Any import numpy as np -import requests import torch import torchvision.transforms as T from diffusers import DiffusionPipeline @@ -22,6 +20,9 @@ from PIL import Image from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.runtime.models.vision_utils import ( + load_image as load_vision_image, +) from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) @@ -337,11 +338,8 @@ class DiffusersExecutionStage(PipelineStage): batch.image_path = batch.image_path[0] try: - if batch.image_path.startswith(("http://", "https://")): - response = requests.get(batch.image_path, timeout=30) - response.raise_for_status() - return Image.open(BytesIO(response.content)).convert("RGB") - return Image.open(batch.image_path).convert("RGB") + image = load_vision_image(batch.image_path) + return image.convert("RGB") except Exception as e: logger.error("Failed to load image from %s: %s", batch.image_path, e) return None