Update handler.py

handler.py

@@ -1,47 +1,74 @@
 import base64
-from io import BytesIO
-from diffusers import StableDiffusionXLPipeline
-import torch
+import io
 from PIL import Image
+import torch
+from diffusers import StableDiffusionXLPipeline
+from typing import Any, Dict
+
 
 class EndpointHandler:
-    def __init__(self, path=""):
-        print("🔧 Loading Juggernaut-SFW model...")
+    def __init__(self, model_dir: str, **kwargs):
+        print("🔥 Initializing Juggernaut XL Handler (Prompt + Optional Image)...")
+
+        # Load XL model from your big repo
         self.pipe = StableDiffusionXLPipeline.from_pretrained(
             "Gjm1234/juggernaut-sfw",
             torch_dtype=torch.float16,
+            use_safetensors=True
         ).to("cuda")
 
-        # Disable controlnet entirely (you are not loading any ControlNet weights)
-        print("⚠️ ControlNet disabled — no weights provided.")
-
-        # Memory optimisation
         self.pipe.enable_attention_slicing()
-        self.pipe.enable_model_cpu_offload()
 
-    def __call__(self, data):
-        prompt = data.get("prompt", None)
-        if prompt is None:
+        print("✅ Pipeline loaded successfully.")
+
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+
+        # Must receive `inputs`
+        if "inputs" not in data:
+            return {"error": "Body must contain 'inputs' object"}
+
+        inputs = data["inputs"]
+
+        prompt = inputs.get("prompt", None)
+        if not prompt:
             return {"error": "prompt is required"}
 
-        num_images = int(data.get("num_images", 4))
-        if num_images < 1:
-            num_images = 1
+        num_images = inputs.get("num_images", 10)
+        image_b64 = inputs.get("image", None)
+
+        init_image = None
+        if image_b64:
+            try:
+                img_bytes = base64.b64decode(image_b64)
+                init_image = Image.open(io.BytesIO(img_bytes)).convert("RGB")
+            except Exception as e:
+                return {"error": f"Invalid image data: {str(e)}"}
 
-        print("🎨 Generating images…")
+        # Run txt2img OR img2img depending on whether image was sent
+        if init_image is None:
+            print("🎨 Running TEXT → IMAGE")
+            output = self.pipe(
+                prompt=prompt,
+                num_images_per_prompt=num_images
+            )
+        else:
+            print("🎨 Running IMAGE → IMAGE")
+            output = self.pipe(
+                prompt=prompt,
+                image=init_image,
+                strength=0.6,
+                num_images_per_prompt=num_images
+            )
 
-        images = self.pipe(
-            prompt=prompt,
-            num_inference_steps=25,
-            guidance_scale=6,
-            num_images_per_prompt=num_images
-        ).images
+        images = output.images
 
-        # Encode all images to Base64
-        encoded = []
+        # convert to base64 array
+        results = []
         for img in images:
-            buf = BytesIO()
-            img.save(buf, format="PNG")
-            encoded.append(base64.b64encode(buf.getvalue()).decode("utf-8"))
+            buffered = io.BytesIO()
+            img.save(buffered, format="PNG")
+            img_b64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
+            results.append(img_b64)
 
-        return {"images": encoded}
+        print(f"✅ Returning {len(results)} images.")
+        return {"images": results}