import gradio as gr
import torchvision
import torch
model = torch.hub.load(
    "ndahlquist/pytorch-hub-stylegan:0.0.1",
    "style_gan",
    pretrained=True)
def iface(name):
  model.eval()
  device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
  model.to(device)
  latents = torch.randn(1, 512, device=device)
  with torch.no_grad():
    img = model(latents)
    img = (img.clamp(-1, 1) + 1) / 2.0
    img2=img.numpy()
    img3=img2.squeeze()
    img4 = img3.swapaxes(0,1)
    img5=img4.swapaxes(1,2)
  out= img5
  return out
#demo = gr.Interface(fn=iface,inputs='text', outputs='image')
with gr.Blocks() as demo:
    output=gr.Image(label="Face")
    run_btn = gr.Button('Run')
    run_btn.click(fn=iface, outputs=output)
demo.launch()