Update app.py

app.py

@@ -1,58 +1,86 @@
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, File, UploadFile, Form
+from fastapi.responses import StreamingResponse
 from PIL import Image
 import torch
 from io import BytesIO
 import requests
-from pydantic import BaseModel
+import logging
 
-app = FastAPI()
+# تنظیم لاگ‌گیری
+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__)
 
-# مدل درخواست ورودی
-class ImageRequest(BaseModel):
-    image_url: str
-    style: str = "face_paint_512_v2"
+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
 
-# پردازش تصویر
-def animegan2_transform(image_url, style_name):
+def animegan2_transform(input_img, style_name):
     try:
-        input_img = Image.open(BytesIO(requests.get(image_url).content)).convert("RGB")
-        input_img = input_img.resize((460, 460))
+        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((128, 128))  # رزولوشن پایین برای سرعت
         model = load_model(style_name)
         face2paint_func = torch.hub.load(
             "bryandlee/animegan2-pytorch:main",
             "face2paint",
-            size=1024,
+            size=256,
             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
+
+@app.post("/transform")
+async def transform_image(file: UploadFile = File(None), style: str = Form(default="face_paint_512_v2"), url: str = Form(None)):
+    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_img = animegan2_transform(image, style)
         output_buffer = BytesIO()
         output_img.save(output_buffer, format="PNG")
         output_buffer.seek(0)
-        return output_buffer
+        logger.info("Returning processed image")
+        return StreamingResponse(output_buffer, media_type="image/png")
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        logger.error(f"API error: {str(e)}")
+        return {"error": str(e)}
 
-# مسیر API برای پردازش تصویر
-@app.post("/process_image")
-async def process_image(request: ImageRequest):
-    valid_styles = ["face_paint_512_v1", "face_paint_512_v2", "paprika", "celeba_distill"]
-    if request.style not in valid_styles:
-        raise HTTPException(status_code=400, detail="استایل نامعتبر!")
-    
-    output_buffer = animegan2_transform(request.image_url, request.style)
-    return {"image": output_buffer.getvalue(), "content_type": "image/png"}
+@app.get("/health")
+async def health_check():
+    return {"status": "healthy"}
 
-# برای تست ساده
-@app.get("/")
-async def root():
-    return {"message": "API برای تبدیل تصویر به انیمه"}
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)