fix everything

handler.py

@@ -2,33 +2,30 @@ 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}
+class EndpointHandler:
+    def __init__(self, path=""):
+        # This runs once when the container starts
+        print("🔹 Loading Swin2SR model...")
+        self.model = pipeline("image-to-image", model="sergeipetrov/swin2SR-classical-sr-x2-64")
+
+    def __call__(self, data):
+        # Handle input data
+        image_input = data.get("inputs")
+        if not image_input:
+            return {"error": "No 'inputs' field provided."}
+
+        # Accept both image URLs and base64 strings
+        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")}