import torch
import gradio as gr
from transformers import AutoProcessor, MusicgenForConditionalGeneration
from TTS.api import TTS
import librosa
import numpy as np
import os

# Establecer variable de entorno para Coqui TTS
os.environ["COQUI_TOS_AGREED"] = "1"

# Cargar el procesador y el modelo de MusicGen
processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
model_musicgen = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")

# Cargar el modelo de TTS
model_tts = TTS(model_name="tts_models/multilingual/multi-dataset/xtts_v2", progress_bar=False).to("cpu")

def generate_music(text_prompt, audio_file=None, tts_text=None):
    # Procesar la descripción de la música
    inputs = processor(text=text_prompt, padding=True, return_tensors="pt")
    
    # Parámetros de duración máxima y mínima para la música generada
    max_duration_sec = 600
    min_duration_sec = 30
    max_new_tokens = int(max_duration_sec * model_musicgen.config.min_length)
    min_new_tokens = int(min_duration_sec * model_musicgen.config.min_length)
    
    # Generar música
    audio_values = model_musicgen.generate(**inputs)
    music = audio_values.cpu().numpy()[0]

    # Si se proporciona un archivo de audio y texto para clonar la voz
    if audio_file and tts_text:
        try:
            # Generar la voz clonada usando TTS
            cloned_audio = model_tts.tts_with_vc(tts_text, speaker_wav=audio_file, language="es")

            # Asegurarse de que la tasa de muestreo sea la correcta (44.1kHz)
            if model_tts.synthesizer.output_sample_rate != 44100:
                cloned_audio, _ = librosa.resample(cloned_audio, model_tts.synthesizer.output_sample_rate, 44100)

            # Convertir música generada en un array de floats
            music_np = librosa.util.buf_to_float(music, n_bytes=2)

            # Ajustar las longitudes de los arrays (rellenar el más corto)
            if len(music_np) > len(cloned_audio):
                cloned_audio = np.pad(cloned_audio, (0, len(music_np) - len(cloned_audio)))
            else:
                music_np = np.pad(music_np, (0, len(cloned_audio) - len(music_np)))

            # Combinar la música y el audio clonado
            combined_audio = music_np + cloned_audio

            # Normalizar y convertir a formato int16
            combined_audio = combined_audio / np.max(np.abs(combined_audio))
            combined_audio_int16 = (combined_audio * 32767).astype(np.int16)

            return (44100, combined_audio_int16)
        except Exception as e:
            print(f"Error combinando audio: {e}")
            return (44100, music)

    return (44100, music)

# Crear la interfaz Gradio
iface = gr.Interface(
    fn=generate_music,
    inputs=[
        gr.Textbox(label="Descripción de la música"),
        gr.Audio(type="filepath", label="Subir audio de voz (Opcional)"),
        gr.Textbox(label="Texto para clonar la voz (Opcional)"),
    ],
    outputs=gr.Audio(label="Música generada", type="numpy"),
    title="Generador de Música con MusicGen y XTTS",
    description="Introduce una descripción de la música que deseas generar y opcionalmente un audio de voz para clonar con texto.",
)

# Lanzar la interfaz
iface.launch()