import gradio as gr
import torch
from torchvision import transforms
from PIL import Image
from animegan2_pytorch import AnimeGAN2

# Load AnimeGAN model
model = AnimeGAN2('face_paint_512_v2')
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)

# Function to process images
def convert_to_anime(image):
    image = Image.fromarray(image)
    image = transforms.ToTensor()(image).unsqueeze(0).to(device)
    with torch.no_grad():
        output = model(image)[0].cpu()
    output = transforms.ToPILImage()(output)
    return output

# Gradio UI
iface = gr.Interface(
    fn=convert_to_anime,
    inputs=gr.Image(type="numpy"),
    outputs=gr.Image(type="pil"),
    title="AnimeGAN Image Converter",
    description="Upload an image to convert it into anime style!",
)

iface.launch()