llm-in-text/backend/tts_asr.py

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

os.environ.setdefault("HF_ENDPOINT", "https://hf-mirror.com")

from fastapi import APIRouter, HTTPException
from pydantic import BaseModel

logger = logging.getLogger(__name__)

try:
    import numpy as np  # type: ignore
except Exception as exc:  # pragma: no cover
    logger.debug("numpy import failed: %s", exc)
    np = None  # type: ignore

try:
    import torch  # type: ignore
except Exception as exc:  # pragma: no cover
    logger.debug("torch import failed: %s", exc)
    torch = None  # type: ignore

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

try:
    from faster_whisper import WhisperModel  # type: ignore
except Exception as exc:  # pragma: no cover
    logger.debug("faster_whisper import failed: %s", exc)
    WhisperModel = None  # type: ignore

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

meta_router = APIRouter()
generation_router = APIRouter()

MODEL_ID_HF = "Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign"
MODEL_ID_MS = "Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign"
ASR_MODEL_ID = os.getenv("ASR_MODEL_ID", "small")
ASR_COMPUTE_TYPE = os.getenv("ASR_COMPUTE_TYPE", "int8")

_tts_model: Optional["Qwen3TTSModel"] = None
_asr_model: Optional["WhisperModel"] = None


def _get_device_map() -> str:
    if torch is None:
        return "cpu"
    if torch.cuda.is_available():
        return "cuda"
    try:
        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            return "mps"
    except Exception as exc:  # pragma: no cover
        logger.debug("MPS check failed: %s", exc)
    return "cpu"


def _download_tts_model_from_modelscope() -> Optional[str]:
    if snapshot_download is None:
        return None
    cache_dir = os.path.join(os.path.dirname(__file__), "models")
    os.makedirs(cache_dir, exist_ok=True)
    try:
        return snapshot_download(MODEL_ID_MS, cache_dir=cache_dir, revision="master")
    except Exception as exc:  # pragma: no cover
        logger.warning("ModelScope TTS download failed: %s", exc)
        return None


def _ensure_tts_model() -> "Qwen3TTSModel":
    global _tts_model
    if _tts_model is not None:
        return _tts_model
    if np is None or torch is None or Qwen3TTSModel is None:
        raise RuntimeError("TTS 依赖未安装完整")

    device_map = _get_device_map()
    dtype = torch.float16 if device_map != "cpu" else torch.float32

    model_path = _download_tts_model_from_modelscope()
    last_error = None

    for candidate in [model_path, MODEL_ID_HF]:
        if not candidate:
            continue
        try:
            _tts_model = Qwen3TTSModel.from_pretrained(  # type: ignore
                candidate,
                device_map=device_map,
                dtype=dtype,
            )
            return _tts_model
        except Exception as exc:
            last_error = exc
            logger.warning("TTS model load failed from %s: %s", candidate, exc)

    raise RuntimeError(f"TTS 模型加载失败: {last_error}") from last_error


def _ensure_asr_model() -> "WhisperModel":
    global _asr_model
    if _asr_model is not None:
        return _asr_model
    if WhisperModel is None:
        raise RuntimeError("faster-whisper 未安装")

    device = "cuda" if _get_device_map() == "cuda" else "cpu"
    compute_type = ASR_COMPUTE_TYPE if device == "cpu" else "float16"
    _asr_model = WhisperModel(ASR_MODEL_ID, device=device, compute_type=compute_type)
    return _asr_model


async def _warmup_tts():
    await asyncio.to_thread(_ensure_tts_model)


async def _warmup_asr():
    await asyncio.to_thread(_ensure_asr_model)


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 _normalize_language(language: Optional[str]) -> Optional[str]:
    if not language:
        return None
    value = language.strip().lower()
    if value in {"auto", ""}:
        return None
    mapping = {
        "zh-cn": "zh",
        "zh-hans": "zh",
        "zh-tw": "zh",
        "en-us": "en",
        "ja-jp": "ja",
        "ko-kr": "ko",
    }
    return mapping.get(value, value.split("-")[0])


@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,
        },
        "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()
    await _warmup_asr()
    return {
        "tts_warmup": _tts_model is not None,
        "asr_warmup": _asr_model is not None,
        "device": _get_device_map(),
    }


async def generate_tts_response(
    text: str,
    instruct: str = "",
    speaker: str = "Vivian",
    output_format: str = "wav",
) -> TTSResponse:
    del speaker
    del output_format
    if np is None:
        raise HTTPException(status_code=501, detail="numpy 未安装，TTS 功能不可用")

    try:
        model = _ensure_tts_model()
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc))

    try:
        wavs, sample_rate = await asyncio.to_thread(
            model.generate_voice_design,  # type: ignore
            text=text,
            language="Chinese",
            instruct=instruct or "",
        )
    except Exception as exc:
        logger.exception("TTS inference failed")
        raise HTTPException(status_code=500, detail=f"TTS 推理失败: {exc}")

    wav_data = wavs[0] if isinstance(wavs, (list, tuple)) else wavs
    if hasattr(wav_data, "cpu"):
        wav_data = wav_data.cpu().numpy()
    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)
        sf.write(tmp_path, wav_data, sample_rate)
        with open(tmp_path, "rb") as handle:
            audio_bytes = handle.read()
    except Exception as exc:
        logger.exception("TTS audio encode failed")
        raise HTTPException(status_code=500, detail=f"音频编码失败: {exc}")
    finally:
        if tmp_path and os.path.exists(tmp_path):
            os.unlink(tmp_path)

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


async def generate_asr_response(audio_bytes: bytes, language: Optional[str] = "zh-CN") -> ASRResponse:
    if not audio_bytes:
        raise HTTPException(status_code=400, detail="音频内容为空")

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

    normalized_language = _normalize_language(language)
    tmp_path = None
    try:
        fd, tmp_path = tempfile.mkstemp(suffix=".wav")
        os.close(fd)
        with open(tmp_path, "wb") as handle:
            handle.write(audio_bytes)

        segments, info = await asyncio.to_thread(
            model.transcribe,
            tmp_path,
            language=normalized_language,
            vad_filter=True,
            beam_size=5,
        )
        text = "".join(segment.text for segment in segments).strip()
        if not text:
            raise RuntimeError("ASR 返回结果为空")
        detected_language = getattr(info, "language", normalized_language or "unknown")
        return ASRResponse(text=text, language=str(detected_language))
    except HTTPException:
        raise
    except Exception as exc:
        logger.exception("ASR inference failed")
        raise HTTPException(status_code=500, detail=f"ASR 推理失败: {exc}")
    finally:
        if tmp_path and os.path.exists(tmp_path):
            os.unlink(tmp_path)


@generation_router.post("/tts", response_model=TTSResponse)
async def tts_endpoint(req: TTSRequest):
    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):
    app.include_router(meta_router, prefix="/v1/tts-asr")
    if include_generation_routes:
        app.include_router(generation_router, prefix="/v1/tts-asr")