Update handler.py

handler.py

@@ -1,62 +1,72 @@
-import io
-import base64
-import torch
-import os
+import io, os, torch, base64
 from PIL import Image
-from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel
+from diffusers import (
+    StableDiffusionXLPipeline,
+    ControlNetModel
+)
 
-BASE_MODEL = "Gjm1234/juggernaut-sfw"
-CONTROLNET = "lllyasviel/controlnet-depth-sdxl-1.0"
+JUGGERNAUT_REPO = "Gjm1234/juggernaut-sfw"
+CONTROLNET_REPO = "thibaud/controlnet-openpose-sdxl-1.0"
 
 class EndpointHandler:
-    def __init__(self, path=""):
-        print("🔧 Initializing handler…")
+    def __init__(self, root=""):
+        print("🔧 Initializing Juggernaut + ControlNet")
 
-        HF_TOKEN = os.environ.get("HF_TOKEN")
-        if not HF_TOKEN:
-            raise RuntimeError("❌ HF_TOKEN not found in environment variables")
+        token = os.environ.get("HF_TOKEN")
+        if not token:
+            raise RuntimeError("❌ Missing HF_TOKEN")
 
-        print("🔧 Loading ControlNet with token…")
-        controlnet = ControlNetModel.from_pretrained(
-            CONTROLNET,
+        print("📥 Loading ControlNet …")
+        self.controlnet = ControlNetModel.from_pretrained(
+            CONTROLNET_REPO,
             torch_dtype=torch.float16,
-            token=HF_TOKEN
-        )
+            use_safetensors=True,
+            token=token
+        ).to("cuda")
 
-        print("🚀 Loading Juggernaut XL model with token…")
-        self.pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
-            BASE_MODEL,
-            controlnet=controlnet,
+        print("📥 Loading Juggernaut …")
+        self.pipe = StableDiffusionXLPipeline.from_pretrained(
+            JUGGERNAUT_REPO,
             torch_dtype=torch.float16,
             use_safetensors=True,
-            token=HF_TOKEN
+            token=token,
+            controlnet=self.controlnet
         ).to("cuda")
 
-        self.pipe.enable_xformers_memory_efficient_attention()
-        print("✅ Pipeline loaded successfully!")
+        # Memory optimizations
+        self.pipe.enable_attention_slicing()
+        self.pipe.enable_vae_slicing()
+        self.pipe.unet.to(memory_format=torch.channels_last)
+
+        print("✅ Ready!")
 
     def __call__(self, data):
         prompt = data.get("inputs", "")
-        img_b64 = data.get("image", None)
-
-        # Decode or create blank input
-        if img_b64:
-            img_bytes = base64.b64decode(img_b64)
-            init = Image.open(io.BytesIO(img_bytes)).convert("RGB")
-        else:
-            init = Image.new("RGB", (1024, 1024), "white")
-
-        outputs = []
-        for _ in range(10):  # ALWAYS generate 10 images
-            result = self.pipe(
+        neg = data.get("negative_prompt", "")
+
+        # Optional: base64 input image for editing
+        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")
+
+        results = []
+        for _ in range(10):
+            output = self.pipe(
                 prompt=prompt,
-                image=init,
-                num_inference_steps=20,
-                guidance_scale=6.0,
+                negative_prompt=neg,
+                image=input_image,
+                controlnet_conditioning_scale=0.7,
+                num_inference_steps=25,
+                guidance_scale=7.5,
+                width=1024,
+                height=1024,
             ).images[0]
 
-            buffer = io.BytesIO()
-            result.save(buffer, format="PNG")
-            outputs.append(base64.b64encode(buffer.getvalue()).decode())
+            buf = io.BytesIO()
+            output.save(buf, format="PNG")
+            results.append(base64.b64encode(buf.getvalue()).decode())
 
-        return { "images": outputs }
+        return { "images": results }