refactor: replace Kokoro-82M with suno/bark for TTS, update HF cache path, and add model warmup on startup.

2026-04-06 13:40:41 +08:00
parent 7985fe9641
commit caf1ac1c01
8 changed files with 143 additions and 18 deletions
@@ -1,4 +1,4 @@
-import asyncio
+import asyncio
 import base64
 import logging
 import os
@@ -37,6 +37,15 @@ def _get_markitdown():

 app = FastAPI()

+@app.on_event("startup")
+async def startup_event():
+    logger.info("Starting blocking preload for TTS and ASR models...")
+    try:
+        from tts_asr import _warmup_all
+        await _warmup_all()
+    except Exception as e:
+        logger.warning(f"Failed to initiate model warmup: {e}")
+
 ACTIVE_COMPLETIONS: dict[str, asyncio.Task] = {}
 ACTIVE_COMPLETIONS_LOCK = asyncio.Lock()

@@ -22,7 +22,7 @@ logging.basicConfig(
 logger = logging.getLogger("api_benchmarker")

 # Constants
-DEFAULT_BASE_URL = "https://api.imageteach.tech:8002"
+DEFAULT_BASE_URL = "http://localhost:8001"
 DEFAULT_API_KEY = "your-secret-key-here"
 CHARS_PER_TOKEN = 4

@@ -31,7 +31,8 @@ def get_dummy_base64_image():
    return "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII="

 def get_dummy_base64_audio():
-    return "UklGRigAAABXQVZFZm10IBAAAAABAAEARKwAAIhYAQACABAAZGF0YQQAAAAAAA=="
+    # A bit longer dummy audio to pass validation (44 bytes header + some data)
+    return "UklGRigAAABXQVZFZm10IBAAAAABAAEARKwAAIhYAQACABAAZGF0YQQAAAAAAA" + "A" * 100 + "=="

 def generate_context_text(tokens: int) -> str:
    """Generate synthetic text of approximately 'tokens' tokens."""
@@ -191,7 +192,7 @@ class ApiBenchmarker:
                elif task_type == "tts":
                    metric = await self._execute_request(client, name, "POST", "/v1/tts-asr/tts", json={
                        "text": "This is a performance benchmark for the text to speech engine.",
-                        "voice": "af_bella",
+                        "voice": "v2/en_speaker_6",
                        "format": "wav"
                    })
                elif task_type == "asr":
@@ -0,0 +1,37 @@
+import asyncio
+import base64
+import os
+import sys
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+from backend.tts_asr import _tts_sync_with_retry
+
+async def play_audio():
+    print("生成测试音频中，请稍候...")
+    test_text = "这是一段用以测试新语音模型音质的中文合成音频。"
+    
+    try:
+        audio_bytes, sr = await _tts_sync_with_retry(test_text, rate=1.0)
+        
+        # 保存到本地文件
+        wav_path = os.path.join(os.path.dirname(__file__), "test_audio.wav")
+        with open(wav_path, "wb") as f:
+            f.write(audio_bytes)
+            
+        print(f"音频已生成并保存到: {wav_path}")
+        print("正在尝试在 macOS 上播放...")
+        
+        # Mac OS 的播放命令
+        os.system(f"afplay '{wav_path}'")
+        
+        print("播放完成。")
+            
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        print(f"音频生成失败: {str(e)}")
+
+if __name__ == "__main__":
+    asyncio.run(play_audio())
@@ -16,8 +16,9 @@ from unittest.mock import Mock, MagicMock, patch
 import tempfile
 import numpy as np

-# 确保可以导入tts_asr模块
-sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
+# 确保可以导入backend和tts_asr模块
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))


 class TestAppleSiliconDetection(unittest.TestCase):
@@ -174,7 +175,7 @@ class TestModelSizeSelection(unittest.TestCase):
            import backend.tts_asr as tts_asr_module
            importlib.reload(tts_asr_module)
            
-            from backend.tts_asr import _get_recommended_model_size
+            from backend.tts_asr import _get_recommended_model_size, WHISPER_MODEL_SIZES
            # 应该回退到推荐大小而不崩溃
            size = _get_recommended_model_size()
            self.assertIn(size, WHISPER_MODEL_SIZES.keys())
@@ -299,8 +300,8 @@ class TestModelCacheCheck(unittest.TestCase):
        """测试离线模式下缺失模型的处理"""
        from backend.tts_asr import _check_model_cached
        
-        # 模拟transformers缓存路径
-        with patch('transformers.file_utils.default_cache_path', '/nonexistent/path'):
+        # 模拟缓存路径
+        with patch('huggingface_hub.constants.HF_HUB_CACHE', '/nonexistent/path'):
            result = _check_model_cached('nonexistent/model')
            # 应该返回False（模型未缓存）
            self.assertFalse(result)
@@ -0,0 +1,56 @@
+import asyncio
+import base64
+import os
+import sys
+
+# 确保能找到backend模块
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
+
+from backend.tts_asr import _tts_sync_with_retry, _load_asr_pipeline_with_retry, _get_asr_pipeline
+
+async def verify_tts_asr_cross():
+    print("====================================")
+    print(" 开始严格交叉验证: TTS 生成 -> ASR 解析")
+    print("====================================")
+    
+    test_text = "苹果设备支持离线大模型运算"
+    print(f"\n[1] 正在调用 TTS 引擎 (suno/bark-small)...")
+    print(f"目标文本: '{test_text}'")
+    
+    try:
+        # TTS生成
+        audio_bytes, sr = await _tts_sync_with_retry(test_text, rate=1.0)
+        print(f"-> TTS 成功生成音频数据，大小: {len(audio_bytes)} Bytes, 采样率: {sr}Hz")
+    except Exception as e:
+        print(f"-> TTS 失败: {str(e)}")
+        sys.exit(1)
+        
+    print("\n[2] 正在调用 ASR 引擎 (Whisper)...")
+    try:
+        loaded = await _load_asr_pipeline_with_retry()
+        if not loaded:
+            print("-> ASR 核心加载失败！")
+            sys.exit(1)
+            
+        print("-> ASR 加载成功，开始解析音频...")
+        
+        # 将生成的wav bytes传递给ASR进行语音识别
+        asr_pipeline = _get_asr_pipeline()
+        result = asr_pipeline(audio_bytes, generate_kwargs={"task": "transcribe"})
+        recognized_text = result.get('text', '')
+        print(f"-> ASR 识别结果: '{recognized_text.strip()}'")
+        
+        if len(recognized_text.strip()) > 0:
+            print("\n结论: ✅ 验证成功！TTS和ASR模块功能链路闭环完成。")
+        else:
+            print("\n结论: ❌ ASR输出为空字符，闭环失败。")
+            sys.exit(1)
+            
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        print(f"-> ASR 分析阶段失败: {str(e)}")
+        sys.exit(1)
+
+if __name__ == "__main__":
+    asyncio.run(verify_tts_asr_cross())
@@ -283,8 +283,9 @@ def _check_model_cached(model_id: str) -> bool:
        return True  # 非离线模式，不检查缓存
    
    try:
-        from transformers import file_utils
-        cache_dir = file_utils.default_cache_path
+        import os
+        from huggingface_hub.constants import HF_HUB_CACHE
+        cache_dir = HF_HUB_CACHE
        
        # 简单的缓存检查：查找模型目录
        model_name = model_id.replace("/", "--")
@@ -350,14 +351,14 @@ async def _load_tts_pipeline_with_retry(max_retries: int = 2) -> bool:
                    device_to_use = _device_arg()
                    torch_dtype = _get_torch_dtype()
                    
-                    model_id = "hexgrad/Kokoro-82M"
+                    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] 加载 Kokoro-82M 模型 (尝试 %d/%d, 设备: %s)...", 
+                    logger.info("[TTS] 加载 suno/bark 模型 (尝试 %d/%d, 设备: %s)...", 
                               attempt + 1, max_retries, device_to_use)

                    _tts_pipeline = await asyncio.to_thread(
@@ -370,7 +371,7 @@ async def _load_tts_pipeline_with_retry(max_retries: int = 2) -> bool:
                        )
                    )

-                    logger.info("[TTS] Kokoro-82M 模型加载完成")
+                    logger.info("[TTS] suno/bark 模型加载完成")
                    return True

                except RuntimeError as e:
@@ -552,7 +553,7 @@ async def _warmup_tts() -> bool:

        def warmup_inference():
            try:
-                result = tts(TTS_WARMUP_TEXT, voice="af_bella")
+                result = tts(TTS_WARMUP_TEXT)
                if isinstance(result, dict):
                    audio = result.get("audio")
                    if hasattr(audio, "cpu"):
@@ -759,7 +760,7 @@ async def _tts_sync_with_retry(text: str, voice: str = "af_bella", rate: float =
    for attempt in range(max_retries):
        try:
            def inference():
-                result = tts(text, voice=voice)
+                result = tts(text)
                audio = None
                sr = sample_rate

@@ -770,7 +771,7 @@ async def _tts_sync_with_retry(text: str, voice: str = "af_bella", rate: float =
                    audio = result[0]

                if audio is None:
-                    raise RuntimeError("Kokoro 未返回音频数据")
+                    raise RuntimeError("TTS模型未返回音频数据")

                if hasattr(audio, "cpu"):
                    audio = audio.cpu().numpy()
@@ -1017,7 +1018,7 @@ async def get_config(api_key: str = Security(get_api_key)):
            "cuda_memory_limit_mb": caps.cuda_memory_limit_mb,
        },
        "model": {
-            "tts": "hexgrad/Kokoro-82M",
+            "tts": "suno/bark",
            "asr_current_size": _asr_model_size,
            "asr_recommended_size": caps.recommended_model_size,
            "available_sizes": list(WHISPER_MODEL_SIZES.keys()),
@@ -0,0 +1,20 @@
+import asyncio
+from backend.tts_asr import _tts_sync_with_retry, _load_asr_pipeline_with_retry, _asr_pipeline
+import base64
+
+async def main():
+    text = "早上好"
+    print(f"Testing TTS with text: {text}")
+    audio_bytes, sr = await _tts_sync_with_retry(text, rate=1.0)
+    print(f"TTS generated {len(audio_bytes)} bytes of audio.")
+    
+    print("Testing ASR...")
+    await _load_asr_pipeline_with_retry()
+    asr = _asr_pipeline
+    
+    # Needs to process audio_bytes. ASR expects float32 numpy array or bytes?
+    # the pipeline takes bytes or dict with raw array
+    result = asr(audio_bytes) 
+    print("ASR output:", result)
+
+asyncio.run(main())