Copy the files in the new repo for public release

pocket_tts/utils/__init__.py

@@ -0,0 +1 @@
+"""Utilities."""

pocket_tts/utils/config.py

@@ -0,0 +1,122 @@
+"""Configuration models for loading YAML config files."""
+
+from pathlib import Path
+
+import yaml
+from pydantic import BaseModel, ConfigDict
+
+
+class StrictModel(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+
+# Flow configuration
+class FlowConfig(StrictModel):
+    dim: int
+    depth: int
+
+
+# Transformer configuration for FlowLM
+class FlowLMTransformerConfig(StrictModel):
+    hidden_scale: int
+    max_period: int
+    d_model: int
+    num_heads: int
+    num_layers: int
+
+
+class LookupTable(StrictModel):
+    dim: int
+    n_bins: int
+    tokenizer: str
+    tokenizer_path: str
+
+
+# Root configuration
+class FlowLMConfig(StrictModel):
+    """Root configuration model for YAML config files."""
+
+    dtype: str
+
+    # Nested configurations
+    flow: FlowConfig
+    transformer: FlowLMTransformerConfig
+
+    # conditioning
+    lookup_table: LookupTable
+    weights_path: str | None = None
+
+
+# SEANet configuration
+class SEANetConfig(StrictModel):
+    dimension: int
+    channels: int
+    n_filters: int
+    n_residual_layers: int
+    ratios: list[int]
+    kernel_size: int
+    residual_kernel_size: int
+    last_kernel_size: int
+    dilation_base: int
+    pad_mode: str
+    compress: int
+
+
+# Transformer configuration for Mimi
+class MimiTransformerConfig(StrictModel):
+    d_model: int
+    input_dimension: int
+    output_dimensions: tuple[int, ...]
+    num_heads: int
+    num_layers: int
+    layer_scale: float
+    context: int
+    max_period: float = 10000.0
+    dim_feedforward: int
+
+
+# Quantizer configuration
+class QuantizerConfig(StrictModel):
+    dimension: int
+    output_dimension: int
+
+
+# Root configuration
+class MimiConfig(StrictModel):
+    """Root configuration model for Mimi YAML config files."""
+
+    dtype: str
+
+    # Sample rate and channels
+    sample_rate: int
+    channels: int
+    frame_rate: float
+
+    # SEANet configurations
+    seanet: SEANetConfig
+
+    # Transformer
+    transformer: MimiTransformerConfig
+
+    # Quantizer
+    quantizer: QuantizerConfig
+    weights_path: str | None = None
+
+
+class Config(StrictModel):
+    flow_lm: FlowLMConfig
+    mimi: MimiConfig
+    weights_path: str | None = None
+    weights_path_without_voice_cloning: str | None = None
+
+
+def load_config(yaml_path: str | Path) -> Config:
+    yaml_path = Path(yaml_path)
+
+    if not yaml_path.exists():
+        raise FileNotFoundError(f"Config file not found: {yaml_path}")
+
+    with open(yaml_path, "r") as f:
+        config_dict = yaml.safe_load(f)
+
+    return Config(**config_dict)

pocket_tts/utils/debugging.py

@@ -0,0 +1,26 @@
+import torch
+from torch.utils._python_dispatch import TorchDispatchMode
+
+
+def to_str(obj):
+    if isinstance(obj, (torch.Tensor, torch.nn.Parameter)):
+        return f"T(s={list(obj.shape)})"
+    elif isinstance(obj, (list, tuple)):
+        return "[" + ", ".join(to_str(o) for o in obj) + "]"
+    elif isinstance(obj, dict):
+        return "{" + ", ".join(f"{to_str(k)}: {to_str(v)}" for k, v in obj.items()) + "}"
+    else:
+        return str(obj)
+
+
+class LoggingMode(TorchDispatchMode):
+    """Useful to check implementation differences."""
+
+    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
+        output = func(*args, **kwargs or {})
+        print(
+            f"Aten function called: {func}, args: "
+            f"{to_str(args)}, kwargs: {to_str(kwargs)} -> "
+            f"output: {to_str(output)}"
+        )
+        return output

pocket_tts/utils/logging_utils.py

@@ -0,0 +1,41 @@
+import logging
+from contextlib import contextmanager
+
+
+class PocketTTSFilter(logging.Filter):
+    def filter(self, record):
+        return record.name.startswith("pocket_tts")
+
+
+@contextmanager
+def enable_logging(library_name, level):
+    # Get the specific logger and its parent
+    logger = logging.getLogger(library_name)
+    parent_logger = logging.getLogger("pocket_tts")
+
+    # Store original configuration
+    old_level = logger.level
+    old_parent_level = parent_logger.level
+    old_handlers = parent_logger.handlers.copy()
+
+    # Configure logging format for pocket_tts logger
+    parent_logger.setLevel(level)
+
+    # Clear existing handlers and add our custom formatter with filter
+    parent_logger.handlers.clear()
+    handler = logging.StreamHandler()
+    formatter = logging.Formatter("%(levelname)s: %(message)s")
+    handler.setFormatter(formatter)
+    handler.addFilter(PocketTTSFilter())
+    parent_logger.addHandler(handler)
+    parent_logger.propagate = False
+
+    try:
+        yield logger
+    finally:
+        # Restore original configuration
+        logger.setLevel(old_level)
+        parent_logger.setLevel(old_parent_level)
+        parent_logger.handlers.clear()
+        for h in old_handlers:
+            parent_logger.addHandler(h)

pocket_tts/utils/utils.py

@@ -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"]

pocket_tts/utils/weights_loading.py

@@ -0,0 +1,35 @@
+from pathlib import Path
+
+import safetensors
+
+
+def get_flow_lm_state_dict(path: Path) -> dict:
+    state_dict = {}
+    with safetensors.safe_open(path, framework="pt", device="cpu") as f:
+        for key in f.keys():
+            if (
+                key.startswith("flow.w_s_t.")
+                or key == "condition_provider.conditioners.transcript_in_segment.learnt_padding"
+                or key == "condition_provider.conditioners.speaker_wavs.learnt_padding"
+            ):
+                # skip lookup table weights
+                continue
+            new_name = key
+            if key == "condition_provider.conditioners.transcript_in_segment.embed.weight":
+                new_name = "conditioner.embed.weight"
+            if key == "condition_provider.conditioners.speaker_wavs.output_proj.weight":
+                new_name = "speaker_proj_weight"
+            state_dict[new_name] = f.get_tensor(key)
+    return state_dict
+
+
+def get_mimi_state_dict(path: Path) -> dict:
+    state_dict = {}
+    with safetensors.safe_open(path, framework="pt", device="cpu") as f:
+        for key in f.keys():
+            if key.startswith("model.quantizer.vq.") or key == "model.quantizer.logvar_proj.weight":
+                # skip vq weights
+                continue
+
+            state_dict[key.removeprefix("model.")] = f.get_tensor(key)
+    return state_dict