import torch
import gradio as gr
import bigvgan
import librosa
import soundfile as sf
from meldataset import get_mel_spectrogram

# Tải mô hình BigVGAN
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = bigvgan.BigVGAN.from_pretrained('nvidia/bigvgan_v2_44khz_128band_512x', use_cuda_kernel=False)
model = torch.nn.utils.parametrizations.weight_norm(model, name="weight_norm")  # Cập nhật weight_norm
model = model.eval().to(device)

# Hàm tổng hợp âm thanh
def synthesize_audio(audio_file):
    # Tải và xử lý tệp âm thanh
    wav, sr = librosa.load(audio_file, sr=model.h.sampling_rate, mono=True)
    wav = torch.FloatTensor(wav).unsqueeze(0).to(device)

    # Tính toán mel spectrogram
    mel = get_mel_spectrogram(wav, model.h).to(device)

    # Tổng hợp âm thanh từ mel spectrogram
    with torch.inference_mode():
        wav_gen = model(mel)
    
    # Chuyển đổi tín hiệu âm thanh sang định dạng PCM 16 bit
    wav_gen_float = wav_gen.squeeze(0).cpu()
    wav_gen_int16 = (wav_gen_float * 32767.0).numpy().astype('int16')
    
    # Lưu tệp âm thanh
    output_path = "output.wav"
    sf.write(output_path, wav_gen_int16, model.h.sampling_rate, subtype='PCM_16')
    
    return output_path

# Tạo giao diện Gradio
iface = gr.Interface(
    fn=synthesize_audio,
    inputs=gr.Audio(type="filepath"),
    outputs=gr.Audio(type="filepath"),
    title="BigVGAN Audio Synthesis",
    description="Upload an audio file to synthesize a new audio waveform using BigVGAN."
)

# Chạy ứng dụng
iface.launch()