import gradio as gr
import torch
# import torchaudio
from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan
from datasets import load_dataset
from transformers import pipeline

def text_to_audio(text):
    # 这里是你的文本转音频的算法实现
    # 加载处理器和模型
    processor = SpeechT5Processor.from_pretrained("/home/tt/speecht5/speecht5_finetuned/checkpoint-19000")
    model = SpeechT5ForTextToSpeech.from_pretrained("/home/tt/speecht5/speecht5_finetuned/checkpoint-19000")
    
    # 加载说话者嵌入数据集
    embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
    speaker_index = 0
    speaker_embeddings = torch.tensor(embeddings_dataset[speaker_index]["xvector"]).unsqueeze(0)

    # 处理输入文本
    inputs = processor(text=text, return_tensors="pt")

    # 加载声码器
    vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan")
    
    # 生成语音
    speech = model.generate_speech(inputs["input_ids"], speaker_embeddings, vocoder=vocoder)

    # 调整音频张量的形状以匹配 torchaudio.save 的要求
    speech_2d = speech.unsqueeze(0)  # 添加一个维度，使其成为二维张量

    # 保存音频为 .wav 文件
    # torchaudio.save("generated_speech1.wav", speech_2d, 16000)

    return speech


def audio_to_text(audio):
    sr, audio_array = audio

    data_dict = {"array" : audio_array, "sampling_rate" : sr}

    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    pipe = pipeline(    
        "automatic-speech-recognition", model="/home/rouvling/temp_lug/gradio/asr/checkpoint-4000-old", device=device
    )

    result=pipe(data_dict, max_new_tokens=256)
    return result['text']


text_to_audio_interface = gr.Interface(
    fn=text_to_audio,
    inputs="text",
    outputs="audio",
    title="Text to Audio",
    description="输入文字，生成对应音频",
)

audio_to_text_interface = gr.Interface(
    fn=audio_to_text,
    inputs="audio",
    outputs="text",
    title="Audio to Text",
    description="上传音频，生成对应文字",
)

demo = gr.TabbedInterface(
    [text_to_audio_interface, audio_to_text_interface],
    ["Text to Audio", "Audio to Text"],
)

demo.launch(server_name="0.0.0.0", server_port=8080)