Update app.py

app.py

@@ -1028,18 +1028,36 @@ networks = {
         "conditional": True
     }
 }
+NETWORK_CHOICES = ["FFHQ (Faces) 256px", "ImageNet 64px"]
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-# Preload all networks
-for network_name, network_info in networks.items():
-    print(f'Loading network {network_name} from "{network_info["path"]}"...')
-    with dnnlib.util.open_url(network_info["path"]) as f:
-        network_info["G"] = legacy.load_network_pkl(f)['G_ema'].to(device)
+# Create separate global variables for each network
+G_ffhq = None
+G_imagenet = None
+
+# Load the FFHQ network
+print(f'Loading network "{NETWORK_CHOICES[0]}" from "{path_ffhq}"...')
+with dnnlib.util.open_url(path_ffhq) as f:
+    G_ffhq = legacy.load_network_pkl(f)['G_ema'].to(device)
+
+# Load the ImageNet network
+print(f'Loading network "{NETWORK_CHOICES[1]}" from "{path_imagenet}"...')
+with dnnlib.util.open_url(path_imagenet) as f:
+    G_imagenet = legacy.load_network_pkl(f)['G_ema'].to(device)
+
 
 @spaces.GPU
 def generate_random_images(network_choice, class_label=None, num_images=2):
     # Get the selected generator
-    G = networks[network_choice]["G"]
+    if network_choice == NETWORK_CHOICES[0]:  # FFHQ
+        G = G_ffhq
+        is_conditional = False
+    elif network_choice == NETWORK_CHOICES[1]:  # ImageNet
+        G = G_imagenet
+        is_conditional = True
+    else:
+        # This case should not be reached with the Gradio Radio component
+        raise ValueError("Invalid network choice selected.")
     
     # Generate random seeds
     seeds = np.random.randint(0, 100000, num_images)