Update app.py

app.py

@@ -1,88 +1,61 @@
-from flask import Flask, request, send_file, jsonify
+from fastapi import FastAPI, File, UploadFile, Form
 from PIL import Image
 import torch
 from io import BytesIO
-import requests
-import logging
+import gradio as gr
 
-# تنظیم لاگ‌گیری
-logging.basicConfig(
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
-    level=logging.INFO,
-    handlers=[logging.StreamHandler(), logging.FileHandler("app.log")]
-)
-logger = logging.getLogger(__name__)
-
-app = Flask(__name__)
+app = FastAPI()
 
+# کدهای قبلی برای بارگذاری مدل و پردازش تصویر
 def load_model(style_name):
-    try:
-        logger.info(f"Loading model for style: {style_name}")
-        model = torch.hub.load(
-            "bryandlee/animegan2-pytorch:main",
-            "generator",
-            pretrained=style_name,
-            verbose=False
-        ).eval()
-        return model
-    except Exception as e:
-        logger.error(f"Error loading model: {str(e)}")
-        raise
+    model = torch.hub.load(
+        "bryandlee/animegan2-pytorch:main",
+        "generator",
+        pretrained=style_name,
+        verbose=False
+    ).eval()
+    return model
 
 def animegan2_transform(input_img, style_name):
-    try:
-        logger.info("Processing image...")
-        if isinstance(input_img, str):
-            input_img = Image.open(BytesIO(requests.get(input_img).content)).convert("RGB")
-        elif isinstance(input_img, Image.Image):
-            input_img = input_img.convert("RGB")
-        else:
-            raise ValueError("فرمت تصویر ورودی صحیح نیست!")
-        
-        input_img = input_img.resize((64, 64))  # رزولوشن بسیار پایین برای سرعت
-        model = load_model(style_name)
-        face2paint_func = torch.hub.load(
-            "bryandlee/animegan2-pytorch:main",
-            "face2paint",
-            size=128,
-            verbose=False
-        )
-        output_img = face2paint_func(model, input_img)
-        logger.info("Image processed successfully")
-        return output_img
-    except Exception as e:
-        logger.error(f"Error processing image: {str(e)}")
-        raise
+    if isinstance(input_img, str):
+        input_img = Image.open(BytesIO(requests.get(input_img).content)).convert("RGB")
+    elif isinstance(input_img, Image.Image):
+        input_img = input_img.convert("RGB")
+    input_img = input_img.resize((460, 460))
+    model = load_model(style_name)
+    face2paint_func = torch.hub.load(
+        "bryandlee/animegan2-pytorch:main",
+        "face2paint",
+        size=1024,
+        verbose=False
+    )
+    output_img = face2paint_func(model, input_img)
+    return output_img
 
-@app.route("/transform", methods=["POST"])
-def transform_image():
-    try:
-        style = request.form.get("style", "face_paint_512_v2")
-        logger.info(f"Received request with style: {style}")
-        
-        if "file" in request.files:
-            file = request.files["file"]
-            image = Image.open(file.stream).convert("RGB")
-        elif "url" in request.form:
-            url = request.form["url"]
-            image = Image.open(BytesIO(requests.get(url).content)).convert("RGB")
-        else:
-            logger.error("No file or URL provided")
-            return jsonify({"error": "لطفاً تصویر یا URL ارائه دهید"}), 400
-        
-        output_img = animegan2_transform(image, style)
-        output_buffer = BytesIO()
-        output_img.save(output_buffer, format="PNG")
-        output_buffer.seek(0)
-        logger.info("Returning processed image")
-        return send_file(output_buffer, mimetype="image/png")
-    except Exception as e:
-        logger.error(f"API error: {str(e)}")
-        return jsonify({"error": str(e)}), 500
+@app.post("/animegan")
+async def process_image(file: UploadFile = File(...), style: str = Form(...)):
+    image = Image.open(BytesIO(await file.read())).convert("RGB")
+    output_img = animegan2_transform(image, style)
+    output_buffer = BytesIO()
+    output_img.save(output_buffer, format="PNG")
+    output_buffer.seek(0)
+    return {"image": output_buffer.getvalue()}
 
-@app.route("/health", methods=["GET"])
-def health_check():
-    return jsonify({"status": "healthy"})
+# برای رابط 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 - تبدیل تصویر به انیمه با انتخاب استایل"
+)
 
 if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=8000)
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)