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
2026-03-24 11:29:14 +01:00
committed by GitHub
parent 0dcf6cb442
commit 670e777260
8 changed files with 1506 additions and 512 deletions
+13 -2
View File
@@ -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(