Update app.py

app.py

@@ -4,36 +4,29 @@ import torchaudio
 from audiocraft.models import MusicGen
 from audiocraft.data.audio import audio_write
 
-# Carrega o modelo pré-treinado
 model = MusicGen.get_pretrained('nateraw/musicgen-songstarter-v0.2')
-model.set_generation_params(duration=8)  # Configura a duração da geração para 8 segundos
+model.set_generation_params(duration=8)  # generate 8 seconds.
 
-@spaces.GPU(duration=120)  # Habilita o uso de GPU
-def generate_audio(prompt, audio_input=None):
-    descriptions = [prompt]  # Usa o prompt como descrição
-    if audio_input is None:
-        # Geração incondicional com descrições
-        wav = model.generate(descriptions)
+@spaces.GPU(duration=120)  # Specify duration if the function is expected to take more than 60s
+def generate_music(description, audio_file):
+    if audio_file is None:
+        wav = model.generate([description])  # generates 1 sample based on the provided description
     else:
-        # Carrega o áudio de entrada e gera com base nele e nas descrições
-        melody, sr = torchaudio.load(audio_input)
-        # Assegura que o tensor está no formato de ponto flutuante de precisão simples
-        melody = melody.float()
-        wav = model.generate_with_chroma(descriptions, melody[None], sr)
+        melody, sr = torchaudio.load(audio_file)
+        wav = model.generate_with_chroma([description], melody[None], sr)  # generates using the melody from the given audio and the provided description
+    
+    audio_write('output', wav[0].cpu(), model.sample_rate, strategy="loudness", loudness_compressor=True)
+    return 'output.wav'
 
-    # Salva o arquivo de áudio gerado
-    output_path = 'generated_audio.wav'
-    audio_write(output_path, wav[0].cpu().float(), model.sample_rate, strategy="loudness", loudness_compressor=True)
-
-    return output_path
-
-# Cria a interface de usuário com Gradio
 iface = gr.Interface(
-    fn=generate_audio,
+    fn=generate_music,
     inputs=[
-        gr.Textbox(label="Prompt", placeholder="Enter a description..."),
-        gr.File(label="Audio Input (optional)", type="filepath")],
-    outputs=gr.File(label="Generated Audio")
+        gr.Text(label="Description"),
+        gr.Audio(type="filepath", label="Audio File (optional)")
+    ],
+    outputs=gr.Audio(type="file"),
+    title="MusicGen",
+    description="Generate music using the MusicGen model. Provide a description and optionally an audio file for melody.",
 )
 
 iface.launch()