Initial GFPGAN custom handler

handler.py

@@ -0,0 +1,115 @@
+import os
+import io
+import torch
+import base64
+import requests
+import numpy as np
+from PIL import Image
+from gfpgan import GFPGANer
+from realesrgan import RealESRGANer
+from basicsr.archs.rrdbnet_arch import RRDBNet
+
+
+class EndpointHandler:
+    def __init__(self, path="."):
+        print("🚀 Initializing GFPGANv1 Face Restoration Pipeline...")
+
+        # ------------------------------------------------------------
+        # Load Real-ESRGAN (background upscaler)
+        # ------------------------------------------------------------
+        self.esrgan_url = (
+            "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/"
+            "RealESRGAN_x4plus.pth"
+        )
+        self.esrgan_path = os.path.join(path, "RealESRGAN_x4plus.pth")
+
+        if not os.path.exists(self.esrgan_path):
+            print("📥 Downloading RealESRGAN_x4plus.pth...")
+            r = requests.get(self.esrgan_url)
+            r.raise_for_status()
+            with open(self.esrgan_path, "wb") as f:
+                f.write(r.content)
+            print("✅ Downloaded Real-ESRGAN model.")
+
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        self.bg_upsampler = RealESRGANer(
+            scale=4,
+            model_path=self.esrgan_path,
+            model=model,
+            half=False,
+            device=device,
+        )
+
+        # ------------------------------------------------------------
+        # Load GFPGAN model
+        # ------------------------------------------------------------
+        self.gfpgan_url = (
+            "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth"
+        )
+        self.gfpgan_path = os.path.join(path, "GFPGANv1.4.pth")
+
+        if not os.path.exists(self.gfpgan_path):
+            print("📥 Downloading GFPGANv1.4.pth...")
+            r = requests.get(self.gfpgan_url)
+            r.raise_for_status()
+            with open(self.gfpgan_path, "wb") as f:
+                f.write(r.content)
+            print("✅ Downloaded GFPGANv1.4.pth.")
+
+        self.face_enhancer = GFPGANer(
+            model_path=self.gfpgan_path,
+            upscale=4,
+            arch="clean",
+            channel_multiplier=2,
+            bg_upsampler=self.bg_upsampler,
+        )
+
+        print("✅ GFPGANv1 Face Restoration Pipeline Ready!")
+
+    # ------------------------------------------------------------------
+    # Main callable
+    # ------------------------------------------------------------------
+    def __call__(self, data):
+        try:
+            image = self._load_image(data)
+            restored = self._restore_face(image)
+            return self._encode_image(restored)
+        except Exception as e:
+            print("💥 Error:", str(e))
+            return {"error": str(e)}
+
+    # ------------------------------------------------------------------
+    # Helper functions
+    # ------------------------------------------------------------------
+    def _load_image(self, data):
+        # Handles raw bytes, base64 JSON, or dict
+        if isinstance(data, (bytes, bytearray)):
+            return Image.open(io.BytesIO(data)).convert("RGB")
+
+        if isinstance(data, dict):
+            field = data.get("inputs") or data.get("image")
+            if isinstance(field, str):
+                field = base64.b64decode(field)
+            return Image.open(io.BytesIO(field)).convert("RGB")
+
+        if isinstance(data, str):
+            decoded = base64.b64decode(data)
+            return Image.open(io.BytesIO(decoded)).convert("RGB")
+
+        raise ValueError("Expected image bytes or base64 string.")
+
+    def _restore_face(self, image):
+        img_np = np.array(image)[:, :, ::-1]  # RGB -> BGR
+        cropped_faces, restored_faces, restored_img = self.face_enhancer.enhance(
+            img_np, has_aligned=False, only_center_face=False, paste_back=True
+        )
+        return Image.fromarray(restored_img[:, :, ::-1])  # back to RGB
+
+    def _encode_image(self, pil_img):
+        buf = io.BytesIO()
+        pil_img.save(buf, format="PNG")
+        encoded = base64.b64encode(buf.getvalue()).decode("utf-8")
+        return {"image": encoded}
+

requirements.txt

@@ -0,0 +1,11 @@
+torch==2.1.2
+torchvision==0.16.2
+gfpgan==1.3.8
+realesrgan==0.3.0
+basicsr==1.4.2
+facexlib==0.3.0
+numpy==1.26.4
+Pillow>=10.0.0
+opencv-python
+requests
+