Update app.py

app.py

@@ -2,64 +2,93 @@ 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
+import gradio as gr
+import logging
 
-app = FastAPI(title="AnimeGANv2 API")
+# تنظیم لاگ‌گیری
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-# بارگذاری مدل‌ها بر اساس استایل انتخابی
+app = FastAPI()
+
+# بارگذاری مدل‌ها
 def load_model(style_name):
-    model = torch.hub.load(
-        "bryandlee/animegan2-pytorch:main",
-        "generator",
-        pretrained=style_name,
-        verbose=False
-    ).eval()
-    return model
+    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
 
-# پردازش تصویر با AnimeGAN
+# پردازش تصویر
 def animegan2_transform(input_img, style_name):
-    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((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
-
-# 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)
+    try:
+        logger.info("Processing image...")
+        if isinstance(input_img, str):
+            logger.info("Input is URL, fetching image...")
+            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((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)
+        logger.info("Image processed successfully")
+        return output_img
+    except Exception as e:
+        logger.error(f"Error processing image: {str(e)}")
+        raise
 
-    # ذخیره تصویر خروجی در بافر
-    img_byte_arr = io.BytesIO()
-    output_img.save(img_byte_arr, format='PNG')
-    img_byte_arr.seek(0)
+# Endpoint برای API
+@app.post("/transform")
+async def transform_image(file: UploadFile = File(...), style: str = Form(default="face_paint_512_v2")):
+    try:
+        logger.info(f"Received image with style: {style}")
+        image = Image.open(file.file).convert("RGB")
+        output_img = animegan2_transform(image, style)
+        output_buffer = BytesIO()
+        output_img.save(output_buffer, format="PNG")
+        output_buffer.seek(0)
+        return StreamingResponse(output_buffer, media_type="image/png")
+    except Exception as e:
+        logger.error(f"API error: {str(e)}")
+        return {"error": str(e)}
 
-    # ارسال تصویر به صورت stream
-    return StreamingResponse(img_byte_arr, media_type="image/png")
+# رابط Gradio (برای تست در Hugging Face)
+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="تصاویر خود را با کیفیت بالا و استایل‌های مختلف به انیمه تبدیل کنید."
+)
 
 if __name__ == "__main__":
     import uvicorn