diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..0ce5574 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,6 @@ +# Keep the Docker build context to only the files copied by Dockerfile. +* + +!export_onnx.py +!CMakeLists.txt +!pocket_tts.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index b906194..419107c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -32,7 +32,7 @@ include(FetchContent) set(FETCHCONTENT_QUIET OFF) # ONNX Runtime (pre-built) - https://github.com/microsoft/onnxruntime/releases -set(ONNXRUNTIME_VERSION "1.23.2") +set(ONNXRUNTIME_VERSION "1.23.2" CACHE STRING "ONNX Runtime version") if(CMAKE_SYSTEM_NAME STREQUAL "Linux") if(CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64|AMD64") diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..8bafc13 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,81 @@ +ARG PYTHON_IMAGE=python:3.12-slim-bookworm +ARG POCKET_TTS_REF=058886528d0b6f2f2d4022de2e244a5260729e6e +ARG EXPORTER_ONNXRUNTIME_VERSION=1.23.2 +ARG ONNXRUNTIME_VERSION=1.23.2 + +FROM ${PYTHON_IMAGE} AS exporter +ARG POCKET_TTS_REF +ARG EXPORTER_ONNXRUNTIME_VERSION + +ENV PYTHONDONTWRITEBYTECODE=1 \ + PYTHONUNBUFFERED=1 \ + PIP_NO_CACHE_DIR=1 + +WORKDIR /src + +RUN apt-get update && \ + apt-get install -y --no-install-recommends git ca-certificates && \ + rm -rf /var/lib/apt/lists/* + +COPY export_onnx.py ./ + +RUN python -m pip install --upgrade pip && \ + python -m pip install --index-url https://download.pytorch.org/whl/cpu torch && \ + python -m pip install \ + "pocket-tts @ git+https://github.com/kyutai-labs/pocket-tts.git@${POCKET_TTS_REF}" \ + onnx \ + "onnxruntime==${EXPORTER_ONNXRUNTIME_VERSION}" && \ + python -m pip freeze | tee /opt/pocket-tts-exporter-requirements.txt + +RUN mkdir -p /opt/pocket-tts/models && \ + python export_onnx.py --output-dir /opt/pocket-tts/models --language english_2026-01 --no-validate + + +FROM ${PYTHON_IMAGE} AS builder +ARG ONNXRUNTIME_VERSION + +ENV PIP_NO_CACHE_DIR=1 + +WORKDIR /src + +RUN apt-get update && \ + apt-get install -y --no-install-recommends \ + build-essential \ + ca-certificates \ + git && \ + rm -rf /var/lib/apt/lists/* + +RUN python -m pip install --upgrade pip cmake + +COPY CMakeLists.txt pocket_tts.cpp ./ + +RUN cmake -S . -B /tmp/build -DCMAKE_BUILD_TYPE=Release -DBUILD_SHARED_LIB=OFF -DONNXRUNTIME_VERSION="${ONNXRUNTIME_VERSION}" && \ + cmake --build /tmp/build -j"$(nproc)" + +RUN mkdir -p /opt/pocket-tts/runtime && \ + install -Dm755 /src/pocket-tts /opt/pocket-tts/runtime/pocket-tts && \ + cp -a /tmp/build/_deps/onnxruntime-src/lib/. /opt/pocket-tts/runtime/ + + +FROM ${PYTHON_IMAGE} AS runtime + +RUN apt-get update && \ + apt-get install -y --no-install-recommends \ + ca-certificates \ + libgomp1 && \ + rm -rf /var/lib/apt/lists/* + +WORKDIR /app + +ENV LD_LIBRARY_PATH=/app + +COPY --from=builder /opt/pocket-tts/runtime/ /app/ +COPY --from=exporter /opt/pocket-tts/models/ /app/models/ +COPY --from=exporter /opt/pocket-tts-exporter-requirements.txt /app/exporter-requirements.txt + +RUN mkdir -p /app/voices/.cache + +EXPOSE 8080 + +ENTRYPOINT ["/app/pocket-tts"] +CMD ["--server", "--port", "8080"] diff --git a/README.md b/README.md index 43cead7..bc18f15 100644 --- a/README.md +++ b/README.md @@ -121,7 +121,7 @@ The `/v1/audio/speech` endpoint is compatible with the OpenAI TTS API. Any clien ```bash curl -X POST http://localhost:8080/v1/audio/speech \ -H "Content-Type: application/json" \ - -d '{"model": "tts-1", "input": "Hello world!", "voice": "voice", "response_format": "wav"}' \ + -d '{"model": "tts-1", "input": "Hello world!", "voice": "voice.wav", "response_format": "wav"}' \ --output speech.wav ``` @@ -145,6 +145,27 @@ int ptt_stream_read(void* stream_ctx, float** out_samples, int* out_len); void ptt_stream_end(void* stream_ctx); ``` +### Docker + +This will build and run `pocket-tts-cpp:latest`: + +``` +docker compose up --build +``` + +After build can also be used like this: +``` +docker run --rm -it -v ./voices:/app/voices -p 8080:8080 pocket-tts-cpp --server --port 8080 +``` + +Usage same as before: +``` +curl -X POST http://localhost:8080/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{"model": "tts-1", "input": "Hello world!", "voice": "voice.wav", "response_format": "wav"}' \ + --output speech.wav +``` + ## Caching PocketTTS uses two layers of disk caching, both stored under `voices/.cache/`: diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..6cdc579 --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,15 @@ +services: + pocket-tts-cpp: + image: pocket-tts-cpp:latest + container_name: pocket-tts-cpp + build: + context: . + args: + POCKET_TTS_REF: ${POCKET_TTS_REF:-058886528d0b6f2f2d4022de2e244a5260729e6e} + EXPORTER_ONNXRUNTIME_VERSION: ${POCKET_TTS_EXPORTER_ONNXRUNTIME_VERSION:-1.23.2} + ONNXRUNTIME_VERSION: ${POCKET_TTS_ONNXRUNTIME_VERSION:-1.23.2} + volumes: + - ./voices:/app/voices + ports: + - "8080:8080" + command: ["--server", "--port", "8080", "--precision", "${POCKET_TTS_PRECISION:-fp32}"] diff --git a/export_onnx.py b/export_onnx.py index 109521d..eda7df0 100644 --- a/export_onnx.py +++ b/export_onnx.py @@ -1198,8 +1198,10 @@ def main(): help="Which models to export") parser.add_argument("--max-seq", type=int, default=MAX_SEQ_LEN, help=f"Max KV cache length for flow_lm (default: {MAX_SEQ_LEN})") + parser.add_argument("--language", default="english_2026-01", + help="Built-in PocketTTS config to use when --config is not set") parser.add_argument("--config", default=None, - help="Path to config YAML (default: use built-in b6369a24)") + help="Path to config YAML (overrides --language)") parser.add_argument("--validate-only", action="store_true", help="Skip export/quantize, only validate existing ONNX files") parser.add_argument("--no-validate", action="store_true", @@ -1247,13 +1249,16 @@ def _download(url: str, dest: Path): _download(f"{HF_BASE}/tokenizer.model", tokenizer_file) print(f" ✓ {tokenizer_file}") - # Build config pointing at local weights + # Build config pointing at local weights/tokenizer. Upstream renamed the + # original b6369a24 English config to english_2026-01 in pocket-tts v2.x. import pocket_tts as _ptt - config_path = Path(_ptt.__file__).parent / "config" / "b6369a24.yaml" if args.config: config_path = Path(args.config) + else: + config_path = Path(_ptt.__file__).parent / "config" / f"{args.language}.yaml" config = load_config(config_path) config.weights_path = str(weights_file) + config.flow_lm.lookup_table.tokenizer_path = str(tokenizer_file) # Load model with voice cloning weights print("Loading model...")