Copy the files in the new repo for public release

2026-01-12 11:25:16 +01:00
commit f5fa328a58
54 changed files with 6411 additions and 0 deletions
@@ -0,0 +1,102 @@
+import hashlib
+import logging
+import time
+from pathlib import Path
+
+import requests
+import safetensors.torch
+import torch
+from huggingface_hub import hf_hub_download
+from torch import nn
+
+PROJECT_ROOT = Path(__file__).parent.parent.parent
+
+PREDEFINED_VOICES = {
+    "marius": "hf://kyutai/pocket-tts-without-voice-cloning/embeddings/marius.safetensors",
+    "javert": "hf://kyutai/pocket-tts-without-voice-cloning/embeddings/javert.safetensors",
+    "jean": "hf://kyutai/pocket-tts-without-voice-cloning/embeddings/jean.safetensors",
+    "fantine": "hf://kyutai/pocket-tts-without-voice-cloning/embeddings/fantine.safetensors",
+    "cosette": "hf://kyutai/pocket-tts-without-voice-cloning/embeddings/cosette.safetensors",
+    "eponine": "hf://kyutai/pocket-tts-without-voice-cloning/embeddings/eponine.safetensors",
+    "azelma": "hf://kyutai/pocket-tts-without-voice-cloning/embeddings/azelma.safetensors",
+}
+
+
+def make_cache_directory() -> Path:
+    cache_dir = Path.home() / ".cache" / "pocket_tts"
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    return cache_dir
+
+
+def print_nb_parameters(model: nn.Module, model_name: str):
+    logger = logging.getLogger(__name__)
+    state_dict = model.state_dict()
+    total = 0
+    for key, value in state_dict.items():
+        logger.info("%s: %,d", key, value.numel())
+        total += value.numel()
+    logger.info("Total number of parameters in %s: %,d", model_name, total)
+
+
+def size_of_dict(state_dict: dict) -> int:
+    total_size = 0
+    for value in state_dict.values():
+        if isinstance(value, torch.Tensor):
+            total_size += value.numel() * value.element_size()
+        elif isinstance(value, dict):
+            total_size += size_of_dict(value)
+    return total_size
+
+
+class display_execution_time:
+    def __init__(self, task_name: str, print_output: bool = True):
+        self.task_name = task_name
+        self.print_output = print_output
+        self.start_time = None
+        self.elapsed_time_ms = None
+        self.logger = logging.getLogger(__name__)
+
+    def __enter__(self):
+        self.start_time = time.monotonic()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        end_time = time.monotonic()
+        self.elapsed_time_ms = int((end_time - self.start_time) * 1000)
+        if self.print_output:
+            self.logger.info("%s took %d ms", self.task_name, self.elapsed_time_ms)
+        return False  # Don't suppress exceptions
+
+
+def download_if_necessary(file_path: str) -> Path:
+    if file_path.startswith("http://") or file_path.startswith("https://"):
+        cache_dir = make_cache_directory()
+        cached_file = cache_dir / (
+            hashlib.sha256(file_path.encode()).hexdigest() + "." + file_path.split(".")[-1]
+        )
+        if not cached_file.exists():
+            response = requests.get(file_path)
+            response.raise_for_status()
+            with open(cached_file, "wb") as f:
+                f.write(response.content)
+        return cached_file
+    elif file_path.startswith("hf://"):
+        file_path = file_path.removeprefix("hf://")
+        splitted = file_path.split("/")
+        repo_id = "/".join(splitted[:2])
+        filename = "/".join(splitted[2:])
+        cached_file = hf_hub_download(repo_id=repo_id, filename=filename)
+        return Path(cached_file)
+    else:
+        return Path(file_path)
+
+
+def load_predefined_voice(voice_name: str) -> torch.Tensor:
+    if voice_name not in PREDEFINED_VOICES:
+        raise ValueError(
+            f"Predefined voice '{voice_name}' not found"
+            f", available voices are {list(PREDEFINED_VOICES)}."
+        )
+    voice_file = download_if_necessary(PREDEFINED_VOICES[voice_name])
+    # There is only one tensor in the file.
+    return safetensors.torch.load_file(voice_file)["audio_prompt"]