Update app.py

app.py

@@ -11,6 +11,7 @@ import numpy as np
 import re
 import requests
 import time
+import gc
 
 # Import necessary functions and classes
 from utils import load_t5, load_clap
@@ -69,9 +70,10 @@ def unload_current_model():
     global global_model, current_model_name
     if global_model is not None:
         del global_model
-        torch.cuda.empty_cache()
         global_model = None
     current_model_name = None
+    torch.cuda.empty_cache()
+    gc.collect()
 
 def load_model(model_name, device, model_url=None):
     global global_model, current_model_name
@@ -121,8 +123,7 @@ def load_model(model_name, device, model_url=None):
         load_time = end_time - start_time
         return f"Successfully loaded model: {model_name} in {load_time:.2f} seconds"
     except Exception as e:
-        global_model = None
-        current_model_name = None
+        unload_current_model()
         print(f"Error loading model {model_name}: {str(e)}")
         return f"Failed to load model: {model_name}. Error: {str(e)}"
 
@@ -229,6 +230,11 @@ def generate_music(prompt, seed, cfg_scale, steps, duration, device, batch_size=
 
         all_waveforms.append(waveform)
 
+        # Clear some memory after each segment
+        del images, latents, mel_spectrogram, x_i
+        torch.cuda.empty_cache()
+        gc.collect()
+
     # Concatenate all waveforms
     final_waveform = np.concatenate(all_waveforms)