Update handler.py

handler.py

@@ -1,89 +1,57 @@
+import io
+import base64
 import torch
-import os
-import gc
-from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel
 from PIL import Image
-import base64
-from io import BytesIO
+from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel
 
+BASE_MODEL = "Gjm1234/juggernaut-sfw"
+CONTROLNET = "lllyasviel/controlnet-depth-sdxl-1.0"
 
 class EndpointHandler:
-    def __init__(self, model_dir):
-        print("🔧 Initializing improved memory-safe handler...")
+    def __init__(self, path=""):
+        print("🔧 Initializing handler — loading remote models...")
 
-        # Prevent cuda fragmentation
-        os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
-
-        # HF token
-        token = os.getenv("HF_TOKEN", None)
-
-        # Load ControlNet
-        print("🔧 Loading ControlNet…")
+        print("🔧 Loading ControlNet...")
         controlnet = ControlNetModel.from_pretrained(
-            os.path.join(model_dir, "controlnet"),
-            torch_dtype=torch.float16,
-            use_safetensors=True,
-            token=token
+            CONTROLNET,
+            torch_dtype=torch.float16
         )
 
-        # Load main model
-        print("🔧 Loading Juggernaut XL…")
+        print("🚀 Loading Juggernaut XL main model...")
         self.pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
-            os.path.join(model_dir, "model"),
+            BASE_MODEL,
             controlnet=controlnet,
             torch_dtype=torch.float16,
-            use_safetensors=True,
-            token=token
-        )
+            use_safetensors=True
+        ).to("cuda")
 
-        # VRAM-saving settings
-        self.pipe.to("cuda")
-        self.pipe.enable_attention_slicing()
-        self.pipe.enable_vae_slicing()
-        self.pipe.enable_sequential_cpu_offload()
-
-        print("✅ Pipeline ready!")
+        self.pipe.enable_xformers_memory_efficient_attention()
+        print("✅ Pipeline ready")
 
     def __call__(self, data):
-        try:
-            prompt = data.get("prompt", "")
-            image_b64 = data.get("image", None)
-
-            if not prompt:
-                return {"error": "Missing prompt"}
-
-            if not image_b64:
-                return {"error": "Missing image input"}
-
-            # Decode ControlNet image
-            try:
-                image_bytes = base64.b64decode(image_b64)
-                control_image = Image.open(BytesIO(image_bytes)).convert("RGB")
-            except:
-                return {"error": "Invalid base64 image"}
-
-            # Run the pipeline
+        prompt = data.get("inputs", "")
+        img_b64 = data.get("image", None)
+
+        # Decode input image OR generate blank white one
+        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 10 variations
             result = self.pipe(
                 prompt=prompt,
-                image=control_image,
-                num_inference_steps=20,
-                guidance_scale=3.0,
-                controlnet_conditioning_scale=1.0,
-                height=768,
-                width=512
+                image=init,
+                num_inference_steps=25,
+                guidance_scale=6.0,
+                width=1024,
+                height=1024,
             ).images[0]
 
-            # Convert output to base64
-            buffered = BytesIO()
-            result.save(buffered, format="JPEG")
-            output_b64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
-
-            return {"image": output_b64}
-
-        except Exception as e:
-            return {"error": str(e)}
+            buf = io.BytesIO()
+            result.save(buf, format="PNG")
+            outputs.append(base64.b64encode(buf.getvalue()).decode())
 
-        finally:
-            # 🔥 Force GPU/CPU memory cleanup
-            torch.cuda.empty_cache()
-            gc.collect()
+        return { "images": outputs }