feat(upscale): realesrgan x4 wrapper with 0.5-resize bridge

tests/test_upscale.py

@@ -0,0 +1,26 @@
+from unittest import mock
+import pytest
+from PIL import Image
+
+import upscale
+
+
+@pytest.fixture
+def small_image():
+    return Image.new("RGB", (256, 256), color=(120, 50, 200))
+
+
+def test_realesrgan_2x_produces_2x_image(small_image, monkeypatch):
+    """RealESRGAN runs 4x then we scale down 0.5 → net 2x."""
+    def fake_run_4x(_model_path, image):
+        w, h = image.size
+        return image.resize((w * 4, h * 4), Image.LANCZOS)
+    monkeypatch.setattr(upscale, "_realesrgan_4x", fake_run_4x)
+
+    out = upscale.realesrgan_2x(small_image, model_path="/dev/null")
+    assert out.size == (512, 512)
+
+
+def test_realesrgan_2x_rejects_none():
+    with pytest.raises(ValueError):
+        upscale.realesrgan_2x(None, model_path="/dev/null")

upscale.py

@@ -0,0 +1,49 @@
+"""RealESRGAN x4plus wrapper + 0.5-resize bridge.
+
+This module only handles the *pixel-space* upscale. The Z-Image-Turbo refinement
+pass (img2img at denoise=0.33) lives in :mod:`modes` since it shares the pipeline.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+from PIL import Image
+
+
+def realesrgan_2x(image: Image.Image | None, model_path: Path | str) -> Image.Image:
+    """RealESRGAN x4plus → ``image.resize(0.5)`` → net 2x upscale."""
+    if image is None:
+        raise ValueError("upscale needs an input image")
+    upscaled = _realesrgan_4x(model_path, image)
+    w, h = upscaled.size
+    return upscaled.resize((w // 2, h // 2), Image.LANCZOS)
+
+
+_MODEL_CACHE: dict[str, Any] = {}
+
+
+def _realesrgan_4x(model_path: Path | str, image: Image.Image) -> Image.Image:
+    """Run RealESRGAN x4plus on ``image``. Caches the model in-process."""
+    import numpy as np
+    from realesrgan import RealESRGANer
+    from basicsr.archs.rrdbnet_arch import RRDBNet
+
+    key = str(model_path)
+    if key not in _MODEL_CACHE:
+        net = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+        _MODEL_CACHE[key] = RealESRGANer(
+            scale=4,
+            model_path=key,
+            model=net,
+            tile=512,        # split into tiles to avoid OOM on large inputs
+            tile_pad=10,
+            pre_pad=0,
+            half=False,      # bf16 elsewhere; keep this fp32 for stability
+            gpu_id=None,
+        )
+
+    upsampler = _MODEL_CACHE[key]
+    arr = np.array(image.convert("RGB"))
+    out_arr, _ = upsampler.enhance(arr, outscale=4)
+    return Image.fromarray(out_arr)