llm-in-text/backend/tts_asr.py

import asyncio
import base64
import logging
from io import BytesIO
from typing import Optional

import torch
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel

logger = logging.getLogger(__name__)

# New TTS model import
try:
    from qwen_tts import Qwen3TTSModel  # type: ignore
except Exception:  # pragma: no cover
    Qwen3TTSModel = None  # type: ignore

router = APIRouter()

# Global TTS model instance
_tts_model: Optional["Qwen3TTSModel"] = None


def _get_device_map() -> str:
    """设备检测逻辑：优先 CUDA，其次 MPS，最后 CPU"""
    if torch.cuda.is_available():
        return "cuda:0"
    try:
        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            return "mps"
    except Exception:
        pass
    return "cpu"


async def _warmup_tts():
    """预热 TTS 模型"""
    await asyncio.to_thread(_load_tts_model_with_retry)


async def _warmup_all():
    """预热所有模型（TTS 和 ASR）"""
    logger.info("[Warmup] 开始预热 TTS 模型...")
    await _warmup_tts()
    logger.info("[Warmup] TTS 模型预热完成")


def _load_tts_model_with_retry(max_retries: int = 3) -> "Qwen3TTSModel":
    """加载 TTS 模型，支持多个镜像源"""
    global _tts_model
    if _tts_model is not None:
        return _tts_model
    if Qwen3TTSModel is None:
        raise RuntimeError("qwen_tts 库未安装，无法加载 TTS 模型")

    candidates = [
        "Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign",
        "ModelScope/Qwen3-TTS-12Hz-1.7B-VoiceDesign",
    ]
    device_map = _get_device_map()
    last_err = None
    for i, model_id in enumerate(candidates, start=1):
        try:
            _tts_model = Qwen3TTSModel.from_pretrained(
                model_id,
                device_map=device_map,
                dtype=torch.float16,
                attn_implementation="flash_attention_2",
            )
            logger.info("Loaded TTS model from %s", model_id)
            return _tts_model
        except Exception as e:
            logger.warning("Failed to load TTS model from %s: %s", model_id, e)
            last_err = e
            if i >= max_retries:
                break
    raise RuntimeError(f"Unable to load TTS model from sources: {candidates}") from last_err


class TTSRequest(BaseModel):
    text: str
    instruct: str = ""
    speaker: str = "Vivian"
    format: str = "wav"


class TTSResponse(BaseModel):
    audio_base64: str
    format: str
    duration_ms: int


class ModelStatus(BaseModel):
    tts_loaded: bool
    asr_loaded: bool = False
    device: str
    tts_last_used: Optional[float] = None
    asr_last_used: Optional[float] = None


def _ensure_model() -> "Qwen3TTSModel":
    """确保模型已加载"""
    global _tts_model
    if _tts_model is None:
        _tts_model = _load_tts_model_with_retry()
    return _tts_model


@router.get("/status", response_model=ModelStatus)
async def get_status():
    """获取模型状态"""
    return ModelStatus(
        tts_loaded=_tts_model is not None,
        asr_loaded=False,
        device=_get_device_map(),
    )


@router.get("/config")
async def get_config():
    """获取配置信息"""
    return {
        "model": {
            "tts": "Qwen3-TTS-12Hz-1.7B-VoiceDesign",
            "asr": None,
        },
        "device": _get_device_map(),
        "status": {
            "tts_loaded": _tts_model is not None,
            "asr_loaded": False,
        }
    }


@router.post("/warmup")
async def warmup_models():
    """手动触发模型预热"""
    await _warmup_tts()
    return {
        "tts_warmup": _tts_model is not None,
        "device": _get_device_map(),
    }


@router.post("/tts", response_model=TTSResponse)
async def tts_endpoint(req: TTSRequest):
    """TTS 文字转语音端点"""
    try:
        model = _ensure_model()
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

    text = req.text
    instruct = req.instruct or ""
    speaker = req.speaker or "Vivian"

    try:
        wavs_sr = model.generate_custom_voice(
            text=text,
            language="Chinese",
            speaker=speaker,
            instruct=instruct,
        )
    except Exception as e:
        logger.exception("TTS 推理失败")
        raise HTTPException(status_code=500, detail=f"TTS 推理失败: {e}")

    # Normalize output
    if isinstance(wavs_sr, tuple) and len(wavs_sr) == 2:
        wav_data, sr = wavs_sr
    else:
        wav_data, sr = wavs_sr[0], wavs_sr[1]  # type: ignore

    # 编码 WAV 到内存
    try:
        import soundfile as sf
        bio = BytesIO()
        sf.write(bio, wav_data, sr, format="WAV")
        audio_bytes = bio.getvalue()
    except Exception as e:
        logger.exception("音频编码失败")
        raise HTTPException(status_code=500, detail=f"音频编码失败: {e}")

    # 计算时长（毫秒）
    duration_ms = int(len(wav_data) / sr * 1000) if sr > 0 else 0

    # 返回 JSON 格式，包含 base64 编码的音频
    audio_base64 = base64.b64encode(audio_bytes).decode("utf-8")
    return TTSResponse(
        audio_base64=audio_base64,
        format="wav",
        duration_ms=duration_ms,
    )


def register_tts_asr_routes(app):
    """注册 TTS/ASR 路由到 FastAPI 应用"""
    app.include_router(router, prefix="/v1/tts-asr")