import torch
import whisperx
import gradio as gr

# 检测设备类型
device = "cuda" if torch.cuda.is_available() else "cpu"

# 设置计算类型
compute_type = "float16" if device == "cuda" else "int8"

# 加载 WhisperX 模型
model = whisperx.load_model("large-v3", device=device, compute_type=compute_type)

def transcribe(audio_path):
    # 使用WhisperX进行转录
    result = model.transcribe(audio_path)
    return result['text']

# 创建Gradio接口
iface = gr.Interface(
    fn=transcribe,
    inputs=gr.Audio(source="upload", type="filepath"),
    outputs="text",
    title="WhisperX 语音转文字",
    description="上传音频文件，使用WhisperX模型进行转录。"
)

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