llm-in-text/backend/tts_asr.py

import asyncio
import base64
import io
import logging
import os
import tempfile
import wave
from typing import Optional

# 设置 Hugging Face / ModelScope 镜像源为国内镜像
os.environ.setdefault("HF_ENDPOINT", "https://hf-mirror.com")

import numpy as np
import torch
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel

logger = logging.getLogger(__name__)

try:
    asyncio.get_running_loop()
except RuntimeError:
    asyncio.set_event_loop(asyncio.new_event_loop())

# New TTS model import
try:
    from qwen_tts import Qwen3TTSModel  # type: ignore
except Exception as e:  # pragma: no cover
    logger.debug("qwen_tts import failed (optional): %s", e)
    Qwen3TTSModel = None  # type: ignore

# ASR model import (MLX-based, Apple Silicon only)
try:
    from mlx_audio.stt.models.qwen3_asr import (  # type: ignore
        ForcedAlignerModel,
        Qwen3ASRModel,
    )
except Exception as e:  # pragma: no cover
    logger.debug("mlx_audio import failed (optional): %s", e)
    Qwen3ASRModel = None  # type: ignore
    ForcedAlignerModel = None  # type: ignore

try:
    from modelscope import snapshot_download  # type: ignore
except Exception as e:  # pragma: no cover
    logger.debug("modelscope import failed (optional): %s", e)

meta_router = APIRouter()
generation_router = APIRouter()

# Global model instances
_tts_model: Optional["Qwen3TTSModel"] = None
_asr_model: Optional[object] = None  # Qwen3ASRModel or ForcedAlignerModel
_align_model: Optional[object] = None  # Qwen3-ForcedAlignerModel

# Model paths for loading
MODEL_ID_HF = "Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign"
MODEL_ID_MS = "Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign"

# ModelScope ASR/ForcedAligner models (MLX 4-bit format)
ASR_MODEL_ID_MS = "aufklarer/Qwen3-ASR-0.6B-MLX-4bit"
ALIGN_MODEL_ID_MS = "aufklarer/Qwen3-ForcedAligner-0.6B-MLX"


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 as e:  # noqa: ANN001
        logger.debug("MPS check failed: %s", e)
    return "cpu"


def _download_model_from_modelscope() -> Optional[str]:
    """从 ModelScope 下载模型到本地缓存目录"""
    try:
        cache_dir = os.path.join(os.path.dirname(__file__), "models")
        os.makedirs(cache_dir, exist_ok=True)
        model_dir = snapshot_download(
            MODEL_ID_MS,
            cache_dir=cache_dir,
            revision="master"
        )
        logger.info("ModelScope 模型下载完成: %s", model_dir)
        return model_dir
    except Exception as e:  # noqa: ANN001
        logger.warning("ModelScope 下载失败: %s", e)
        return None


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


async def _warmup_asr():
    """预热 ASR 模型（从 ModelScope 下载并加载）"""
    await asyncio.to_thread(_load_asr_models)


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

    if Qwen3ASRModel is not None:
        logger.info("[Warmup] 开始预热 ASR 模型...")
        await _warmup_asr()
        logger.info("[Warmup] ASR 模型预热完成")


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 模型")

    device_map = _get_device_map()
    last_err = None

    # 策略1: 尝试从 ModelScope 下载后加载
    for attempt in range(max_retries):
        try:
            logger.info("尝试从 ModelScope 下载 TTS 模型...")
            model_path = _download_model_from_modelscope()
            if model_path and os.path.isdir(model_path):
                _tts_model = Qwen3TTSModel.from_pretrained(  # type: ignore
                    model_path,
                    device_map=device_map,
                    dtype=torch.float16,
                )
                logger.info("ModelScope TTS 模型加载成功: %s", model_path)
                return _tts_model
        except Exception as e:  # noqa: ANN001
            logger.warning("ModelScope TTS 加载失败 (尝试 %d/%d): %s", attempt + 1, max_retries, e)
            last_err = e

    # 策略2: 尝试从 HuggingFace 镜像加载
    for attempt in range(max_retries):
        try:
            logger.info("尝试从 HuggingFace 镜像加载 TTS...")
            _tts_model = Qwen3TTSModel.from_pretrained(  # type: ignore
                MODEL_ID_HF,
                device_map=device_map,
                dtype=torch.float16,
            )
            logger.info("HuggingFace TTS 模型加载成功")
            return _tts_model
        except Exception as e:  # noqa: ANN001
            logger.warning("HuggingFace TTS 加载失败 (尝试 %d/%d): %s", attempt + 1, max_retries, e)
            last_err = e

    raise RuntimeError(f"无法加载 TTS 模型: {last_err}") from last_err


def _load_asr_models() -> None:
    """从 ModelScope 下载并加载 ASR/ForcedAligner MLX 模型"""
    global _asr_model, _align_model

    if snapshot_download is None:
        logger.warning("modelscope 未安装，跳过 ASR 模型加载")
        return

    if Qwen3ASRModel is None:
        logger.warning("mlx_audio 未安装，跳过 ASR 模型加载")
        return

    # Download and load ASR model from ModelScope
    try:
        logger.info("从 ModelScope 下载 ASR 模型...")
        asr_cache_dir = os.path.join(os.path.dirname(__file__), "models", "asr")
        asr_model_dir = snapshot_download(ASR_MODEL_ID_MS, cache_dir=asr_cache_dir)
        _load_asr_from_path(asr_model_dir)
    except Exception as e:  # noqa: ANN001
        logger.warning("ASR ModelScope 下载失败，尝试 hf-mirror: %s", e)
        try:
            _load_asr_from_hf_mirror()
        except Exception as e2:  # noqa: ANN001
            logger.warning("ASR hf-mirror 加载失败，跳过 ASR: %s", e2)

    # Download and load ForcedAligner model from ModelScope
    try:
        logger.info("从 ModelScope 下载 ForcedAligner 模型...")
        align_cache_dir = os.path.join(os.path.dirname(__file__), "models", "aligner")
        align_model_dir = snapshot_download(ALIGN_MODEL_ID_MS, cache_dir=align_cache_dir)
        _load_align_from_path(align_model_dir)
    except Exception as e:  # noqa: ANN001
        logger.warning("ForcedAligner ModelScope 下载失败，尝试 hf-mirror: %s", e)
        try:
            _load_align_from_hf_mirror()
        except Exception as e2:  # noqa: ANN001
            logger.warning("ForcedAligner hf-mirror 加载失败，跳过: %s", e2)


def _load_asr_from_path(model_dir: str) -> None:
    """从本地路径加载 ASR MLX 模型"""
    global _asr_model
    try:
        from mlx_audio.stt.utils import load as stt_load  # type: ignore

        model = stt_load(model_dir)
        _asr_model = model
        logger.info("ASR 模型加载成功 (路径: %s)", model_dir)
    except Exception as e:  # noqa: ANN001
        logger.warning("ASR MLX 加载失败，尝试直接构建: %s", e)
        try:
            from mlx.core import load as mx_load  # type: ignore

            weights = mx_load(os.path.join(model_dir, "model.safetensors"))
            from mlx_lm import load as lm_load  # type: ignore

            model = lm_load(model_dir, model_cls=Qwen3ASRModel)
            _asr_model = model
        except Exception as e2:  # noqa: ANN001
            raise RuntimeError(f"无法加载 ASR MLX 模型: {e2}") from e


def _load_asr_from_hf_mirror() -> None:
    """从 hf-mirror 加载 ASR MLX 模型"""
    global _asr_model
    try:
        from mlx_audio.stt.utils import load as stt_load  # type: ignore

        model = stt_load("mlx-community/Qwen3-ASR-0.6B-4bit")
        _asr_model = model
    except Exception as e:  # noqa: ANN001
        raise RuntimeError(f"无法从 hf-mirror 加载 ASR MLX: {e}") from e


def _load_align_from_path(model_dir: str) -> None:
    """从本地路径加载 ForcedAligner MLX 模型"""
    global _align_model
    try:
        from mlx_audio.stt.utils import load as stt_load  # type: ignore

        model = stt_load(model_dir)
        _align_model = model
    except Exception as e:  # noqa: ANN001
        raise RuntimeError(f"无法加载 ForcedAligner MLX 模型 (路径: {model_dir}): {e}") from e


def _load_align_from_hf_mirror() -> None:
    """从 hf-mirror 加载 ForcedAligner MLX 模型"""
    global _align_model
    try:
        from mlx_audio.stt.utils import load as stt_load  # type: ignore

        model = stt_load("mlx-community/Qwen3-ForcedAligner-0.6B-4bit")
        _align_model = model
    except Exception as e:  # noqa: ANN001
        raise RuntimeError(f"无法从 hf-mirror 加载 ForcedAligner MLX: {e}") from e


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 ASRRequest(BaseModel):
    audio_base64: str
    language: Optional[str] = "zh-CN"


class ASRResponse(BaseModel):
    text: str
    language: Optional[str] = None


class ModelStatus(BaseModel):
    tts_loaded: bool
    asr_loaded: bool = False
    device: str


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


def _ensure_asr_model():
    """确保 ASR 模型已加载（懒加载）"""
    global _asr_model
    if _asr_model is None:
        try:
            from mlx_audio.stt.utils import load as stt_load  # type: ignore

            _asr_model = stt_load(ASR_MODEL_ID_MS)
        except Exception as e:  # noqa: ANN001
            raise RuntimeError(f"无法加载 ASR MLX 模型 (路径: {ASR_MODEL_ID_MS}): {e}") from e
    return _asr_model


def _ensure_align_model():
    """确保 ForcedAligner 模型已加载（懒加载）"""
    global _align_model
    if _align_model is None:
        try:
            from mlx_audio.stt.utils import load as stt_load  # type: ignore

            _align_model = stt_load(ALIGN_MODEL_ID_MS)
        except Exception as e:  # noqa: ANN001
            raise RuntimeError(f"无法加载 ForcedAligner MLX 模型 (路径: {ALIGN_MODEL_ID_MS}): {e}") from e
    return _align_model


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


@meta_router.get("/config")
async def get_config():
    """获取配置信息"""
    return {
        "model": {
            "tts": MODEL_ID_MS,
            "asr": ASR_MODEL_ID_MS if Qwen3ASRModel is not None else None,
        },
        "device": _get_device_map(),
        "status": {
            "tts_loaded": _tts_model is not None,
            "asr_loaded": _asr_model is not None,
        }
    }


@meta_router.post("/warmup")
async def warmup_models():
    """手动触发模型预热"""
    await _warmup_tts()

    if Qwen3ASRModel is not None:
        await _warmup_asr()

    return {
        "tts_warmup": _tts_model is not None,
        "asr_warmup": _asr_model is not None if Qwen3ASRModel else False,
        "device": _get_device_map(),
    }


async def generate_tts_response(
    text: str,
    instruct: str = "",
    speaker: str = "Vivian",
    output_format: str = "wav",
) -> TTSResponse:
    del speaker  # current model path does not expose multi-speaker routing
    del output_format  # current implementation always returns wav
    try:
        model = _ensure_tts_model()
    except Exception as e:  # noqa: ANN001
        raise HTTPException(status_code=500, detail=str(e))

    try:
        wavs, sr = model.generate_voice_design(  # type: ignore
            text=text,
            language="Chinese",
            instruct=instruct or "",
        )
    except Exception as e:  # noqa: ANN001
        logger.exception("TTS 推理失败")
        raise HTTPException(status_code=500, detail=f"TTS 推理失败: {e}")

    wav_data = wavs[0] if isinstance(wavs, (list, tuple)) else wavs
    if hasattr(wav_data, 'numpy'):  # type: ignore
        wav_data = wav_data.cpu().numpy()  # type: ignore
    wav_data = np.asarray(wav_data, dtype=np.float32)

    tmp_path = None
    try:
        import soundfile as sf  # type: ignore

        fd, tmp_path = tempfile.mkstemp(suffix=".wav")
        os.close(fd)  # type: ignore
        sf.write(tmp_path, wav_data, sr)
        with open(tmp_path, "rb") as f:  # noqa: SIM115
            audio_bytes = f.read()
    except Exception as e:  # noqa: ANN001
        logger.exception("音频编码失败")
        raise HTTPException(status_code=500, detail=f"音频编码失败: {e}")
    finally:
        if tmp_path and os.path.exists(tmp_path):  # noqa: SIM201
            try:
                os.unlink(tmp_path)
            except Exception:
                pass

    duration_ms = int(len(wav_data) / sr * 1000) if sr > 0 else 0
    audio_base64 = base64.b64encode(audio_bytes).decode("utf-8")
    return TTSResponse(audio_base64=audio_base64, format="wav", duration_ms=duration_ms)


async def generate_asr_response(audio_bytes: bytes, language: Optional[str] = "zh-CN") -> ASRResponse:
    if Qwen3ASRModel is None:
        raise HTTPException(status_code=501, detail="mlx_audio 未安装，ASR 功能不可用")

    try:
        model = _ensure_asr_model()
    except Exception as e:  # noqa: ANN001
        raise HTTPException(status_code=500, detail=f"ASR 模型加载失败: {e}")

    try:
        wav_buffer = io.BytesIO(audio_bytes)
        with wave.open(wav_buffer, 'rb') as wf:  # noqa: SIM115
            n_channels = wf.getnchannels()
            sampwidth = wf.getsampwidth()
            framerate = wf.getframerate()
            n_frames = wf.getnframes()

            raw_data = wf.readframes(n_frames)
            audio_array = np.frombuffer(raw_data, dtype=np.int16 if sampwidth == 2 else np.float32)

            if n_channels > 1:
                audio_array = np.mean(audio_array.reshape(-1, n_channels), axis=1)

            if framerate != 16000:
                try:
                    import scipy.signal as signal  # type: ignore

                    n_samples = int(len(audio_array) * 16000 / framerate)
                    audio_array = signal.resample(audio_array, n_samples)  # type: ignore
                except Exception as e2:  # noqa: ANN001
                    logger.warning("重采样失败，使用原始音频: %s", e2)

            if audio_array.dtype == np.int16:
                audio_array = audio_array.astype(np.float32) / 32768.0

        result = model.generate(  # type: ignore
            audio_array,
            language=language if language else None,
        )

        recognized_text = getattr(result, 'text', str(result)) if hasattr(result, 'text') else str(result)
        detected_lang = getattr(result, 'language', language or "zh-CN")
        if isinstance(detected_lang, list) and len(detected_lang) > 0:
            detected_lang = detected_lang[0]

        return ASRResponse(text=recognized_text, language=str(detected_lang))
    except HTTPException:
        raise
    except Exception as e:  # noqa: ANN001
        logger.exception("ASR 推理失败")
        raise HTTPException(status_code=500, detail=f"ASR 推理失败: {e}")


@generation_router.post("/tts", response_model=TTSResponse)
async def tts_endpoint(req: TTSRequest):
    """TTS 文字转语音端点"""
    return await generate_tts_response(
        text=req.text,
        instruct=req.instruct or "",
        speaker=req.speaker,
        output_format=req.format,
    )


@generation_router.post("/asr", response_model=ASRResponse)
async def asr_endpoint(req: ASRRequest):
    """语音识别端点（非流式）"""
    audio_bytes = base64.b64decode(req.audio_base64)
    return await generate_asr_response(audio_bytes, req.language if req.language else None)


def register_tts_asr_routes(app, include_generation_routes: bool = True):
    """注册 TTS/ASR 路由到 FastAPI 应用"""
    app.include_router(meta_router, prefix="/v1/tts-asr")
    if include_generation_routes:
        app.include_router(generation_router, prefix="/v1/tts-asr")


router = APIRouter()
router.include_router(meta_router)
router.include_router(generation_router)