Update app.py

app.py

@@ -1,39 +1,65 @@
-# ----- Patch برای torchvision -----
-import sys
-import torchvision.transforms.functional as F
-sys.modules['torchvision.transforms.functional_tensor'] = F
-# --------------------------------
-
+import gradio as gr
 from gfpgan import GFPGANer
 from huggingface_hub import hf_hub_download
-from PIL import Image
-import gradio as gr
+import cv2
+import numpy as np
 
-# دانلود مدل v1.4 (سازگار با arch='original')
-model_path = hf_hub_download(
-    repo_id="leonelhs/gfpgan",
-    filename="GFPGANv1.4.pth"
-)
+# کش مدل‌ها
+loaded_models = {}
 
-restorer = GFPGANer(
-    model_path=model_path,
-    upscale=1,
-    arch='original',   # 🔑 تغییر دادیم
-    channel_multiplier=2,
-    bg_upsampler=None
-)
+def load_model(version):
+    if version == "v1.4 (original)":
+        if version not in loaded_models:
+            model_path = hf_hub_download("leonelhs/gfpgan", "GFPGANv1.4.pth")
+            loaded_models[version] = GFPGANer(
+                model_path=model_path,
+                upscale=1,
+                arch='original',
+                channel_multiplier=2,
+                bg_upsampler=None
+            )
+        return loaded_models[version]
+
+    elif version == "v1.3 (clean)":
+        if version not in loaded_models:
+            model_path = hf_hub_download("leonelhs/gfpgan", "GFPGANv1.3.pth")
+            loaded_models[version] = GFPGANer(
+                model_path=model_path,
+                upscale=1,
+                arch='clean',
+                channel_multiplier=2,
+                bg_upsampler=None
+            )
+        return loaded_models[version]
+
+
+def enhance_face(image, version):
+    restorer = load_model(version)
+
+    # تبدیل تصویر PIL/numpy به BGR
+    if isinstance(image, np.ndarray):
+        img = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+    else:
+        raise ValueError("Invalid image format")
+
+    _, _, restored_img = restorer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
+
+    # تبدیل دوباره به RGB برای نمایش
+    restored_img = cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB)
+
+    return restored_img
 
-def restore_cartoon(image):
-    image = image.convert("RGB")
-    restored_image, _ = restorer.enhance(image, has_aligned=False)
-    return restored_image
 
 iface = gr.Interface(
-    fn=restore_cartoon,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Image(type="pil"),
-    title="Cartoon Face Restorer",
-    description="Upload your AnimeGAN cartoon → get a modern restored cartoon."
+    fn=enhance_face,
+    inputs=[
+        gr.Image(type="numpy", label="Upload your Anime/Cartoon image"),
+        gr.Radio(["v1.3 (clean)", "v1.4 (original)"], value="v1.4 (original)", label="Choose GFPGAN Version")
+    ],
+    outputs=gr.Image(type="numpy", label="Restored Cartoon Face"),
+    title="Cartoon Face Restoration with GFPGAN",
+    description="Upload a cartoonized image (e.g. from AnimeGAN) and restore it with GFPGAN."
 )
 
-iface.launch()
+if __name__ == "__main__":
+    iface.launch()