pikasTech-unidesk/deploy/selfmedia-voice/api.py

import io
import os
import threading
import time
from contextlib import asynccontextmanager

import torch
import torchaudio
import uvicorn
from cosyvoice.cli.cosyvoice import AutoModel
from fastapi import FastAPI, HTTPException
from fastapi.responses import Response
from pydantic import BaseModel, Field


class SpeechRequest(BaseModel):
    model: str | None = None
    input: str = Field(min_length=1)
    voice: str | None = None
    response_format: str = "wav"
    speed: float = 1.0
    instructions: str | None = None
    instruct: str | None = None


class Runtime:
    def __init__(self) -> None:
        self.model = None
        self.loaded_at = None
        self.onnx_providers = None
        self.error = None
        self.lock = threading.Lock()

    def load(self) -> None:
        try:
            if not torch.cuda.is_available():
                raise RuntimeError("CUDA is unavailable")
            device = os.environ["VOICE_DEVICE"]
            if not device.startswith("cuda:"):
                raise RuntimeError("VOICE_DEVICE must select a CUDA device")
            torch.cuda.set_device(int(device.split(":", 1)[1]))
            self.model = AutoModel(
                model_dir=os.environ["MODEL_DIR"],
                load_trt=env_bool("VOICE_LOAD_TRT"),
                load_vllm=env_bool("VOICE_LOAD_VLLM"),
                fp16=env_bool("VOICE_FP16"),
            )
            self.onnx_providers = self.model.frontend.speech_tokenizer_session.get_providers()
            if "CUDAExecutionProvider" not in self.onnx_providers:
                raise RuntimeError(f"CUDAExecutionProvider was not instantiated: {self.onnx_providers}")
            self.loaded_at = time.time()
        except Exception as error:
            self.error = f"{type(error).__name__}: {error}"
            raise


def env_bool(name: str) -> bool:
    value = os.environ[name].lower()
    if value not in {"true", "false"}:
        raise RuntimeError(f"{name} must be true or false")
    return value == "true"


runtime = Runtime()


@asynccontextmanager
async def lifespan(_: FastAPI):
    runtime.load()
    yield


app = FastAPI(title="SelfMedia Voice", version="1", lifespan=lifespan)


@app.get("/health")
def health() -> dict:
    return {
        "ok": runtime.model is not None and runtime.error is None,
        "model": os.environ["MODEL_ID"],
        "modelRef": os.environ["MODEL_REF"],
        "sourceRef": os.environ["SOURCE_REF"],
        "torch": torch.__version__,
        "torchaudio": torchaudio.__version__,
        "cuda": torch.version.cuda,
        "gpu": torch.cuda.get_device_name(torch.cuda.current_device()) if torch.cuda.is_available() else None,
        "device": os.environ["VOICE_DEVICE"],
        "fp16": env_bool("VOICE_FP16"),
        "loadTrt": env_bool("VOICE_LOAD_TRT"),
        "loadVllm": env_bool("VOICE_LOAD_VLLM"),
        "workers": int(os.environ["VOICE_WORKERS"]),
        "concurrency": os.environ["VOICE_CONCURRENCY"],
        "onnxProviders": runtime.onnx_providers,
        "loadedAt": runtime.loaded_at,
        "error": runtime.error,
    }


@app.post("/v1/audio/speech")
def speech(request: SpeechRequest) -> Response:
    if request.response_format.lower() != "wav":
        raise HTTPException(status_code=400, detail="response_format must be wav")
    if not float(os.environ["VOICE_SPEED_MIN"]) <= request.speed <= float(os.environ["VOICE_SPEED_MAX"]):
        raise HTTPException(status_code=400, detail="speed is outside the configured range")
    if runtime.model is None:
        raise HTTPException(status_code=503, detail=runtime.error or "model is not ready")
    if request.voice not in {None, os.environ["REFERENCE_VOICE_ID"]}:
        raise HTTPException(status_code=400, detail="voice must select the configured reference voice")

    started = time.perf_counter()
    instruction = request.instructions or request.instruct
    with runtime.lock:
        inference = (
            runtime.model.inference_instruct2(
                request.input,
                instruction,
                os.environ["PROMPT_WAV"],
                stream=False,
                speed=request.speed,
            )
            if instruction
            else runtime.model.inference_zero_shot(
                request.input,
                os.environ["PROMPT_TEXT"],
                os.environ["PROMPT_WAV"],
                stream=False,
                speed=request.speed,
            )
        )
        chunks = [
            output["tts_speech"].cpu()
            for output in inference
        ]
    if not chunks:
        raise HTTPException(status_code=500, detail="empty inference result")
    waveform = torch.cat(chunks, dim=1)
    output = io.BytesIO()
    torchaudio.save(output, waveform, runtime.model.sample_rate, format="wav")
    return Response(
        content=output.getvalue(),
        media_type="audio/wav",
        headers={
            "X-Inference-Seconds": f"{time.perf_counter() - started:.6f}",
            "X-Sample-Rate": str(runtime.model.sample_rate),
            "X-Inference-Mode": "instruct2" if instruction else "zero-shot",
        },
    )


if __name__ == "__main__":
    workers = int(os.environ["VOICE_WORKERS"])
    if workers != 1:
        raise RuntimeError("VOICE_WORKERS must be 1 for the single-GPU model runtime")
    uvicorn.run(app, host="0.0.0.0", port=int(os.environ["VOICE_PORT"]), workers=workers)