Update handler.py

handler.py

@@ -1,77 +1,89 @@
-from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel
 import torch
+import os
+import gc
+from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel
 from PIL import Image
 import base64
 from io import BytesIO
-import os
+
 
 class EndpointHandler:
     def __init__(self, model_dir):
-        print("🔑 Loading HF token...")
-        hf_token = os.getenv("HF_TOKEN")
+        print("🔧 Initializing improved memory-safe handler...")
+
+        # Prevent cuda fragmentation
+        os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
 
-        print("🔧 Loading ControlNet...")
+        # HF token
+        token = os.getenv("HF_TOKEN", None)
+
+        # Load ControlNet
+        print("🔧 Loading ControlNet…")
         controlnet = ControlNetModel.from_pretrained(
-            "diffusers/controlnet-depth-sdxl-1.0",
+            os.path.join(model_dir, "controlnet"),
             torch_dtype=torch.float16,
             use_safetensors=True,
-            token=hf_token
+            token=token
         )
 
-        print("🧠 Loading Juggernaut XL...")
-        base_model = "Gjm1234/juggernaut-sfw"
-
+        # Load main model
+        print("🔧 Loading Juggernaut XL…")
         self.pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
-            base_model,
+            os.path.join(model_dir, "model"),
             controlnet=controlnet,
             torch_dtype=torch.float16,
             use_safetensors=True,
-            token=hf_token
-        ).to("cuda")
+            token=token
+        )
 
-        # IMPORTANT FIX — remove xformers, use PyTorch attention instead
-        if hasattr(self.pipe, "enable_model_cpu_offload"):
-            self.pipe.enable_model_cpu_offload()
+        # 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 loaded successfully!")
+        print("✅ Pipeline ready!")
 
     def __call__(self, data):
-        print("📥 Received request...")
-
-        prompt = data.get("prompt", "")
-        negative_prompt = data.get("negative_prompt", "blurry, bad quality, distorted, extra limbs")
-
-        num_images = 10
-
-        # Decode uploaded image
-        image_b64 = data.get("image")
-        if image_b64:
-            print("🖼️ Decoding input image...")
-            image_data = base64.b64decode(image_b64)
-            init_image = Image.open(BytesIO(image_data)).convert("RGB")
-        else:
-            print("⚠️ No image uploaded — generating blank control input.")
-            init_image = Image.new("RGB", (1024, 1024), "white")
-
-        print("🎨 Generating images...")
-        output = self.pipe(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            image=init_image,
-            num_inference_steps=30,
-            num_images_per_prompt=num_images,
-            guidance_scale=7.0,
-        )
+        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
+            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
+            ).images[0]
+
+            # Convert output to base64
+            buffered = BytesIO()
+            result.save(buffered, format="JPEG")
+            output_b64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
 
-        images = output.images
+            return {"image": output_b64}
 
-        print("📤 Encoding output images...")
-        result = []
-        for img in images:
-            buffer = BytesIO()
-            img.save(buffer, format="PNG")
-            b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
-            result.append(b64)
+        except Exception as e:
+            return {"error": str(e)}
 
-        print("✅ Returning images...")
-        return {"images": result}
+        finally:
+            # 🔥 Force GPU/CPU memory cleanup
+            torch.cuda.empty_cache()
+            gc.collect()