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

Fix support for devices that are not cpu (#165)

Browse Source 
Fixes https://github.com/kyutai-labs/pocket-tts/issues/164
This commit is contained in:
Gabriel de Marmiesse
2026-04-22 19:18:48 +02:00
committed by GitHub
parent 5304da453b
commit 13542b4f9d
1 changed files with 10 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pocket_tts/models/tts_model.py
+10 -7
View File
@@ -90,8 +90,8 @@ class TTSModel(nn.Module):
self.remove_semicolons = remove_semicolons
@property
def device(self) -> str:
return next(self.parameters()).device.type
def device(self) -> torch.device:
return next(self.parameters()).device
@property
def sample_rate(self) -> int:
@@ -848,7 +848,7 @@ class TTSModel(nn.Module):
if isinstance(audio_conditioning, str):
audio_conditioning = download_if_necessary(audio_conditioning)
return _import_model_state(audio_conditioning)
return _import_model_state(audio_conditioning, self.device)
elif (
isinstance(audio_conditioning, str)
@@ -864,7 +864,8 @@ class TTSModel(nn.Module):
return _import_model_state(
download_if_necessary(
get_predefined_voice(language=self.origin.stem, name=audio_conditioning)
)
),
self.device,
)
if not self.has_voice_cloning and isinstance(audio_conditioning, (str, Path)):
@@ -1051,7 +1052,9 @@ def export_model_state(model_state: dict[str, dict[str, torch.Tensor]], dest: st
safetensors.torch.save_file(dict_to_store, dest)
def _import_model_state(source: str | Path) -> dict[str, dict[str, torch.Tensor]]:
def _import_model_state(
source: str | Path, device: torch.device
) -> dict[str, dict[str, torch.Tensor]]:
result = {}
with safetensors.safe_open(source, framework="pt") as f:
for key in f.keys():
@@ -1062,8 +1065,8 @@ def _import_model_state(source: str | Path) -> dict[str, dict[str, torch.Tensor]
# but it's not needed anymore
tensor = f.get_tensor(key)
result[module_name]["offset"] = torch.full(
(1,), fill_value=tensor.shape[0], dtype=torch.long, device=tensor.device
(1,), fill_value=tensor.shape[0], dtype=torch.long, device=device
)
else:
result[module_name][tensor_key] = f.get_tensor(key)
result[module_name][tensor_key] = f.get_tensor(key).to(device)
return result