Update handler.py

handler.py

@@ -7,41 +7,39 @@ import base64
 
 class EndpointHandler:
     def __init__(self, path: str = ""):
-        # Load base FLUX pipeline
         self.pipe = DiffusionPipeline.from_pretrained(
             "black-forest-labs/FLUX.1-dev",
             torch_dtype=torch.float16,
             use_auth_token=True
         )
 
-        # Load your LoRA weights from the repo
         self.pipe.load_lora_weights(path, weight_name="c1t3_v1.safetensors")
 
-        # Move to GPU if available
         if torch.cuda.is_available():
             self.pipe.to("cuda")
         else:
             self.pipe.to("cpu")
 
-        # Optional: enable memory optimization
         self.pipe.enable_model_cpu_offload()
 
-        # Initialize Compel (prompt parser for FLUX)
-        self.compel = Compel(tokenizer=self.pipe.tokenizer, text_encoder=self.pipe.text_encoder)
+        self.compel = Compel(
+            tokenizer=self.pipe.tokenizer,
+            text_encoder=self.pipe.text_encoder
+        )
 
-    def __call__(self, data: Dict[str, str]) -> Dict:
-    inputs = data.get("inputs", {})
-    prompt = inputs.get("prompt", "")
-    if not prompt:
-        return {"error": "No prompt provided."}
+    def __call__(self, data: Dict[str, Dict[str, str]]) -> Dict:
+        inputs = data.get("inputs", {})
+        prompt = inputs.get("prompt", "")
+        if not prompt:
+            return {"error": "No prompt provided."}
 
-    print(f"Received prompt: {prompt}")  # helpful logging
+        print(f"Received prompt: {prompt}")
 
-    conditioning = self.compel(prompt)
-    image = self.pipe(prompt_embeds=conditioning).images[0]
+        conditioning = self.compel(prompt)
+        image = self.pipe(prompt_embeds=conditioning).images[0]
 
-    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}