Add GFPGAN + RealESRGAN hybrid handler

handler.py

@@ -1,134 +1,138 @@
-import io
 import os
-import cv2
+import io
 import torch
 import base64
-import logging
 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
 
-# ======================================================
-#  LOGGING CONFIGURATION
-# ======================================================
-logging.basicConfig(level=logging.DEBUG)
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.DEBUG)
-logger.debug("📦 [INIT] Importing GFPGAN handler module...")
-
-# ======================================================
-#  GFPGAN MODEL URL
-# ======================================================
-MODEL_URL = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth"
-MODEL_NAME = "GFPGANv1.4.pth"
 
-
-# ======================================================
-#  ENDPOINT HANDLER
-# ======================================================
 class EndpointHandler:
     def __init__(self, path="."):
-        logger.debug("🚀 [INIT] Starting GFPGAN EndpointHandler initialization...")
-        logger.debug(f"📁 Working directory: {os.getcwd()}")
-        logger.debug(f"📂 Handler path argument: {path}")
-
-        model_path = os.path.join(path, MODEL_NAME)
-        logger.debug(f"🔗 [MODEL] Expected model path: {model_path}")
-
-        # ------------------------------------------------------
-        # Download model if missing
-        # ------------------------------------------------------
-        if not os.path.exists(model_path):
-            try:
-                logger.debug(f"📥 [DOWNLOAD] Model not found locally — fetching from {MODEL_URL}")
-                r = requests.get(MODEL_URL, stream=True)
-                r.raise_for_status()
-                with open(model_path, "wb") as f:
-                    for chunk in r.iter_content(chunk_size=8192):
-                        if chunk:
-                            f.write(chunk)
-                logger.debug("✅ [MODEL] Downloaded GFPGAN weights successfully.")
-            except Exception as e:
-                logger.error(f"💥 [ERROR] Failed to download GFPGAN weights: {e}")
-                raise
-
-        # ------------------------------------------------------
-        # Initialize GFPGANer (same as official Gradio demo)
-        # ------------------------------------------------------
-        try:
-            logger.debug("🧠 [MODEL] Initializing GFPGANer (upscale=2, arch='clean')...")
-            self.restorer = GFPGANer(
-                model_path=model_path,
-                upscale=2,                   # Rescaling factor = 2
-                arch="clean",
-                channel_multiplier=2,
-                bg_upsampler=None
-            )
-            logger.debug("✅ [MODEL] GFPGAN model initialized successfully.")
-        except Exception as e:
-            logger.error(f"💥 [ERROR] Model initialization failed: {e}")
-            raise
+        print("🚀 [INIT] Starting GFPGAN + RealESRGAN hybrid handler initialization...")
+        print(f"📂 Working directory: {os.getcwd()}")
+        print(f"📁 Handler path argument: {path}")
+
+        # ------------------------------
+        # Download GFPGAN v1.4 weights
+        # ------------------------------
+        self.model_url = "https://github.com/TencentARC/GFPGAN/releases/download/v1.4.0/GFPGANv1.4.pth"
+        self.model_path = os.path.join(path, "GFPGANv1.4.pth")
+        if not os.path.exists(self.model_path):
+            print(f"📥 [DOWNLOAD] Fetching GFPGAN v1.4 weights...")
+            r = requests.get(self.model_url)
+            r.raise_for_status()
+            with open(self.model_path, "wb") as f:
+                f.write(r.content)
+            print(f"✅ [MODEL] Downloaded GFPGAN model to {self.model_path}")
+        else:
+            print("✅ [MODEL] GFPGAN weights already exist locally.")
+
+        # ------------------------------
+        # Setup background upsampler (Real-ESRGAN)
+        # ------------------------------
+        print("🧠 [INIT] Setting up Real-ESRGAN background upsampler...")
+        rrdbnet = RRDBNet(
+            num_in_ch=3, num_out_ch=3,
+            num_feat=64, num_block=23,
+            num_grow_ch=32, scale=4
+        )
+
+        self.bg_upsampler = RealESRGANer(
+            scale=2,
+            model_path=None,  # auto-download model weights
+            model=rrdbnet,
+            tile=400,
+            tile_pad=10,
+            pre_pad=0,
+            half=False,
+            device="cuda" if torch.cuda.is_available() else "cpu",
+        )
+        print("✅ [INIT] Real-ESRGAN background upsampler ready.")
+
+        # ------------------------------
+        # Setup GFPGANer
+        # ------------------------------
+        print("🧠 [INIT] Setting up GFPGANer (v1.4)...")
+        self.restorer = GFPGANer(
+            model_path=self.model_path,
+            upscale=2,
+            arch="clean",
+            channel_multiplier=2,
+            bg_upsampler=self.bg_upsampler,
+            device="cuda" if torch.cuda.is_available() else "cpu",
+        )
+        print("✅ [INIT DONE] GFPGAN + RealESRGAN hybrid handler ready.")
 
-    # ======================================================
-    #  INFERENCE CALL
-    # ======================================================
+    # ----------------------------------------------------------
+    # Main inference entry point
+    # ----------------------------------------------------------
     def __call__(self, data):
-        logger.debug("⚙️ [INFER] Starting inference...")
-        logger.debug(f"📥 Incoming data type: {type(data)}")
+        print("🛰️ [CALL] Endpoint invoked!")
+        print(f"📦 [CALL] Raw input type: {type(data)}")
 
-        # ------------------------------------------------------
-        # Handle both JSON base64 and raw bytes
-        # ------------------------------------------------------
         try:
-            if isinstance(data, dict) and "inputs" in data:
-                logger.debug("📦 Detected JSON base64 input")
-                image_bytes = base64.b64decode(data["inputs"])
-            elif isinstance(data, (bytes, bytearray)):
-                logger.debug("📦 Detected raw bytes input")
-                image_bytes = data
-            else:
-                raise ValueError("Unsupported input format — expected bytes or base64 data")
-
-            logger.debug(f"🧾 [BYTES] Received {len(image_bytes)} bytes")
+            image = self.preprocess(data)
+            print("🧩 [STEP] Image preprocessed successfully.")
+            restored = self.inference(image)
+            print("🎨 [STEP] Inference completed successfully.")
+            return self.postprocess(restored)
         except Exception as e:
-            logger.error(f"💥 [ERROR] Input parsing failed: {e}")
-            return {"error": f"Invalid input: {e}"}
+            print(f"💥 [ERROR] Exception during call: {str(e)}")
+            return {"error": str(e)}
 
-        # ------------------------------------------------------
-        # Decode image
-        # ------------------------------------------------------
-        try:
-            img_np = np.array(Image.open(io.BytesIO(image_bytes)).convert("RGB"))
-            logger.debug(f"🖼️ [IMAGE] Loaded image of shape: {img_np.shape}")
-        except Exception as e:
-            logger.error(f"💥 [ERROR] Failed to load image: {e}")
-            return {"error": f"Image loading failed: {e}"}
+    # ----------------------------------------------------------
+    # Preprocessing
+    # ----------------------------------------------------------
+    def preprocess(self, data):
+        print("🔧 [PREPROCESS] Starting...")
+        if isinstance(data, (bytes, bytearray)):
+            print("🖼️ [PREPROCESS] Raw bytes detected.")
+            return Image.open(io.BytesIO(data)).convert("RGB")
 
-        # ------------------------------------------------------
-        # Run GFPGAN restoration
-        # ------------------------------------------------------
-        try:
-            cropped_faces, restored_faces, restored_img = self.restorer.enhance(
-                img_np,
-                has_aligned=False,
-                only_center_face=False,
-                paste_back=True  # Matches GFPGAN web demo
-            )
-            logger.debug("✅ [RESTORE] Face restoration completed successfully.")
-        except Exception as e:
-            logger.error(f"💥 [ERROR] GFPGAN enhancement failed: {e}")
-            return {"error": f"Enhancement failed: {e}"}
+        if isinstance(data, dict):
+            img_field = data.get("inputs") or data.get("image")
+            if isinstance(img_field, str):
+                print("🧬 [PREPROCESS] Base64 string detected.")
+                decoded = base64.b64decode(img_field)
+                return Image.open(io.BytesIO(decoded)).convert("RGB")
+            if isinstance(img_field, (bytes, bytearray)):
+                print("🧩 [PREPROCESS] Byte array detected.")
+                return Image.open(io.BytesIO(img_field)).convert("RGB")
 
-        # ------------------------------------------------------
-        # Encode result as base64 PNG
-        # ------------------------------------------------------
-        try:
-            _, buffer = cv2.imencode(".png", restored_img[:, :, ::-1])  # BGR→RGB
-            img_base64 = base64.b64encode(buffer).decode("utf-8")
-            logger.debug("📤 [ENCODE] Encoded restored image successfully.")
-            return {"image": img_base64}
-        except Exception as e:
-            logger.error(f"💥 [ERROR] Failed to encode image: {e}")
-            return {"error": f"Encoding failed: {e}"}
+        raise ValueError("Unsupported input format — expected bytes or base64 data.")
+
+    # ----------------------------------------------------------
+    # Inference
+    # ----------------------------------------------------------
+    def inference(self, image):
+        print("⚙️ [INFERENCE] Running GFPGAN + RealESRGAN enhancement...")
+        cropped_faces, restored_faces, restored_img = self.restorer.enhance(
+            image,
+            has_aligned=False,
+            only_center_face=False,
+            paste_back=True,
+        )
+        print(f"✅ [INFERENCE] Restored image size: {restored_img.shape}")
+        return restored_img
+
+    # ----------------------------------------------------------
+    # Postprocess
+    # ----------------------------------------------------------
+    def postprocess(self, restored_img):
+        print("📤 [POSTPROCESS] Encoding restored image...")
+        if isinstance(restored_img, torch.Tensor):
+            restored_img = restored_img.detach().cpu().numpy()
+
+        # Convert numpy to PIL if needed
+        if not isinstance(restored_img, Image.Image):
+            restored_img = Image.fromarray(restored_img[..., ::-1])  # BGR -> RGB
+
+        buf = io.BytesIO()
+        restored_img.save(buf, format="PNG")
+        encoded = base64.b64encode(buf.getvalue()).decode("utf-8")
+        print("✅ [POSTPROCESS] Image encoding complete.")
+        return {"image": encoded}
 

requirements.txt

@@ -1,6 +1,7 @@
 torch==2.1.0
 torchvision==0.16.0
 gfpgan==1.3.8
+realesrgan==0.3.0
 basicsr==1.4.2
 facexlib==0.3.0
 opencv-python