Update app.py

app.py

@@ -125,18 +125,23 @@ def load_model(model_name, device, model_url=None):
 def load_resources(device):
     global global_t5, global_clap, global_vae, global_vocoder, global_diffusion
     
-    print("Loading T5 and CLAP models...")
-    global_t5 = load_t5(device, max_length=256)
-    global_clap = load_clap(device, max_length=256)
-    
-    print("Loading VAE and vocoder...")
-    global_vae = AutoencoderKL.from_pretrained('cvssp/audioldm2', subfolder="vae").to(device)
-    global_vocoder = SpeechT5HifiGan.from_pretrained('cvssp/audioldm2', subfolder="vocoder").to(device)
-    
-    print("Initializing diffusion...")
-    global_diffusion = RF()
-    
-    print("Base resources loaded successfully!")
+    try:
+        print("Loading T5 and CLAP models...")
+        global_t5 = load_t5(device, max_length=256)
+        global_clap = load_clap(device, max_length=256)
+        
+        print("Loading VAE and vocoder...")
+        global_vae = AutoencoderKL.from_pretrained('cvssp/audioldm2', subfolder="vae").to(device)
+        global_vocoder = SpeechT5HifiGan.from_pretrained('cvssp/audioldm2', subfolder="vocoder").to(device)
+        
+        print("Initializing diffusion...")
+        global_diffusion = RF()
+        
+        print("Base resources loaded successfully!")
+        return "Resources loaded successfully!"
+    except Exception as e:
+        print(f"Error loading resources: {str(e)}")
+        return f"Failed to load resources. Error: {str(e)}"
 
 def generate_music(prompt, seed, cfg_scale, steps, duration, device, batch_size=4, progress=gr.Progress()):
     global global_model, global_t5, global_clap, global_vae, global_vocoder, global_diffusion
@@ -144,6 +149,9 @@ def generate_music(prompt, seed, cfg_scale, steps, duration, device, batch_size=
     if global_model is None:
         return "Please select and load a model first.", None
     
+    if global_t5 is None or global_clap is None or global_vae is None or global_vocoder is None or global_diffusion is None:
+        return "Resources not properly loaded. Please reload the page and try again.", None
+    
     if seed == 0:
         seed = random.randint(1, 1000000)
     print(f"Using seed: {seed}")
@@ -293,6 +301,9 @@ with gr.Blocks(theme=theme) as iface:
     output_audio = gr.Audio(type="filepath")
 
     def on_load_model_click(model_name, device, url):
+        resource_status = load_resources(device)
+        if "Failed" in resource_status:
+            return resource_status
         if url:
             result = load_model(None, device, model_url=url)
         else:
@@ -303,7 +314,7 @@ with gr.Blocks(theme=theme) as iface:
     generate_button.click(generate_music, inputs=[prompt, seed, cfg_scale, steps, duration, device_choice], outputs=[output_status, output_audio])
 
     # Load default model and resources on startup
-    iface.load(lambda: (load_resources("cpu"), on_load_model_click(default_model, "cpu", None)), inputs=None, outputs=None)
+    iface.load(lambda: on_load_model_click(default_model, "cpu", None), inputs=None, outputs=None)
 
 # Launch the interface
 iface.launch()