Fix GFPGAN weights URL (v1.3.0 release tag)

handler.py

@@ -1,174 +1,137 @@
-import os
 import io
-import base64
+import os
+import cv2
 import torch
+import base64
+import logging
 import requests
 import numpy as np
 from PIL import Image
 from gfpgan import GFPGANer
 
+# ======================================================
+#  LOGGING CONFIGURATION
+# ======================================================
+logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+
+logger.debug("📦 [INIT] Importing GFPGAN handler module...")
+
 
+# ======================================================
+#  GFPGAN MODEL DOWNLOAD CONFIG
+# ======================================================
+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="."):
-        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.3.0/GFPGANv1.4.pth"
-        )
-        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...")
+        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 not found
+        # ------------------------------------------------------
+        if not os.path.exists(model_path):
             try:
-                r = requests.get(self.model_url, timeout=60)
+                logger.debug(f"📥 [DOWNLOAD] Model not found locally. Downloading from {MODEL_URL}")
+                r = requests.get(MODEL_URL, stream=True)
                 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}")
+                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 model successfully.")
             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}")
+                logger.error(f"💥 [ERROR] Failed to download GFPGAN weights: {e}")
+                raise
 
+        # ------------------------------------------------------
+        # Initialize GFPGANer just like the official Gradio demo
+        # ------------------------------------------------------
         try:
-            print("🧠 [MODEL INIT] Initializing GFPGANer...")
+            logger.debug("🧠 [MODEL] Initializing GFPGANer (v1.4, upscale=2, no background upsampler)...")
             self.restorer = GFPGANer(
-                model_path=self.model_path,
-                upscale=2,
+                model_path=model_path,
+                upscale=2,                   # Matches Gradio demo “Rescaling factor: 2”
                 arch="clean",
                 channel_multiplier=2,
                 bg_upsampler=None,
-                device=device,
+                version="1.4",
             )
-            print("✅ [MODEL INIT] GFPGANer successfully initialized.")
+            logger.debug("✅ [MODEL] GFPGAN model initialized successfully.")
         except Exception as e:
-            print(f"💥 [ERROR] Failed to initialize GFPGANer: {e}")
-            raise e
+            logger.error(f"💥 [ERROR] Model initialization failed: {e}")
+            raise
 
-        print("🎉 [INIT DONE] GFPGAN model ready for inference!")
-
-    # ----------------------------
-    # CALLABLE ENTRY POINT
-    # ----------------------------
+    # ======================================================
+    #  INFERENCE CALL
+    # ======================================================
     def __call__(self, data):
-        print("\n🛰️ [CALL] ======= Incoming request =======")
-        print(f"🔍 [CALL] Input type: {type(data)}")
-        print(f"🔍 [CALL] Raw preview: {str(data)[:300]}")
+        logger.debug("⚙️ [INFER] Starting inference...")
+        logger.debug(f"📥 Incoming data keys: {list(data.keys())}")
 
+        # ------------------------------------------------------
+        # Handle both raw bytes and JSON base64 inputs
+        # ------------------------------------------------------
         try:
-            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
-
+            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")
         except Exception as e:
-            print("💥 [CALL ERROR] Exception occurred during inference!")
-            print(f"❗ [TRACEBACK]: {repr(e)}")
-            return {"error": str(e)}
-
-    # ----------------------------
-    # 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)):
-            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):
-            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}")
+            logger.error(f"💥 [ERROR] Input parsing failed: {e}")
+            return {"error": f"Invalid input: {e}"}
 
+        # ------------------------------------------------------
+        # Load and convert image
+        # ------------------------------------------------------
         try:
-            np_img = np.array(image)
-            print(f"🔢 [INFERENCE] NumPy array shape: {np_img.shape}, dtype: {np_img.dtype}")
-
-            # Convert RGB → BGR
-            np_img = np_img[:, :, ::-1]
-            print("🎨 [INFERENCE] Converted RGB to BGR for OpenCV-style processing.")
+            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}"}
 
-            # 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
+        # ------------------------------------------------------
+        # Perform face restoration
+        # ------------------------------------------------------
+        try:
+            cropped_faces, restored_faces, restored_img = self.restorer.enhance(
+                img_np,
+                has_aligned=False,
+                only_center_face=False,
+                paste_back=True,  # ⬅️ Matches web demo behavior
             )
-
-            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
-
+            logger.debug("✅ [RESTORE] Face restoration complete.")
         except Exception as e:
-            print(f"💥 [INFERENCE ERROR] {e}")
-            raise e
-
-    # ----------------------------
-    # POSTPROCESS STEP
-    # ----------------------------
-    def postprocess(self, restored_img):
-        print("📦 [POSTPROCESS] Starting encoding...")
+            logger.error(f"💥 [ERROR] GFPGAN enhancement failed: {e}")
+            return {"error": f"Enhancement failed: {e}"}
+
+        # ------------------------------------------------------
+        # Encode restored image to base64 for response
+        # ------------------------------------------------------
         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}
+            _, buffer = cv2.imencode(".png", restored_img[:, :, ::-1])  # BGR → RGB
+            img_base64 = base64.b64encode(buffer).decode("utf-8")
+            logger.debug("📤 [ENCODE] Successfully encoded restored image.")
+            return {"image": img_base64}
         except Exception as e:
-            print(f"💥 [POSTPROCESS ERROR] {e}")
-            raise e
+            logger.error(f"💥 [ERROR] Failed to encode restored image: {e}")
+            return {"error": f"Encoding failed: {e}"}