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

This commit is contained in:
Mick
2026-04-20 23:29:02 +08:00
committed by GitHub
parent 0be6ab04dd
commit 9a0fd2ff0c
5 changed files with 132 additions and 57 deletions
@@ -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:
@@ -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:
@@ -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(
@@ -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)
@@ -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