ydy0615/llm-in-text
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(tts): add api endpoints and optimization for apple silicon

Browse Source 
Introduce a comprehensive TTS/ASR module that:
- Adds /v1/tts-asr/config, /status, /warmup, /tts, /asr endpoints with detailed JSON responses
- Implements Apple‑Silicon detection, device selection (MPS/CUDA/CPU), and memory limiting logic
- Supports selectable model size, quantization, and offline mode via environment variables
- Adds robust audio validation and multi‑path resampling fallback
- Provides new README sections for API usage, device detection, and performance benchmarking
- Includes a full testing suite: unit tests, integration tests, macOS simulation and performance reports
- Updates backend dependencies and CI scripts
- Adds new front‑end views and components for Univer editor integration

All changes are backward compatible; new features are exposed through environment variables and new API routes.
This commit is contained in:
ydy0615
2026-04-06 11:14:09 +08:00
parent c70cb2a9f0
commit 7985fe9641
27 changed files with 9304 additions and 260 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/tts_asr.py
+473 -116
View File
@@ -1,12 +1,16 @@
# 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 typing import Optional
from dataclasses import dataclass
from pathlib import Path
from typing import Optional, Dict, Any
from fastapi import APIRouter, HTTPException, Security
from pydantic import BaseModel
@@ -21,15 +25,48 @@ 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
_device = None
_device_tested = False
_asr_model_size: Optional[str] = None
_device_caps: Optional[DeviceCapabilities] = None
_tts_last_used = 0.0
_asr_last_used = 0.0
_tts_loading = False
@@ -38,83 +75,187 @@ _tts_lock = asyncio.Lock()
_asr_lock = asyncio.Lock()
def _test_device_capability(device_str: str) -> tuple[bool, str]:
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
if device_str == "cpu":
return True, ""
if device_str == "mps":
if not hasattr(torch.backends, "mps") or not torch.backends.mps.is_available():
return False, "MPS 不可用"
if not torch.backends.mps.is_built():
return False, "MPS 未编译"
test_tensor = torch.randn(2, 2, device="mps")
_ = test_tensor @ test_tensor
del test_tensor
torch.mps.empty_cache()
return True, ""
if device_str.startswith("cuda"):
if not torch.cuda.is_available():
return False, "CUDA 不可用"
torch.cuda.empty_cache()
return True, ""
return False, f"未知设备类型: {device_str}"
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:
return False, f"设备测试失败: {str(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:
"""
获取最佳计算设备，支持环境变量覆盖和降级策略
"""
global _device, _device_tested
if _device is not None and _device_tested:
return _device
import torch
device_preference = []
caps = _detect_device_capabilities()
# 环境变量强制指定
if TTS_ASR_DEVICE == "cpu":
_device = "cpu"
_device_tested = True
logger.info("[Device] 强制使用 CPU (环境变量)")
return _device
elif TTS_ASR_DEVICE in ("mps", "cuda", "auto"):
if TTS_ASR_DEVICE != "auto":
device_preference = [TTS_ASR_DEVICE, "cpu"]
return "cpu"
elif TTS_ASR_DEVICE == "mps":
if caps.mps_available:
logger.info("[Device] 强制使用 MPS (环境变量)")
return "mps"
else:
if platform.system() == "Darwin":
device_preference = ["mps", "cpu"]
else:
device_preference = ["cuda", "cpu"]
else:
device_preference = ["mps", "cuda", "cpu"]
for dev in device_preference:
ok, err = _test_device_capability(dev)
if ok:
_device = dev
_device_tested = True
logger.info("[Device] 使用 %s 加速", dev.upper() if dev != "cpu" else "CPU")
return _device
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] %s 不可用: %s", dev.upper() if dev != "cpu" else "CPU", err)
logger.warning("[Device] CUDA不可用，降级到CPU")
return "cpu"
# 自动选择
return caps.device
_device = "cpu"
_device_tested = True
logger.info("[Device] 降级使用 CPU")
return _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:
@@ -127,13 +268,47 @@ def _device_arg() -> str:
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
if _device and _device.startswith("cuda"):
caps = _detect_device_capabilities()
if caps.cuda_available:
torch.cuda.empty_cache()
except Exception:
pass
@@ -142,7 +317,8 @@ def _clear_cuda_cache():
def _clear_mps_cache():
try:
import torch
if _device == "mps":
caps = _detect_device_capabilities()
if caps.mps_available:
torch.mps.empty_cache()
except Exception:
pass
@@ -173,6 +349,13 @@ async def _load_tts_pipeline_with_retry(max_retries: int = 2) -> bool:
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)
@@ -180,7 +363,7 @@ async def _load_tts_pipeline_with_retry(max_retries: int = 2) -> bool:
_tts_pipeline = await asyncio.to_thread(
lambda: pipeline(
"text-to-speech",
model="hexgrad/Kokoro-82M",
model=model_id,
trust_remote_code=True,
device=device_to_use,
torch_dtype=torch_dtype,
@@ -194,13 +377,16 @@ async def _load_tts_pipeline_with_retry(max_retries: int = 2) -> bool:
error_str = str(e)
if "MPS" in error_str or "mps" in error_str:
logger.warning("[TTS] MPS 推理失败，尝试降级到 CPU: %s", error_str)
global _device
_device = "cpu"
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)
_device = "cpu"
caps = _detect_device_capabilities()
caps.cuda_available = False
caps.device = "cpu"
_clear_cuda_cache()
continue
else:
@@ -219,9 +405,9 @@ async def _load_tts_pipeline_with_retry(max_retries: int = 2) -> bool:
async def _load_asr_pipeline_with_retry(max_retries: int = 2) -> bool:
"""
加载ASR管道，支持重试和降级
加载ASR管道，支持重试、降级、模型大小选择和量化
"""
global _asr_pipeline, _asr_loading
global _asr_pipeline, _asr_loading, _asr_model_size
async with _asr_lock:
if _asr_pipeline is not None:
@@ -236,48 +422,87 @@ async def _load_asr_pipeline_with_retry(max_retries: int = 2) -> bool:
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 large-v3-turbo 模型 (尝试 %d/%d, 设备: %s)...",
attempt + 1, max_retries, device_to_use)
model_id = "openai/whisper-large-v3-turbo"
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,
torch_dtype=torch_dtype,
low_cpu_mem_usage=True,
use_safetensors=True,
**load_kwargs
)
processor = AutoProcessor.from_pretrained(model_id)
return pipeline(
"automatic-speech-recognition",
model=model,
tokenizer=processor.tokenizer,
feature_extractor=processor.feature_extractor,
torch_dtype=torch_dtype,
device=device_to_use,
)
# 如果使用了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 large-v3-turbo 模型加载完成")
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)
global _device
_device = "cpu"
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)
_device = "cpu"
caps = _detect_device_capabilities()
caps.cuda_available = False
caps.device = "cpu"
_clear_cuda_cache()
continue
else:
@@ -458,6 +683,52 @@ def _check_and_unload_idle_models():
_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
@@ -521,34 +792,37 @@ async def _tts_sync_with_retry(text: str, voice: str = "af_bella", rate: float =
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)
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)
global _device
_device = "cpu"
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)
_device = "cpu"
caps = _detect_device_capabilities()
caps.cuda_available = False
caps.device = "cpu"
_clear_cuda_cache()
if attempt < max_retries - 1:
continue
@@ -570,6 +844,10 @@ async def _asr_sync_with_retry(audio_data: bytes, language: str = "zh", max_retr
_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 模型加载失败")
@@ -578,17 +856,25 @@ async def _asr_sync_with_retry(audio_data: bytes, language: str = "zh", max_retr
try:
import soundfile as sf
audio_array, sample_rate = await asyncio.to_thread(lambda: sf.read(audio_path))
# 健壮的音频读取
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:
import librosa
audio_array = await asyncio.to_thread(
lambda: librosa.resample(audio_array, orig_sr=sample_rate, target_sr=16000)
)
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)
@@ -614,15 +900,18 @@ async def _asr_sync_with_retry(audio_data: bytes, language: str = "zh", max_retr
if "MPS" in error_str or "mps" in error_str:
logger.warning("[ASR] MPS 推理错误，尝试降级重试 (尝试 %d/%d): %s",
attempt + 1, max_retries, error_str)
global _device
_device = "cpu"
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)
_device = "cpu"
caps = _detect_device_capabilities()
caps.cuda_available = False
caps.device = "cpu"
_clear_cuda_cache()
if attempt < max_retries - 1:
continue
@@ -684,9 +973,13 @@ class ASRResponse(BaseModel):
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):
@@ -697,18 +990,70 @@ def get_api_key(api_key: str):
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,
)
@@ -718,10 +1063,22 @@ 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,
},
}