Update handler.py

handler.py

@@ -4,16 +4,19 @@ from diffusers import DiffusionPipeline
 from compel import Compel
 from io import BytesIO
 import base64
+import os
 
 class EndpointHandler:
     def __init__(self, path: str = ""):
+        print(f"Initializing model from: {path}")
         self.pipe = DiffusionPipeline.from_pretrained(
             "black-forest-labs/FLUX.1-dev",
             torch_dtype=torch.float16,
             use_auth_token=True
         )
-
-        self.pipe.load_lora_weights(path, weight_name="c1t3_v1.safetensors")
+        lora_path = os.path.join(path, "c1t3_v1.safetensors")
+        print(f"Loading LoRA weights from: {lora_path}")
+        self.pipe.load_lora_weights(lora_path)
 
         if torch.cuda.is_available():
             self.pipe.to("cuda")
@@ -26,20 +29,27 @@ class EndpointHandler:
             tokenizer=self.pipe.tokenizer,
             text_encoder=self.pipe.text_encoder
         )
+        print("Model initialized.")
+
+    def __call__(self, data: Dict) -> Dict:
+        print("Received data:", data)
 
-    def __call__(self, data: Dict[str, Dict[str, str]]) -> Dict:
         inputs = data.get("inputs", {})
         prompt = inputs.get("prompt", "")
+        print("Extracted prompt:", prompt)
+
         if not prompt:
             return {"error": "No prompt provided."}
 
-        print(f"Received prompt: {prompt}")
-
         conditioning = self.compel(prompt)
+        print("Conditioning complete.")
+
         image = self.pipe(prompt_embeds=conditioning).images[0]
+        print("Image generated.")
 
         buffer = BytesIO()
         image.save(buffer, format="PNG")
         base64_image = base64.b64encode(buffer.getvalue()).decode("utf-8")
+        print("Returning image.")
 
-        return {"image": base64_image}
+        return {"image": base64_image}