Update app.py

app.py

@@ -4,44 +4,34 @@ import torchaudio
 from audiocraft.models import MusicGen
 from audiocraft.data.audio import audio_write
 
-# Importação necessária para o funcionamento do modelo e manipulação de áudio
+# Carrega o modelo pré-treinado
 model = MusicGen.get_pretrained('nateraw/musicgen-songstarter-v0.2')
-model.set_generation_params(duration=8)  # Duração de 8 segundos para a geração
+model.set_generation_params(duration=8)  # Configura a duração da geração para 8 segundos
 
-@spaces.GPU(duration=120)  # Decorador para habilitar o uso de GPU
-def generate_audio(mode, descriptions=None, melody_path=None):
-    if mode == 'unconditional':
-        wav = model.generate_unconditional(4)
-    elif mode == 'descriptions':
-        if descriptions is None:
-            descriptions = ['acoustic, guitar, melody, trap, d minor, 90 bpm']
-        descriptions = descriptions.split(",")  # Converte a string em lista
-        wav = model.generate(descriptions * 3)  # Gera 3 amostras
-    elif mode == 'melody':
-        if melody_path is None:
-            return "Melody path cannot be empty for melody mode."
-        melody, sr = torchaudio.load(melody_path)
-        if descriptions is None:
-            descriptions = ['acoustic, guitar, melody, trap, d minor, 90 bpm']
-        descriptions = descriptions.split(",")  # Converte a string em lista
-        wav = model.generate_with_chroma(descriptions, melody[None].expand(3, -1, -1), sr)
+@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)
+    else:
+        # Carrega o áudio de entrada e gera com base nele e nas descrições
+        melody, sr = torchaudio.load(audio_input)
+        wav = model.generate_with_chroma(descriptions, melody[None], sr)
 
-    # Salva os arquivos de áudio gerados
-    for idx, one_wav in enumerate(wav):
-        audio_write(f'output_{idx}.wav', one_wav.cpu(), model.sample_rate,
-                    strategy="loudness", loudness_compressor=True)
+    # Salva o arquivo de áudio gerado
+    output_path = 'generated_audio.wav'
+    audio_write(output_path, wav[0].cpu(), model.sample_rate, strategy="loudness", loudness_compressor=True)
 
-    # Retorna os caminhos dos arquivos de áudio gerados
-    return [f"output_{idx}.wav" for idx in range(len(wav))]
+    return output_path
 
-# Define a interface de usuário com Gradio
+# Cria a interface de usuário com Gradio
 iface = gr.Interface(
     fn=generate_audio,
     inputs=[
-        gr.Dropdown(['unconditional', 'descriptions', 'melody'], label="Generation Mode"),
-        gr.Textbox(label="Descriptions (comma-separated, optional)", optional=True),
-        gr.File(label="Melody File Path (.mp3, optional)", optional=True, type="filepath")],
-    outputs=gr.File(label="Generated Audio", type="file", multiple=True)
+        gr.Textbox(label="Prompt", placeholder="Enter a description..."),
+        gr.File(label="Audio Input (optional)", type="filepath", optional=True)],
+    outputs=gr.File(label="Generated Audio")
 )
 
 iface.launch()