add custom handler for swin2sr

handler.py

@@ -0,0 +1,34 @@
+from transformers import pipeline
+from PIL import Image
+import io, base64, requests
+
+# Load model once when the container starts
+def load_model():
+    print("Loading Swin2SR model...")
+    return pipeline("image-to-image", model="sergeipetrov/swin2SR-classical-sr-x2-64")
+
+model = load_model()
+
+# Hugging Face Inference Endpoint entrypoint
+def __call__(data):
+    # data can be dict with "inputs" (URL or base64)
+    image_input = data.get("inputs")
+
+    # If it's a URL
+    if isinstance(image_input, str) and image_input.startswith("http"):
+        image = Image.open(requests.get(image_input, stream=True).raw).convert("RGB")
+    else:
+        # assume it's base64
+        image_bytes = base64.b64decode(image_input)
+        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+
+    # Run inference
+    result = model(image)[0]
+
+    # Convert to base64 for API response
+    buffered = io.BytesIO()
+    result.save(buffered, format="PNG")
+    encoded_image = base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+    return {"image_base64": encoded_image}
+

requirements.txt

@@ -0,0 +1,5 @@
+transformers>=4.38.0
+torch
+Pillow
+requests
+