Update app.py

app.py

@@ -1,9 +1,13 @@
-import gradio as gr
-import torch
+from fastapi import FastAPI, File, UploadFile, Form
+from fastapi.responses import StreamingResponse
 from PIL import Image
+import torch
+import io
 import requests
 from io import BytesIO
 
+app = FastAPI(title="AnimeGANv2 API")
+
 # بارگذاری مدل‌ها بر اساس استایل انتخابی
 def load_model(style_name):
     model = torch.hub.load(
@@ -23,39 +27,40 @@ def animegan2_transform(input_img, style_name):
     else:
         raise ValueError("فرمت تصویر ورودی صحیح نیست!")
     
-    # تغییر رزولوشن به 1024 برای کیفیت بالا
     input_img = input_img.resize((460, 460))
 
-    # بارگذاری مدل انتخابی
     model = load_model(style_name)
-
-    # بارگذاری تابع face2paint
     face2paint_func = torch.hub.load(
         "bryandlee/animegan2-pytorch:main",
         "face2paint",
-        size=1024,  # رزولوشن بالا
+        size=1024,
         verbose=False
     )
 
-    # اجرای مدل
     output_img = face2paint_func(model, input_img)
     return output_img
 
-# رابط کاربری Gradio
-iface = gr.Interface(
-    fn=animegan2_transform,
-    inputs=[
-        gr.Image(type="pil", label="آپلود تصویر یا وارد کردن URL"),
-        gr.Dropdown(
-            ["face_paint_512_v1", "face_paint_512_v2", "paprika", "celeba_distill"],
-            value="face_paint_512_v2",
-            label="انتخاب استایل"
-        )
-    ],
-    outputs=gr.Image(type="pil", label="تصویر انیمه (کیفیت بالا)"),
-    title="AnimeGANv2 - تبدیل تصویر به انیمه با انتخاب استایل",
-    description="تصاویر خود را با کیفیت بالا و استایل‌های مختلف به انیمه تبدیل کنید."
-)
+# Endpoint برای تبدیل تصویر
+@app.post("/animegan2/")
+async def transform_image(file: UploadFile = File(None), url: str = Form(None), style_name: str = Form("face_paint_512_v2")):
+    if file:
+        input_img = Image.open(file.file).convert("RGB")
+    elif url:
+        input_img = url
+    else:
+        return {"error": "لطفاً تصویر یا URL ارائه دهید!"}
+
+    # پردازش تصویر
+    output_img = animegan2_transform(input_img, style_name)
+
+    # ذخیره تصویر خروجی در بافر
+    img_byte_arr = io.BytesIO()
+    output_img.save(img_byte_arr, format='PNG')
+    img_byte_arr.seek(0)
+
+    # ارسال تصویر به صورت stream
+    return StreamingResponse(img_byte_arr, media_type="image/png")
 
 if __name__ == "__main__":
-    iface.launch()
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)