Fix handler to include EndpointHandler class

handler.py

@@ -8,96 +8,76 @@ from realesrgan import RealESRGANer
 from basicsr.archs.rrdbnet_arch import RRDBNet
 
 
-# ==========================================================
-# MODEL INITIALIZATION
-# ==========================================================
-print("🔹 Initializing Real-ESRGAN x4 model...")
-
-MODEL_URL = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"
-MODEL_PATH = "/repository/RealESRGAN_x4plus.pth"
-
-# Download model if missing
-if not os.path.exists(MODEL_PATH):
-    print(f"📥 Downloading RealESRGAN_x4plus.pth from {MODEL_URL} ...")
-    r = requests.get(MODEL_URL)
-    r.raise_for_status()
-    with open(MODEL_PATH, "wb") as f:
-        f.write(r.content)
-    print(f"✅ Download complete: {MODEL_PATH}")
-
-# Create model architecture
-model = RRDBNet(
-    num_in_ch=3,
-    num_out_ch=3,
-    num_feat=64,
-    num_block=23,
-    num_grow_ch=32,
-    scale=4
-)
-
-# Load Real-ESRGAN model
-upsampler = RealESRGANer(
-    scale=4,
-    model_path=MODEL_PATH,
-    model=model,
-    half=False,
-    device="cuda" if torch.cuda.is_available() else "cpu"
-)
-
-print("✅ Real-ESRGAN model initialized and ready.")
-
-
-# ==========================================================
-# HANDLER FUNCTIONS
-# ==========================================================
-def preprocess(request):
-    """
-    Converts request body (bytes) into a PIL image.
-    """
-    if isinstance(request, (bytes, bytearray)):
-        image = Image.open(io.BytesIO(request)).convert("RGB")
-        return image
-    else:
-        raise ValueError("Request body must be raw image bytes.")
-
-
-def inference(image: Image.Image):
-    """
-    Runs Real-ESRGAN on the input PIL image.
-    """
-    try:
-        output, _ = upsampler.enhance(image, outscale=4)
+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_path = os.path.join(path, "RealESRGAN_x4plus.pth")
+
+        # Download model if not present
+        if not os.path.exists(self.model_path):
+            print(f"📥 Downloading RealESRGAN_x4plus.pth from {self.model_url} ...")
+            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}")
+
+        # Build Real-ESRGAN model
+        model = RRDBNet(
+            num_in_ch=3,
+            num_out_ch=3,
+            num_feat=64,
+            num_block=23,
+            num_grow_ch=32,
+            scale=4,
+        )
+
+        self.upsampler = RealESRGANer(
+            scale=4,
+            model_path=self.model_path,
+            model=model,
+            half=False,
+            device="cuda" if torch.cuda.is_available() else "cpu",
+        )
+
+        print("✅ Real-ESRGAN model initialized and ready.")
+
+    # ==========================================================
+    # MAIN CALL METHOD
+    # ==========================================================
+    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)
+            return self.postprocess(output)
+        except Exception as e:
+            return {"error": str(e)}
+
+    # ==========================================================
+    # PREPROCESS / INFERENCE / POSTPROCESS
+    # ==========================================================
+    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 inference(self, image):
+        """Run Real-ESRGAN model."""
+        output, _ = self.upsampler.enhance(image, outscale=4)
         return output
-    except Exception as e:
-        raise RuntimeError(f"Inference failed: {e}")
-
-
-def postprocess(output_image: Image.Image):
-    """
-    Converts a PIL Image to base64-encoded PNG string for JSON output.
-    """
-    buffer = io.BytesIO()
-    output_image.save(buffer, format="PNG")
-    image_bytes = buffer.getvalue()
-    buffer.close()
-
-    encoded = base64.b64encode(image_bytes).decode("utf-8")
-    return {"image": encoded}
-
 
-# ==========================================================
-# MAIN ENTRY POINT
-# ==========================================================
-def __call__(self, data):
-    """
-    This method is called automatically by Hugging Face Inference Toolkit.
-    `data` is the raw HTTP request body (image bytes).
-    """
-    try:
-        image = preprocess(data)
-        output = inference(image)
-        result = postprocess(output)
-        return result
-    except Exception as e:
-        return {"error": str(e)}
+    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()
+        return {"image": encoded}