valenti/spark-pocket-tts
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix voice cloning with quantization (#181)

Browse Source 
Fix https://github.com/kyutai-labs/pocket-tts/issues/178
This commit is contained in:
Gabriel de Marmiesse
2026-05-04 15:50:42 +02:00
committed by GitHub
parent 7db278ee57
commit a882e0680e
1 changed files with 9 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pocket_tts/modules/transformer.py
+9 -2
View File
@@ -118,8 +118,15 @@ class StreamingMultiheadAttention(StatefulModule):
self.out_proj = nn.Linear(embed_dim, mult * embed_dim, bias=False) self.out_proj = nn.Linear(embed_dim, mult * embed_dim, bias=False)
def init_state(self, batch_size: int, sequence_length: int) -> dict[str, torch.Tensor]: def init_state(self, batch_size: int, sequence_length: int) -> dict[str, torch.Tensor]:
device = self.in_proj.weight.device weight = self.in_proj.weight
dtype = self.in_proj.weight.dtype if callable(weight):
# torch.ao dynamic-quantized Linear: weight is a method returning a
# qint8 tensor, but activations stay float32 — that's what the cache holds.
device = weight().device
dtype = torch.float32
else:
device = weight.device
dtype = weight.dtype
return self._cache_backend.init_state(batch_size, sequence_length, device, dtype) return self._cache_backend.init_state(batch_size, sequence_length, device, dtype)
def increment_step(self, state: dict, increment: int = 1): def increment_step(self, state: dict, increment: int = 1):