Update handler.py

handler.py

@@ -1,53 +1,26 @@
-import torch
 from diffusers import DiffusionPipeline
-from io import BytesIO
-import base64
+import torch
 
 class EndpointHandler:
-    def __init__(self, path="."):
-        # Load model
+    def __init__(self, path=""):
+        print("Loading Juggernaut XL…")
         self.pipe = DiffusionPipeline.from_pretrained(
             path,
-            torch_dtype=torch.float16,
-            use_safetensors=True
+            torch_dtype=torch.float16
         ).to("cuda")
 
-        self.pipe.enable_model_cpu_offload()
-        self.pipe.enable_vae_tiling()
-
-    def img_to_b64(self, img):
-        buffer = BytesIO()
-        img.save(buffer, format="JPEG", quality=90)
-        return base64.b64encode(buffer.getvalue()).decode("utf-8")
-
     def __call__(self, data):
-        # HF ALWAYS sends "inputs"
         prompt = data.get("inputs", "")
-
         params = data.get("parameters", {})
-        negative = params.get("negative_prompt", "")
+
         steps = params.get("num_inference_steps", 28)
         cfg = params.get("guidance_scale", 4.5)
-        width = params.get("width", 1024)
-        height = params.get("height", 1024)
-        num_images = params.get("num_images", 1)
 
-        # Clamp for safety
-        if num_images > 10:
-            num_images = 10
-
-        # Run inference
         result = self.pipe(
-            prompt=prompt,
-            negative_prompt=negative,
+            prompt,
             num_inference_steps=steps,
-            guidance_scale=cfg,
-            width=width,
-            height=height,
-            num_images_per_prompt=num_images
-        ).images
-
-        # Convert images → base64
-        output = [self.img_to_b64(img) for img in result]
+            guidance_scale=cfg
+        )
 
-        return {"images": output}
+        # MUST return a PIL image
+        return result.images[0]