import os
import sys
import types
import cv2
import gradio as gr
import numpy as np
from gfpgan import GFPGANer
from huggingface_hub import hf_hub_download

# -------- Patch برای جلوگیری از ارور torchvision.transforms.functional_tensor --------
try:
    import torchvision.transforms.functional_tensor
except ModuleNotFoundError:
    # ساخت یک ماژول خالی به جای functional_tensor
    sys.modules['torchvision.transforms.functional_tensor'] = types.ModuleType('functional_tensor')

# -------- کش مدل‌ها --------
loaded_models = {}

def load_model(version):
    if version not in loaded_models:
        if version == "v1.4 (original)":
            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
            )
        elif version == "v1.3 (clean)":
            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)

    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
    )

    restored_img = cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB)
    return restored_img

# -------- رابط کاربری Gradio --------
iface = gr.Interface(
    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."
)

if __name__ == "__main__":
    iface.launch()