Views
No views yet
1from lightweight_gan import Generator
2import torch
3from matplotlib import pyplot as plt
4from huggingface_hub import PyTorchModelHubMixin
5
6# Initialize a generator model
7gan_new = Generator(latent_dim=256, image_size=256, attn_res_layers = [32])
8
9# Load from local saved state dict
10# gan_new.load_state_dict(torch.load('/content/orbgan_e3_state_dict.pt'))
11
12# Load from model hub:
13class GeneratorWithPyTorchModelHubMixin(gan_new.__class__, PyTorchModelHubMixin):
14 pass
15gan_new.__class__ = GeneratorWithPyTorchModelHubMixin
16gan_new = gan_new.from_pretrained('johnowhitaker/orbgan_e1', latent_dim=256, image_size=256, attn_res_layers = [32])
17
18# View some examples
19n_rows = 3
20ims = gan_new(torch.randn(n_rows**2, 256)).clamp_(0., 1.)
21fig, axs = plt.subplots(n_rows, n_rows, figsize=(9, 9))
22for i, ax in enumerate(axs.flatten()):
23 ax.imshow(ims[i].permute(1, 2, 0).detach().cpu().numpy())
24plt.tight_layout()