Update main.py

main.py

@@ -118,53 +118,73 @@ async def generate_image(request: ImageRequest):
         - **Color Theme**: Follow the **{request.color_theme}** to set the mood and tone for the entire scene.
         - Negative Prompt: Avoid grainy, blurry, or deformed outputs.
         - **Text to Include in Image**: Clearly display the text **"{request.text}"** as part of the composition (e.g., on a card, badge, or banner) attached to the subject in a realistic and contextually appropriate way.
-        - Write the prompt only in the output. Do not include anything else except the prompt. Do not write Generated Image Prompt: in the output.
-
-        Chat History:
-        {chat_history}
-
-        User Prompt:
-        {request.user_prompt}
-
-        Generated Image Prompt:
         """
         refined_prompt = llm.invoke(prompt).content.strip()
+
+        # Save the prompt in MongoDB
         collection.insert_one({"session_id": request.chat_id, "role": "user", "content": request.user_prompt})
         collection.insert_one({"session_id": request.chat_id, "role": "ai", "content": refined_prompt})
 
-        # Simulate image generation
-        response = requests.post(
-            "https://api.bfl.ml/v1/flux-pro-1.1",
-            json={"prompt": refined_prompt}
-        ).json()
-        image_url = response["result"]["sample"]
-
-        # Download and save the image locally
-        image_response = requests.get(image_url)
-        img = Image.open(BytesIO(image_response.content))
-        filename = f"generated_{uuid.uuid4()}.png"
-        filepath = save_image_locally(img, filename)
-        return filepath
+        # API call to image generation service
+        url = "https://api.bfl.ml/v1/flux-pro-1.1"
+        headers = {
+            "accept": "application/json",
+            "x-key": "4f69d408-4979-4812-8ad2-ec6d232c9ddf",
+            "Content-Type": "application/json"
+        }
+        payload = {
+            "prompt": refined_prompt,
+            "width": 1024,
+            "height": 1024,
+            "guidance_scale": 1,
+            "num_inference_steps": 50,
+            "max_sequence_length": 512,
+        }
+
+        # Initial request to generate the image
+        response = requests.post(url, headers=headers, json=payload).json()
+        if "id" not in response:
+            raise HTTPException(status_code=500, detail="Error generating image: ID missing from response")
+
+        request_id = response["id"]
+
+        # Poll for the image result
+        while True:
+            time.sleep(0.5)
+            result = requests.get(
+                "https://api.bfl.ml/v1/get_result",
+                headers=headers,
+                params={"id": request_id},
+            ).json()
+
+            if result["status"] == "Ready":
+                if "result" in result and "sample" in result["result"]:
+                    image_url = result["result"]["sample"]
+
+                    # Download the image
+                    image_response = requests.get(image_url)
+                    if image_response.status_code == 200:
+                        img = Image.open(BytesIO(image_response.content))
+                        filename = f"generated_{uuid.uuid4()}.png"
+                        filepath = save_image_locally(img, filename)
+                        return filepath
+                    else:
+                        raise HTTPException(status_code=500, detail="Failed to download the image")
+                else:
+                    raise HTTPException(status_code=500, detail="Expected 'sample' key not found in the result")
+            elif result["status"] == "Error":
+                raise HTTPException(status_code=500, detail=f"Image generation failed: {result.get('error', 'Unknown error')}")
 
     # Run the request processing in a thread
     future = executor.submit(process_request)
-    filepath = await run_in_threadpool(future.result)  # Wait for the task to complete
-
-    # Load the image to return as a response
-    with open(filepath, "rb") as f:
-        image_data = f.read()
+    filepath = await run_in_threadpool(future.result)
 
-    # Convert the file path into a downloadable format
-    file_url = f"/images/{os.path.basename(filepath)}"
-
-    # Return the response
     return {
         "status": "Image generated successfully",
         "file_path": filepath,
-        "file_url": file_url,
-        "image": image_data,
     }
 
+
 @app.post("/upscale-image", response_model=dict)
 async def upscale_image(image_url: str, background_tasks: BackgroundTasks):
     def process_image():