Update main.py

main.py

@@ -1,5 +1,6 @@
 from fastapi import FastAPI, HTTPException, BackgroundTasks, Depends
 from pydantic import BaseModel
+from fastapi.staticfiles import StaticFiles
 from fastapi.concurrency import run_in_threadpool
 from fastapi.middleware.cors import CORSMiddleware
 import uuid
@@ -42,6 +43,9 @@ chat_sessions = {}
 image_storage_dir = tempfile.mkdtemp()
 print(f"Temporary directory for images: {image_storage_dir}")
 
+# Serve the temporary image directory as a static file directory
+app.mount("/images", StaticFiles(directory=image_storage_dir), name="images")
+
 # Dictionary to store images temporarily
 image_store = {}
 
@@ -142,7 +146,6 @@ async def generate_image(request: ImageRequest):
             "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")
@@ -168,7 +171,10 @@ async def generate_image(request: ImageRequest):
                         img = Image.open(BytesIO(image_response.content))
                         filename = f"generated_{uuid.uuid4()}.png"
                         filepath = save_image_locally(img, filename)
-                        return filepath
+
+                        # Generate a URL for the image
+                        file_url = f"/images/{filename}"
+                        return filepath, file_url
                     else:
                         raise HTTPException(status_code=500, detail="Failed to download the image")
                 else:
@@ -176,16 +182,15 @@ async def generate_image(request: ImageRequest):
             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)
+    filepath, file_url = await run_in_threadpool(future.result)
 
     return {
         "status": "Image generated successfully",
         "file_path": filepath,
+        "file_url": file_url,
     }
 
-
 @app.post("/upscale-image", response_model=dict)
 async def upscale_image(image_url: str, background_tasks: BackgroundTasks):
     def process_image():