Initial Real-ESRGAN custom handler

README.md

@@ -0,0 +1,5 @@
+# Real-ESRGAN Custom Inference Handler
+
+Upscales images x4 using the Real-ESRGAN model.
+Provides a Hugging Face Inference Endpoint compatible `EndpointHandler`.
+

handler.py

@@ -0,0 +1,56 @@
+# handler.py
+from base64 import b64encode, b64decode
+from io import BytesIO
+from pathlib import Path
+from PIL import Image
+import numpy as np
+from realesrgan import RealESRGANer
+from basicsr.archs.rrdbnet_arch import RRDBNet
+
+class EndpointHandler:
+    def __init__(self, model_dir: str = "", **kwargs):
+        """
+        Called once when the endpoint starts.
+        Loads the Real-ESRGAN model weights.
+        """
+        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)
+        self.upsampler = RealESRGANer(
+            scale=4,
+            model_path=model_path,
+            model=rrdbnet,
+            tile=0,
+            pre_pad=0,
+            half=True,
+        )
+
+    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")
+
+        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)
+
+        buf = BytesIO()
+        out_img.save(buf, format="PNG")
+        encoded = b64encode(buf.getvalue()).decode("utf-8")
+
+        return {"image": encoded}
+

requirements.txt

@@ -0,0 +1,7 @@
+torch>=2.0.0
+torchvision>=0.15.0
+realesrgan>=0.3.0
+basicsr>=1.4.2
+numpy
+Pillow
+