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)