fix1.4

handler.py

@@ -1,138 +1,189 @@
-import os
-import io
-import torch
 import base64
+import io
+import json
+import logging
+import os
+from typing import Any, Dict
+
+import numpy as np
 import requests
 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("🚀 [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.3.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)
+# -----------------------------------------------------------------------------
+# Logging setup
+# -----------------------------------------------------------------------------
+logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+
+# -----------------------------------------------------------------------------
+# Model paths and URLs
+# -----------------------------------------------------------------------------
+REPO_DIR = os.environ.get("HF_HOME", "/repository")
+
+GFPGAN_WEIGHTS_PATH = os.path.join(REPO_DIR, "GFPGANv1.4.pth")
+REAL_ESRGAN_WEIGHTS_PATH = os.path.join(REPO_DIR, "realesr-general-x4v3.pth")
+
+GFPGAN_URLS = [
+    # ✅ working file (GFPGANv1.4.pth is hosted under v1.3.0 release tag)
+    "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth",
+    "https://github.com/TencentARC/GFPGAN/releases/download/v1.4.0/GFPGANv1.4.pth",
+]
+
+# ✅ updated Real-ESRGAN v0.2.5.0 URLs
+REAL_ESRGAN_URLS = [
+    "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth",
+    "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth",
+]
+
+# -----------------------------------------------------------------------------
+# Helpers
+# -----------------------------------------------------------------------------
+def _ensure_file(path: str, urls) -> None:
+    if os.path.exists(path) and os.path.getsize(path) > 0:
+        logger.debug(f"✅ File exists: {path}")
+        return
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    last_err = None
+    for u in urls:
+        try:
+            logger.debug(f"⬇️  Downloading {u}")
+            r = requests.get(u, timeout=60)
             r.raise_for_status()
-            with open(self.model_path, "wb") as f:
+            with open(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
+            logger.debug(f"✅ Saved to {path}")
+            return
+        except Exception as e:
+            last_err = e
+            logger.warning(f"⚠️ Failed from {u}: {e}")
+    raise RuntimeError(f"❌ Could not download required file: {last_err}")
+
+def _to_bgr(image_bytes: bytes) -> np.ndarray:
+    pil = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+    return np.array(pil)[:, :, ::-1].copy()
+
+def _encode_bgr(bgr_img: np.ndarray) -> str:
+    import cv2
+    success, buf = cv2.imencode(".png", bgr_img)
+    if not success:
+        raise RuntimeError("Failed to encode image.")
+    return base64.b64encode(buf.tobytes()).decode("utf-8")
+
+# -----------------------------------------------------------------------------
+# EndpointHandler
+# -----------------------------------------------------------------------------
+class EndpointHandler:
+    """
+    Custom handler for GFPGAN v1.4 + Real-ESRGAN (realesr-general-x4v3.pth)
+    Emulates the behavior of the official Gradio demo.
+    """
+
+    def __init__(self, path: str = REPO_DIR):
+        logger.debug("🚀 [INIT] Starting GFPGAN + Real-ESRGAN handler...")
+        logger.debug(f"📂 Repository path: {path}")
+
+        # 1️⃣ Ensure model weights
+        _ensure_file(GFPGAN_WEIGHTS_PATH, GFPGAN_URLS)
+        _ensure_file(REAL_ESRGAN_WEIGHTS_PATH, REAL_ESRGAN_URLS)
+
+        # 2️⃣ Device setup
+        import torch
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.half = torch.cuda.is_available()
+        logger.debug(f"🧠 Device: {self.device}, half precision: {self.half}")
+
+        # 3️⃣ Build Real-ESRGAN upsampler (x4)
+        logger.debug("🧩 Initializing Real-ESRGAN background upsampler...")
+        rrdb = 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,
+            scale=4,
+            model_path=REAL_ESRGAN_WEIGHTS_PATH,
+            model=rrdb,
+            tile=0,
             tile_pad=10,
             pre_pad=0,
-            half=False,
-            device="cuda" if torch.cuda.is_available() else "cpu",
+            half=self.half,
+            device=self.device,
         )
-        print("✅ [INIT] Real-ESRGAN background upsampler ready.")
+        logger.debug("✅ Real-ESRGAN upsampler ready (x4).")
 
-        # ------------------------------
-        # Setup GFPGANer
-        # ------------------------------
-        print("🧠 [INIT] Setting up GFPGANer (v1.4)...")
+        # 4️⃣ Build GFPGAN restorer (v1.4)
+        logger.debug("🧩 Initializing GFPGAN v1.4 restorer...")
         self.restorer = GFPGANer(
-            model_path=self.model_path,
+            model_path=GFPGAN_WEIGHTS_PATH,
             upscale=2,
             arch="clean",
             channel_multiplier=2,
             bg_upsampler=self.bg_upsampler,
-            device="cuda" if torch.cuda.is_available() else "cpu",
+            device=self.device,
         )
-        print("✅ [INIT DONE] GFPGAN + RealESRGAN hybrid handler ready.")
-
-    # ----------------------------------------------------------
-    # Main inference entry point
-    # ----------------------------------------------------------
-    def __call__(self, data):
-        print("🛰️ [CALL] Endpoint invoked!")
-        print(f"📦 [CALL] Raw input type: {type(data)}")
+        logger.debug("✅ GFPGAN v1.4 initialized.")
 
+    # -------------------------------------------------------------------------
+    # Inference
+    # -------------------------------------------------------------------------
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
         try:
-            image = self.preprocess(data)
-            print("🧩 [STEP] Image preprocessed successfully.")
-            restored = self.inference(image)
-            print("🎨 [STEP] Inference completed successfully.")
-            return self.postprocess(restored)
+            logger.debug(f"🌀 Received data type: {type(data)}")
+
+            # --- Parse input ---
+            image_bytes = None
+            parameters = {}
+
+            if isinstance(data, (bytes, bytearray)):
+                image_bytes = bytes(data)
+            elif isinstance(data, dict):
+                b64 = data.get("inputs") or data.get("image")
+                if b64:
+                    image_bytes = base64.b64decode(b64)
+                parameters = data.get("parameters") or {}
+            elif isinstance(data, str):
+                try:
+                    parsed = json.loads(data)
+                    b64 = parsed.get("inputs") or parsed.get("image")
+                    if b64:
+                        image_bytes = base64.b64decode(b64)
+                    parameters = parsed.get("parameters") or {}
+                except Exception as e:
+                    logger.warning(f"⚠️ JSON parse error: {e}")
+
+            if not image_bytes:
+                return {"error": "Unsupported input format — expected bytes or base64 data"}
+
+            scale = int(parameters.get("scale", 2))
+            logger.debug(f"🔧 Using scale factor: {scale}")
+
+            # Convert to BGR
+            bgr_input = _to_bgr(image_bytes)
+            logger.debug(f"📏 Input image shape: {bgr_input.shape}")
+
+            # Enhance
+            logger.debug("✨ Running GFPGAN restoration...")
+            _, _, restored_img = self.restorer.enhance(
+                bgr_input, has_aligned=False, only_center_face=False, paste_back=True
+            )
+
+            if restored_img is None:
+                return {"error": "Restoration failed — no output."}
+
+            b64_img = _encode_bgr(restored_img)
+            logger.debug("✅ Restoration complete.")
+            return {"image": b64_img}
+
         except Exception as e:
-            print(f"💥 [ERROR] Exception during call: {str(e)}")
-            return {"error": str(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")
-
-        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")
-
-        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}
+            logger.exception("🔥 Inference error")
+            return {"error": f"{type(e).__name__}: {e}"}