Update handler.py

handler.py

@@ -30,20 +30,18 @@ class EndpointHandler:
         self.compel = Compel(tokenizer=self.pipe.tokenizer, text_encoder=self.pipe.text_encoder)
 
     def __call__(self, data: Dict[str, str]) -> Dict:
-        # Get prompt from request data
-        prompt = data.get("prompt", "")
-        if not prompt:
-            return {"error": "No prompt provided."}
+    inputs = data.get("inputs", {})
+    prompt = inputs.get("prompt", "")
+    if not prompt:
+        return {"error": "No prompt provided."}
 
-        # Generate prompt conditioning using Compel
-        conditioning = self.compel(prompt)
+    print(f"Received prompt: {prompt}")  # helpful logging
 
-        # Generate image using FLUX + LoRA
-        image = self.pipe(prompt_embeds=conditioning).images[0]
+    conditioning = self.compel(prompt)
+    image = self.pipe(prompt_embeds=conditioning).images[0]
 
-        # Convert image to base64 string for API response
-        buffer = BytesIO()
-        image.save(buffer, format="PNG")
-        base64_image = base64.b64encode(buffer.getvalue()).decode("utf-8")
+    buffer = BytesIO()
+    image.save(buffer, format="PNG")
+    base64_image = base64.b64encode(buffer.getvalue()).decode("utf-8")
 
-        return {"image": base64_image}
+    return {"image": base64_image}