Update api.py

api.py

@@ -1,12 +1,18 @@
 import io
 import random
 import os
+import uuid
 import requests
 from fastapi import FastAPI, HTTPException
-from fastapi.responses import StreamingResponse
+from fastapi.responses import JSONResponse
 from fastapi.concurrency import run_in_threadpool
+from fastapi.staticfiles import StaticFiles
 from PIL import Image
 
+# Directory to store generated images
+IMAGE_DIR = "generated_images"
+os.makedirs(IMAGE_DIR, exist_ok=True)
+
 # API Configuration for Hugging Face
 API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
 headers = {"Authorization": f"Bearer {os.getenv('HF')}"}
@@ -18,7 +24,7 @@ def query(payload):
     return response.content
 
 def generate_image(prompt: str) -> Image.Image:
-    # Use a random seed so even identical prompts produce different images.
+    # Use a random seed so identical prompts produce different images.
     random_seed = random.randint(0, 999999)
     payload = {
         "inputs": prompt,
@@ -29,19 +35,32 @@ def generate_image(prompt: str) -> Image.Image:
 
 app = FastAPI()
 
+# Mount the static file directory to serve generated images.
+app.mount("/generated_images", StaticFiles(directory=IMAGE_DIR), name="generated_images")
+
+# Define your base URL (adjust as needed or set via environment variable)
+BASE_URL = os.getenv("BASE_URL", "https://bowozzz-gen.hf.space")
+
 @app.get("/")
 async def get_generated_image(text: str):
     """
-    API endpoint that waits for the image to be generated and returns it.
-    Example usage:
-        https://bowozzz-gen.hf.space/api?text=jokowi
+    API endpoint that generates an image for the given text prompt,
+    saves it, and returns a link to download the image.
+    Example: https://bowozzz-gen.hf.space/api?text=jokowi
     """
     try:
-        # Run the generate_image function in a threadpool so that this async endpoint waits for it.
+        # Run the blocking image generation in a threadpool
         image = await run_in_threadpool(generate_image, text)
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
-    img_buffer = io.BytesIO()
-    image.save(img_buffer, format="PNG")
-    img_buffer.seek(0)
-    return StreamingResponse(img_buffer, media_type="image/png")
+    
+    # Generate a unique filename and save the image to the static folder
+    filename = f"{uuid.uuid4().hex}.png"
+    file_path = os.path.join(IMAGE_DIR, filename)
+    image.save(file_path, format="PNG")
+    
+    # Construct the download URL using the base URL and static folder mount
+    download_url = f"{BASE_URL}/generated_images/{filename}"
+    
+    # Return the download URL in a JSON response
+    return JSONResponse(content={"download_url": download_url})