Update app.py

app.py

@@ -1,48 +1,36 @@
-import gradio as gr
+from gfpgan import GFPGANer
 from PIL import Image
-import torch
-from torchvision import transforms
-from huggingface_hub import hf_hub_download
+import gradio as gr
 
 # -----------------------------
-# بارگذاری مدل CodeFormer مستقیم از Hugging Face
+# بارگذاری مدل GFPGAN (عمومی)
 # -----------------------------
-model_repo = "sczhou/CodeFormer"
-model_file = "CodeFormer.pth"  # نام مدل داخل repo
-model_path = hf_hub_download(repo_id=model_repo, filename=model_file)
-
-# فرض: کلاس CodeFormer موجود باشد
-from codeformer import CodeFormer
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model = CodeFormer(pretrained=model_path, device=device)
+# مدل به صورت خودکار از HF دانلود می‌شود و عمومی است
+restorer = GFPGANer(
+    model_path=None,    # None → مدل عمومی
+    upscale=1,          # بزرگنمایی 1x
+    arch='clean',       # معماری مدل
+    channel_multiplier=2,
+    bg_upsampler=None
+)
 
 # -----------------------------
-# تابع تبدیل تصویر
+# تابع ترمیم کارتونی
 # -----------------------------
-def modern_cartoon(image):
+def restore_cartoon(image):
     image = image.convert("RGB")
-    transform = transforms.Compose([
-        transforms.Resize((512, 512)),
-        transforms.ToTensor()
-    ])
-    img_tensor = transform(image).unsqueeze(0).to(device)
-
-    with torch.no_grad():
-        output_tensor = model.enhance(img_tensor, fidelity=0.8)
-
-    output_image = transforms.ToPILImage()(output_tensor.squeeze(0).cpu())
-    return output_image
+    restored_image, _ = restorer.enhance(image, has_aligned=False)
+    return restored_image
 
 # -----------------------------
-# رابط Gradio
+# رابط کاربری Gradio
 # -----------------------------
 iface = gr.Interface(
-    fn=modern_cartoon,
+    fn=restore_cartoon,
     inputs=gr.Image(type="pil"),
     outputs=gr.Image(type="pil"),
-    title="AnimeGAN → Modern Cartoon",
-    description="Upload your AnimeGAN output → get a modern cartoon with restored face and background."
+    title="Cartoon Face Restorer",
+    description="Upload your AnimeGAN output → get a restored modern cartoon with face and background fixed."
 )
 
 iface.launch()