Add num_punks selector

app.py

@@ -1,8 +1,11 @@
+import gradio as gr
+import torch
+from huggingface_hub import hf_hub_download
 from torch import nn
+from torchvision.utils import save_image
+
 
 class Generator(nn.Module):
-    # Refer to the link below for explanations about nc, nz, and ngf
-    # https://pytorch.org/tutorials/beginner/dcgan_faces_tutorial.html#inputs
     def __init__(self, nc=4, nz=100, ngf=64):
         super(Generator, self).__init__()
         self.network = nn.Sequential(
@@ -23,29 +26,29 @@ class Generator(nn.Module):
         output = self.network(input)
         return output
 
-from huggingface_hub import hf_hub_download
-import torch
 
 model = Generator()
 weights_path = hf_hub_download('nateraw/cryptopunks-gan', 'generator.pth')
-model.load_state_dict(torch.load(weights_path, map_location=torch.device('cpu'))) # Use 'cuda' if you have a GPU available
+model.load_state_dict(torch.load(weights_path, map_location=torch.device('cpu')))
 
-from torchvision.utils import save_image
 
-def predict(seed):
-    num_punks = 4
+def predict(seed, num_punks):
     torch.manual_seed(seed)
     z = torch.randn(num_punks, 100, 1, 1)
     punks = model(z)
     save_image(punks, "punks.png", normalize=True)
     return 'punks.png'
 
-import gradio as gr
 
 gr.Interface(
     predict,
     inputs=[
         gr.inputs.Slider(label='Seed', minimum=0, maximum=1000, default=42),
+        gr.inputs.Slider(label='Number of Punks', minimum=4, maximum=64, step=1, default=10),
     ],
     outputs="image",
-).launch()
+    title="Cryptopunks GAN",
+    description="These CryptoPunks do not exist. Generate random punks with an initial seed!",
+    article="<p style='text-align: center'><a href='https://arxiv.org/pdf/1511.06434.pdf'>Unsupervised Representation Learning with Deep Convolutional Generative Adversarial Networks</a> | <a href='https://github.com/teddykoker/cryptopunks-gan'>Github Repo</a></p>",
+    examples=[[123], [42], [456], [1337]],
+).launch(cache_examples=True)