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:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user