Update app.py

app.py

@@ -9,7 +9,7 @@ from huggingface_hub import hf_hub_download
 from PIL import Image
 from torch import nn
 from torchvision.utils import save_image
-
+hfapi = HfApi()
 
 class Generator(nn.Module):
     def __init__(self, num_channels=4, latent_dim=100, hidden_size=64):
@@ -42,10 +42,6 @@ class Generator(nn.Module):
 
         return pixel_values
 
-model = Generator()
-weights_path = hf_hub_download('huggingnft/dooggies', 'pytorch_model.bin')
-model.load_state_dict(torch.load(weights_path, map_location=torch.device('cpu')))
-
 
 @torch.no_grad()
 def interpolate(save_dir='./lerp/', frames=100, rows=8, cols=8):
@@ -76,7 +72,10 @@ def interpolate(save_dir='./lerp/', frames=100, rows=8, cols=8):
                    save_all=True, duration=100, loop=1)
 
 
-def predict(choice, seed):
+def predict(model_name, choice, seed):
+    model = Generator()
+    weights_path = hf_hub_download(f'huggingnft/{model_name}', 'pytorch_model.bin')
+    model.load_state_dict(torch.load(weights_path, map_location=torch.device('cpu')))
     torch.manual_seed(seed)
 
     if choice == 'interpolation':
@@ -92,9 +91,11 @@ def predict(choice, seed):
         return 'image.png'
 
 
+models = [model.modelId[model.modelId.index("/") + 1:] for model in hfapi.list_models(author="huggingnft")]
 gr.Interface(
     predict,
     inputs=[
+        gr.inputs.Dropdown(models, label='Model'),
         gr.inputs.Dropdown(['image', 'interpolation'], label='Output Type'),
         gr.inputs.Slider(label='Seed', minimum=0, maximum=1000, default=42),
     ],
@@ -102,5 +103,5 @@ gr.Interface(
     title="Cryptopunks GAN",
     description="These CryptoPunks do not exist. You have the choice of either generating random punks, or a gif showing the interpolation between two random punk grids.",
     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=[["interpolation", 123], ["interpolation", 42], ["image", 456], ["image", 42]],
+    examples=[["interpolation", 100], ["interpolation", 500], ["image", 100], ["image", 500]],
 ).launch(cache_examples=True)