Add extreme logging to GFPGAN handler for debugging

handler.py

@@ -1,115 +1,174 @@
 import os
 import io
-import torch
 import base64
+import torch
 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"
+        print("🚀 [INIT] Starting GFPGAN EndpointHandler initialization...")
+        print(f"📁 Working directory: {os.getcwd()}")
+        print(f"📂 Handler path argument: {path}")
+
+        # ----------------------------
+        # MODEL DOWNLOAD
+        # ----------------------------
+        self.model_url = (
+            "https://github.com/TencentARC/GFPGAN/releases/download/v1.4.0/GFPGANv1.4.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)
+        self.model_path = os.path.join(path, "GFPGANv1.4.pth")
+        print(f"🔗 [MODEL] Model URL: {self.model_url}")
+        print(f"📦 [MODEL] Local model path: {self.model_path}")
+
+        if not os.path.exists(self.model_path):
+            print("📥 [DOWNLOAD] Model not found locally — starting download...")
+            try:
+                r = requests.get(self.model_url, timeout=60)
+                r.raise_for_status()
+                with open(self.model_path, "wb") as f:
+                    f.write(r.content)
+                print(f"✅ [DOWNLOAD] Model successfully saved to {self.model_path}")
+            except Exception as e:
+                print(f"💥 [ERROR] Failed to download GFPGAN weights: {e}")
+                raise e
+        else:
+            print("📚 [CACHE] Found existing model file, skipping download.")
+
+        # ----------------------------
+        # INITIALIZE RESTORER
+        # ----------------------------
         device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"💻 [DEVICE] Using device: {device}")
 
-        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,
-        )
+        try:
+            print("🧠 [MODEL INIT] Initializing GFPGANer...")
+            self.restorer = GFPGANer(
+                model_path=self.model_path,
+                upscale=2,
+                arch="clean",
+                channel_multiplier=2,
+                bg_upsampler=None,
+                device=device,
+            )
+            print("✅ [MODEL INIT] GFPGANer successfully initialized.")
+        except Exception as e:
+            print(f"💥 [ERROR] Failed to initialize GFPGANer: {e}")
+            raise e
 
-        print("✅ GFPGANv1 Face Restoration Pipeline Ready!")
+        print("🎉 [INIT DONE] GFPGAN model ready for inference!")
 
-    # ------------------------------------------------------------------
-    # Main callable
-    # ------------------------------------------------------------------
+    # ----------------------------
+    # CALLABLE ENTRY POINT
+    # ----------------------------
     def __call__(self, data):
+        print("\n🛰️ [CALL] ======= Incoming request =======")
+        print(f"🔍 [CALL] Input type: {type(data)}")
+        print(f"🔍 [CALL] Raw preview: {str(data)[:300]}")
+
         try:
-            image = self._load_image(data)
-            restored = self._restore_face(image)
-            return self._encode_image(restored)
+            image = self.preprocess(data)
+            print(f"📸 [PREPROCESS DONE] Type: {type(image)}, Size: {image.size}")
+
+            restored = self.inference(image)
+            print(f"🎨 [INFERENCE DONE] Output type: {type(restored)}, Size: {restored.size}")
+
+            response = self.postprocess(restored)
+            print(f"📦 [POSTPROCESS DONE] Keys: {list(response.keys())}")
+            print("✅ [CALL COMPLETE] Successfully processed request.\n")
+            return response
+
         except Exception as e:
-            print("💥 Error:", str(e))
+            print("💥 [CALL ERROR] Exception occurred during inference!")
+            print(f"❗ [TRACEBACK]: {repr(e)}")
             return {"error": str(e)}
 
-    # ------------------------------------------------------------------
-    # Helper functions
-    # ------------------------------------------------------------------
-    def _load_image(self, data):
-        # Handles raw bytes, base64 JSON, or dict
+    # ----------------------------
+    # PREPROCESS STEP
+    # ----------------------------
+    def preprocess(self, data):
+        print("🔧 [PREPROCESS] Starting preprocessing step...")
+        print(f"🔧 [PREPROCESS] Received type: {type(data)}")
+
+        # Case 1: raw bytes directly
         if isinstance(data, (bytes, bytearray)):
-            return Image.open(io.BytesIO(data)).convert("RGB")
+            print("📥 [PREPROCESS] Detected raw bytes input.")
+            img = Image.open(io.BytesIO(data)).convert("RGB")
+            print(f"📸 [PREPROCESS] Loaded image — size: {img.size}, mode: {img.mode}")
+            return img
 
+        # Case 2: dict payload (like {'inputs': base64_string})
         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")
+            print("📦 [PREPROCESS] Detected dict input.")
+            keys = list(data.keys())
+            print(f"📦 [PREPROCESS] Dict keys: {keys}")
+            img_field = data.get("inputs") or data.get("image") or None
+
+            if img_field is None:
+                raise ValueError("Missing 'inputs' or 'image' key in JSON payload.")
+
+            # Base64 string case
+            if isinstance(img_field, str):
+                print("🧬 [PREPROCESS] Input is a base64 string — decoding...")
+                img_field = base64.b64decode(img_field)
+
+            # Raw bytes inside dict
+            if isinstance(img_field, (bytes, bytearray)):
+                print("📥 [PREPROCESS] Input is bytes inside dict — converting to image...")
+                img = Image.open(io.BytesIO(img_field)).convert("RGB")
+                print(f"📸 [PREPROCESS] Image loaded from dict — size: {img.size}, mode: {img.mode}")
+                return img
+
+        raise ValueError("Unsupported input format — expected bytes or base64 data")
+
+    # ----------------------------
+    # INFERENCE STEP
+    # ----------------------------
+    def inference(self, image):
+        print("🧠 [INFERENCE] Starting restoration process...")
+        print(f"🧠 [INFERENCE] Input image type: {type(image)}, size: {image.size}, mode: {image.mode}")
 
-        if isinstance(data, str):
-            decoded = base64.b64decode(data)
-            return Image.open(io.BytesIO(decoded)).convert("RGB")
+        try:
+            np_img = np.array(image)
+            print(f"🔢 [INFERENCE] NumPy array shape: {np_img.shape}, dtype: {np_img.dtype}")
 
-        raise ValueError("Expected image bytes or base64 string.")
+            # Convert RGB → BGR
+            np_img = np_img[:, :, ::-1]
+            print("🎨 [INFERENCE] Converted RGB to BGR for OpenCV-style processing.")
 
-    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
+            # Run the GFPGAN restoration
+            print("🚀 [INFERENCE] Running GFPGANer.enhance()...")
+            _, _, restored_img = self.restorer.enhance(
+                np_img, has_aligned=False, only_center_face=False, paste_back=True
+            )
 
-    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}
+            if restored_img is None:
+                raise RuntimeError("Restoration failed — GFPGAN returned None")
+
+            print("✅ [INFERENCE] Restoration successful, converting to PIL.Image...")
+            restored_pil = Image.fromarray(restored_img[:, :, ::-1])  # Convert BGR → RGB
+            print(f"🖼️ [INFERENCE] Restored image size: {restored_pil.size}, mode: {restored_pil.mode}")
+            return restored_pil
+
+        except Exception as e:
+            print(f"💥 [INFERENCE ERROR] {e}")
+            raise e
+
+    # ----------------------------
+    # POSTPROCESS STEP
+    # ----------------------------
+    def postprocess(self, restored_img):
+        print("📦 [POSTPROCESS] Starting encoding...")
+        try:
+            buf = io.BytesIO()
+            restored_img.save(buf, format="PNG")
+            encoded = base64.b64encode(buf.getvalue()).decode("utf-8")
+            print("✅ [POSTPROCESS] Image successfully encoded to base64 (len=%d)" % len(encoded))
+            return {"image": encoded}
+        except Exception as e:
+            print(f"💥 [POSTPROCESS ERROR] {e}")
+            raise e
 

requirements.txt

@@ -1,9 +1,6 @@
-torch==2.1.2
-torchvision==0.16.2
+torch>=2.1.0
+torchvision>=0.16.0
 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