import gradio as gr
import whisper
import os

# 加载模型 (base 模型速度快且支持多语种)
# 如果希望更精准但速度慢，可以将 "base" 改为 "small" 或 "medium"
model = whisper.load_model("base")

def speech_to_text(audio, language):
    if audio is None:
        return "请先上传音频文件或录音。"
    
    # 语言选择逻辑
    # Whisper 支持自动识别，也可手动指定提高准确率
    lang_map = {
        "自动检测": None,
        "中文 (Chinese)": "zh",
        "英语 (English)": "en",
        "日语 (Japanese)": "ja",
        "法语 (French)": "fr",
        "德语 (German)": "de"
    }
    
    selected_lang = lang_map.get(language)
    
    # 执行识别
    # task="transcribe" 为转录，如果想翻译成英文可以选 "translate"
    result = model.transcribe(audio, language=selected_lang)
    
    return result["text"]

# 创建 Gradio 界面
iface = gr.Interface(
    fn=speech_to_text,
    inputs=[
        gr.Audio(type="filepath", label="上传音频或录音"),
        gr.Dropdown(
            choices=["自动检测", "中文 (Chinese)", "英语 (English)", "日语 (Japanese)", "法语 (French)", "德语 (German)"],
            label="音频语言",
            value="自动检测"
        )
    ],
    outputs=gr.Textbox(label="转换出的文本"),
    title="AI 智能语音转文字",
    description="上传音频文件或直接录音，AI 将自动识别并转换为文字。基于 OpenAI Whisper base 模型。"
)

if __name__ == "__main__":
    iface.launch()