import os import tempfile import torch import whisperx from flask import Flask, request, jsonify, render_template from waitress import serve import logging import webbrowser from threading import Timer import shutil import sys import ffmpeg try: from whisperx.diarize import DiarizationPipeline except Exception: DiarizationPipeline = None # --- 全局配置与初始化 --- logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') def _configure_torch_load_compat(): """ PyTorch >=2.6 changed torch.load default `weights_only=True`, which can break some third-party checkpoints (e.g. pyannote VAD models used by whisperx). We only allowlist OmegaConf config objects required by those checkpoints. """ try: import torch.serialization from omegaconf import DictConfig, ListConfig from omegaconf.base import ContainerMetadata torch.serialization.add_safe_globals([DictConfig, ListConfig, ContainerMetadata]) logging.info("已配置 torch.load 安全全局白名单 (omegaconf DictConfig/ListConfig)。") except Exception as e: logging.warning(f"torch.load 兼容性配置跳过: {e}") def _env_bool(name: str, default: bool) -> bool: val = os.environ.get(name) if val is None: return default return val.strip().lower() in {"1", "true", "yes", "y", "on"} def get_hf_token(): """ 获取 Hugging Face 令牌。 优先从当前目录的 'token.txt' 文件读取，如果失败则从环境变量 'HUGGING_FACE_TOKEN' 读取。 """ token = None token_file = 'token.txt' if os.path.exists(token_file): try: with open(token_file, 'r', encoding='utf-8') as f: token = f.read().strip() if token: logging.info(f"成功从 {token_file} 文件中读取 Hugging Face 令牌。") return token except Exception as e: logging.warning(f"无法从 {token_file} 读取令牌: {e}") token = os.environ.get("HUGGING_FACE_TOKEN") if token: logging.info("成功从环境变量中读取 Hugging Face 令牌。") else: logging.warning("在 token.txt 或环境变量中均未找到 Hugging Face 令牌。说话人分离功能将被禁用。") return token HF_TOKEN = get_hf_token() _configure_torch_load_compat() # 设备和计算类型配置 DEVICE = "cuda" if torch.cuda.is_available() else "cpu" COMPUTE_TYPE = "float16" if torch.cuda.is_available() else "int8" BATCH_SIZE = 16 if DEVICE == "cuda" else 8 VAD_METHOD = os.environ.get("VAD_METHOD") or ("silero" if DEVICE == "cpu" else "pyannote") logging.info(f"使用设备: {DEVICE}，计算类型: {COMPUTE_TYPE}") logging.info(f"VAD 方法: {VAD_METHOD}") # 模型配置 ALLOWED_MODELS = ['tiny', 'base', 'small', 'medium', 'large-v1', 'large-v2', 'large-v3', 'large-v3-turbo'] DEFAULT_MODEL = os.environ.get("DEFAULT_MODEL") or ("small" if DEVICE == "cpu" else "large-v3") ALLOW_LARGE_ON_CPU = _env_bool("ALLOW_LARGE_ON_CPU", False) # 模型缓存 whisper_models_cache = {} diarize_model = None diarize_model_loaded = False align_models_cache = {} def get_whisper_model(model_name: str): if model_name not in whisper_models_cache: logging.info(f"正在加载 Whisper 模型 '{model_name}'...") try: try: model = whisperx.load_model(model_name, DEVICE, compute_type=COMPUTE_TYPE, vad_method=VAD_METHOD) except TypeError: model = whisperx.load_model(model_name, DEVICE, compute_type=COMPUTE_TYPE) whisper_models_cache[model_name] = model logging.info(f"模型 '{model_name}' 加载成功。") except Exception as e: logging.error(f"加载 Whisper 模型 '{model_name}' 失败: {e}") if str(e).find('huggingface'): print(f"\n\n=======可能模型下载失败，请尝试科学上网后再次重试=======\n\n") raise return whisper_models_cache[model_name] def get_align_model(language_code: str): if language_code not in align_models_cache: logging.info(f"正在加载对齐模型 (language={language_code})...") model_a, metadata = whisperx.load_align_model(language_code=language_code, device=DEVICE) align_models_cache[language_code] = (model_a, metadata) logging.info("对齐模型加载成功。") return align_models_cache[language_code] def get_diarize_model(): global diarize_model, diarize_model_loaded if not diarize_model_loaded: logging.info("正在尝试加载说话人分离模型...") if DiarizationPipeline is None: logging.warning("未检测到说话人分离依赖 (DiarizationPipeline)，此功能将被禁用。") diarize_model_loaded = True return None if not HF_TOKEN: diarize_model_loaded = True return None try: diarize_model = DiarizationPipeline(use_auth_token=HF_TOKEN, device=DEVICE) diarize_model_loaded = True logging.info("说话人分离模型加载成功。") except Exception as e: logging.error(f"严重错误: 说话人分离模型加载失败。此功能将被禁用。错误信息: {e}") diarize_model = None diarize_model_loaded = True return diarize_model # --- Flask 应用 --- app = Flask(__name__, template_folder='.') @app.route('/', methods=['GET']) def index(): return render_template('index.html') @app.route('/v1/audio/transcriptions', methods=['POST']) def audio_transcriptions(): if 'file' not in request.files: return jsonify({"error": "请求中未包含文件部分"}), 400 file = request.files['file'] if file.filename == '': return jsonify({"error": "未选择任何文件"}), 400 print(request.form) model_id = request.form.get('model', DEFAULT_MODEL) model_name = 'large-v3' if model_id == 'large-v3-turbo' else model_id if model_name not in ALLOWED_MODELS: model_name = DEFAULT_MODEL if DEVICE == "cpu" and (model_name.startswith("large-") or model_name == "large") and not ALLOW_LARGE_ON_CPU: logging.warning(f"CPU 环境下请求大模型 '{model_name}'，将自动降级为 'small' (可通过 ALLOW_LARGE_ON_CPU=1 关闭降级)。") model_name = "small" language = request.form.get('language') or None prompt = request.form.get('prompt') max_speakers=int(request.form.get('max_speakers',-1)) min_speakers=int(request.form.get('min_speakers',0)) logging.info(f"收到请求: 模型='{model_id}', 语言='{language or '自动检测'}', 提示词='{'有' if prompt else '无'}'") input_file_path = None processed_wav_path = None try: suffix = os.path.splitext(file.filename)[1] with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp: file.save(tmp.name) input_file_path = tmp.name logging.info(f"正在将上传的文件 '{file.filename}' 转换为标准的 16kHz 单声道 WAV 格式...") with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_wav: processed_wav_path = tmp_wav.name try: ( ffmpeg .input(input_file_path) .output(processed_wav_path, ac=1, ar=16000, acodec='pcm_s16le', vn=None) .run(capture_stdout=True, capture_stderr=True, overwrite_output=True) ) logging.info("文件格式转换成功。") except ffmpeg.Error as e: error_details = e.stderr.decode('utf-8', errors='ignore') logging.error(f"FFmpeg 文件转换失败: {error_details}") return jsonify({"error": f"音频/视频文件处理失败，可能是文件已损坏或格式不受支持。"}), 400 audio = whisperx.load_audio(processed_wav_path) model = get_whisper_model(model_name) # --- # *** FIX IS HERE *** # --- transcribe_options = {} if language: transcribe_options['language'] = language if prompt: # 使用正确的参数名 'prompt' transcribe_options['prompt'] = prompt print('开始转录') result = model.transcribe(audio, batch_size=BATCH_SIZE, **transcribe_options) print('转录结束，准备对齐') model_a, metadata = get_align_model(result["language"]) result = whisperx.align(result["segments"], model_a, metadata, audio, DEVICE, return_char_alignments=False) if max_speakers>-1: print('进入说话人识别') diar_model = get_diarize_model() if diar_model: try: diarize_segments = diar_model(audio,max_speakers=max_speakers if max_speakers>0 else None,min_speakers=min_speakers if min_speakers>0 else None) result = whisperx.assign_word_speakers(diarize_segments, result) except Exception as e: logging.error(f"说话人分离运行时失败: {e}。将回退到单说话人模式。") speakers = {segment.get('speaker') for segment in result["segments"] if 'speaker' in segment} is_single_speaker = len(speakers) <= 1 logging.info(f"检测到的说话人: {speakers}。单说话人模式: {'是' if is_single_speaker else '否'}") speaker_mapping = {f"SPEAKER_{i:02d}": f"Speaker{i+1}" for i in range(20)} print(result) formatted_segments = [] for segment in result["segments"]: speaker_raw = segment.get("speaker", "SPEAKER_00") speaker_name = speaker_mapping.get(speaker_raw, speaker_raw) text = segment['text'].strip() if not text: continue tmp={ "start": segment['start'], "end": segment['end'], "text": text } segment_speaker = speaker_name if not is_single_speaker else None if segment_speaker: tmp['speaker']=segment_speaker formatted_segments.append(tmp) response_data = {"segments": formatted_segments} return jsonify(response_data) except Exception as e: logging.error(f"处理流程中发生未知错误: {e}", exc_info=True) return jsonify({"error": "处理过程中发生内部错误。"}), 500 finally: if input_file_path and os.path.exists(input_file_path): os.remove(input_file_path) logging.info(f"已清理临时上传文件: {input_file_path}") if processed_wav_path and os.path.exists(processed_wav_path): os.remove(processed_wav_path) logging.info(f"已清理临时WAV文件: {processed_wav_path}") # --- 启动服务 --- def check_ffmpeg(): if not shutil.which("ffmpeg"): logging.error("错误: 系统 PATH 中未找到 FFmpeg。") print("\n错误: 系统 PATH 中未找到 FFmpeg。") print("请确保您已安装 FFmpeg 并且其路径已添加到系统环境变量中。") print("Windows 安装指南: https://www.wikihow.com/Install-FFmpeg-on-Windows") print("macOS (使用 Homebrew): brew install ffmpeg") print("Linux (Ubuntu/Debian): sudo apt update && sudo apt install ffmpeg") sys.exit(1) logging.info("FFmpeg 环境检查通过。") def open_browser(url): webbrowser.open_new(url) if __name__ == '__main__': check_ffmpeg() host = os.environ.get("HOST", "127.0.0.1") port = int(os.environ.get("PORT", "9092")) url = f"http://{host}:{port}" running_in_space = bool(os.environ.get("SPACE_ID")) or bool(os.environ.get("HF_SPACE")) or bool(os.environ.get("SYSTEM") == "spaces") if _env_bool("OPEN_BROWSER", True) and not running_in_space: Timer(1, lambda: open_browser(url)).start() logging.info(f"服务已启动，正在监听 http://{host}:{port}") serve(app, host=host, port=port, threads=10)