Update handler.py

handler.py

@@ -11,12 +11,14 @@ class EndpointHandler:
         self.pipe = DiffusionPipeline.from_pretrained(
             "black-forest-labs/FLUX.1-dev",
             torch_dtype=torch.float16,
-            use_auth_token=True
+            use_auth_token=True  # Required for gated base model
         )
 
+        # Load LoRA weights from your Hugging Face repo
         print("Loading LoRA weights from: Texttra/Cityscape_Studio")
         self.pipe.load_lora_weights("Texttra/Cityscape_Studio", weight_name="c1t3_v1.safetensors")
 
+        # Send to GPU if available
         if torch.cuda.is_available():
             self.pipe.to("cuda")
         else:
@@ -24,6 +26,7 @@ class EndpointHandler:
 
         self.pipe.enable_model_cpu_offload()
 
+        # Initialize Compel for prompt conditioning
         self.compel = Compel(
             tokenizer=self.pipe.tokenizer,
             text_encoder=self.pipe.text_encoder
@@ -33,32 +36,29 @@ class EndpointHandler:
     def __call__(self, data: Dict) -> Dict:
         print("Received data:", data)
 
-        try:
-            inputs = data.get("inputs", {})
-            if isinstance(inputs, str):
-                # In case the input comes in raw string form (e.g., Postman tests)
-                prompt = inputs
-            else:
-                prompt = inputs.get("prompt", "")
+        inputs = data.get("inputs", {})
+        prompt = inputs.get("prompt", "")
+        print("Extracted prompt:", prompt)
 
-            print("Extracted prompt:", prompt)
+        if not prompt:
+            return {"error": "No prompt provided"}
 
-            if not prompt:
-                return {"error": "No prompt provided"}
+        # Generate both prompt and pooled embeddings
+        conditioning, pooled = self.compel(prompt, return_pooled=True)
+        print("Conditioning complete.")
 
-            conditioning = self.compel(prompt)
-            print("Conditioning complete.")
+        # Run the model
+        image = self.pipe(
+            prompt_embeds=conditioning,
+            pooled_prompt_embeds=pooled
+        ).images[0]
+        print("Image generated.")
 
-            image = self.pipe(prompt_embeds=conditioning).images[0]
-            print("Image generated.")
+        # Encode image to base64
+        buffer = BytesIO()
+        image.save(buffer, format="PNG")
+        base64_image = base64.b64encode(buffer.getvalue()).decode("utf-8")
+        print("Returning image.")
 
-            buffer = BytesIO()
-            image.save(buffer, format="PNG")
-            base64_image = base64.b64encode(buffer.getvalue()).decode("utf-8")
+        return {"image": base64_image}
 
-            print("Returning image.")
-            return {"image": base64_image}
-
-        except Exception as e:
-            print(f"Error occurred: {str(e)}")
-            return {"error": str(e)}