Add verified Real-ESRGAN handler with auto-download

handler.py

@@ -1,56 +1,78 @@
-# handler.py
-from base64 import b64encode, b64decode
-from io import BytesIO
-from pathlib import Path
-from PIL import Image
+import os
+import io
+import torch
 import numpy as np
+import requests
+import cv2
+from PIL import Image
 from realesrgan import RealESRGANer
 from basicsr.archs.rrdbnet_arch import RRDBNet
 
+# ======================================================
+# CONFIG
+# ======================================================
+MODEL_URL = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"
+MODEL_PATH = "/repository/RealESRGAN_x4plus.pth"
+
+# ======================================================
+# HANDLER CLASS
+# ======================================================
 class EndpointHandler:
-    def __init__(self, model_dir: str = "", **kwargs):
-        """
-        Called once when the endpoint starts.
-        Loads the Real-ESRGAN model weights.
-        """
+    def __init__(self, model_dir):
         print("🔹 Initializing Real-ESRGAN x4 model...")
-        model_path = str(Path(model_dir) / "RealESRGAN_x4plus.pth")
 
-        # Build model
-        rrdbnet = RRDBNet(num_in_ch=3, num_out_ch=3)
+        # Ensure the model weights exist
+        if not os.path.exists(MODEL_PATH):
+            print(f"📥 Downloading RealESRGAN_x4plus.pth from {MODEL_URL} ...")
+            response = requests.get(MODEL_URL, stream=True)
+            response.raise_for_status()
+            with open(MODEL_PATH, "wb") as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    f.write(chunk)
+            print("✅ Download complete:", MODEL_PATH)
+        else:
+            print("✅ Model file already exists:", MODEL_PATH)
+
+        # Define RRDBNet model
+        self.model = RRDBNet(
+            num_in_ch=3, num_out_ch=3, num_feat=64,
+            num_block=23, num_grow_ch=32, scale=4
+        )
+
+        # Initialize Real-ESRGAN upsampler
         self.upsampler = RealESRGANer(
             scale=4,
-            model_path=model_path,
-            model=rrdbnet,
-            tile=0,
+            model_path=MODEL_PATH,
+            model=self.model,
+            tile=0,              # Disable tile mode for simplicity
+            tile_pad=10,
             pre_pad=0,
-            half=True,
+            half=False           # Disable FP16 for stability
         )
 
+        print("✅ Real-ESRGAN model initialized and ready.")
+
     def __call__(self, data):
-        """
-        Called for each request.
-        Expects a dict with 'inputs' = base64-encoded image or bytes.
-        Returns base64-encoded upscaled image.
-        """
-        image = data.get("inputs")
-
-        if isinstance(image, str):
-            image = Image.open(BytesIO(b64decode(image)))
-        elif isinstance(image, bytes):
-            image = Image.open(BytesIO(image))
-        else:
-            raise ValueError("Input must be base64 string or bytes")
+        print("🚀 Received inference request...")
+
+        # Get input image
+        image_bytes = data.get("inputs") or data.get("image")
+        if image_bytes is None:
+            raise ValueError("❌ No image data found in request payload.")
+
+        # Convert bytes to RGB numpy array
+        if isinstance(image_bytes, list):
+            image_bytes = bytes(image_bytes)
+        img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        img = np.array(img)
 
-        image = np.array(image)
-        image = image[:, :, ::-1]  # RGB→BGR
-        output, _ = self.upsampler.enhance(image, outscale=4)
-        output = output[:, :, ::-1]  # BGR→RGB
-        out_img = Image.fromarray(output)
+        # Run Real-ESRGAN enhancement
+        output, _ = self.upsampler.enhance(img, outscale=4)
+        print("✨ Image enhancement complete.")
 
-        buf = BytesIO()
-        out_img.save(buf, format="PNG")
-        encoded = b64encode(buf.getvalue()).decode("utf-8")
+        # Convert to PNG bytes for return
+        _, buffer = cv2.imencode(".png", cv2.cvtColor(output, cv2.COLOR_RGB2BGR))
+        print("📦 Returning processed image bytes.")
 
-        return {"image": encoded}
+        return {"image": buffer.tobytes()}