Copy the files in the new repo for public release
This commit is contained in:
gabrieldemarmiesse
@@ -0,0 +1,19 @@
|
||||
name: Upload Python Package
|
||||
|
||||
on:
|
||||
release:
|
||||
types: [created]
|
||||
|
||||
jobs:
|
||||
deploy:
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
id-token: write
|
||||
steps:
|
||||
- uses: actions/checkout@v2
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
- name: Build wheels
|
||||
run: uv build
|
||||
- name: Publish package distributions to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
@@ -0,0 +1,39 @@
|
||||
name: Run Tests
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [ main ]
|
||||
pull_request:
|
||||
|
||||
jobs:
|
||||
pre-commit:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
- name: Install dependencies
|
||||
run: uv sync --dev
|
||||
- name: Run pre-commit
|
||||
uses: pre-commit/action@v3.0.0
|
||||
|
||||
pytest:
|
||||
runs-on: ubuntu-latest
|
||||
needs: pre-commit
|
||||
strategy:
|
||||
matrix:
|
||||
python-version: ['3.10', '3.14']
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- name: Set up Python ${{ matrix.python-version }}
|
||||
uses: actions/setup-python@v4
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
- name: Install dependencies
|
||||
run: uv sync --dev -p ${{ matrix.python-version }}
|
||||
- name: Run pytest
|
||||
env:
|
||||
HF_TOKEN: ${{ secrets.HF_TOKEN }}
|
||||
run: uv run -p ${{ matrix.python-version }} pytest ./tests
|
||||
Executable
+88
@@ -0,0 +1,88 @@
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
|
||||
# C extensions
|
||||
*.so
|
||||
|
||||
# macOS dir files
|
||||
.DS_Store
|
||||
|
||||
# Distribution / packaging
|
||||
.Python
|
||||
env/
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
wheels/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
.ipynb_checkpoints
|
||||
|
||||
# Tests and linter
|
||||
.pytest_cache/
|
||||
.mypy_cache/
|
||||
.coverage
|
||||
coverage.json
|
||||
|
||||
# docs
|
||||
/api_docs
|
||||
|
||||
# dotenv
|
||||
.env
|
||||
.envrc
|
||||
|
||||
# virtualenv
|
||||
.venv
|
||||
venv/
|
||||
ENV/
|
||||
|
||||
# egs with manifest files
|
||||
egs/*
|
||||
!egs/example
|
||||
# local datasets
|
||||
dataset/*
|
||||
!dataset/example
|
||||
|
||||
# personal notebooks & scripts
|
||||
*/local_scripts
|
||||
*/notes
|
||||
.vscode/
|
||||
/notebooks
|
||||
/local_scripts
|
||||
/notes
|
||||
|
||||
# cython output
|
||||
/audiocraft/metrics/abx/*.c
|
||||
|
||||
# rust-analyzer
|
||||
Cargo.toml
|
||||
Cargo.lock
|
||||
target/
|
||||
|
||||
# nfs files
|
||||
.nfs*
|
||||
.coverage.*
|
||||
|
||||
/weights
|
||||
test_en_audio_4000.model
|
||||
profile_output*
|
||||
*.safetensors
|
||||
*.wav
|
||||
*.model
|
||||
*.pt
|
||||
|
||||
# Agents files, use a symlink to AGENTS.md instead.
|
||||
CLAUDE.md
|
||||
QWEN.md
|
||||
.claude/
|
||||
@@ -0,0 +1,19 @@
|
||||
repos:
|
||||
- repo: https://github.com/astral-sh/ruff-pre-commit
|
||||
rev: v0.12.7
|
||||
hooks:
|
||||
- id: ruff-check
|
||||
args: [ --fix ]
|
||||
- id: ruff-check
|
||||
args: [ --select=I, --fix ]
|
||||
- id: ruff-format
|
||||
# Use ruff format on markdown files
|
||||
- id: ruff-format
|
||||
types_or: [markdown]
|
||||
require_serial: true
|
||||
entry: doccmd --language python --no-pad-file --command "ruff format"
|
||||
additional_dependencies: ["doccmd"]
|
||||
- repo: https://github.com/astral-sh/uv-pre-commit
|
||||
rev: 0.8.22
|
||||
hooks:
|
||||
- id: uv-lock
|
||||
@@ -0,0 +1 @@
|
||||
3.10
|
||||
@@ -0,0 +1,132 @@
|
||||
# AGENTS.md
|
||||
|
||||
This file provides guidance to AI agents when working with code in this repository.
|
||||
|
||||
## Project Overview
|
||||
|
||||
pocket-tts is a CPU-based text-to-speech (TTS) model. The project uses a flow-based language model architecture with a neural audio codec (Mimi) for efficient speech synthesis.
|
||||
|
||||
**Key Architecture Components:**
|
||||
- **FlowLMModel**: Transformer-based flow language model that generates latent representations from text using Lagrangian Self Distillation (LSD)
|
||||
- **MimiModel**: Neural audio codec (from the `moshi` package) that compresses/decompresses audio to/from latent representations
|
||||
- **Conditioners**: Text processing via SentencePiece tokenizer and lookup table embeddings
|
||||
- **Streaming Architecture**: The entire pipeline supports streaming generation via stateful modules
|
||||
|
||||
## Common Commands
|
||||
|
||||
### Setup and Development
|
||||
```bash
|
||||
# Install pre-commit hooks
|
||||
uvx pre-commit install
|
||||
|
||||
# Run tests (3 parallel workers)
|
||||
uv run pytest -n 3 -v
|
||||
|
||||
# Run a single test
|
||||
uv run pytest tests/test_python_api.py -v
|
||||
|
||||
# Run CLI locally (editable install)
|
||||
uv run pocket-tts generate
|
||||
uv run pocket-tts serve
|
||||
```
|
||||
|
||||
### Linting and Formatting
|
||||
Pre-commit handles this automatically, but you can run manually:
|
||||
```bash
|
||||
# Ruff will run automatically on commit via pre-commit
|
||||
# Includes: ruff-check, ruff-format (with --fix), and import sorting
|
||||
```
|
||||
|
||||
### Building (No Build Step)
|
||||
This is a pure Python package with Rust extensions in `training/rust_exts/audio_ds/` for training-time audio processing. The main package does not require building.
|
||||
|
||||
## Code Structure
|
||||
|
||||
### Main Package (`pocket_tts
|
||||
|
||||
**Entry Points:**
|
||||
- `main.py`: CLI implementation with Typer (two commands: `generate` and `serve`)
|
||||
- `__init__.py`: Public API exports only `TTSModel`
|
||||
|
||||
**Core Models (`models/`):**
|
||||
- `tts_model.py`: Main `TTSModel` class - orchestrates the entire TTS pipeline
|
||||
- `load_model()`: Downloads weights from HuggingFace and initializes models
|
||||
- `get_state_for_audio_prompt()`: Encodes audio prompt (voice) into model state
|
||||
- `generate_audio_stream()`: Streaming generation that yields audio chunks
|
||||
- Uses LRU cache for voice prompts to avoid reprocessing
|
||||
- `flow_lm.py`: `FlowLMModel` - transformer that generates latent audio codes from text
|
||||
|
||||
**Modules (`modules/`):**
|
||||
- `transformer.py`: `StreamingTransformer` and `StreamingMultiheadAttention` with RoPE embeddings
|
||||
- `stateful_module.py`: Base class for streaming support (maintains KV cache and state)
|
||||
- `rope.py`: Rotary Position Embeddings
|
||||
- `mlp.py`: `SimpleMLPAdaLN` (AdaLN-conditioned MLP for flow prediction)
|
||||
- `conv.py`: Convolution utilities
|
||||
- `seanet.py`: SEANet encoder/decoder (imported from moshi)
|
||||
|
||||
**Conditioners (`conditioners/`):**
|
||||
- `text.py`: `LUTConditioner` - SentencePiece tokenizer + embedding lookup table for text
|
||||
|
||||
**Data (`data/`):**
|
||||
- `audio.py`: Audio I/O utilities (reading, writing WAV, streaming)
|
||||
- `audio_utils.py`: Audio processing (resampling, conversion)
|
||||
|
||||
**Utils (`utils/`):**
|
||||
- `config.py`: Pydantic config models for FlowLM and Mimi
|
||||
- `utils.py`: HuggingFace downloads, timing utilities
|
||||
|
||||
**Configuration (`config/`):**
|
||||
- `b6369a24.yaml`: Model configuration (transformer dims, layers, vocab size, etc.)
|
||||
|
||||
### Testing (`tests/`)
|
||||
- `test_python_api.py`: Tests for public Python API
|
||||
- `test_cli_generate.py`: Tests for CLI generate command
|
||||
- `test_documentation_examples.py`: Ensures docs examples work
|
||||
|
||||
## Development Workflow
|
||||
|
||||
### Key Patterns
|
||||
|
||||
1. **Streaming Generation**: The model generates audio frame-by-frame (12.5 Hz frame rate, 80ms per frame). All modules inherit from `StatefulModule` to maintain internal state.
|
||||
|
||||
2. **Voice Cloning**: Audio prompts are encoded via Mimi encoder to create "voice state" (latent representations + speaker embeddings). This state is cached via `lru_cache` on `_cached_get_state_for_audio_prompt()`.
|
||||
|
||||
3. **Flow Matching**: Uses Lagrangian Self Distillation (LSD) with configurable decode steps. Fewer steps = faster but lower quality.
|
||||
|
||||
4. **EOS Detection**: Model predicts end-of-speech via an EOS head. Generation continues for `frames_after_eos` frames after EOS is detected.
|
||||
|
||||
5. **Config-Driven**: Model architecture is defined in YAML configs. Weights are loaded from HuggingFace Hub via `safetensors`.
|
||||
|
||||
### Important Implementation Details
|
||||
|
||||
- **Thread Safety**: The code is NOT thread-safe. Server mode does not support concurrent requests.
|
||||
- **Batching**: Batch size is always 1. No batching support currently.
|
||||
- **Device**: Defaults to CPU. GPU does not provide speedup for this small model.
|
||||
- **Torch Threads**: `torch.set_num_threads(1)` in `tts_model.py` for optimal CPU performance
|
||||
- **dtype**: Models use float32 by default (configurable in YAML)
|
||||
- **Beartype**: Runtime type checking is enabled via beartype claw in `__init__.py`
|
||||
|
||||
### Adding Features
|
||||
|
||||
When adding features, be aware of:
|
||||
- The streaming architecture: any changes to model forward passes need to maintain state correctly
|
||||
- The config system: new model parameters must be added to config classes in `utils/config.py`
|
||||
- The public API: only `TTSModel` is exported; keep implementation details internal
|
||||
- Ruff formatting: line length 100, LF line endings, skip magic trailing comma
|
||||
|
||||
### Model Weights
|
||||
|
||||
Weights are downloaded from HuggingFace Hub on first use:
|
||||
- Model weights: `hf://kyutai/pocket-tts/tts_b6369a24.safetensors`
|
||||
- Tokenizer: `hf://kyutai/pocket-tts/tokenizer.model`
|
||||
- Voice prompts: `hf://kyutai/tts-voices/<speaker>/<style>.wav`
|
||||
|
||||
The `download_if_necessary()` utility handles `hf://` URLs and caches locally.
|
||||
|
||||
## Common Gotchas
|
||||
|
||||
1. **PyTorch Version**: Requires PyTorch 2.5+. Version 2.4.0 produces incorrect audio.
|
||||
2. **Python Version**: Supports 3.11, 3.12, 3.13 only.
|
||||
3. **uv Python Preference**: Set to "only-managed" in pyproject.toml because system Python may lack headers.
|
||||
4. **Moshi Dependency**: The package depends on `moshi` for the MimiModel audio codec.
|
||||
5. **CPU-Only PyTorch**: Uses PyTorch CPU index from `download.pytorch.org/whl/cpu` in uv config.
|
||||
@@ -0,0 +1,60 @@
|
||||
# Contributing to Kyutai Pocket TTS
|
||||
|
||||
We welcome contributions from the community!
|
||||
|
||||
We recommend using `uv` to work on this project. While you can use `pip` too, `uv` is less error-prone and faster.
|
||||
The instructions below assume you are using `uv`.
|
||||
|
||||
## Installing the pre-commit
|
||||
|
||||
```bash
|
||||
uvx pre-commit install
|
||||
```
|
||||
This will handle file formatting and linting.
|
||||
|
||||
If you want to manually run the pre-commit hooks on all files, use:
|
||||
```bash
|
||||
uvx pre-commit run --all-files
|
||||
```
|
||||
|
||||
## Running tests
|
||||
|
||||
```bash
|
||||
uv run pytest -n 3 -v
|
||||
```
|
||||
This will run the test suite with 3 parallel workers.
|
||||
|
||||
## Running the CLI locally
|
||||
|
||||
You can run the CLI commands with:
|
||||
|
||||
```bash
|
||||
uv run pocket-tts generate
|
||||
```
|
||||
|
||||
## Coding agents
|
||||
We use `AGENTS.md` to manage the coding agents context. If your coding agent does
|
||||
not supports reading from this file directly, you can just
|
||||
use a symlink. The most common ones are already added in the `.gitignore` file.
|
||||
|
||||
```bash
|
||||
ln -s AGENTS.md CLAUDE.md
|
||||
# or
|
||||
ln -s AGENTS.md QWEN.md
|
||||
```
|
||||
|
||||
## How does this model work?
|
||||
|
||||
Here is a high-level overview of the architecture:
|
||||
|
||||
|
||||
|
||||
Overall the model has four main components:
|
||||
* The mimi vae encoder, to encode the audio prompts into a latent representation.
|
||||
* The text "encoder" which is just a simple tokenizer + embedding layer.
|
||||
* The calm model (`flow_lm` in the codebase) to generate autoregegressive the audio latents.
|
||||
* The mimi vae decoder to decode the audio latents into waveform.
|
||||
|
||||
Note that two threads run in parallel in the current implementation:
|
||||
* One with the calm model generating the latents.
|
||||
* One with the mimi vae decoder decoding the latents into audio.
|
||||
+14
@@ -0,0 +1,14 @@
|
||||
FROM ghcr.io/astral-sh/uv:debian
|
||||
|
||||
WORKDIR /app
|
||||
COPY ./pyproject.toml .
|
||||
COPY ./uv.lock .
|
||||
COPY ./README.md .
|
||||
COPY ./.python-version .
|
||||
COPY ./pocket_tts ./pocket_tts
|
||||
|
||||
RUN uv run pocket-tts serve --help && \
|
||||
rm -rf /root/.cache/uv && \
|
||||
uv run pocket-tts serve --help
|
||||
|
||||
CMD ["uv", "run", "pocket-tts", "serve"]
|
||||
@@ -0,0 +1,23 @@
|
||||
Permission is hereby granted, free of charge, to any
|
||||
person obtaining a copy of this software and associated
|
||||
documentation files (the "Software"), to deal in the
|
||||
Software without restriction, including without
|
||||
limitation the rights to use, copy, modify, merge,
|
||||
publish, distribute, sublicense, and/or sell copies of
|
||||
the Software, and to permit persons to whom the Software
|
||||
is furnished to do so, subject to the following
|
||||
conditions:
|
||||
|
||||
The above copyright notice and this permission notice
|
||||
shall be included in all copies or substantial portions
|
||||
of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF
|
||||
ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
|
||||
TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
|
||||
PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
|
||||
SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
|
||||
CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
|
||||
OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
|
||||
IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
|
||||
DEALINGS IN THE SOFTWARE.
|
||||
@@ -0,0 +1,115 @@
|
||||
# Pocket TTS
|
||||
|
||||
|
||||
|
||||
A lightweight text-to-speech (TTS) application designed to run efficiently on CPUs.
|
||||
Forget about the hassle of using GPUs and web APIs serving TTS models. With Kyutai's Pocket TTS, generating audio is just a pip install and a function call away.
|
||||
|
||||
Supports Python 3.10, 3.11, 3.12, 3.13 and 3.14. Requires PyTorch 2.5+. Does not require the gpu version of PyTorch.
|
||||
|
||||
[🔊 Demo](https://kyutai.org/tts) |
|
||||
[🐱💻GitHub Repository](https://github.com/kyutai-labs/pocket-tts) |
|
||||
[🤗 Hugging Face Model Card](https://huggingface.co/kyutai/pocket-tts) |
|
||||
[📄 Paper](https://arxiv.org/abs/2509.06926) |
|
||||
[📚 Documentation](https://github.com/kyutai-labs/pocket-tts/tree/main/docs)
|
||||
|
||||
|
||||
## Features
|
||||
* Runs on CPU
|
||||
* Small model size (100M parameters)
|
||||
* Audio streaming
|
||||
* Low latency (~150ms for 2-3 sentences)
|
||||
* Faster than real-time (~6x real-time on a CPU of MacBook Air M4)
|
||||
* Uses only 2 CPU cores
|
||||
* Very simple codebase and easy to use
|
||||
* Voice cloning
|
||||
* English only at the moment
|
||||
|
||||
## Trying it from the website, without installing anything
|
||||
|
||||
Navigate to the [Kyutai website](https://kyutai.org/tts) to try it out directly in your browser. You can input text, select different voices, and generate speech without any installation.
|
||||
|
||||
## Trying it with the CLI
|
||||
|
||||
### The `generate` command
|
||||
You can use pocket-tts directly from the command line. We recommend using
|
||||
`uv` as it installs any dependencies on the fly in an isolated environment (uv installation instructions [here](https://docs.astral.sh/uv/getting-started/installation/#standalone-installer)).
|
||||
You can also use `pip install pocket-tts` to install it manually.
|
||||
|
||||
This will generate a wav file `./tts_output.wav` saying the default text with the default voice, and display some speed statistics.
|
||||
```bash
|
||||
uvx pocket-tts generate
|
||||
# or if you installed it manually with pip:
|
||||
pocket-tts generate
|
||||
```
|
||||
Modify the voice with `--voice` and the text with `--text`. We provide a small catalog of voices:
|
||||
```
|
||||
"marius", "javert", "jean", "fantine", "cosette", "eponine", "azelma"
|
||||
```
|
||||
|
||||
Feel free to check out the [generate documentation](https://github.com/kyutai-labs/pocket-tts/tree/main/docs/generate.md) for more details and examples.
|
||||
|
||||
### The `serve` command
|
||||
|
||||
You can also run a local server to generate audio via HTTP requests.
|
||||
```bash
|
||||
uvx pocket-tts serve
|
||||
# or if you installed it manually with pip:
|
||||
pocket-tts serve
|
||||
```
|
||||
Navigate to `http://localhost:8000` to try the web interface, it's faster than the command line as the model is kept in memory between requests.
|
||||
|
||||
You can check out the [serve documentation](https://github.com/kyutai-labs/pocket-tts/tree/main/docs/serve.md) for more details and examples.
|
||||
|
||||
## Using it as a Python library
|
||||
|
||||
Install the package with
|
||||
```bash
|
||||
pip install pocket-tts
|
||||
# or
|
||||
uv add pocket-tts
|
||||
```
|
||||
|
||||
You can use this package as a simple Python library to generate audio from text.
|
||||
```python
|
||||
from pocket_tts import TTSModel
|
||||
import scipy.io.wavfile
|
||||
|
||||
tts_model = TTSModel.load_model()
|
||||
voice_state = tts_model.get_state_for_audio_prompt(
|
||||
"hf://kyutai/tts-voices/alba-mackenna/casual.wav"
|
||||
)
|
||||
audio = tts_model.generate_audio(voice_state, "Hello world, this is a test.")
|
||||
# Audio is a 1D torch tensor containing PCM data.
|
||||
scipy.io.wavfile.write("output.wav", tts_model.sample_rate, audio.numpy())
|
||||
```
|
||||
|
||||
You can have multiple voice states around if
|
||||
you have multiple voices you want to use. `load_model()`
|
||||
and `get_state_for_audio_prompt()` are relatively slow operations,
|
||||
so we recommend to keep the model and voice states in memory if you can.
|
||||
|
||||
You can check out the [Python API documentation](https://github.com/kyutai-labs/pocket-tts/tree/main/docs/python-api.md) for more details and examples.
|
||||
|
||||
## Unsupported features
|
||||
|
||||
At the moment, we do not support (but would love pull requests adding):
|
||||
- [Running the TTS inside a web browser (WebAssembly)](https://github.com/kyutai-labs/pocket-tts/issues/1)
|
||||
- [A compiled version with for example `torch.compile()` or `candle`.](https://github.com/kyutai-labs/pocket-tts/issues/2)
|
||||
- [Handling long prompts by splitting them in a smart way.](https://github.com/kyutai-labs/pocket-tts/issues/5)
|
||||
- [Adding silence in the text input to generate pauses.](https://github.com/kyutai-labs/pocket-tts/issues/6)
|
||||
- [Quantization to run the computation in int8.](https://github.com/kyutai-labs/pocket-tts/issues/7)
|
||||
|
||||
We tried running this TTS model on the GPU but did not observe a speedup compared to CPU execution,
|
||||
notably because we use a batch size of 1 and a very small model.
|
||||
|
||||
## Development and local setup
|
||||
|
||||
We accept contributions! Feel free to open issues or pull requests on GitHub.
|
||||
|
||||
You can find development instructions in the [CONTRIBUTING.md](https://github.com/kyutai-labs/pocket-tts/tree/main/CONTRIBUTING.md) file. You'll also find there how to have an editable install of the package for local development.
|
||||
|
||||
|
||||
## Prohibited use.
|
||||
|
||||
Use of our model must comply with all applicable laws and regulations and must not result in, involve, or facilitate any illegal, harmful, deceptive, fraudulent, or unauthorized activity. Prohibited uses include, without limitation, voice impersonation or cloning without explicit and lawful consent; misinformation, disinformation, or deception (including fake news, fraudulent calls, or presenting generated content as genuine recordings of real people or events); and the generation of unlawful, harmful, libelous, abusive, harassing, discriminatory, hateful, or privacy-invasive content. We disclaim all liability for any non-compliant use.
|
||||
@@ -0,0 +1,83 @@
|
||||
# Generate Command Documentation
|
||||
|
||||
The `generate` command allows you to generate speech from text directly from the command line using Kyutai Pocket TTS.
|
||||
|
||||
## Basic Usage
|
||||
|
||||
```bash
|
||||
uvx pocket-tts generate
|
||||
# or if installed manually:
|
||||
pocket-tts generate
|
||||
```
|
||||
|
||||
This will generate a WAV file `./tts_output.wav` with the default text and voice.
|
||||
|
||||
## Command Options
|
||||
|
||||
### Core Options
|
||||
|
||||
- `--text TEXT`: Text to generate (default: "Hello world! I am Kyutai Pocket TTS. I'm fast enough to run on small CPUs. I hope you'll like me.")
|
||||
- `--voice VOICE`: Path to audio conditioning file (voice to clone) (default: "hf://kyutai/tts-voices/alba-mackenna/casual.wav"). Urls and local paths are supported.
|
||||
- `--output-path OUTPUT_PATH`: Output path for generated audio (default: "./tts_output.wav")
|
||||
|
||||
### Generation Parameters
|
||||
|
||||
- `--variant VARIANT`: Model signature (default: "b6369a24")
|
||||
- `--lsd-decode-steps LSD_DECODE_STEPS`: Number of generation steps (default: 1)
|
||||
- `--temperature TEMPERATURE`: Temperature for generation (default: 0.7)
|
||||
- `--noise-clamp NOISE_CLAMP`: Noise clamp value (default: None)
|
||||
- `--eos-threshold EOS_THRESHOLD`: EOS threshold (default: -4.0)
|
||||
- `--frames-after-eos FRAMES_AFTER_EOS`: Number of frames to generate after EOS (default: None, auto-calculated based on the text length). Each frame is 80ms.
|
||||
|
||||
### Performance Options
|
||||
|
||||
- `--device DEVICE`: Device to use (default: "cpu", you may not get a speedup by using a gpu since it's a small model)
|
||||
- `--quiet`, `-q`: Disable logging output
|
||||
|
||||
## Examples
|
||||
|
||||
### Basic Generation
|
||||
|
||||
```bash
|
||||
# Generate with default settings
|
||||
pocket-tts generate
|
||||
|
||||
# Custom text
|
||||
pocket-tts generate --text "Hello, this is a custom message."
|
||||
|
||||
# Custom output path
|
||||
pocket-tts generate --output-path "./my_audio.wav"
|
||||
```
|
||||
|
||||
### Voice Selection
|
||||
|
||||
```bash
|
||||
# Use different voice from HuggingFace
|
||||
pocket-tts generate --voice "hf://kyutai/tts-voices/jessica-jian/casual.wav"
|
||||
|
||||
# Use local voice file
|
||||
pocket-tts generate --voice "./my_voice.wav"
|
||||
```
|
||||
|
||||
### Quality Tuning
|
||||
|
||||
```bash
|
||||
# Higher quality (more steps)
|
||||
pocket-tts generate --lsd-decode-steps 5 --temperature 0.5
|
||||
|
||||
# More expressive (higher temperature)
|
||||
pocket-tts generate --temperature 1.0
|
||||
|
||||
# Adjust EOS threshold, smaller means finishing earlier.
|
||||
pocket-tts generate --eos-threshold -3.0
|
||||
```
|
||||
|
||||
## Output Format
|
||||
|
||||
The generate command always outputs WAV files in the following format:
|
||||
- **Sample Rate**: 24kHz
|
||||
- **Channels**: Mono
|
||||
- **Bit Depth**: 16-bit PCM
|
||||
- **Format**: Standard WAV file
|
||||
|
||||
For more advanced usage, see the [Python API documentation](python-api.md) or consider using the [serve command](serve.md) for web-based generation and quick iteration.
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 175 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 206 KiB |
@@ -0,0 +1,223 @@
|
||||
# Python API Documentation
|
||||
|
||||
Kyutai Pocket TTS provides a Python API for integrating text-to-speech capabilities into your applications.
|
||||
|
||||
## Installation
|
||||
|
||||
```bash
|
||||
pip install pocket-tts
|
||||
```
|
||||
|
||||
## Quick Start
|
||||
|
||||
```python
|
||||
from pocket_tts import TTSModel
|
||||
import scipy.io.wavfile
|
||||
|
||||
# Load the model
|
||||
tts_model = TTSModel.load_model()
|
||||
|
||||
# Get voice state from an audio file
|
||||
voice_state = tts_model.get_state_for_audio_prompt(
|
||||
"hf://kyutai/tts-voices/alba-mackenna/casual.wav"
|
||||
)
|
||||
|
||||
# Generate audio
|
||||
audio = tts_model.generate_audio(voice_state, "Hello world, this is a test.")
|
||||
|
||||
# Save to file
|
||||
scipy.io.wavfile.write("output.wav", tts_model.sample_rate, audio.numpy())
|
||||
```
|
||||
|
||||
## Core Classes
|
||||
|
||||
### TTSModel
|
||||
|
||||
The main class for text-to-speech generation.
|
||||
|
||||
#### Class Methods
|
||||
|
||||
##### `load_model(variant="b6369a24", temp=0.7, lsd_decode_steps=1, noise_clamp=None, eos_threshold=-4.0)`
|
||||
|
||||
Load and return a TTSModel instance with pre-trained weights.
|
||||
|
||||
**Parameters:**
|
||||
- `variant` (str): Model variant identifier (default: "b6369a24")
|
||||
- `temp` (float): Sampling temperature for generation (default: 0.7)
|
||||
- `lsd_decode_steps` (int): Number of generation steps (default: 1)
|
||||
- `noise_clamp` (float | None): Maximum value for noise sampling (default: None)
|
||||
- `eos_threshold` (float): Threshold for end-of-sequence detection (default: -4.0)
|
||||
|
||||
**Returns:**
|
||||
- `TTSModel`: Loaded model instance on CPU
|
||||
|
||||
**Example:**
|
||||
```python
|
||||
from pocket_tts import TTSModel
|
||||
|
||||
# Load with default settings
|
||||
model = TTSModel.load_model()
|
||||
|
||||
# Load with custom parameters
|
||||
model = TTSModel.load_model(variant="b6369a24", temp=0.5, lsd_decode_steps=5, eos_threshold=-3.0)
|
||||
```
|
||||
|
||||
#### Properties
|
||||
|
||||
##### `device` (str)
|
||||
|
||||
Returns the device type where the model is running ("cpu" or "cuda").
|
||||
By default, the model runs on CPU.
|
||||
|
||||
```python
|
||||
from pocket_tts import TTSModel
|
||||
|
||||
model = TTSModel.load_model()
|
||||
print(f"Model running on: {model.device}")
|
||||
```
|
||||
|
||||
##### `sample_rate` (int)
|
||||
|
||||
Returns the generated audio sample rate (typically 24000 Hz).
|
||||
|
||||
```python
|
||||
from pocket_tts import TTSModel
|
||||
|
||||
model = TTSModel.load_model()
|
||||
print(f"Sample rate: {model.sample_rate} Hz")
|
||||
```
|
||||
|
||||
#### Methods
|
||||
|
||||
##### `get_state_for_audio_prompt(audio_conditioning, truncate=False)`
|
||||
|
||||
Extract model state for a given audio file or URL (voice cloning).
|
||||
|
||||
**Parameters:**
|
||||
- `audio_conditioning` (Path | str | torch.Tensor): Audio file path, URL, or tensor
|
||||
- `truncate` (bool): Whether to truncate the audio (default: False)
|
||||
|
||||
**Returns:**
|
||||
- `dict`: Model state dictionary containing hidden states and positional information
|
||||
|
||||
**Example:**
|
||||
```python
|
||||
from pocket_tts import TTSModel
|
||||
|
||||
model = TTSModel.load_model()
|
||||
# From HuggingFace URL
|
||||
voice_state = model.get_state_for_audio_prompt("hf://kyutai/tts-voices/alba-mackenna/casual.wav")
|
||||
|
||||
# From local file
|
||||
voice_state = model.get_state_for_audio_prompt("./my_voice.wav")
|
||||
|
||||
# From HTTP URL
|
||||
voice_state = model.get_state_for_audio_prompt(
|
||||
"https://huggingface.co/kyutai/tts-voices/resolve"
|
||||
"/main/expresso/ex01-ex02_default_001_channel1_168s.wav"
|
||||
)
|
||||
```
|
||||
|
||||
##### `generate_audio(model_state, text_to_generate, frames_after_eos=None, copy_state=True)`
|
||||
|
||||
Generate complete audio tensor from text input.
|
||||
|
||||
**Parameters:**
|
||||
- `model_state` (dict): Model state from `get_state_for_audio_prompt()`
|
||||
- `text_to_generate` (str): Text to convert to speech
|
||||
- `frames_after_eos` (int | None): Frames to generate after EOS detection (default: None)
|
||||
- `copy_state` (bool): Whether to copy the state (default: True)
|
||||
|
||||
**Returns:**
|
||||
- `torch.Tensor`: Audio 1D tensor with shape [samples]
|
||||
|
||||
**Example:**
|
||||
```python
|
||||
from pocket_tts import TTSModel
|
||||
|
||||
model = TTSModel.load_model()
|
||||
|
||||
voice_state = model.get_state_for_audio_prompt("hf://kyutai/tts-voices/alba-mackenna/casual.wav")
|
||||
|
||||
# Generate audio
|
||||
audio = model.generate_audio(voice_state, "Hello world!", frames_after_eos=2, copy_state=True)
|
||||
|
||||
print(f"Generated audio shape: {audio.shape}")
|
||||
print(f"Audio duration: {audio.shape[-1] / model.sample_rate:.2f} seconds")
|
||||
```
|
||||
|
||||
##### `generate_audio_stream(model_state, text_to_generate, frames_after_eos=None, copy_state=True)`
|
||||
|
||||
Generate audio streaming chunks from text input.
|
||||
|
||||
**Parameters:** Same as `generate_audio()`
|
||||
|
||||
**Yields:**
|
||||
- `torch.Tensor`: Audio chunks with shape [samples]
|
||||
|
||||
**Example:**
|
||||
```python
|
||||
from pocket_tts import TTSModel
|
||||
|
||||
model = TTSModel.load_model()
|
||||
|
||||
voice_state = model.get_state_for_audio_prompt("hf://kyutai/tts-voices/alba-mackenna/casual.wav")
|
||||
# Stream generation
|
||||
for chunk in model.generate_audio_stream(voice_state, "Long text content..."):
|
||||
# Process each chunk as it's generated
|
||||
print(f"Generated chunk: {chunk.shape[0]} samples")
|
||||
# Could save chunks to file or play in real-time
|
||||
```
|
||||
|
||||
## Advanced Usage
|
||||
|
||||
### Voice Management
|
||||
|
||||
```python
|
||||
from pocket_tts import TTSModel
|
||||
|
||||
model = TTSModel.load_model()
|
||||
# Preload multiple voices
|
||||
voices = {
|
||||
"casual": model.get_state_for_audio_prompt("hf://kyutai/tts-voices/alba-mackenna/casual.wav"),
|
||||
"funny": model.get_state_for_audio_prompt(
|
||||
"https://huggingface.co/kyutai/tts-voices/resolve/main/expresso/ex01-ex02_default_001_channel1_168s.wav"
|
||||
),
|
||||
}
|
||||
|
||||
# Generate with different voices
|
||||
casual_audio = model.generate_audio(voices["casual"], "Hey there!")
|
||||
funny_audio = model.generate_audio(voices["funny"], "Good morning.")
|
||||
```
|
||||
|
||||
### Batch Processing
|
||||
|
||||
```python
|
||||
from pocket_tts import TTSModel
|
||||
import scipy.io.wavfile
|
||||
import torch
|
||||
|
||||
model = TTSModel.load_model()
|
||||
|
||||
voice_state = model.get_state_for_audio_prompt("hf://kyutai/tts-voices/alba-mackenna/casual.wav")
|
||||
# Process multiple texts efficiently by re-using the same voice state
|
||||
texts = [
|
||||
"First sentence to generate.",
|
||||
"Second sentence to generate.",
|
||||
"Third sentence to generate.",
|
||||
]
|
||||
|
||||
audios = []
|
||||
for text in texts:
|
||||
audio = model.generate_audio(voice_state, text)
|
||||
audios.append(audio)
|
||||
|
||||
# Concatenate all audio
|
||||
full_audio = torch.cat(audios, dim=0)
|
||||
scipy.io.wavfile.write("batch_output.wav", model.sample_rate, full_audio.numpy())
|
||||
```
|
||||
|
||||
### Streaming to File
|
||||
You can refer to our CLI implementation which can stream audio to a wav file.
|
||||
|
||||
For more information about the command-line interface, see the [Generate Documentation](generate.md) or [Serve Documentation](serve.md).
|
||||
@@ -0,0 +1,48 @@
|
||||
# Serve Command Documentation
|
||||
|
||||
The `serve` command starts a FastAPI web server that provides both a web interface and HTTP API for text-to-speech generation.
|
||||
|
||||
## Basic Usage
|
||||
|
||||
```bash
|
||||
uvx pocket-tts serve
|
||||
# or if installed manually:
|
||||
pocket-tts serve
|
||||
```
|
||||
|
||||
This starts a server on `http://localhost:8000` with the default voice model.
|
||||
|
||||
## Command Options
|
||||
|
||||
- `--voice VOICE`: Path to voice prompt audio file (voice to clone) (default: "hf://kyutai/tts-voices/alba-mackenna/casual.wav")
|
||||
- `--host HOST`: Host to bind to (default: "localhost")
|
||||
- `--port PORT`: Port to bind to (default: 8000)
|
||||
- `--reload`: Enable auto-reload for development
|
||||
|
||||
## Examples
|
||||
|
||||
### Basic Server
|
||||
|
||||
```bash
|
||||
# Start with default settings
|
||||
pocket-tts serve
|
||||
|
||||
# Custom host and port
|
||||
pocket-tts serve --host "localhost" --port 8080
|
||||
```
|
||||
|
||||
### Custom Voice
|
||||
|
||||
```bash
|
||||
# Use different voice
|
||||
pocket-tts serve --default-voice "hf://kyutai/tts-voices/jessica-jian/casual.wav"
|
||||
|
||||
# Use local voice file
|
||||
pocket-tts serve --default-voice "./my_voice.wav"
|
||||
```
|
||||
|
||||
## Web Interface
|
||||
|
||||
Once the server is running, navigate to `http://localhost:8000` to access the web interface.
|
||||
|
||||
For more advanced usage, see the [Python API documentation](python-api.md) for direct integration with the TTS model.
|
||||
@@ -0,0 +1,16 @@
|
||||
from beartype import BeartypeConf
|
||||
from beartype.claw import beartype_this_package
|
||||
|
||||
beartype_this_package(conf=BeartypeConf(is_color=False))
|
||||
|
||||
from pocket_tts.models.tts_model import TTSModel # noqa: E402
|
||||
|
||||
# Public methods:
|
||||
# TTSModel.device
|
||||
# TTSModel.sample_rate
|
||||
# TTSModel.load_model
|
||||
# TTSModel.generate_audio
|
||||
# TTSModel.generate_audio_stream
|
||||
# TTSModel.get_state_for_audio_prompt
|
||||
|
||||
__all__ = ["TTSModel"]
|
||||
@@ -0,0 +1,6 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
from pocket_tts.main import cli_app
|
||||
|
||||
if __name__ == "__main__":
|
||||
cli_app()
|
||||
@@ -0,0 +1,38 @@
|
||||
import logging
|
||||
from typing import Generic, NamedTuple, TypeVar
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
Prepared = TypeVar("Prepared") # represents the prepared condition input type.
|
||||
|
||||
|
||||
class TokenizedText(NamedTuple):
|
||||
tokens: torch.Tensor # should be long tensor.
|
||||
|
||||
|
||||
class BaseConditioner(nn.Module, Generic[Prepared]):
|
||||
"""Base model for all conditioner modules.
|
||||
|
||||
Args:
|
||||
dim (int): internal dim of the model.
|
||||
output_dim (int): Output dim of the conditioner.
|
||||
force_linear (bool, optional): Force linear projection even when `dim == output_dim`.
|
||||
output_bias (bool): if True, the output projection will have a bias.
|
||||
learn_padding (bool): if True, the padding value will be learnt, zero otherwise.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, dim: int, output_dim: int, output_bias: bool = False, force_linear: bool = True
|
||||
):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.output_dim = output_dim
|
||||
assert force_linear or dim != output_dim
|
||||
assert not output_bias
|
||||
|
||||
def forward(self, inputs: TokenizedText) -> torch.Tensor:
|
||||
return self._get_condition(inputs)
|
||||
@@ -0,0 +1,61 @@
|
||||
import logging
|
||||
|
||||
import sentencepiece
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from pocket_tts.conditioners.base import BaseConditioner, TokenizedText
|
||||
from pocket_tts.utils.utils import download_if_necessary
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class SentencePieceTokenizer:
|
||||
"""This tokenizer should be used for natural language descriptions.
|
||||
For example:
|
||||
["he didn't, know he's going home.", 'shorter sentence'] =>
|
||||
[[78, 62, 31, 4, 78, 25, 19, 34],
|
||||
[59, 77, PAD, PAD, PAD, PAD, PAD, PAD]]
|
||||
|
||||
Args:
|
||||
n_bins (int): should be equal to the number of elements in the sentencepiece tokenizer.
|
||||
tokenizer_path (str): path to the sentencepiece tokenizer model.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, nbins: int, tokenizer_path: str) -> None:
|
||||
logger.info("Loading sentencepiece tokenizer from %s", tokenizer_path)
|
||||
tokenizer_path = download_if_necessary(tokenizer_path)
|
||||
self.sp = sentencepiece.SentencePieceProcessor(str(tokenizer_path))
|
||||
assert nbins == self.sp.vocab_size(), (
|
||||
f"sentencepiece tokenizer has vocab size={self.sp.vocab_size()} but nbins={nbins} was specified"
|
||||
)
|
||||
|
||||
def __call__(self, text: str) -> TokenizedText:
|
||||
return TokenizedText(torch.tensor(self.sp.encode(text, out_type=int))[None, :])
|
||||
|
||||
|
||||
class LUTConditioner(BaseConditioner):
|
||||
"""Lookup table TextConditioner.
|
||||
|
||||
Args:
|
||||
n_bins (int): Number of bins.
|
||||
dim (int): Hidden dim of the model (text-encoder/LUT).
|
||||
output_dim (int): Output dim of the conditioner.
|
||||
tokenizer (str): Name of the tokenizer.
|
||||
possible_values (list[str] or None): list of possible values for the tokenizer.
|
||||
"""
|
||||
|
||||
def __init__(self, n_bins: int, tokenizer_path: str, dim: int, output_dim: int):
|
||||
super().__init__(dim=dim, output_dim=output_dim)
|
||||
self.tokenizer = SentencePieceTokenizer(n_bins, tokenizer_path)
|
||||
self.embed = nn.Embedding(n_bins + 1, self.dim) # n_bins + 1 for padding.
|
||||
|
||||
def prepare(self, x: str) -> TokenizedText:
|
||||
tokens = self.tokenizer(x)
|
||||
tokens = tokens[0].to(self.embed.weight.device)
|
||||
return TokenizedText(tokens)
|
||||
|
||||
def _get_condition(self, inputs: TokenizedText) -> torch.Tensor:
|
||||
embeds = self.embed(inputs[0])
|
||||
return embeds
|
||||
@@ -0,0 +1,57 @@
|
||||
# sig: b6369a24
|
||||
|
||||
weights_path: hf://kyutai/pocket-tts/tts_b6369a24.safetensors
|
||||
weights_path_without_voice_cloning: hf://kyutai/pocket-tts-without-voice-cloning/tts_b6369a24.safetensors
|
||||
|
||||
flow_lm:
|
||||
dtype: float32
|
||||
flow:
|
||||
depth: 6
|
||||
dim: 512
|
||||
transformer:
|
||||
d_model: 1024
|
||||
hidden_scale: 4
|
||||
max_period: 10000
|
||||
num_heads: 16
|
||||
num_layers: 6
|
||||
lookup_table:
|
||||
dim: 1024
|
||||
n_bins: 4000
|
||||
tokenizer: sentencepiece
|
||||
tokenizer_path: hf://kyutai/pocket-tts-without-voice-cloning/tokenizer.model
|
||||
#weights_path: flow_lm_b6369a24.safetensors
|
||||
|
||||
mimi:
|
||||
dtype: float32
|
||||
sample_rate: 24000
|
||||
channels: 1
|
||||
frame_rate: 12.5
|
||||
seanet:
|
||||
dimension: 512
|
||||
channels: 1
|
||||
n_filters: 64
|
||||
n_residual_layers: 1
|
||||
ratios:
|
||||
- 6
|
||||
- 5
|
||||
- 4
|
||||
kernel_size: 7
|
||||
residual_kernel_size: 3
|
||||
last_kernel_size: 3
|
||||
dilation_base: 2
|
||||
pad_mode: constant
|
||||
compress: 2
|
||||
transformer:
|
||||
d_model: 512
|
||||
num_heads: 8
|
||||
num_layers: 2
|
||||
layer_scale: 0.01
|
||||
context: 250
|
||||
dim_feedforward: 2048
|
||||
input_dimension: 512
|
||||
output_dimensions:
|
||||
- 512
|
||||
quantizer:
|
||||
dimension: 32
|
||||
output_dimension: 512
|
||||
#weights_path: mimi_b6369a24.safetensors
|
||||
@@ -0,0 +1,2 @@
|
||||
"""Audio loading and writing support. Datasets for raw audio
|
||||
or also including some metadata."""
|
||||
@@ -0,0 +1,126 @@
|
||||
"""
|
||||
Audio IO methods are defined in this module (info, read, write),
|
||||
We rely on av library for faster read when possible, otherwise on torchaudio.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import wave
|
||||
from contextlib import nullcontext
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from beartype.typing import Iterator
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
FIRST_CHUNK_LENGTH_SECONDS = float(os.environ.get("FIRST_CHUNK_LENGTH_SECONDS", "0"))
|
||||
|
||||
|
||||
def audio_read(filepath: str | Path) -> tuple[torch.Tensor, int]:
|
||||
"""Read audio using Python's wave module."""
|
||||
with wave.open(str(filepath), "rb") as wav_file:
|
||||
sample_rate = wav_file.getframerate()
|
||||
|
||||
# Read all audio data as 16-bit signed integers
|
||||
raw_data = wav_file.readframes(-1)
|
||||
samples = np.frombuffer(raw_data, dtype=np.int16).astype(np.float32) / 32768.0
|
||||
|
||||
# Return as mono tensor (channels, samples)
|
||||
wav = torch.from_numpy(samples.reshape(1, -1))
|
||||
return wav, sample_rate
|
||||
|
||||
|
||||
class StreamingWAVWriter:
|
||||
"""WAV writer using Python's standard library wave module."""
|
||||
|
||||
def __init__(self, output_stream, sample_rate: int):
|
||||
self.output_stream = output_stream
|
||||
self.sample_rate = sample_rate
|
||||
self.wave_writer = None
|
||||
self.first_chunk_buffer = []
|
||||
|
||||
def write_header(self, sample_rate: int):
|
||||
"""Initialize WAV writer with header."""
|
||||
# For stdout streaming, we need to handle the unseekable stream case
|
||||
# The wave module supports unseekable streams since Python 3.4
|
||||
self.wave_writer = wave.open(self.output_stream, "wb")
|
||||
self.wave_writer.setnchannels(1) # Mono
|
||||
self.wave_writer.setsampwidth(2) # 16-bit
|
||||
self.wave_writer.setframerate(sample_rate)
|
||||
self.wave_writer.setnframes(1_000_000_000)
|
||||
|
||||
def write_pcm_data(self, audio_chunk: torch.Tensor):
|
||||
"""Write PCM data using wave module."""
|
||||
# Convert to int16 PCM bytes
|
||||
chunk_int16 = (audio_chunk.clamp(-1, 1) * 32767).short()
|
||||
chunk_bytes = chunk_int16.detach().cpu().numpy().tobytes()
|
||||
|
||||
if self.first_chunk_buffer is not None:
|
||||
self.first_chunk_buffer.append(chunk_bytes)
|
||||
total_length = sum(len(c) for c in self.first_chunk_buffer)
|
||||
target_length = (
|
||||
int(self.sample_rate * FIRST_CHUNK_LENGTH_SECONDS) * 2
|
||||
) # 2 bytes per sample
|
||||
if total_length < target_length:
|
||||
return
|
||||
self._flush()
|
||||
return
|
||||
|
||||
# Use writeframesraw to avoid frame count validation for streaming
|
||||
self.wave_writer.writeframesraw(chunk_bytes)
|
||||
|
||||
def _flush(self):
|
||||
if self.first_chunk_buffer is not None:
|
||||
self.wave_writer.writeframesraw(b"".join(self.first_chunk_buffer))
|
||||
self.first_chunk_buffer = None
|
||||
|
||||
def finalize(self):
|
||||
"""Close the wave writer."""
|
||||
self._flush()
|
||||
|
||||
# Let's add 200ms of silence to ensure proper playback
|
||||
silence_duration_sec = 0.2
|
||||
num_silence_samples = int(self.sample_rate * silence_duration_sec)
|
||||
|
||||
self.wave_writer.writeframesraw(bytes(num_silence_samples * 2))
|
||||
|
||||
if self.wave_writer:
|
||||
# do not update the header for unseekable streams
|
||||
self.wave_writer._patchheader = lambda: None
|
||||
self.wave_writer.close()
|
||||
|
||||
|
||||
def is_file_like(obj):
|
||||
"""Check if object has basic file-like methods."""
|
||||
return all(hasattr(obj, attr) for attr in ["write", "close"])
|
||||
|
||||
|
||||
def stream_audio_chunks(
|
||||
path: str | Path | None | Any, audio_chunks: Iterator[torch.Tensor], sample_rate: int
|
||||
):
|
||||
"""Stream audio chunks to a WAV file or stdout, optionally playing them."""
|
||||
if path == "-":
|
||||
f = sys.stdout.buffer
|
||||
elif path is None:
|
||||
f = nullcontext()
|
||||
elif is_file_like(path):
|
||||
f = path
|
||||
else:
|
||||
f = open(path, "wb")
|
||||
|
||||
with f:
|
||||
if path is not None:
|
||||
writer = StreamingWAVWriter(f, sample_rate)
|
||||
writer.write_header(sample_rate)
|
||||
|
||||
for chunk in audio_chunks:
|
||||
# Then write to file
|
||||
if path is not None:
|
||||
writer.write_pcm_data(chunk)
|
||||
|
||||
if path is not None:
|
||||
writer.finalize()
|
||||
@@ -0,0 +1,28 @@
|
||||
"""Various utilities for audio convertion (pcm format, sample rate and channels),
|
||||
and volume normalization."""
|
||||
|
||||
import torch
|
||||
from scipy.signal import resample_poly
|
||||
|
||||
|
||||
def convert_audio(
|
||||
wav: torch.Tensor, from_rate: int | float, to_rate: int | float, to_channels: int
|
||||
) -> torch.Tensor:
|
||||
"""Convert audio to new sample rate and number of audio channels."""
|
||||
if from_rate != to_rate:
|
||||
# Convert to numpy for scipy resampling
|
||||
wav_np = wav.detach().cpu().numpy()
|
||||
|
||||
# Calculate resampling parameters
|
||||
gcd = int(torch.gcd(torch.tensor(from_rate), torch.tensor(to_rate)).item())
|
||||
up = int(to_rate // gcd)
|
||||
down = int(from_rate // gcd)
|
||||
|
||||
# Resample using scipy
|
||||
resampled_np = resample_poly(wav_np, up, down, axis=-1)
|
||||
|
||||
# Convert back to torch tensor
|
||||
wav = torch.from_numpy(resampled_np).to(wav.device).to(wav.dtype)
|
||||
|
||||
assert wav.shape[-2] == to_channels
|
||||
return wav
|
||||
@@ -0,0 +1,7 @@
|
||||
DEFAULT_AUDIO_PROMPT = "marius"
|
||||
DEFAULT_VARIANT = "b6369a24"
|
||||
DEFAULT_TEMPERATURE = 0.7
|
||||
DEFAULT_LSD_DECODE_STEPS = 1
|
||||
DEFAULT_NOISE_CLAMP = None
|
||||
DEFAULT_EOS_THRESHOLD = -4.0
|
||||
DEFAULT_FRAMES_AFTER_EOS = None
|
||||
@@ -0,0 +1,254 @@
|
||||
import io
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from queue import Queue
|
||||
|
||||
import typer
|
||||
import uvicorn
|
||||
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import FileResponse, StreamingResponse
|
||||
from typing_extensions import Annotated
|
||||
|
||||
from pocket_tts.data.audio import stream_audio_chunks
|
||||
from pocket_tts.default_parameters import (
|
||||
DEFAULT_AUDIO_PROMPT,
|
||||
DEFAULT_EOS_THRESHOLD,
|
||||
DEFAULT_FRAMES_AFTER_EOS,
|
||||
DEFAULT_LSD_DECODE_STEPS,
|
||||
DEFAULT_NOISE_CLAMP,
|
||||
DEFAULT_TEMPERATURE,
|
||||
DEFAULT_VARIANT,
|
||||
)
|
||||
from pocket_tts.models.tts_model import TTSModel
|
||||
from pocket_tts.utils.logging_utils import enable_logging
|
||||
from pocket_tts.utils.utils import PREDEFINED_VOICES, size_of_dict
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
cli_app = typer.Typer(
|
||||
help="Kyutai Pocket TTS - Text-to-Speech generation tool", pretty_exceptions_show_locals=False
|
||||
)
|
||||
|
||||
|
||||
# ------------------------------------------------------
|
||||
# The pocket-tts server implementation
|
||||
# ------------------------------------------------------
|
||||
|
||||
# Global model instance
|
||||
tts_model = None
|
||||
global_model_state = None
|
||||
|
||||
web_app = FastAPI(
|
||||
title="Kyutai Pocket TTS API", description="Text-to-Speech generation API", version="1.0.0"
|
||||
)
|
||||
web_app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=[
|
||||
"http://localhost:3000",
|
||||
"https://pod1-10007.internal.kyutai.org",
|
||||
"https://kyutai.org",
|
||||
],
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
|
||||
@web_app.get("/")
|
||||
async def root():
|
||||
"""Serve the frontend."""
|
||||
static_path = Path(__file__).parent / "static" / "index.html"
|
||||
return FileResponse(static_path)
|
||||
|
||||
|
||||
@web_app.get("/health")
|
||||
async def health():
|
||||
return {"status": "healthy"}
|
||||
|
||||
|
||||
def write_to_queue(queue, text_to_generate, model_state):
|
||||
"""Allows writing to the StreamingResponse as if it were a file."""
|
||||
|
||||
class FileLikeToQueue(io.IOBase):
|
||||
def __init__(self, queue):
|
||||
self.queue = queue
|
||||
|
||||
def write(self, data):
|
||||
self.queue.put(data)
|
||||
|
||||
def flush(self):
|
||||
pass
|
||||
|
||||
def close(self):
|
||||
self.queue.put(None)
|
||||
|
||||
audio_chunks = tts_model.generate_audio_stream(
|
||||
model_state=model_state, text_to_generate=text_to_generate
|
||||
)
|
||||
stream_audio_chunks(FileLikeToQueue(queue), audio_chunks, tts_model.config.mimi.sample_rate)
|
||||
|
||||
|
||||
def generate_data_with_state(text_to_generate: str, model_state: dict):
|
||||
queue = Queue()
|
||||
|
||||
# Run your function in a thread
|
||||
thread = threading.Thread(target=write_to_queue, args=(queue, text_to_generate, model_state))
|
||||
thread.start()
|
||||
|
||||
# Yield data as it becomes available
|
||||
i = 0
|
||||
while True:
|
||||
data = queue.get()
|
||||
if data is None:
|
||||
break
|
||||
i += 1
|
||||
yield data
|
||||
|
||||
thread.join()
|
||||
|
||||
|
||||
@web_app.post("/tts")
|
||||
def text_to_speech(
|
||||
text: str = Form(...),
|
||||
voice_url: str | None = Form(None),
|
||||
voice_wav: UploadFile | None = File(None),
|
||||
):
|
||||
"""
|
||||
Generate speech from text using the pre-loaded voice prompt or a custom voice.
|
||||
|
||||
Args:
|
||||
text: Text to convert to speech
|
||||
voice_url: Optional voice URL (http://, https://, or hf://)
|
||||
voice_wav: Optional uploaded voice file (mutually exclusive with voice_url)
|
||||
"""
|
||||
if not text.strip():
|
||||
raise HTTPException(status_code=400, detail="Text cannot be empty")
|
||||
|
||||
if voice_url is not None and voice_wav is not None:
|
||||
raise HTTPException(status_code=400, detail="Cannot provide both voice_url and voice_wav")
|
||||
|
||||
# Use the appropriate model state
|
||||
if voice_url is not None:
|
||||
if not (
|
||||
voice_url.startswith("http://")
|
||||
or voice_url.startswith("https://")
|
||||
or voice_url.startswith("hf://")
|
||||
or voice_url in PREDEFINED_VOICES
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=400, detail="voice_url must start with http://, https://, or hf://"
|
||||
)
|
||||
model_state = tts_model._cached_get_state_for_audio_prompt(voice_url, truncate=True)
|
||||
logging.warning("Using voice from URL: %s", voice_url)
|
||||
elif voice_wav is not None:
|
||||
# Use uploaded voice file
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_file:
|
||||
content = voice_wav.file.read()
|
||||
temp_file.write(content)
|
||||
temp_file.flush()
|
||||
|
||||
try:
|
||||
model_state = tts_model.get_state_for_audio_prompt(
|
||||
Path(temp_file.name), truncate=True
|
||||
)
|
||||
finally:
|
||||
os.unlink(temp_file.name)
|
||||
else:
|
||||
# Use default global model state
|
||||
model_state = global_model_state
|
||||
|
||||
return StreamingResponse(
|
||||
generate_data_with_state(text, model_state),
|
||||
media_type="audio/wav",
|
||||
headers={
|
||||
"Content-Disposition": "attachment; filename=generated_speech.wav",
|
||||
"Transfer-Encoding": "chunked",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@cli_app.command()
|
||||
def serve(
|
||||
voice: Annotated[
|
||||
str, typer.Option(help="Path to voice prompt audio file (voice to clone)")
|
||||
] = DEFAULT_AUDIO_PROMPT,
|
||||
host: Annotated[str, typer.Option(help="Host to bind to")] = "localhost",
|
||||
port: Annotated[int, typer.Option(help="Port to bind to")] = 8000,
|
||||
reload: Annotated[bool, typer.Option(help="Enable auto-reload")] = False,
|
||||
):
|
||||
"""Start the FastAPI server."""
|
||||
|
||||
global tts_model, global_model_state
|
||||
tts_model = TTSModel.load_model(DEFAULT_VARIANT)
|
||||
|
||||
# Pre-load the voice prompt
|
||||
global_model_state = tts_model.get_state_for_audio_prompt(voice)
|
||||
logger.info(f"The size of the model state is {size_of_dict(global_model_state) // 1e6} MB")
|
||||
|
||||
uvicorn.run("pocket_tts.main:web_app", host=host, port=port, reload=reload)
|
||||
|
||||
|
||||
# ------------------------------------------------------
|
||||
# The pocket-tts single generation CLI implementation
|
||||
# ------------------------------------------------------
|
||||
|
||||
|
||||
@cli_app.command()
|
||||
def generate(
|
||||
text: Annotated[
|
||||
str, typer.Option(help="Text to generate")
|
||||
] = "Hello world. I am Kyutai's Pocket TTS. I'm fast enough to run on small CPUs. I hope you'll like me.",
|
||||
voice: Annotated[
|
||||
str, typer.Option(help="Path to audio conditioning file (voice to clone)")
|
||||
] = DEFAULT_AUDIO_PROMPT,
|
||||
quiet: Annotated[bool, typer.Option("-q", "--quiet", help="Disable logging output")] = False,
|
||||
variant: Annotated[str, typer.Option(help="Model signature")] = DEFAULT_VARIANT,
|
||||
lsd_decode_steps: Annotated[
|
||||
int, typer.Option(help="Number of generation steps")
|
||||
] = DEFAULT_LSD_DECODE_STEPS,
|
||||
temperature: Annotated[
|
||||
float, typer.Option(help="Temperature for generation")
|
||||
] = DEFAULT_TEMPERATURE,
|
||||
noise_clamp: Annotated[float, typer.Option(help="Noise clamp value")] = DEFAULT_NOISE_CLAMP,
|
||||
eos_threshold: Annotated[float, typer.Option(help="EOS threshold")] = DEFAULT_EOS_THRESHOLD,
|
||||
frames_after_eos: Annotated[
|
||||
int, typer.Option(help="Number of frames to generate after EOS")
|
||||
] = DEFAULT_FRAMES_AFTER_EOS,
|
||||
output_path: Annotated[
|
||||
str, typer.Option(help="Output path for generated audio")
|
||||
] = "./tts_output.wav",
|
||||
device: Annotated[str, typer.Option(help="Device to use")] = "cpu",
|
||||
):
|
||||
"""Generate speech using Kyutai Pocket TTS."""
|
||||
if "cuda" in device:
|
||||
# Cuda graphs capturing does not play nice with multithreading.
|
||||
os.environ["NO_CUDA_GRAPH"] = "1"
|
||||
|
||||
log_level = logging.ERROR if quiet else logging.INFO
|
||||
with enable_logging("pocket_tts", log_level):
|
||||
tts_model = TTSModel.load_model(
|
||||
variant, temperature, lsd_decode_steps, noise_clamp, eos_threshold
|
||||
)
|
||||
tts_model.to(device)
|
||||
|
||||
model_state_for_voice = tts_model.get_state_for_audio_prompt(voice)
|
||||
# Stream audio generation directly to file or stdout
|
||||
audio_chunks = tts_model.generate_audio_stream(
|
||||
model_state=model_state_for_voice,
|
||||
text_to_generate=text,
|
||||
frames_after_eos=frames_after_eos,
|
||||
)
|
||||
|
||||
stream_audio_chunks(output_path, audio_chunks, tts_model.config.mimi.sample_rate)
|
||||
|
||||
# Only print the result message if not writing to stdout
|
||||
if output_path != "-":
|
||||
logger.info("Results written in %s", output_path)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
cli_app()
|
||||
@@ -0,0 +1,3 @@
|
||||
"""
|
||||
Models for EnCodec (and Mimi), and FlowLMModel.
|
||||
"""
|
||||
@@ -0,0 +1,209 @@
|
||||
import logging
|
||||
from functools import partial
|
||||
|
||||
import torch
|
||||
import torch.export
|
||||
from beartype.typing import Callable
|
||||
from torch import nn
|
||||
from typing_extensions import Self
|
||||
|
||||
from pocket_tts.conditioners.text import LUTConditioner
|
||||
from pocket_tts.modules.mimi_transformer import StreamingTransformer
|
||||
from pocket_tts.modules.mlp import SimpleMLPAdaLN
|
||||
from pocket_tts.utils.config import FlowLMConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
FlowNet2 = Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor]
|
||||
|
||||
|
||||
def lsd_decode(v_t: FlowNet2, x_0: torch.Tensor, num_steps: int = 1) -> torch.Tensor:
|
||||
"""Rebuilds the data sample from starting point x_0.
|
||||
|
||||
Lagrangian Self Distillation (https://arxiv.org/pdf/2505.18825)
|
||||
|
||||
Args:
|
||||
v_t: Function taking t and x_t as input and returning the flow.
|
||||
x_0: Starting point from the known distribution.
|
||||
num_steps: Number of steps to take.
|
||||
|
||||
Returns:
|
||||
x_1_hat: (B, D) Reconstructed data sample.
|
||||
"""
|
||||
current = x_0
|
||||
for i in range(num_steps):
|
||||
s = i / num_steps
|
||||
t = (i + 1) / num_steps
|
||||
flow_dir = v_t(
|
||||
s * torch.ones_like(x_0[..., :1]), t * torch.ones_like(x_0[..., :1]), current
|
||||
)
|
||||
current += flow_dir / num_steps
|
||||
return current
|
||||
|
||||
|
||||
class FlowLMModel(nn.Module):
|
||||
"""Transformer-based flow language model on multiple streams of latents.
|
||||
|
||||
Args:
|
||||
conditioner (LUTConditioner): Text conditioner for processing text inputs.
|
||||
flow: Flow module that defines the flow loss and sampling strategy.
|
||||
flow_net: Trainable function (cond, t, x_t) -> u_t.
|
||||
dim (int): Dimension of the transformer encoder.
|
||||
norm (str): Normalization method.
|
||||
attribute_dropouts (dict): Attribute dropout probabilities.
|
||||
ldim (int): Latent dimension.
|
||||
stats_ema_decay (float): Decay for the EMA of the latent statistics.
|
||||
**kwargs: Additional parameters for the transformer encoder.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
conditioner: LUTConditioner,
|
||||
flow_net: SimpleMLPAdaLN,
|
||||
transformer: StreamingTransformer,
|
||||
dim: int = 128,
|
||||
ldim: int = 64,
|
||||
stats_ema_decay: float = 0.999,
|
||||
text_padding_weight: float = 1.0,
|
||||
dtype=None,
|
||||
):
|
||||
super().__init__()
|
||||
self.conditioner = conditioner
|
||||
self.ldim = ldim
|
||||
self.stats_ema_decay = stats_ema_decay
|
||||
self.dim = dim
|
||||
self.text_padding_weight = text_padding_weight
|
||||
self.dtype = dtype
|
||||
|
||||
self.flow_net = flow_net
|
||||
self.register_buffer("emb_std", torch.ones(ldim, dtype=dtype))
|
||||
self.register_buffer("emb_mean", torch.zeros(ldim, dtype=dtype))
|
||||
self.bos_emb = torch.nn.Parameter(torch.randn(ldim, dtype=dtype))
|
||||
|
||||
self.input_linear = nn.Linear(self.ldim, dim, bias=False, dtype=dtype)
|
||||
self.transformer = transformer
|
||||
self.out_norm = nn.LayerNorm(dim, eps=1e-5)
|
||||
self.out_eos = nn.Linear(dim, 1, dtype=dtype)
|
||||
|
||||
@property
|
||||
def device(self) -> str:
|
||||
return next(self.parameters()).device.type
|
||||
|
||||
def forward(
|
||||
self,
|
||||
sequence: torch.Tensor,
|
||||
text_embeddings: torch.Tensor,
|
||||
model_state: dict,
|
||||
lsd_decode_steps: int,
|
||||
temp: float,
|
||||
noise_clamp: float | None,
|
||||
eos_threshold: float,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Apply language model on sequence and conditions.
|
||||
Given a tensor of sequence of shape [B, S, ldim], returns the loss in training mode
|
||||
or the reconstructed latent in generation mode.
|
||||
|
||||
Args:
|
||||
sequence (torch.Tensor): Latents to model.
|
||||
text_embeddings (torch.Tensor): Pre-computed conditioning
|
||||
tensor.
|
||||
lsd_decode_steps (int): Number of steps to decode when generating audio.
|
||||
If zero, the model computes the loss.
|
||||
Returns:
|
||||
(output, eos_output, metrics). If `lsd_decode_steps` is zero, `output` is the loss tensor of shape [B, S],
|
||||
otherwise it is the reconstructed latent.
|
||||
"""
|
||||
# NaN values signal a BOS position.
|
||||
sequence = torch.where(torch.isnan(sequence), self.bos_emb, sequence)
|
||||
input_ = self.input_linear(sequence)
|
||||
|
||||
transformer_out = self.backbone(input_, text_embeddings, sequence, model_state=model_state)
|
||||
transformer_out = transformer_out.to(torch.float32)
|
||||
assert lsd_decode_steps > 0
|
||||
|
||||
transformer_out = transformer_out[:, -1]
|
||||
out_eos = self.out_eos(transformer_out) > eos_threshold
|
||||
|
||||
noise_shape = transformer_out.shape[:-1] + (self.ldim,)
|
||||
std = temp**0.5
|
||||
noise = torch.empty(noise_shape, dtype=transformer_out.dtype, device=transformer_out.device)
|
||||
if noise_clamp is None:
|
||||
torch.nn.init.normal_(noise, mean=0.0, std=std)
|
||||
else:
|
||||
torch.nn.init.trunc_normal_(noise, mean=0.0, std=std, a=-noise_clamp, b=noise_clamp)
|
||||
conditioned_flow = partial(self.flow_net, transformer_out)
|
||||
return lsd_decode(conditioned_flow, noise, lsd_decode_steps), out_eos
|
||||
|
||||
def backbone(
|
||||
self, input_, text_embeddings: torch.Tensor, sequence, model_state: dict
|
||||
) -> torch.Tensor:
|
||||
# Most of the time, one of those two tensors is empty, it allows us
|
||||
# to input text or audio embeddings into the model without adding an
|
||||
# if-else branch.
|
||||
# print("text_embeddings shape:", text_embeddings.shape)
|
||||
# if text_embeddings.numel() != 0:
|
||||
# torch.save(text_embeddings, "debug_flow_lm_text_embeddings.pt")
|
||||
input_ = torch.cat([text_embeddings, input_], dim=1)
|
||||
# transformer_out = self.transformer(input_, model_state=model_state)
|
||||
transformer_out = self.transformer(input_, model_state)
|
||||
if self.out_norm:
|
||||
transformer_out = self.out_norm(transformer_out)
|
||||
# remove the prefix from the model outputs (condition is prepended)
|
||||
transformer_out = transformer_out[:, -sequence.shape[1] :]
|
||||
return transformer_out
|
||||
|
||||
def _sample_next_latent(
|
||||
self,
|
||||
sequence: torch.Tensor,
|
||||
text_embeddings: torch.Tensor,
|
||||
model_state: dict,
|
||||
lsd_decode_steps: int,
|
||||
temp: float,
|
||||
noise_clamp: float | None,
|
||||
eos_threshold: float,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Sample next latent from the model given a sequence and a set of conditions.
|
||||
Args:
|
||||
sequence (torch.Tensor): Current sequence of shape [B, K, S]
|
||||
with K corresponding to the number of codebooks and S the number of sequence steps.
|
||||
S = 1 in streaming mode, except for the first step that contains a bigger prompt.
|
||||
text_embeddings (torch.Tensor): Condition tensor.
|
||||
n_steps (int): Number of flow steps to decode when generating audio.
|
||||
Returns:
|
||||
next_latent (torch.Tensor), is_eos (torch.Tensor): Next latent tensor of shape [B, 1, ldim]
|
||||
and is_eos tensor of shape [B, 1] with 1 on EOS positions.
|
||||
"""
|
||||
result = self(
|
||||
sequence=sequence,
|
||||
text_embeddings=text_embeddings,
|
||||
lsd_decode_steps=lsd_decode_steps,
|
||||
temp=temp,
|
||||
noise_clamp=noise_clamp,
|
||||
eos_threshold=eos_threshold,
|
||||
model_state=model_state,
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
@classmethod
|
||||
def from_pydantic_config(cls, config: FlowLMConfig, latent_dim: int) -> Self:
|
||||
d_model = config.transformer.d_model
|
||||
flow_mlp = SimpleMLPAdaLN.from_pydantic_config(config, latent_dim, d_model)
|
||||
|
||||
conditioner = LUTConditioner(
|
||||
n_bins=config.lookup_table.n_bins,
|
||||
tokenizer_path=str(config.lookup_table.tokenizer_path),
|
||||
dim=config.lookup_table.dim,
|
||||
output_dim=d_model,
|
||||
)
|
||||
|
||||
transformer = StreamingTransformer.from_pydantic_config(config.transformer)
|
||||
|
||||
return cls(
|
||||
flow_net=flow_mlp,
|
||||
transformer=transformer,
|
||||
dim=d_model,
|
||||
conditioner=conditioner,
|
||||
ldim=latent_dim,
|
||||
dtype=getattr(torch, config.dtype),
|
||||
)
|
||||
@@ -0,0 +1,111 @@
|
||||
import logging
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from pocket_tts.modules.conv import pad_for_conv1d
|
||||
from pocket_tts.modules.dummy_quantizer import DummyQuantizer
|
||||
from pocket_tts.modules.mimi_transformer import ProjectedTransformer
|
||||
from pocket_tts.modules.resample import ConvDownsample1d, ConvTrUpsample1d
|
||||
from pocket_tts.modules.seanet import SEANetDecoder, SEANetEncoder
|
||||
|
||||
logger = logging.getLogger()
|
||||
|
||||
|
||||
class MimiModel(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
encoder: SEANetEncoder,
|
||||
decoder: SEANetDecoder,
|
||||
quantizer: DummyQuantizer,
|
||||
frame_rate: float,
|
||||
encoder_frame_rate: float,
|
||||
sample_rate: int,
|
||||
channels: int,
|
||||
encoder_transformer: ProjectedTransformer,
|
||||
decoder_transformer: ProjectedTransformer,
|
||||
):
|
||||
super().__init__()
|
||||
self.encoder = encoder
|
||||
self.decoder = decoder
|
||||
self.encoder_transformer = encoder_transformer
|
||||
self.decoder_transformer = decoder_transformer
|
||||
self.quantizer = quantizer
|
||||
self.frame_rate = frame_rate
|
||||
self.sample_rate = sample_rate
|
||||
self.channels = channels
|
||||
self.encoder_frame_rate = encoder_frame_rate
|
||||
|
||||
# We will need the dimension for the resampling. In general the encoder will be a SeanetEncoder
|
||||
# which exposes a `dimension` attribute.
|
||||
dimension = encoder.dimension
|
||||
assert isinstance(dimension, int), (
|
||||
f"Dimension should be int, got {dimension} of type {type(dimension)}."
|
||||
)
|
||||
self.dimension = dimension
|
||||
|
||||
if encoder_frame_rate != frame_rate:
|
||||
assert self.encoder_frame_rate > self.frame_rate, "Cannot upsample with conv."
|
||||
downsample_stride = self.encoder_frame_rate / self.frame_rate
|
||||
assert downsample_stride == int(downsample_stride), (
|
||||
f"Only integer strides are supported, got {downsample_stride}"
|
||||
)
|
||||
self.downsample = ConvDownsample1d(int(downsample_stride), dimension=dimension)
|
||||
self.upsample = ConvTrUpsample1d(int(downsample_stride), dimension=dimension)
|
||||
|
||||
@property
|
||||
def frame_size(self) -> int:
|
||||
return int(self.sample_rate / self.frame_rate)
|
||||
|
||||
def _to_framerate(self, x: torch.Tensor):
|
||||
# Convert from the encoder frame rate to the overall framerate.
|
||||
_, _, length = x.shape
|
||||
frame_rate = self.encoder_frame_rate
|
||||
new_frame_rate = self.frame_rate
|
||||
if frame_rate == new_frame_rate:
|
||||
return x
|
||||
return self.downsample(x, model_state=None)
|
||||
|
||||
def _to_encoder_framerate(self, x: torch.Tensor, mimi_state) -> torch.Tensor:
|
||||
# Convert from overall framerate to the encoder frame rate.
|
||||
_, _, length = x.shape
|
||||
frame_rate = self.encoder_frame_rate
|
||||
new_frame_rate = self.frame_rate
|
||||
if frame_rate == new_frame_rate:
|
||||
return x
|
||||
return self.upsample(x, mimi_state)
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
raise NotImplementedError()
|
||||
|
||||
def decode_from_latent(self, latent: torch.Tensor, mimi_state) -> torch.Tensor:
|
||||
emb = self._to_encoder_framerate(latent, mimi_state)
|
||||
(emb,) = self.decoder_transformer(emb, mimi_state)
|
||||
out = self.decoder(emb, mimi_state)
|
||||
# out contains extra padding added by the encoder and decoder
|
||||
return out
|
||||
|
||||
def encode_to_latent(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""Projects a batch of waveforms to unquantized latent space.
|
||||
|
||||
Args:
|
||||
x (torch.Tensor): Float tensor of shape [B, C, T].
|
||||
|
||||
Returns:
|
||||
Unquantized embeddings.
|
||||
"""
|
||||
assert x.dim() == 3, (
|
||||
f"CompressionModel._encode_to_unquantized_latent expects audio of shape [B, C, T] but got {x.shape}"
|
||||
)
|
||||
|
||||
frame_size = self.frame_size
|
||||
|
||||
# The underlying convolutions no longer accept partial inputs,
|
||||
# `x` needs to be exactly a multiple of the frame size,
|
||||
# reproducing the previous padding behavior here.
|
||||
x = pad_for_conv1d(x, frame_size, frame_size)
|
||||
emb = self.encoder(x, model_state=None)
|
||||
|
||||
(emb,) = self.encoder_transformer(emb, model_state=None)
|
||||
emb = self._to_framerate(emb)
|
||||
return emb
|
||||
@@ -0,0 +1,719 @@
|
||||
import copy
|
||||
import logging
|
||||
import os
|
||||
import queue
|
||||
import statistics
|
||||
import threading
|
||||
import time
|
||||
from functools import lru_cache
|
||||
from pathlib import Path
|
||||
|
||||
import safetensors
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
from typing_extensions import Self
|
||||
|
||||
from pocket_tts.conditioners.base import TokenizedText
|
||||
from pocket_tts.data.audio import audio_read
|
||||
from pocket_tts.data.audio_utils import convert_audio
|
||||
from pocket_tts.default_parameters import (
|
||||
DEFAULT_EOS_THRESHOLD,
|
||||
DEFAULT_LSD_DECODE_STEPS,
|
||||
DEFAULT_NOISE_CLAMP,
|
||||
DEFAULT_TEMPERATURE,
|
||||
DEFAULT_VARIANT,
|
||||
)
|
||||
from pocket_tts.models.flow_lm import FlowLMModel
|
||||
from pocket_tts.models.mimi import MimiModel
|
||||
from pocket_tts.modules import mimi_transformer
|
||||
from pocket_tts.modules.dummy_quantizer import DummyQuantizer
|
||||
from pocket_tts.modules.seanet import SEANetDecoder, SEANetEncoder
|
||||
from pocket_tts.modules.stateful_module import increment_steps, init_states
|
||||
from pocket_tts.utils.config import Config, load_config
|
||||
from pocket_tts.utils.utils import (
|
||||
PREDEFINED_VOICES,
|
||||
display_execution_time,
|
||||
download_if_necessary,
|
||||
load_predefined_voice,
|
||||
size_of_dict,
|
||||
)
|
||||
from pocket_tts.utils.weights_loading import get_flow_lm_state_dict, get_mimi_state_dict
|
||||
|
||||
torch.set_num_threads(1)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
KEYS_TO_MODIFY = [
|
||||
"mimi.decoder.model.0.conv.conv.bias",
|
||||
"mimi.decoder.model.0.conv.conv.weight",
|
||||
"mimi.decoder.model.11.conv.conv.bias",
|
||||
"mimi.decoder.model.11.conv.conv.weight",
|
||||
"mimi.decoder.model.2.convtr.convtr.bias",
|
||||
"mimi.decoder.model.2.convtr.convtr.weight",
|
||||
"mimi.decoder.model.3.block.1.conv.conv.bias",
|
||||
"mimi.decoder.model.3.block.1.conv.conv.weight",
|
||||
"mimi.decoder.model.3.block.3.conv.conv.bias",
|
||||
"mimi.decoder.model.3.block.3.conv.conv.weight",
|
||||
"mimi.decoder.model.5.convtr.convtr.bias",
|
||||
"mimi.decoder.model.5.convtr.convtr.weight",
|
||||
"mimi.decoder.model.6.block.1.conv.conv.bias",
|
||||
"mimi.decoder.model.6.block.1.conv.conv.weight",
|
||||
"mimi.decoder.model.6.block.3.conv.conv.bias",
|
||||
"mimi.decoder.model.6.block.3.conv.conv.weight",
|
||||
"mimi.decoder.model.8.convtr.convtr.bias",
|
||||
"mimi.decoder.model.8.convtr.convtr.weight",
|
||||
"mimi.decoder.model.9.block.1.conv.conv.bias",
|
||||
"mimi.decoder.model.9.block.1.conv.conv.weight",
|
||||
"mimi.decoder.model.9.block.3.conv.conv.bias",
|
||||
"mimi.decoder.model.9.block.3.conv.conv.weight",
|
||||
"mimi.downsample.conv.conv.conv.weight",
|
||||
"mimi.encoder.model.0.conv.conv.bias",
|
||||
"mimi.encoder.model.0.conv.conv.weight",
|
||||
"mimi.encoder.model.1.block.1.conv.conv.bias",
|
||||
"mimi.encoder.model.1.block.1.conv.conv.weight",
|
||||
"mimi.encoder.model.1.block.3.conv.conv.bias",
|
||||
"mimi.encoder.model.1.block.3.conv.conv.weight",
|
||||
"mimi.encoder.model.11.conv.conv.bias",
|
||||
"mimi.encoder.model.11.conv.conv.weight",
|
||||
"mimi.encoder.model.3.conv.conv.bias",
|
||||
"mimi.encoder.model.3.conv.conv.weight",
|
||||
"mimi.encoder.model.4.block.1.conv.conv.bias",
|
||||
"mimi.encoder.model.4.block.1.conv.conv.weight",
|
||||
"mimi.encoder.model.4.block.3.conv.conv.bias",
|
||||
"mimi.encoder.model.4.block.3.conv.conv.weight",
|
||||
"mimi.encoder.model.6.conv.conv.bias",
|
||||
"mimi.encoder.model.6.conv.conv.weight",
|
||||
"mimi.encoder.model.7.block.1.conv.conv.bias",
|
||||
"mimi.encoder.model.7.block.1.conv.conv.weight",
|
||||
"mimi.encoder.model.7.block.3.conv.conv.bias",
|
||||
"mimi.encoder.model.7.block.3.conv.conv.weight",
|
||||
"mimi.encoder.model.9.conv.conv.bias",
|
||||
"mimi.encoder.model.9.conv.conv.weight",
|
||||
"mimi.upsample.convtr.convtr.convtr.weight",
|
||||
]
|
||||
|
||||
KEY_WITH_REMOVED_ZERO = [
|
||||
"mimi.encoder_transformer.transformer.layers.0.self_attn.in_projs.0.weight",
|
||||
"mimi.encoder_transformer.transformer.layers.0.self_attn.out_projs.0.weight",
|
||||
"mimi.encoder_transformer.transformer.layers.1.self_attn.in_projs.0.weight",
|
||||
"mimi.encoder_transformer.transformer.layers.1.self_attn.out_projs.0.weight",
|
||||
"mimi.decoder_transformer.transformer.layers.0.self_attn.in_projs.0.weight",
|
||||
"mimi.decoder_transformer.transformer.layers.0.self_attn.out_projs.0.weight",
|
||||
"mimi.decoder_transformer.transformer.layers.1.self_attn.in_projs.0.weight",
|
||||
"mimi.decoder_transformer.transformer.layers.1.self_attn.out_projs.0.weight",
|
||||
]
|
||||
|
||||
|
||||
def adapt_state_dict_of_TTSModel(state_dict: dict) -> dict:
|
||||
# should be removed before release
|
||||
new_state_dict = {}
|
||||
for key, value in state_dict.items():
|
||||
if key == "mimi.quantizer.input_proj.weight":
|
||||
continue
|
||||
if key in KEYS_TO_MODIFY:
|
||||
new_key = key.replace(".conv.conv.", ".conv.").replace(".convtr.convtr.", ".convtr.")
|
||||
new_state_dict[new_key] = value
|
||||
elif key in KEY_WITH_REMOVED_ZERO:
|
||||
new_key = key.replace("s.0.weight", ".weight")
|
||||
new_state_dict[new_key] = value
|
||||
else:
|
||||
new_state_dict[key] = value
|
||||
return new_state_dict
|
||||
|
||||
|
||||
class TTSModel(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
flow_lm: FlowLMModel,
|
||||
temp: float,
|
||||
lsd_decode_steps: int,
|
||||
noise_clamp: float | None,
|
||||
eos_threshold,
|
||||
config: Config,
|
||||
):
|
||||
super().__init__()
|
||||
self.flow_lm = flow_lm
|
||||
self.temp = temp
|
||||
self.lsd_decode_steps = lsd_decode_steps
|
||||
self.noise_clamp = noise_clamp
|
||||
self.eos_threshold = eos_threshold
|
||||
self.config = config
|
||||
self.has_voice_cloning = True
|
||||
|
||||
@property
|
||||
def device(self) -> str:
|
||||
return next(self.parameters()).device.type
|
||||
|
||||
@property
|
||||
def sample_rate(self) -> int:
|
||||
return self.config.mimi.sample_rate
|
||||
|
||||
@classmethod
|
||||
def _from_pydantic_config(
|
||||
cls, config: Config, temp, lsd_decode_steps, noise_clamp: float | None, eos_threshold
|
||||
) -> Self:
|
||||
flow_lm = FlowLMModel.from_pydantic_config(
|
||||
config.flow_lm, latent_dim=config.mimi.quantizer.dimension
|
||||
)
|
||||
tts_model = cls(flow_lm, temp, lsd_decode_steps, noise_clamp, eos_threshold, config)
|
||||
return tts_model
|
||||
|
||||
@classmethod
|
||||
def _from_pydantic_config_with_weights(
|
||||
cls, config: Config, temp, lsd_decode_steps, noise_clamp: float | None, eos_threshold
|
||||
) -> Self:
|
||||
tts_model = cls._from_pydantic_config(
|
||||
config, temp, lsd_decode_steps, noise_clamp, eos_threshold
|
||||
)
|
||||
tts_model.flow_lm.speaker_proj_weight = torch.nn.Parameter(
|
||||
torch.zeros((1024, 512), dtype=torch.float32)
|
||||
)
|
||||
if config.flow_lm.weights_path is not None:
|
||||
if config.mimi.weights_path is None:
|
||||
raise ValueError(
|
||||
"If you specify flow_lm.weights_path you should specify mimi.weights_path"
|
||||
)
|
||||
logger.info(f"Loading FlowLM weights from {config.flow_lm.weights_path}")
|
||||
state_dict_flowlm = get_flow_lm_state_dict(
|
||||
download_if_necessary(config.flow_lm.weights_path)
|
||||
)
|
||||
tts_model.flow_lm.load_state_dict(state_dict_flowlm, strict=True)
|
||||
|
||||
# safetensors.torch.save_file(tts_model.state_dict(), "7442637a.safetensors")
|
||||
# Create mimi config directly from the provided config using model_dump
|
||||
mimi_config = config.mimi.model_dump()
|
||||
|
||||
# Build mimi model from config
|
||||
encoder = SEANetEncoder(**mimi_config["seanet"])
|
||||
decoder = SEANetDecoder(**mimi_config["seanet"])
|
||||
|
||||
encoder_transformer = mimi_transformer.ProjectedTransformer(**mimi_config["transformer"])
|
||||
decoder_transformer = mimi_transformer.ProjectedTransformer(**mimi_config["transformer"])
|
||||
quantizer = DummyQuantizer(**mimi_config["quantizer"])
|
||||
|
||||
tts_model.mimi = MimiModel(
|
||||
encoder,
|
||||
decoder,
|
||||
quantizer,
|
||||
channels=mimi_config["channels"],
|
||||
sample_rate=mimi_config["sample_rate"],
|
||||
frame_rate=mimi_config["frame_rate"],
|
||||
encoder_frame_rate=mimi_config["sample_rate"] / encoder.hop_length,
|
||||
encoder_transformer=encoder_transformer,
|
||||
decoder_transformer=decoder_transformer,
|
||||
).to(device="cpu")
|
||||
|
||||
# Load mimi weights from the config safetensors file with complete mapping for strict loading
|
||||
|
||||
if config.mimi.weights_path is not None:
|
||||
if config.flow_lm.weights_path is None:
|
||||
raise ValueError(
|
||||
"If you specify mimi.weights_path you should specify flow_lm.weights_path"
|
||||
)
|
||||
logger.info(f"Loading Mimi weights from {config.mimi.weights_path}")
|
||||
mimi_state = get_mimi_state_dict(download_if_necessary(config.mimi.weights_path))
|
||||
tts_model.mimi.load_state_dict(mimi_state, strict=True)
|
||||
|
||||
tts_model.mimi.eval()
|
||||
# tts_model.to(dtype=torch.float32)
|
||||
|
||||
# uncomment to save the weights
|
||||
# tts_model = tts_model.to(dtype=torch.bfloat16)
|
||||
# safetensors.torch.save_file(tts_model.state_dict(), "tts_b6369a24.safetensors")
|
||||
if config.weights_path is not None:
|
||||
logger.info(f"Loading TTSModel weights from {config.weights_path}")
|
||||
try:
|
||||
weights_file = download_if_necessary(config.weights_path)
|
||||
except Exception:
|
||||
tts_model.has_voice_cloning = False
|
||||
weights_file = download_if_necessary(config.weights_path_without_voice_cloning)
|
||||
|
||||
state_dict = safetensors.torch.load_file(weights_file)
|
||||
state_dict = adapt_state_dict_of_TTSModel(state_dict)
|
||||
tts_model.load_state_dict(state_dict, strict=True)
|
||||
|
||||
if config.flow_lm.weights_path is None and config.weights_path is None:
|
||||
logger.warning(
|
||||
"No weights_path specified for FlowLM or TTSModel, model is uninitialized!"
|
||||
)
|
||||
size_in_mb = size_of_dict(tts_model.state_dict()) // 1e6
|
||||
logging.info(f"TTS Model loaded successfully. Its size is {size_in_mb} MB")
|
||||
|
||||
return tts_model
|
||||
|
||||
def load_model(
|
||||
variant: str = DEFAULT_VARIANT,
|
||||
temp: float | int = DEFAULT_TEMPERATURE,
|
||||
lsd_decode_steps: int = DEFAULT_LSD_DECODE_STEPS,
|
||||
noise_clamp: float | int | None = DEFAULT_NOISE_CLAMP,
|
||||
eos_threshold: float = DEFAULT_EOS_THRESHOLD,
|
||||
) -> Self:
|
||||
"""Load a pre-trained TTS model with specified configuration.
|
||||
|
||||
This class method loads a complete TTS model including the flow language model
|
||||
and Mimi compression model from pre-trained weights. The model is initialized
|
||||
with the specified generation parameters and ready for inference.
|
||||
|
||||
Args:
|
||||
variant: Model variant identifier corresponding to a config file name
|
||||
(e.g., '610b0b2c'). Must match a YAML file in the config directory.
|
||||
temp: Sampling temperature for generation. Higher values produce more
|
||||
diverse but potentially lower quality output.
|
||||
lsd_decode_steps: Number of steps for Lagrangian Self Distillation
|
||||
decoding. More steps can improve quality but increase computation.
|
||||
noise_clamp: Maximum value for noise sampling. If None, no clamping
|
||||
is applied. Helps prevent extreme values in generation.
|
||||
eos_threshold: Threshold for end-of-sequence detection. Higher values
|
||||
make the model more likely to continue generating.
|
||||
|
||||
Returns:
|
||||
TTSModel: Fully initialized model with loaded weights on cpu, ready for
|
||||
text-to-speech generation.
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If the specified config file or model weights
|
||||
are not found.
|
||||
ValueError: If the configuration is invalid or incompatible.
|
||||
"""
|
||||
config = load_config(Path(__file__).parents[1] / f"config/{variant}.yaml")
|
||||
tts_model = TTSModel._from_pydantic_config_with_weights(
|
||||
config, temp, lsd_decode_steps, noise_clamp, eos_threshold
|
||||
)
|
||||
return tts_model
|
||||
|
||||
def _run_flow_lm_and_increment_step(
|
||||
self,
|
||||
model_state: dict,
|
||||
text_tokens: torch.Tensor | None = None,
|
||||
backbone_input_latents: torch.Tensor | None = None,
|
||||
audio_conditioning: torch.Tensor | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""First one is the backbone output, second one is the audio decoding output."""
|
||||
if text_tokens is None:
|
||||
text_tokens = torch.zeros((1, 0), dtype=torch.int64, device=self.flow_lm.device)
|
||||
if backbone_input_latents is None:
|
||||
backbone_input_latents = torch.empty(
|
||||
(1, 0, self.flow_lm.ldim), dtype=self.flow_lm.dtype, device=self.flow_lm.device
|
||||
)
|
||||
if audio_conditioning is None:
|
||||
audio_conditioning = torch.empty(
|
||||
(1, 0, self.flow_lm.dim), dtype=self.flow_lm.dtype, device=self.flow_lm.device
|
||||
)
|
||||
|
||||
output = self._run_flow_lm(
|
||||
text_tokens=text_tokens,
|
||||
backbone_input_latents=backbone_input_latents,
|
||||
model_state=model_state,
|
||||
audio_conditioning=audio_conditioning,
|
||||
)
|
||||
increment_by = (
|
||||
text_tokens.shape[1] + backbone_input_latents.shape[1] + audio_conditioning.shape[1]
|
||||
)
|
||||
increment_steps(self.flow_lm, model_state, increment=increment_by)
|
||||
return output
|
||||
|
||||
def _run_flow_lm(
|
||||
self,
|
||||
model_state: dict,
|
||||
text_tokens: torch.Tensor,
|
||||
backbone_input_latents: torch.Tensor,
|
||||
audio_conditioning: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
text_embeddings = self.flow_lm.conditioner(TokenizedText(text_tokens))
|
||||
text_embeddings = torch.cat([text_embeddings, audio_conditioning], dim=1)
|
||||
|
||||
output_embeddings, is_eos = self.flow_lm._sample_next_latent(
|
||||
backbone_input_latents,
|
||||
text_embeddings,
|
||||
model_state=model_state,
|
||||
lsd_decode_steps=self.lsd_decode_steps,
|
||||
temp=self.temp,
|
||||
noise_clamp=self.noise_clamp,
|
||||
eos_threshold=self.eos_threshold,
|
||||
)
|
||||
return output_embeddings[:, None, :], is_eos
|
||||
|
||||
def _encode_audio(self, audio: torch.Tensor) -> torch.Tensor:
|
||||
encoded = self.mimi.encode_to_latent(audio)
|
||||
latents = encoded.transpose(-1, -2).to(torch.float32)
|
||||
conditioning = F.linear(latents, self.flow_lm.speaker_proj_weight)
|
||||
return conditioning
|
||||
|
||||
def _decode_audio_worker(self, latents_queue: queue.Queue, result_queue: queue.Queue):
|
||||
"""Worker thread function for decoding audio latents from queue with immediate streaming."""
|
||||
try:
|
||||
audio_chunks = []
|
||||
mimi_state = init_states(self.mimi, batch_size=1, sequence_length=1000)
|
||||
while True:
|
||||
latent = latents_queue.get()
|
||||
if latent is None:
|
||||
break
|
||||
mimi_decoding_input = latent * self.flow_lm.emb_std + self.flow_lm.emb_mean
|
||||
transposed = mimi_decoding_input.transpose(-1, -2)
|
||||
quantized = self.mimi.quantizer(transposed)
|
||||
|
||||
t = time.monotonic()
|
||||
audio_frame = self.mimi.decode_from_latent(quantized, mimi_state)
|
||||
increment_steps(self.mimi, mimi_state, increment=16)
|
||||
audio_frame_duration = audio_frame.shape[2] / self.config.mimi.sample_rate
|
||||
# We could log the timings here.
|
||||
logger.debug(
|
||||
" " * 30 + "Decoded %d ms of audio with mimi in %d ms",
|
||||
int(audio_frame_duration * 1000),
|
||||
int((time.monotonic() - t) * 1000),
|
||||
)
|
||||
audio_chunks.append(audio_frame)
|
||||
|
||||
result_queue.put(("chunk", audio_frame))
|
||||
|
||||
latents_queue.task_done()
|
||||
|
||||
# Signal completion
|
||||
result_queue.put(("done", None))
|
||||
|
||||
except Exception as e:
|
||||
# Put error in result queue
|
||||
result_queue.put(("error", e))
|
||||
|
||||
@torch.no_grad
|
||||
def generate_audio(
|
||||
self,
|
||||
model_state: dict,
|
||||
text_to_generate: str,
|
||||
frames_after_eos: int | None = None,
|
||||
copy_state: bool = True,
|
||||
) -> torch.Tensor:
|
||||
"""Generate complete audio tensor from text input.
|
||||
|
||||
This method generates the full audio output for the given text prompt
|
||||
and returns it as a single tensor. It internally uses the streaming
|
||||
generation method but collects all chunks before returning.
|
||||
|
||||
This method is NOT thread-safe; separate model instances should be used
|
||||
for concurrent generation.
|
||||
|
||||
Args:
|
||||
model_state: Model state dictionary containing hidden states and
|
||||
positional information. Can be obtained from get_state_for_audio_prompt()
|
||||
or init_states(). The state may be modified during generation.
|
||||
text_to_generate: Input text to convert to speech. The text will be
|
||||
automatically formatted (capitalization, punctuation) for optimal
|
||||
generation quality.
|
||||
frames_after_eos: Number of additional frames to generate after
|
||||
detecting end-of-sequence. If None, automatically determined
|
||||
based on text length (1-3 frames).
|
||||
copy_state: Whether to create a deep copy of the model state before
|
||||
generation. If True, preserves the original state for reuse.
|
||||
If False, modifies the input state in-place. Defaults to True.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Generated audio tensor with shape [channels, samples]
|
||||
at the model's sample rate (typically 24kHz). The audio is
|
||||
normalized and ready for playback or saving.
|
||||
You can get the sample rate from the `sample_rate` attribute.
|
||||
|
||||
Raises:
|
||||
ValueError: If text_to_generate is empty or invalid.
|
||||
RuntimeError: If generation fails due to model errors.
|
||||
"""
|
||||
audio_chunks = []
|
||||
for chunk in self.generate_audio_stream(
|
||||
model_state=model_state,
|
||||
text_to_generate=text_to_generate,
|
||||
frames_after_eos=frames_after_eos,
|
||||
copy_state=copy_state,
|
||||
):
|
||||
audio_chunks.append(chunk)
|
||||
return torch.cat(audio_chunks, dim=0)
|
||||
|
||||
@torch.no_grad
|
||||
def generate_audio_stream(
|
||||
self,
|
||||
model_state: dict,
|
||||
text_to_generate: str,
|
||||
frames_after_eos: int | None = None,
|
||||
copy_state: bool = True,
|
||||
):
|
||||
"""Generate audio streaming chunks from text input.
|
||||
|
||||
This method generates audio from text and yields chunks as they become
|
||||
available, enabling real-time playback or processing. It uses multithreading
|
||||
to parallelize generation and decoding for optimal performance.
|
||||
This method is NOT thread-safe; separate model instances should be used
|
||||
for concurrent generation.
|
||||
|
||||
Args:
|
||||
model_state: Model state dictionary containing hidden states and
|
||||
positional information. Can be obtained from get_state_for_audio_prompt()
|
||||
or init_states(). The state may be modified during generation.
|
||||
text_to_generate: Input text to convert to speech. The text will be
|
||||
automatically formatted (capitalization, punctuation) for optimal
|
||||
generation quality.
|
||||
frames_after_eos: Number of additional frames to generate after
|
||||
detecting end-of-sequence. If None, automatically determined
|
||||
based on text length (1-3 frames). Defaults to None.
|
||||
copy_state: Whether to create a deep copy of the model state before
|
||||
generation. If True, preserves the original state for reuse.
|
||||
If False, modifies the input state in-place. Defaults to True.
|
||||
|
||||
Yields:
|
||||
torch.Tensor: Audio chunks with shape [samples] at the model's
|
||||
sample rate (typically 24kHz). Chunks are yielded as soon as
|
||||
they are decoded, enabling real-time streaming.
|
||||
|
||||
Raises:
|
||||
ValueError: If text_to_generate is empty or invalid.
|
||||
RuntimeError: If generation fails due to model errors or threading issues.
|
||||
|
||||
Note:
|
||||
This method uses multithreading to parallelize latent generation
|
||||
and audio decoding. Generation performance is logged including
|
||||
real-time factor (RTF) metrics.
|
||||
"""
|
||||
text_to_generate, frames_after_eos_guess = prepare_text_prompt(text_to_generate)
|
||||
if frames_after_eos is None:
|
||||
frames_after_eos = frames_after_eos_guess
|
||||
|
||||
if copy_state:
|
||||
model_state = copy.deepcopy(model_state)
|
||||
|
||||
# Set up multithreaded generation and decoding
|
||||
latents_queue = queue.Queue()
|
||||
result_queue = queue.Queue()
|
||||
|
||||
# Start decoder worker thread
|
||||
decoder_thread = threading.Thread(
|
||||
target=self._decode_audio_worker, args=(latents_queue, result_queue), daemon=True
|
||||
)
|
||||
logger.info("starting timer now!")
|
||||
t_generating = time.monotonic()
|
||||
decoder_thread.start()
|
||||
|
||||
# Generate latents and add them to queue (decoder processes them in parallel)
|
||||
self._generate(
|
||||
model_state=model_state,
|
||||
text_to_generate=text_to_generate,
|
||||
frames_after_eos=frames_after_eos,
|
||||
latents_queue=latents_queue,
|
||||
result_queue=result_queue,
|
||||
)
|
||||
|
||||
# Stream audio chunks as they become available
|
||||
total_generated_samples = 0
|
||||
while True:
|
||||
result = result_queue.get()
|
||||
if result[0] == "chunk":
|
||||
# Audio chunk available immediately for streaming/playback
|
||||
audio_chunk = result[1]
|
||||
total_generated_samples += audio_chunk.shape[-1]
|
||||
yield audio_chunk[0, 0] # Remove batch, channel
|
||||
elif result[0] == "done":
|
||||
# Generation complete
|
||||
break
|
||||
elif result[0] == "error":
|
||||
# Wait for decoder thread to finish cleanly before propagating error
|
||||
with display_execution_time("Waiting for mimi decoder to finish"):
|
||||
decoder_thread.join()
|
||||
# Propagate error
|
||||
raise result[1]
|
||||
|
||||
# Wait for decoder thread to finish cleanly
|
||||
with display_execution_time("Waiting for mimi decoder to finish"):
|
||||
decoder_thread.join()
|
||||
|
||||
# Print timing information
|
||||
duration_generated_audio = int(
|
||||
total_generated_samples * 1000 / self.config.mimi.sample_rate
|
||||
)
|
||||
generation_time = int((time.monotonic() - t_generating) * 1000)
|
||||
real_time_factor = duration_generated_audio / generation_time
|
||||
|
||||
logger.info(
|
||||
"Generated: %d ms of audio in %d ms so %.2fx faster than real-time",
|
||||
duration_generated_audio,
|
||||
generation_time,
|
||||
real_time_factor,
|
||||
)
|
||||
|
||||
@torch.no_grad
|
||||
def _generate(
|
||||
self,
|
||||
model_state: dict,
|
||||
text_to_generate: str,
|
||||
frames_after_eos: int,
|
||||
latents_queue: queue.Queue,
|
||||
result_queue: queue.Queue,
|
||||
):
|
||||
gen_len_sec = len(text_to_generate.split()) * 1 + 2.0
|
||||
max_gen_len = int(gen_len_sec * 12.5)
|
||||
prepared = self.flow_lm.conditioner.prepare(text_to_generate)
|
||||
|
||||
with display_execution_time("Prompting text"):
|
||||
self._run_flow_lm_and_increment_step(
|
||||
model_state=model_state, text_tokens=prepared.tokens
|
||||
)
|
||||
|
||||
def run_generation():
|
||||
try:
|
||||
self._autoregressive_generation(
|
||||
model_state, max_gen_len, frames_after_eos, latents_queue
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error in autoregressive generation: {e}")
|
||||
# Signal decoder to stop by putting None (completion sentinel)
|
||||
if latents_queue is not None:
|
||||
latents_queue.put(None)
|
||||
# Report error to main thread
|
||||
if result_queue is not None:
|
||||
result_queue.put(("error", e))
|
||||
|
||||
generation_thread = threading.Thread(target=run_generation, daemon=True)
|
||||
generation_thread.start()
|
||||
|
||||
def _autoregressive_generation(
|
||||
self, model_state: dict, max_gen_len: int, frames_after_eos: int, latents_queue: queue.Queue
|
||||
):
|
||||
backbone_input = torch.full(
|
||||
(1, 1, self.flow_lm.ldim),
|
||||
fill_value=float("NaN"),
|
||||
device=next(iter(self.flow_lm.parameters())).device,
|
||||
dtype=self.flow_lm.dtype,
|
||||
)
|
||||
steps_times = []
|
||||
eos_step = None
|
||||
for generation_step in range(max_gen_len):
|
||||
with display_execution_time("Generating latent", print_output=False) as timer:
|
||||
next_latent, is_eos = self._run_flow_lm_and_increment_step(
|
||||
model_state=model_state, backbone_input_latents=backbone_input
|
||||
)
|
||||
if is_eos.item() and eos_step is None:
|
||||
eos_step = generation_step
|
||||
if eos_step is not None and generation_step >= eos_step + frames_after_eos:
|
||||
break
|
||||
|
||||
# Add generated latent to queue for immediate decoding
|
||||
latents_queue.put(next_latent)
|
||||
backbone_input = next_latent
|
||||
steps_times.append(timer.elapsed_time_ms)
|
||||
else:
|
||||
if os.environ.get("KPOCKET_TTS_ERROR_WITHOUT_EOS", "0") == "1":
|
||||
raise RuntimeError("Generation reached maximum length without EOS!")
|
||||
logger.warning(
|
||||
"Maximum generation length reached without EOS, this very often indicates an error."
|
||||
)
|
||||
|
||||
# Add sentinel value to signal end of generation
|
||||
latents_queue.put(None)
|
||||
logger.info("Average generation step time: %d ms", int(statistics.mean(steps_times)))
|
||||
|
||||
@lru_cache(maxsize=2)
|
||||
def _cached_get_state_for_audio_prompt(
|
||||
self, audio_conditioning: Path | str | torch.Tensor, truncate: bool = False
|
||||
) -> dict:
|
||||
return self.get_state_for_audio_prompt(audio_conditioning, truncate)
|
||||
|
||||
@torch.no_grad
|
||||
def get_state_for_audio_prompt(
|
||||
self, audio_conditioning: Path | str | torch.Tensor, truncate: bool = False
|
||||
) -> dict:
|
||||
"""Create model state conditioned on audio prompt for continuation.
|
||||
|
||||
This method processes an audio prompt and creates a model state that
|
||||
captures the acoustic characteristics (speaker voice, style, prosody)
|
||||
for use in subsequent text-to-speech generation. The resulting state
|
||||
enables voice cloning and audio continuation with speaker consistency.
|
||||
|
||||
Args:
|
||||
audio_conditioning: Audio prompt to condition on. Can be:
|
||||
- Path: Local file path to audio file
|
||||
- str: URL to download audio file from
|
||||
- torch.Tensor: Pre-loaded audio tensor with shape [channels, samples]
|
||||
truncate: Whether to truncate long audio prompts to 30 seconds.
|
||||
Helps prevent memory issues with very long inputs. Defaults to False.
|
||||
|
||||
Returns:
|
||||
dict: Model state dictionary containing hidden states and positional
|
||||
information conditioned on the audio prompt. This state can be
|
||||
passed to `generate_audio()` or `generate_audio_stream()` for
|
||||
voice-consistent generation.
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If audio file path doesn't exist.
|
||||
ValueError: If audio tensor is invalid or empty.
|
||||
RuntimeError: If audio processing or encoding fails.
|
||||
|
||||
Note:
|
||||
- Audio is automatically resampled to the model's sample rate (24kHz)
|
||||
- The audio is encoded using the Mimi compression model and projected
|
||||
to the flow model's latent space
|
||||
- Processing time is logged for performance monitoring
|
||||
- The state preserves speaker characteristics for voice cloning
|
||||
"""
|
||||
if isinstance(audio_conditioning, str) and audio_conditioning in PREDEFINED_VOICES:
|
||||
# We get the audio conditioning directly from the safetensors file.
|
||||
prompt = load_predefined_voice(audio_conditioning)
|
||||
else:
|
||||
if not self.has_voice_cloning and isinstance(audio_conditioning, (str, Path)):
|
||||
raise ValueError(
|
||||
f"We could not download the weights for the model with voice cloning, "
|
||||
f"but you're trying to use voice cloning. "
|
||||
f"Without voice cloning, you can use our catalog of voices {list(PREDEFINED_VOICES)}. "
|
||||
f"If you want access to the model with voice cloning, go to "
|
||||
f"https://huggingface.co/kyutai/pocket-tts and accept the terms, "
|
||||
f"then make sure you're logged in locally with `uvx hf auth login`."
|
||||
)
|
||||
if isinstance(audio_conditioning, str):
|
||||
audio_conditioning = download_if_necessary(audio_conditioning)
|
||||
|
||||
if isinstance(audio_conditioning, Path):
|
||||
audio, conditioning_sample_rate = audio_read(audio_conditioning)
|
||||
|
||||
if truncate:
|
||||
max_samples = int(30 * conditioning_sample_rate) # 30 seconds of audio
|
||||
if audio.shape[-1] > max_samples:
|
||||
audio = audio[..., :max_samples]
|
||||
logger.info(f"Audio truncated to first 30 seconds ({max_samples} samples)")
|
||||
|
||||
audio_conditioning = convert_audio(
|
||||
audio, conditioning_sample_rate, self.config.mimi.sample_rate, 1
|
||||
)
|
||||
|
||||
with display_execution_time("Encoding audio prompt"):
|
||||
prompt = self._encode_audio(audio_conditioning.unsqueeze(0).to(self.device))
|
||||
# import safetensors.torch
|
||||
# safetensors.torch.save_file({"audio_prompt": prompt}, "embeddings/eponine.safetensors")
|
||||
|
||||
model_state = init_states(self.flow_lm, batch_size=1, sequence_length=1000)
|
||||
|
||||
with display_execution_time("Prompting audio"):
|
||||
self._run_flow_lm_and_increment_step(model_state=model_state, audio_conditioning=prompt)
|
||||
|
||||
return model_state
|
||||
|
||||
|
||||
def prepare_text_prompt(text: str) -> tuple[str, int]:
|
||||
text = text.strip()
|
||||
if text == "":
|
||||
raise ValueError("Text prompt cannot be empty")
|
||||
number_of_words = len(text.split())
|
||||
if number_of_words <= 4:
|
||||
frames_after_eos_guess = 3
|
||||
else:
|
||||
frames_after_eos_guess = 1
|
||||
|
||||
# Make sure it starts with an uppercase letter
|
||||
if not text[0].isupper():
|
||||
text = text[0].upper() + text[1:]
|
||||
|
||||
# Let's make sure it ends with some kind of punctuation
|
||||
# If it ends with a letter or digit, we add a period.
|
||||
if text[-1].isalnum():
|
||||
text = text + "."
|
||||
|
||||
# The model does not perform well when there are very few tokens, so
|
||||
# we can add empty spaces at the beginning to increase the token count.
|
||||
if len(text.split()) < 5:
|
||||
text = " " * 8 + text
|
||||
|
||||
return text, frames_after_eos_guess
|
||||
@@ -0,0 +1 @@
|
||||
"""Modules used for building the models."""
|
||||
@@ -0,0 +1,161 @@
|
||||
import math
|
||||
import warnings
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
from pocket_tts.modules.stateful_module import StatefulModule
|
||||
|
||||
|
||||
def get_extra_padding_for_conv1d(
|
||||
x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0
|
||||
) -> int:
|
||||
"""See `pad_for_conv1d`."""
|
||||
length = x.shape[-1]
|
||||
n_frames = (length - kernel_size + padding_total) / stride + 1
|
||||
ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total)
|
||||
return ideal_length - length
|
||||
|
||||
|
||||
def pad_for_conv1d(x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0):
|
||||
"""Pad for a convolution to make sure that the last window is full.
|
||||
Extra padding is added at the end. This is required to ensure that we can rebuild
|
||||
an output of the same length, as otherwise, even with padding, some time steps
|
||||
might get removed.
|
||||
For instance, with total padding = 4, kernel size = 4, stride = 2:
|
||||
0 0 1 2 3 4 5 0 0 # (0s are padding)
|
||||
1 2 3 # (output frames of a convolution, last 0 is never used)
|
||||
0 0 1 2 3 4 5 0 # (output of tr. conv., but pos. 5 is going to get removed as padding)
|
||||
1 2 3 4 # once you removed padding, we are missing one time step !
|
||||
"""
|
||||
extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total)
|
||||
return F.pad(x, (0, extra_padding))
|
||||
|
||||
|
||||
class StreamingConv1d(StatefulModule):
|
||||
"""Conv1d with some builtin handling of asymmetric or causal padding
|
||||
and normalization.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
kernel_size: int,
|
||||
stride: int = 1,
|
||||
dilation: int = 1,
|
||||
groups: int = 1,
|
||||
bias: bool = True,
|
||||
pad_mode: str = "constant",
|
||||
):
|
||||
super().__init__()
|
||||
assert pad_mode in ["constant", "replicate"], pad_mode
|
||||
self.pad_mode = pad_mode
|
||||
# warn user on unusual setup between dilation and stride
|
||||
if stride > 1 and dilation > 1:
|
||||
warnings.warn(
|
||||
"StreamingConv1d has been initialized with stride > 1 and dilation > 1"
|
||||
f" (kernel_size={kernel_size} stride={stride}, dilation={dilation})."
|
||||
)
|
||||
self.conv = nn.Conv1d(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride,
|
||||
dilation=dilation,
|
||||
groups=groups,
|
||||
bias=bias,
|
||||
)
|
||||
|
||||
@property
|
||||
def _stride(self) -> int:
|
||||
return self.conv.stride[0]
|
||||
|
||||
@property
|
||||
def _kernel_size(self) -> int:
|
||||
return self.conv.kernel_size[0]
|
||||
|
||||
@property
|
||||
def _effective_kernel_size(self) -> int:
|
||||
dilation = self.conv.dilation[0]
|
||||
return (self._kernel_size - 1) * dilation + 1 # effective kernel size with dilations
|
||||
|
||||
def init_state(self, batch_size: int, sequence_length: int) -> dict[str, torch.Tensor]:
|
||||
stride = self._stride
|
||||
# Effective kernel size accounting for dilation.
|
||||
kernel = self._effective_kernel_size
|
||||
previous = torch.zeros(batch_size, self.conv.in_channels, kernel - stride)
|
||||
first = torch.ones(batch_size, dtype=torch.bool)
|
||||
return dict(previous=previous, first=first)
|
||||
|
||||
def forward(self, x, model_state: dict | None):
|
||||
B, C, T = x.shape
|
||||
S = self._stride
|
||||
assert T > 0 and T % S == 0, "Steps must be multiple of stride"
|
||||
if model_state is None:
|
||||
state = self.init_state(B, 0)
|
||||
else:
|
||||
state = self.get_state(model_state)
|
||||
TP = state["previous"].shape[-1]
|
||||
if TP and self.pad_mode == "replicate":
|
||||
assert T >= TP, "Not enough content to pad streaming."
|
||||
init = x[..., :1]
|
||||
state["previous"][:] = torch.where(
|
||||
state["first"].view(-1, 1, 1), init, state["previous"]
|
||||
)
|
||||
if TP:
|
||||
x = torch.cat([state["previous"], x], dim=-1)
|
||||
y = self.conv(x)
|
||||
if TP:
|
||||
state["previous"][:] = x[..., -TP:]
|
||||
if self.pad_mode == "replicate":
|
||||
state["first"] = torch.zeros_like(state["first"])
|
||||
return y
|
||||
|
||||
|
||||
class StreamingConvTranspose1d(StatefulModule):
|
||||
"""ConvTranspose1d with some builtin handling of asymmetric or causal padding
|
||||
and normalization.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
kernel_size: int,
|
||||
stride: int = 1,
|
||||
groups: int = 1,
|
||||
bias: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
self.convtr = nn.ConvTranspose1d(
|
||||
in_channels, out_channels, kernel_size, stride, groups=groups, bias=bias
|
||||
)
|
||||
|
||||
@property
|
||||
def _stride(self) -> int:
|
||||
return self.convtr.stride[0]
|
||||
|
||||
@property
|
||||
def _kernel_size(self) -> int:
|
||||
return self.convtr.kernel_size[0]
|
||||
|
||||
def init_state(self, batch_size: int, sequence_length: int) -> dict[str, torch.Tensor]:
|
||||
K = self._kernel_size
|
||||
S = self._stride
|
||||
return dict(partial=torch.zeros(batch_size, self.convtr.out_channels, K - S))
|
||||
|
||||
def forward(self, x, mimi_state: dict):
|
||||
layer_state = self.get_state(mimi_state)["partial"]
|
||||
y = self.convtr(x)
|
||||
PT = layer_state.shape[-1]
|
||||
if PT > 0:
|
||||
y[..., :PT] += layer_state
|
||||
bias = self.convtr.bias
|
||||
for_partial = y[..., -PT:]
|
||||
if bias is not None:
|
||||
for_partial -= bias[:, None]
|
||||
layer_state[:] = for_partial
|
||||
y = y[..., :-PT]
|
||||
return y
|
||||
@@ -0,0 +1,18 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
class DummyQuantizer(nn.Module):
|
||||
"""Simplified quantizer that only provides output projection for TTS.
|
||||
|
||||
This removes all unnecessary quantization logic since we don't use actual quantization.
|
||||
"""
|
||||
|
||||
def __init__(self, dimension: int, output_dimension: int):
|
||||
super().__init__()
|
||||
self.dimension = dimension
|
||||
self.output_dimension = output_dimension
|
||||
self.output_proj = torch.nn.Conv1d(self.dimension, self.output_dimension, 1, bias=False)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.output_proj(x)
|
||||
@@ -0,0 +1,11 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class LayerScale(nn.Module):
|
||||
def __init__(self, channels: int, init: float):
|
||||
super().__init__()
|
||||
self.scale = nn.Parameter(torch.full((channels,), init))
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
return self.scale * x
|
||||
@@ -0,0 +1,285 @@
|
||||
from typing import NamedTuple
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from einops import rearrange
|
||||
from torch.nn import functional as F
|
||||
from typing_extensions import Self
|
||||
|
||||
from pocket_tts.modules.layer_scale import LayerScale
|
||||
from pocket_tts.modules.rope import RotaryEmbedding
|
||||
from pocket_tts.modules.stateful_module import StatefulModule
|
||||
from pocket_tts.modules.transformer import StreamingMultiheadAttention
|
||||
from pocket_tts.utils.config import FlowLMTransformerConfig
|
||||
|
||||
|
||||
class KVCacheResult(NamedTuple):
|
||||
keys: torch.Tensor
|
||||
values: torch.Tensor
|
||||
positions: torch.Tensor
|
||||
|
||||
@staticmethod
|
||||
def from_kv(keys: torch.Tensor, values: torch.Tensor) -> "KVCacheResult":
|
||||
B, H, T, D = keys.shape
|
||||
assert tuple(values.shape[:-1]) == (B, H, T)
|
||||
positions = torch.arange(T, device=keys.device, dtype=torch.long)
|
||||
return KVCacheResult(keys, values, positions.expand(B, -1))
|
||||
|
||||
|
||||
def complete(
|
||||
cache: torch.Tensor, end_offset: torch.Tensor, k: torch.Tensor, v: torch.Tensor
|
||||
) -> KVCacheResult:
|
||||
capacity = cache.shape[3]
|
||||
assert k.shape[:-1] == v.shape[:-1], (k.shape, v.shape)
|
||||
B, H, T, D = k.shape
|
||||
assert T > 0
|
||||
indexes = torch.arange(T, device=end_offset.device, dtype=end_offset.dtype)
|
||||
indexes = indexes + end_offset.view(-1, 1)
|
||||
indexes = indexes % capacity
|
||||
# indexes is [B, T]
|
||||
# k is [B, H, T, D]
|
||||
# cache is [B, H, T', D]
|
||||
this_indexes = indexes.view(B, 1, T, 1)
|
||||
this_indexes = this_indexes.expand(-1, H, T, D)
|
||||
cache[0].scatter_(2, this_indexes, k)
|
||||
cache[1].scatter_(2, this_indexes, v)
|
||||
|
||||
keys = cache[0]
|
||||
values = cache[1]
|
||||
|
||||
indexes = torch.arange(capacity, device=end_offset.device, dtype=torch.long)
|
||||
|
||||
# end_index correspond to the actual index where the last value was written.
|
||||
last_offset = end_offset.view(-1, 1) + T - 1
|
||||
end_index = last_offset % capacity
|
||||
delta = indexes - end_index
|
||||
|
||||
positions = torch.where(delta <= 0, last_offset + delta, last_offset + delta - capacity)
|
||||
end_offset[:] = end_offset + T
|
||||
invalid = indexes >= end_offset.view(-1, 1)
|
||||
positions = torch.where(invalid, torch.full_like(positions, -1), positions)
|
||||
|
||||
return KVCacheResult(keys, values, positions)
|
||||
|
||||
|
||||
class MimiStreamingMultiheadAttention(StatefulModule):
|
||||
def __init__(self, embed_dim: int, num_heads: int, context: int, rope: RotaryEmbedding):
|
||||
super().__init__()
|
||||
|
||||
self.embed_dim = embed_dim
|
||||
self.context = context
|
||||
self.rope = rope
|
||||
self.num_heads = num_heads
|
||||
out_dim = 3 * embed_dim
|
||||
|
||||
self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False)
|
||||
self.in_proj = nn.Linear(embed_dim, out_dim, bias=False)
|
||||
|
||||
def init_state(self, batch_size: int, sequence_length: int) -> dict[str, torch.Tensor]:
|
||||
dim_per_head = self.embed_dim // self.num_heads
|
||||
|
||||
state = {}
|
||||
state["offset"] = torch.zeros(batch_size, dtype=torch.long)
|
||||
state["cache"] = torch.zeros((2, batch_size, self.num_heads, sequence_length, dim_per_head))
|
||||
state["end_offset"] = torch.zeros(batch_size, dtype=torch.long)
|
||||
return state
|
||||
|
||||
def increment_step(self, state, increment: int = 1):
|
||||
state["offset"] += increment
|
||||
|
||||
def _complete_kv(self, k, v, model_state: dict | None) -> KVCacheResult:
|
||||
if model_state is None:
|
||||
return KVCacheResult.from_kv(k, v)
|
||||
else:
|
||||
layer_state = self.get_state(model_state)
|
||||
return complete(layer_state["cache"], layer_state["end_offset"], k, v)
|
||||
|
||||
def forward(self, query: torch.Tensor, model_state: dict | None) -> torch.Tensor:
|
||||
B, T = query.shape[:2]
|
||||
|
||||
if model_state is None:
|
||||
offset = torch.zeros(B, device=query.device, dtype=torch.long)
|
||||
else:
|
||||
offset = self.get_state(model_state)["offset"]
|
||||
|
||||
projected = self.in_proj(query)
|
||||
|
||||
q, k, v = rearrange(projected, "b t (p h d) -> p b h t d", p=3, h=self.num_heads)
|
||||
|
||||
# Permute from [b, h, t, d] to [b, t, h, d] for rope
|
||||
q = q.permute(0, 2, 1, 3)
|
||||
k = k.permute(0, 2, 1, 3)
|
||||
q, k = self.rope(q, k, offset)
|
||||
# Permute back from [b, t, h, d] to [b, h, t, d]
|
||||
q = q.permute(0, 2, 1, 3)
|
||||
k = k.permute(0, 2, 1, 3)
|
||||
|
||||
k, v, pos_k = self._complete_kv(k, v, model_state)
|
||||
pos_k = pos_k[:, None]
|
||||
pos_q = offset.view(-1, 1, 1) + torch.arange(T, device=q.device, dtype=torch.long).view(
|
||||
-1, 1
|
||||
)
|
||||
delta = pos_q - pos_k
|
||||
attn_bias = (pos_k >= 0) & (delta >= 0)
|
||||
attn_bias = attn_bias & (delta < self.context)
|
||||
attn_bias = attn_bias[:, None]
|
||||
|
||||
x = F.scaled_dot_product_attention(q, k, v, attn_bias, dropout_p=0.0)
|
||||
|
||||
x = rearrange(x, "b h t d -> b t (h d)")
|
||||
x = self.out_proj(x)
|
||||
return x
|
||||
|
||||
|
||||
class StreamingTransformerLayer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
d_model: int,
|
||||
num_heads: int,
|
||||
dim_feedforward: int,
|
||||
context: int | None,
|
||||
rope: RotaryEmbedding,
|
||||
layer_scale: float | None = None,
|
||||
attention_kind: str = "mimi",
|
||||
):
|
||||
super().__init__()
|
||||
# Redefine self_attn to our streaming multi-head attention
|
||||
if attention_kind == "mimi":
|
||||
# TODO: we should actually use StreamingMultiheadAttention here and add context window
|
||||
# support. And we should then delete MimiStreamingMultiheadAttention.
|
||||
# The implementation is really close.
|
||||
self.self_attn = MimiStreamingMultiheadAttention(
|
||||
context=context, rope=rope, embed_dim=d_model, num_heads=num_heads
|
||||
)
|
||||
else:
|
||||
self.self_attn = StreamingMultiheadAttention(
|
||||
rope=rope, embed_dim=d_model, num_heads=num_heads
|
||||
)
|
||||
self.norm1 = nn.LayerNorm(d_model, eps=1e-5)
|
||||
self.norm2 = nn.LayerNorm(d_model, eps=1e-5)
|
||||
|
||||
self.linear1 = nn.Linear(d_model, dim_feedforward, bias=False)
|
||||
self.linear2 = nn.Linear(dim_feedforward, d_model, bias=False)
|
||||
|
||||
if layer_scale is None:
|
||||
self.layer_scale_1 = nn.Identity()
|
||||
self.layer_scale_2 = nn.Identity()
|
||||
else:
|
||||
self.layer_scale_1 = LayerScale(d_model, layer_scale)
|
||||
self.layer_scale_2 = LayerScale(d_model, layer_scale)
|
||||
|
||||
def _ff_block(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x_orig = x
|
||||
x = self.norm2(x)
|
||||
update = self.linear2(F.gelu(self.linear1(x)))
|
||||
return x_orig.to(update) + self.layer_scale_2(update)
|
||||
|
||||
def _sa_block(self, x: torch.Tensor, model_state: dict | None) -> torch.Tensor:
|
||||
x_orig = x
|
||||
x = self.norm1(x)
|
||||
update = self.self_attn(x, model_state)
|
||||
return x_orig.to(update) + self.layer_scale_1(update)
|
||||
|
||||
def forward(self, x: torch.Tensor, model_state: dict | None) -> torch.Tensor:
|
||||
x = self._sa_block(x, model_state)
|
||||
x = self._ff_block(x)
|
||||
return x
|
||||
|
||||
|
||||
class StreamingTransformer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
d_model: int,
|
||||
num_heads: int,
|
||||
num_layers: int,
|
||||
layer_scale: float | None = None,
|
||||
dim_feedforward: int | list[int] = 2048,
|
||||
context: int | None = None,
|
||||
max_period: float = 10_000.0,
|
||||
kind: str = "mimi",
|
||||
):
|
||||
super().__init__()
|
||||
assert d_model % num_heads == 0
|
||||
self.max_period = max_period
|
||||
|
||||
self.rope = RotaryEmbedding(max_period=max_period)
|
||||
|
||||
self.layers = nn.ModuleList()
|
||||
for _ in range(num_layers):
|
||||
self.layers.append(
|
||||
StreamingTransformerLayer(
|
||||
d_model=d_model,
|
||||
num_heads=num_heads,
|
||||
dim_feedforward=dim_feedforward,
|
||||
context=context,
|
||||
rope=self.rope,
|
||||
layer_scale=layer_scale,
|
||||
attention_kind=kind,
|
||||
)
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_pydantic_config(cls, config: FlowLMTransformerConfig) -> Self:
|
||||
dim_feedforward = int(config.d_model * config.hidden_scale)
|
||||
return cls(
|
||||
d_model=config.d_model,
|
||||
num_heads=config.num_heads,
|
||||
num_layers=config.num_layers,
|
||||
dim_feedforward=dim_feedforward,
|
||||
max_period=float(config.max_period),
|
||||
kind="flow_lm",
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor, model_state: dict | None):
|
||||
for layer in self.layers:
|
||||
x = layer(x, model_state)
|
||||
return x
|
||||
|
||||
|
||||
class ProjectedTransformer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
input_dimension: int,
|
||||
output_dimensions: tuple[int, ...],
|
||||
d_model: int,
|
||||
num_heads: int,
|
||||
num_layers: int,
|
||||
layer_scale: float,
|
||||
context: int,
|
||||
max_period: float,
|
||||
dim_feedforward: int,
|
||||
):
|
||||
super().__init__()
|
||||
self.transformer = StreamingTransformer(
|
||||
d_model=d_model,
|
||||
num_heads=num_heads,
|
||||
num_layers=num_layers,
|
||||
layer_scale=layer_scale,
|
||||
context=context,
|
||||
max_period=max_period,
|
||||
dim_feedforward=dim_feedforward,
|
||||
)
|
||||
self.input_dimension = input_dimension
|
||||
self.output_dimensions = output_dimensions
|
||||
self.input_proj = None
|
||||
if d_model != input_dimension:
|
||||
self.input_proj = nn.Linear(input_dimension, d_model, bias=False)
|
||||
|
||||
self.output_projs = nn.ModuleList()
|
||||
for output_dimension in output_dimensions:
|
||||
if d_model == output_dimension:
|
||||
self.output_projs.append(nn.Identity())
|
||||
else:
|
||||
self.output_projs.append(nn.Linear(d_model, output_dimension, bias=False))
|
||||
|
||||
def forward(self, x, model_state: dict | None):
|
||||
x = x.transpose(1, 2)
|
||||
if self.input_proj is not None:
|
||||
x = self.input_proj(x)
|
||||
z = self.transformer(x, model_state)
|
||||
ys = []
|
||||
for output_proj in self.output_projs:
|
||||
y = output_proj(z)
|
||||
y = y.transpose(1, 2)
|
||||
ys.append(y)
|
||||
return ys
|
||||
@@ -0,0 +1,215 @@
|
||||
"""
|
||||
Taken from
|
||||
https://github.com/LTH14/mar/blob/fe470ac24afbee924668d8c5c83e9fec60af3a73/models/diffloss.py
|
||||
|
||||
"""
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from typing_extensions import Self
|
||||
|
||||
from pocket_tts.utils.config import FlowLMConfig
|
||||
|
||||
|
||||
def modulate(x, shift, scale):
|
||||
return x * (1 + scale) + shift
|
||||
|
||||
|
||||
def _rms_norm(x: torch.Tensor, alpha: torch.Tensor, eps: float):
|
||||
assert x.dim() >= alpha.dim()
|
||||
x_dtype = x.dtype
|
||||
var = eps + x.var(dim=-1, keepdim=True)
|
||||
y = (x * (alpha.to(var) * torch.rsqrt(var))).to(x_dtype)
|
||||
return y
|
||||
|
||||
|
||||
class RMSNorm(nn.Module):
|
||||
def __init__(self, dim: int, eps: float = 1e-5):
|
||||
super().__init__()
|
||||
self.eps = eps
|
||||
alpha_shape = (dim,)
|
||||
self.alpha = nn.Parameter(torch.full(alpha_shape, 1.0, requires_grad=True))
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
return _rms_norm(x, self.alpha, self.eps)
|
||||
|
||||
|
||||
class LayerNorm(nn.Module):
|
||||
"""Reimplementation of LayerNorm because the default one doesn't support jvp."""
|
||||
|
||||
def __init__(self, channels, eps=1e-6, elementwise_affine=True):
|
||||
super().__init__()
|
||||
self.eps = eps
|
||||
if elementwise_affine:
|
||||
self.weight = nn.Parameter(torch.ones(channels))
|
||||
self.bias = nn.Parameter(torch.zeros(channels))
|
||||
|
||||
def forward(self, x):
|
||||
mean = x.mean(dim=-1, keepdim=True)
|
||||
var = x.var(dim=-1, unbiased=False, keepdim=True)
|
||||
x = (x - mean) / torch.sqrt(var + self.eps)
|
||||
if hasattr(self, "weight"):
|
||||
x = x * self.weight + self.bias
|
||||
return x
|
||||
|
||||
|
||||
class TimestepEmbedder(nn.Module):
|
||||
"""Embeds scalar timesteps into vector representations."""
|
||||
|
||||
def __init__(
|
||||
self, hidden_size: int, frequency_embedding_size: int = 256, max_period: int = 10000
|
||||
):
|
||||
super().__init__()
|
||||
blocks = [
|
||||
nn.Linear(frequency_embedding_size, hidden_size, bias=True),
|
||||
nn.SiLU(),
|
||||
nn.Linear(hidden_size, hidden_size, bias=True),
|
||||
]
|
||||
blocks.append(RMSNorm(hidden_size))
|
||||
self.mlp = nn.Sequential(*blocks)
|
||||
self.frequency_embedding_size = frequency_embedding_size
|
||||
half = frequency_embedding_size // 2
|
||||
self.register_buffer(
|
||||
"freqs", torch.exp(-math.log(max_period) * torch.arange(start=0, end=half) / half)
|
||||
)
|
||||
|
||||
def forward(self, t):
|
||||
args = t * self.freqs.to(t.dtype)
|
||||
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
|
||||
assert not (self.frequency_embedding_size % 2)
|
||||
t_emb = self.mlp(embedding)
|
||||
return t_emb
|
||||
|
||||
|
||||
class ResBlock(nn.Module):
|
||||
"""
|
||||
A residual block that can optionally change the number of channels.
|
||||
:param channels: the number of input channels.
|
||||
"""
|
||||
|
||||
def __init__(self, channels):
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
|
||||
self.in_ln = LayerNorm(channels, eps=1e-6)
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(channels, channels, bias=True),
|
||||
nn.SiLU(),
|
||||
nn.Linear(channels, channels, bias=True),
|
||||
)
|
||||
|
||||
self.adaLN_modulation = nn.Sequential(
|
||||
nn.SiLU(), nn.Linear(channels, 3 * channels, bias=True)
|
||||
)
|
||||
|
||||
def forward(self, x, y):
|
||||
shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(y).chunk(3, dim=-1)
|
||||
h = modulate(self.in_ln(x), shift_mlp, scale_mlp)
|
||||
h = self.mlp(h)
|
||||
return x + gate_mlp * h
|
||||
|
||||
|
||||
class FinalLayer(nn.Module):
|
||||
"""
|
||||
The final layer adopted from DiT.
|
||||
"""
|
||||
|
||||
def __init__(self, model_channels, out_channels):
|
||||
super().__init__()
|
||||
self.norm_final = LayerNorm(model_channels, elementwise_affine=False, eps=1e-6)
|
||||
self.linear = nn.Linear(model_channels, out_channels, bias=True)
|
||||
self.adaLN_modulation = nn.Sequential(
|
||||
nn.SiLU(), nn.Linear(model_channels, 2 * model_channels, bias=True)
|
||||
)
|
||||
|
||||
def forward(self, x, c):
|
||||
shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1)
|
||||
x = modulate(self.norm_final(x), shift, scale)
|
||||
x = self.linear(x)
|
||||
return x
|
||||
|
||||
|
||||
class SimpleMLPAdaLN(nn.Module):
|
||||
"""Taken from https://arxiv.org/abs/2406.11838.
|
||||
|
||||
The MLP for Diffusion Loss.
|
||||
:param in_channels: channels in the input Tensor.
|
||||
:param model_channels: base channel count for the model.
|
||||
:param out_channels: channels in the output Tensor.
|
||||
:param cond_channels: channels in the condition.
|
||||
:param num_res_blocks: number of residual blocks per downsample.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
model_channels,
|
||||
out_channels,
|
||||
cond_channels,
|
||||
num_res_blocks,
|
||||
num_time_conds=1,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.in_channels = in_channels
|
||||
self.model_channels = model_channels
|
||||
self.out_channels = out_channels
|
||||
self.num_res_blocks = num_res_blocks
|
||||
self.num_time_conds = num_time_conds
|
||||
|
||||
assert num_time_conds != 1
|
||||
self.time_embed = nn.ModuleList(
|
||||
[TimestepEmbedder(model_channels) for _ in range(num_time_conds)]
|
||||
)
|
||||
self.cond_embed = nn.Linear(cond_channels, model_channels)
|
||||
|
||||
self.input_proj = nn.Linear(in_channels, model_channels)
|
||||
|
||||
res_blocks = []
|
||||
for i in range(num_res_blocks):
|
||||
res_blocks.append(ResBlock(model_channels))
|
||||
|
||||
self.res_blocks = nn.ModuleList(res_blocks)
|
||||
self.final_layer = FinalLayer(model_channels, out_channels)
|
||||
|
||||
@classmethod
|
||||
def from_pydantic_config(cls, cfg: FlowLMConfig, latent_dim: int, cond_dim: int) -> Self:
|
||||
config = cfg.flow
|
||||
|
||||
flow_dim = config.dim
|
||||
flow_depth = config.depth
|
||||
num_time_conds = 2
|
||||
return SimpleMLPAdaLN(
|
||||
latent_dim, flow_dim, latent_dim, cond_dim, flow_depth, num_time_conds=num_time_conds
|
||||
)
|
||||
|
||||
def forward(
|
||||
self, c: torch.Tensor, s: torch.Tensor, t: torch.Tensor, x: torch.Tensor
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Apply the model to an input batch.
|
||||
:param c: conditioning from AR transformer.
|
||||
:param s: start time tensor.
|
||||
:param t: target time tensor.
|
||||
:param x: an [N x C] Tensor of inputs.
|
||||
:return: an [N x C] Tensor of outputs.
|
||||
"""
|
||||
# Combine time conditions
|
||||
ts = [s, t]
|
||||
x = self.input_proj(x)
|
||||
assert len(ts) == self.num_time_conds, (
|
||||
f"Expected {self.num_time_conds} time conditions, got {len(ts)}"
|
||||
)
|
||||
assert self.num_time_conds != 1
|
||||
t_combined = (
|
||||
sum(self.time_embed[i](ts[i]) for i in range(self.num_time_conds)) / self.num_time_conds
|
||||
)
|
||||
c = self.cond_embed(c)
|
||||
y = t_combined + c
|
||||
|
||||
for block in self.res_blocks:
|
||||
x = block(x, y)
|
||||
|
||||
return self.final_layer(x, y)
|
||||
@@ -0,0 +1,46 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from pocket_tts.modules.conv import StreamingConv1d, StreamingConvTranspose1d
|
||||
|
||||
|
||||
class ConvDownsample1d(nn.Module):
|
||||
"""
|
||||
Downsampling by some integer amount `stride` using convolutions
|
||||
with a kernel size of twice the stride.
|
||||
"""
|
||||
|
||||
def __init__(self, stride: int, dimension: int):
|
||||
super().__init__()
|
||||
self.conv = StreamingConv1d(
|
||||
dimension,
|
||||
dimension,
|
||||
kernel_size=2 * stride,
|
||||
stride=stride,
|
||||
groups=1,
|
||||
bias=False,
|
||||
pad_mode="replicate",
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor, model_state: dict | None):
|
||||
return self.conv(x, model_state)
|
||||
|
||||
|
||||
class ConvTrUpsample1d(nn.Module):
|
||||
"""
|
||||
Upsample by some integer amount `stride` using transposed convolutions.
|
||||
"""
|
||||
|
||||
def __init__(self, stride: int, dimension: int):
|
||||
super().__init__()
|
||||
self.convtr = StreamingConvTranspose1d(
|
||||
dimension,
|
||||
dimension,
|
||||
kernel_size=2 * stride,
|
||||
stride=stride,
|
||||
groups=dimension,
|
||||
bias=False,
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor, model_state: dict | None):
|
||||
return self.convtr(x, model_state)
|
||||
@@ -0,0 +1,74 @@
|
||||
import math
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
def apply_rope(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
offset: int | torch.Tensor = 0,
|
||||
max_period: int | float = 10_000,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
q (torch.Tensor): Queries, shape `[B, T, H, D]`.
|
||||
k (torch.Tensor): Keys, shape `[B, T, H, D]`.
|
||||
offset (int): Current offset, e.g. when streaming.
|
||||
max_period (float): Maximum period for the cos and sin.
|
||||
"""
|
||||
|
||||
B, T, H, D = q.shape
|
||||
Bk, Tk, Hk, Dk = k.shape
|
||||
assert (B, T, D) == (Bk, Tk, Dk)
|
||||
assert D > 0
|
||||
assert D % 2 == 0
|
||||
assert max_period > 0
|
||||
|
||||
ds = torch.arange(D // 2, device=q.device, dtype=torch.float32)
|
||||
freqs = torch.exp(ds * (-math.log(max_period) * 2 / D))
|
||||
|
||||
# could be optimized in one call
|
||||
ts = torch.arange(T, device=q.device, dtype=torch.float32)
|
||||
ts += offset
|
||||
ts = ts.view(-1, 1, 1)
|
||||
|
||||
q = q.view(B, T, H, D // 2, 2)
|
||||
k = k.view(B, T, Hk, D // 2, 2)
|
||||
|
||||
# convention is `r` suffix is real part, `i` is imaginary.
|
||||
qr = q[..., 0].float()
|
||||
qi = q[..., 1].float()
|
||||
|
||||
kr = k[..., 0].float()
|
||||
ki = k[..., 1].float()
|
||||
|
||||
rotr = torch.cos(freqs * ts)
|
||||
roti = torch.sin(freqs * ts)
|
||||
qor = qr * rotr - qi * roti
|
||||
qoi = qr * roti + qi * rotr
|
||||
|
||||
kor = kr * rotr - ki * roti
|
||||
koi = kr * roti + ki * rotr
|
||||
|
||||
dtype = q.dtype
|
||||
qo = torch.stack([qor.to(dtype), qoi.to(dtype)], dim=-1)
|
||||
ko = torch.stack([kor.to(dtype), koi.to(dtype)], dim=-1)
|
||||
|
||||
return qo.view(B, T, H, D), ko.view(B, T, Hk, D)
|
||||
|
||||
|
||||
class RotaryEmbedding(nn.Module):
|
||||
"""Rotary positional embedding (RoPE) from [Su et al 2022](https://arxiv.org/abs/2104.09864).
|
||||
|
||||
Args:
|
||||
max_period (float): Maximum period of the rotation frequencies.
|
||||
"""
|
||||
|
||||
def __init__(self, max_period: float | int = 10000.0):
|
||||
super().__init__()
|
||||
self.max_period = max_period
|
||||
|
||||
def forward(self, q: torch.Tensor, k: torch.Tensor, offset: torch.Tensor | int):
|
||||
"""Apply rope rotation to query or key tensor."""
|
||||
return apply_rope(q, k, offset, self.max_period)
|
||||
@@ -0,0 +1,180 @@
|
||||
import numpy as np
|
||||
import torch.nn as nn
|
||||
|
||||
from .conv import StreamingConv1d, StreamingConvTranspose1d
|
||||
|
||||
|
||||
class SEANetResnetBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
kernel_sizes: list[int] = [3, 1],
|
||||
dilations: list[int] = [1, 1],
|
||||
pad_mode: str = "reflect",
|
||||
compress: int = 2,
|
||||
):
|
||||
super().__init__()
|
||||
assert len(kernel_sizes) == len(dilations), (
|
||||
"Number of kernel sizes should match number of dilations"
|
||||
)
|
||||
hidden = dim // compress
|
||||
block = nn.ModuleList([])
|
||||
for i, (kernel_size, dilation) in enumerate(zip(kernel_sizes, dilations)):
|
||||
in_chs = dim if i == 0 else hidden
|
||||
out_chs = dim if i == len(kernel_sizes) - 1 else hidden
|
||||
block += [
|
||||
nn.ELU(alpha=1.0),
|
||||
StreamingConv1d(
|
||||
in_chs, out_chs, kernel_size=kernel_size, dilation=dilation, pad_mode=pad_mode
|
||||
),
|
||||
]
|
||||
self.block = block
|
||||
|
||||
def forward(self, x, model_state: dict | None):
|
||||
v = x
|
||||
for layer in self.block:
|
||||
if isinstance(layer, StreamingConv1d):
|
||||
v = layer(v, model_state)
|
||||
else:
|
||||
v = layer(v)
|
||||
assert x.shape == v.shape, (x.shape, v.shape, x.shape)
|
||||
return x + v
|
||||
|
||||
|
||||
class SEANetEncoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
channels: int = 1,
|
||||
dimension: int = 128,
|
||||
n_filters: int = 32,
|
||||
n_residual_layers: int = 3,
|
||||
ratios: list[int] = [8, 5, 4, 2],
|
||||
kernel_size: int = 7,
|
||||
last_kernel_size: int = 7,
|
||||
residual_kernel_size: int = 3,
|
||||
dilation_base: int = 2,
|
||||
pad_mode: str = "reflect",
|
||||
compress: int = 2,
|
||||
):
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.dimension = dimension
|
||||
self.n_filters = n_filters
|
||||
self.ratios = list(reversed(ratios))
|
||||
del ratios
|
||||
self.n_residual_layers = n_residual_layers
|
||||
self.hop_length = int(np.prod(self.ratios))
|
||||
self.n_blocks = len(self.ratios) + 2 # first and last conv + residual blocks
|
||||
|
||||
mult = 1
|
||||
model = nn.ModuleList(
|
||||
[StreamingConv1d(channels, mult * n_filters, kernel_size, pad_mode=pad_mode)]
|
||||
)
|
||||
# Downsample to raw audio scale
|
||||
for i, ratio in enumerate(self.ratios):
|
||||
# Add residual layers
|
||||
for j in range(n_residual_layers):
|
||||
model += [
|
||||
SEANetResnetBlock(
|
||||
mult * n_filters,
|
||||
kernel_sizes=[residual_kernel_size, 1],
|
||||
dilations=[dilation_base**j, 1],
|
||||
pad_mode=pad_mode,
|
||||
compress=compress,
|
||||
)
|
||||
]
|
||||
|
||||
# Add downsampling layers
|
||||
model += [
|
||||
nn.ELU(alpha=1.0),
|
||||
StreamingConv1d(
|
||||
mult * n_filters,
|
||||
mult * n_filters * 2,
|
||||
kernel_size=ratio * 2,
|
||||
stride=ratio,
|
||||
pad_mode=pad_mode,
|
||||
),
|
||||
]
|
||||
mult *= 2
|
||||
|
||||
model += [
|
||||
nn.ELU(alpha=1.0),
|
||||
StreamingConv1d(mult * n_filters, dimension, last_kernel_size, pad_mode=pad_mode),
|
||||
]
|
||||
|
||||
self.model = model
|
||||
|
||||
def forward(self, x, model_state: dict | None):
|
||||
for layer in self.model:
|
||||
if isinstance(layer, (StreamingConv1d, SEANetResnetBlock)):
|
||||
x = layer(x, model_state)
|
||||
else:
|
||||
x = layer(x)
|
||||
return x
|
||||
|
||||
|
||||
class SEANetDecoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
channels: int = 1,
|
||||
dimension: int = 128,
|
||||
n_filters: int = 32,
|
||||
n_residual_layers: int = 3,
|
||||
ratios: list[int] = [8, 5, 4, 2],
|
||||
kernel_size: int = 7,
|
||||
last_kernel_size: int = 7,
|
||||
residual_kernel_size: int = 3,
|
||||
dilation_base: int = 2,
|
||||
pad_mode: str = "reflect",
|
||||
compress: int = 2,
|
||||
):
|
||||
super().__init__()
|
||||
self.dimension = dimension
|
||||
self.channels = channels
|
||||
self.n_filters = n_filters
|
||||
self.ratios = ratios
|
||||
del ratios
|
||||
self.n_residual_layers = n_residual_layers
|
||||
self.hop_length = int(np.prod(self.ratios))
|
||||
self.n_blocks = len(self.ratios) + 2 # first and last conv + residual blocks
|
||||
mult = int(2 ** len(self.ratios))
|
||||
model = nn.ModuleList(
|
||||
[StreamingConv1d(dimension, mult * n_filters, kernel_size, pad_mode=pad_mode)]
|
||||
)
|
||||
# Upsample to raw audio scale
|
||||
for _, ratio in enumerate(self.ratios):
|
||||
# Add upsampling layers
|
||||
model += [
|
||||
nn.ELU(alpha=1.0),
|
||||
StreamingConvTranspose1d(
|
||||
mult * n_filters, mult * n_filters // 2, kernel_size=ratio * 2, stride=ratio
|
||||
),
|
||||
]
|
||||
# Add residual layers
|
||||
for j in range(n_residual_layers):
|
||||
model += [
|
||||
SEANetResnetBlock(
|
||||
mult * n_filters // 2,
|
||||
kernel_sizes=[residual_kernel_size, 1],
|
||||
dilations=[dilation_base**j, 1],
|
||||
pad_mode=pad_mode,
|
||||
compress=compress,
|
||||
)
|
||||
]
|
||||
|
||||
mult //= 2
|
||||
|
||||
# Add final layers
|
||||
model += [
|
||||
nn.ELU(alpha=1.0),
|
||||
StreamingConv1d(n_filters, channels, last_kernel_size, pad_mode=pad_mode),
|
||||
]
|
||||
self.model = model
|
||||
|
||||
def forward(self, z, model_state: dict | None):
|
||||
for layer in self.model:
|
||||
if isinstance(layer, (StreamingConvTranspose1d, SEANetResnetBlock, StreamingConv1d)):
|
||||
z = layer(z, model_state)
|
||||
else:
|
||||
z = layer(z)
|
||||
return z
|
||||
@@ -0,0 +1,45 @@
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
def init_states(
|
||||
model: nn.Module, batch_size: int, sequence_length: int
|
||||
) -> dict[str, dict[str, torch.Tensor]]:
|
||||
result = {}
|
||||
for module_name, module in model.named_modules():
|
||||
if not isinstance(module, StatefulModule):
|
||||
continue
|
||||
module._module_absolute_name = module_name
|
||||
module_state = module.init_state(batch_size, sequence_length=sequence_length)
|
||||
result[module_name] = module_state
|
||||
return result
|
||||
|
||||
|
||||
def increment_steps(
|
||||
module: nn.Module, model_state: dict[str, dict[str, torch.Tensor]], increment: int = 1
|
||||
):
|
||||
# print("incrementing steps by", increment)
|
||||
for module_name, module in module.named_modules():
|
||||
if not isinstance(module, StatefulModule):
|
||||
continue
|
||||
module.increment_step(model_state[module_name], increment)
|
||||
|
||||
|
||||
class StatefulModule(ABC, nn.Module):
|
||||
def __init__(self, *args, **kwds):
|
||||
self._module_absolute_name = None
|
||||
return super().__init__(*args, **kwds)
|
||||
|
||||
@abstractmethod
|
||||
def init_state(self, batch_size: int, sequence_length: int):
|
||||
"""Initialize the state."""
|
||||
raise NotImplementedError
|
||||
|
||||
def increment_step(self, state: dict, increment: int = 1):
|
||||
pass
|
||||
|
||||
def get_state(self, model_state: dict[str, dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
|
||||
"""Get the state for this module from the model state."""
|
||||
return model_state[self._module_absolute_name]
|
||||
@@ -0,0 +1,126 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
from pocket_tts.modules.rope import RotaryEmbedding
|
||||
from pocket_tts.modules.stateful_module import StatefulModule
|
||||
|
||||
|
||||
def complete_kv(
|
||||
cache: torch.Tensor, current_end: torch.Tensor, k: torch.Tensor, v: torch.Tensor
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
current_end = current_end.shape[0]
|
||||
|
||||
cache[0, :, current_end : current_end + k.shape[1]] = k
|
||||
cache[1, :, current_end : current_end + v.shape[1]] = v
|
||||
valid = cache[:, :, : current_end + k.shape[1]]
|
||||
return valid[0], valid[1]
|
||||
|
||||
|
||||
def _materialize_causal_mask(
|
||||
shape: tuple[int, ...], shift: int, device: str | torch.device = "cpu"
|
||||
) -> torch.Tensor:
|
||||
dtype = torch.float32
|
||||
|
||||
num_queries, num_keys = shape[-2:]
|
||||
shift = num_keys - num_queries
|
||||
|
||||
tensor = torch.full(shape, dtype=dtype, fill_value=1, device=device)
|
||||
mask = torch.tril(tensor, diagonal=shift).to(dtype)
|
||||
mask = torch.log(mask)
|
||||
return mask.to(dtype)
|
||||
|
||||
|
||||
class StreamingMultiheadAttention(StatefulModule):
|
||||
"""Similar to `nn.MultiheadAttention` but with support for streaming.
|
||||
|
||||
Args:
|
||||
embed_dim (int): Dimension to project to.
|
||||
num_heads (int): Number of heads.
|
||||
context (int, optional): Number of time steps the attention can access to.
|
||||
Can access `context` time steps into the past.
|
||||
rope (`RotaryEmbedding`, optional): Rope embedding to use.
|
||||
device (torch.device, optional): Device on which to initialize.
|
||||
dtype (torch.dtype, optional): dtype to use.
|
||||
"""
|
||||
|
||||
def __init__(self, embed_dim: int, num_heads: int, rope: RotaryEmbedding):
|
||||
super().__init__()
|
||||
|
||||
self.embed_dim = embed_dim
|
||||
self.rope = rope
|
||||
self.num_heads = num_heads
|
||||
|
||||
out_dim = embed_dim
|
||||
num_kv = num_heads
|
||||
kv_dim = (embed_dim // num_heads) * num_kv
|
||||
out_dim += 2 * kv_dim
|
||||
mult = 1
|
||||
in_proj = nn.Linear(embed_dim, mult * out_dim, bias=False)
|
||||
# We try to follow the default PyTorch MHA convention, to easily compare results.
|
||||
self.in_proj_weight = in_proj.weight
|
||||
self.out_proj = nn.Linear(embed_dim, mult * embed_dim, bias=False)
|
||||
|
||||
def _get_mask(self, shape: tuple[int, int], shift: int, device: torch.device) -> torch.Tensor:
|
||||
return _materialize_causal_mask(shape, shift=shift, device=device)
|
||||
|
||||
def init_state(self, batch_size: int, sequence_length: int) -> dict[str, torch.Tensor]:
|
||||
dim_per_head = self.embed_dim // self.num_heads
|
||||
initial_current_end = torch.zeros((0,)).to(self.in_proj_weight.device)
|
||||
return dict(
|
||||
current_end=initial_current_end,
|
||||
cache=torch.full(
|
||||
(2, batch_size, sequence_length, self.num_heads, dim_per_head),
|
||||
float("NaN"),
|
||||
device=self.in_proj_weight.device,
|
||||
dtype=self.in_proj_weight.dtype,
|
||||
),
|
||||
)
|
||||
|
||||
def increment_step(self, state: dict, increment: int = 1):
|
||||
new_size = state["current_end"].shape[0] + increment
|
||||
state["current_end"] = torch.zeros((new_size,)).to(state["current_end"].device)
|
||||
|
||||
def _complete_kv(self, k, v, state: dict | None):
|
||||
k, v = complete_kv(state["cache"], state["current_end"], k, v)
|
||||
return k, v
|
||||
|
||||
def _apply_rope(self, query: torch.Tensor, key: torch.Tensor, state: dict | None):
|
||||
# Apply rope embeddings to query and key tensors.
|
||||
streaming_offset = self._streaming_offset(state)
|
||||
return self.rope(query, key, offset=streaming_offset)
|
||||
|
||||
def _streaming_offset(self, state: dict | None) -> torch.Tensor | int:
|
||||
return state["current_end"].shape[0]
|
||||
|
||||
def check_model_state(self, model_state: dict):
|
||||
if model_state is None:
|
||||
raise ValueError("model_state must be provided")
|
||||
return self.get_state(model_state)
|
||||
|
||||
def forward(self, query: torch.Tensor, model_state: dict | None):
|
||||
state = self.check_model_state(model_state)
|
||||
|
||||
projected = nn.functional.linear(query, self.in_proj_weight)
|
||||
# Reshape from (b, t, p*h*d) to (b, t, p, h, d) where p=3, h=num_heads
|
||||
b, t, _ = projected.shape
|
||||
d = self.embed_dim // self.num_heads
|
||||
packed = projected.view(b, t, 3, self.num_heads, d)
|
||||
q, k, v = torch.unbind(packed, dim=2)
|
||||
q, k = self._apply_rope(q, k, state)
|
||||
k, v = self._complete_kv(k, v, state)
|
||||
|
||||
mask_shape = (query.shape[1], query.shape[1] + state["current_end"].shape[0])
|
||||
shift = state["current_end"].shape[0]
|
||||
|
||||
attn_mask = self._get_mask(mask_shape, shift=shift, device=q.device)
|
||||
|
||||
q, k, v = [x.transpose(1, 2) for x in (q, k, v)]
|
||||
x = F.scaled_dot_product_attention(q, k, v, attn_mask)
|
||||
x = x.transpose(1, 2)
|
||||
# Reshape from (b, t, h, d) to (b, t, h*d)
|
||||
b, t, h, d = x.shape
|
||||
x = x.reshape(b, t, h * d)
|
||||
x = self.out_proj(x)
|
||||
|
||||
return x
|
||||
@@ -0,0 +1,374 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>Pocket TTS - Streaming</title>
|
||||
<script src="https://cdn.tailwindcss.com"></script>
|
||||
<style>
|
||||
body {
|
||||
font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
|
||||
}
|
||||
.spinner {
|
||||
animation: spin 1s linear infinite;
|
||||
}
|
||||
@keyframes spin {
|
||||
from { transform: rotate(0deg); }
|
||||
to { transform: rotate(360deg); }
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body class="bg-gray-50 min-h-screen">
|
||||
<div class="max-w-xl mx-auto p-4 space-y-4">
|
||||
<h1 class="text-xl font-bold text-center">Pocket TTS</h1>
|
||||
|
||||
<div class="bg-white p-4 rounded-lg shadow">
|
||||
<textarea
|
||||
id="text-input"
|
||||
class="w-full border rounded p-2"
|
||||
placeholder="Enter your text here..."
|
||||
rows="4"
|
||||
>Hello world. I am Kyutai's Pocket TTS. I'm fast enough to run on small CPUs. I hope you'll like me.</textarea>
|
||||
</div>
|
||||
|
||||
<div class="bg-white p-4 rounded-lg shadow">
|
||||
<label for="voice-url-input" class="block text-sm font-medium text-gray-700 mb-2">
|
||||
Optional voice URL (leave empty to use default voice):
|
||||
</label>
|
||||
<input
|
||||
type="text"
|
||||
id="voice-url-input"
|
||||
class="w-full border rounded p-2"
|
||||
placeholder="hf://kyutai/tts-voices/alba-mackenna/casual.wav"
|
||||
value="marius"
|
||||
/>
|
||||
<p class="text-xs text-gray-500 mt-1">
|
||||
Supports: http://, https://, or hf:// URLs.<br>
|
||||
You can also use predefined voices:<br>
|
||||
"marius", "javert", "jean", "fantine", "cosette", "eponine", "azelma".
|
||||
</p>
|
||||
</div>
|
||||
|
||||
<div class="bg-white p-4 rounded-lg shadow">
|
||||
<label for="voice-wav-input" class="block text-sm font-medium text-gray-700 mb-2">
|
||||
Or upload a WAV file for voice cloning:
|
||||
</label>
|
||||
<input
|
||||
type="file"
|
||||
id="voice-wav-input"
|
||||
class="w-full border rounded p-2"
|
||||
accept=".wav,audio/wav"
|
||||
/>
|
||||
<p class="text-xs text-gray-500 mt-1">
|
||||
Upload a WAV file to use as voice reference. Takes precedence over voice URL.
|
||||
</p>
|
||||
</div>
|
||||
|
||||
|
||||
<button
|
||||
id="generate-btn"
|
||||
class="w-full px-4 py-2 rounded text-white bg-blue-600 hover:bg-blue-700 disabled:bg-gray-400 disabled:cursor-not-allowed"
|
||||
>
|
||||
<span id="generate-text">Generate audio</span>
|
||||
</button>
|
||||
|
||||
<div id="status" class="hidden">
|
||||
<div class="flex items-center space-x-2">
|
||||
<div id="status-spinner" class="spinner w-4 h-4 border-2 border-blue-600 border-t-transparent rounded-full"></div>
|
||||
<span id="status-text"></span>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div id="audio-section" class="hidden">
|
||||
<div class="bg-white p-4 rounded-lg shadow">
|
||||
<h3 class="text-lg font-semibold mb-2">Generated Audio</h3>
|
||||
<audio id="audio" controls class="w-full"></audio>
|
||||
<div class="mt-2">
|
||||
<a id="download-link" class="text-blue-600 text-sm hover:underline">Download</a>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<script>
|
||||
class StreamingWavPlayer {
|
||||
constructor() {
|
||||
this.audioContext = new (window.AudioContext || window.webkitAudioContext)();
|
||||
this.sampleRate = 0;
|
||||
this.numChannels = 0;
|
||||
this.headerParsed = false;
|
||||
this.headerBuffer = new Uint8Array(44);
|
||||
this.headerBytesReceived = 0;
|
||||
this.nextStartTime = 0;
|
||||
this.isPlaying = false;
|
||||
this.minBufferSize = 16384;
|
||||
this.pcmData = new Uint8Array(0);
|
||||
}
|
||||
|
||||
parseWavHeader(header) {
|
||||
const view = new DataView(header.buffer);
|
||||
|
||||
const riff = String.fromCharCode.apply(null, Array.from(header.slice(0, 4)));
|
||||
const wave = String.fromCharCode.apply(null, Array.from(header.slice(8, 12)));
|
||||
|
||||
if (riff !== 'RIFF' || wave !== 'WAVE') {
|
||||
throw new Error('Invalid WAV file');
|
||||
}
|
||||
|
||||
this.numChannels = view.getUint16(22, true);
|
||||
this.sampleRate = view.getUint32(24, true);
|
||||
const bitsPerSample = view.getUint16(34, true);
|
||||
|
||||
console.log(`WAV Format: ${this.sampleRate}Hz, ${this.numChannels} channels, ${bitsPerSample} bits`);
|
||||
|
||||
this.headerParsed = true;
|
||||
}
|
||||
|
||||
appendPcmData(newData) {
|
||||
const newBuffer = new Uint8Array(this.pcmData.length + newData.length);
|
||||
newBuffer.set(this.pcmData);
|
||||
newBuffer.set(newData, this.pcmData.length);
|
||||
this.pcmData = newBuffer;
|
||||
}
|
||||
|
||||
async tryPlayBuffer() {
|
||||
if (!this.headerParsed || this.pcmData.length < this.minBufferSize) {
|
||||
return;
|
||||
}
|
||||
|
||||
const bytesPerSample = this.numChannels * 2;
|
||||
const samplesToPlay = Math.floor(this.pcmData.length / bytesPerSample);
|
||||
const bytesToPlay = samplesToPlay * bytesPerSample;
|
||||
|
||||
if (bytesToPlay === 0) return;
|
||||
|
||||
const dataToPlay = this.pcmData.slice(0, bytesToPlay);
|
||||
this.pcmData = this.pcmData.slice(bytesToPlay);
|
||||
|
||||
const audioBuffer = this.audioContext.createBuffer(
|
||||
this.numChannels,
|
||||
samplesToPlay,
|
||||
this.sampleRate
|
||||
);
|
||||
|
||||
const int16Data = new Int16Array(dataToPlay.buffer, dataToPlay.byteOffset, samplesToPlay * this.numChannels);
|
||||
|
||||
for (let channel = 0; channel < this.numChannels; channel++) {
|
||||
const channelData = audioBuffer.getChannelData(channel);
|
||||
for (let i = 0; i < samplesToPlay; i++) {
|
||||
channelData[i] = int16Data[i * this.numChannels + channel] / 32768;
|
||||
}
|
||||
}
|
||||
|
||||
const source = this.audioContext.createBufferSource();
|
||||
source.buffer = audioBuffer;
|
||||
source.connect(this.audioContext.destination);
|
||||
|
||||
const currentTime = this.audioContext.currentTime;
|
||||
const startTime = Math.max(currentTime, this.nextStartTime);
|
||||
|
||||
source.start(startTime);
|
||||
|
||||
// Track first audio playback
|
||||
if (!this.firstAudioPlayed && window.firstAudioCallback) {
|
||||
this.firstAudioPlayed = true;
|
||||
window.firstAudioCallback();
|
||||
}
|
||||
|
||||
this.nextStartTime = startTime + audioBuffer.duration;
|
||||
this.isPlaying = true;
|
||||
|
||||
if (this.pcmData.length >= this.minBufferSize) {
|
||||
setTimeout(() => this.tryPlayBuffer(), 10);
|
||||
}
|
||||
}
|
||||
|
||||
addChunk(chunk) {
|
||||
if (!this.headerParsed) {
|
||||
const headerBytesNeeded = 44 - this.headerBytesReceived;
|
||||
const bytesToCopy = Math.min(headerBytesNeeded, chunk.length);
|
||||
|
||||
this.headerBuffer.set(
|
||||
chunk.slice(0, bytesToCopy),
|
||||
this.headerBytesReceived
|
||||
);
|
||||
|
||||
this.headerBytesReceived += bytesToCopy;
|
||||
|
||||
if (this.headerBytesReceived >= 44) {
|
||||
this.parseWavHeader(this.headerBuffer);
|
||||
|
||||
if (chunk.length > bytesToCopy) {
|
||||
this.appendPcmData(chunk.slice(bytesToCopy));
|
||||
}
|
||||
}
|
||||
} else {
|
||||
this.appendPcmData(chunk);
|
||||
}
|
||||
|
||||
this.tryPlayBuffer();
|
||||
}
|
||||
|
||||
stop() {
|
||||
this.audioContext.close();
|
||||
this.isPlaying = false;
|
||||
}
|
||||
}
|
||||
|
||||
// Application state
|
||||
let streamingPlayer = null;
|
||||
let currentAudioBlob = null;
|
||||
|
||||
// DOM elements
|
||||
const textInput = document.getElementById('text-input');
|
||||
const generateBtn = document.getElementById('generate-btn');
|
||||
const generateText = document.getElementById('generate-text');
|
||||
const status = document.getElementById('status');
|
||||
const statusText = document.getElementById('status-text');
|
||||
const statusSpinner = document.getElementById('status-spinner');
|
||||
const audioSection = document.getElementById('audio-section');
|
||||
const audio = document.getElementById('audio');
|
||||
const downloadLink = document.getElementById('download-link');
|
||||
|
||||
// Event listeners
|
||||
generateBtn.addEventListener('click', generateAudio);
|
||||
|
||||
function showStatus(message, isLoading = false) {
|
||||
statusText.textContent = message;
|
||||
statusSpinner.style.display = isLoading ? 'block' : 'none';
|
||||
status.classList.remove('hidden');
|
||||
}
|
||||
|
||||
function hideStatus() {
|
||||
status.classList.add('hidden');
|
||||
}
|
||||
|
||||
async function generateAudio() {
|
||||
const text = textInput.value.trim();
|
||||
|
||||
if (!text) {
|
||||
showStatus('Please enter some text to generate speech.', false);
|
||||
setTimeout(hideStatus, 3000);
|
||||
return;
|
||||
}
|
||||
|
||||
// Stop any currently playing audio
|
||||
if (streamingPlayer) {
|
||||
streamingPlayer.stop();
|
||||
streamingPlayer = null;
|
||||
}
|
||||
|
||||
// Track timing
|
||||
const startTime = performance.now();
|
||||
let firstAudioTime = null;
|
||||
|
||||
// Set callback for first audio
|
||||
window.firstAudioCallback = () => {
|
||||
if (!firstAudioTime) {
|
||||
firstAudioTime = performance.now();
|
||||
const timeToFirstAudio = ((firstAudioTime - startTime) / 1000).toFixed(2);
|
||||
showStatus(`First audio in ${timeToFirstAudio}s...`, true);
|
||||
}
|
||||
};
|
||||
|
||||
// Update UI
|
||||
generateBtn.disabled = true;
|
||||
generateText.textContent = 'Generating...';
|
||||
showStatus('Generating speech...', true);
|
||||
audioSection.classList.add('hidden');
|
||||
|
||||
try {
|
||||
const formData = new FormData();
|
||||
formData.append('text', text);
|
||||
|
||||
// Add voice URL if provided
|
||||
const voiceUrl = document.getElementById('voice-url-input').value.trim();
|
||||
if (voiceUrl) {
|
||||
formData.append('voice_url', voiceUrl);
|
||||
}
|
||||
|
||||
// Add voice WAV file if uploaded
|
||||
const voiceWavFile = document.getElementById('voice-wav-input').files[0];
|
||||
if (voiceWavFile) {
|
||||
formData.append('voice_wav', voiceWavFile);
|
||||
}
|
||||
|
||||
const response = await fetch('/tts', {
|
||||
method: 'POST',
|
||||
body: formData
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
throw new Error(`Server error: ${response.status}`);
|
||||
}
|
||||
|
||||
// Clone the response for both streaming and blob collection
|
||||
const responseForPlayback = response.clone();
|
||||
const responseForHistory = response.clone();
|
||||
|
||||
// Start streaming playback
|
||||
const reader = responseForPlayback.body.getReader();
|
||||
streamingPlayer = new StreamingWavPlayer();
|
||||
|
||||
const processStream = async () => {
|
||||
try {
|
||||
while (true) {
|
||||
const { done, value } = await reader.read();
|
||||
if (done) break;
|
||||
|
||||
if (value) {
|
||||
streamingPlayer.addChunk(value);
|
||||
}
|
||||
}
|
||||
} catch (e) {
|
||||
console.error('Error processing stream:', e);
|
||||
}
|
||||
};
|
||||
|
||||
// Start processing the stream
|
||||
processStream();
|
||||
|
||||
// Also collect the full blob for download/audio element
|
||||
const blob = await responseForHistory.blob();
|
||||
const totalTime = ((performance.now() - startTime) / 1000).toFixed(2);
|
||||
currentAudioBlob = blob;
|
||||
|
||||
// Update UI with audio
|
||||
const audioUrl = URL.createObjectURL(blob);
|
||||
audio.src = audioUrl;
|
||||
|
||||
// Wait for audio metadata to load to get duration
|
||||
audio.addEventListener('loadedmetadata', () => {
|
||||
const audioDuration = audio.duration;
|
||||
const speedRatio = (audioDuration / parseFloat(totalTime)).toFixed(1);
|
||||
|
||||
const timeToFirst = firstAudioTime ? ((firstAudioTime - startTime) / 1000).toFixed(2) : 'N/A';
|
||||
|
||||
showStatus(
|
||||
`✨ First audio: ${timeToFirst}s | Total: ${totalTime}s | ${speedRatio}x faster than real-time`,
|
||||
false
|
||||
);
|
||||
});
|
||||
downloadLink.href = audioUrl;
|
||||
downloadLink.download = `tts-audio.wav`;
|
||||
|
||||
audioSection.classList.remove('hidden');
|
||||
|
||||
} catch (error) {
|
||||
console.error('Error generating audio:', error);
|
||||
showStatus(`Error: ${error.message}`, false);
|
||||
setTimeout(hideStatus, 3000);
|
||||
} finally {
|
||||
// Re-enable button
|
||||
generateBtn.disabled = false;
|
||||
generateText.textContent = 'Generate Audio';
|
||||
}
|
||||
}
|
||||
|
||||
// Focus on text input when page loads
|
||||
window.addEventListener('load', () => {
|
||||
textInput.focus();
|
||||
});
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
@@ -0,0 +1 @@
|
||||
"""Utilities."""
|
||||
@@ -0,0 +1,122 @@
|
||||
"""Configuration models for loading YAML config files."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
|
||||
class StrictModel(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
|
||||
# Flow configuration
|
||||
class FlowConfig(StrictModel):
|
||||
dim: int
|
||||
depth: int
|
||||
|
||||
|
||||
# Transformer configuration for FlowLM
|
||||
class FlowLMTransformerConfig(StrictModel):
|
||||
hidden_scale: int
|
||||
max_period: int
|
||||
d_model: int
|
||||
num_heads: int
|
||||
num_layers: int
|
||||
|
||||
|
||||
class LookupTable(StrictModel):
|
||||
dim: int
|
||||
n_bins: int
|
||||
tokenizer: str
|
||||
tokenizer_path: str
|
||||
|
||||
|
||||
# Root configuration
|
||||
class FlowLMConfig(StrictModel):
|
||||
"""Root configuration model for YAML config files."""
|
||||
|
||||
dtype: str
|
||||
|
||||
# Nested configurations
|
||||
flow: FlowConfig
|
||||
transformer: FlowLMTransformerConfig
|
||||
|
||||
# conditioning
|
||||
lookup_table: LookupTable
|
||||
weights_path: str | None = None
|
||||
|
||||
|
||||
# SEANet configuration
|
||||
class SEANetConfig(StrictModel):
|
||||
dimension: int
|
||||
channels: int
|
||||
n_filters: int
|
||||
n_residual_layers: int
|
||||
ratios: list[int]
|
||||
kernel_size: int
|
||||
residual_kernel_size: int
|
||||
last_kernel_size: int
|
||||
dilation_base: int
|
||||
pad_mode: str
|
||||
compress: int
|
||||
|
||||
|
||||
# Transformer configuration for Mimi
|
||||
class MimiTransformerConfig(StrictModel):
|
||||
d_model: int
|
||||
input_dimension: int
|
||||
output_dimensions: tuple[int, ...]
|
||||
num_heads: int
|
||||
num_layers: int
|
||||
layer_scale: float
|
||||
context: int
|
||||
max_period: float = 10000.0
|
||||
dim_feedforward: int
|
||||
|
||||
|
||||
# Quantizer configuration
|
||||
class QuantizerConfig(StrictModel):
|
||||
dimension: int
|
||||
output_dimension: int
|
||||
|
||||
|
||||
# Root configuration
|
||||
class MimiConfig(StrictModel):
|
||||
"""Root configuration model for Mimi YAML config files."""
|
||||
|
||||
dtype: str
|
||||
|
||||
# Sample rate and channels
|
||||
sample_rate: int
|
||||
channels: int
|
||||
frame_rate: float
|
||||
|
||||
# SEANet configurations
|
||||
seanet: SEANetConfig
|
||||
|
||||
# Transformer
|
||||
transformer: MimiTransformerConfig
|
||||
|
||||
# Quantizer
|
||||
quantizer: QuantizerConfig
|
||||
weights_path: str | None = None
|
||||
|
||||
|
||||
class Config(StrictModel):
|
||||
flow_lm: FlowLMConfig
|
||||
mimi: MimiConfig
|
||||
weights_path: str | None = None
|
||||
weights_path_without_voice_cloning: str | None = None
|
||||
|
||||
|
||||
def load_config(yaml_path: str | Path) -> Config:
|
||||
yaml_path = Path(yaml_path)
|
||||
|
||||
if not yaml_path.exists():
|
||||
raise FileNotFoundError(f"Config file not found: {yaml_path}")
|
||||
|
||||
with open(yaml_path, "r") as f:
|
||||
config_dict = yaml.safe_load(f)
|
||||
|
||||
return Config(**config_dict)
|
||||
@@ -0,0 +1,26 @@
|
||||
import torch
|
||||
from torch.utils._python_dispatch import TorchDispatchMode
|
||||
|
||||
|
||||
def to_str(obj):
|
||||
if isinstance(obj, (torch.Tensor, torch.nn.Parameter)):
|
||||
return f"T(s={list(obj.shape)})"
|
||||
elif isinstance(obj, (list, tuple)):
|
||||
return "[" + ", ".join(to_str(o) for o in obj) + "]"
|
||||
elif isinstance(obj, dict):
|
||||
return "{" + ", ".join(f"{to_str(k)}: {to_str(v)}" for k, v in obj.items()) + "}"
|
||||
else:
|
||||
return str(obj)
|
||||
|
||||
|
||||
class LoggingMode(TorchDispatchMode):
|
||||
"""Useful to check implementation differences."""
|
||||
|
||||
def __torch_dispatch__(self, func, types, args=(), kwargs=None):
|
||||
output = func(*args, **kwargs or {})
|
||||
print(
|
||||
f"Aten function called: {func}, args: "
|
||||
f"{to_str(args)}, kwargs: {to_str(kwargs)} -> "
|
||||
f"output: {to_str(output)}"
|
||||
)
|
||||
return output
|
||||
@@ -0,0 +1,41 @@
|
||||
import logging
|
||||
from contextlib import contextmanager
|
||||
|
||||
|
||||
class PocketTTSFilter(logging.Filter):
|
||||
def filter(self, record):
|
||||
return record.name.startswith("pocket_tts")
|
||||
|
||||
|
||||
@contextmanager
|
||||
def enable_logging(library_name, level):
|
||||
# Get the specific logger and its parent
|
||||
logger = logging.getLogger(library_name)
|
||||
parent_logger = logging.getLogger("pocket_tts")
|
||||
|
||||
# Store original configuration
|
||||
old_level = logger.level
|
||||
old_parent_level = parent_logger.level
|
||||
old_handlers = parent_logger.handlers.copy()
|
||||
|
||||
# Configure logging format for pocket_tts logger
|
||||
parent_logger.setLevel(level)
|
||||
|
||||
# Clear existing handlers and add our custom formatter with filter
|
||||
parent_logger.handlers.clear()
|
||||
handler = logging.StreamHandler()
|
||||
formatter = logging.Formatter("%(levelname)s: %(message)s")
|
||||
handler.setFormatter(formatter)
|
||||
handler.addFilter(PocketTTSFilter())
|
||||
parent_logger.addHandler(handler)
|
||||
parent_logger.propagate = False
|
||||
|
||||
try:
|
||||
yield logger
|
||||
finally:
|
||||
# Restore original configuration
|
||||
logger.setLevel(old_level)
|
||||
parent_logger.setLevel(old_parent_level)
|
||||
parent_logger.handlers.clear()
|
||||
for h in old_handlers:
|
||||
parent_logger.addHandler(h)
|
||||
@@ -0,0 +1,102 @@
|
||||
import hashlib
|
||||
import logging
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import requests
|
||||
import safetensors.torch
|
||||
import torch
|
||||
from huggingface_hub import hf_hub_download
|
||||
from torch import nn
|
||||
|
||||
PROJECT_ROOT = Path(__file__).parent.parent.parent
|
||||
|
||||
PREDEFINED_VOICES = {
|
||||
"marius": "hf://kyutai/pocket-tts-without-voice-cloning/embeddings/marius.safetensors",
|
||||
"javert": "hf://kyutai/pocket-tts-without-voice-cloning/embeddings/javert.safetensors",
|
||||
"jean": "hf://kyutai/pocket-tts-without-voice-cloning/embeddings/jean.safetensors",
|
||||
"fantine": "hf://kyutai/pocket-tts-without-voice-cloning/embeddings/fantine.safetensors",
|
||||
"cosette": "hf://kyutai/pocket-tts-without-voice-cloning/embeddings/cosette.safetensors",
|
||||
"eponine": "hf://kyutai/pocket-tts-without-voice-cloning/embeddings/eponine.safetensors",
|
||||
"azelma": "hf://kyutai/pocket-tts-without-voice-cloning/embeddings/azelma.safetensors",
|
||||
}
|
||||
|
||||
|
||||
def make_cache_directory() -> Path:
|
||||
cache_dir = Path.home() / ".cache" / "pocket_tts"
|
||||
cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
return cache_dir
|
||||
|
||||
|
||||
def print_nb_parameters(model: nn.Module, model_name: str):
|
||||
logger = logging.getLogger(__name__)
|
||||
state_dict = model.state_dict()
|
||||
total = 0
|
||||
for key, value in state_dict.items():
|
||||
logger.info("%s: %,d", key, value.numel())
|
||||
total += value.numel()
|
||||
logger.info("Total number of parameters in %s: %,d", model_name, total)
|
||||
|
||||
|
||||
def size_of_dict(state_dict: dict) -> int:
|
||||
total_size = 0
|
||||
for value in state_dict.values():
|
||||
if isinstance(value, torch.Tensor):
|
||||
total_size += value.numel() * value.element_size()
|
||||
elif isinstance(value, dict):
|
||||
total_size += size_of_dict(value)
|
||||
return total_size
|
||||
|
||||
|
||||
class display_execution_time:
|
||||
def __init__(self, task_name: str, print_output: bool = True):
|
||||
self.task_name = task_name
|
||||
self.print_output = print_output
|
||||
self.start_time = None
|
||||
self.elapsed_time_ms = None
|
||||
self.logger = logging.getLogger(__name__)
|
||||
|
||||
def __enter__(self):
|
||||
self.start_time = time.monotonic()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
end_time = time.monotonic()
|
||||
self.elapsed_time_ms = int((end_time - self.start_time) * 1000)
|
||||
if self.print_output:
|
||||
self.logger.info("%s took %d ms", self.task_name, self.elapsed_time_ms)
|
||||
return False # Don't suppress exceptions
|
||||
|
||||
|
||||
def download_if_necessary(file_path: str) -> Path:
|
||||
if file_path.startswith("http://") or file_path.startswith("https://"):
|
||||
cache_dir = make_cache_directory()
|
||||
cached_file = cache_dir / (
|
||||
hashlib.sha256(file_path.encode()).hexdigest() + "." + file_path.split(".")[-1]
|
||||
)
|
||||
if not cached_file.exists():
|
||||
response = requests.get(file_path)
|
||||
response.raise_for_status()
|
||||
with open(cached_file, "wb") as f:
|
||||
f.write(response.content)
|
||||
return cached_file
|
||||
elif file_path.startswith("hf://"):
|
||||
file_path = file_path.removeprefix("hf://")
|
||||
splitted = file_path.split("/")
|
||||
repo_id = "/".join(splitted[:2])
|
||||
filename = "/".join(splitted[2:])
|
||||
cached_file = hf_hub_download(repo_id=repo_id, filename=filename)
|
||||
return Path(cached_file)
|
||||
else:
|
||||
return Path(file_path)
|
||||
|
||||
|
||||
def load_predefined_voice(voice_name: str) -> torch.Tensor:
|
||||
if voice_name not in PREDEFINED_VOICES:
|
||||
raise ValueError(
|
||||
f"Predefined voice '{voice_name}' not found"
|
||||
f", available voices are {list(PREDEFINED_VOICES)}."
|
||||
)
|
||||
voice_file = download_if_necessary(PREDEFINED_VOICES[voice_name])
|
||||
# There is only one tensor in the file.
|
||||
return safetensors.torch.load_file(voice_file)["audio_prompt"]
|
||||
@@ -0,0 +1,35 @@
|
||||
from pathlib import Path
|
||||
|
||||
import safetensors
|
||||
|
||||
|
||||
def get_flow_lm_state_dict(path: Path) -> dict:
|
||||
state_dict = {}
|
||||
with safetensors.safe_open(path, framework="pt", device="cpu") as f:
|
||||
for key in f.keys():
|
||||
if (
|
||||
key.startswith("flow.w_s_t.")
|
||||
or key == "condition_provider.conditioners.transcript_in_segment.learnt_padding"
|
||||
or key == "condition_provider.conditioners.speaker_wavs.learnt_padding"
|
||||
):
|
||||
# skip lookup table weights
|
||||
continue
|
||||
new_name = key
|
||||
if key == "condition_provider.conditioners.transcript_in_segment.embed.weight":
|
||||
new_name = "conditioner.embed.weight"
|
||||
if key == "condition_provider.conditioners.speaker_wavs.output_proj.weight":
|
||||
new_name = "speaker_proj_weight"
|
||||
state_dict[new_name] = f.get_tensor(key)
|
||||
return state_dict
|
||||
|
||||
|
||||
def get_mimi_state_dict(path: Path) -> dict:
|
||||
state_dict = {}
|
||||
with safetensors.safe_open(path, framework="pt", device="cpu") as f:
|
||||
for key in f.keys():
|
||||
if key.startswith("model.quantizer.vq.") or key == "model.quantizer.logvar_proj.weight":
|
||||
# skip vq weights
|
||||
continue
|
||||
|
||||
state_dict[key.removeprefix("model.")] = f.get_tensor(key)
|
||||
return state_dict
|
||||
@@ -0,0 +1,71 @@
|
||||
[project]
|
||||
name = "pocket-tts"
|
||||
requires-python = ">= 3.10,<3.15"
|
||||
description = "Kyutai's pocket-sized TTS!"
|
||||
version = "0.1.0"
|
||||
readme = "README.md"
|
||||
dependencies = [
|
||||
"numpy>=2",
|
||||
# I tried 2.4.0 but the generated audio is wrong, so 2.5+ is needed.
|
||||
"torch>=2.5.0",
|
||||
"pydantic>=2",
|
||||
"sentencepiece>=0.2.1",
|
||||
"beartype>=0.22.5",
|
||||
"safetensors>=0.4.0",
|
||||
"typer>=0.10.0",
|
||||
"typing_extensions>=4.0.0",
|
||||
"fastapi>=0.100",
|
||||
"uvicorn>=0.13.0",
|
||||
"python-multipart>=0.0.21",
|
||||
"scipy>=1.5.0",
|
||||
"einops>=0.4.0",
|
||||
"huggingface_hub>=0.10",
|
||||
"requests>=2.20.0",
|
||||
]
|
||||
|
||||
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
"coverage>=7.6.12",
|
||||
"line-profiler>=5.0.0",
|
||||
"pytest>=9.0.2",
|
||||
"pytest-xdist>=3.8.0",
|
||||
]
|
||||
|
||||
|
||||
[tool.uv.sources]
|
||||
torch = [
|
||||
{ index = "pytorch-cpu" },
|
||||
]
|
||||
|
||||
[[tool.uv.index]]
|
||||
name = "pytorch-cpu"
|
||||
url = "https://download.pytorch.org/whl/cpu"
|
||||
explicit = true
|
||||
|
||||
[tool.ruff]
|
||||
line-length = 100
|
||||
|
||||
[tool.ruff.format]
|
||||
line-ending = "lf"
|
||||
# avoid having a trailing comma changing
|
||||
# the behavior of the formatter
|
||||
skip-magic-trailing-comma = true
|
||||
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[tool.hatch.build.targets.wheel]
|
||||
packages = ["pocket_tts"]
|
||||
|
||||
[project.scripts]
|
||||
pocket-tts = "pocket_tts.main:cli_app"
|
||||
|
||||
[tool.uv]
|
||||
# System python distribution can have header files missing, so we use only the python
|
||||
# distributions provided by uv which have everything included.
|
||||
python-preference = "only-managed"
|
||||
|
||||
[tool.ruff.lint.flake8-tidy-imports]
|
||||
ban-relative-imports = "all"
|
||||
@@ -0,0 +1,3 @@
|
||||
import os
|
||||
|
||||
os.environ["POCKET_TTS_ERROR_WITHOUT_EOS"] = "1"
|
||||
@@ -0,0 +1,180 @@
|
||||
"""Integration tests for the CLI generate command using real implementation."""
|
||||
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from pocket_tts.data.audio import audio_read
|
||||
from pocket_tts.default_parameters import DEFAULT_VARIANT
|
||||
from pocket_tts.main import cli_app
|
||||
|
||||
other_voice = "https://huggingface.co/kyutai/tts-voices/resolve/main/expresso/ex01-ex02_default_001_channel1_168s.wav"
|
||||
|
||||
runner = CliRunner()
|
||||
|
||||
|
||||
def test_generate_basic_usage(tmp_path):
|
||||
"""Test basic generate command with default parameters."""
|
||||
output_file = tmp_path / "test_output.wav"
|
||||
|
||||
result = runner.invoke(
|
||||
cli_app,
|
||||
["generate", "--text", "Hello world, this is a test.", "--output-path", str(output_file)],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0
|
||||
# Verify output file was created and contains audio
|
||||
assert output_file.exists()
|
||||
assert output_file.stat().st_size > 0
|
||||
|
||||
# Verify it's a valid audio file
|
||||
audio, sample_rate = audio_read(str(output_file))
|
||||
assert audio.shape[0] == 1 # Mono channel
|
||||
assert audio.shape[1] > 0 # Has audio samples
|
||||
assert sample_rate == 24000 # Expected sample rate
|
||||
|
||||
|
||||
def test_generate_with_custom_voice(tmp_path):
|
||||
"""Test generate command with custom voice prompt."""
|
||||
output_file = tmp_path / "custom_voice_test.wav"
|
||||
|
||||
result = runner.invoke(
|
||||
cli_app,
|
||||
[
|
||||
"generate",
|
||||
"--text",
|
||||
"Testing custom voice.",
|
||||
"--voice",
|
||||
other_voice,
|
||||
"--output-path",
|
||||
str(output_file),
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert output_file.exists()
|
||||
|
||||
# Verify audio content
|
||||
audio, sample_rate = audio_read(str(output_file))
|
||||
assert audio.shape[0] == 1 # Mono channel
|
||||
assert audio.shape[1] > 0 # Has audio samples
|
||||
assert sample_rate == 24000
|
||||
|
||||
|
||||
def test_generate_with_custom_parameters(tmp_path):
|
||||
"""Test generate command with custom generation parameters."""
|
||||
output_file = tmp_path / "custom_params_test.wav"
|
||||
|
||||
result = runner.invoke(
|
||||
cli_app,
|
||||
[
|
||||
"generate",
|
||||
"--text",
|
||||
"Testing custom parameters.",
|
||||
"--variant",
|
||||
DEFAULT_VARIANT,
|
||||
"--temperature",
|
||||
"0.8",
|
||||
"--lsd-decode-steps",
|
||||
"2",
|
||||
"--eos-threshold",
|
||||
"-3.0",
|
||||
"--frames-after-eos",
|
||||
"7",
|
||||
"--output-path",
|
||||
str(output_file),
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert output_file.exists()
|
||||
|
||||
audio, sample_rate = audio_read(str(output_file))
|
||||
assert audio.shape[0] == 1 # Mono channel
|
||||
assert audio.shape[1] > 0 # Has audio samples
|
||||
assert sample_rate == 24000
|
||||
|
||||
|
||||
def test_generate_verbose_mode(tmp_path):
|
||||
"""Test generate command with verbose logging."""
|
||||
output_file = tmp_path / "verbose_test.wav"
|
||||
|
||||
result = runner.invoke(
|
||||
cli_app,
|
||||
["generate", "--text", "Testing verbose mode.", "-q", "--output-path", str(output_file)],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert output_file.exists()
|
||||
|
||||
|
||||
def test_generate_default_text(tmp_path):
|
||||
"""Test generate command with default text when no text provided."""
|
||||
output_file = tmp_path / "default_text_test.wav"
|
||||
|
||||
result = runner.invoke(cli_app, ["generate", "--output-path", str(output_file)])
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert output_file.exists()
|
||||
|
||||
audio, sample_rate = audio_read(str(output_file))
|
||||
assert audio.shape[0] == 1 # Mono channel
|
||||
assert audio.shape[1] > 0 # Has audio samples
|
||||
assert sample_rate == 24000
|
||||
|
||||
|
||||
def test_generate_long_text(tmp_path):
|
||||
"""Test generate command with longer text."""
|
||||
long_text = "This is a longer text to test the TTS system. " * 5
|
||||
output_file = tmp_path / "long_text_test.wav"
|
||||
|
||||
result = runner.invoke(
|
||||
cli_app, ["generate", "--text", long_text, "--output-path", str(output_file)]
|
||||
)
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert output_file.exists()
|
||||
|
||||
audio, sample_rate = audio_read(str(output_file))
|
||||
assert audio.shape[0] == 1 # Mono channel
|
||||
assert audio.shape[1] > 0 # Has audio samples
|
||||
assert sample_rate == 24000
|
||||
# Longer text should produce longer audio
|
||||
assert audio.shape[1] > 24000 * 10 # At least 10 second of audio
|
||||
|
||||
|
||||
def test_generate_command_help():
|
||||
"""Test that generate command help is displayed correctly."""
|
||||
result = runner.invoke(cli_app, ["generate", "--help"])
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert "Generate speech using Kyutai Pocket TTS" in result.output
|
||||
assert "--text" in result.output
|
||||
assert "--voice" in result.output
|
||||
assert "--output-path" in result.output
|
||||
assert "--variant" in result.output
|
||||
assert "--temperature" in result.output
|
||||
assert "--device" in result.output
|
||||
|
||||
|
||||
def test_generate_multiple_runs(tmp_path):
|
||||
"""Test multiple consecutive generate commands."""
|
||||
for i in range(3):
|
||||
output_file = tmp_path / f"test_run_{i + 1}.wav"
|
||||
|
||||
result = runner.invoke(
|
||||
cli_app,
|
||||
[
|
||||
"generate",
|
||||
"--text",
|
||||
f"This is test run number {i + 1}.",
|
||||
"--output-path",
|
||||
str(output_file),
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert output_file.exists()
|
||||
|
||||
audio, sample_rate = audio_read(str(output_file))
|
||||
assert audio.shape[0] == 1 # Mono channel
|
||||
assert audio.shape[1] > 0 # Has audio samples
|
||||
assert sample_rate == 24000
|
||||
@@ -0,0 +1,172 @@
|
||||
"""We also add the imports in the functions to make sure we didn't forget them."""
|
||||
# ruff: noqa: F841
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def test_readme_example():
|
||||
import scipy.io.wavfile
|
||||
|
||||
from pocket_tts import TTSModel
|
||||
|
||||
tts_model = TTSModel.load_model()
|
||||
voice_state = tts_model.get_state_for_audio_prompt(
|
||||
"hf://kyutai/tts-voices/alba-mackenna/casual.wav"
|
||||
)
|
||||
audio = tts_model.generate_audio(voice_state, "Hello world, this is a test.")
|
||||
# Audio is a torch tensor containing PCM data.
|
||||
scipy.io.wavfile.write("output.wav", tts_model.sample_rate, audio.numpy())
|
||||
|
||||
|
||||
def test_quick_start():
|
||||
import scipy.io.wavfile
|
||||
|
||||
from pocket_tts import TTSModel
|
||||
|
||||
# Load the model
|
||||
tts_model = TTSModel.load_model()
|
||||
|
||||
# Get voice state from an audio file
|
||||
voice_state = tts_model.get_state_for_audio_prompt(
|
||||
"hf://kyutai/tts-voices/alba-mackenna/casual.wav"
|
||||
)
|
||||
|
||||
# Generate audio
|
||||
audio = tts_model.generate_audio(voice_state, "Hello world, this is a test.")
|
||||
|
||||
# Save to file
|
||||
scipy.io.wavfile.write("output.wav", tts_model.sample_rate, audio.numpy())
|
||||
|
||||
|
||||
def test_load_model():
|
||||
from pocket_tts import TTSModel
|
||||
|
||||
# Load with default settings
|
||||
model = TTSModel.load_model()
|
||||
|
||||
# Load with custom parameters
|
||||
model = TTSModel.load_model(
|
||||
variant="b6369a24", temp=0.5, lsd_decode_steps=5, eos_threshold=-3.0
|
||||
)
|
||||
|
||||
|
||||
def test_device():
|
||||
from pocket_tts import TTSModel
|
||||
|
||||
model = TTSModel.load_model()
|
||||
print(f"Model running on: {model.device}")
|
||||
|
||||
|
||||
def test_sample_rate():
|
||||
from pocket_tts import TTSModel
|
||||
|
||||
model = TTSModel.load_model()
|
||||
print(f"Sample rate: {model.sample_rate} Hz")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def make_my_voice_file():
|
||||
import requests
|
||||
|
||||
url = "https://huggingface.co/kyutai/tts-voices/resolve/main/expresso/ex01-ex02_default_001_channel1_168s.wav"
|
||||
response = requests.get(url)
|
||||
with open("my_voice.wav", "wb") as f:
|
||||
f.write(response.content)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("make_my_voice_file")
|
||||
def test_get_state_for_audio_prompt():
|
||||
from pocket_tts import TTSModel
|
||||
|
||||
model = TTSModel.load_model()
|
||||
# From HuggingFace URL
|
||||
voice_state = model.get_state_for_audio_prompt(
|
||||
"hf://kyutai/tts-voices/alba-mackenna/casual.wav"
|
||||
)
|
||||
|
||||
# From local file
|
||||
voice_state = model.get_state_for_audio_prompt("./my_voice.wav")
|
||||
|
||||
# From HTTP URL
|
||||
voice_state = model.get_state_for_audio_prompt(
|
||||
"https://huggingface.co/kyutai/tts-voices/resolve"
|
||||
"/main/expresso/ex01-ex02_default_001_channel1_168s.wav"
|
||||
)
|
||||
|
||||
|
||||
def test_generate_audio():
|
||||
from pocket_tts import TTSModel
|
||||
|
||||
model = TTSModel.load_model()
|
||||
|
||||
voice_state = model.get_state_for_audio_prompt(
|
||||
"hf://kyutai/tts-voices/alba-mackenna/casual.wav"
|
||||
)
|
||||
|
||||
# Generate audio
|
||||
audio = model.generate_audio(voice_state, "Hello world!", frames_after_eos=2, copy_state=True)
|
||||
|
||||
print(f"Generated audio shape: {audio.shape}")
|
||||
print(f"Audio duration: {audio.shape[-1] / model.sample_rate:.2f} seconds")
|
||||
|
||||
|
||||
def test_generate_audio_stream():
|
||||
from pocket_tts import TTSModel
|
||||
|
||||
model = TTSModel.load_model()
|
||||
|
||||
voice_state = model.get_state_for_audio_prompt(
|
||||
"hf://kyutai/tts-voices/alba-mackenna/casual.wav"
|
||||
)
|
||||
# Stream generation
|
||||
for chunk in model.generate_audio_stream(voice_state, "Long text content..."):
|
||||
# Process each chunk as it's generated
|
||||
print(f"Generated chunk: {chunk.shape[0]} samples")
|
||||
# Could save chunks to file or play in real-time
|
||||
|
||||
|
||||
def test_voice_management():
|
||||
from pocket_tts import TTSModel
|
||||
|
||||
model = TTSModel.load_model()
|
||||
# Preload multiple voices
|
||||
voices = {
|
||||
"casual": model.get_state_for_audio_prompt(
|
||||
"hf://kyutai/tts-voices/alba-mackenna/casual.wav"
|
||||
),
|
||||
"funny": model.get_state_for_audio_prompt(
|
||||
"https://huggingface.co/kyutai/tts-voices/resolve/main/expresso/ex01-ex02_default_001_channel1_168s.wav"
|
||||
),
|
||||
}
|
||||
|
||||
# Generate with different voices
|
||||
casual_audio = model.generate_audio(voices["casual"], "Hey there!")
|
||||
funny_audio = model.generate_audio(voices["funny"], "Good morning.")
|
||||
|
||||
|
||||
def test_batch_processing():
|
||||
import scipy.io.wavfile
|
||||
import torch
|
||||
|
||||
from pocket_tts import TTSModel
|
||||
|
||||
model = TTSModel.load_model()
|
||||
|
||||
voice_state = model.get_state_for_audio_prompt(
|
||||
"hf://kyutai/tts-voices/alba-mackenna/casual.wav"
|
||||
)
|
||||
# Process multiple texts efficiently by re-using the same voice state
|
||||
texts = [
|
||||
"First sentence to generate.",
|
||||
"Second sentence to generate.",
|
||||
"Third sentence to generate.",
|
||||
]
|
||||
|
||||
audios = []
|
||||
for text in texts:
|
||||
audio = model.generate_audio(voice_state, text)
|
||||
audios.append(audio)
|
||||
|
||||
# Concatenate all audio
|
||||
full_audio = torch.cat(audios, dim=0)
|
||||
scipy.io.wavfile.write("batch_output.wav", model.sample_rate, full_audio.numpy())
|
||||
Reference in New Issue
Block a user