Update app.py

app.py

@@ -75,9 +75,14 @@ def interpolate(model, save_dir='./lerp/', frames=100, rows=8, cols=8):
 
 
 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')))
+    try:
+        model = Generator(3)
+        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')))
+    except:
+        model = Generator(4)
+        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':