Fix handler to include EndpointHandler class

handler.py

@@ -12,19 +12,21 @@ class EndpointHandler:
     def __init__(self, path="."):
         print("🔹 Initializing Real-ESRGAN x4 model...")
 
-        self.model_url = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"
+        self.model_url = (
+            "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/"
+            "RealESRGAN_x4plus.pth"
+        )
         self.model_path = os.path.join(path, "RealESRGAN_x4plus.pth")
 
-        # Download model if not present
+        # Download model weights if missing
         if not os.path.exists(self.model_path):
-            print(f"📥 Downloading RealESRGAN_x4plus.pth from {self.model_url} ...")
+            print(f"📥 Downloading RealESRGAN_x4plus.pth ...")
             r = requests.get(self.model_url)
             r.raise_for_status()
             with open(self.model_path, "wb") as f:
                 f.write(r.content)
-            print(f"✅ Download complete: {self.model_path}")
+            print(f"✅ Downloaded model to {self.model_path}")
 
-        # Build Real-ESRGAN model
         model = RRDBNet(
             num_in_ch=3,
             num_out_ch=3,
@@ -45,13 +47,9 @@ class EndpointHandler:
         print("✅ Real-ESRGAN model initialized and ready.")
 
     # ==========================================================
-    # MAIN CALL METHOD
+    # Main callable
     # ==========================================================
     def __call__(self, data):
-        """
-        This is called automatically by Hugging Face Inference Toolkit.
-        It receives the raw image bytes from the POST body.
-        """
         try:
             image = self.preprocess(data)
             output = self.inference(image)
@@ -60,24 +58,39 @@ class EndpointHandler:
             return {"error": str(e)}
 
     # ==========================================================
-    # PREPROCESS / INFERENCE / POSTPROCESS
+    # Steps
     # ==========================================================
-    def preprocess(self, request):
-        """Convert raw bytes into a PIL image."""
-        if isinstance(request, (bytes, bytearray)):
-            return Image.open(io.BytesIO(request)).convert("RGB")
-        raise ValueError("Expected raw image bytes.")
+    def preprocess(self, data):
+        """Accept raw bytes OR dict-style payloads."""
+        # case 1: raw image bytes
+        if isinstance(data, (bytes, bytearray)):
+            return Image.open(io.BytesIO(data)).convert("RGB")
+
+        # case 2: dict with "inputs" key (HF default)
+        if isinstance(data, dict) and "inputs" in data:
+            img_field = data["inputs"]
+
+            # base64-encoded image
+            if isinstance(img_field, str):
+                try:
+                    return Image.open(io.BytesIO(base64.b64decode(img_field))).convert("RGB")
+                except Exception:
+                    raise ValueError("Invalid base64 image string in 'inputs'.")
+
+            # already bytes
+            if isinstance(img_field, (bytes, bytearray)):
+                return Image.open(io.BytesIO(img_field)).convert("RGB")
+
+        raise ValueError("Expected raw image bytes or {'inputs': <bytes/base64>}.")
 
     def inference(self, image):
-        """Run Real-ESRGAN model."""
         output, _ = self.upsampler.enhance(image, outscale=4)
         return output
 
     def postprocess(self, output_image):
-        """Convert output image to base64-encoded PNG for JSON."""
-        buffer = io.BytesIO()
-        output_image.save(buffer, format="PNG")
-        encoded = base64.b64encode(buffer.getvalue()).decode("utf-8")
-        buffer.close()
+        buf = io.BytesIO()
+        output_image.save(buf, format="PNG")
+        encoded = base64.b64encode(buf.getvalue()).decode("utf-8")
+        buf.close()
         return {"image": encoded}