Update handler.py

handler.py

@@ -6,26 +6,35 @@ import base64
 
 class EndpointHandler:
     def __init__(self, path: str = ""):
-        print(f"Initializing SDXL model from: {path}")
+        print(f"🚀 Initializing Bh0r with Juggernaut-XL v9 as base model...")
 
-        # Load the base SDXL model
+        # Load Juggernaut-XL v9 instead of SDXL base
         self.pipe = StableDiffusionXLPipeline.from_pretrained(
-            "stabilityai/stable-diffusion-xl-base-1.0",
+            "RunDiffusion/Juggernaut-XL-v9",
             torch_dtype=torch.float16,
             variant="fp16"
         )
 
-        print("Loading LoRA weights from: Texttra/Bh0r")
+        print("✅ Juggernaut-XL v9 base model loaded successfully.")
+
+        # Load Bh0r LoRA
+        print("🧩 Loading Bh0r LoRA weights...")
         self.pipe.load_lora_weights(
             "Texttra/Bh0r",
             weight_name="Bh0r-10.safetensors",
             adapter_name="bh0r_lora"
         )
         self.pipe.set_adapters(["bh0r_lora"], adapter_weights=[0.9])
+
+        print("✅ Bh0r LoRA loaded with 0.9 weight.")
+
+        # Fuse LoRA into base model
         self.pipe.fuse_lora()
+        print("🔗 Fused LoRA into base model.")
 
+        # Move to GPU if available
         self.pipe.to("cuda" if torch.cuda.is_available() else "cpu")
-        print("Model ready.")
+        print("🎯 Model ready on device:", "cuda" if torch.cuda.is_available() else "cpu")
 
     def __call__(self, data: Dict) -> Dict:
         print("Received data:", data)
@@ -37,6 +46,7 @@ class EndpointHandler:
         if not prompt:
             return {"error": "No prompt provided."}
 
+        # Generate the image
         image = self.pipe(
             prompt,
             num_inference_steps=45,
@@ -45,6 +55,7 @@ class EndpointHandler:
 
         print("Image generated.")
 
+        # Convert to base64
         buffer = BytesIO()
         image.save(buffer, format="PNG")
         base64_image = base64.b64encode(buffer.getvalue()).decode("utf-8")