import torch
import gradio as gr
import matplotlib.pyplot as plt
import torchvision
use_gpu = True if torch.cuda.is_available() else False

model = torch.hub.load('facebookresearch/pytorch_GAN_zoo:hub', 'DCGAN', pretrained=True, useGPU=use_gpu)

def dcgan(num_images):
  noise, _ = model.buildNoiseData(int(num_images))
  with torch.no_grad():
      generated_images = model.test(noise)
  plt.imshow(torchvision.utils.make_grid(generated_images).permute(1, 2, 0).cpu().numpy())
  plt.axis("off")
  return plt


inputs = gr.inputs.Number(label="number of images")
outputs = gr.outputs.Image(label="Output Image")

title = "DCGAN"
description = "demo for DCGAN. To use it, simply add the number of images to generate or click on the examples. Read more below."
article = "<p style='text-align: center'><a href='https://arxiv.org/abs/1511.06434'>Unsupervised Representation Learning with Deep Convolutional Generative Adversarial Networks</a> | <a href='https://github.com/facebookresearch/pytorch_GAN_zoo/blob/master/models/DCGAN.py'>Github Repo</a></p>"
examples = [
            [1], 
            [2],
            [3],
            [4], 
            [64]
]


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