Update handler.py

handler.py

@@ -1,74 +1,74 @@
-import io, os, torch, base64
-from PIL import Image
+import base64
+import io
+import torch
 from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel
+from PIL import Image
+import os
 
-BASE_MODEL = "Gjm1234/juggernaut-sfw"
-CONTROLNET_MODEL = "diffusers/controlnet-depth-sdxl-1.0"   # ✅ RECOMMENDED WORKING MODEL
+HF_TOKEN = os.getenv("HF_TOKEN", None)
 
 class EndpointHandler:
-    def __init__(self, root=""):
-        print("🔧 Initializing Juggernaut + ControlNet")
-
-        token = os.environ.get("HF_TOKEN")
-        if not token:
-            raise RuntimeError("❌ Missing HF_TOKEN")
+    def __init__(self, path=""):
+        print("🔧 Initializing Juggernaut + ControlNet…")
 
-        # 🚫 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 …")
+        # Load ControlNet
+        print("📥 Loading ControlNet…")
         self.controlnet = ControlNetModel.from_pretrained(
-            CONTROLNET_MODEL,
+            "lllyasviel/controlnet-depth-sdxl-1.0",
             torch_dtype=torch.float16,
             use_safetensors=True,
-            token=token
+            token=HF_TOKEN
         )
 
-        print("📥 Loading Juggernaut XL (base)…")
+        # Load your big base model repo
+        print("📥 Loading Base Model (juggernaut-sfw)…")
         self.pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
-            BASE_MODEL,
+            "Gjm1234/juggernaut-sfw",
             controlnet=self.controlnet,
             torch_dtype=torch.float16,
             use_safetensors=True,
-            token=token
+            token=HF_TOKEN
         ).to("cuda")
 
-        self.pipe.enable_vae_slicing()
-        self.pipe.enable_attention_slicing()
-        self.pipe.unet.to(memory_format=torch.channels_last)
+        # prevent OOM
+        self.pipe.enable_model_cpu_offload()
 
-        print("✅ Juggernaut + ControlNet ready!")
+        print("🚀 Pipeline Loaded Successfully!")
 
     def __call__(self, data):
-
-        prompt = data.get("prompt", "")
-        negative = data.get("negative_prompt", "")
-
-        image_b64 = data.get("image")
-        input_image = None
-
-        if image_b64:
-            decoded = base64.b64decode(image_b64)
-            input_image = Image.open(io.BytesIO(decoded)).convert("RGB")
-
-        results = []
-
-        # Run 10 variations
-        for _ in range(10):
-            out = self.pipe(
+        try:
+            prompt = data.get("prompt", "")
+            negative_prompt = data.get("negative_prompt", "")
+            num_images = data.get("num_images", 10)
+
+            # Handle optional image
+            image_data = data.get("image", None)
+
+            if image_data:
+                # Base64 → PIL
+                image_bytes = base64.b64decode(image_data)
+                image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+            else:
+                # ControlNet requires something — generate depth = blank
+                image = Image.new("RGB", (1024, 1024), "white")
+
+            # Run generation
+            output = self.pipe(
                 prompt=prompt,
-                negative_prompt=negative,
-                image=input_image,
-                num_inference_steps=25,
-                guidance_scale=5.5,
-                width=1024,
-                height=1024
-            ).images[0]
-
-            buf = io.BytesIO()
-            out.save(buf, format="PNG")
-            results.append(base64.b64encode(buf.getvalue()).decode())
-
-        return { "images": results }
+                negative_prompt=negative_prompt,
+                image=image,
+                num_inference_steps=28,
+                guidance_scale=6.5,
+                num_images_per_prompt=num_images,
+            )
+
+            images = []
+            for img in output.images:
+                buf = io.BytesIO()
+                img.save(buf, format="PNG")
+                images.append(base64.b64encode(buf.getvalue()).decode("utf-8"))
+
+            return {"images": images}
+
+        except Exception as e:
+            return {"error": str(e)}