From 2b1cf209d662f61fbafdcb623c94f55ae2e23834 Mon Sep 17 00:00:00 2001 From: WildAi <2853742+wildminder@users.noreply.github.com> Date: Wed, 7 Oct 2026 13:22:03 +0300 Subject: [PATCH] feat(loader): dedupe shared weights across packs - ve/s3gen/s3gen_v3 safetensors + grapheme tokenizer are byte-identical across every repo that ships them (sha256-verified): fetched at most once into a reserved _shared store, hardlinked into each pack - downloads carry pack-unique files only (allow_patterns); legacy full-repo folders satisfy the store before anything is fetched - vendored MTL loads ve/s3gen from either serialization, safetensors first; upstream's strict load refused every SLP pack (P-036) - .pt twins (s3gen.pt, s3gen_v3.pt, t3_cfg.pt, ve.pt) are never fetched --- README.md | 39 ++++++- modules/loader.py | 186 +++++++++++++++++++++++------- modules/model_discovery.py | 7 +- modules/model_info.py | 59 ++++++---- src/PATCHES.md | 2 +- src/chatterbox/mtl_tts.py | 65 ++++++++--- tests/test_local_checkpoint.py | 18 ++- tests/test_patcher.py | 11 +- tests/test_shared_assets.py | 200 +++++++++++++++++++++++++++++++++ tests/test_tts_routing.py | 4 +- 10 files changed, 498 insertions(+), 93 deletions(-) create mode 100644 tests/test_shared_assets.py diff --git a/README.md b/README.md index d54d214..cd7bbfa 100644 --- a/README.md +++ b/README.md @@ -230,14 +230,41 @@ downloaded: - **Single Language Packs** (the six `ResembleAI/Chatterbox-Multilingual-` finetunes) load through the **Chatterbox Multilingual TTS 🌍** node. Each ships its own per-language t3 checkpoint (`t3_zh_cmn`, `t3_es_mx_latam`, - `t3_es_es`, `t3_pt_br`, `t3_pt_pt`, `t3_hi`) plus the multilingual s3gen and - grapheme tokenizer, so the pack folder is self-describing — the pack decides - what loads, not the variant widget. Only zh-cmn ships a built-in voice - (`conds.pt`); the other five fetch the shared `ve.pt` from - `ResembleAI/chatterbox` on first load and **require a reference audio clip**. - The `language` widget is ignored — every pack speaks its own language. + `t3_es_es`, `t3_pt_br`, `t3_pt_pt`, `t3_hi`); the decoder, grapheme + tokenizer and voice encoder are byte-identical across all six and come from + the shared store instead of being downloaded six times (see below). Only + zh-cmn ships a built-in voice (`conds.pt`); the other five **require a + reference audio clip**. The `language` widget is ignored — every pack + speaks its own language. - Classic packs (`t3_cfg.safetensors`) are the voice-conversion packs too. +#### How packs share files + +Four files are byte-identical across every repo that ships them (verified by +sha256): `ve.safetensors`, `s3gen.safetensors`, `s3gen_v3.safetensors` and the +grapheme tokenizer `grapheme_mtl_merged_expanded_v1.json`. The node therefore +fetches each of them at most once, into a reserved store at + +```text +ComfyUI/models/tts/chatterbox/_shared/ +``` + +and hardlinks it into every pack that needs it — a hardlink costs no extra +disk, and each pack folder stays self-contained. When one of these files +already exists anywhere on disk (a full `chatterbox/` download from an older +version counts), it is linked from there and nothing is fetched at all. + +A pack download consequently carries only its unique files: an extra Single +Language Pack is its ~2 GB t3 checkpoint and nothing else. The one exception +is Flash, whose `s3gen.safetensors` is a different decoder despite the shared +name — it always downloads its own. + +Downloads never fetch the legacy `.pt` twins (`s3gen.pt`, `s3gen_v3.pt`, +`t3_cfg.pt`, `ve.pt`): they hold the same weights as their `.safetensors` +siblings and nothing loads them. Existing folders that carry the twins keep +working — you can delete those four files from a full `chatterbox/` download +to reclaim ~3 GB. + Folders from earlier builds named like a pack plus a variant (`chatterbox-nano__turbo`) are ordinary packs as far as the scanner is concerned and keep loading; rename or delete them at will. diff --git a/modules/loader.py b/modules/loader.py index 388bbe1..e3623a3 100644 --- a/modules/loader.py +++ b/modules/loader.py @@ -33,16 +33,24 @@ from __future__ import annotations import inspect import logging import os +import shutil from typing import Any, Optional import torch -from .constants import DEFAULT_MODEL_PACK_NAME, MTL_T3_MODELS +from .constants import ( + CHATTERBOX_FILES_TO_DOWNLOAD, + DEFAULT_MODEL_PACK_NAME, + MTL_T3_MODELS, +) from .model_discovery import get_model_pack_path from .model_info import ( AVAILABLE_CHATTERBOX_MODELS, CHATTERBOX_WEIGHT_FILES, MODEL_CONFIGS, + MTL_TOKENIZER_FILE, + SHARED_ASSET_DIR_NAME, + SHARED_ASSET_FILES, is_chatterbox_model_dir, ) @@ -76,15 +84,24 @@ _WEIGHT_ATTRIBUTES = ("t3", "s3gen", "ve") # checkpoint is appended by the caller, because its name is not fixed -- and # so is the s3gen decoder's: the Single Language Packs ship ``s3gen_v3.pt`` # (P-036), the shared repo ships ``s3gen.pt``. -MTL_CHECKPOINT_FILES: tuple[str, ...] = ( - "ve.pt", - "grapheme_mtl_merged_expanded_v1.json", -) +#: Files every multilingual pack holds *besides* the resolved T3, voice +#: encoder and s3gen checkpoints (P-036 resolves those per pack): just the +#: grapheme tokenizer. +MTL_CHECKPOINT_FILES: tuple[str, ...] = (MTL_TOKENIZER_FILE,) #: The s3gen filenames a multilingual pack may carry, in resolution order: -#: the shared repo's canonical name first, the Single Language Packs' v3 -#: retrain second. -MTL_S3GEN_CHECKPOINTS: tuple[str, ...] = ("s3gen.pt", "s3gen_v3.pt") +#: the Single Language Packs' v3 retrain (safetensors first), then the shared +#: repo's canonical decoder. +MTL_S3GEN_CHECKPOINTS: tuple[str, ...] = ( + "s3gen_v3.safetensors", + "s3gen_v3.pt", + "s3gen.safetensors", + "s3gen.pt", +) + +#: Pack folders that historically hold the shared repo's weights. Read-only +#: search candidates for shared assets; downloads never write into them. +LEGACY_SHARED_DIRS: tuple[str, ...] = ("chatterbox", DEFAULT_MODEL_PACK_NAME) #: Turbo and Nano share one layout, and so share one completeness list. The #: tokenizer is checked separately: Turbo ships ``tokenizer.json``, Nano ships @@ -533,7 +550,12 @@ class ChatterboxLoader: # -- download ------------------------------------------------------ @staticmethod - def _ensure_downloaded(repo_id: str, local_dir: str, model_name: str = "") -> None: + def _ensure_downloaded( + repo_id: str, + local_dir: str, + model_name: str = "", + allow_patterns: Optional[list[str]] = None, + ) -> None: """Fetch an official pack from ``repo_id`` when it is not on disk. The completeness check is :meth:`is_pack_complete` rather than a @@ -566,6 +588,8 @@ class ChatterboxLoader: from huggingface_hub import snapshot_download kwargs: dict[str, Any] = {"repo_id": repo_id, "local_dir": local_dir} + if allow_patterns: + kwargs["allow_patterns"] = list(allow_patterns) try: accepts_symlinks = ( "local_dir_use_symlinks" in inspect.signature(snapshot_download).parameters @@ -583,6 +607,7 @@ class ChatterboxLoader: repo_id: str, local_dir: str, client_id: Optional[str] = None, + allow_patterns: Optional[list[str]] = None, ) -> None: """Ensure the pack exists, preferring the progress-tracking manager. @@ -591,12 +616,25 @@ class ChatterboxLoader: blocks the prompt with no way out. The manager is therefore the path taken, and :meth:`_ensure_downloaded` is the fallback for when it cannot be imported (or has been torn down). + + ``allow_patterns`` narrows the fetch to the named files -- the shared + weights are deduplicated away (see ``SHARED_ASSET_FILES``), so a pack + download must not drag every byte of the repo along. The manager's API + has no patterns parameter, so a filtered fetch goes straight to + ``snapshot_download`` (which resumes and skips complete files); the + manager keeps serving the unfiltered legacy path. """ if ChatterboxLoader.is_pack_complete(local_dir): return if not repo_id: return + if allow_patterns: + ChatterboxLoader._ensure_downloaded( + repo_id, local_dir, model_name, allow_patterns + ) + return + manager = None try: from .downloader import get_download_manager @@ -675,7 +713,22 @@ class ChatterboxLoader: model_name, pack_dir, ) - ChatterboxLoader.download_pack(model_name, repo_id, pack_dir, client_id) + files = list(entry.get("files") or CHATTERBOX_FILES_TO_DOWNLOAD) + unique = set(entry.get("unique_files") or ()) + ChatterboxLoader.download_pack( + model_name, + repo_id, + pack_dir, + client_id, + allow_patterns=[ + name + for name in files + if name not in SHARED_ASSET_FILES or name in unique + ], + ) + ChatterboxLoader._ensure_shared_assets( + pack_dir, [name for name in files if name in SHARED_ASSET_FILES] + ) try: os.makedirs(pack_dir, exist_ok=True) except OSError as exc: @@ -734,6 +787,59 @@ class ChatterboxLoader: """ return [name for name in required if not os.path.isfile(os.path.join(pack_dir, name))] + @staticmethod + def _shared_store_dir() -> str: + """The reserved ``_shared`` directory under the Chatterbox models root.""" + from .model_discovery import get_models_root + + return os.path.join(get_models_root(), SHARED_ASSET_DIR_NAME) + + @staticmethod + def _ensure_shared_assets(pack_dir: str, names) -> None: + """Materialize shared weights into ``pack_dir`` under canonical names. + + The shared files (``SHARED_ASSET_FILES``) are byte-identical across + every repo that ships them, so each is fetched at most once. Search + order: the pack itself (nothing to do), the reserved ``_shared`` + store, then the legacy full-repo folders a user may already have; + only when the file exists nowhere does one ``hf_hub_download`` from + the shared Chatterbox repo populate the store. The store copy is then + hardlinked into the pack -- both live under the Chatterbox models + root, so they share a volume and the link costs no extra disk -- with + a plain copy as the fallback for exotic filesystems. + """ + from huggingface_hub import hf_hub_download + + from .constants import CHATTERBOX_REPO_ID + + store = ChatterboxLoader._shared_store_dir() + root = os.path.dirname(store) + for name in names: + target = os.path.join(pack_dir, name) + if os.path.isfile(target): + continue + source = None + candidates = [store] + candidates.extend(os.path.join(root, d) for d in LEGACY_SHARED_DIRS) + for candidate in candidates: + probe = os.path.join(candidate, name) + if os.path.isfile(probe): + source = probe + break + if source is None: + os.makedirs(store, exist_ok=True) + source = hf_hub_download( + repo_id=CHATTERBOX_REPO_ID, filename=name, local_dir=store + ) + try: + os.link(source, target) + except OSError: + try: + os.makedirs(pack_dir, exist_ok=True) + except OSError: + pass + shutil.copyfile(source, target) + @staticmethod def _ensure_entry_point_dir( pack_dir: str, @@ -741,13 +847,12 @@ class ChatterboxLoader: required: tuple[str, ...], model_name: str, client_id: Optional[str] = None, - companion_files: Optional[dict[str, str]] = None, ) -> None: """Fetch ``repo_id`` into ``pack_dir`` when a required file is absent. - The shared prologue for the multilingual and Turbo entry points. Both - are ordinary Hugging Face snapshots with their own file layouts, so the - completeness test is ``required`` -- not + The shared prologue for the multilingual, Turbo and Flash entry + points. All are ordinary Hugging Face snapshots with their own file + layouts, so the completeness test is ``required`` -- not :meth:`is_pack_complete`, which is written for the *classic* layout and would therefore re-download a perfectly good multilingual pack on every run. @@ -756,10 +861,12 @@ class ChatterboxLoader: drops under ``models/tts/chatterbox``) satisfy the same test and take the same path: there is no special case for them anywhere in this class. - ``companion_files`` maps a required file that lives in a *different* - repository to that repository's id -- the Single Language Packs carry - no ``ve.pt`` of their own; it is fetched per-file from the shared - Chatterbox repo (a ~17 MB voice encoder, no download manager needed). + Downloads are narrowed by the registry entry's ``files`` list minus + the shared assets (``SHARED_ASSET_FILES``) -- those are byte-identical + across the family and are materialized afterwards by + :meth:`_ensure_shared_assets`, hardlinked out of the ``_shared`` + store. An entry-less (community) pack falls back to a whole-repo + fetch: its layout is unknown. """ if os.path.isdir(pack_dir) and not ChatterboxLoader._missing_checkpoint_files( pack_dir, required @@ -767,29 +874,29 @@ class ChatterboxLoader: return if not repo_id: return + entry = ChatterboxLoader._registry_entry(model_name) or {} + unique = set(entry.get("unique_files") or ()) + shared_needed = [ + name for name in required if name in SHARED_ASSET_FILES and name not in unique + ] + patterns = None + if entry.get("files"): + patterns = [ + name + for name in entry["files"] + if name not in SHARED_ASSET_FILES or name in unique + ] logging.info( "[Chatterbox TTS] Fetching '%s' from %s into '%s'...", model_name, repo_id, pack_dir, ) - ChatterboxLoader.download_pack(model_name, repo_id, pack_dir, client_id) - if companion_files: - from huggingface_hub import hf_hub_download + ChatterboxLoader.download_pack( + model_name, repo_id, pack_dir, client_id, allow_patterns=patterns + ) + ChatterboxLoader._ensure_shared_assets(pack_dir, shared_needed) - for filename, companion_repo in companion_files.items(): - if not companion_repo or os.path.isfile(os.path.join(pack_dir, filename)): - continue - logging.info( - "[Chatterbox TTS] Fetching companion file '%s' from %s...", - filename, - companion_repo, - ) - hf_hub_download( - repo_id=companion_repo, filename=filename, local_dir=pack_dir - ) - - @staticmethod @staticmethod def _registry_entry(model_name: str) -> Optional[dict]: """The registry config for ``model_name``, official or local, or ``None``.""" @@ -835,10 +942,10 @@ class ChatterboxLoader: """The s3gen decoder filename a multilingual pack actually holds. The directory wins (a local community pack is on disk when this is - called with a ``model_path``); a registry entry whose ``files`` list - names the v3 decoder follows; the shared repo's canonical - ``s3gen.pt`` is the default, so the missing-file message keeps - naming it. + called with a ``model_path``), safetensors before any legacy ``.pt`` + twin; a registry entry whose ``files`` list names a decoder follows; + the Single Language Packs' ``s3gen_v3.safetensors`` is the default, + so the missing-file message keeps naming it. """ if os.path.isdir(pack_dir): for name in MTL_S3GEN_CHECKPOINTS: @@ -881,7 +988,7 @@ class ChatterboxLoader: _refuse_wrong_variant_pack(pack_dir, model_name, entry_point="multilingual") t3_model = ChatterboxLoader._resolve_mtl_t3_file(pack_dir, model_name, t3_model) s3gen_file = ChatterboxLoader._resolve_mtl_s3gen_file(pack_dir, model_name) - required = (*MTL_CHECKPOINT_FILES, s3gen_file, t3_model) + required = (*MTL_CHECKPOINT_FILES, "ve.safetensors", s3gen_file, t3_model) from .constants import CHATTERBOX_REPO_ID @@ -892,7 +999,6 @@ class ChatterboxLoader: required, model_name, client_id, - companion_files=entry.get("companion_files"), ) missing = ChatterboxLoader._missing_checkpoint_files(pack_dir, required) diff --git a/modules/model_discovery.py b/modules/model_discovery.py index 1d71b62..a91541f 100644 --- a/modules/model_discovery.py +++ b/modules/model_discovery.py @@ -87,13 +87,18 @@ def _scan(root: str) -> List[Tuple[str, str]]: Uses ``os.scandir`` rather than ``os.listdir`` so the directory is read once and each entry's type comes from the readdir data already in hand. Every subdirectory counts as a pack: a partially downloaded one must stay - selectable so the download can finish it. + selectable so the download can finish it. The reserved shared-weight + store (``SHARED_ASSET_DIR_NAME``) is not a pack and never appears. """ + from .model_info import SHARED_ASSET_DIR_NAME + packs: List[Tuple[str, str]] = [] try: with os.scandir(root) as entries: for entry in entries: try: + if entry.name == SHARED_ASSET_DIR_NAME: + continue if entry.is_dir(): packs.append((entry.name, entry.path)) except OSError as exc: diff --git a/modules/model_info.py b/modules/model_info.py index a30c57e..afe8182 100644 --- a/modules/model_info.py +++ b/modules/model_info.py @@ -68,6 +68,21 @@ __all__ = [ # grapheme map the vendored ``MTLTokenizer`` loads instead of a fast tokenizer. MTL_TOKENIZER_FILE = "grapheme_mtl_merged_expanded_v1.json" +# ── Shared assets ───────────────────────────────────────────────────── +# Byte-identical across every repo that ships them (verified 2026-10-07 by +# sha256 and HF LFS OIDs), so they are fetched at most once into the +# reserved `_shared` store and hardlinked into each pack that needs them. +# Downloads never fetch them into a pack directly. +SHARED_ASSET_DIR_NAME = "_shared" +SHARED_ASSET_FILES = frozenset( + { + "ve.safetensors", + "s3gen.safetensors", + "s3gen_v3.safetensors", + MTL_TOKENIZER_FILE, + } +) + MODEL_CONFIGS: Dict[str, Dict[str, Any]] = { # The pack a saved workflow names by default. Monolingual English, voice @@ -87,7 +102,7 @@ MODEL_CONFIGS: Dict[str, Dict[str, Any]] = { "architecture": "chatterbox_multilingual", "sample_rate": 24000, "size_gb": 2.7, - "files": ["t3_cfg.safetensors", "ve.safetensors", "s3gen.safetensors", "tokenizer.json", "conds.pt"], + "files": ["t3_cfg.safetensors", "ve.safetensors", "s3gen.safetensors", MTL_TOKENIZER_FILE, "conds.pt"], "description": "Original single-language T3 weights, used by the multilingual loader as its v1.", }, "T3_v2": { @@ -95,7 +110,7 @@ MODEL_CONFIGS: Dict[str, Dict[str, Any]] = { "architecture": "chatterbox_multilingual", "sample_rate": 24000, "size_gb": 2.7, - "files": ["t3_mtl23ls_v2.safetensors", "ve.safetensors", "s3gen.safetensors", "tokenizer.json", "conds.pt"], + "files": ["t3_mtl23ls_v2.safetensors", "ve.safetensors", "s3gen.safetensors", MTL_TOKENIZER_FILE, "conds.pt"], "description": "23-language multilingual T3 (upstream default).", }, "T3_v3": { @@ -103,21 +118,22 @@ MODEL_CONFIGS: Dict[str, Dict[str, Any]] = { "architecture": "chatterbox_multilingual", "sample_rate": 24000, "size_gb": 2.7, - "files": ["t3_mtl23ls_v3.safetensors", "ve.safetensors", "s3gen.safetensors", "tokenizer.json", "conds.pt"], + "files": ["t3_mtl23ls_v3.safetensors", "ve.safetensors", "s3gen.safetensors", MTL_TOKENIZER_FILE, "conds.pt"], "description": "23-language multilingual T3, newer checkpoint: better similarity, fewer hallucinations.", }, # The Single Language Pack: six dedicated finetunes of the multilingual # V3 family, one per priority language (2026-06 release). Each ships its - # own per-language t3 checkpoint ("t3_checkpoint") and s3gen_v3.pt; only - # zh-cmn bundles ve.pt and conds.pt -- the rest take ve.pt from the shared - # repo at load time ("companion_files") and have NO built-in voice, so a - # reference clip is required. + # own per-language t3 checkpoint ("t3_checkpoint"); the s3gen_v3 decoder, + # grapheme tokenizer and voice encoder are byte-identical across all six + # and come from the shared store (SHARED_ASSET_FILES). Only zh-cmn + # bundles a conds.pt (built-in voice); the rest have none, so a reference + # clip is required. "chatterbox-multilingual-zh-cmn": { "repo_id": "ResembleAI/Chatterbox-Multilingual-zh-cmn", "architecture": "chatterbox_multilingual", "sample_rate": 24000, "size_gb": 2.8, - "files": ["t3_zh_cmn.safetensors", "s3gen_v3.pt", MTL_TOKENIZER_FILE, "ve.pt", "conds.pt"], + "files": ["t3_zh_cmn.safetensors", "s3gen_v3.safetensors", MTL_TOKENIZER_FILE, "conds.pt"], "t3_checkpoint": "t3_zh_cmn.safetensors", "language": "zh", "description": "Single Language Pack: Chinese (zh). Ships a built-in voice.", @@ -127,10 +143,9 @@ MODEL_CONFIGS: Dict[str, Dict[str, Any]] = { "architecture": "chatterbox_multilingual", "sample_rate": 24000, "size_gb": 2.6, - "files": ["t3_es_mx_latam.safetensors", "s3gen_v3.pt", MTL_TOKENIZER_FILE], + "files": ["t3_es_mx_latam.safetensors", "s3gen_v3.safetensors", MTL_TOKENIZER_FILE], "t3_checkpoint": "t3_es_mx_latam.safetensors", "language": "es", - "companion_files": {"ve.pt": CHATTERBOX_REPO_ID}, "description": "Single Language Pack: Latin-American Spanish (es). Needs a reference clip.", }, "chatterbox-multilingual-es-es": { @@ -138,10 +153,9 @@ MODEL_CONFIGS: Dict[str, Dict[str, Any]] = { "architecture": "chatterbox_multilingual", "sample_rate": 24000, "size_gb": 2.6, - "files": ["t3_es_es.safetensors", "s3gen_v3.pt", MTL_TOKENIZER_FILE], + "files": ["t3_es_es.safetensors", "s3gen_v3.safetensors", MTL_TOKENIZER_FILE], "t3_checkpoint": "t3_es_es.safetensors", "language": "es", - "companion_files": {"ve.pt": CHATTERBOX_REPO_ID}, "description": "Single Language Pack: Spain Spanish (es). Needs a reference clip.", }, "chatterbox-multilingual-pt-br": { @@ -149,10 +163,9 @@ MODEL_CONFIGS: Dict[str, Dict[str, Any]] = { "architecture": "chatterbox_multilingual", "sample_rate": 24000, "size_gb": 2.6, - "files": ["t3_pt_br.safetensors", "s3gen_v3.pt", MTL_TOKENIZER_FILE], + "files": ["t3_pt_br.safetensors", "s3gen_v3.safetensors", MTL_TOKENIZER_FILE], "t3_checkpoint": "t3_pt_br.safetensors", "language": "pt", - "companion_files": {"ve.pt": CHATTERBOX_REPO_ID}, "description": "Single Language Pack: Brazilian Portuguese (pt). Needs a reference clip.", }, "chatterbox-multilingual-pt-pt": { @@ -160,10 +173,9 @@ MODEL_CONFIGS: Dict[str, Dict[str, Any]] = { "architecture": "chatterbox_multilingual", "sample_rate": 24000, "size_gb": 2.6, - "files": ["t3_pt_pt.safetensors", "s3gen_v3.pt", MTL_TOKENIZER_FILE], + "files": ["t3_pt_pt.safetensors", "s3gen_v3.safetensors", MTL_TOKENIZER_FILE], "t3_checkpoint": "t3_pt_pt.safetensors", "language": "pt", - "companion_files": {"ve.pt": CHATTERBOX_REPO_ID}, "description": "Single Language Pack: Portugal Portuguese (pt). Needs a reference clip.", }, "chatterbox-multilingual-hi": { @@ -171,10 +183,9 @@ MODEL_CONFIGS: Dict[str, Dict[str, Any]] = { "architecture": "chatterbox_multilingual", "sample_rate": 24000, "size_gb": 2.6, - "files": ["t3_hi.safetensors", "s3gen_v3.pt", MTL_TOKENIZER_FILE], + "files": ["t3_hi.safetensors", "s3gen_v3.safetensors", MTL_TOKENIZER_FILE], "t3_checkpoint": "t3_hi.safetensors", "language": "hi", - "companion_files": {"ve.pt": CHATTERBOX_REPO_ID}, "description": "Single Language Pack: Hindi (hi). Needs a reference clip.", }, # The Turbo family, in ONE folder. Turbo is the family's name upstream @@ -190,7 +201,7 @@ MODEL_CONFIGS: Dict[str, Dict[str, Any]] = { "architecture": "chatterbox_turbo", "sample_rate": 24000, "size_gb": 1.4, - "files": ["t3_turbo_v1.safetensors", "ve.safetensors", "s3gen.safetensors", "s3gen_meanflow.safetensors", "tokenizer.json", "conds.pt"], + "files": ["t3_turbo_v1.safetensors", "t3_turbo_v1.yaml", "ve.safetensors", "s3gen.safetensors", "s3gen_meanflow.safetensors", "tokenizer.json", "conds.pt", "vocab.json", "merges.txt", "tokenizer_config.json", "added_tokens.json", "special_tokens_map.json"], "description": "350M English model with paralinguistic tags; lower VRAM.", }, # A quarter of Turbo again; runs on CPU at realtime. Shares the Turbo @@ -203,12 +214,15 @@ MODEL_CONFIGS: Dict[str, Dict[str, Any]] = { "size_gb": 2.9, "files": [ "t3_nano_v1.safetensors", + "t3_nano_v1.yaml", "ve.safetensors", "s3gen.safetensors", "s3gen_meanflow.safetensors", "vocab.json", "merges.txt", "tokenizer_config.json", + "added_tokens.json", + "special_tokens_map.json", "conds.pt", ], "description": "Turbo-family 110M model for CPU and tight latency budgets. " @@ -227,6 +241,10 @@ MODEL_CONFIGS: Dict[str, Dict[str, Any]] = { "sample_rate": 24000, "size_gb": 3.1, "files": ["t3_flash.safetensors", "s3gen.safetensors", "ve.safetensors", "tokenizer.json"], + # Flash's s3gen.safetensors is a DIFFERENT decoder than the one the + # other families share despite the same filename -- it must never be + # resolved from, or linked into, the shared store. + "unique_files": ["s3gen.safetensors"], "description": "Block-diffusion Chatterbox: fast parallel decoding. Needs a " "reference audio clip.", }, @@ -260,7 +278,8 @@ CHATTERBOX_WEIGHT_FILES = ( "t3_flash.safetensors", "ve.safetensors", "s3gen.safetensors", - "s3gen_v3.pt", # Single Language Pack decoder (P-036) + "s3gen_v3.safetensors", # Single Language Pack decoder (P-036) + "s3gen_v3.pt", # legacy serialization, still loaded when present "s3gen.pt", ) diff --git a/src/PATCHES.md b/src/PATCHES.md index 7952d0a..92ecc9f 100644 --- a/src/PATCHES.md +++ b/src/PATCHES.md @@ -58,7 +58,7 @@ ascending order. | P-034 | `chatterbox/models/s3gen/s3gen.py` | `S3Token2Mel.forward` | compat patch | The ref-dict casting loop no longer dtype-casts integer tensors: `prompt_token` (token ids) and `prompt_token_len` (lengths) are integers that must move between devices but never between dtypes. Under a bf16 flow the old blanket cast rounded the combined token length 635 down to 632 (bf16's 7-bit mantissa), truncating the padding mask and crashing `flow.inference` with a shape mismatch — and quietly corrupted token ids above 256 that bf16 cannot represent exactly. In fp32 the cast was an exact no-op, so upstream never noticed; device movement is preserved for both branches. | | P-019 | `chatterbox/models/t3/t3.py` | `T3.inference` | ComfyUI-API patch | Accepts `pbar` and advances it once per sampled token. `max_new_tokens` keeps upstream's default of `None` but now falls back to `self.hp.max_speech_tokens` instead of raising, so the node may omit it. | | P-025 | `chatterbox/mtl_tts.py` | `_apply_watermark` | feature addition | Same optional-`perth` patch as P-023 for the multilingual entry point. Also applies the `str(device)` comparison of P-003 and P-007 in `ChatterboxMultilingualTTS.from_local`, which upstream still compares raw. | -| P-036 | `chatterbox/mtl_tts.py` | `_resolve_s3gen_checkpoint` | compat patch | The Single Language Pack finetunes (`ResembleAI/Chatterbox-Multilingual-zh-cmn` and its five siblings) ship their S3Gen decoder as `s3gen_v3.pt`, while the shared multilingual repo ships `s3gen.pt` — upstream's hard-coded `s3gen.pt` load made every SLP pack fail at the file check. `_resolve_s3gen_checkpoint` picks whichever the pack holds (preferring `s3gen.pt`), the completeness check tests that name, and `from_local` loads it. `REQUIRED_CHECKPOINT_FILES` drops `s3gen.pt` (now pack-dependent, resolved by the caller alongside the T3 filename), and `modules/loader.py`'s `MTL_CHECKPOINT_FILES` follows. The s3gen state-dict load is non-strict for the same reason: the SLP checkpoints predate the persistent `tokenizer._mel_filters`/`tokenizer.window` buffers (see `S3Token2Wav.ignore_state_dict_missing`), which are recomputed in `__init__` and never learned — a strict load refuses every SLP pack. The classic entry point already loads s3gen non-strict. | +| P-036 | `chatterbox/mtl_tts.py` | `_resolve_s3gen_checkpoint` | compat patch | The Single Language Pack finetunes (`ResembleAI/Chatterbox-Multilingual-zh-cmn` and its five siblings) ship their S3Gen decoder as `s3gen_v3.safetensors` (plus a legacy `.pt` twin), while the shared multilingual repo ships `s3gen.safetensors` — upstream's hard-coded `s3gen.pt` load made every SLP pack fail at the file check. `_resolve_s3gen_checkpoint` picks whichever serialization the pack holds (safetensors first), `_resolve_ve_checkpoint` does the same for the voice encoder, and `_load_checkpoint` dispatches `load_file`/`torch.load` by extension, so a pack may carry either form. `REQUIRED_CHECKPOINT_FILES` shrinks to the grapheme tokenizer; the resolved ve/s3gen/T3 names join the required set at call time. The s3gen state-dict load is non-strict for a related reason: the SLP checkpoints predate the persistent `tokenizer._mel_filters`/`tokenizer.window` buffers (see `S3Token2Wav.ignore_state_dict_missing`), which are recomputed in `__init__` and never learned — a strict load refuses every SLP pack. The classic entry point already loads s3gen non-strict. `from_pretrained`'s `allow_patterns` fetch the safetensors serializations only — the repo's `.pt` twins are the same weights and are never loaded. | | P-029 | `chatterbox/mtl_tts.py` | `_require_checkpoint_files` | compat patch | Same patch as P-027 for `ChatterboxMultilingualTTS.from_local`. Its set differs: `ve.pt`, `s3gen.pt` and `grapheme_mtl_merged_expanded_v1.json`, plus the T3 checkpoint for the selected version. That last name is not fixed, so `from_local` appends the already-resolved `t3_model` to the tuple rather than the helper hard-coding it. | | P-026 | `chatterbox/tts_turbo.py` | `_apply_watermark` | feature addition | Same optional-`perth` patch as P-023 for the turbo entry point. Also applies the `str(device)` comparison of P-003 and P-007 in `ChatterboxTurboTTS.from_local`, which upstream still compares raw, and makes `pyloudnorm` an optional import — it is used only by `norm_loudness`, which already degrades gracefully, so the module-level hard import was making the whole entry point unimportable without it. | | P-030 | `chatterbox/tts_turbo.py` | `_require_checkpoint_files` | compat patch | Same patch as P-027 for `ChatterboxTurboTTS.from_local`, covering Turbo and Nano together — there is no separate Nano module, so there is no separate `from_local`. The variant's T3 checkpoint (`t3_turbo_v1.safetensors` or `t3_nano_v1.safetensors`) is appended by the caller for the same reason as P-029. | diff --git a/src/chatterbox/mtl_tts.py b/src/chatterbox/mtl_tts.py index e0b8ade..aba143b 100644 --- a/src/chatterbox/mtl_tts.py +++ b/src/chatterbox/mtl_tts.py @@ -37,26 +37,48 @@ MULTILINGUAL_T3_MODELS = { } # PATCHED (src/PATCHES.md P-027): the file set a multilingual pack has to hold -# *besides* the version-selected T3 checkpoint, and the up-front check for it. -# The T3 filename is version-dependent and the s3gen filename is -# pack-dependent (P-036); both are resolved by the caller, so neither is in -# this tuple. +# *besides* the resolved T3, voice-encoder and s3gen checkpoints, and the +# up-front check for it. The other three filenames are pack-dependent (P-036) +# and resolved by the caller, so none of them is in this tuple. REQUIRED_CHECKPOINT_FILES = ( - "ve.pt", "grapheme_mtl_merged_expanded_v1.json", ) # PATCHED (src/PATCHES.md P-036): the Single Language Pack finetunes ship -# their S3Gen decoder as ``s3gen_v3.pt``; the shared multilingual repo ships -# ``s3gen.pt``. Resolve whichever the pack actually holds -- ``s3gen.pt`` -# stays the fallback so the missing-file message names the canonical file -# when neither exists. +# their S3Gen decoder as ``s3gen_v3.safetensors`` (plus a legacy ``.pt`` +# twin), while the shared multilingual repo ships ``s3gen.safetensors``. +# Resolve whichever serialization the pack actually holds, preferring +# safetensors (no pickle, what every other entry point loads); the canonical +# ``s3gen.safetensors`` stays the fallback so the missing-file message names +# it when nothing exists. def _resolve_s3gen_checkpoint(ckpt_dir): - for name in ("s3gen.pt", "s3gen_v3.pt"): + for name in ( + "s3gen_v3.safetensors", + "s3gen_v3.pt", + "s3gen.safetensors", + "s3gen.pt", + ): if (Path(ckpt_dir) / name).is_file(): return name - return "s3gen.pt" + return "s3gen.safetensors" + + +def _resolve_ve_checkpoint(ckpt_dir): + """The voice-encoder file the pack holds, safetensors preferred (P-036).""" + for name in ("ve.safetensors", "ve.pt"): + if (Path(ckpt_dir) / name).is_file(): + return name + return "ve.safetensors" + + +def _load_checkpoint(path, map_location=None): + """Read a state dict from either serialization, by extension (P-036).""" + if str(path).endswith(".safetensors"): + from safetensors.torch import load_file + + return load_file(str(path)) + return torch.load(path, map_location=map_location, weights_only=True) def _require_checkpoint_files(ckpt_dir, required): @@ -272,10 +294,15 @@ class ChatterboxMultilingualTTS: # PATCHED (src/PATCHES.md P-027): an incomplete directory is refused # here, by name, instead of at the first `torch.load`. - # PATCHED (src/PATCHES.md P-036): the s3gen filename follows the pack. + # PATCHED (src/PATCHES.md P-036): the s3gen and ve filenames follow + # the pack (safetensors preferred); the resolved names join the + # required set so the missing-file message names what is actually + # expected here. + ve_checkpoint = _resolve_ve_checkpoint(ckpt_dir) s3gen_checkpoint = _resolve_s3gen_checkpoint(ckpt_dir) _require_checkpoint_files( - ckpt_dir, (*REQUIRED_CHECKPOINT_FILES, s3gen_checkpoint, t3_model) + ckpt_dir, + (*REQUIRED_CHECKPOINT_FILES, ve_checkpoint, s3gen_checkpoint, t3_model), ) # Always load to CPU first for non-CUDA devices to handle CUDA-saved models @@ -289,8 +316,11 @@ class ChatterboxMultilingualTTS: map_location = None ve = VoiceEncoder() + # PATCHED (src/PATCHES.md P-036): resolve the voice-encoder file the + # pack actually holds (safetensors preferred) instead of hard-coding + # the legacy ``ve.pt`` name. ve.load_state_dict( - torch.load(ckpt_dir / "ve.pt", map_location=map_location, weights_only=True) + _load_checkpoint(ckpt_dir / _resolve_ve_checkpoint(ckpt_dir), map_location) ) ve.to(device).eval() @@ -308,7 +338,7 @@ class ChatterboxMultilingualTTS: # recomputed in __init__, never learned. Load non-strict like the # classic path in tts.py does; a strict load refuses every SLP pack. s3gen.load_state_dict( - torch.load(ckpt_dir / s3gen_checkpoint, map_location=map_location, weights_only=True), + _load_checkpoint(ckpt_dir / s3gen_checkpoint, map_location), strict=False, ) s3gen.to(device).eval() @@ -343,7 +373,10 @@ class ChatterboxMultilingualTTS: repo_id=REPO_ID, repo_type="model", revision="main", - allow_patterns=["ve.pt", t3_model, "s3gen.pt", "grapheme_mtl_merged_expanded_v1.json", "conds.pt", "Cangjie5_TC.json"], + # PATCHED (src/PATCHES.md P-036): fetch the safetensors + # serializations only -- the repo's ``.pt`` twins are the + # same weights and never loaded. + allow_patterns=["ve.safetensors", t3_model, "s3gen.safetensors", "grapheme_mtl_merged_expanded_v1.json", "conds.pt", "Cangjie5_TC.json"], token=os.getenv("HF_TOKEN"), ) ) diff --git a/tests/test_local_checkpoint.py b/tests/test_local_checkpoint.py index 513abb1..24d492f 100644 --- a/tests/test_local_checkpoint.py +++ b/tests/test_local_checkpoint.py @@ -366,12 +366,18 @@ class TestIncompleteDirectoryIsRefused: assert "_require_checkpoint_files(" in source, key def test_the_multilingual_set_is_the_pt_layout(self, real_vendor) -> None: - """The two fixed files; the s3gen + t3 names are resolved per pack (P-036).""" - assert MTL_CHECKPOINT_FILES == ( - "ve.pt", - "grapheme_mtl_merged_expanded_v1.json", + """The grapheme tokenizer is fixed; ve + s3gen + t3 resolve per pack (P-036). + + The resolution order prefers the safetensors serializations -- the + ``.pt`` twins hold the same weights and are never fetched. + """ + assert MTL_CHECKPOINT_FILES == ("grapheme_mtl_merged_expanded_v1.json",) + assert MTL_S3GEN_CHECKPOINTS == ( + "s3gen_v3.safetensors", + "s3gen_v3.pt", + "s3gen.safetensors", + "s3gen.pt", ) - assert MTL_S3GEN_CHECKPOINTS == ("s3gen.pt", "s3gen_v3.pt") assert real_vendor["mtl_tts"].REQUIRED_CHECKPOINT_FILES == MTL_CHECKPOINT_FILES def test_the_turbo_set_is_shared_by_turbo_and_nano(self, real_vendor) -> None: @@ -410,7 +416,7 @@ class TestIncompleteDirectoryIsRefused: ) message = str(caught.value) assert "t3_mtl23ls_v3.safetensors" in message - assert "ve.pt" in message and "s3gen.pt" in message + assert "ve.safetensors" in message and "s3gen_v3.safetensors" in message assert str(empty) in message def test_the_turbo_loader_refusal_names_the_missing_files(self, tmp_path, monkeypatch) -> None: diff --git a/tests/test_patcher.py b/tests/test_patcher.py index a6b93ad..d378a99 100644 --- a/tests/test_patcher.py +++ b/tests/test_patcher.py @@ -733,8 +733,13 @@ class TestLoaderSurface: empty.mkdir() attempts: list[tuple] = [] - def record_download_pack(model_name, repo_id, local_dir, client_id=None): + patterns_seen: list = [] + + def record_download_pack( + model_name, repo_id, local_dir, client_id=None, allow_patterns=None + ): attempts.append((model_name, repo_id, local_dir, client_id)) + patterns_seen.append(allow_patterns) monkeypatch.setattr( ChatterboxLoader, "download_pack", staticmethod(record_download_pack) @@ -752,6 +757,10 @@ class TestLoaderSurface: "load_model must attempt a download before giving up; the failure " "message is only truthful if it did" ) + # The classic fetch is filtered: shared weights come from the store. + assert patterns_seen == [ + ["t3_cfg.safetensors", "tokenizer.json", "conds.pt"] + ] def test_download_pack_is_skipped_for_a_complete_pack(self, pack_dir) -> None: called: list[tuple] = [] diff --git a/tests/test_shared_assets.py b/tests/test_shared_assets.py new file mode 100644 index 0000000..6f37c8a --- /dev/null +++ b/tests/test_shared_assets.py @@ -0,0 +1,200 @@ +"""The shared-weight store (2026-10-07 dedupe). + +``ve.safetensors``, ``s3gen.safetensors``, ``s3gen_v3.safetensors`` and the +grapheme tokenizer are byte-identical across every repo that ships them, so +they are fetched at most once into the reserved ``_shared`` directory and +hardlinked into each pack. Downloads carry only the pack-unique files. +""" + +from __future__ import annotations + +import os + +import pytest + +from ComfyUI_Chatterbox.modules import model_discovery +from ComfyUI_Chatterbox.modules.loader import ChatterboxLoader +from ComfyUI_Chatterbox.modules.model_info import ( + CHATTERBOX_WEIGHT_FILES, + MTL_TOKENIZER_FILE, + SHARED_ASSET_DIR_NAME, + SHARED_ASSET_FILES, +) + +SHARED_STORE = ("ve.safetensors", "s3gen.safetensors", "s3gen_v3.safetensors", MTL_TOKENIZER_FILE) + + +@pytest.fixture() +def models_root(tmp_path, monkeypatch): + """A private Chatterbox models root per test.""" + root = tmp_path / "models" / "tts" / "chatterbox" + monkeypatch.setattr(model_discovery, "get_models_root", lambda: str(root)) + return root + + +class TestTheScanSkipsTheSharedStore: + def test_the_reserved_store_is_never_a_pack(self, models_root) -> None: + store = models_root / SHARED_ASSET_DIR_NAME + pack = models_root / "some-pack" + store.mkdir(parents=True) + pack.mkdir(parents=True) + assert model_discovery._scan(str(models_root)) == [("some-pack", str(pack))] + + +class TestSharedAssetMaterialization: + def test_a_shared_file_is_linked_from_the_store(self, models_root) -> None: + store = models_root / SHARED_ASSET_DIR_NAME + store.mkdir(parents=True) + (store / "ve.safetensors").write_bytes(b"ve-weights") + pack = models_root / "fresh-pack" + pack.mkdir(parents=True) + + ChatterboxLoader._ensure_shared_assets(pack, ["ve.safetensors"]) + + linked = pack / "ve.safetensors" + assert linked.read_bytes() == b"ve-weights" + assert os.stat(store / "ve.safetensors").st_nlink == 2 + + def test_a_legacy_full_repo_folder_is_used_before_any_fetch( + self, models_root, monkeypatch + ) -> None: + legacy = models_root / "chatterbox" + legacy.mkdir(parents=True) + (legacy / "s3gen_v3.safetensors").write_bytes(b"v3-decoder") + pack = models_root / "slp" + pack.mkdir(parents=True) + + fetched: list = [] + + def refuse_fetch(**kwargs): + fetched.append(kwargs) + + monkeypatch.setattr("huggingface_hub.hf_hub_download", refuse_fetch) + + ChatterboxLoader._ensure_shared_assets(pack, ["s3gen_v3.safetensors"]) + + assert (pack / "s3gen_v3.safetensors").read_bytes() == b"v3-decoder" + assert not fetched, "a byte-identical legacy copy must not be re-fetched" + + def test_a_missing_shared_file_is_fetched_into_the_store( + self, models_root, monkeypatch + ) -> None: + pack = models_root / "slp" + pack.mkdir(parents=True) + + def fake_fetch(*, repo_id, filename, local_dir): + path = os.path.join(local_dir, filename) + with open(path, "wb") as handle: + handle.write(b"fresh") + return path + + monkeypatch.setattr("huggingface_hub.hf_hub_download", fake_fetch) + + ChatterboxLoader._ensure_shared_assets(pack, ["ve.safetensors"]) + + assert (pack / "ve.safetensors").read_bytes() == b"fresh" + assert ( + models_root / SHARED_ASSET_DIR_NAME / "ve.safetensors" + ).read_bytes() == b"fresh" + + def test_a_pack_that_already_holds_the_file_is_left_alone(self, models_root) -> None: + pack = models_root / "slp" + pack.mkdir(parents=True) + (pack / "ve.safetensors").write_bytes(b"pack-own") + + ChatterboxLoader._ensure_shared_assets(pack, ["ve.safetensors"]) + + assert (pack / "ve.safetensors").read_bytes() == b"pack-own" + + +class TestFilteredDownloads: + """The fetch carries pack-unique files only; shared ones never download.""" + + def _run(self, monkeypatch, models_root, model_name, required, entry): + from ComfyUI_Chatterbox.modules import model_info + + monkeypatch.setitem(model_info.MODEL_CONFIGS, model_name, entry) + seen: dict = {} + + def fake_download(model_name, repo_id, local_dir, client_id=None, allow_patterns=None): + seen["patterns"] = allow_patterns + seen["repo_id"] = repo_id + + monkeypatch.setattr(ChatterboxLoader, "download_pack", fake_download) + + def fake_fetch(*, repo_id, filename, local_dir): + path = os.path.join(local_dir, filename) + with open(path, "wb") as handle: + handle.write(b"shared") + return path + + monkeypatch.setattr("huggingface_hub.hf_hub_download", fake_fetch) + + ChatterboxLoader._ensure_entry_point_dir( + str(models_root / "pack-dir"), + entry["repo_id"], + required, + model_name, + ) + return seen + + def test_an_slp_download_carries_only_its_unique_files( + self, models_root, monkeypatch + ) -> None: + entry = { + "repo_id": "ResembleAI/Chatterbox-Multilingual-zh-cmn", + "files": [ + "t3_zh_cmn.safetensors", + "s3gen_v3.safetensors", + MTL_TOKENIZER_FILE, + "conds.pt", + ], + } + seen = self._run( + monkeypatch, + models_root, + "chatterbox-multilingual-zh-cmn", + ("grapheme_mtl_merged_expanded_v1.json", "ve.safetensors", + "s3gen_v3.safetensors", "t3_zh_cmn.safetensors"), + entry, + ) + assert seen["patterns"] == ["t3_zh_cmn.safetensors", "conds.pt"] + + def test_flash_keeps_its_own_s3gen_decoder(self, models_root, monkeypatch) -> None: + from ComfyUI_Chatterbox.modules import model_info + + entry = dict(model_info.MODEL_CONFIGS["chatterbox-flash"]) + seen = self._run( + monkeypatch, + models_root, + "chatterbox-flash", + ("ve.safetensors", "s3gen.safetensors", "t3_flash.safetensors", + "tokenizer.json"), + entry, + ) + assert "s3gen.safetensors" in seen["patterns"] + assert "ve.safetensors" not in seen["patterns"] + + def test_the_turbo_fetch_keeps_the_meanflow_decoder( + self, models_root, monkeypatch + ) -> None: + from ComfyUI_Chatterbox.modules import model_info + + entry = dict(model_info.MODEL_CONFIGS["chatterbox-turbo"]) + seen = self._run( + monkeypatch, + models_root, + "chatterbox-turbo", + ("ve.safetensors", "s3gen.safetensors", "s3gen_meanflow.safetensors", + "t3_turbo_v1.safetensors"), + entry, + ) + assert "s3gen_meanflow.safetensors" in seen["patterns"] + assert "ve.safetensors" not in seen["patterns"] + assert "s3gen.safetensors" not in seen["patterns"] + + +class TestTheSharedSetIsRealWeights: + def test_every_shared_asset_is_a_recognised_weight(self) -> None: + for name in SHARED_ASSET_FILES: + assert name in CHATTERBOX_WEIGHT_FILES or name == MTL_TOKENIZER_FILE diff --git a/tests/test_tts_routing.py b/tests/test_tts_routing.py index 1c1e878..21907e7 100644 --- a/tests/test_tts_routing.py +++ b/tests/test_tts_routing.py @@ -601,13 +601,13 @@ class TestSingleLanguagePackRouting: ChatterboxLoader._resolve_mtl_s3gen_file( str(tmp_path / "missing"), "chatterbox-multilingual-zh-cmn" ) - == "s3gen_v3.pt" + == "s3gen_v3.safetensors" ) assert ( ChatterboxLoader._resolve_mtl_s3gen_file( str(tmp_path / "missing"), "unknown-local-pack" ) - == "s3gen.pt" + == "s3gen_v3.safetensors" ) def test_a_slp_dir_sniffs_as_multilingual(self, tmp_path) -> None: