Update api.py

api.py

@@ -1,5 +1,5 @@
 from fastapi import FastAPI, HTTPException, Request
-from fastapi.responses import HTMLResponse, JSONResponse
+from fastapi.responses import HTMLResponse, JSONResponse, FileResponse
 from fastapi.middleware.cors import CORSMiddleware
 import threading
 import gradio as gr
@@ -7,6 +7,8 @@ import torch
 import numpy as np
 from diffusers import DiffusionPipeline
 from transformers import pipeline
+from PIL import Image
+import os
 
 app = FastAPI()
 
@@ -15,7 +17,7 @@ origins = [
     "http://localhost",
     "http://localhost:8000",
     "http://localhost:7860",
-    "https://nikajoon-test1.hf.space",
+    "https://nikajoon-test1.hf.space",  # URL فضای شما در Hugging Face
     "http://tomko.ir",  # دامنه سایت شما
 ]
 
@@ -67,9 +69,18 @@ async def generate_image(request: Request):
         raise HTTPException(status_code=400, detail="Prompt is required")
     
     image, _ = genie(prompt)
-    # ذخیره‌سازی تصویر و بازگرداندن URL یا داده‌های تصویر به عنوان پاسخ
-    image.save("output.png")
-    return JSONResponse(content={"image_url": "output.png"})
+    image_path = "images/output.png"
+    os.makedirs(os.path.dirname(image_path), exist_ok=True)
+    image.save(image_path)
+    
+    return JSONResponse(content={"image_url": f"/images/output.png"})
+
+@app.get("/images/{image_name}", response_class=FileResponse)
+async def get_image(image_name: str):
+    image_path = os.path.join("images", image_name)
+    if not os.path.exists(image_path):
+        raise HTTPException(status_code=404, detail="Image not found")
+    return FileResponse(image_path)
 
 @app.get("/", response_class=HTMLResponse)
 async def read_root():