Update handler.py

handler.py

@@ -11,13 +11,14 @@ class EndpointHandler:
         self.pipe = DiffusionPipeline.from_pretrained(
             "black-forest-labs/FLUX.1-dev",
             torch_dtype=torch.float16,
-            use_auth_token=True  # Safe since base model is gated
+            use_auth_token=True  # Required for gated base model
         )
 
-        # ✅ Load LoRA directly from Hugging Face Hub
+        # ✅ Load LoRA from Hugging Face Hub (or local if path is set)
         print("Loading LoRA weights from: Texttra/Cityscape_Studio")
         self.pipe.load_lora_weights("Texttra/Cityscape_Studio", weight_name="c1t3_v1.safetensors")
 
+        # Move to GPU if available
         if torch.cuda.is_available():
             self.pipe.to("cuda")
         else:
@@ -41,10 +42,14 @@ class EndpointHandler:
         if not prompt:
             return {"error": "No prompt provided."}
 
-        conditioning = self.compel(prompt)
+        # ✅ FLUX requires both prompt_embeds and pooled_prompt_embeds
+        conditioning, pooled = self.compel(prompt, return_pooled=True)
         print("Conditioning complete.")
 
-        image = self.pipe(prompt_embeds=conditioning).images[0]
+        image = self.pipe(
+            prompt_embeds=conditioning,
+            pooled_prompt_embeds=pooled
+        ).images[0]
         print("Image generated.")
 
         buffer = BytesIO()