import gradio as gr
import torch, librosa, librosa.display, soundfile as sf
import numpy as np, io, matplotlib.pyplot as plt
from diffusers import AudioLDM2Pipeline

# Load model
model_id = "cvssp/audioldm2-large"
pipe = AudioLDM2Pipeline.from_pretrained(model_id)
pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")

def convert_voice(prompt, input_audio):
    audio, sr = librosa.load(input_audio, sr=16000)
    dur = librosa.get_duration(y=audio, sr=sr)

    result = pipe(
        prompt=prompt,
        num_inference_steps=10,
        audio_length_in_s=dur
    ).audios[0]

    out_path = "converted.wav"
    sf.write(out_path, result, 16000)

    # Spectrogram
    fig, ax = plt.subplots(figsize=(8,3))
    spec = librosa.amplitude_to_db(
        np.abs(librosa.stft(result)), ref=np.max
    )
    img = librosa.display.specshow(
        spec, sr=16000, x_axis="time", y_axis="log", ax=ax
    )
    plt.colorbar(img, ax=ax, format="%+2.0f dB")
    plt.title("Generated Spectrogram")

    buf = io.BytesIO()
    plt.savefig(buf, format="png")
    buf.seek(0)

    return out_path, buf

gr.Interface(
    fn=convert_voice,
    inputs=[
        gr.Textbox(label="Describe new voice style"),
        gr.Audio(type="filepath")
    ],
    outputs=[
        gr.Audio(label="Generated Voice"),
        gr.Image(label="Spectrogram")
    ],
    title="🎵 VoiceMorphAI",
    description="Upload voice and transform style using AudioLDM2"
).launch()