Refactor code structure for improved readability and maintainability

2026-07-22 22:58:28 +05:30
parent d1308bf149
commit bc8bf2007d
99 changed files with 4784 additions and 111 deletions
@@ -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})