Update app.py

app.py

@@ -1,4 +1,8 @@
+import torch
 from torch import nn
+from huggingface_hub import hf_hub_download
+from torchvision.utils import save_image
+import gradio as gr
 
 class Generator(nn.Module):
     # Refer to the link below for explanations about nc, nz, and ngf
@@ -23,35 +27,23 @@ 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
 
-
-
-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.Slider(0, 1000, label='Seed', default=42),
+        gr.Slider(4, 64, label='Number of Punks', step=1, default=10),
     ],
     outputs="image",
-).launch()
+    examples=[[123, 15], [42, 29], [456, 8], [1337, 35]],
+).launch(cache_examples=True)