Koko-TTS: Ultra-Lightweight Flow-Matching Text-to-Speech

Koko-TTS is an ultra-lightweight, high-fidelity Text-to-Speech (TTS) model with ~24.6M parameters. Built upon a Flow-Matching framework, it combines a Matcha-style UNet architecture conditioned on RoPE-based text representations with the highly efficient Vocos 24kHz neural vocoder. The model is designed for ultra-fast, real-time speech synthesis without compromising on audio quality.

Model Details

Metric / Parameter	Value
Model Size	~24.6M parameters
Sample Rate	24,000 Hz
Vocoder	Vocos Mel 24kHz (`charactr/vocos-mel-24khz`)
Available Voices	128 Speaker IDs (0 to 127)
Training Dataset	jano3/libritts-r-128spk-vocos-mel
Language	English (Natively extensible to multi-script via tokenizer)

Intended Use & Limitations

Intended Use: Fast, real-time text-to-speech generation for applications requiring lightweight models (e.g., edge devices, games, voice assistants).
Extensibility: The built-in character tokenizer natively supports Latin, Cyrillic, Arabic/Urdu, and IPA, making it highly adaptable for fine-tuning on other languages.
Limitations: Due to its ultra-lightweight architecture (~24.6M parameters), while the audio quality is highly reasonable and generally outperforms older architectures like VITS, it may not match the absolute fidelity of massive, parameter-heavy TTS models. Additionally, the acoustic weights are currently optimized for English, and the quality heavily depends on the provided speaker ID and inference parameters. It is not designed for zero-shot voice cloning without fine-tuning.

Audio Samples

Speaker	Transcript	Audio Sample
Speaker 0	"The morning was quiet, and a gentle breeze moved through the trees. Somewhere in the distance, birds were singing, while the first light of day slowly filled the sky."
Speaker 46	"Well, here we are, take a breath, relax, and listen, sometimes, a quiet moment is all we need."
Speaker 123	"The quick brown fox jumps over the lazy dog. This is a demonstration of koko, an ultra-lightweight, high-quality text-to-speech model designed to combine fast inference with exceptional audio quality."

Quickstart

1. Installation

Ensure you have the required libraries installed:

pip install -q torch torchaudio transformers vocos tokenizers

2. Inference

Generating speech is straightforward using the transformers library.

import torch
import torchaudio
from transformers import AutoModel

device = "cuda" if torch.cuda.is_available() else "cpu"

# Load the model with custom code execution enabled
model = AutoModel.from_pretrained(
    "saki22/koko-tts", 
    trust_remote_code=True
).to(device)

# Generate speech waveform
audio = model.inference(
    text="Hello! This is Koko-TTS running fast and smooth.",
    spk_id=0,               # Choose a speaker between 0 and 127
    temperature=0.667,      # Controls variance/expressiveness
    cfg_strength=1.5,       # Classifier-Free Guidance strength
    n_steps=16,             # Number of ODE solver steps
    solver="euler"          # ODE solver type (euler or midpoint)
)

# Save the generated audio to a .wav file
torchaudio.save("output.wav", audio.unsqueeze(0), sample_rate=24000)

# Optional: Play directly if using Jupyter Notebook / Google Colab
# from IPython.display import Audio
# Audio(audio.numpy(), rate=24000, autoplay=True)

Fine-Tuning Guide

Koko-TTS is highly modular and designed to be easily fine-tuned on custom datasets.

1. Preprocessing (Mel Extraction & Alignment)

You will need to extract 100-dimensional Mel-spectrograms using Vocos to match the target representation Koko-TTS was trained on.

import torch
from vocos import Vocos

device = "cuda" if torch.cuda.is_available() else "cpu"
vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz").to(device)

def process_audio_sample(audio_tensor_24k, input_ids):
    with torch.no_grad():
        # Extract 100-dim Mel Spectrogram
        mel = vocos.feature_extractor(audio_tensor_24k.unsqueeze(0).to(device)).squeeze(0)
    
    # Note: You must compute token durations matching len(input_ids)
    # durations = get_alignment(audio_tensor_24k, input_ids) 
    
    return input_ids, durations, mel

2. Training Loop Example

import torch
from transformers import AutoModel

# Load model for training
model = AutoModel.from_pretrained("saki22/koko-tts", trust_remote_code=True).cuda()
model.train()

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5, weight_decay=1e-2)

for batch in dataloader:
    optimizer.zero_grad()

    # Forward pass calculates Flow Matching loss & Duration predictor loss
    outputs = model(
        input_ids=batch["input_ids"].cuda(),
        durations=batch["durations"].cuda(),
        mel_target=batch["mel_target"].cuda(),
        mel_lengths=batch["mel_lengths"].cuda(),
        spk_id=batch["speaker_ids"].cuda()
    )

    loss = outputs["loss"]
    loss.backward()

    # Gradient clipping is recommended for stable training
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    optimizer.step()

    print(f"Total Loss: {loss.item():.4f} | Mel Flow: {outputs['flow_loss'].item():.4f} | Duration: {outputs['duration_loss'].item():.4f}")

Architecture Highlights

Text Encoder: A RoPE-based Transformer encoder that processes characters natively.
Duration Predictor: A robust convolution-based module conditioned on speaker embeddings.
Decoder: A Matcha-style 1D UNet utilizing SnakeBeta activations, ResNet blocks, and Multi-head Self-Attention, trained via continuous normalizing flows (Flow-Matching).

License & Citation

This project is open-sourced under the Apache-2.0 License. If you use Koko-TTS in your research or project, please cite it as:

@misc{koko_tts_2026,
  author = {saki22},
  title = {Koko-TTS: Lightweight Flow-Matching Text-to-Speech},
  year = {2026},
  publisher = {Hugging Face},
  howpublished = {\url{https://huggingface.co/saki22/koko-tts}}
}

Downloads last month: 66

Safetensors

Model size

24.6M params

Tensor type

F32

saki22
/

koko-tts