Update handler.py

handler.py

@@ -1,6 +1,7 @@
 import io
 import base64
 import torch
+import os
 from PIL import Image
 from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel
 
@@ -9,30 +10,36 @@ CONTROLNET = "lllyasviel/controlnet-depth-sdxl-1.0"
 
 class EndpointHandler:
     def __init__(self, path=""):
-        print("🔧 Initializing handler — loading remote models...")
+        print("🔧 Initializing handler…")
 
-        print("🔧 Loading ControlNet...")
+        HF_TOKEN = os.environ.get("HF_TOKEN")
+        if not HF_TOKEN:
+            raise RuntimeError("❌ HF_TOKEN not found in environment variables")
+
+        print("🔧 Loading ControlNet with token…")
         controlnet = ControlNetModel.from_pretrained(
             CONTROLNET,
-            torch_dtype=torch.float16
+            torch_dtype=torch.float16,
+            token=HF_TOKEN
         )
 
-        print("🚀 Loading Juggernaut XL main model...")
+        print("🚀 Loading Juggernaut XL model with token…")
         self.pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
             BASE_MODEL,
             controlnet=controlnet,
             torch_dtype=torch.float16,
-            use_safetensors=True
+            use_safetensors=True,
+            token=HF_TOKEN
         ).to("cuda")
 
         self.pipe.enable_xformers_memory_efficient_attention()
-        print("✅ Pipeline ready")
+        print("✅ Pipeline loaded successfully!")
 
     def __call__(self, data):
         prompt = data.get("inputs", "")
         img_b64 = data.get("image", None)
 
-        # Decode input image OR generate blank white one
+        # Decode or create blank input
         if img_b64:
             img_bytes = base64.b64decode(img_b64)
             init = Image.open(io.BytesIO(img_bytes)).convert("RGB")
@@ -40,18 +47,16 @@ class EndpointHandler:
             init = Image.new("RGB", (1024, 1024), "white")
 
         outputs = []
-        for _ in range(10):  # always 10 variations
+        for _ in range(10):  # ALWAYS generate 10 images
             result = self.pipe(
                 prompt=prompt,
                 image=init,
-                num_inference_steps=25,
+                num_inference_steps=20,
                 guidance_scale=6.0,
-                width=1024,
-                height=1024,
             ).images[0]
 
-            buf = io.BytesIO()
-            result.save(buf, format="PNG")
-            outputs.append(base64.b64encode(buf.getvalue()).decode())
+            buffer = io.BytesIO()
+            result.save(buffer, format="PNG")
+            outputs.append(base64.b64encode(buffer.getvalue()).decode())
 
         return { "images": outputs }