sgl-project--sglang/python/sglang/multimodal_gen/runtime/postprocess/realesrgan_upscaler.py

# 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
import time
from hashlib import sha256
from typing import Optional
from urllib.parse import unquote, urlparse

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"
_DEFAULT_REALESRGAN_FILENAMES_BY_SCALE = {
    2: "RealESRGAN_x2.pth",
    4: "RealESRGAN_x4.pth",
    8: "RealESRGAN_x8.pth",
}
_LOW_MEMORY_TILED_UPSCALE_FREE_BYTES = 2 * 1024**3
_REALESRGAN_TILE_SIZE = 256
_REALESRGAN_TILE_PAD = 32

# Module-level cache: model_path -> UpscalerModel instance
_MODEL_CACHE: dict[str, "UpscalerModel"] = {}
_RESOLVED_MODEL_PATH_CACHE: dict[str, str] = {}


def _default_model_path_for_scale(scale: int) -> str:
    filename = _DEFAULT_REALESRGAN_FILENAMES_BY_SCALE.get(
        int(scale),
        _DEFAULT_REALESRGAN_FILENAME,
    )
    return f"{_DEFAULT_REALESRGAN_HF_REPO}:{filename}"


# ---------------------------------------------------------------------------
# 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]
        in_channels = state_dict["conv_first.weight"].shape[1]
        if in_channels == 3:
            scale = 4
        elif in_channels == 12:
            scale = 2
        elif in_channels == 48:
            scale = 1
        else:
            raise ValueError(
                f"Unsupported RRDBNet conv_first input channels: {in_channels}"
            )
        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, scale=%d",
            num_feat,
            num_block,
            num_grow_ch,
            scale,
        )
        return RRDBNet(
            num_in_ch=3,
            num_out_ch=3,
            scale=scale,
            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

    @property
    def dtype(self) -> torch.dtype:
        return next(self.net.parameters()).dtype

    def _copy_input_to_device(self, frames: np.ndarray) -> torch.Tensor:
        return torch.from_numpy(frames).to(self.device)

    def _preprocess_input_tensor(self, imgs_t: torch.Tensor) -> torch.Tensor:
        imgs_t = imgs_t.permute(0, 3, 1, 2).to(dtype=self.dtype).mul_(1.0 / 255.0)
        if self.device.type == "cuda":
            imgs_t = imgs_t.contiguous(memory_format=torch.channels_last)
        return imgs_t

    @staticmethod
    def _postprocess_output_tensor(out: torch.Tensor) -> torch.Tensor:
        out = out.permute(0, 2, 3, 1).clamp(0.0, 1.0).mul_(255.0)
        return out.to(torch.uint8).contiguous()

    @staticmethod
    def _copy_output_to_host(out: torch.Tensor) -> np.ndarray:
        return out.cpu().numpy()

    def _start_cuda_timer(self):
        if self.device.type != "cuda":
            return None
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        start.record()
        return start, end

    @staticmethod
    def _stop_cuda_timer(timer) -> None:
        if timer is not None:
            timer[1].record()

    @staticmethod
    def _cuda_elapsed_s(timer, fallback_s: float) -> float:
        if timer is None:
            return fallback_s
        timer[1].synchronize()
        return timer[0].elapsed_time(timer[1]) / 1000.0

    def _should_use_tiled_upscale(self, h: int, w: int) -> bool:
        if self.device.type != "cuda":
            return False
        free_bytes, _ = torch.cuda.mem_get_info(self.device)
        output_bytes = h * w * self.scale * self.scale * 3 * 4
        required_free_bytes = max(
            _LOW_MEMORY_TILED_UPSCALE_FREE_BYTES,
            output_bytes * 4,
        )
        return free_bytes < required_free_bytes

    def _upscale_tiled_to_cpu(
        self,
        img_t: torch.Tensor,
        tile_size: int = _REALESRGAN_TILE_SIZE,
        tile_pad: int = _REALESRGAN_TILE_PAD,
    ) -> torch.Tensor:
        _, channels, h, w = img_t.shape
        scale = self.scale
        output = torch.empty(
            (1, channels, h * scale, w * scale),
            dtype=torch.float32,
            device="cpu",
        )

        for y in range(0, h, tile_size):
            tile_h = min(tile_size, h - y)
            in_y0 = max(y - tile_pad, 0)
            in_y1 = min(y + tile_h + tile_pad, h)
            out_y0 = y * scale
            out_y1 = (y + tile_h) * scale
            crop_y0 = (y - in_y0) * scale
            crop_y1 = crop_y0 + tile_h * scale

            for x in range(0, w, tile_size):
                tile_w = min(tile_size, w - x)
                in_x0 = max(x - tile_pad, 0)
                in_x1 = min(x + tile_w + tile_pad, w)
                out_x0 = x * scale
                out_x1 = (x + tile_w) * scale
                crop_x0 = (x - in_x0) * scale
                crop_x1 = crop_x0 + tile_w * scale

                tile = img_t[..., in_y0:in_y1, in_x0:in_x1]
                out_tile = self.net(tile)
                out_tile = out_tile[..., crop_y0:crop_y1, crop_x0:crop_x1].float()
                output[..., out_y0:out_y1, out_x0:out_x1].copy_(out_tile.cpu())

        return output

    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(device=self.device, dtype=self.dtype)
        )
        with torch.no_grad():
            if self._should_use_tiled_upscale(h, w):
                logger.info(
                    "Using tiled Real-ESRGAN upscale for low GPU memory: "
                    "frame=%dx%d, tile_size=%d, tile_pad=%d",
                    w,
                    h,
                    _REALESRGAN_TILE_SIZE,
                    _REALESRGAN_TILE_PAD,
                )
                out = self._upscale_tiled_to_cpu(img_t)
            else:
                try:
                    out = self.net(img_t)
                except torch.cuda.OutOfMemoryError:
                    if self.device.type != "cuda":
                        raise
                    torch.cuda.empty_cache()
                    logger.warning(
                        "Real-ESRGAN full-frame upscale OOM; retrying with tiled upscale"
                    )
                    out = self._upscale_tiled_to_cpu(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)

    def upscale_batch(
        self, frames: list[np.ndarray], outscale: float | None = None
    ) -> list[np.ndarray]:
        """Upscale same-resolution HWC uint8 frames in one batched forward pass."""
        if not frames:
            return []

        h, w = frames[0].shape[:2]
        if any(frame.shape[:2] != (h, w) for frame in frames):
            raise ValueError("All frames in a batch must have the same resolution")

        total_start_time = time.perf_counter()

        start_time = time.perf_counter()
        imgs = np.stack(frames, axis=0)
        stack_duration_s = time.perf_counter() - start_time

        start_time = time.perf_counter()
        h2d_timer = self._start_cuda_timer()
        imgs_t = self._copy_input_to_device(imgs)
        self._stop_cuda_timer(h2d_timer)
        h2d_wall_duration_s = time.perf_counter() - start_time

        start_time = time.perf_counter()
        input_preprocess_timer = self._start_cuda_timer()
        imgs_t = self._preprocess_input_tensor(imgs_t)
        self._stop_cuda_timer(input_preprocess_timer)
        input_preprocess_wall_duration_s = time.perf_counter() - start_time

        start_time = time.perf_counter()
        forward_timer = self._start_cuda_timer()
        with torch.inference_mode():
            out = self.net(imgs_t)
        self._stop_cuda_timer(forward_timer)
        forward_wall_duration_s = time.perf_counter() - start_time

        resize_timer = None
        resize_wall_duration_s = 0.0
        if outscale is not None and outscale != self.scale:
            start_time = time.perf_counter()
            resize_timer = self._start_cuda_timer()
            target_h = int(h * outscale)
            target_w = int(w * outscale)
            out = F.interpolate(
                out, size=(target_h, target_w), mode="bicubic", align_corners=False
            )
            self._stop_cuda_timer(resize_timer)
            resize_wall_duration_s = time.perf_counter() - start_time

        start_time = time.perf_counter()
        output_postprocess_timer = self._start_cuda_timer()
        out = self._postprocess_output_tensor(out)
        self._stop_cuda_timer(output_postprocess_timer)
        output_postprocess_wall_duration_s = time.perf_counter() - start_time

        start_time = time.perf_counter()
        output_d2h_timer = self._start_cuda_timer()
        out_np = self._copy_output_to_host(out)
        self._stop_cuda_timer(output_d2h_timer)
        output_d2h_wall_duration_s = time.perf_counter() - start_time

        start_time = time.perf_counter()
        outputs = [frame for frame in out_np]
        post_duration_s = time.perf_counter() - start_time

        h2d_duration_s = self._cuda_elapsed_s(h2d_timer, h2d_wall_duration_s)
        input_preprocess_duration_s = self._cuda_elapsed_s(
            input_preprocess_timer, input_preprocess_wall_duration_s
        )
        forward_duration_s = self._cuda_elapsed_s(
            forward_timer, forward_wall_duration_s
        )
        resize_duration_s = self._cuda_elapsed_s(resize_timer, resize_wall_duration_s)
        output_postprocess_duration_s = self._cuda_elapsed_s(
            output_postprocess_timer, output_postprocess_wall_duration_s
        )
        output_d2h_duration_s = self._cuda_elapsed_s(
            output_d2h_timer, output_d2h_wall_duration_s
        )
        total_duration_s = time.perf_counter() - total_start_time
        timing_source = "cuda_event" if self.device.type == "cuda" else "wall"
        logger.info(
            "RealESRGAN batch upscale: batch=%d input=%dx%d native_scale=%dx outscale=%s "
            "dtype=%s timing=%s total=%.3fs stack=%.3fs input_h2d=%.3fs "
            "input_pre=%.3fs forward=%.3fs resize=%.3fs output_post=%.3fs "
            "output_d2h=%.3fs python_post=%.3fs",
            len(frames),
            w,
            h,
            self.scale,
            outscale if outscale is not None else self.scale,
            self.dtype,
            timing_source,
            total_duration_s,
            stack_duration_s,
            h2d_duration_s,
            input_preprocess_duration_s,
            forward_duration_s,
            resize_duration_s,
            output_postprocess_duration_s,
            output_d2h_duration_s,
            post_duration_s,
        )
        return outputs


# ---------------------------------------------------------------------------
# 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_model_path_for_scale(self._scale)

        # 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]

    def upscale_batched(self, frames: list[np.ndarray]) -> list[np.ndarray]:
        """Upscale HWC uint8 frames with batched forwards grouped by resolution."""
        if not frames:
            return frames
        total_start_time = time.perf_counter()
        model = self._ensure_model_loaded()
        outscale = self._scale if self._scale != model.scale else None
        output_frames: list[np.ndarray | None] = [None] * len(frames)
        groups: dict[tuple[int, ...], list[int]] = {}
        for idx, frame in enumerate(frames):
            groups.setdefault(tuple(frame.shape), []).append(idx)

        for shape, indices in groups.items():
            logger.info(
                "RealESRGAN upscale group: frames=%d shape=%s indices=%s",
                len(indices),
                shape,
                indices,
            )
            group_frames = [frames[idx] for idx in indices]
            group_outputs = model.upscale_batch(group_frames, outscale=outscale)
            for idx, output in zip(indices, group_outputs):
                output_frames[idx] = output

        if any(frame is None for frame in output_frames):
            raise RuntimeError("RealESRGAN batch upscale did not produce all frames")

        total_duration_s = time.perf_counter() - total_start_time
        logger.info(
            "RealESRGAN batch_upscale_frames completed in %.3f seconds for %d frames across %d groups",
            total_duration_s,
            len(frames),
            len(groups),
        )
        return [frame for frame in output_frames if frame is not None]


# ---------------------------------------------------------------------------
# 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).
    - An http(s) URL to a .pth file, downloaded into the local cache.
    - 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.
    """
    cached_path = _RESOLVED_MODEL_PATH_CACHE.get(model_path)
    if cached_path is not None:
        return cached_path

    if os.path.isfile(model_path):
        _RESOLVED_MODEL_PATH_CACHE[model_path] = model_path
        return model_path

    parsed_url = urlparse(model_path)
    if parsed_url.scheme in ("http", "https"):
        filename = (
            os.path.basename(unquote(parsed_url.path)) or _DEFAULT_REALESRGAN_FILENAME
        )
        cache_dir = os.path.join(
            os.path.expanduser("~"), ".cache", "sglang", "realesrgan"
        )
        os.makedirs(cache_dir, exist_ok=True)
        cache_key = sha256(model_path.encode("utf-8")).hexdigest()[:12]
        local_path = os.path.join(cache_dir, f"{cache_key}-{filename}")
        if not os.path.isfile(local_path):
            tmp_path = f"{local_path}.tmp"
            logger.info("Downloading Real-ESRGAN weights from URL %s", model_path)
            torch.hub.download_url_to_file(model_path, tmp_path, progress=False)
            os.replace(tmp_path, local_path)
        _RESOLVED_MODEL_PATH_CACHE[model_path] = local_path
        return local_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
    _RESOLVED_MODEL_PATH_CACHE[model_path] = local_path
    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)


def batch_upscale_frames(
    frames: list[np.ndarray],
    model_path: Optional[str] = None,
    scale: int = 4,
) -> list[np.ndarray]:
    """
    Batched Real-ESRGAN upscaling for realtime video paths.

    The default ``upscale_frames`` API intentionally keeps its original
    per-frame behavior. Call this helper only when the caller can tolerate
    batched execution and same-shape grouping semantics.
    """
    upscaler = ImageUpscaler(
        model_path=model_path,
        scale=scale,
        half_precision=current_platform.is_cuda(),
    )
    return upscaler.upscale_batched(frames)