Update app.py

app.py

@@ -1,83 +1,61 @@
 import gradio as gr
 import torch
-import numpy as np
 from PIL import Image
 import requests
-import io
-
-# ---- دانلود و آماده‌سازی AnimeGAN ----
-def load_animegan(style):
-    return torch.hub.load(
-        'bryandlee/animegan2-pytorch:main',
-        style,
-        pretrained=True,
-        device=device
-    )
-
-# ---- نسخه ساده rgb_to_grayscale (بدون torchvision) ----
-def rgb_to_grayscale(img):
-    if img.ndim != 4 or img.shape[1] != 3:
-        raise ValueError("Expected (B, 3, H, W)")
-    r, g, b = img[:, 0:1], img[:, 1:2], img[:, 2:3]
-    return (0.2989 * r + 0.5870 * g + 0.1140 * b)
-
-# ---- بارگذاری GFPGAN (برای بهبود چهره) ----
-from gfpgan import GFPGANer
-restorer = GFPGANer(
-    model_path='https://github.com/TencentARC/GFPGAN/releases/download/v1.4/GFPGANv1.4.pth',
-    upscale=1,
-    arch='clean',
-    channel_multiplier=2
-)
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# ---- پردازش تصویر ----
-def process_image(image, style_name, face_enhance):
-    if image is None:
-        return None
-
-    # تغییر رزولوشن بالا
-    image = image.convert("RGB")
-    image = image.resize((1024, 1024), Image.BICUBIC)
-
-    # انتخاب مدل AnimeGAN
-    model = load_animegan(style_name)
-    face2paint = torch.hub.load(
-        'bryandlee/animegan2-pytorch:main',
-        'face2paint',
-        pretrained=True,
-        device=device
+from io import BytesIO
+
+# بارگذاری مدل‌ها بر اساس استایل انتخابی
+def load_model(style_name):
+    model = torch.hub.load(
+        "bryandlee/animegan2-pytorch:main",
+        "generator",
+        pretrained=style_name,
+        verbose=False
+    ).eval()
+    return model
+
+# پردازش تصویر با 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("فرمت تصویر ورودی صحیح نیست!")
+    
+    # تغییر رزولوشن به 1024 برای کیفیت بالا
+    input_img = input_img.resize((1024, 1024))
+
+    # بارگذاری مدل انتخابی
+    model = load_model(style_name)
+
+    # بارگذاری تابع face2paint
+    face2paint_func = torch.hub.load(
+        "bryandlee/animegan2-pytorch:main",
+        "face2paint",
+        size=1024,  # رزولوشن بالا
+        verbose=False
     )
 
-    # تبدیل به انیمه
-    anime_img = face2paint(model, image)
-
-    # بهبود چهره با GFPGAN (اختیاری)
-    if face_enhance:
-        _, _, restored_img = restorer.enhance(
-            np.array(anime_img),
-            has_aligned=False,
-            only_center_face=False,
-            paste_back=True
-        )
-        anime_img = Image.fromarray(restored_img)
+    # اجرای مدل
+    output_img = face2paint_func(model, input_img)
+    return output_img
 
-    return anime_img
-
-# ---- رابط کاربری ----
-styles = ["face_paint_512_v1", "face_paint_512_v2", "paprika", "celeba_distill"]
-
-demo = gr.Interface(
-    fn=process_image,
+# رابط کاربری Gradio
+iface = gr.Interface(
+    fn=animegan2_transform,
     inputs=[
-        gr.Image(type="pil", label="عکس ورودی"),
-        gr.Dropdown(choices=styles, value="face_paint_512_v2", label="انتخاب استایل"),
-        gr.Checkbox(label="بهبود چهره با GFPGAN", value=False)
+        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(label="خروجی", type="pil"),
-    title="AnimeGAN + GFPGAN",
-    description="تبدیل عکس به انیمه با انتخاب استایل و بهبود چهره"
+    outputs=gr.Image(type="pil", label="تصویر انیمه (کیفیت بالا)"),
+    title="AnimeGANv2 - تبدیل تصویر به انیمه با انتخاب استایل",
+    description="تصاویر خود را با کیفیت بالا و استایل‌های مختلف به انیمه تبدیل کنید."
 )
 
-demo.launch()
+if __name__ == "__main__":
+    iface.launch()