Update app.py

app.py

@@ -1,56 +1,47 @@
 import gradio as gr
+import spaces
 import torchaudio
 from audiocraft.models import MusicGen
-import spaces
-import logging
-import tempfile
-import numpy as np
-
-# Configure logging
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
-
-@spaces.GPU(duration=120)
-def generate_music(description, melody_audio):
-    logging.info("Starting music generation.")
-
-    # Load the pre-trained model
-    logging.info("Loading pre-trained model.")
-    model = MusicGen.get_pretrained('nateraw/musicgen-songstarter-v0.2')
-    model.set_generation_params(duration=8)
-
-    if description:
-        description = [description]
-        if melody_audio:
-            logging.info(f"Loading audio melody from: {melody_audio}")
-            melody, sr = torchaudio.load(melody_audio)
-            logging.info("Generating music with description and melody.")
-            wav = model.generate_with_chroma(description, melody[None], sr)
-        else:
-            logging.info("Generating music with description only.")
-            wav = model.generate(description)
-    else:
-        logging.info("Generating music unconditionally.")
-        wav = model.generate_unconditional(1)
-
-    # Check the shape of the generated audio tensor
-    logging.info(f"The shape of the generated audio tensor: {wav[0].shape}")
-
-    # Convert the generated audio to the format expected by gr.Audio
-    audio_data = (model.sample_rate, (wav[0].cpu().numpy() * 32767).astype(np.int16))
-    logging.info("Music generated successfully.")
-
-    return audio_data
-
-# Create the Gradio interface
-with gr.Blocks() as demo:
-    with gr.Row():
-        with gr.Column():
-            description = gr.Textbox(label="Music Description")
-            melody_audio = gr.Audio(label="Melody Audio")
-        with gr.Column():
-            output_audio = gr.Audio(label="Generated Music", interactive=False)
-
-    generate_btn = gr.Button("Generate Music")
-    generate_btn.click(generate_music, inputs=[description, melody_audio], outputs=[output_audio])
-
-demo.launch()
+from audiocraft.data.audio import audio_write
+
+# Importação necessária para o funcionamento do modelo e manipulação de áudio
+model = MusicGen.get_pretrained('nateraw/musicgen-songstarter-v0.2')
+model.set_generation_params(duration=8)  # Duração de 8 segundos para a geração
+
+@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)
+
+    # 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)
+
+    # Retorna os caminhos dos arquivos de áudio gerados
+    return [f"output_{idx}.wav" for idx in range(len(wav))]
+
+# Define 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)
+)
+
+iface.launch()