llm-in-text/backend/tts_asr.py

# 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:
        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:
        from transformers import file_utils
        cache_dir = file_utils.default_cache_path
        
        # 简单的缓存检查：查找模型目录
        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:
    """
    加载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 = "hexgrad/Kokoro-82M"
                    
                    # 离线模式检查
                    if TTS_ASR_OFFLINE_MODE and not _check_model_cached(model_id):
                        logger.error("[TTS] 离线模式下模型 %s 未缓存", model_id)
                        return False

                    logger.info("[TTS] 加载 Kokoro-82M 模型 (尝试 %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] Kokoro-82M 模型加载完成")
                    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:
    """
    加载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():
    """同步获取TTS管道（已弃用，保留兼容性）"""
    if _tts_pipeline is not None:
        return _tts_pipeline
    raise RuntimeError("TTS 管道未加载，请使用 _load_tts_pipeline_with_retry()")


def _get_asr_pipeline():
    """同步获取ASR管道（已弃用，保留兼容性）"""
    if _asr_pipeline is not None:
        return _asr_pipeline
    raise RuntimeError("ASR 管道未加载，请使用 _load_asr_pipeline_with_retry()")


async def _warmup_tts() -> bool:
    """
    预热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, voice="af_bella")
                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:
    """
    预热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]:
    """
    预热所有模型
    返回: (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():
    """
    检查并卸载空闲超过阈值的模型
    """
    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:
    """
    验证音频数据的有效性
    """
    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]:
    """
    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, voice=voice)
                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("Kokoro 未返回音频数据")

                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:
    """
    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]:
    raise RuntimeError("请使用 _tts_sync_with_retry()")


def _asr_sync(audio_data: bytes, language: str = "zh") -> str:
    raise RuntimeError("请使用 _asr_sync_with_retry()")


async def _text_to_speech(text: str, voice: str = "af_bella", rate: float = 1.0) -> tuple[bytes, int]:
    return await _tts_sync_with_retry(text, voice, rate)


async def _speech_to_text(audio_data: bytes, language: str = "zh") -> str:
    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": "hexgrad/Kokoro-82M",
            "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)
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)
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())