Update handler.py

handler.py

@@ -1,12 +1,9 @@
 import io, os, torch, base64
 from PIL import Image
-from diffusers import (
-    StableDiffusionXLPipeline,
-    ControlNetModel
-)
+from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel
 
-JUGGERNAUT_REPO = "Gjm1234/juggernaut-sfw"
-CONTROLNET_REPO = "thibaud/controlnet-openpose-sdxl-1.0"
+BASE_MODEL = "Gjm1234/juggernaut-sfw"
+CONTROLNET_MODEL = "diffusers/controlnet-depth-sdxl-1.0"   # ✅ RECOMMENDED WORKING MODEL
 
 class EndpointHandler:
     def __init__(self, root=""):
@@ -16,57 +13,62 @@ class EndpointHandler:
         if not token:
             raise RuntimeError("❌ Missing HF_TOKEN")
 
+        # 🚫 Disable flash/xformers/SDP
+        torch.backends.cuda.enable_flash_sdp(False)
+        torch.backends.cuda.enable_mem_efficient_sdp(False)
+        torch.backends.cuda.enable_math_sdp(True)
+
         print("📥 Loading ControlNet …")
         self.controlnet = ControlNetModel.from_pretrained(
-            CONTROLNET_REPO,
+            CONTROLNET_MODEL,
             torch_dtype=torch.float16,
             use_safetensors=True,
             token=token
-        ).to("cuda")
+        )
 
-        print("📥 Loading Juggernaut …")
-        self.pipe = StableDiffusionXLPipeline.from_pretrained(
-            JUGGERNAUT_REPO,
+        print("📥 Loading Juggernaut XL (base)…")
+        self.pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
+            BASE_MODEL,
+            controlnet=self.controlnet,
             torch_dtype=torch.float16,
             use_safetensors=True,
-            token=token,
-            controlnet=self.controlnet
+            token=token
         ).to("cuda")
 
-        # Memory optimizations
-        self.pipe.enable_attention_slicing()
         self.pipe.enable_vae_slicing()
+        self.pipe.enable_attention_slicing()
         self.pipe.unet.to(memory_format=torch.channels_last)
 
-        print("✅ Ready!")
+        print("✅ Juggernaut + ControlNet ready!")
 
     def __call__(self, data):
-        prompt = data.get("inputs", "")
-        neg = data.get("negative_prompt", "")
 
-        # Optional: base64 input image for editing
+        prompt = data.get("prompt", "")
+        negative = data.get("negative_prompt", "")
+
         image_b64 = data.get("image")
         input_image = None
 
         if image_b64:
-            img_bytes = base64.b64decode(image_b64)
-            input_image = Image.open(io.BytesIO(img_bytes)).convert("RGB")
+            decoded = base64.b64decode(image_b64)
+            input_image = Image.open(io.BytesIO(decoded)).convert("RGB")
 
         results = []
+
+        # Run 10 variations
         for _ in range(10):
-            output = self.pipe(
+            out = self.pipe(
                 prompt=prompt,
-                negative_prompt=neg,
+                negative_prompt=negative,
                 image=input_image,
-                controlnet_conditioning_scale=0.7,
                 num_inference_steps=25,
-                guidance_scale=7.5,
+                guidance_scale=5.5,
                 width=1024,
-                height=1024,
+                height=1024
             ).images[0]
 
             buf = io.BytesIO()
-            output.save(buf, format="PNG")
+            out.save(buf, format="PNG")
             results.append(base64.b64encode(buf.getvalue()).decode())
 
         return { "images": results }