Update app.py

app.py

@@ -12,9 +12,14 @@ import re
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 # Cargar los modelos una sola vez
+def load_model(model_name):
+    """Cargar el modelo especificado."""
+    return MusicGen.get_pretrained(model_name)
+
+# Inicializar los modelos
 models = {
-    'Large': MusicGen.get_pretrained('nateraw/musicgen-songstarter-v0.2').to(device),
-    'Small': MusicGen.get_pretrained('facebook/musicgen-small').to(device)
+    'Large': load_model('nateraw/musicgen-songstarter-v0.2'),
+    'Small': load_model('facebook/musicgen-small')
 }
 
 def get_model(model_choice):
@@ -33,12 +38,12 @@ def generate_music(description, melody_audio, duration, model_choice):
                 description = [description]
                 if melody_audio:
                     melody, sr = torchaudio.load(melody_audio, normalize=True)
-                    melody = melody.to(device)
+                    melody = melody.to(device) if torch.cuda.is_available() else melody
                     wav = model.generate_with_chroma(description, melody[None], sr)
                 else:
-                    wav = model.generate(description, use_amp=torch.cuda.is_available())
+                    wav = model.generate(description)
             else:
-                wav = model.generate_unconditional(1, use_amp=torch.cuda.is_available())
+                wav = model.generate_unconditional(1)
 
             filename = f'{str(uuid.uuid4())}.wav'
             path = audio_write(filename, wav[0].cpu().to(torch.float32), model.sample_rate, strategy="loudness", loudness_compressor=True)