Add int8 dynamic quantization support (#147)
# Add int8 dynamic quantization Closes #7 ## Summary This PR implements mixed-precision int8 dynamic quantization for pocket-tts, as requested in issue #7. It adds a `quantize=True` flag to `TTSModel.load_model()` that reduces runtime memory usage by ~48% and improves inference speed by ~27% on x86 and up to 23% on ARM, with no measurable impact on audio quality.
This commit is contained in:
nabil-tazi
@@ -0,0 +1,76 @@
|
||||
# Quantization
|
||||
|
||||
Pocket TTS supports dynamic int8 quantization to reduce runtime memory usage and improve inference speed on x86 CPUs.
|
||||
|
||||
## Quick Start
|
||||
|
||||
### CLI
|
||||
|
||||
```bash
|
||||
pocket-tts generate --quantize --text "Hello world"
|
||||
pocket-tts serve --quantize
|
||||
```
|
||||
|
||||
### Python API
|
||||
|
||||
```python
|
||||
from pocket_tts import TTSModel
|
||||
|
||||
model = TTSModel.load_model(quantize=True)
|
||||
voice_state = model.get_state_for_audio_prompt("alba")
|
||||
audio = model.generate_audio(voice_state, "Hello world!")
|
||||
```
|
||||
|
||||
## Installation
|
||||
|
||||
Quantization works out of the box on any supported PyTorch version (2.5+) using `torch.ao`.
|
||||
|
||||
For optimized performance, install `torchao` (requires torch 2.10+):
|
||||
|
||||
```bash
|
||||
pip install pocket-tts[quantize]
|
||||
```
|
||||
|
||||
The quantization module automatically selects the best available backend:
|
||||
|
||||
- **torchao** (torch 2.10+): optimized C++ kernels, faster on both ARM and x86
|
||||
- **torch.ao** (torch 2.5-2.9): deprecated but functional fallback
|
||||
|
||||
## Performance
|
||||
|
||||
Benchmarks on the full eval paragraph across 8 voices, 5 isolated runs per config.
|
||||
|
||||
### x86 (FBGEMM, ubuntu-latest GitHub Actions runner)
|
||||
|
||||
| Config | Runtime Memory | RTS | Speedup vs Baseline |
|
||||
| --------------------------- | -------------- | --------- | ------------------- |
|
||||
| baseline | 450 MB | 3.17x | -- |
|
||||
| **attention_ffn (default)** | **234 MB** | **4.04x** | **1.27x** |
|
||||
| all | 206 MB | 4.01x | 1.26x |
|
||||
|
||||
### ARM (QNNPACK, Apple M4 MacBook Air, torchao backend)
|
||||
|
||||
| Config | Runtime Memory | RTS | Speedup vs Baseline |
|
||||
| --------------------------- | -------------- | --------- | ------------------- |
|
||||
| baseline | 450 MB | 6.33x | -- |
|
||||
| **attention_ffn (default)** | **234 MB** | **7.76x** | **1.23x** |
|
||||
|
||||
With the `torch.ao` fallback (torch 2.5-2.9), ARM performance is ~16% slower than baseline rather than faster. Upgrading to torch 2.10+ with torchao is recommended for ARM users.
|
||||
|
||||
## Quality
|
||||
|
||||
Quantization has no measurable impact on speech quality:
|
||||
|
||||
- **WER (Word Error Rate)**: WER delta for the default attention_ffn config is −0.022 ±0.032 - the ± range crosses zero, meaning the delta is indistinguishable from measurement noise.
|
||||
- **Subjective listening**: no audible difference across all 8 voices
|
||||
|
||||
## What gets quantized
|
||||
|
||||
When `quantize=True`, the following layer groups in the FlowLM transformer are quantized to int8:
|
||||
|
||||
| Group | Params | Description |
|
||||
| ----------- | ------ | ------------------------------------------------ |
|
||||
| `attention` | ~25M | Q/K/V/output projections in transformer layers |
|
||||
| `ffn` | ~50M | Feed-forward linear layers in transformer layers |
|
||||
|
||||
The flow matching network (`flow_net`, ~7M params) and the Mimi VAE decoder (convolutional) remain in float32.
|
||||
+8
-2
@@ -188,11 +188,14 @@ def serve(
|
||||
help="Path to locally-saved model config .yaml file or model variant signature"
|
||||
),
|
||||
] = DEFAULT_VARIANT,
|
||||
quantize: Annotated[
|
||||
bool, typer.Option(help="Apply int8 quantization to reduce memory usage")
|
||||
] = False,
|
||||
):
|
||||
"""Start the FastAPI server."""
|
||||
|
||||
global tts_model, global_model_state
|
||||
tts_model = TTSModel.load_model(config)
|
||||
tts_model = TTSModel.load_model(config, quantize=quantize)
|
||||
|
||||
# Pre-load the voice prompt
|
||||
global_model_state = tts_model.get_state_for_audio_prompt(voice)
|
||||
@@ -236,6 +239,9 @@ def generate(
|
||||
max_tokens: Annotated[
|
||||
int, typer.Option(help="Maximum number of tokens per chunk.")
|
||||
] = MAX_TOKEN_PER_CHUNK,
|
||||
quantize: Annotated[
|
||||
bool, typer.Option(help="Apply int8 quantization to reduce memory usage")
|
||||
] = False,
|
||||
):
|
||||
"""Generate speech using Kyutai Pocket TTS."""
|
||||
log_level = logging.ERROR if quiet else logging.INFO
|
||||
@@ -248,7 +254,7 @@ def generate(
|
||||
logger.error("No input received from stdin.")
|
||||
raise typer.Exit(code=1)
|
||||
tts_model = TTSModel.load_model(
|
||||
config, temperature, lsd_decode_steps, noise_clamp, eos_threshold
|
||||
config, temperature, lsd_decode_steps, noise_clamp, eos_threshold, quantize=quantize
|
||||
)
|
||||
tts_model.to(device)
|
||||
|
||||
|
||||
@@ -33,6 +33,7 @@ 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 StatefulModule, increment_steps, init_states
|
||||
from pocket_tts.quantization import RECOMMENDED_CONFIG, apply_dynamic_int8
|
||||
from pocket_tts.utils.config import Config, load_config
|
||||
from pocket_tts.utils.utils import (
|
||||
PREDEFINED_VOICES,
|
||||
@@ -192,6 +193,7 @@ class TTSModel(nn.Module):
|
||||
lsd_decode_steps: int = DEFAULT_LSD_DECODE_STEPS,
|
||||
noise_clamp: float | int | None = DEFAULT_NOISE_CLAMP,
|
||||
eos_threshold: float = DEFAULT_EOS_THRESHOLD,
|
||||
quantize: bool = False,
|
||||
) -> Self:
|
||||
"""Load a pre-trained TTS model with specified configuration.
|
||||
|
||||
@@ -210,6 +212,11 @@ class TTSModel(nn.Module):
|
||||
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.
|
||||
quantize: If True, apply dynamic int8 quantization to the transformer's
|
||||
attention and FFN layers. Reduces runtime memory by ~48% and improves
|
||||
inference speed by ~27% on x86 (FBGEMM).
|
||||
No measurable impact on speech quality (WER unchanged).
|
||||
For optimized performance, install torchao: ``pip install pocket-tts[quantize]``
|
||||
|
||||
Returns:
|
||||
TTSModel: Fully initialized model with loaded weights on cpu, ready for
|
||||
@@ -227,8 +234,8 @@ class TTSModel(nn.Module):
|
||||
# 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)
|
||||
# Load with int8 quantization
|
||||
model = TTSModel.load_model(quantize=True)
|
||||
```
|
||||
"""
|
||||
if str(config).endswith(".yaml"):
|
||||
@@ -241,6 +248,10 @@ class TTSModel(nn.Module):
|
||||
tts_model = TTSModel._from_pydantic_config_with_weights(
|
||||
config, temp, lsd_decode_steps, noise_clamp, eos_threshold
|
||||
)
|
||||
|
||||
if quantize:
|
||||
apply_dynamic_int8(tts_model.flow_lm, RECOMMENDED_CONFIG)
|
||||
|
||||
return tts_model
|
||||
|
||||
def _run_flow_lm_and_increment_step(
|
||||
|
||||
@@ -0,0 +1,128 @@
|
||||
"""
|
||||
Dynamic int8 quantization for pocket-tts.
|
||||
|
||||
Uses torchao if available (torch 2.10+ with C++ extensions), otherwise
|
||||
falls back to torch.ao.quantization (deprecated but functional on torch 2.5-2.9).
|
||||
|
||||
Quantizes attention (Q/K/V/output projections) and FFN (linear1/linear2) layers
|
||||
in the FlowLM transformer. The flow matching network and Mimi VAE decoder
|
||||
remain in float32.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import platform
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Used by load_model(quantize=True) to specify which layer groups to quantize.
|
||||
RECOMMENDED_CONFIG = {"attention", "ffn"}
|
||||
|
||||
|
||||
def _get_backend():
|
||||
"""Detect the best available quantization backend.
|
||||
|
||||
Returns "torchao" if torchao is installed with working C++ extensions,
|
||||
otherwise returns "torch.ao".
|
||||
"""
|
||||
try:
|
||||
import importlib.util
|
||||
|
||||
if importlib.util.find_spec("torchao") is None:
|
||||
return "torch.ao"
|
||||
|
||||
import torchao
|
||||
|
||||
if hasattr(torchao, "_C") or not getattr(torchao, "_SKIPPED_CPP_EXTENSIONS", False):
|
||||
return "torchao"
|
||||
except ImportError:
|
||||
pass
|
||||
return "torch.ao"
|
||||
|
||||
|
||||
def _quantize_module_torchao(module: nn.Module):
|
||||
"""Apply int8 dynamic quantization using torchao."""
|
||||
from torchao.quantization import Int8DynamicActivationInt8WeightConfig, quantize_
|
||||
|
||||
quantize_(module, Int8DynamicActivationInt8WeightConfig())
|
||||
|
||||
|
||||
def _ensure_quantization_engine():
|
||||
"""Set the quantization engine for torch.ao (QNNPACK for ARM, FBGEMM for x86)."""
|
||||
if platform.machine() in ("arm64", "aarch64"):
|
||||
torch.backends.quantized.engine = "qnnpack"
|
||||
elif torch.backends.quantized.engine == "none":
|
||||
torch.backends.quantized.engine = "fbgemm"
|
||||
|
||||
|
||||
def apply_dynamic_int8(flow_lm: nn.Module, quantize_groups: set[str]) -> nn.Module:
|
||||
"""
|
||||
Apply dynamic int8 quantization to the specified layer groups of a FlowLM model.
|
||||
|
||||
Automatically selects the best available backend:
|
||||
- torchao (torch 2.10+): optimized C++ kernels, faster on both ARM and x86
|
||||
- torch.ao (torch 2.5-2.9): deprecated but functional fallback
|
||||
|
||||
Args:
|
||||
flow_lm: The FlowLM model (model.flow_lm)
|
||||
quantize_groups: Set of group keys to quantize.
|
||||
Valid keys: "attention", "ffn", "flow_net"
|
||||
|
||||
Returns:
|
||||
The quantized model (modified in-place).
|
||||
"""
|
||||
if not quantize_groups:
|
||||
logger.info("No quantization groups specified, returning model unchanged.")
|
||||
return flow_lm
|
||||
|
||||
backend = _get_backend()
|
||||
logger.info("Using quantization backend: %s", backend)
|
||||
|
||||
if backend == "torchao":
|
||||
_apply_torchao(flow_lm, quantize_groups)
|
||||
else:
|
||||
_apply_torch_ao(flow_lm, quantize_groups)
|
||||
|
||||
return flow_lm
|
||||
|
||||
|
||||
def _apply_torchao(flow_lm: nn.Module, quantize_groups: set[str]) -> None:
|
||||
"""Apply quantization using torchao backend."""
|
||||
if "flow_net" in quantize_groups:
|
||||
_quantize_module_torchao(flow_lm.flow_net)
|
||||
|
||||
for layer in flow_lm.transformer.layers:
|
||||
if "attention" in quantize_groups:
|
||||
_quantize_module_torchao(layer.self_attn)
|
||||
|
||||
if "ffn" in quantize_groups:
|
||||
wrapper1 = nn.Sequential(layer.linear1)
|
||||
wrapper2 = nn.Sequential(layer.linear2)
|
||||
_quantize_module_torchao(wrapper1)
|
||||
_quantize_module_torchao(wrapper2)
|
||||
layer.linear1 = wrapper1[0]
|
||||
layer.linear2 = wrapper2[0]
|
||||
|
||||
|
||||
def _apply_torch_ao(flow_lm: nn.Module, quantize_groups: set[str]) -> None:
|
||||
"""Apply quantization using deprecated torch.ao backend (fallback)."""
|
||||
from torch.ao.quantization import quantize_dynamic
|
||||
|
||||
_ensure_quantization_engine()
|
||||
|
||||
if "flow_net" in quantize_groups:
|
||||
quantize_dynamic(flow_lm.flow_net, {nn.Linear}, dtype=torch.qint8, inplace=True)
|
||||
|
||||
for layer in flow_lm.transformer.layers:
|
||||
if "attention" in quantize_groups:
|
||||
quantize_dynamic(layer.self_attn, {nn.Linear}, dtype=torch.qint8, inplace=True)
|
||||
|
||||
if "ffn" in quantize_groups:
|
||||
layer.linear1 = quantize_dynamic(
|
||||
nn.Sequential(layer.linear1), {nn.Linear}, dtype=torch.qint8
|
||||
)[0]
|
||||
layer.linear2 = quantize_dynamic(
|
||||
nn.Sequential(layer.linear2), {nn.Linear}, dtype=torch.qint8
|
||||
)[0]
|
||||
@@ -25,6 +25,7 @@ dependencies = [
|
||||
|
||||
[project.optional-dependencies]
|
||||
audio = ["soundfile>=0.12.0"]
|
||||
quantize = ["torchao>=0.16.0"]
|
||||
|
||||
|
||||
[dependency-groups]
|
||||
|
||||
@@ -0,0 +1,648 @@
|
||||
"""
|
||||
Evaluation harness for pocket-tts int8 quantization strategy.
|
||||
|
||||
For each quantization config:
|
||||
- Generates audio for eval sentences per voice
|
||||
- Measures Real-Time Speed (RTS): audio_duration / wall_clock_time
|
||||
- Measures quality: SNR vs baseline, PESQ, Whisper WER
|
||||
- Saves audio files for manual listening
|
||||
- Outputs a summary CSV and markdown report
|
||||
|
||||
Usage:
|
||||
uv run python scripts/evaluate_quantization.py --config all
|
||||
uv run python scripts/evaluate_quantization.py --all-configs
|
||||
uv run python scripts/evaluate_quantization.py --all-configs --skip-quality
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import json
|
||||
import logging
|
||||
import statistics
|
||||
import time
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
import scipy.io.wavfile
|
||||
import torch
|
||||
|
||||
from pocket_tts import TTSModel
|
||||
from pocket_tts.quantization import apply_dynamic_int8
|
||||
|
||||
# Quantization configs for benchmarking. Each maps to a set of layer group keys.
|
||||
CONFIGS = {
|
||||
"baseline": set(),
|
||||
"all": {"attention", "ffn", "flow_net"},
|
||||
"attention_ffn": {"attention", "ffn"},
|
||||
"attention": {"attention"},
|
||||
"ffn": {"ffn"},
|
||||
"flow_net": {"flow_net"},
|
||||
"flow_net_attention": {"flow_net", "attention"},
|
||||
"ffn_flow_net": {"ffn", "flow_net"},
|
||||
}
|
||||
|
||||
|
||||
def get_model_size_mb(model: torch.nn.Module) -> float:
|
||||
"""Returns model size in MB by serializing to a byte buffer."""
|
||||
import io
|
||||
|
||||
buf = io.BytesIO()
|
||||
torch.save(model.state_dict(), buf)
|
||||
return buf.tell() / (1024**2)
|
||||
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
VOICES = ["alba", "marius", "javert", "jean", "fantine", "cosette", "eponine", "azelma"]
|
||||
|
||||
# Main eval paragraph (prosody / naturalness)
|
||||
EVAL_TEXT = (
|
||||
"Sarah had always wondered whether the old lighthouse still worked. "
|
||||
"On clear mornings, she'd walk the narrow coastal path — past the fishing boats, "
|
||||
"the salt-weathered fences, the occasional curious gull — and stop at the cliff's "
|
||||
"edge to look. Three years ago, she'd climbed the iron staircase herself: one hundred "
|
||||
"and twelve steps, each one groaning under her boots. At the top, the lens was still "
|
||||
"there, enormous and dusty, its glass facets catching the pale winter light. Nobody "
|
||||
"had switched it on in decades. Still, on stormy nights, she sometimes thought she "
|
||||
"could see a faint glow from her bedroom window. She never mentioned this to anyone."
|
||||
)
|
||||
|
||||
# Diverse sentences for WER stress testing
|
||||
QUALITY_SENTENCES = [
|
||||
# Normal prose
|
||||
"Sarah had always wondered whether the old lighthouse still worked.",
|
||||
# Numbers
|
||||
"The package weighs one hundred and twelve grams and costs forty-three dollars.",
|
||||
# Rare words
|
||||
"The archaeologist meticulously catalogued the palaeolithic artefacts.",
|
||||
# Punctuation-heavy (prosody stress test)
|
||||
"Wait — are you certain? Because if not, we should stop, reconsider, and try again.",
|
||||
# Short sentence (trailing artifact detector)
|
||||
"Yes.",
|
||||
]
|
||||
|
||||
|
||||
@dataclass
|
||||
class GenerationResult:
|
||||
voice: str
|
||||
config_id: str
|
||||
audio_duration_sec: float
|
||||
wall_clock_sec: float
|
||||
rts: float
|
||||
load_time_sec: float
|
||||
model_size_mb: float
|
||||
output_path: str
|
||||
error: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class QualityResult:
|
||||
config_id: str
|
||||
voice: str
|
||||
sentence_idx: int
|
||||
sentence: str
|
||||
snr_db: Optional[float] = None
|
||||
pesq_score: Optional[float] = None
|
||||
wer: Optional[float] = None
|
||||
baseline_transcript: Optional[str] = None
|
||||
quantized_transcript: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ConfigSummary:
|
||||
config_id: str
|
||||
quantize_groups: list
|
||||
model_size_mb: float
|
||||
load_time_sec: float
|
||||
mean_rts: float
|
||||
min_rts: float
|
||||
max_rts: float
|
||||
rts_vs_baseline: Optional[float] = None
|
||||
mean_snr_db: Optional[float] = None
|
||||
mean_pesq: Optional[float] = None
|
||||
mean_wer_baseline: Optional[float] = None
|
||||
mean_wer_quantized: Optional[float] = None
|
||||
results: list = field(default_factory=list)
|
||||
quality_results: list = field(default_factory=list)
|
||||
|
||||
|
||||
def load_and_quantize_model(config_id: str):
|
||||
"""Load and quantize model. Returns (model, load_time, model_size_mb)."""
|
||||
quantize_groups = CONFIGS[config_id]
|
||||
|
||||
t0 = time.perf_counter()
|
||||
tts_model = TTSModel.load_model(temp=0.0)
|
||||
|
||||
if quantize_groups:
|
||||
logger.info("[%s] Applying int8 quantization: %s", config_id, quantize_groups)
|
||||
apply_dynamic_int8(tts_model.flow_lm, quantize_groups)
|
||||
|
||||
load_time = time.perf_counter() - t0
|
||||
model_size = get_model_size_mb(tts_model.flow_lm)
|
||||
|
||||
logger.info("[%s] Loaded in %.2fs, FlowLM size ~%.1fMB", config_id, load_time, model_size)
|
||||
return tts_model, load_time, model_size
|
||||
|
||||
|
||||
def save_audio(audio_tensor, sample_rate, output_path):
|
||||
"""Save audio tensor to WAV file."""
|
||||
audio_np = audio_tensor.numpy() if hasattr(audio_tensor, "numpy") else np.array(audio_tensor)
|
||||
audio_int16 = (audio_np * 32767).clip(-32768, 32767).astype(np.int16)
|
||||
scipy.io.wavfile.write(str(output_path), sample_rate, audio_int16)
|
||||
|
||||
|
||||
def generate_for_voice(tts_model, voice, config_id, output_dir):
|
||||
"""Generate audio for a single voice and return metrics."""
|
||||
logger.info("[%s] Generating for voice: %s", config_id, voice)
|
||||
|
||||
try:
|
||||
voice_state = tts_model.get_state_for_audio_prompt(voice)
|
||||
|
||||
t0 = time.perf_counter()
|
||||
audio = tts_model.generate_audio(model_state=voice_state, text_to_generate=EVAL_TEXT)
|
||||
wall_clock = time.perf_counter() - t0
|
||||
|
||||
audio_duration = len(audio) / tts_model.sample_rate
|
||||
rts = audio_duration / wall_clock
|
||||
|
||||
output_path = output_dir / f"{config_id}__{voice}.wav"
|
||||
save_audio(audio, tts_model.sample_rate, output_path)
|
||||
|
||||
logger.info(
|
||||
"[%s][%s] audio=%.1fs wall=%.1fs RTS=%.2fx",
|
||||
config_id,
|
||||
voice,
|
||||
audio_duration,
|
||||
wall_clock,
|
||||
rts,
|
||||
)
|
||||
|
||||
return GenerationResult(
|
||||
voice=voice,
|
||||
config_id=config_id,
|
||||
audio_duration_sec=round(audio_duration, 2),
|
||||
wall_clock_sec=round(wall_clock, 2),
|
||||
rts=round(rts, 3),
|
||||
load_time_sec=0.0,
|
||||
model_size_mb=0.0,
|
||||
output_path=str(output_path),
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("[%s][%s] Failed: %s", config_id, voice, e)
|
||||
return GenerationResult(
|
||||
voice=voice,
|
||||
config_id=config_id,
|
||||
audio_duration_sec=0,
|
||||
wall_clock_sec=0,
|
||||
rts=0,
|
||||
load_time_sec=0,
|
||||
model_size_mb=0,
|
||||
output_path="",
|
||||
error=str(e),
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Quality metrics
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def compute_snr(baseline_audio: torch.Tensor, quantized_audio: torch.Tensor) -> float:
|
||||
"""Compute Signal-to-Noise Ratio in dB between baseline and quantized audio."""
|
||||
# Align lengths (generation is non-deterministic, lengths may differ slightly)
|
||||
min_len = min(len(baseline_audio), len(quantized_audio))
|
||||
baseline = baseline_audio[:min_len].float()
|
||||
quantized = quantized_audio[:min_len].float()
|
||||
|
||||
noise = baseline - quantized
|
||||
signal_power = baseline.pow(2).mean()
|
||||
noise_power = noise.pow(2).mean()
|
||||
|
||||
if noise_power == 0:
|
||||
return float("inf")
|
||||
return (10 * torch.log10(signal_power / noise_power)).item()
|
||||
|
||||
|
||||
def compute_pesq(
|
||||
baseline_audio: np.ndarray, quantized_audio: np.ndarray, sample_rate: int
|
||||
) -> Optional[float]:
|
||||
"""Compute PESQ score. Returns None if pesq is not installed."""
|
||||
try:
|
||||
from pesq import pesq
|
||||
except ImportError:
|
||||
return None
|
||||
|
||||
# PESQ requires 16kHz or 8kHz
|
||||
target_sr = 16000
|
||||
if sample_rate != target_sr:
|
||||
import scipy.signal
|
||||
|
||||
baseline_audio = scipy.signal.resample(
|
||||
baseline_audio, int(len(baseline_audio) * target_sr / sample_rate)
|
||||
)
|
||||
quantized_audio = scipy.signal.resample(
|
||||
quantized_audio, int(len(quantized_audio) * target_sr / sample_rate)
|
||||
)
|
||||
|
||||
# Align lengths
|
||||
min_len = min(len(baseline_audio), len(quantized_audio))
|
||||
baseline_audio = baseline_audio[:min_len].astype(np.float32)
|
||||
quantized_audio = quantized_audio[:min_len].astype(np.float32)
|
||||
|
||||
# Normalize to [-1, 1]
|
||||
max_val = max(np.abs(baseline_audio).max(), np.abs(quantized_audio).max(), 1e-8)
|
||||
baseline_audio = baseline_audio / max_val
|
||||
quantized_audio = quantized_audio / max_val
|
||||
|
||||
try:
|
||||
return pesq(target_sr, baseline_audio, quantized_audio, "wb")
|
||||
except Exception as e:
|
||||
logger.warning("PESQ failed: %s", e)
|
||||
return None
|
||||
|
||||
|
||||
def compute_wer(audio_path: str, reference_text: str, whisper_model) -> tuple[float, str]:
|
||||
"""Compute Word Error Rate using Whisper. Returns (wer, transcript)."""
|
||||
try:
|
||||
from jiwer import wer
|
||||
except ImportError:
|
||||
return None, ""
|
||||
|
||||
result = whisper_model.transcribe(str(audio_path), language="en")
|
||||
transcript = result["text"].strip()
|
||||
error_rate = wer(reference_text, transcript)
|
||||
return error_rate, transcript
|
||||
|
||||
|
||||
def run_quality_eval(
|
||||
baseline_model, quantized_model, config_id, voice, output_dir, whisper_model=None
|
||||
):
|
||||
"""Run quality evaluation for a single voice across all quality sentences."""
|
||||
results = []
|
||||
voice_state_b = baseline_model.get_state_for_audio_prompt(voice)
|
||||
voice_state_q = quantized_model.get_state_for_audio_prompt(voice)
|
||||
|
||||
for i, sentence in enumerate(QUALITY_SENTENCES):
|
||||
logger.info(
|
||||
"[%s][%s] Quality sentence %d/%d", config_id, voice, i + 1, len(QUALITY_SENTENCES)
|
||||
)
|
||||
|
||||
baseline_audio = baseline_model.generate_audio(
|
||||
model_state=voice_state_b, text_to_generate=sentence
|
||||
)
|
||||
quantized_audio = quantized_model.generate_audio(
|
||||
model_state=voice_state_q, text_to_generate=sentence
|
||||
)
|
||||
|
||||
sr = baseline_model.sample_rate
|
||||
qr = QualityResult(config_id=config_id, voice=voice, sentence_idx=i, sentence=sentence)
|
||||
|
||||
# SNR
|
||||
qr.snr_db = round(compute_snr(baseline_audio, quantized_audio), 2)
|
||||
|
||||
# PESQ
|
||||
qr.pesq_score = compute_pesq(baseline_audio.numpy(), quantized_audio.numpy(), sr)
|
||||
if qr.pesq_score is not None:
|
||||
qr.pesq_score = round(qr.pesq_score, 3)
|
||||
|
||||
# WER (save temp files for Whisper)
|
||||
if whisper_model is not None:
|
||||
baseline_path = output_dir / f"quality_{config_id}__{voice}__s{i}_baseline.wav"
|
||||
quantized_path = output_dir / f"quality_{config_id}__{voice}__s{i}_quantized.wav"
|
||||
save_audio(baseline_audio, sr, baseline_path)
|
||||
save_audio(quantized_audio, sr, quantized_path)
|
||||
|
||||
wer_b, transcript_b = compute_wer(baseline_path, sentence, whisper_model)
|
||||
wer_q, transcript_q = compute_wer(quantized_path, sentence, whisper_model)
|
||||
|
||||
qr.wer = wer_q
|
||||
qr.baseline_transcript = transcript_b
|
||||
qr.quantized_transcript = transcript_q
|
||||
|
||||
logger.info(
|
||||
"[%s][%s][s%d] SNR=%.1fdB PESQ=%s WER_b=%.2f WER_q=%.2f",
|
||||
config_id,
|
||||
voice,
|
||||
i,
|
||||
qr.snr_db,
|
||||
qr.pesq_score or "N/A",
|
||||
wer_b or 0,
|
||||
wer_q or 0,
|
||||
)
|
||||
else:
|
||||
logger.info(
|
||||
"[%s][%s][s%d] SNR=%.1fdB PESQ=%s",
|
||||
config_id,
|
||||
voice,
|
||||
i,
|
||||
qr.snr_db,
|
||||
qr.pesq_score or "N/A",
|
||||
)
|
||||
|
||||
results.append(qr)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Run configs
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def run_config(config_id, output_dir, voices):
|
||||
"""Run full evaluation for a single quantization config."""
|
||||
logger.info("=" * 60)
|
||||
logger.info("Evaluating config: %s (groups: %s)", config_id, CONFIGS[config_id] or "none")
|
||||
logger.info("=" * 60)
|
||||
|
||||
tts_model, load_time, model_size = load_and_quantize_model(config_id)
|
||||
|
||||
results = []
|
||||
for voice in voices:
|
||||
result = generate_for_voice(tts_model, voice, config_id, output_dir)
|
||||
result.load_time_sec = round(load_time, 2)
|
||||
result.model_size_mb = round(model_size, 1)
|
||||
results.append(result)
|
||||
|
||||
valid_rts = [r.rts for r in results if not r.error and r.rts > 0]
|
||||
mean_rts = statistics.mean(valid_rts) if valid_rts else 0
|
||||
min_rts = min(valid_rts) if valid_rts else 0
|
||||
max_rts = max(valid_rts) if valid_rts else 0
|
||||
|
||||
return ConfigSummary(
|
||||
config_id=config_id,
|
||||
quantize_groups=sorted(CONFIGS[config_id]),
|
||||
model_size_mb=round(model_size, 1),
|
||||
load_time_sec=round(load_time, 2),
|
||||
mean_rts=round(mean_rts, 3),
|
||||
min_rts=round(min_rts, 3),
|
||||
max_rts=round(max_rts, 3),
|
||||
results=results,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Reporting
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def write_csv(summaries, output_dir):
|
||||
csv_path = output_dir / "results.csv"
|
||||
with open(csv_path, "w", newline="") as f:
|
||||
writer = csv.DictWriter(
|
||||
f,
|
||||
fieldnames=[
|
||||
"config_id",
|
||||
"quantize_groups",
|
||||
"voice",
|
||||
"audio_duration_sec",
|
||||
"wall_clock_sec",
|
||||
"rts",
|
||||
"load_time_sec",
|
||||
"model_size_mb",
|
||||
"error",
|
||||
"output_path",
|
||||
],
|
||||
)
|
||||
writer.writeheader()
|
||||
for summary in summaries:
|
||||
for r in summary.results:
|
||||
writer.writerow(
|
||||
{
|
||||
"config_id": r.config_id,
|
||||
"quantize_groups": "|".join(summary.quantize_groups) or "none",
|
||||
"voice": r.voice,
|
||||
"audio_duration_sec": r.audio_duration_sec,
|
||||
"wall_clock_sec": r.wall_clock_sec,
|
||||
"rts": r.rts,
|
||||
"load_time_sec": r.load_time_sec,
|
||||
"model_size_mb": r.model_size_mb,
|
||||
"error": r.error or "",
|
||||
"output_path": r.output_path,
|
||||
}
|
||||
)
|
||||
logger.info("CSV written: %s", csv_path)
|
||||
|
||||
|
||||
def write_quality_csv(summaries, output_dir):
|
||||
csv_path = output_dir / "quality_results.csv"
|
||||
rows = []
|
||||
for s in summaries:
|
||||
for qr in s.quality_results:
|
||||
rows.append(asdict(qr))
|
||||
if not rows:
|
||||
return
|
||||
with open(csv_path, "w", newline="") as f:
|
||||
writer = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
|
||||
writer.writeheader()
|
||||
writer.writerows(rows)
|
||||
logger.info("Quality CSV written: %s", csv_path)
|
||||
|
||||
|
||||
def write_markdown_report(summaries, output_dir):
|
||||
report_path = output_dir / "report.md"
|
||||
lines = [
|
||||
"# pocket-tts int8 Quantization Evaluation Report",
|
||||
f"\nGenerated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
|
||||
f"\nEval text length: {len(EVAL_TEXT.split())} words",
|
||||
f"\nVoices tested: {', '.join(VOICES)}",
|
||||
"\n---\n",
|
||||
"## Speed Summary\n",
|
||||
"| Config | Quantized Groups | Model Size | Load Time | Mean RTS | Speedup vs Baseline |",
|
||||
"|--------|-----------------|------------|-----------|----------|---------------------|",
|
||||
]
|
||||
|
||||
for s in summaries:
|
||||
groups_str = ", ".join(s.quantize_groups) if s.quantize_groups else "none"
|
||||
speedup = f"{s.rts_vs_baseline:.2f}x" if s.rts_vs_baseline is not None else "---"
|
||||
lines.append(
|
||||
f"| `{s.config_id}` | {groups_str} | {s.model_size_mb:.1f} MB "
|
||||
f"| {s.load_time_sec:.1f}s | {s.mean_rts:.2f}x | {speedup} |"
|
||||
)
|
||||
|
||||
# Quality summary
|
||||
has_quality = any(s.quality_results for s in summaries)
|
||||
if has_quality:
|
||||
lines += ["\n---\n", "## Quality Summary\n"]
|
||||
lines.append(
|
||||
"| Config | Mean SNR (dB) | Mean PESQ | Mean WER (baseline) | Mean WER (quantized) |"
|
||||
)
|
||||
lines.append(
|
||||
"|--------|--------------|-----------|--------------------|--------------------|"
|
||||
)
|
||||
for s in summaries:
|
||||
if not s.quality_results:
|
||||
continue
|
||||
snr = f"{s.mean_snr_db:.1f}" if s.mean_snr_db is not None else "N/A"
|
||||
pesq_s = f"{s.mean_pesq:.3f}" if s.mean_pesq is not None else "N/A"
|
||||
wer_b = f"{s.mean_wer_baseline:.3f}" if s.mean_wer_baseline is not None else "N/A"
|
||||
wer_q = f"{s.mean_wer_quantized:.3f}" if s.mean_wer_quantized is not None else "N/A"
|
||||
lines.append(f"| `{s.config_id}` | {snr} | {pesq_s} | {wer_b} | {wer_q} |")
|
||||
|
||||
lines += ["\n---\n", "## Per-Voice Speed Results\n"]
|
||||
|
||||
for s in summaries:
|
||||
lines.append(f"### `{s.config_id}`\n")
|
||||
lines.append("| Voice | Audio Duration | Wall Clock | RTS | Notes |")
|
||||
lines.append("|-------|---------------|------------|-----|-------|")
|
||||
for r in s.results:
|
||||
note = f"ERROR: {r.error}" if r.error else "(listen)"
|
||||
lines.append(
|
||||
f"| {r.voice} | {r.audio_duration_sec:.1f}s "
|
||||
f"| {r.wall_clock_sec:.1f}s | {r.rts:.2f}x | {note} |"
|
||||
)
|
||||
lines.append("")
|
||||
|
||||
report_path.write_text("\n".join(lines))
|
||||
logger.info("Report written: %s", report_path)
|
||||
|
||||
|
||||
def write_json_summary(summaries, output_dir):
|
||||
json_path = output_dir / "summary.json"
|
||||
json_path.write_text(json.dumps([asdict(s) for s in summaries], indent=2))
|
||||
logger.info("JSON summary written: %s", json_path)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Entry point
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Evaluate pocket-tts quantization configs")
|
||||
parser.add_argument("--config", nargs="+", choices=list(CONFIGS.keys()))
|
||||
parser.add_argument("--all-configs", action="store_true")
|
||||
parser.add_argument("--voices", nargs="+", default=VOICES, choices=VOICES)
|
||||
parser.add_argument("--output-dir", default="./eval_output")
|
||||
parser.add_argument(
|
||||
"--skip-quality", action="store_true", help="Skip quality metrics (SNR/PESQ/WER)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--quality-voices",
|
||||
nargs="+",
|
||||
default=["alba", "marius"],
|
||||
help="Voices to run quality eval on (default: alba, marius)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
configs_to_run = (
|
||||
list(CONFIGS.keys()) if args.all_configs else (args.config or ["baseline", "all"])
|
||||
)
|
||||
|
||||
# Ensure baseline is first
|
||||
configs_to_run = ["baseline"] + [c for c in configs_to_run if c != "baseline"]
|
||||
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
output_dir = Path(args.output_dir) / timestamp
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
logger.info("Output directory: %s", output_dir)
|
||||
|
||||
# --- Speed evaluation ---
|
||||
summaries = []
|
||||
baseline_rts = None
|
||||
|
||||
for config_id in configs_to_run:
|
||||
summary = run_config(config_id, output_dir, args.voices)
|
||||
|
||||
if config_id == "baseline":
|
||||
baseline_rts = summary.mean_rts
|
||||
elif baseline_rts and baseline_rts > 0:
|
||||
summary.rts_vs_baseline = round(summary.mean_rts / baseline_rts, 3)
|
||||
|
||||
summaries.append(summary)
|
||||
|
||||
# --- Quality evaluation ---
|
||||
if not args.skip_quality:
|
||||
# Load Whisper model once
|
||||
whisper_model = None
|
||||
try:
|
||||
import whisper
|
||||
|
||||
logger.info("Loading Whisper model for WER evaluation...")
|
||||
whisper_model = whisper.load_model("base")
|
||||
logger.info("Whisper model loaded.")
|
||||
except ImportError:
|
||||
logger.warning(
|
||||
"openai-whisper not installed, skipping WER. Install with: pip install openai-whisper jiwer"
|
||||
)
|
||||
|
||||
# Load baseline model once for quality comparison
|
||||
baseline_model = TTSModel.load_model()
|
||||
|
||||
for summary in summaries:
|
||||
if summary.config_id == "baseline":
|
||||
continue # no need to compare baseline against itself
|
||||
|
||||
quantized_model, _, _ = load_and_quantize_model(summary.config_id)
|
||||
|
||||
all_quality = []
|
||||
for voice in args.quality_voices:
|
||||
if voice not in args.voices:
|
||||
continue
|
||||
quality_results = run_quality_eval(
|
||||
baseline_model,
|
||||
quantized_model,
|
||||
summary.config_id,
|
||||
voice,
|
||||
output_dir,
|
||||
whisper_model=whisper_model,
|
||||
)
|
||||
all_quality.extend(quality_results)
|
||||
|
||||
summary.quality_results = all_quality
|
||||
|
||||
# Aggregate quality metrics
|
||||
snrs = [q.snr_db for q in all_quality if q.snr_db is not None]
|
||||
pesqs = [q.pesq_score for q in all_quality if q.pesq_score is not None]
|
||||
wers_q = [q.wer for q in all_quality if q.wer is not None]
|
||||
|
||||
if snrs:
|
||||
summary.mean_snr_db = round(statistics.mean(snrs), 2)
|
||||
if pesqs:
|
||||
summary.mean_pesq = round(statistics.mean(pesqs), 3)
|
||||
if wers_q:
|
||||
summary.mean_wer_quantized = round(statistics.mean(wers_q), 3)
|
||||
|
||||
# Also compute baseline WER for comparison
|
||||
if whisper_model is not None:
|
||||
from jiwer import wer as compute_wer_score
|
||||
|
||||
baseline_wers = []
|
||||
for qr in all_quality:
|
||||
if qr.baseline_transcript is not None:
|
||||
baseline_wers.append(compute_wer_score(qr.sentence, qr.baseline_transcript))
|
||||
if baseline_wers:
|
||||
summary.mean_wer_baseline = round(statistics.mean(baseline_wers), 3)
|
||||
|
||||
# --- Write reports ---
|
||||
write_csv(summaries, output_dir)
|
||||
write_quality_csv(summaries, output_dir)
|
||||
write_markdown_report(summaries, output_dir)
|
||||
write_json_summary(summaries, output_dir)
|
||||
|
||||
# --- Print summary ---
|
||||
print("\n" + "=" * 60)
|
||||
print("EVALUATION COMPLETE")
|
||||
print("=" * 60)
|
||||
print(
|
||||
f"{'Config':<20} {'Mean RTS':>10} {'Speedup':>10} {'Size MB':>10} {'SNR dB':>10} {'PESQ':>8} {'WER':>8}"
|
||||
)
|
||||
print("-" * 86)
|
||||
for s in summaries:
|
||||
speedup = f"{s.rts_vs_baseline:.2f}x" if s.rts_vs_baseline else "---"
|
||||
snr = f"{s.mean_snr_db:.1f}" if s.mean_snr_db is not None else "---"
|
||||
pesq_s = f"{s.mean_pesq:.3f}" if s.mean_pesq is not None else "---"
|
||||
wer_s = f"{s.mean_wer_quantized:.3f}" if s.mean_wer_quantized is not None else "---"
|
||||
print(
|
||||
f"{s.config_id:<20} {s.mean_rts:>10.2f}x {speedup:>10} {s.model_size_mb:>10.1f} {snr:>10} {pesq_s:>8} {wer_s:>8}"
|
||||
)
|
||||
print(f"\nAudio files + report: {output_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,59 @@
|
||||
"""
|
||||
Tests for int8 quantization.
|
||||
|
||||
Verifies:
|
||||
1. Quantized model produces valid audio (not silence, not NaN)
|
||||
2. load_model(quantize=True) applies quantization
|
||||
3. CLI --quantize flag is accepted
|
||||
4. Backend detection works
|
||||
"""
|
||||
|
||||
import torch
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from pocket_tts import TTSModel
|
||||
from pocket_tts.main import cli_app
|
||||
from pocket_tts.quantization import _get_backend
|
||||
|
||||
SHORT_TEXT = "Hello, this is a test."
|
||||
TEST_VOICE = "alba"
|
||||
|
||||
runner = CliRunner()
|
||||
|
||||
|
||||
def test_quantized_model_produces_audio():
|
||||
model = TTSModel.load_model(quantize=True)
|
||||
voice_state = model.get_state_for_audio_prompt(TEST_VOICE)
|
||||
audio = model.generate_audio(voice_state, SHORT_TEXT)
|
||||
assert audio is not None
|
||||
assert len(audio) > 0
|
||||
assert not torch.isnan(audio).any(), "Audio contains NaN values"
|
||||
assert not torch.isinf(audio).any(), "Audio contains Inf values"
|
||||
assert audio.abs().max() > 0, "Audio is silent"
|
||||
|
||||
|
||||
def test_quantize_flag_applies_quantization():
|
||||
model_q = TTSModel.load_model(quantize=True)
|
||||
model_b = TTSModel.load_model(quantize=False)
|
||||
# Quantized model's Linear layers should have different weight types than baseline
|
||||
layer_q = model_q.flow_lm.transformer.layers[0]
|
||||
layer_b = model_b.flow_lm.transformer.layers[0]
|
||||
weight_type_q = type(layer_q.self_attn.in_proj.weight).__name__
|
||||
weight_type_b = type(layer_b.self_attn.in_proj.weight).__name__
|
||||
assert weight_type_q != weight_type_b, (
|
||||
f"Attention weights should differ after quantization, got {weight_type_q} for both"
|
||||
)
|
||||
|
||||
|
||||
def test_cli_quantize_flag(tmp_path):
|
||||
output_file = tmp_path / "quantized_output.wav"
|
||||
result = runner.invoke(
|
||||
cli_app,
|
||||
["generate", "--quantize", "--text", SHORT_TEXT, "--output-path", str(output_file), "-q"],
|
||||
)
|
||||
assert result.exit_code == 0, f"CLI failed: {result.output}"
|
||||
|
||||
|
||||
def test_backend_detection():
|
||||
backend = _get_backend()
|
||||
assert backend in ("torchao", "torch.ao")
|
||||
Reference in New Issue
Block a user