test

handler.py

@@ -1,31 +1,28 @@
+from typing import Dict, Any
 from transformers import pipeline
 from PIL import Image
 import io, base64, requests
 
 class EndpointHandler:
-    def __init__(self, path=""):
-        # This runs once when the container starts
-        print("🔹 Loading Swin2SR model...")
+    def __init__(self, model_dir: str, **kwargs: Any):
+        print("🔹 Loading Swin2SR model ...")
         self.model = pipeline("image-to-image", model="sergeipetrov/swin2SR-classical-sr-x2-64")
 
-    def __call__(self, data):
-        # Handle input data
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
         image_input = data.get("inputs")
         if not image_input:
-            return {"error": "No 'inputs' field provided."}
+            return {"error": "Missing 'inputs' field"}
 
-        # Accept both image URLs and base64 strings
+        # Accept either a URL or a base64 string
         if isinstance(image_input, str) and image_input.startswith("http"):
             image = Image.open(requests.get(image_input, stream=True).raw).convert("RGB")
         else:
             image_bytes = base64.b64decode(image_input)
             image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
 
-        # Run inference
         result = self.model(image)[0]
 
-        # Encode to base64 for response
-        buffer = io.BytesIO()
-        result.save(buffer, format="PNG")
-        return {"image_base64": base64.b64encode(buffer.getvalue()).decode("utf-8")}
+        buf = io.BytesIO()
+        result.save(buf, format="PNG")
+        return {"image_base64": base64.b64encode(buf.getvalue()).decode("utf-8")}