File size: 1,224 Bytes
b8d5ba9 b495775 084ed1c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 |
import torch
import nltk
nltk.download('wordnet')
from pytorch_pretrained_biggan import (BigGAN, one_hot_from_names, truncated_noise_sample,
save_as_images, display_in_terminal)
initial_archi = 'biggan-deep-128' #@param ['biggan-deep-128', 'biggan-deep-256', 'biggan-deep-512'] {allow-input: true}
gan_model = BigGAN.from_pretrained(initial_archi).cuda().eval()
# Prepare a input
truncation = 0.4
class_vector = one_hot_from_names(initial_class, batch_size=1)
noise_vector = truncated_noise_sample(truncation=truncation, batch_size=1)
# All in tensors
noise_vector = torch.from_numpy(noise_vector)
class_vector = torch.from_numpy(class_vector)
# If you have a GPU, put everything on cuda
noise_vector = noise_vector.to('cuda')
class_vector = class_vector.to('cuda')
gan_model.to('cuda')
# Generate an image
with torch.no_grad():
output = gan_model(noise_vector, class_vector, truncation)
# If you have a GPU put back on CPU
output = output.to('cpu')
# If you have a sixtel compatible terminal you can display the images in the terminal
# (see https://github.com/saitoha/libsixel for details)
#display_in_terminal(output)
# Save results as png images
save_as_images(output) |