chore: import upstream snapshot with attribution

python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py

@@ -0,0 +1,452 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+import asyncio
+import inspect
+import json
+import os
+import shutil
+import tempfile
+import time
+from contextlib import contextmanager
+from typing import Any, Generator, List, Optional, Union
+
+import httpx
+from fastapi import HTTPException, UploadFile
+
+from sglang.multimodal_gen.configs.sample.sampling_params import (
+    DataType,
+    SamplingParams,
+)
+from sglang.multimodal_gen.runtime.entrypoints.utils import (
+    ListLorasReq,
+    MergeLoraWeightsReq,
+    SetLoraReq,
+    ShutdownReq,
+    UnmergeLoraWeightsReq,
+    format_lora_message,
+    save_outputs,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
+from sglang.multimodal_gen.runtime.scheduler_client import AsyncSchedulerClient
+from sglang.multimodal_gen.runtime.server_args import get_global_server_args
+from sglang.multimodal_gen.runtime.utils.common import parse_size
+from sglang.multimodal_gen.runtime.utils.image_io import save_base64_image_to_path
+from sglang.multimodal_gen.runtime.utils.logging_utils import (
+    init_logger,
+    log_batch_completion,
+    log_generation_timer,
+)
+from sglang.multimodal_gen.runtime.utils.trace_wrapper import trace_req
+
+# re-export LoRA protocol types for backward compatibility
+__all__ = [
+    "SetLoraReq",
+    "MergeLoraWeightsReq",
+    "UnmergeLoraWeightsReq",
+    "ListLorasReq",
+    "ShutdownReq",
+    "format_lora_message",
+]
+
+logger = init_logger(__name__)
+
+OUTPUT_QUALITY_MAPPER = {"maximum": 100, "high": 90, "medium": 55, "low": 35}
+DEFAULT_FPS = 24
+DEFAULT_VIDEO_SECONDS = 4
+
+
+def _bad_request(message: str) -> HTTPException:
+    return HTTPException(status_code=400, detail=message)
+
+
+def _parse_size_or_raise(size: str) -> tuple[int, int]:
+    width, height = parse_size(size)
+    if width is None or height is None or width <= 0 or height <= 0:
+        raise _bad_request("size must be formatted as positive WIDTHxHEIGHT")
+    return width, height
+
+
+def _validate_positive_int(kwargs: dict[str, Any], name: str) -> None:
+    value = kwargs.get(name)
+    if value is not None and int(value) <= 0:
+        raise _bad_request(f"{name} must be positive")
+
+
+def flatten_extra_params(payload: Any) -> dict[str, Any]:
+    """Promote vLLM-Omni-style extra_params into regular request fields."""
+    if not isinstance(payload, dict):
+        return {}
+
+    extra_params = payload.pop("extra_params", None)
+    if isinstance(extra_params, str):
+        try:
+            extra_params = json.loads(extra_params)
+        except Exception:
+            extra_params = None
+    if not isinstance(extra_params, dict):
+        if "guardrails" in payload:
+            payload.setdefault("use_guardrails", payload["guardrails"])
+        return payload
+
+    for key, value in extra_params.items():
+        payload.setdefault(key, value)
+    if "guardrails" in extra_params:
+        payload.setdefault("use_guardrails", extra_params["guardrails"])
+
+    return payload
+
+
+@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 choose_output_image_ext(
+    output_format: Optional[str], background: Optional[str]
+) -> str:
+    fmt = (output_format or "").lower()
+    if fmt in {"png", "webp", "jpeg", "jpg"}:
+        return "jpg" if fmt == "jpeg" else fmt
+    if (background or "auto").lower() == "transparent":
+        return "png"
+    return "jpg"
+
+
+def build_sampling_params(request_id: str, **kwargs) -> SamplingParams:
+    """Build SamplingParams from request parameters.
+
+    Handles size parsing, output_quality resolution, and None filtering before
+    delegating to SamplingParams.from_user_sampling_params_args. Callers pass
+    only the parameters they have; None values are stripped automatically so
+    that SamplingParams defaults apply.
+    """
+    server_args = get_global_server_args()
+
+    # pop HTTP-layer params that aren't SamplingParams fields
+    output_quality = kwargs.pop("output_quality", None)
+
+    has_explicit_compression = kwargs.get("output_compression") is not None
+
+    # parse "WxH" size string if provided
+    size = kwargs.pop("size", None)
+    if size:
+        w, h = _parse_size_or_raise(size)
+        # treat None dimensions as unset so parsed size can fill them
+        if kwargs.get("width") is None:
+            kwargs["width"] = w
+        if kwargs.get("height") is None:
+            kwargs["height"] = h
+
+    for name in (
+        "width",
+        "height",
+        "num_frames",
+        "num_inference_steps",
+        "num_outputs_per_prompt",
+    ):
+        _validate_positive_int(kwargs, name)
+
+    # filter out None values to let SamplingParams defaults apply
+    kwargs = {k: v for k, v in kwargs.items() if v is not None}
+    kwargs.setdefault("save_output", True)
+
+    sampling_params = SamplingParams.from_user_sampling_params_args(
+        model_path=server_args.model_path,
+        server_args=server_args,
+        request_id=request_id,
+        **kwargs,
+    )
+
+    # resolve output_quality → output_compression with the correct data_type.
+    # SamplingParams.__post_init__ may have resolved with the wrong data_type
+    # (default VIDEO) before _adjust() set the correct one.
+    if not has_explicit_compression and output_quality is not None:
+        resolved = adjust_output_quality(output_quality, sampling_params.data_type)
+        if resolved is not None:
+            sampling_params.output_compression = resolved
+
+    return sampling_params
+
+
+async def save_image_to_path(
+    image: Union[UploadFile, bytes, 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
+
+
+# Helpers
+async def _save_upload_to_path(
+    upload: Union[UploadFile, bytes], target_path: str
+) -> str:
+    os.makedirs(os.path.dirname(target_path), exist_ok=True)
+    if isinstance(upload, bytes):
+        content = upload
+    elif isinstance(upload, (bytearray, memoryview)):
+        content = bytes(upload)
+    else:
+        read = getattr(upload, "read", None)
+        if not callable(read):
+            raise TypeError(f"Unsupported image upload type: {type(upload).__name__}")
+        content = read()
+        if inspect.isawaitable(content):
+            content = await content
+        if isinstance(content, (bytearray, memoryview)):
+            content = bytes(content)
+        if not isinstance(content, bytes):
+            raise TypeError(
+                f"Image upload read() returned {type(content).__name__}, expected bytes"
+            )
+    with open(target_path, "wb") as f:
+        f.write(content)
+    return target_path
+
+
+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://")):
+        # 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"):
+        if prefer_remote_source:
+            return img_url
+        # encode image base64 url and persist on disk
+        input_path = save_base64_image_to_path(img_url, target_path)
+        return input_path
+    else:
+        raise ValueError("Unsupported image url format")
+
+
+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:
+            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()
+
+                        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}"
+                            )
+                        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,
+                    )
+                    await asyncio.sleep(wait_s)
+    except Exception as e:
+        final_error = last_error or e
+        raise Exception(
+            f"Failed to download image from URL {image_url}: {str(final_error)}"
+        )
+
+
+async def process_generation_batch(
+    scheduler_client: AsyncSchedulerClient,
+    batch,
+) -> tuple[list[str], OutputBatch]:
+    total_start_time = time.perf_counter()
+    with trace_req(batch.trace_ctx), log_generation_timer(logger, batch.prompt):
+        result = await scheduler_client.forward([batch])
+
+        if (
+            result.output is None
+            and result.output_file_paths is None
+            and result.raw_frame_batches is None
+        ):
+            error_msg = result.error or "Unknown error"
+            raise RuntimeError(
+                f"Model generation returned no output. Error from scheduler: {error_msg}"
+            )
+
+        save_file_path_list = []
+        if result.output_file_paths:
+            save_file_path_list = result.output_file_paths
+        elif result.output is not None:
+            num_outputs = len(result.output)
+            save_file_path_list = save_outputs(
+                result.output,
+                batch.data_type,
+                batch.fps,
+                batch.save_output,
+                lambda idx: str(batch.output_file_path(num_outputs, idx)),
+                audio=result.audio,
+                audio_sample_rate=result.audio_sample_rate,
+                output_compression=batch.output_compression,
+                enable_frame_interpolation=batch.enable_frame_interpolation,
+                frame_interpolation_exp=batch.frame_interpolation_exp,
+                frame_interpolation_scale=batch.frame_interpolation_scale,
+                frame_interpolation_model_path=batch.frame_interpolation_model_path,
+                enable_upscaling=batch.enable_upscaling,
+                upscaling_model_path=batch.upscaling_model_path,
+                upscaling_scale=batch.upscaling_scale,
+            )
+
+    total_time = time.perf_counter() - total_start_time
+    if get_global_server_args().batching_max_size > 1:
+        log_batch_completion(
+            logger,
+            len(save_file_path_list),
+            total_time,
+        )
+
+    if result.peak_memory_mb and result.peak_memory_mb > 0:
+        logger.info(f"Peak memory usage: {result.peak_memory_mb:.2f} MB")
+
+    return save_file_path_list, result
+
+
+def merge_image_input_list(*inputs: Union[List, Any, None]) -> List:
+    """
+    Merge multiple image input sources into a single list.
+
+    This function handles both single items and lists of items, merging them
+    into a single flattened list. Useful for processing images, URLs, or other
+    multimedia inputs that can come as either single items or lists.
+
+    Args:
+        *inputs: Variable number of inputs, each can be None, single item, or list
+
+    Returns:
+        List: Flattened list of all non-None inputs
+
+    Example:
+        >>> merge_image_input_list(["img1", "img2"], "img3", None)
+        ["img1", "img2", "img3"]
+    """
+    result = []
+    for input_item in inputs:
+        if input_item is not None:
+            if isinstance(input_item, list):
+                result.extend(input_item)
+            else:
+                result.append(input_item)
+    return result
+
+
+def add_common_data_to_response(
+    response: dict, request_id: str, result: OutputBatch
+) -> dict:
+    if result.peak_memory_mb and result.peak_memory_mb > 0:
+        response["peak_memory_mb"] = result.peak_memory_mb
+
+    if result.metrics and result.metrics.total_duration_s > 0:
+        response["inference_time_s"] = result.metrics.total_duration_s
+
+    response["id"] = request_id
+
+    if result.action_pred is not None:
+        t = result.action_pred
+        response["action"] = {
+            "data": t.tolist(),
+            "shape": list(t.shape),
+            "dtype": str(t.dtype).replace("torch.", ""),
+            "raw_action_dim": result.action_raw_action_dim,
+            "action_mode": result.action_mode,
+            "domain_id": result.action_domain_id,
+        }
+
+    return response
+
+
+def adjust_output_quality(output_quality: str, data_type: DataType = None) -> int:
+    if output_quality == "default":
+        return 50 if data_type == DataType.VIDEO else 75
+    return OUTPUT_QUALITY_MAPPER.get(output_quality, None)