Update app.py

app.py

@@ -8,30 +8,22 @@ import uuid
 import torch
 import re
 
+# Configuración del dispositivo
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-def load_model(model_name):
-    return MusicGen.get_pretrained(model_name)
+# Cargar el modelo `musicgen-small` una única vez
+model = MusicGen.get_pretrained("facebook/musicgen-small")
 
-models = {
-    'Large': load_model('nateraw/musicgen-songstarter-v0.2'),
-    'Small': load_model('facebook/musicgen-small')
-}
-
-def get_model(model_choice):
-    return models.get(model_choice, models['Large'])
-
-@spaces.GPU(duration=0)
-def generate_music(description, melody_audio, duration, model_choice):
-    model = get_model(model_choice)
+def generate_music(description, melody_audio, duration):
     description = clean_text(description)
-    model.set_generation_params(duration=int(duration * 1000))
+    model.set_generation_params(duration=int(duration * 1000))  # Convertir segundos a milisegundos
 
     try:
         with torch.no_grad():
             if description:
                 description = [description]
                 if melody_audio:
+                    # Cargar el archivo de audio para remixar
                     melody, sr = torchaudio.load(melody_audio, normalize=True)
                     melody = melody.to(device) if torch.cuda.is_available() else melody
                     wav = model.generate_with_chroma(description, melody[None], sr)
@@ -56,20 +48,21 @@ def clean_text(text):
     text = re.sub(r'[^a-zA-Z0-9\s]', '', text)
     return text
 
+# Definir la interfaz de Gradio
 description = gr.Textbox(label="Description", placeholder="Acoustic, guitar, melody, trap, D minor, 90 bpm")
 melody_audio = gr.Audio(label="Melody Audio (optional)", type="filepath")
 duration = gr.Number(label="Duration (seconds)", value=10, precision=0)
-model_choice = gr.Radio(choices=['Small', 'Large'], label="Model Choice", value='Large')
-output_path = gr.Audio(label="Generated Music", type="filepath")
+output_path = gr.File(label="Generated Music")
 
 gr.Interface(
     fn=generate_music,
-    inputs=[description, melody_audio, duration, model_choice],
+    inputs=[description, melody_audio, duration],
     outputs=output_path,
-    title="MusicGen Demo",
-    description="Generate music using the MusicGen model. Choose between the Small or Large model.",
+    title="MusicGen Small Demo",
+    description="Generate music using the MusicGen small model. Optionally remix with an audio file. Download the generated audio file.",
     examples=[
-        ["trap, synthesizer, music starters, dark, G# minor, 140 bpm", "./assets/kalhonaho.mp3", 20, 'Large'],
-        ["upbeat, electronic, synthesizer, dance, 120 bpm", None, 30, 'Small']
+        ["happy rock", None, 8],
+        ["energetic EDM", None, 8],
+        ["chillwave", "./assets/example_melody.mp3", 10]
     ]
 ).launch()