nabil-tazi
670e777260
# 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.
60 lines
1.9 KiB
Python
60 lines
1.9 KiB
Python
"""
|
|
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")
|