Update main.py

main.py

@@ -99,7 +99,7 @@ async def new_chat():
     return {"chat_id": chat_id}
 
 @app.post("/generate-image", response_model=dict)
-async def generate_image(request: ImageRequest, background_tasks: BackgroundTasks):
+async def generate_image(request: ImageRequest):
     def process_request():
         chat_history = get_chat_history(request.chat_id)
         prompt = f"""
@@ -145,9 +145,24 @@ async def generate_image(request: ImageRequest, background_tasks: BackgroundTask
         filepath = save_image_locally(img, filename)
         return filepath
 
-    # Add the wrapper function to background tasks
-    background_tasks.add_task(lambda: executor.submit(process_request))
-    return {"status": "Processing"}
+    # 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()
+
+    # 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):