Refactor code structure for improved readability and maintainability

workers/upscale/Dockerfile

@@ -0,0 +1,34 @@
+FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
+
+ENV DEBIAN_FRONTEND=noninteractive
+ENV PIP_NO_CACHE_DIR=1
+ENV PYTHONUNBUFFERED=1
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        python3 python3-pip python3-dev \
+        libgl1 libglib2.0-0 wget ca-certificates \
+    && rm -rf /var/lib/apt/lists/* \
+    && ln -sf /usr/bin/python3 /usr/bin/python
+
+WORKDIR /app
+
+# Newer CUDA-12.1 torch so the compiled kernels cover current GPUs (fixes
+# "no kernel image is available for execution on the device"). The basicsr
+# functional_tensor import (removed in torchvision>=0.16) is handled by the
+# sed-patch below instead of pinning an old torchvision.
+RUN pip install --upgrade pip && \
+    pip install torch==2.1.2 torchvision==0.16.2 \
+        --index-url https://download.pytorch.org/whl/cu121
+
+COPY workers/upscale/requirements.txt .
+RUN pip install -r requirements.txt
+
+# Robustly fix basicsr's removed torchvision import (the #1 Real-ESRGAN crash).
+# Works whether or not the torch/torchvision pin held — rewrites the bad import.
+RUN F=$(python3 -c 'import basicsr,os;print(os.path.join(os.path.dirname(basicsr.__file__),"data","degradations.py"))') && \
+    sed -i 's/from torchvision.transforms.functional_tensor import rgb_to_grayscale/from torchvision.transforms.functional import rgb_to_grayscale/' "$F" && \
+    echo "patched basicsr degradations.py"
+
+COPY workers/upscale/handler.py .
+
+CMD ["python3","-u","handler.py"]

workers/upscale/README.md

@@ -0,0 +1,35 @@
+# Real-ESRGAN Upscale Worker (#7)
+
+RunPod Serverless worker for Real-ESRGAN image super-resolution with optional GFPGAN face enhancement.
+
+- **Dockerfile path:** `/workers/upscale/Dockerfile`
+- **Build context:** repository ROOT (COPY paths are prefixed with `workers/upscale/`).
+- **GPU:** medium (RTX 3090 / RTX 4090, ~24 GB VRAM).
+- **App env var to set:** `RUNPOD_UPSCALE_URL` (point your app at this endpoint's RunPod URL).
+
+## Contract
+
+Input:
+```json
+{"image": "<base64 PNG/JPEG, optional data: prefix>", "scale": 2 or 4 (default 4), "face_enhance": false}
+```
+Output (success):
+```json
+{"image": "<base64 PNG>"}
+```
+Output (error):
+```json
+{"error": "SomeError: message"}
+```
+
+## Model weights
+
+Downloaded on first request and cached to the network volume when mounted:
+- `RealESRGAN_x4plus.pth` (RRDBNet x4) -> `/runpod-volume/weights/`
+- `GFPGANv1.4.pth` (only when `face_enhance=true`) -> `/runpod-volume/weights/`
+
+Attach a network volume mounted at `/runpod-volume` so weights (and `HF_HOME`) persist between cold starts. Without a volume, weights download to `/app/weights` per container.
+
+## Notes
+
+The base model is x4; `scale=2` is produced via RealESRGANer's `outscale` parameter. For `face_enhance`, GFPGAN upscales at its fixed factor and the result is resized to the requested scale.

workers/upscale/handler.py

@@ -0,0 +1,113 @@
+_V = "/runpod-volume"
+import os
+if os.path.isdir(_V):
+    os.environ.setdefault("HF_HOME", os.path.join(_V, "huggingface"))
+
+import io
+import base64
+import traceback
+import urllib.request
+
+import numpy as np
+from PIL import Image
+
+_WEIGHTS_DIR = os.path.join(_V, "weights") if os.path.isdir(_V) else "/app/weights"
+os.makedirs(_WEIGHTS_DIR, exist_ok=True)
+
+_RRDB_URL = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"
+_GFPGAN_URL = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth"
+
+_UPSAMPLER = None
+_FACE_ENHANCER = None
+
+
+def _download(url, dst):
+    if not os.path.isfile(dst):
+        tmp = dst + ".tmp"
+        urllib.request.urlretrieve(url, tmp)
+        os.replace(tmp, dst)
+    return dst
+
+
+def _get_upsampler():
+    global _UPSAMPLER
+    if _UPSAMPLER is None:
+        import torch
+        from basicsr.archs.rrdbnet_arch import RRDBNet
+        from realesrgan import RealESRGANer
+
+        model_path = _download(_RRDB_URL, os.path.join(_WEIGHTS_DIR, "RealESRGAN_x4plus.pth"))
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+        _UPSAMPLER = RealESRGANer(
+            scale=4,
+            model_path=model_path,
+            model=model,
+            tile=0,
+            tile_pad=10,
+            pre_pad=0,
+            half=torch.cuda.is_available(),
+            gpu_id=None,
+        )
+    return _UPSAMPLER
+
+
+def _get_face_enhancer():
+    global _FACE_ENHANCER
+    if _FACE_ENHANCER is None:
+        from gfpgan import GFPGANer
+
+        gfpgan_path = _download(_GFPGAN_URL, os.path.join(_WEIGHTS_DIR, "GFPGANv1.4.pth"))
+        _FACE_ENHANCER = GFPGANer(
+            model_path=gfpgan_path,
+            upscale=4,
+            arch="clean",
+            channel_multiplier=2,
+            bg_upsampler=_get_upsampler(),
+        )
+    return _FACE_ENHANCER
+
+
+def handler(event):
+    try:
+        data = (event or {}).get("input") or {}
+        b64 = data.get("image")
+        if not b64:
+            return {"error": "ValueError: 'image' is required"}
+        if "," in b64:
+            b64 = b64.split(",", 1)[1]
+
+        scale = int(data.get("scale", 4))
+        if scale not in (2, 4):
+            scale = 4
+        face_enhance = bool(data.get("face_enhance", False))
+
+        raw = base64.b64decode(b64)
+        pil = Image.open(io.BytesIO(raw)).convert("RGB")
+        rgb = np.array(pil)
+        bgr = rgb[:, :, ::-1].copy()  # RealESRGANer expects BGR (cv2 convention)
+
+        if face_enhance:
+            enhancer = _get_face_enhancer()
+            _, _, out_bgr = enhancer.enhance(
+                bgr, has_aligned=False, only_center_face=False, paste_back=True
+            )
+            # GFPGANer upscales by its fixed factor; resize to the requested scale.
+            target = (rgb.shape[1] * scale, rgb.shape[0] * scale)
+            out_rgb = out_bgr[:, :, ::-1]
+            out_pil = Image.fromarray(out_rgb).resize(target, Image.LANCZOS)
+        else:
+            upsampler = _get_upsampler()
+            out_bgr, _ = upsampler.enhance(bgr, outscale=scale)
+            out_rgb = out_bgr[:, :, ::-1]
+            out_pil = Image.fromarray(out_rgb)
+
+        buf = io.BytesIO()
+        out_pil.save(buf, format="PNG")
+        return {"image": base64.b64encode(buf.getvalue()).decode("utf-8")}
+    except Exception as exc:
+        traceback.print_exc()
+        return {"error": f"{type(exc).__name__}: {exc}"}
+
+
+import runpod
+runpod.serverless.start({"handler": handler})

workers/upscale/requirements.txt

@@ -0,0 +1,6 @@
+runpod==1.7.9
+realesrgan==0.3.0
+basicsr==1.4.2
+gfpgan==1.3.8
+numpy==1.24.4
+pillow==10.4.0

workers/upscale/test.json

@@ -0,0 +1 @@
+{"input": {"image": "iVBORw0KGgoAAAANSUhEUgAAAgAAAAIACAIAAAB7GkOtAAAFp0lEQVR42u3VMREAMAgAMfRVRO2woK1GagIJCCB3UfDLR74PwEIhAYABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAIABqABgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgACoAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAwDiAWweAhQwAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAwAAkADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADADAAFQAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAMQAUAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADADAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADADAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAYNQ/e2eZ8ltazgAAAABJRU5ErkJggg==", "scale": 4}}