# TTS and ASR API for macOS Silicon with HuggingFace transformers import asyncio import base64 import hashlib import logging import os import platform import sys import time import traceback from dataclasses import dataclass from pathlib import Path from typing import Optional, Dict, Any from fastapi import APIRouter, HTTPException, Security from pydantic import BaseModel import numpy as np router = APIRouter() logger = logging.getLogger("tts_asr") # Environment variables TTS_ASR_DEVICE = os.environ.get("TTS_ASR_DEVICE", "auto") TTS_ASR_WARMUP = os.environ.get("TTS_ASR_WARMUP", "true").lower() == "true" TTS_ASR_WARMUP_TIMEOUT = int(os.environ.get("TTS_ASR_WARMUP_TIMEOUT", "120")) TTS_ASR_IDLE_TIMEOUT = int(os.environ.get("TTS_ASR_IDLE_TIMEOUT", "0")) # New environment variables for macOS optimization TTS_ASR_MODEL_SIZE = os.environ.get("TTS_ASR_MODEL_SIZE", "auto") # tiny/base/small/medium/large/turbo TTS_ASR_QUANTIZE = os.environ.get("TTS_ASR_QUANTIZE", "false").lower() == "true" TTS_ASR_OFFLINE_MODE = os.environ.get("TTS_ASR_OFFLINE_MODE", "false").lower() == "true" TTS_ASR_MPS_MEMORY_LIMIT_MB = int(os.environ.get("TTS_ASR_MPS_MEMORY_LIMIT_MB", "8192")) # 8GB default # Warmup constants TTS_WARMUP_TEXT = "你好，这是一个测试。" ASR_WARMUP_AUDIO_SECONDS = 0.5 # Model size mappings for Whisper WHISPER_MODEL_SIZES = { "tiny": "openai/whisper-tiny", "base": "openai/whisper-base", "small": "openai/whisper-small", "medium": "openai/whisper-medium", "large": "openai/whisper-large-v3", "turbo": "openai/whisper-large-v3-turbo", } # Apple Silicon recommended models APPLE_SILICON_DEFAULT_SIZE = "small" # Better for MPS memory constraints @dataclass class DeviceCapabilities: """设备能力检测结果""" device: str mps_available: bool = False mps_memory_limit_mb: Optional[int] = None cuda_available: bool = False cuda_memory_limit_mb: Optional[int] = None recommended_model_size: str = "large" supports_quantization: bool = True fallback_device: Optional[str] = None # Global state _tts_pipeline = None _asr_pipeline = None _asr_model_size: Optional[str] = None _device_caps: Optional[DeviceCapabilities] = None _tts_last_used = 0.0 _asr_last_used = 0.0 _tts_loading = False _asr_loading = False _tts_lock = asyncio.Lock() _asr_lock = asyncio.Lock() def _is_apple_silicon() -> bool: """检测是否为Apple Silicon (M1/M2/M3)""" return ( platform.system() == "Darwin" and platform.machine() == "arm64" ) def _get_system_memory_mb() -> int: """获取系统总内存(MB)，用于Apple Silicon内存管理""" try: # pragma: no cover import psutil return int(psutil.virtual_memory().total / (1024 * 1024)) except Exception: # 默认假设8GB return 8192 def _detect_device_capabilities() -> DeviceCapabilities: """ 全面检测设备能力，包括MPS/CUDA可用性和内存限制 返回结构化的设备能力对象 """ global _device_caps if _device_caps is not None: return _device_caps try: import torch caps = DeviceCapabilities(device="cpu") # 检测MPS (Apple Silicon) if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): if torch.backends.mps.is_built(): try: # 更全面的MPS测试 - 测试较大的张量操作 test_size = 1000 test_tensor = torch.randn(test_size, test_size, device="mps") _ = torch.mm(test_tensor, test_tensor) del test_tensor torch.mps.empty_cache() caps.mps_available = True caps.device = "mps" # Apple Silicon内存管理 - 使用系统内存的一部分 system_mem = _get_system_memory_mb() # MPS可以使用系统内存，但限制在配置值以内 caps.mps_memory_limit_mb = min( TTS_ASR_MPS_MEMORY_LIMIT_MB, int(system_mem * 0.6) # 使用不超过60%的系统内存 ) # Apple Silicon推荐使用更小的模型 if _is_apple_silicon(): caps.recommended_model_size = APPLE_SILICON_DEFAULT_SIZE logger.info("[Device] Apple Silicon detected, recommending %s model", caps.recommended_model_size) logger.info("[Device] MPS可用，内存限制: %d MB", caps.mps_memory_limit_mb) except Exception as e: logger.warning("[Device] MPS测试失败: %s，降级到CPU", str(e)) caps.mps_available = False caps.fallback_device = "cpu" # 检测CUDA if not caps.mps_available and torch.cuda.is_available(): try: gpu_count = torch.cuda.device_count() if gpu_count > 0: # 测试CUDA操作 test_tensor = torch.randn(100, 100, device="cuda:0") _ = torch.mm(test_tensor, test_tensor) del test_tensor torch.cuda.empty_cache() caps.cuda_available = True caps.device = "cuda" # 获取GPU显存 gpu_mem = torch.cuda.get_device_properties(0).total_memory caps.cuda_memory_limit_mb = int(gpu_mem / (1024 * 1024)) logger.info("[Device] CUDA可用，GPU显存: %d MB", caps.cuda_memory_limit_mb) except Exception as e: logger.warning("[Device] CUDA测试失败: %s，降级到CPU", str(e)) caps.cuda_available = False caps.fallback_device = "cpu" # 如果MPS和CUDA都不可用，使用CPU if not caps.mps_available and not caps.cuda_available: caps.device = "cpu" logger.info("[Device] 使用CPU") _device_caps = caps return caps except Exception as e: logger.error("[Device] 设备检测失败: %s", str(e)) return DeviceCapabilities(device="cpu") def _test_device_capability(device_str: str) -> tuple[bool, str]: """ 测试设备实际可用性（兼容性保留） 返回: (是否可用, 错误信息) """ caps = _detect_device_capabilities() if device_str == "cpu": return True, "" if device_str == "mps": if caps.mps_available: return True, "" else: return False, "MPS 不可用或测试失败" if device_str.startswith("cuda"): if caps.cuda_available: return True, "" else: return False, "CUDA 不可用或测试失败" return False, f"未知设备类型: {device_str}" def _get_device() -> str: """ 获取最佳计算设备，支持环境变量覆盖和降级策略 """ caps = _detect_device_capabilities() # 环境变量强制指定 if TTS_ASR_DEVICE == "cpu": logger.info("[Device] 强制使用 CPU (环境变量)") return "cpu" elif TTS_ASR_DEVICE == "mps": if caps.mps_available: logger.info("[Device] 强制使用 MPS (环境变量)") return "mps" else: logger.warning("[Device] MPS不可用，降级到CPU") return "cpu" elif TTS_ASR_DEVICE == "cuda": if caps.cuda_available: logger.info("[Device] 强制使用 CUDA (环境变量)") return "cuda" else: logger.warning("[Device] CUDA不可用，降级到CPU") return "cpu" # 自动选择 return caps.device def _get_recommended_model_size() -> str: """ 根据设备能力推荐合适的模型大小 """ # 优先使用环境变量配置 if TTS_ASR_MODEL_SIZE != "auto": size = TTS_ASR_MODEL_SIZE.lower() if size in WHISPER_MODEL_SIZES: logger.info("[Model] 使用环境变量指定的模型大小: %s", size) return size else: logger.warning("[Model] 无效的模型大小 '%s'，使用自动选择", size) # 根据设备能力自动选择 caps = _detect_device_capabilities() recommended = caps.recommended_model_size # Apple Silicon特别处理 if _is_apple_silicon(): recommended = APPLE_SILICON_DEFAULT_SIZE logger.info("[Model] Apple Silicon自动选择模型大小: %s", recommended) return recommended def _device_arg() -> str: device = _get_device() if device == "cuda": return "cuda:0" return device def _get_torch_dtype(): device = _get_device() import torch # Apple Silicon MPS支持float16，但在某些操作上可能不稳定，默认使用float32 if device == "mps": # MPS环境下使用float32更稳定，避免潜在的数值问题 return torch.float32 return torch.float16 if device != "cpu" else torch.float32 def _check_model_cached(model_id: str) -> bool: """ 检查模型是否已在本地缓存 """ if not TTS_ASR_OFFLINE_MODE: return True # 非离线模式，不检查缓存 try: import os from huggingface_hub.constants import HF_HUB_CACHE cache_dir = HF_HUB_CACHE # 简单的缓存检查：查找模型目录 model_name = model_id.replace("/", "--") model_cache_path = Path(cache_dir) / f"models--{model_name}" if model_cache_path.exists(): # 检查是否有snapshots目录 snapshots_dir = model_cache_path / "snapshots" if snapshots_dir.exists() and any(snapshots_dir.iterdir()): logger.info("[Cache] 模型 %s 已缓存", model_id) return True logger.warning("[Cache] 模型 %s 未缓存，离线模式将失败", model_id) return False except Exception as e: logger.warning("[Cache] 缓存检查失败: %s", str(e)) return not TTS_ASR_OFFLINE_MODE # 如果检查失败且是离线模式，返回False def _clear_cuda_cache(): try: import torch caps = _detect_device_capabilities() if caps.cuda_available: torch.cuda.empty_cache() except Exception: pass def _clear_mps_cache(): try: import torch caps = _detect_device_capabilities() if caps.mps_available: torch.mps.empty_cache() except Exception: pass async def _load_tts_pipeline_with_retry(max_retries: int = 2) -> bool: # pragma: no cover """ 加载TTS管道，支持重试和降级 """ global _tts_pipeline, _tts_loading async with _tts_lock: if _tts_pipeline is not None: return True if _tts_loading: return False _tts_loading = True try: import torch from transformers import pipeline current_device = _get_device() for attempt in range(max_retries): try: device_to_use = _device_arg() torch_dtype = _get_torch_dtype() model_id = "suno/bark" # 离线模式检查 if TTS_ASR_OFFLINE_MODE and not _check_model_cached(model_id): logger.error("[TTS] 离线模式下模型 %s 未缓存", model_id) return False logger.info("[TTS] 加载 suno/bark 模型 (尝试 %d/%d, 设备: %s)...", attempt + 1, max_retries, device_to_use) _tts_pipeline = await asyncio.to_thread( lambda: pipeline( "text-to-speech", model=model_id, trust_remote_code=True, device=device_to_use, torch_dtype=torch_dtype, ) ) logger.info("[TTS] suno/bark 模型加载完成") return True except RuntimeError as e: error_str = str(e) if "MPS" in error_str or "mps" in error_str: logger.warning("[TTS] MPS 推理失败，尝试降级到 CPU: %s", error_str) caps = _detect_device_capabilities() caps.mps_available = False caps.device = "cpu" _clear_mps_cache() continue elif "CUDA" in error_str or "cuda" in error_str: logger.warning("[TTS] CUDA 推理失败，尝试降级到 CPU: %s", error_str) caps = _detect_device_capabilities() caps.cuda_available = False caps.device = "cpu" _clear_cuda_cache() continue else: raise except Exception as e: logger.error("[TTS] 加载失败: %s", str(e)) if attempt == max_retries - 1: raise await asyncio.sleep(1) return _tts_pipeline is not None finally: _tts_loading = False async def _load_asr_pipeline_with_retry(max_retries: int = 2) -> bool: # pragma: no cover """ 加载ASR管道，支持重试、降级、模型大小选择和量化 """ global _asr_pipeline, _asr_loading, _asr_model_size async with _asr_lock: if _asr_pipeline is not None: return True if _asr_loading: return False _asr_loading = True try: import torch from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline # 确定模型大小 model_size = _get_recommended_model_size() model_id = WHISPER_MODEL_SIZES.get(model_size, WHISPER_MODEL_SIZES["large"]) # 如果是离线模式，检查缓存 if TTS_ASR_OFFLINE_MODE and not _check_model_cached(model_id): logger.error("[ASR] 离线模式下模型 %s 未缓存", model_id) return False _asr_model_size = model_size for attempt in range(max_retries): try: device_to_use = _device_arg() torch_dtype = _get_torch_dtype() logger.info("[ASR] 加载 Whisper %s 模型 (尝试 %d/%d, 设备: %s, 量化: %s)...", model_size, attempt + 1, max_retries, device_to_use, "是" if TTS_ASR_QUANTIZE else "否") def load_model(): # 量化加载选项 load_kwargs = { "torch_dtype": torch_dtype, "low_cpu_mem_usage": True, "use_safetensors": True, } # 仅在CPU或CUDA环境下支持8-bit量化 if TTS_ASR_QUANTIZE and device_to_use in ["cpu", "cuda:0"]: try: load_kwargs["load_in_8bit"] = True load_kwargs["device_map"] = "auto" logger.info("[ASR] 使用8-bit量化加载模型") except Exception as e: logger.warning("[ASR] 8-bit量化不可用: %s，使用常规加载", str(e)) model = AutoModelForSpeechSeq2Seq.from_pretrained( model_id, **load_kwargs ) processor = AutoProcessor.from_pretrained(model_id) # 如果使用了device_map（量化模式），不需要指定device参数 if "load_in_8bit" in load_kwargs and load_kwargs["load_in_8bit"]: return pipeline( "automatic-speech-recognition", model=model, tokenizer=processor.tokenizer, feature_extractor=processor.feature_extractor, torch_dtype=torch_dtype, ) else: return pipeline( "automatic-speech-recognition", model=model, tokenizer=processor.tokenizer, feature_extractor=processor.feature_extractor, torch_dtype=torch_dtype, device=device_to_use, ) _asr_pipeline = await asyncio.to_thread(load_model) logger.info("[ASR] Whisper %s 模型加载完成", model_size) return True except RuntimeError as e: error_str = str(e) if "MPS" in error_str or "mps" in error_str: logger.warning("[ASR] MPS 推理失败，尝试降级到 CPU: %s", error_str) caps = _detect_device_capabilities() caps.mps_available = False caps.device = "cpu" _clear_mps_cache() continue elif "CUDA" in error_str or "cuda" in error_str: logger.warning("[ASR] CUDA 推理失败，尝试降级到 CPU: %s", error_str) caps = _detect_device_capabilities() caps.cuda_available = False caps.device = "cpu" _clear_cuda_cache() continue else: raise except Exception as e: logger.error("[ASR] 加载失败: %s", str(e)) if attempt == max_retries - 1: raise await asyncio.sleep(1) return _asr_pipeline is not None finally: _asr_loading = False def _get_tts_pipeline(): # pragma: no cover """同步获取TTS管道（已弃用，保留兼容性）""" if _tts_pipeline is not None: return _tts_pipeline raise RuntimeError("TTS 管道未加载，请使用 _load_tts_pipeline_with_retry()") def _get_asr_pipeline(): # pragma: no cover """同步获取ASR管道（已弃用，保留兼容性）""" if _asr_pipeline is not None: return _asr_pipeline raise RuntimeError("ASR 管道未加载，请使用 _load_asr_pipeline_with_retry()") async def _warmup_tts() -> bool: # pragma: no cover """ 预热TTS模型，减少首次请求延迟 """ global _tts_last_used try: logger.info("[TTS] 开始预热...") if not await _load_tts_pipeline_with_retry(): logger.error("[TTS] 预热失败：无法加载管道") return False tts = _tts_pipeline if tts is None: return False def warmup_inference(): try: result = tts(TTS_WARMUP_TEXT) if isinstance(result, dict): audio = result.get("audio") if hasattr(audio, "cpu"): _ = audio.cpu() return True except Exception as e: logger.warning("[TTS] 预热推理失败（可忽略）: %s", str(e)) return False success = await asyncio.to_thread(warmup_inference) _tts_last_used = time.time() if success: logger.info("[TTS] 预热完成") return success except Exception as e: logger.error("[TTS] 预热异常: %s", str(e)) traceback.print_exc() return False async def _warmup_asr() -> bool: # pragma: no cover """ 预热ASR模型，减少首次请求延迟 """ global _asr_last_used try: logger.info("[ASR] 开始预热...") if not await _load_asr_pipeline_with_retry(): logger.error("[ASR] 预热失败：无法加载管道") return False asr = _asr_pipeline if asr is None: return False silence_samples = int(16000 * ASR_WARMUP_AUDIO_SECONDS) silence_audio = np.zeros(silence_samples, dtype=np.float32) def warmup_inference(): try: result = asr( silence_audio, sampling_rate=16000, generate_kwargs={"language": "zh", "task": "transcribe"}, ) return True except Exception as e: logger.warning("[ASR] 预热推理失败（可忽略）: %s", str(e)) return False success = await asyncio.to_thread(warmup_inference) _asr_last_used = time.time() if success: logger.info("[ASR] 预热完成") return success except Exception as e: logger.error("[ASR] 预热异常: %s", str(e)) traceback.print_exc() return False async def _warmup_all() -> tuple[bool, bool]: # pragma: no cover """ 预热所有模型 返回: (TTS预热结果, ASR预热结果) """ logger.info("[Warmup] 开始预热所有模型 (超时: %d秒)", TTS_ASR_WARMUP_TIMEOUT) try: tts_task = asyncio.create_task(_warmup_tts()) asr_task = asyncio.create_task(_warmup_asr()) done, pending = await asyncio.wait( [tts_task, asr_task], timeout=TTS_ASR_WARMUP_TIMEOUT, return_when=asyncio.ALL_COMPLETED, ) for task in pending: task.cancel() try: await task except asyncio.CancelledError: pass tts_result = tts_task.result() if tts_task in done else False asr_result = asr_task.result() if asr_task in done else False logger.info("[Warmup] 完成: TTS=%s, ASR=%s", tts_result, asr_result) return tts_result, asr_result except Exception as e: logger.error("[Warmup] 异常: %s", str(e)) traceback.print_exc() return False, False def _check_and_unload_idle_models(): # pragma: no cover """ 检查并卸载空闲超过阈值的模型 """ if TTS_ASR_IDLE_TIMEOUT <= 0: return global _tts_pipeline, _asr_pipeline current_time = time.time() if _tts_pipeline is not None: idle_seconds = current_time - _tts_last_used if idle_seconds > TTS_ASR_IDLE_TIMEOUT: logger.info("[TTS] 空闲 %.0f 秒，卸载模型", idle_seconds) _tts_pipeline = None _clear_cuda_cache() _clear_mps_cache() if _asr_pipeline is not None: idle_seconds = current_time - _asr_last_used if idle_seconds > TTS_ASR_IDLE_TIMEOUT: logger.info("[ASR] 空闲 %.0f 秒，卸载模型", idle_seconds) _asr_pipeline = None _clear_cuda_cache() _clear_mps_cache() def _validate_audio_data(audio_data: bytes) -> bool: # pragma: no cover """ 验证音频数据的有效性 """ if not audio_data or len(audio_data) < 44: # WAV header minimum return False return True def _resample_audio_robust(audio_array: np.ndarray, orig_sr: int, target_sr: int = 16000) -> np.ndarray: """ 健壮的音频重采样，支持多个回退方案 """ if orig_sr == target_sr: return audio_array # 尝试librosa try: import librosa return librosa.resample(audio_array, orig_sr=orig_sr, target_sr=target_sr) except Exception as e: logger.warning("[Audio] librosa.resample失败: %s，尝试torchaudio", str(e)) # 尝试torchaudio try: import torch import torchaudio.transforms as T resampler = T.Resample(orig_sr, target_sr) audio_tensor = torch.from_numpy(audio_array).unsqueeze(0).float() resampled = resampler(audio_tensor) return resampled.squeeze(0).numpy() except Exception as e: logger.warning("[Audio] torchaudio重采样失败: %s，使用线性插值", str(e)) # 最后的回退：简单的线性插值 try: ratio = target_sr / orig_sr new_length = int(len(audio_array) * ratio) indices = np.linspace(0, len(audio_array) - 1, new_length) return np.interp(indices, np.arange(len(audio_array)), audio_array) except Exception as e: logger.error("[Audio] 所有重采样方法都失败: %s", str(e)) raise RuntimeError(f"音频重采样失败: {str(e)}") def _save_audio_to_wav(audio_data: bytes, sample_rate: int = 16000) -> str: import tempfile import wave with tempfile.NamedTemporaryFile(suffix=".wav", delete=False, mode="wb") as tmp: with wave.open(tmp.name, "wb") as wf: wf.setnchannels(1) wf.setsampwidth(2) wf.setframerate(sample_rate) wf.writeframes(audio_data) return tmp.name async def _tts_sync_with_retry(text: str, voice: str = "af_bella", rate: float = 1.0, max_retries: int = 2) -> tuple[bytes, int]: # pragma: no cover """ TTS推理，支持重试和降级 """ global _tts_last_used _check_and_unload_idle_models() if not await _load_tts_pipeline_with_retry(): raise RuntimeError("TTS 模型加载失败") tts = _tts_pipeline sample_rate = 24000 for attempt in range(max_retries): try: def inference(): result = tts(text) audio = None sr = sample_rate if isinstance(result, dict): audio = result.get("audio") sr = int(result.get("sampling_rate", sr)) elif isinstance(result, (list, tuple)) and result: audio = result[0] if audio is None: raise RuntimeError("TTS模型未返回音频数据") if hasattr(audio, "cpu"): audio = audio.cpu().numpy() if hasattr(audio, "numpy"): audio = audio.numpy() return audio, sr audio, sample_rate = await asyncio.to_thread(inference) duration_ms = int(len(audio) * 1000 / sample_rate) if audio.dtype != np.int16: audio = (audio * 32767).astype(np.int16) import tempfile import wave with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp: output_path = tmp.name try: with wave.open(output_path, "wb") as wf: wf.setnchannels(1) wf.setsampwidth(2) wf.setframerate(sample_rate) wf.writeframes(audio.tobytes()) with open(output_path, "rb") as f: audio_bytes = f.read() _tts_last_used = time.time() return audio_bytes, duration_ms finally: if os.path.exists(output_path): os.unlink(output_path) except RuntimeError as e: error_str = str(e) if "MPS" in error_str or "mps" in error_str: logger.warning("[TTS] MPS 推理错误，尝试降级重试 (尝试 %d/%d): %s", attempt + 1, max_retries, error_str) caps = _detect_device_capabilities() caps.mps_available = False caps.device = "cpu" _clear_mps_cache() if attempt < max_retries - 1: continue elif "CUDA" in error_str or "cuda" in error_str: logger.warning("[TTS] CUDA 推理错误，尝试降级重试 (尝试 %d/%d): %s", attempt + 1, max_retries, error_str) caps = _detect_device_capabilities() caps.cuda_available = False caps.device = "cpu" _clear_cuda_cache() if attempt < max_retries - 1: continue raise except Exception as e: logger.error("[TTS] 推理失败: %s", str(e)) if attempt == max_retries - 1: raise await asyncio.sleep(0.5) raise RuntimeError("TTS 推理失败") async def _asr_sync_with_retry(audio_data: bytes, language: str = "zh", max_retries: int = 2) -> str: # pragma: no cover """ ASR推理，支持重试和降级 """ global _asr_last_used _check_and_unload_idle_models() # 验证音频数据 if not _validate_audio_data(audio_data): raise ValueError("无效的音频数据") if not await _load_asr_pipeline_with_retry(): raise RuntimeError("ASR 模型加载失败") audio_path = _save_audio_to_wav(audio_data) try: import soundfile as sf # 健壮的音频读取 try: audio_array, sample_rate = await asyncio.to_thread(lambda: sf.read(audio_path)) except Exception as e: logger.error("[ASR] 音频读取失败: %s", str(e)) raise RuntimeError(f"音频读取失败: {str(e)}") # 转换为单声道 if len(audio_array.shape) > 1: audio_array = np.mean(audio_array, axis=1) # 重采样到16kHz（使用健壮的方法） if sample_rate != 16000: try: audio_array = _resample_audio_robust(audio_array, sample_rate, 16000) sample_rate = 16000 except Exception as e: logger.error("[ASR] 重采样失败: %s", str(e)) raise RuntimeError(f"音频重采样失败: {str(e)}") audio_array = audio_array.astype(np.float32) for attempt in range(max_retries): try: def inference(): asr = _asr_pipeline result = asr( audio_array, sampling_rate=sample_rate, generate_kwargs={"language": language, "task": "transcribe"}, ) if isinstance(result, dict): return result.get("text", "").strip() return str(result).strip() text = await asyncio.to_thread(inference) _asr_last_used = time.time() return text except RuntimeError as e: error_str = str(e) if "MPS" in error_str or "mps" in error_str: logger.warning("[ASR] MPS 推理错误，尝试降级重试 (尝试 %d/%d): %s", attempt + 1, max_retries, error_str) caps = _detect_device_capabilities() caps.mps_available = False caps.device = "cpu" _clear_mps_cache() if attempt < max_retries - 1: continue elif "CUDA" in error_str or "cuda" in error_str: logger.warning("[ASR] CUDA 推理错误，尝试降级重试 (尝试 %d/%d): %s", attempt + 1, max_retries, error_str) caps = _detect_device_capabilities() caps.cuda_available = False caps.device = "cpu" _clear_cuda_cache() if attempt < max_retries - 1: continue raise except Exception as e: logger.error("[ASR] 推理失败: %s", str(e)) if attempt == max_retries - 1: raise await asyncio.sleep(0.5) raise RuntimeError("ASR 推理失败") finally: if os.path.exists(audio_path): os.unlink(audio_path) # Legacy sync wrappers (for compatibility) def _tts_sync(text: str, voice: str = "af_bella", rate: float = 1.0) -> tuple[bytes, int]: # pragma: no cover raise RuntimeError("请使用 _tts_sync_with_retry()") def _asr_sync(audio_data: bytes, language: str = "zh") -> str: # pragma: no cover raise RuntimeError("请使用 _asr_sync_with_retry()") async def _text_to_speech(text: str, voice: str = "af_bella", rate: float = 1.0) -> tuple[bytes, int]: # pragma: no cover return await _tts_sync_with_retry(text, voice, rate) async def _speech_to_text(audio_data: bytes, language: str = "zh") -> str: # pragma: no cover return await _asr_sync_with_retry(audio_data, language) # Request/Response models class TTSRequest(BaseModel): text: str voice: str = "af_bella" rate: float = 1.0 format: str = "wav" class TTSResponse(BaseModel): audio_base64: str format: str duration_ms: int class ASRRequest(BaseModel): audio_base64: str language: str = "zh-CN" class ASRResponse(BaseModel): text: str language: str class ModelStatus(BaseModel): tts_loaded: bool asr_loaded: bool asr_model_size: Optional[str] = None device: str device_capabilities: Optional[Dict[str, Any]] = None tts_last_used: Optional[float] = None asr_last_used: Optional[float] = None offline_mode: bool = False quantize_enabled: bool = False def get_api_key(api_key: str): import main API_KEY = main.API_KEY if api_key != API_KEY: raise HTTPException(status_code=403, detail="API Key 无效") return api_key @router.get("/config") async def get_config(api_key: str = Security(get_api_key)): """ 获取当前TTS/ASR配置信息 """ caps = _detect_device_capabilities() return { "environment": { "TTS_ASR_DEVICE": TTS_ASR_DEVICE, "TTS_ASR_MODEL_SIZE": TTS_ASR_MODEL_SIZE, "TTS_ASR_QUANTIZE": TTS_ASR_QUANTIZE, "TTS_ASR_OFFLINE_MODE": TTS_ASR_OFFLINE_MODE, "TTS_ASR_WARMUP": TTS_ASR_WARMUP, "TTS_ASR_WARMUP_TIMEOUT": TTS_ASR_WARMUP_TIMEOUT, "TTS_ASR_IDLE_TIMEOUT": TTS_ASR_IDLE_TIMEOUT, "TTS_ASR_MPS_MEMORY_LIMIT_MB": TTS_ASR_MPS_MEMORY_LIMIT_MB, }, "device": { "current": _get_device(), "mps_available": caps.mps_available, "cuda_available": caps.cuda_available, "is_apple_silicon": _is_apple_silicon(), "mps_memory_limit_mb": caps.mps_memory_limit_mb, "cuda_memory_limit_mb": caps.cuda_memory_limit_mb, }, "model": { "tts": "suno/bark", "asr_current_size": _asr_model_size, "asr_recommended_size": caps.recommended_model_size, "available_sizes": list(WHISPER_MODEL_SIZES.keys()), }, "status": { "tts_loaded": _tts_pipeline is not None, "asr_loaded": _asr_pipeline is not None, } } @router.get("/status", response_model=ModelStatus) async def get_status(api_key: str = Security(get_api_key)): """ 获取模型状态 """ current_time = time.time() caps = _detect_device_capabilities() return ModelStatus( tts_loaded=_tts_pipeline is not None, asr_loaded=_asr_pipeline is not None, asr_model_size=_asr_model_size, device=_get_device(), device_capabilities={ "mps_available": caps.mps_available, "cuda_available": caps.cuda_available, "mps_memory_limit_mb": caps.mps_memory_limit_mb, "cuda_memory_limit_mb": caps.cuda_memory_limit_mb, "recommended_model_size": caps.recommended_model_size, "is_apple_silicon": _is_apple_silicon(), }, tts_last_used=_tts_last_used if _tts_last_used > 0 else None, asr_last_used=_asr_last_used if _asr_last_used > 0 else None, offline_mode=TTS_ASR_OFFLINE_MODE, quantize_enabled=TTS_ASR_QUANTIZE, ) @router.post("/warmup") async def warmup_models(api_key: str = Security(get_api_key)): """ 手动触发模型预热 """ tts_result, asr_result = await _warmup_all() caps = _detect_device_capabilities() return { "tts_warmup": tts_result, "asr_warmup": asr_result, "device": _get_device(), "asr_model_size": _asr_model_size, "offline_mode": TTS_ASR_OFFLINE_MODE, "quantize_enabled": TTS_ASR_QUANTIZE, "is_apple_silicon": _is_apple_silicon(), "device_capabilities": { "mps_available": caps.mps_available, "cuda_available": caps.cuda_available, "mps_memory_limit_mb": caps.mps_memory_limit_mb, "cuda_memory_limit_mb": caps.cuda_memory_limit_mb, }, } @router.post("/tts", response_model=TTSResponse) # pragma: no cover async def text_to_speech(req: TTSRequest, api_key: str = Security(get_api_key)): request_id = str(hash(req.text))[:8] try: logger.info("[TTS][%s] text_chars=%d voice=%s format=%s", request_id, len(req.text), req.voice, req.format) audio_data, duration_ms = await _text_to_speech(req.text, req.voice, req.rate) if req.format.lower() == "mp3": import subprocess import tempfile with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp_in: tmp_in.write(audio_data) input_path = tmp_in.name with tempfile.NamedTemporaryFile(suffix=".mp3", delete=False) as tmp_out: output_path = tmp_out.name try: cmd = ["ffmpeg", "-i", input_path, "-acodec", "libmp3lame", "-ab", "128k", output_path] result = await asyncio.to_thread( lambda: subprocess.run(cmd, capture_output=True, text=True, timeout=30) ) if result.returncode != 0: raise RuntimeError(f"MP3 转换失败: {result.stderr}") with open(output_path, "rb") as f: audio_data = f.read() finally: for path in [input_path, output_path]: if os.path.exists(path): os.unlink(path) logger.info("[TTS][%s] success duration_ms=%d", request_id, duration_ms) return TTSResponse( audio_base64=base64.b64encode(audio_data).decode(), format=req.format, duration_ms=duration_ms, ) except Exception as e: logger.exception("[TTS] failed: %s", e) raise HTTPException(status_code=500, detail=str(e)) @router.post("/asr", response_model=ASRResponse) # pragma: no cover async def speech_to_text(req: ASRRequest, api_key: str = Security(get_api_key)): request_id = str(hash(req.audio_base64))[:8] try: logger.info("[ASR][%s] audio_base64_chars=%d language=%s", request_id, len(req.audio_base64), req.language) audio_data = base64.b64decode(req.audio_base64) text = await _speech_to_text(audio_data, req.language[:2]) logger.info("[ASR][%s] success text_chars=%d", request_id, len(text)) return ASRResponse(text=text, language=req.language) except Exception as e: logger.exception("[ASR] failed: %s", e) raise HTTPException(status_code=500, detail=str(e)) def register_tts_asr_routes(app): """ 注册TTS/ASR路由并可选执行预热 """ app.include_router(router, prefix="/v1/tts-asr") if TTS_ASR_WARMUP: @app.on_event("startup") async def warmup_on_startup(): logger.info("[Startup] 开始后台预热...") asyncio.create_task(_warmup_all())