Update app.py

app.py

@@ -1,83 +1,66 @@
-import os
-import uuid
-import torch
-import re
+import spaces
 import gradio as gr
 import torchaudio
 from audiocraft.models import MusicGen
 from audiocraft.data.audio import audio_write
-import spaces  # Importar spaces
-
-# Decorador para gestionar el uso de GPU
-def gpu_decorator(duration):
-    def decorator(func):
-        def wrapper(*args, **kwargs):
-            with spaces.GPU(duration=duration):  # Solicitar GPU por el tiempo especificado
-                device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-                return func(*args, device=device, **kwargs)  # Pasar el dispositivo a la función
-        return wrapper
-    return decorator
+import logging
+import os 
+import uuid
+from torch.cuda.amp import autocast
+import torch
 
-# Cargar el modelo `musicgen-small` una única vez
-model = MusicGen.get_pretrained("facebook/musicgen-small")
-model.to(torch.device('cpu'))  # Inicialmente configurar el modelo para CPU
+# Configura o logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
-@gpu_decorator(duration=120)  # Decorar la función con el uso de GPU
-def generate_music(description, melody_audio, duration, device):
-    # Limpiar el texto de la descripción
-    description = clean_text(description)
-    model.set_generation_params(duration=int(duration * 1000))  # Convertir segundos a milisegundos
+logging.info("Carregando o modelo pré-treinado.")
+model = MusicGen.get_pretrained('nateraw/musicgen-songstarter-v0.2')
 
-    try:
-        # Cambiar el modelo a GPU si está disponible
-        model.to(device)
+@spaces.GPU(duration=120)
+def generate_music(description, melody_audio, duration):
+    with autocast():
+        logging.info("Iniciando a geração de música.")
+        model.set_generation_params(duration=duration)
         
-        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)
-                    wav = model.generate_with_chroma(description, melody[None], sr)
-                else:
-                    wav = model.generate(description)
+        if description:
+            description = [description]
+            if melody_audio:
+                logging.info(f"Carregando a melodia de áudio de: {melody_audio}")
+                melody, sr = torchaudio.load(melody_audio)
+                logging.info("Gerando música com descrição e melodia.")
+                wav = model.generate_with_chroma(description, melody[None], sr)
             else:
-                wav = model.generate_unconditional(1)
-
-            # Guardar el archivo de música generado
-            filename = f'{str(uuid.uuid4())}.wav'
-            path = audio_write(filename, wav[0].cpu().to(torch.float32), model.sample_rate, strategy="loudness", loudness_compressor=True)
-
-            if not os.path.exists(path):
-                raise ValueError(f'Failed to save audio to {path}')
-            
-            return path
-
-    except Exception as e:
-        return str(e)
-
-def clean_text(text):
-    text = re.sub(r'http\S+', '', text)
-    text = re.sub(r'[^a-zA-Z0-9\s]', '', text)
-    return text
+                logging.info("Gerando música apenas com descrição.")
+                wav = model.generate(description)
+        else:
+            logging.info("Gerando música de forma incondicional.")
+            wav = model.generate_unconditional(1)
+        
+        filename = f'{str(uuid.uuid4())}.wav'
+        logging.info(f"Salvando a música gerada com o nome: {filename}")
+        path = audio_write(filename, wav[0].cpu().to(torch.float32), model.sample_rate, strategy="loudness", loudness_compressor=True)
+        print("Música salva em", path, ".")
+        # Verifica a forma do tensor de áudio e se foi salvo corretamente
+        logging.info(f"A forma do tensor de áudio gerado: {wav[0].shape}")
+        logging.info("Música gerada e salva com sucesso.")
+        if not os.path.exists(path):
+            raise ValueError(f'Failed to save audio to {path}')
 
-# Definir la interfaz de Gradio
-description = gr.Textbox(label="Description", placeholder="Acoustic, guitar, melody, trap, D minor, 90 bpm")
+        return path
+    
+# Define a interface 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)
-output_path = gr.File(label="Generated Music")
+duration = gr.Slider(label="Duration (seconds)", minimum=10, maximum=600, step=10, value=30)  # Máximo 10 minutos (600 segundos)
+output_path = gr.Audio(label="Generated Music", type="filepath")
 
 gr.Interface(
     fn=generate_music,
     inputs=[description, melody_audio, duration],
     outputs=output_path,
-    title="MusicGen Melody Demo",
-    description="Generate music using the MusicGen melody model. Optionally remix with an audio file. Download the generated audio file.",
+    title="MusicGen Demo",
+    description="Generate music using the MusicGen model.",
     examples=[
-        ["happy rock", None, 8],
-        ["energetic EDM", None, 8],
-        ["chillwave", "./assets/example_melody.mp3", 10]
+        ["trap, synthesizer, songstarters, dark, G# minor, 140 bpm", "./assets/kalhonaho.mp3", 30],
+        ["upbeat, electronic, synth, dance, 120 bpm", None, 60]
     ]
 ).launch()
-