from huggingface_hub import hf_hub_url, hf_hub_download

import gradio as gr
import numpy as np
import requests

import torch
from torchvision import transforms
from torch.autograd import Variable

from PIL import Image

import warnings
warnings.filterwarnings('ignore')

# モデルのダウンロード
path_to_model = hf_hub_download(
    repo_id="opetrova/face-frontalization",
    filename="generator_v0.pt"
)

# network.py をカレントディレクトリにダウンロード
network_url = hf_hub_url(repo_id="opetrova/face-frontalization", filename="network.py")
r = requests.get(network_url, allow_redirects=True)
open('network.py', 'wb').write(r.content)

# PyTorch 2.6 以降は weights_only=False を指定しないとエラーになる
saved_model = torch.load(path_to_model, map_location=torch.device("cpu"), weights_only=False)

def frontalize(image):
    # 画像を [1, 3, 128, 128] tensor に変換
    preprocess = transforms.Compose((
        transforms.ToPILImage(),
        transforms.Resize(size=(128, 128)),
        transforms.ToTensor(),
    ))
    input_tensor = torch.unsqueeze(preprocess(image), 0)

    # 推論
    generated_image = saved_model(Variable(input_tensor.type(torch.FloatTensor)))
    generated_image = generated_image.detach().squeeze().permute(1, 2, 0).numpy()
    generated_image = (generated_image + 1.0) / 2.0  # [-1,1] → [0,1]

    return generated_image

# Gradio インターフェース
iface = gr.Interface(
    fn=frontalize,
    inputs=gr.Image(type="numpy"),
    outputs="image",
    title="Face Frontalization",
    description=(
        'PyTorch implementation of a supervised GAN '
        '(see <a href="https://blog.scaleway.com/gpu-instances-using-deep-learning-to-obtain-frontal-rendering-of-facial-images/">blog post</a>)'
    ),
    examples=["amos.png", "clarissa.png"],
)

iface.launch()