import gradio as gr
from huggingface_hub import PyTorchModelHubMixin
import torch
import matplotlib.pyplot as plt
import torchvision
from networks_fastgan import MyGenerator
import click
import PIL
from image_generator import generate_images

def image_generation(model, number_of_images=1):
    G = MyGenerator.from_pretrained("Cropinky/projected_gan_impressionism")
    img = generate_images(model)
    #return f"generating {number_of_images} images from {model}"
    return img
if __name__ == "__main__":
    description = "TODO: when generating only 1 image use an esrgan to increase its resolution \n TODO: allow generation of multiple images"
    inputs = gr.inputs.Radio(["Abstract Expressionism", "Impressionism", "Cubism", "Pop Art", "Color Field", "Hana Hanak houses", "Hana Hanak houses - abstract expressionism", "Hana Hanak houses - color field"])
    outputs = gr.outputs.Image(label="Generated Image", type="pil")
    #outputs = "text"
    title = "Projected GAN for painting generation v0.1"
    article = "<p style='text-align: center'><a href='https://github.com/autonomousvision/projected_gan'>Official projected GAN github repo + paper</a></p>"



    gr.Interface(image_generation, inputs, outputs, title=title, article = article, description = description, 
    analytics_enabled=False).launch(debug=True)

    app, local_url, share_url = iface.launch()