Update handler.py

handler.py

@@ -11,7 +11,7 @@ class EndpointHandler:
         self.pipe = DiffusionPipeline.from_pretrained(
             "black-forest-labs/FLUX.1-dev",
             torch_dtype=torch.float16,
-            use_auth_token=True  # Required for gated base model
+            use_auth_token=True
         )
 
         print("Loading LoRA weights from: Texttra/Cityscape_Studio")
@@ -33,9 +33,32 @@ class EndpointHandler:
     def __call__(self, data: Dict) -> Dict:
         print("Received data:", data)
 
-        inputs = data.get("inputs", {})
-        prompt = inputs.get("prompt", "")
-        print("Extracted prompt:", prompt)
+        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", "")
 
-        if not prompt:
-            return {"error": "No prompt provided"}
+            print("Extracted prompt:", prompt)
+
+            if not prompt:
+                return {"error": "No prompt provided"}
+
+            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}
+
+        except Exception as e:
+            print(f"Error occurred: {str(e)}")
+            return {"error": str(e)}