import gradio as gr
from huggingface_hub import hf_hub_download
import torch
import matplotlib.pyplot as plt
import numpy as np

path = hf_hub_download('huggan/ArtGAN', 'ArtGAN.pt')
model = torch.load(path)
device = 'cuda' if torch.cuda.is_available() else 'cpu'

def generate(seed):
    with torch.no_grad():
        noise = torch.randn(seed, 100, 1, 1, device=device)
    with torch.no_grad():
        art = model(noise).detach().cpu()
    gen = np.transpose(art[-1], (1, 2, 0))
    fig = plt.figure(figsize=(5, 5))
    plt.imshow(gen)
    plt.axis('off')
    return fig

gr.Interface(
    fn=generate,
    inputs=[
        gr.inputs.Slider
        (
            label='noise',
            minimum=10,
            maximum=100,
            step=1,
            default=25
        )
    ],  
    outputs=gr.outputs.Image(type='plot'),
    title='ArtGAN',
    description='Generate A Abract Art Using ArtGAN',
).launch()