StyleGAN / app.py
sudhir2016's picture
Update app.py
bdeec8d
raw
history blame
612 Bytes
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')
demo.launch()