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

[diffusion] feat: implement upscaling (#19723)

Browse Source
This commit is contained in:
Junhao Liu
2026-03-08 11:06:40 -07:00
committed by GitHub
parent b77dd41db0
commit 7662b8b919
17 changed files with 1005 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
python/sglang/multimodal_gen/configs/sample/sampling_params.py
View File

@@ -113,6 +113,13 @@ class SamplingParams:
None # local dir or HF repo ID with flownet.pkl (default: elfgum/RIFE-4.22.lite)
)
# Upscaling
enable_upscaling: bool = False
upscaling_model_path: str | None = (
None # local .pth, HF repo ID, or repo_id:filename (default: ai-forever/Real-ESRGAN)
)
upscaling_scale: int = 4
# Batch info
num_outputs_per_prompt: int = 1
seed: int = 42
@@ -875,7 +882,30 @@ class SamplingParams:
type=str,
default=SamplingParams.frame_interpolation_model_path,
help="Local directory or HuggingFace repo ID containing RIFE flownet.pkl weights "
"(default: elfgum/RIFE-4.22.lite, downloaded automatically).",
"(default: elfgum/RIFE-4.22.lite, downloaded automatically). "
"Only RIFE 4.22.lite architecture is supported; other RIFE versions or "
"frame interpolation models are not compatible.",
)
parser.add_argument(
"--enable-upscaling",
action="store_true",
help="Enable post-generation upscaling using Real-ESRGAN.",
)
parser.add_argument(
"--upscaling-model-path",
type=str,
default=SamplingParams.upscaling_model_path,
help="Local .pth file, HuggingFace repo ID, or repo_id:filename for Real-ESRGAN weights "
"(default: ai-forever/Real-ESRGAN with RealESRGAN_x4.pth). "
"Only RRDBNet (e.g. RealESRGAN_x4plus) and SRVGGNetCompact (e.g. realesr-animevideov3) "
"architectures are supported; other super-resolution models are not compatible. "
"Use 'repo_id:filename' to specify a custom weight file from a HF repo.",
)
parser.add_argument(
"--upscaling-scale",
type=int,
default=SamplingParams.upscaling_scale,
help="Upscaling factor (default: 4).",
)
return parser

3
python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py
View File

@@ -277,6 +277,9 @@ class DiffGenerator:
frame_interpolation_exp=req.frame_interpolation_exp,
frame_interpolation_scale=req.frame_interpolation_scale,
frame_interpolation_model_path=req.frame_interpolation_model_path,
enable_upscaling=req.enable_upscaling,
upscaling_model_path=req.upscaling_model_path,
upscaling_scale=req.upscaling_scale,
)
for idx in range(len(samples_out)):

3
python/sglang/multimodal_gen/runtime/entrypoints/http_server.py
View File

@@ -203,6 +203,9 @@ async def forward_to_scheduler(
frame_interpolation_exp=sp.frame_interpolation_exp,
frame_interpolation_scale=sp.frame_interpolation_scale,
frame_interpolation_model_path=sp.frame_interpolation_model_path,
enable_upscaling=sp.enable_upscaling,
upscaling_model_path=sp.upscaling_model_path,
upscaling_scale=sp.upscaling_scale,
)
if hasattr(response, "model_dump"):

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

@@ -133,6 +133,9 @@ async def generations(
enable_teacache=request.enable_teacache,
output_compression=request.output_compression,
output_quality=request.output_quality,
enable_upscaling=request.enable_upscaling,
upscaling_model_path=request.upscaling_model_path,
upscaling_scale=request.upscaling_scale,
)
batch = prepare_request(
server_args=server_args,
@@ -207,6 +210,9 @@ async def edits(
output_quality: Optional[str] = Form("default"),
output_compression: Optional[int] = Form(None),
enable_teacache: Optional[bool] = Form(False),
enable_upscaling: Optional[bool] = Form(False),
upscaling_model_path: Optional[str] = Form(None),
upscaling_scale: Optional[int] = Form(4),
num_frames: int = Form(1),
):
request_id = generate_request_id()
@@ -262,6 +268,9 @@ async def edits(
num_frames=num_frames,
output_compression=output_compression,
output_quality=output_quality,
enable_upscaling=enable_upscaling,
upscaling_model_path=upscaling_model_path,
upscaling_scale=upscaling_scale,
)
batch = prepare_request(
server_args=server_args,

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

@@ -48,6 +48,10 @@ class ImageGenerationsRequest(BaseModel):
output_quality: Optional[str] = "default"
output_compression: Optional[int] = None
enable_teacache: Optional[bool] = False
# Upscaling
enable_upscaling: Optional[bool] = False
upscaling_model_path: Optional[str] = None
upscaling_scale: Optional[int] = 4
diffusers_kwargs: Optional[Dict[str, Any]] = None # kwargs for diffusers backend
@@ -97,6 +101,10 @@ class VideoGenerationsRequest(BaseModel):
frame_interpolation_exp: Optional[int] = 1 # 1=2×, 2=4×
frame_interpolation_scale: Optional[float] = 1.0
frame_interpolation_model_path: Optional[str] = None
# Upscaling
enable_upscaling: Optional[bool] = False
upscaling_model_path: Optional[str] = None
upscaling_scale: Optional[int] = 4
output_quality: Optional[str] = "default"
output_compression: Optional[int] = None
output_path: Optional[str] = None

3
python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py
View File

@@ -286,6 +286,9 @@ async def process_generation_batch(
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

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

@@ -75,6 +75,9 @@ def _build_video_sampling_params(request_id: str, request: VideoGenerationsReque
frame_interpolation_exp=request.frame_interpolation_exp,
frame_interpolation_scale=request.frame_interpolation_scale,
frame_interpolation_model_path=request.frame_interpolation_model_path,
enable_upscaling=request.enable_upscaling,
upscaling_model_path=request.upscaling_model_path,
upscaling_scale=request.upscaling_scale,
output_path=request.output_path,
output_compression=request.output_compression,
output_quality=request.output_quality,
@@ -181,6 +184,9 @@ async def create_video(
frame_interpolation_exp: Optional[int] = Form(1),
frame_interpolation_scale: Optional[float] = Form(1.0),
frame_interpolation_model_path: Optional[str] = Form(None),
enable_upscaling: Optional[bool] = Form(False),
upscaling_model_path: Optional[str] = Form(None),
upscaling_scale: Optional[int] = Form(4),
output_quality: Optional[str] = Form("default"),
output_compression: Optional[int] = Form(None),
extra_body: Optional[str] = Form(None),
@@ -256,6 +262,9 @@ async def create_video(
frame_interpolation_exp=frame_interpolation_exp,
frame_interpolation_scale=frame_interpolation_scale,
frame_interpolation_model_path=frame_interpolation_model_path,
enable_upscaling=enable_upscaling,
upscaling_model_path=upscaling_model_path,
upscaling_scale=upscaling_scale,
output_compression=output_compression,
output_quality=output_quality,
**(

21
python/sglang/multimodal_gen/runtime/entrypoints/utils.py
View File

@@ -348,6 +348,9 @@ def save_outputs(
frame_interpolation_exp: int = 1,
frame_interpolation_scale: float = 1.0,
frame_interpolation_model_path: Optional[str] = None,
enable_upscaling: bool = False,
upscaling_model_path: Optional[str] = None,
upscaling_scale: int = 4,
) -> list[str]:
"""Save outputs to files and return the list of file paths."""
output_paths: list[str] = []
@@ -369,6 +372,9 @@ def save_outputs(
frame_interpolation_exp=frame_interpolation_exp,
frame_interpolation_scale=frame_interpolation_scale,
frame_interpolation_model_path=frame_interpolation_model_path,
enable_upscaling=enable_upscaling,
upscaling_model_path=upscaling_model_path,
upscaling_scale=upscaling_scale,
)
if samples_out is not None:
@@ -401,6 +407,9 @@ def post_process_sample(
frame_interpolation_exp: int = 1,
frame_interpolation_scale: float = 1.0,
frame_interpolation_model_path: Optional[str] = None,
enable_upscaling: bool = False,
upscaling_model_path: Optional[str] = None,
upscaling_scale: int = 4,
):
"""
Process sample output, optionally interpolate video frames, and save.
@@ -456,7 +465,17 @@ def post_process_sample(
)
fps = fps * multiplier
# 3. Save outputs if requested
# 3. Upscaling (images and videos)
if enable_upscaling and frames:
from sglang.multimodal_gen.runtime.postprocess import upscale_frames
frames = upscale_frames(
frames,
model_path=upscaling_model_path,
scale=upscaling_scale,
)
# 4. Save outputs if requested
if save_output:
if save_file_path:
os.makedirs(os.path.dirname(save_file_path), exist_ok=True)

3
python/sglang/multimodal_gen/runtime/managers/gpu_worker.py
View File

@@ -269,6 +269,9 @@ class GPUWorker:
frame_interpolation_exp=req.frame_interpolation_exp,
frame_interpolation_scale=req.frame_interpolation_scale,
frame_interpolation_model_path=req.frame_interpolation_model_path,
enable_upscaling=req.enable_upscaling,
upscaling_model_path=req.upscaling_model_path,
upscaling_scale=req.upscaling_scale,
)
output_batch.output_file_paths = output_paths
output_batch.output = None

13
python/sglang/multimodal_gen/runtime/postprocess/__init__.py
View File

@@ -1,9 +1,18 @@
# SPDX-License-Identifier: Apache-2.0
"""Frame interpolation support for SGLang diffusion pipelines."""
"""Frame interpolation and upscaling support for SGLang diffusion pipelines."""
from sglang.multimodal_gen.runtime.postprocess.realesrgan_upscaler import (
ImageUpscaler,
upscale_frames,
)
from sglang.multimodal_gen.runtime.postprocess.rife_interpolator import (
FrameInterpolator,
interpolate_video_frames,
)
__all__ = ["FrameInterpolator", "interpolate_video_frames"]
__all__ = [
"FrameInterpolator",
"interpolate_video_frames",
"ImageUpscaler",
"upscale_frames",
]

484
python/sglang/multimodal_gen/runtime/postprocess/realesrgan_upscaler.py Normal file
View File

@@ -0,0 +1,484 @@
# SPDX-License-Identifier: Apache-2.0
"""
Real-ESRGAN upscaling for SGLang diffusion pipelines.
Real-ESRGAN model code is vendored and adapted from:
- https://github.com/xinntao/Real-ESRGAN (BSD-3-Clause License)
Copyright (c) 2021 xinntao
The ImageUpscaler wrapper and integration code are original work.
"""
import math
import os
from typing import Optional
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
# Default HuggingFace repo and filename for Real-ESRGAN weights
_DEFAULT_REALESRGAN_HF_REPO = "ai-forever/Real-ESRGAN"
_DEFAULT_REALESRGAN_FILENAME = "RealESRGAN_x4.pth"
# Module-level cache: model_path -> UpscalerModel instance
_MODEL_CACHE: dict[str, "UpscalerModel"] = {}
# ---------------------------------------------------------------------------
# Vendored Real-ESRGAN architecture code
# (SRVGGNetCompact, ResidualDenseBlock, RRDB, RRDBNet)
# ---------------------------------------------------------------------------
class SRVGGNetCompact(nn.Module):
"""Compact VGG-style network for super resolution.
Corresponds to ``realesr-animevideov3`` and ``realesr-general-x4v3``.
Reference: xinntao/Real-ESRGAN (BSD-3-Clause).
"""
def __init__(
self,
num_in_ch: int = 3,
num_out_ch: int = 3,
num_feat: int = 64,
num_conv: int = 16,
upscale: int = 4,
act_type: str = "prelu",
):
super().__init__()
self.num_in_ch = num_in_ch
self.num_out_ch = num_out_ch
self.num_feat = num_feat
self.num_conv = num_conv
self.upscale = upscale
self.act_type = act_type
self.body = nn.ModuleList()
# first conv
self.body.append(nn.Conv2d(num_in_ch, num_feat, 3, 1, 1))
# first activation
self.body.append(self._make_act(act_type, num_feat))
# body convs + activations
for _ in range(num_conv):
self.body.append(nn.Conv2d(num_feat, num_feat, 3, 1, 1))
self.body.append(self._make_act(act_type, num_feat))
# last conv: maps to out_ch * upscale^2 for pixel shuffle
self.body.append(nn.Conv2d(num_feat, num_out_ch * upscale * upscale, 3, 1, 1))
self.upsampler = nn.PixelShuffle(upscale)
@staticmethod
def _make_act(act_type: str, num_feat: int) -> nn.Module:
if act_type == "relu":
return nn.ReLU(inplace=True)
elif act_type == "prelu":
return nn.PReLU(num_parameters=num_feat)
elif act_type == "leakyrelu":
return nn.LeakyReLU(negative_slope=0.1, inplace=True)
else:
raise ValueError(f"Unsupported activation type: {act_type}")
def forward(self, x: torch.Tensor) -> torch.Tensor:
out = x
for layer in self.body:
out = layer(out)
out = self.upsampler(out)
# residual addition with nearest upsampled input
base = F.interpolate(x, scale_factor=self.upscale, mode="nearest")
return out + base
class ResidualDenseBlock(nn.Module):
"""Residual Dense Block used in RRDB (RealESRGAN_x4plus)."""
def __init__(self, num_feat: int = 64, num_grow_ch: int = 32):
super().__init__()
self.conv1 = nn.Conv2d(num_feat, num_grow_ch, 3, 1, 1)
self.conv2 = nn.Conv2d(num_feat + num_grow_ch, num_grow_ch, 3, 1, 1)
self.conv3 = nn.Conv2d(num_feat + 2 * num_grow_ch, num_grow_ch, 3, 1, 1)
self.conv4 = nn.Conv2d(num_feat + 3 * num_grow_ch, num_grow_ch, 3, 1, 1)
self.conv5 = nn.Conv2d(num_feat + 4 * num_grow_ch, num_feat, 3, 1, 1)
self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x1 = self.lrelu(self.conv1(x))
x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
return x5 * 0.2 + x
class RRDB(nn.Module):
"""Residual in Residual Dense Block."""
def __init__(self, num_feat: int, num_grow_ch: int = 32):
super().__init__()
self.rdb1 = ResidualDenseBlock(num_feat, num_grow_ch)
self.rdb2 = ResidualDenseBlock(num_feat, num_grow_ch)
self.rdb3 = ResidualDenseBlock(num_feat, num_grow_ch)
def forward(self, x: torch.Tensor) -> torch.Tensor:
out = self.rdb1(x)
out = self.rdb2(out)
out = self.rdb3(out)
return out * 0.2 + x
class RRDBNet(nn.Module):
"""RRDB network for RealESRGAN_x4plus (heavier, higher quality for photos)."""
def __init__(
self,
num_in_ch: int = 3,
num_out_ch: int = 3,
scale: int = 4,
num_feat: int = 64,
num_block: int = 23,
num_grow_ch: int = 32,
):
super().__init__()
self.scale = scale
in_ch = num_in_ch
if scale == 2:
in_ch = num_in_ch * 4
elif scale == 1:
in_ch = num_in_ch * 16
self.conv_first = nn.Conv2d(in_ch, num_feat, 3, 1, 1)
self.body = nn.Sequential(
*[RRDB(num_feat, num_grow_ch) for _ in range(num_block)]
)
self.conv_body = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
# upsample
self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
def forward(self, x: torch.Tensor) -> torch.Tensor:
if self.scale == 2:
feat = F.pixel_unshuffle(x, 2)
elif self.scale == 1:
feat = F.pixel_unshuffle(x, 4)
else:
feat = x
feat = self.conv_first(feat)
body_feat = self.conv_body(self.body(feat))
feat = feat + body_feat
feat = self.lrelu(
self.conv_up1(F.interpolate(feat, scale_factor=2, mode="nearest"))
)
feat = self.lrelu(
self.conv_up2(F.interpolate(feat, scale_factor=2, mode="nearest"))
)
return self.conv_last(self.lrelu(self.conv_hr(feat)))
# ---------------------------------------------------------------------------
# Architecture auto-detection
# ---------------------------------------------------------------------------
def _build_net_from_state_dict(state_dict: dict) -> nn.Module:
"""Detect architecture from checkpoint keys and return an unloaded network."""
if "conv_first.weight" in state_dict:
# RRDBNet (e.g., RealESRGAN_x4plus)
num_feat = state_dict["conv_first.weight"].shape[0]
num_block = sum(
1
for k in state_dict
if k.startswith("body.") and k.endswith(".rdb1.conv1.weight")
)
num_grow_ch = state_dict["body.0.rdb1.conv1.weight"].shape[0]
logger.info(
"Detected RRDBNet: num_feat=%d, num_block=%d, num_grow_ch=%d",
num_feat,
num_block,
num_grow_ch,
)
return RRDBNet(
num_in_ch=3,
num_out_ch=3,
scale=4,
num_feat=num_feat,
num_block=num_block,
num_grow_ch=num_grow_ch,
)
else:
# SRVGGNetCompact (e.g., realesr-animevideov3)
num_feat = state_dict["body.0.weight"].shape[0]
# body layout: [first_conv, first_act, (conv, act)*num_conv, last_conv]
# count 4-D weight tensors = first_conv + loop_convs + last_conv = num_conv + 2
conv_keys = sorted(
[
k
for k in state_dict
if k.startswith("body.")
and k.endswith(".weight")
and state_dict[k].dim() == 4
],
key=lambda k: int(k.split(".")[1]),
)
num_conv = len(conv_keys) - 2 # subtract first and last
# upscale from last conv output channels: out_ch = num_out_ch * upscale^2
last_out_ch = state_dict[conv_keys[-1]].shape[0]
upscale = int(math.sqrt(last_out_ch / 3))
logger.info(
"Detected SRVGGNetCompact: num_feat=%d, num_conv=%d, upscale=%d",
num_feat,
num_conv,
upscale,
)
return SRVGGNetCompact(
num_in_ch=3,
num_out_ch=3,
num_feat=num_feat,
num_conv=num_conv,
upscale=upscale,
act_type="prelu",
)
# ---------------------------------------------------------------------------
# UpscalerModel
# ---------------------------------------------------------------------------
class UpscalerModel:
"""Wraps a Real-ESRGAN network, provides load() and upscale() API."""
def __init__(self, net: nn.Module, scale: int):
self.net = net
self.scale = scale # the model's native upscaling factor (e.g. 4)
@property
def device(self) -> torch.device:
return next(self.net.parameters()).device
def upscale(self, frame: np.ndarray, outscale: float | None = None) -> np.ndarray:
"""Upscale a single HWC uint8 frame → HWC uint8 frame.
Args:
frame: Input HWC uint8 numpy array.
outscale: Desired final upscaling factor. If different from the
model's native scale, a cheap resize is applied after
the network output (same approach as the official
Real-ESRGAN ``inference_realesrgan.py --outscale``).
``None`` means use the model's native scale as-is.
"""
h, w = frame.shape[:2]
img = frame.astype(np.float32) / 255.0
img_t = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0).to(self.device)
with torch.no_grad():
out = self.net(img_t)
# If the desired outscale differs from the model's native scale,
# resize to (h * outscale, w * outscale).
if outscale is not None and outscale != self.scale:
target_h = int(h * outscale)
target_w = int(w * outscale)
out = F.interpolate(
out, size=(target_h, target_w), mode="bicubic", align_corners=False
)
out_np = out.squeeze(0).permute(1, 2, 0).clamp(0.0, 1.0).cpu().numpy()
return (out_np * 255.0).astype(np.uint8)
# ---------------------------------------------------------------------------
# ImageUpscaler public class
# ---------------------------------------------------------------------------
class ImageUpscaler:
"""
Lazy-loaded Real-ESRGAN upscaler.
Weights are downloaded and cached on first call to `.upscale()`.
Supports both SRVGGNetCompact (lightweight, default) and RRDBNet (heavier).
"""
def __init__(
self,
model_path: Optional[str] = None,
scale: int = 4,
half_precision: bool = False,
):
self._model_path = model_path
self._scale = scale
self._half_precision = half_precision
def _ensure_model_loaded(self) -> UpscalerModel:
"""Download/load Real-ESRGAN weights, detect arch, and cache globally."""
model_path = self._model_path or _DEFAULT_REALESRGAN_HF_REPO
# Resolve: local .pth pass-through, or HF repo → download single file
resolved_path = _resolve_model_path(model_path)
if resolved_path in _MODEL_CACHE:
return _MODEL_CACHE[resolved_path]
logger.info("Loading Real-ESRGAN weights from %s", resolved_path)
try:
state_dict = torch.load(
resolved_path, map_location="cpu", weights_only=True
)
except Exception as e:
raise RuntimeError(
f"Failed to load Real-ESRGAN checkpoint from '{resolved_path}'. "
f"The file may be corrupted or not a valid PyTorch checkpoint. "
f"Original error: {e}"
) from e
# Some checkpoints wrap weights under a 'params' or 'params_ema' key
if "params_ema" in state_dict:
state_dict = state_dict["params_ema"]
elif "params" in state_dict:
state_dict = state_dict["params"]
try:
net = _build_net_from_state_dict(state_dict)
net.load_state_dict(state_dict, strict=True)
except (RuntimeError, KeyError) as e:
raise RuntimeError(
f"Real-ESRGAN weight file '{resolved_path}' is not compatible "
f"with the supported architectures (SRVGGNetCompact / RRDBNet). "
f"Please ensure you are using a valid Real-ESRGAN checkpoint. "
f"Original error: {e}"
) from e
net.eval()
device = current_platform.get_local_torch_device()
if self._half_precision:
net = net.half()
net = net.to(device)
# Detect the model's native scale from network architecture
native_scale = 4 # sensible default
if hasattr(net, "upscale"):
native_scale = net.upscale
elif hasattr(net, "scale"):
native_scale = net.scale
model = UpscalerModel(net=net, scale=native_scale)
_MODEL_CACHE[resolved_path] = model
logger.info(
"Real-ESRGAN model loaded on device: %s (native_scale=%dx, outscale=%s)",
device,
native_scale,
f"{self._scale}x" if self._scale != native_scale else "native",
)
return model
def upscale(self, frames: list[np.ndarray]) -> list[np.ndarray]:
"""Upscale a list of HWC uint8 frames.
Uses the model's native scale for super-resolution, then resizes to
the desired ``outscale`` if it differs (cheap bicubic resize).
"""
if not frames:
return frames
model = self._ensure_model_loaded()
outscale = self._scale if self._scale != model.scale else None
return [model.upscale(frame, outscale=outscale) for frame in frames]
# ---------------------------------------------------------------------------
# HF download helper
# ---------------------------------------------------------------------------
def _resolve_model_path(model_path: str) -> str:
"""Return a local .pth file path.
Accepts:
- An existing local file path (pass-through).
- A HuggingFace ``repo_id`` → downloads the default weight file
(``RealESRGAN_x4.pth``).
- A HuggingFace ``repo_id:filename`` → downloads *filename* from *repo_id*,
allowing users to specify custom weight files hosted on HF.
"""
if os.path.isfile(model_path):
return model_path
# Parse optional "repo_id:filename" syntax; fall back to default filename.
if ":" in model_path and not model_path.startswith("/"):
repo_id, filename = model_path.split(":", 1)
else:
repo_id = model_path
filename = _DEFAULT_REALESRGAN_FILENAME
try:
from huggingface_hub import hf_hub_download
except ImportError as e:
raise ImportError(
"huggingface_hub is required to download Real-ESRGAN weights. "
"Install it with: pip install huggingface_hub"
) from e
logger.info(
"Downloading Real-ESRGAN weights from HF repo %s (file: %s)",
repo_id,
filename,
)
try:
local_path = hf_hub_download(
repo_id=repo_id,
filename=filename,
)
except Exception as e:
raise FileNotFoundError(
f"Failed to download Real-ESRGAN weights from HuggingFace repo "
f"'{repo_id}' (file: '{filename}'). If you are using a custom "
f"model, provide either a local .pth file path or use the "
f"'repo_id:filename' format (e.g. 'my-org/my-esrgan:weights.pth'). "
f"Original error: {e}"
) from e
return local_path
# ---------------------------------------------------------------------------
# Module-level convenience function
# ---------------------------------------------------------------------------
def upscale_frames(
frames: list[np.ndarray],
model_path: Optional[str] = None,
scale: int = 4,
half_precision: bool = False,
) -> list[np.ndarray]:
"""
Convenience wrapper around ImageUpscaler.
The model always runs at its native resolution (e.g. 4× for
``RealESRGAN_x4.pth``). If *scale* differs from the native factor,
a cheap bicubic resize is applied after the network output – the same
approach used by the official Real-ESRGAN ``--outscale`` flag.
Args:
frames: List of uint8 HWC numpy frames.
model_path: Local .pth file, HuggingFace repo ID, or
``repo_id:filename`` for a custom weight file.
None → default ``ai-forever/Real-ESRGAN`` with
``RealESRGAN_x4.pth``.
scale: Desired final upscaling factor (e.g. 2, 3, 4).
The 4× model is used internally; the output is
resized to match *scale* when it differs.
half_precision: Use fp16 inference (faster on supported GPUs).
Returns:
List of upscaled uint8 HWC numpy frames.
"""
upscaler = ImageUpscaler(
model_path=model_path, scale=scale, half_precision=half_precision
)
return upscaler.upscale(frames)

196
python/sglang/multimodal_gen/test/server/perf_baselines.json
View File

@@ -2136,6 +2136,202 @@
"expected_e2e_ms": 716.81,
"expected_avg_denoise_ms": 67.32,
"expected_median_denoise_ms": 74.1
},
"flux_2_image_t2i_upscaling_4x": {
"stages_ms": {
"InputValidationStage": 0.15,
"TextEncodingStage": 537.4,
"ImageVAEEncodingStage": 0.01,
"LatentPreparationStage": 1.11,
"TimestepPreparationStage": 39.57,
"DenoisingStage": 24738.35,
"DecodingStage": 14.04
},
"denoise_step_ms": {
"0": 70.71,
"1": 476.99,
"2": 502.01,
"3": 483.39,
"4": 500.65,
"5": 487.43,
"6": 502.39,
"7": 484.76,
"8": 498.5,
"9": 489.98,
"10": 499.22,
"11": 490.47,
"12": 498.98,
"13": 491.42,
"14": 495.17,
"15": 492.24,
"16": 494.69,
"17": 491.68,
"18": 497.13,
"19": 493.63,
"20": 495.29,
"21": 496.19,
"22": 496.52,
"23": 496.31,
"24": 493.64,
"25": 494.35,
"26": 493.27,
"27": 495.52,
"28": 493.06,
"29": 494.66,
"30": 494.08,
"31": 496.01,
"32": 494.79,
"33": 495.81,
"34": 493.94,
"35": 495.56,
"36": 493.5,
"37": 495.98,
"38": 495.82,
"39": 496.7,
"40": 495.29,
"41": 496.84,
"42": 495.67,
"43": 495.32,
"44": 496.48,
"45": 496.03,
"46": 495.65,
"47": 498.27,
"48": 496.44,
"49": 496.79
},
"expected_e2e_ms": 25735.08,
"expected_avg_denoise_ms": 486.1,
"expected_median_denoise_ms": 495.42
},
"wan2_1_t2v_1.3b_upscaling_4x": {
"stages_ms": {
"InputValidationStage": 0.08,
"TextEncodingStage": 1164.21,
"LatentPreparationStage": 0.26,
"TimestepPreparationStage": 3.86,
"DenoisingStage": 10234.35,
"DecodingStage": 499.44
},
"denoise_step_ms": {
"0": 200.87,
"1": 202.85,
"2": 203.32,
"3": 206.02,
"4": 205.75,
"5": 204.14,
"6": 205.12,
"7": 204.65,
"8": 203.99,
"9": 204.96,
"10": 204.35,
"11": 206.89,
"12": 200.69,
"13": 209.67,
"14": 204.91,
"15": 203.5,
"16": 206.73,
"17": 202.43,
"18": 205.92,
"19": 204.61,
"20": 211.47,
"21": 197.43,
"22": 203.58,
"23": 205.82,
"24": 204.01,
"25": 205.06,
"26": 204.86,
"27": 206.03,
"28": 200.78,
"29": 206.99,
"30": 206.58,
"31": 202.84,
"32": 204.51,
"33": 204.19,
"34": 202.89,
"35": 204.55,
"36": 205.03,
"37": 204.2,
"38": 203.92,
"39": 204.9,
"40": 203.24,
"41": 204.21,
"42": 205.76,
"43": 205.32,
"44": 202.63,
"45": 205.67,
"46": 204.55,
"47": 202.89,
"48": 205.29,
"49": 203.87
},
"expected_e2e_ms": 12021.58,
"expected_avg_denoise_ms": 204.49,
"expected_median_denoise_ms": 204.55
},
"wan2_1_t2v_1.3b_frame_interp_2x_upscaling_4x": {
"stages_ms": {
"InputValidationStage": 0.03,
"TextEncodingStage": 1089.94,
"LatentPreparationStage": 0.12,
"TimestepPreparationStage": 1.99,
"DenoisingStage": 8617.56,
"DecodingStage": 469.68
},
"denoise_step_ms": {
"0": 122.51,
"1": 172.31,
"2": 169.58,
"3": 171.76,
"4": 171.5,
"5": 174.24,
"6": 175.85,
"7": 171.2,
"8": 172.0,
"9": 172.13,
"10": 171.22,
"11": 174.8,
"12": 174.01,
"13": 172.28,
"14": 172.77,
"15": 173.52,
"16": 172.19,
"17": 175.12,
"18": 172.78,
"19": 175.1,
"20": 171.54,
"21": 173.38,
"22": 171.61,
"23": 174.14,
"24": 174.06,
"25": 172.32,
"26": 173.08,
"27": 173.94,
"28": 173.32,
"29": 174.3,
"30": 173.63,
"31": 172.21,
"32": 174.4,
"33": 173.25,
"34": 173.54,
"35": 175.12,
"36": 172.93,
"37": 172.76,
"38": 174.73,
"39": 174.46,
"40": 172.66,
"41": 174.58,
"42": 173.9,
"43": 174.88,
"44": 172.35,
"45": 173.52,
"46": 175.94,
"47": 172.88,
"48": 174.97,
"49": 172.94
},
"expected_e2e_ms": 10425.77,
"expected_avg_denoise_ms": 172.28,
"expected_median_denoise_ms": 173.29
}
}
}

42
python/sglang/multimodal_gen/test/server/test_server_utils.py
View File

@@ -871,9 +871,7 @@ def get_generate_fn(
req_background = None # Not specified in current request
# Build extra_body for optional features
extra_body = {}
if sampling_params.enable_teacache:
extra_body["enable_teacache"] = True
extra_body = dict(sampling_params.extras)
response = client.images.with_raw_response.generate(
model=model_path,
@@ -898,6 +896,13 @@ def get_generate_fn(
# Validate output file
expected_width, expected_height = parse_dimensions(output_size)
if (
sampling_params.extras.get("enable_upscaling")
and expected_width
and expected_height
):
expected_width *= sampling_params.extras.get("upscaling_scale", 4)
expected_height *= sampling_params.extras.get("upscaling_scale", 4)
validate_image_file(
tmp_path,
expected_filename,
@@ -947,8 +952,7 @@ def get_generate_fn(
# Build extra_body for optional features
extra_body = {"num_frames": sampling_params.num_frames}
if sampling_params.enable_teacache:
extra_body["enable_teacache"] = True
extra_body.update(sampling_params.extras)
images = [open(image_path, "rb") for image_path in image_paths]
try:
@@ -1076,22 +1080,18 @@ def get_generate_fn(
pytest.skip(f"{case_id}: no text prompt configured")
# Build extra_body for optional features
extra_body = {}
if sampling_params.enable_teacache:
extra_body["enable_teacache"] = True
extra_body = dict(sampling_params.extras)
if sampling_params.num_frames:
extra_body["num_frames"] = sampling_params.num_frames
if sampling_params.enable_frame_interpolation:
extra_body["enable_frame_interpolation"] = True
extra_body["frame_interpolation_exp"] = (
sampling_params.frame_interpolation_exp
)
# Compute expected output frame count for validation
expected_frame_count = None
if sampling_params.enable_frame_interpolation and sampling_params.num_frames:
if (
sampling_params.extras.get("enable_frame_interpolation")
and sampling_params.num_frames
):
n = sampling_params.num_frames
exp = sampling_params.frame_interpolation_exp
exp = sampling_params.extras.get("frame_interpolation_exp", 1)
expected_frame_count = (n - 1) * (2**exp) + 1
return _create_and_download_video(
@@ -1118,9 +1118,7 @@ def get_generate_fn(
pytest.skip(f"{case_id}: file missing: {image_path}")
# Build extra_body for optional features
extra_body = {}
if sampling_params.enable_teacache:
extra_body["enable_teacache"] = True
extra_body = dict(sampling_params.extras)
with image_path.open("rb") as fh:
return _create_and_download_video(
@@ -1140,8 +1138,7 @@ def get_generate_fn(
# Build extra_body for optional features
extra_body = {"reference_url": sampling_params.image_path}
if sampling_params.enable_teacache:
extra_body["enable_teacache"] = True
extra_body.update(sampling_params.extras)
return _create_and_download_video(
client,
@@ -1170,9 +1167,7 @@ def get_generate_fn(
pytest.skip(f"{case_id}: file missing: {image_path}")
# Build extra_body for optional features
extra_body = {}
if sampling_params.enable_teacache:
extra_body["enable_teacache"] = True
extra_body = dict(sampling_params.extras)
with image_path.open("rb") as fh:
return _create_and_download_video(
@@ -1186,6 +1181,7 @@ def get_generate_fn(
extra_body={
"fps": sampling_params.fps,
"num_frames": sampling_params.num_frames,
**extra_body,
},
)

62
python/sglang/multimodal_gen/test/server/testcase_configs.py
View File

@@ -233,12 +233,9 @@ class DiffusionSamplingParams:
num_outputs_per_prompt: int = 1
# TeaCache acceleration
enable_teacache: bool = False
# Frame interpolation
enable_frame_interpolation: bool = False
frame_interpolation_exp: int = 1 # 1 = 2×, 2 = 4×
# Additional request-level parameters (e.g. enable_teacache, enable_upscaling, …)
# merged directly into the OpenAI extra_body dict.
extras: dict = field(default_factory=dict)
@dataclass(frozen=True)
@@ -494,6 +491,19 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [
),
MULTI_FRAME_I2I_sampling_params,
),
# Upscaling (Real-ESRGAN 4×) for T2I
DiffusionTestCase(
"flux_2_image_t2i_upscaling_4x",
DiffusionServerArgs(
model_path="black-forest-labs/FLUX.2-dev",
modality="image",
),
DiffusionSamplingParams(
prompt="Doraemon is eating dorayaki",
output_size="1024x1024",
extras={"enable_upscaling": True, "upscaling_scale": 4},
),
),
]
HUNYUAN3D_SHAPE_sampling_params = DiffusionSamplingParams(
@@ -536,10 +546,10 @@ ONE_GPU_CASES_B: list[DiffusionTestCase] = [
),
DiffusionSamplingParams(
prompt=T2V_PROMPT,
enable_teacache=True,
extras={"enable_teacache": True},
),
),
# Frame interpolation correctness (2× / exp=1)
# Frame interpolation (2× / exp=1)
# Uses the same 1.3B model already in the suite;
DiffusionTestCase(
"wan2_1_t2v_1.3b_frame_interp_2x",
@@ -550,8 +560,40 @@ ONE_GPU_CASES_B: list[DiffusionTestCase] = [
),
DiffusionSamplingParams(
prompt=T2V_PROMPT,
enable_frame_interpolation=True,
frame_interpolation_exp=1,
extras={"enable_frame_interpolation": True, "frame_interpolation_exp": 1},
),
),
# Upscaling (Real-ESRGAN 4×)
# Uses the same 1.3B model already in the suite;
DiffusionTestCase(
"wan2_1_t2v_1.3b_upscaling_4x",
DiffusionServerArgs(
model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
modality="video",
custom_validator="video",
),
DiffusionSamplingParams(
prompt=T2V_PROMPT,
extras={"enable_upscaling": True, "upscaling_scale": 4},
),
),
# Combined: Frame interpolation (2×) + Upscaling (4×)
# Verifies that both post-processing steps compose correctly.
DiffusionTestCase(
"wan2_1_t2v_1.3b_frame_interp_2x_upscaling_4x",
DiffusionServerArgs(
model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
modality="video",
custom_validator="video",
),
DiffusionSamplingParams(
prompt=T2V_PROMPT,
extras={
"enable_frame_interpolation": True,
"frame_interpolation_exp": 1,
"enable_upscaling": True,
"upscaling_scale": 4,
},
),
),
# LoRA test case for single transformer + merge/unmerge API test