Added example usage

README.md

@@ -1 +1,28 @@
-lightweight GAN trained on glid-3 orbs for demo I'm working on.
+lightweight GAN trained on glid-3 orbs for demo I'm working on.
+
+Demo:
+
+from lightweight_gan import Generator
+import torch
+from matplotlib import pyplot as plt
+from huggingface_hub import PyTorchModelHubMixin
+class GeneratorWithPyTorchModelHubMixin(gan.__class__, PyTorchModelHubMixin):
+    pass
+
+# Initialize a generator model
+gan_new = Generator(latent_dim=256, image_size=256, attn_res_layers = [32])
+
+# Load from local saved state dict
+# gan_new.load_state_dict(torch.load('/content/orbgan_e3_state_dict.pt')) 
+
+# Load from model hub:
+gan_new.__class__ = GeneratorWithPyTorchModelHubMixin
+gan_new = gan_new.from_pretrained('johnowhitaker/orbgan_e1', latent_dim=256, image_size=256, attn_res_layers = [32])
+
+# View some examples
+n_rows = 3
+ims = gan_new(torch.randn(n_rows**2, 256)).clamp_(0., 1.)
+fig, axs = plt.subplots(n_rows, n_rows, figsize=(9, 9))
+for i, ax in enumerate(axs.flatten()):
+  ax.imshow(ims[i].permute(1, 2, 0).detach().cpu().numpy())
+plt.tight_layout()