Update app.py

app.py

@@ -1,28 +1,19 @@
-from fastapi import FastAPI, File, UploadFile, Form
-from fastapi.responses import StreamingResponse
-from PIL import Image
+import gradio as gr
 import torch
-from io import BytesIO
+from PIL import Image
 import requests
-import logging
-
-# تنظیم لاگ‌گیری
-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__)
+from io import BytesIO
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import StreamingResponse
+import uvicorn
+import tempfile
 
+# Initialize FastAPI app
 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",
@@ -31,65 +22,84 @@ def load_model(style_name):
         ).eval()
         return model
     except Exception as e:
-        logger.error(f"Error loading model: {str(e)}")
-        raise
+        raise HTTPException(status_code=500, detail=f"Error loading model: {str(e)}")
+
+# پردازش تصویر با AnimeGAN
+def animegan2_transform(image_url, style_name):
+    # بررسی استایل‌های معتبر
+    valid_styles = ["face_paint_512_v1", "face_paint_512_v2", "paprika", "celeba_distill"]
+    if style_name not in valid_styles:
+        raise HTTPException(status_code=400, detail="Invalid style name!")
+
+    # دانلود و پردازش تصویر
+    try:
+        response = requests.get(image_url)
+        response.raise_for_status()
+        input_img = Image.open(BytesIO(response.content)).convert("RGB")
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Error loading image: {str(e)}")
+
+    # تغییر رزولوشن
+    input_img = input_img.resize((460, 460))
 
-# پردازش تصویر
-def animegan2_transform(input_img, style_name):
+    # بارگذاری مدل
+    model = load_model(style_name)
+
+    # بارگذاری تابع face2paint
     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((256, 256))  # رزولوشن کمتر برای تست سریع‌تر
-        model = load_model(style_name)
         face2paint_func = torch.hub.load(
             "bryandlee/animegan2-pytorch:main",
             "face2paint",
-            size=512,  # رزولوشن کمتر برای تست
+            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
+        raise HTTPException(status_code=500, detail=f"Error loading face2paint: {str(e)}")
+
+    # اجرای مدل
+    output_img = face2paint_func(model, input_img)
+
+    # ذخیره تصویر خروجی به‌صورت موقت
+    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
+    output_img.save(temp_file.name, format="PNG")
+    
+    return temp_file.name
 
-# Endpoint برای API
-@app.post("/transform")
-async def transform_image(file: UploadFile = File(None), style: str = Form(default="face_paint_512_v2"), url: str = Form(None)):
+# FastAPI endpoint برای تبدیل تصویر
+@app.post("/api/animegan2")
+async def animegan2_api(image_url: str, style: str = "face_paint_512_v2"):
     try:
-        logger.info(f"Received request with style: {style}")
-        if file:
-            image = Image.open(file.file).convert("RGB")
-        elif url:
-            image = Image.open(BytesIO(requests.get(url).content)).convert("RGB")
-        else:
-            logger.error("No file or URL provided")
-            return {"error": "لطفاً تصویر یا URL ارائه دهید"}
+        # پردازش تصویر
+        output_path = animegan2_transform(image_url, style)
         
-        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 StreamingResponse(output_buffer, media_type="image/png")
+        # ارسال تصویر به‌صورت استریم
+        with open(output_path, "rb") as f:
+            return StreamingResponse(f, media_type="image/png")
+    except HTTPException as e:
+        raise e
     except Exception as e:
-        logger.error(f"API error: {str(e)}")
-        return {"error": str(e)}
+        raise HTTPException(status_code=500, detail=f"Internal error: {str(e)}")
 
-# سلامت‌سنجی سرور
-@app.get("/health")
-async def health_check():
-    return {"status": "healthy"}
+# رابط کاربری Gradio (اختیاری، اگر هنوز می‌خوای رابط گرافیکی داشته باشی)
+def gradio_interface():
+    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="تصاویر خود را با کیفیت بالا و استایل‌های مختلف به انیمه تبدیل کنید."
+    )
+    return iface
 
+# راه‌اندازی Gradio در کنار FastAPI
 if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    # لانچ کردن Gradio (اختیاری)
+    iface = gradio_interface()
+    iface.launch(server_name="0.0.0.0", server_port=7860, share=False)