[diffusion] optimize: default to in-memory loading for URL/base64 image inputs (#23118)

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:

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:
+async def save_image_to_path(
-    input_path = await _maybe_url_image(image, target_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)
+            for attempt in range(1, max_attempts + 1):
-            response.raise_for_status()
+                try:
+                    response = await client.get(image_url, timeout=10.0)
+                    response.raise_for_status()
 
-            # Determine file extension from content type or URL after downloading
+                    # Determine file extension from content type or URL after downloading
-            if not os.path.splitext(target_path)[1]:
+                    if not os.path.splitext(target_path)[1]:
-                content_type = response.headers.get("content-type", "").lower()
+                        content_type = response.headers.get("content-type", "").lower()
 
-                url_path = image_url.split("?")[0]
+                        url_path = image_url.split("?")[0]
-                _, url_ext = os.path.splitext(url_path)
+                        _, url_ext = os.path.splitext(url_path)
-                url_ext = url_ext.lower()
+                        url_ext = url_ext.lower()
 
-                if url_ext in {".jpg", ".jpeg", ".png", ".webp", ".gif", ".bmp"}:
+                        if url_ext in {
-                    ext = ".jpg" if url_ext == ".jpeg" else url_ext
+                            ".jpg",
-                elif content_type.startswith("image/"):
+                            ".jpeg",
-                    if "jpeg" in content_type or "jpg" in content_type:
+                            ".png",
-                        ext = ".jpg"
+                            ".webp",
-                    elif "png" in content_type:
+                            ".gif",
-                        ext = ".png"
+                            ".bmp",
-                    elif "webp" in content_type:
+                        }:
-                        ext = ".webp"
+                            ext = ".jpg" if url_ext == ".jpeg" else url_ext
-                    else:
+                        elif content_type.startswith("image/"):
-                        ext = ".jpg"  # Default to jpg
+                            if "jpeg" in content_type or "jpg" in content_type:
-                elif content_type == "application/octet-stream":
+                                ext = ".jpg"
-                    # for octet-stream, if we couldn't get it from URL, default to jpg
+                            elif "png" in content_type:
-                    ext = ".jpg"
+                                ext = ".png"
-                else:
+                            elif "webp" in content_type:
-                    raise ValueError(
+                                ext = ".webp"
-                        f"URL does not point to an image. Content-Type: {content_type}"
+                            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}"
+                    await asyncio.sleep(wait_s)
-
-            with open(target_path, "wb") as f:
-                f.write(response.content)
-
-            return target_path
     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:

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(

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 isinstance(image, (str, bytes)):
-        if image.startswith("http://") or image.startswith("https://"):
+        if isinstance(image, str) and os.path.isfile(image):
-            image = PIL.Image.open(requests.get(image, stream=True).raw)
-        elif os.path.isfile(image):
             image = PIL.Image.open(image)
         else:
-            raise ValueError(
+            # in-memory loading path
-                f"Incorrect path or URL. URLs must start with `http://` or `https://`, and {image} is not a valid 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)

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://")):
+            image = load_vision_image(batch.image_path)
-                response = requests.get(batch.image_path, timeout=30)
+            return image.convert("RGB")
-                response.raise_for_status()
-                return Image.open(BytesIO(response.content)).convert("RGB")
-            return Image.open(batch.image_path).convert("RGB")
         except Exception as e:
             logger.error("Failed to load image from %s: %s", batch.image_path, e)
             return None