feat: LLM 应用网页开发及内联建议功能实现
This commit is contained in:
+641
-100
@@ -4,6 +4,9 @@ import base64
|
||||
import logging
|
||||
import os
|
||||
import platform
|
||||
import time
|
||||
import traceback
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import APIRouter, HTTPException, Security
|
||||
from pydantic import BaseModel
|
||||
@@ -12,84 +15,447 @@ 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"))
|
||||
|
||||
# Warmup constants
|
||||
TTS_WARMUP_TEXT = "你好，这是一个测试。"
|
||||
ASR_WARMUP_AUDIO_SECONDS = 0.5
|
||||
|
||||
# Global state
|
||||
_tts_pipeline = None
|
||||
_asr_pipeline = None
|
||||
_device = None
|
||||
_device_tested = False
|
||||
_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 _get_device():
|
||||
global _device
|
||||
if _device is not None:
|
||||
def _test_device_capability(device_str: str) -> tuple[bool, str]:
|
||||
"""
|
||||
测试设备实际可用性
|
||||
返回: (是否可用, 错误信息)
|
||||
"""
|
||||
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}"
|
||||
except Exception as e:
|
||||
return False, f"设备测试失败: {str(e)}"
|
||||
|
||||
|
||||
def _get_device() -> str:
|
||||
"""
|
||||
获取最佳计算设备，支持环境变量覆盖和降级策略
|
||||
"""
|
||||
global _device, _device_tested
|
||||
|
||||
if _device is not None and _device_tested:
|
||||
return _device
|
||||
|
||||
import torch
|
||||
|
||||
if platform.system() == "Darwin" and hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
||||
_device = "mps"
|
||||
logger.info("[Device] 使用 MPS 加速")
|
||||
elif torch.cuda.is_available():
|
||||
_device = "cuda"
|
||||
logger.info("[Device] 使用 CUDA 加速")
|
||||
else:
|
||||
device_preference = []
|
||||
|
||||
if TTS_ASR_DEVICE == "cpu":
|
||||
_device = "cpu"
|
||||
logger.info("[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"]
|
||||
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
|
||||
else:
|
||||
logger.warning("[Device] %s 不可用: %s", dev.upper() if dev != "cpu" else "CPU", err)
|
||||
|
||||
_device = "cpu"
|
||||
_device_tested = True
|
||||
logger.info("[Device] 降级使用 CPU")
|
||||
return _device
|
||||
|
||||
|
||||
def _device_arg():
|
||||
def _device_arg() -> str:
|
||||
device = _get_device()
|
||||
if device == "cuda":
|
||||
return "cuda:0"
|
||||
return device
|
||||
|
||||
|
||||
def _get_torch_dtype():
|
||||
device = _get_device()
|
||||
import torch
|
||||
return torch.float16 if device != "cpu" else torch.float32
|
||||
|
||||
|
||||
def _clear_cuda_cache():
|
||||
try:
|
||||
import torch
|
||||
if _device and _device.startswith("cuda"):
|
||||
torch.cuda.empty_cache()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _clear_mps_cache():
|
||||
try:
|
||||
import torch
|
||||
if _device == "mps":
|
||||
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()
|
||||
|
||||
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="hexgrad/Kokoro-82M",
|
||||
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)
|
||||
global _device
|
||||
_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"
|
||||
_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
|
||||
|
||||
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
|
||||
|
||||
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"
|
||||
|
||||
def load_model():
|
||||
model = AutoModelForSpeechSeq2Seq.from_pretrained(
|
||||
model_id,
|
||||
torch_dtype=torch_dtype,
|
||||
low_cpu_mem_usage=True,
|
||||
use_safetensors=True,
|
||||
)
|
||||
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,
|
||||
)
|
||||
|
||||
_asr_pipeline = await asyncio.to_thread(load_model)
|
||||
logger.info("[ASR] Whisper large-v3-turbo 模型加载完成")
|
||||
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"
|
||||
_clear_mps_cache()
|
||||
continue
|
||||
elif "CUDA" in error_str or "cuda" in error_str:
|
||||
logger.warning("[ASR] CUDA 推理失败，尝试降级到 CPU: %s", error_str)
|
||||
_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():
|
||||
global _tts_pipeline
|
||||
"""同步获取TTS管道（已弃用，保留兼容性）"""
|
||||
if _tts_pipeline is not None:
|
||||
return _tts_pipeline
|
||||
|
||||
import torch
|
||||
from transformers import pipeline
|
||||
|
||||
logger.info("[TTS] 加载 Kokoro-82M 模型...")
|
||||
_tts_pipeline = pipeline(
|
||||
"text-to-speech",
|
||||
model="hexgrad/Kokoro-82M",
|
||||
trust_remote_code=True,
|
||||
device=_device_arg(),
|
||||
torch_dtype=torch.float16 if _get_device() != "cpu" else torch.float32,
|
||||
)
|
||||
logger.info("[TTS] Kokoro-82M 模型加载完成")
|
||||
return _tts_pipeline
|
||||
raise RuntimeError("TTS 管道未加载，请使用 _load_tts_pipeline_with_retry()")
|
||||
|
||||
|
||||
def _get_asr_pipeline():
|
||||
global _asr_pipeline
|
||||
"""同步获取ASR管道（已弃用，保留兼容性）"""
|
||||
if _asr_pipeline is not None:
|
||||
return _asr_pipeline
|
||||
raise RuntimeError("ASR 管道未加载，请使用 _load_asr_pipeline_with_retry()")
|
||||
|
||||
import torch
|
||||
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
|
||||
|
||||
logger.info("[ASR] 加载 Whisper large-v3-turbo 模型...")
|
||||
model_id = "openai/whisper-large-v3-turbo"
|
||||
model = AutoModelForSpeechSeq2Seq.from_pretrained(
|
||||
model_id,
|
||||
torch_dtype=torch.float16 if _get_device() != "cpu" else torch.float32,
|
||||
low_cpu_mem_usage=True,
|
||||
use_safetensors=True,
|
||||
)
|
||||
processor = AutoProcessor.from_pretrained(model_id)
|
||||
_asr_pipeline = pipeline(
|
||||
"automatic-speech-recognition",
|
||||
model=model,
|
||||
tokenizer=processor.tokenizer,
|
||||
feature_extractor=processor.feature_extractor,
|
||||
torch_dtype=torch.float16 if _get_device() != "cpu" else torch.float32,
|
||||
device=_device_arg(),
|
||||
)
|
||||
logger.info("[ASR] Whisper large-v3-turbo 模型加载完成")
|
||||
return _asr_pipeline
|
||||
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 _save_audio_to_wav(audio_data: bytes, sample_rate: int = 16000) -> str:
|
||||
@@ -105,74 +471,193 @@ def _save_audio_to_wav(audio_data: bytes, sample_rate: int = 16000) -> str:
|
||||
return tmp.name
|
||||
|
||||
|
||||
def _tts_sync(text: str, voice: str = "af_bella", rate: float = 1.0) -> tuple[bytes, int]:
|
||||
tts = _get_tts_pipeline()
|
||||
result = tts(text, voice=voice)
|
||||
audio = None
|
||||
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
|
||||
if isinstance(result, dict):
|
||||
audio = result.get("audio")
|
||||
sample_rate = int(result.get("sampling_rate", sample_rate))
|
||||
elif isinstance(result, (list, tuple)) and result:
|
||||
audio = result[0]
|
||||
|
||||
if audio is None:
|
||||
raise RuntimeError("Kokoro 未返回音频数据")
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
def inference():
|
||||
result = tts(text, voice=voice)
|
||||
audio = None
|
||||
sr = sample_rate
|
||||
|
||||
if hasattr(audio, "cpu"):
|
||||
audio = audio.cpu().numpy()
|
||||
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]
|
||||
|
||||
duration_ms = int(len(audio) * 1000 / sample_rate)
|
||||
if audio is None:
|
||||
raise RuntimeError("Kokoro 未返回音频数据")
|
||||
|
||||
if audio.dtype != np.int16:
|
||||
audio = (audio * 32767).astype(np.int16)
|
||||
if hasattr(audio, "cpu"):
|
||||
audio = audio.cpu().numpy()
|
||||
|
||||
import tempfile
|
||||
import wave
|
||||
if hasattr(audio, "numpy"):
|
||||
audio = audio.numpy()
|
||||
|
||||
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:
|
||||
return f.read(), duration_ms
|
||||
finally:
|
||||
if os.path.exists(output_path):
|
||||
os.unlink(output_path)
|
||||
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)
|
||||
global _device
|
||||
_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"
|
||||
_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 _text_to_speech(text: str, voice: str = "af_bella", rate: float = 1.0) -> tuple[bytes, int]:
|
||||
return await asyncio.to_thread(_tts_sync, text, voice, rate)
|
||||
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()
|
||||
|
||||
def _asr_sync(audio_data: bytes, language: str = "zh") -> str:
|
||||
import soundfile as sf
|
||||
if not await _load_asr_pipeline_with_retry():
|
||||
raise RuntimeError("ASR 模型加载失败")
|
||||
|
||||
asr = _get_asr_pipeline()
|
||||
audio_path = _save_audio_to_wav(audio_data)
|
||||
|
||||
try:
|
||||
audio_array, sample_rate = sf.read(audio_path)
|
||||
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()
|
||||
import soundfile as sf
|
||||
|
||||
audio_array, sample_rate = await asyncio.to_thread(lambda: sf.read(audio_path))
|
||||
|
||||
if len(audio_array.shape) > 1:
|
||||
audio_array = np.mean(audio_array, axis=1)
|
||||
|
||||
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
|
||||
|
||||
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)
|
||||
global _device
|
||||
_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"
|
||||
_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 asyncio.to_thread(_asr_sync, audio_data, language)
|
||||
return await _asr_sync_with_retry(audio_data, language)
|
||||
|
||||
|
||||
# Request/Response models
|
||||
class TTSRequest(BaseModel):
|
||||
text: str
|
||||
voice: str = "af_bella"
|
||||
@@ -196,21 +681,58 @@ class ASRResponse(BaseModel):
|
||||
language: str
|
||||
|
||||
|
||||
class ModelStatus(BaseModel):
|
||||
tts_loaded: bool
|
||||
asr_loaded: bool
|
||||
device: str
|
||||
tts_last_used: Optional[float] = None
|
||||
asr_last_used: Optional[float] = None
|
||||
|
||||
|
||||
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("/status", response_model=ModelStatus)
|
||||
async def get_status(api_key: str = Security(get_api_key)):
|
||||
"""
|
||||
获取模型状态
|
||||
"""
|
||||
current_time = time.time()
|
||||
return ModelStatus(
|
||||
tts_loaded=_tts_pipeline is not None,
|
||||
asr_loaded=_asr_pipeline is not None,
|
||||
device=_get_device(),
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
@router.post("/warmup")
|
||||
async def warmup_models(api_key: str = Security(get_api_key)):
|
||||
"""
|
||||
手动触发模型预热
|
||||
"""
|
||||
tts_result, asr_result = await _warmup_all()
|
||||
return {
|
||||
"tts_warmup": tts_result,
|
||||
"asr_warmup": asr_result,
|
||||
"device": _get_device(),
|
||||
}
|
||||
|
||||
|
||||
@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)
|
||||
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
|
||||
@@ -222,7 +744,9 @@ async def text_to_speech(req: TTSRequest, api_key: str = Security(get_api_key)):
|
||||
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))
|
||||
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:
|
||||
@@ -231,8 +755,14 @@ async def text_to_speech(req: TTSRequest, api_key: str = Security(get_api_key)):
|
||||
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)
|
||||
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))
|
||||
@@ -242,15 +772,26 @@ async def text_to_speech(req: TTSRequest, api_key: str = Security(get_api_key)):
|
||||
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)
|
||||
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())
|
||||
|
||||
Reference in New Issue
Block a user