g

handler.py

@@ -2,6 +2,7 @@ import os
 import io
 import torch
 import logging
+import base64
 import requests
 import numpy as np
 import cv2
@@ -20,7 +21,7 @@ class EndpointHandler:
         self.half = self.device == "cuda"
         self.path = path
 
-        # Model URLs
+        # URLs
         self.gfpgan_model_url = (
             "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth"
         )
@@ -28,31 +29,30 @@ class EndpointHandler:
             "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth"
         )
 
-        # Local paths
+        # Local model paths
         self.gfpgan_model_path = os.path.join(path, "GFPGANv1.4.pth")
         self.realesr_model_path = os.path.join(path, "realesr-general-x4v3.pth")
 
-        # Lazy init
         self.bg_upsampler = None
         self.restorer = None
 
-        # Ensure model files exist
         self._ensure_model(self.gfpgan_model_url, self.gfpgan_model_path)
         self._ensure_model(self.realesr_model_url, self.realesr_model_path)
-
         logger.info(f"🧠 Device: {self.device}, half precision: {self.half}")
 
-    def _ensure_model(self, url, local_path):
-        if not os.path.exists(local_path):
-            logger.info(f"⬇️ Downloading {url}")
+    def _ensure_model(self, url, path):
+        if not os.path.exists(path):
+            logger.info(f"⬇️ Downloading model from {url}")
             r = requests.get(url, timeout=60)
             r.raise_for_status()
-            with open(local_path, "wb") as f:
+            with open(path, "wb") as f:
                 f.write(r.content)
-            logger.info(f"✅ Saved to {local_path}")
+            logger.info(f"✅ Model saved to {path}")
+        else:
+            logger.info(f"📁 Found cached model: {path}")
 
     def _init_models(self):
-        """Lazy-load GFPGAN and Real-ESRGAN models."""
+        """Lazy-load models"""
         if self.bg_upsampler is None:
             logger.info("🧩 Initializing Real-ESRGAN upsampler...")
             model = SRVGGNetCompact(
@@ -81,35 +81,56 @@ class EndpointHandler:
             )
             logger.info("✅ Models ready!")
 
-    def __call__(self, data):
-        """Restore a face photo."""
-        self._init_models()
-
+    def _load_image(self, data):
+        """Handle different input formats."""
         if isinstance(data, dict) and "inputs" in data:
             data = data["inputs"]
 
-        # Load image
         if isinstance(data, (bytes, bytearray)):
-            image = Image.open(io.BytesIO(data)).convert("RGB")
-        elif isinstance(data, str) and data.startswith("http"):
-            resp = requests.get(data)
-            image = Image.open(io.BytesIO(resp.content)).convert("RGB")
-        elif isinstance(data, Image.Image):
-            image = data
-        else:
-            raise ValueError("Unsupported input type")
+            logger.info("📦 Received raw bytes input")
+            return Image.open(io.BytesIO(data)).convert("RGB")
+
+        if isinstance(data, str):
+            if data.startswith("http"):
+                logger.info(f"🌐 Downloading image from URL: {data}")
+                resp = requests.get(data)
+                return Image.open(io.BytesIO(resp.content)).convert("RGB")
+            else:
+                # assume base64
+                logger.info("🧬 Decoding base64 image input")
+                try:
+                    decoded = base64.b64decode(data)
+                    return Image.open(io.BytesIO(decoded)).convert("RGB")
+                except Exception as e:
+                    logger.error(f"❌ Failed to decode base64: {e}")
+                    raise ValueError("Invalid base64 image input")
+
+        raise ValueError("Unsupported input type")
 
+    def __call__(self, data):
+        self._init_models()
+        logger.info("⚙️ Starting inference...")
+
+        # Load input
+        image = self._load_image(data)
         input_img = np.array(image, dtype=np.uint8)
 
+        logger.info(f"📏 Input image shape: {input_img.shape}")
+
         cropped_faces, restored_faces, restored_img = self.restorer.enhance(
             input_img, has_aligned=False, only_center_face=False, paste_back=True
         )
 
+        logger.info("🖼️ Restoration complete, encoding output...")
+
+        # Encode result as base64
         _, buffer = cv2.imencode(".jpg", restored_img)
-        output_bytes = io.BytesIO(buffer.tobytes())
+        b64_output = base64.b64encode(buffer).decode("utf-8")
 
+        logger.info("✅ Returning base64 image JSON")
         return {
-            "output": output_bytes,
+            "image": b64_output,
+            "status": "success",
             "info": "Restored with GFPGAN v1.4 + Real-ESRGAN x4v3"
         }