From 735e28ed82bd65097cffbd0d501e4d48e3f6bbea Mon Sep 17 00:00:00 2001 From: james-song Date: Sun, 28 Dec 2025 19:04:09 +0900 Subject: [PATCH] fix: use ComfyUI models folder for model downloads instead of HF cache - Change node.py to use chatterbox_handler's download function - Models now download to models/tts/chatterbox/ directory - Remove hf_hub_download from tokenizer, use local file path instead - Bump version to 1.0.2 --- node.py | 9 +- pyproject.toml | 2 +- src/chatterbox/models/tokenizers/tokenizer.py | 172 +++++++++++------- 3 files changed, 113 insertions(+), 70 deletions(-) diff --git a/node.py b/node.py index ebbf0c6..1d75d99 100644 --- a/node.py +++ b/node.py @@ -61,10 +61,15 @@ def get_or_load_model(device=None): cache_key = f"chatterbox_mtl_{device}" if cache_key not in _MODEL_CACHE: - from .src.chatterbox.mtl_tts import ChatterboxMultilingualTTS + from .chatterbox_handler import ( + load_chatterbox_multilingual_tts_model, + DEFAULT_MTL_MODEL_PACK_NAME, + ) print(f"[Chatterbox] Loading multilingual model on {device}...") - _MODEL_CACHE[cache_key] = ChatterboxMultilingualTTS.from_pretrained(device) + _MODEL_CACHE[cache_key] = load_chatterbox_multilingual_tts_model( + DEFAULT_MTL_MODEL_PACK_NAME, device + ) print("[Chatterbox] Model loaded successfully.") return _MODEL_CACHE[cache_key] diff --git a/pyproject.toml b/pyproject.toml index 35550f4..06b701a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "comfyui-for-chatterbox" -version = "1.0.1" +version = "1.0.2" description = "ComfyUI custom nodes for Chatterbox TTS with multilingual support (23 languages)" readme = "README.md" requires-python = ">=3.10" diff --git a/src/chatterbox/models/tokenizers/tokenizer.py b/src/chatterbox/models/tokenizers/tokenizer.py index 84d45d3..f9505b0 100644 --- a/src/chatterbox/models/tokenizers/tokenizer.py +++ b/src/chatterbox/models/tokenizers/tokenizer.py @@ -5,7 +5,6 @@ import torch from pathlib import Path from unicodedata import category, normalize from tokenizers import Tokenizer -from huggingface_hub import hf_hub_download # Special tokens @@ -17,6 +16,7 @@ SPECIAL_TOKENS = [SOT, EOT, UNK, SPACE, "[PAD]", "[SEP]", "[CLS]", "[MASK]"] logger = logging.getLogger(__name__) + class EnTokenizer: def __init__(self, vocab_file_path): self.tokenizer: Tokenizer = Tokenizer.from_file(vocab_file_path) @@ -36,7 +36,7 @@ class EnTokenizer: """ clean_text > (append `lang_id`) > replace SPACE > encode text using Tokenizer """ - txt = txt.replace(' ', SPACE) + txt = txt.replace(" ", SPACE) code = self.tokenizer.encode(txt) ids = code.ids return ids @@ -46,16 +46,13 @@ class EnTokenizer: seq = seq.cpu().numpy() txt: str = self.tokenizer.decode(seq, skip_special_tokens=False) - txt = txt.replace(' ', '') - txt = txt.replace(SPACE, ' ') - txt = txt.replace(EOT, '') - txt = txt.replace(UNK, '') + txt = txt.replace(" ", "") + txt = txt.replace(SPACE, " ") + txt = txt.replace(EOT, "") + txt = txt.replace(UNK, "") return txt -# Model repository -REPO_ID = "ResembleAI/chatterbox" - # Global instances for optional dependencies _kakasi = None _dicta = None @@ -75,17 +72,18 @@ def is_katakana(c: str) -> bool: def hiragana_normalize(text: str) -> str: """Japanese text normalization: converts kanji to hiragana; katakana remains the same.""" global _kakasi - + try: if _kakasi is None: import pykakasi + _kakasi = pykakasi.kakasi() - + result = _kakasi.convert(text) out = [] - + for r in result: - inp = r['orig'] + inp = r["orig"] hira = r["hira"] # Any kanji in the phrase @@ -95,20 +93,23 @@ def hiragana_normalize(text: str) -> str: out.append(hira) # All katakana - elif all([is_katakana(c) for c in inp]) if inp else False: # Safety check for empty inp - out.append(r['orig']) + elif ( + all([is_katakana(c) for c in inp]) if inp else False + ): # Safety check for empty inp + out.append(r["orig"]) else: out.append(inp) - + normalized_text = "".join(out) - + # Decompose Japanese characters for tokenizer compatibility import unicodedata - normalized_text = unicodedata.normalize('NFKD', normalized_text) - + + normalized_text = unicodedata.normalize("NFKD", normalized_text) + return normalized_text - + except ImportError: logger.warning("pykakasi not available - Japanese text processing skipped") return text @@ -117,14 +118,15 @@ def hiragana_normalize(text: str) -> str: def add_hebrew_diacritics(text: str) -> str: """Hebrew text normalization: adds diacritics to Hebrew text.""" global _dicta - + try: if _dicta is None: from dicta_onnx import Dicta + _dicta = Dicta() - + return _dicta.add_diacritics(text) - + except ImportError: logger.warning("dicta_onnx not available - Hebrew text processing skipped") return text @@ -135,47 +137,51 @@ def add_hebrew_diacritics(text: str) -> str: def korean_normalize(text: str) -> str: """Korean text normalization: decompose syllables into Jamo for tokenization.""" - + def decompose_hangul(char): """Decompose Korean syllable into Jamo components.""" - if not ('\uac00' <= char <= '\ud7af'): + if not ("\uac00" <= char <= "\ud7af"): return char - + # Hangul decomposition formula base = ord(char) - 0xAC00 initial = chr(0x1100 + base // (21 * 28)) medial = chr(0x1161 + (base % (21 * 28)) // 28) - final = chr(0x11A7 + base % 28) if base % 28 > 0 else '' - + final = chr(0x11A7 + base % 28) if base % 28 > 0 else "" + return initial + medial + final - + # Decompose syllables and normalize punctuation - result = ''.join(decompose_hangul(char) for char in text) + result = "".join(decompose_hangul(char) for char in text) return result.strip() class ChineseCangjieConverter: """Converts Chinese characters to Cangjie codes for tokenization.""" - + def __init__(self, model_dir=None): self.word2cj = {} self.cj2word = {} self.segmenter = None self._load_cangjie_mapping(model_dir) self._init_segmenter() - + def _load_cangjie_mapping(self, model_dir=None): - """Load Cangjie mapping from HuggingFace model repository.""" + """Load Cangjie mapping from local model directory.""" try: - cangjie_file = hf_hub_download( - repo_id=REPO_ID, - filename="Cangjie5_TC.json", - cache_dir=model_dir - ) - + if model_dir is None: + logger.warning("No model directory provided for Cangjie mapping") + return + + cangjie_file = Path(model_dir) / "Cangjie5_TC.json" + + if not cangjie_file.exists(): + logger.warning(f"Cangjie mapping file not found: {cangjie_file}") + return + with open(cangjie_file, "r", encoding="utf-8") as fp: data = json.load(fp) - + for entry in data: word, code = entry.split("\t")[:2] self.word2cj[word] = code @@ -183,19 +189,22 @@ class ChineseCangjieConverter: self.cj2word[code] = [word] else: self.cj2word[code].append(word) - + except Exception as e: logger.warning(f"Could not load Cangjie mapping: {e}") - + def _init_segmenter(self): """Initialize pkuseg segmenter.""" try: from spacy_pkuseg import pkuseg + self.segmenter = pkuseg() except ImportError: - logger.warning("pkuseg not available - Chinese segmentation will be skipped") + logger.warning( + "pkuseg not available - Chinese segmentation will be skipped" + ) self.segmenter = None - + def _cangjie_encode(self, glyph: str): """Encode a single Chinese glyph to Cangjie code.""" normed_glyph = glyph @@ -205,9 +214,7 @@ class ChineseCangjieConverter: index = self.cj2word[code].index(normed_glyph) index = str(index) if index > 0 else "" return code + str(index) - - def __call__(self, text): """Convert Chinese characters in text to Cangjie tokens.""" output = [] @@ -216,7 +223,7 @@ class ChineseCangjieConverter: full_text = " ".join(segmented_words) else: full_text = text - + for t in full_text: if category(t) == "Lo": cangjie = self._cangjie_encode(t) @@ -237,16 +244,19 @@ class ChineseCangjieConverter: def add_russian_stress(text: str) -> str: """Russian text normalization: adds stress marks to Russian text.""" global _russian_stresser - + try: if _russian_stresser is None: from russian_text_stresser.text_stresser import RussianTextStresser + _russian_stresser = RussianTextStresser() - + return _russian_stresser.stress_text(text) - + except ImportError: - logger.warning("russian_text_stresser not available - Russian stress labeling skipped") + logger.warning( + "russian_text_stresser not available - Russian stress labeling skipped" + ) return text except Exception as e: logger.warning(f"Russian stress labeling failed: {e}") @@ -265,7 +275,13 @@ class MTLTokenizer: assert SOT in voc assert EOT in voc - def preprocess_text(self, raw_text: str, language_id: str = None, lowercase: bool = True, nfkd_normalize: bool = True): + def preprocess_text( + self, + raw_text: str, + language_id: str = None, + lowercase: bool = True, + nfkd_normalize: bool = True, + ): """ Text preprocessor that handles lowercase conversion and NFKD normalization. """ @@ -274,34 +290,56 @@ class MTLTokenizer: preprocessed_text = preprocessed_text.lower() if nfkd_normalize: preprocessed_text = normalize("NFKD", preprocessed_text) - + return preprocessed_text - def text_to_tokens(self, text: str, language_id: str = None, lowercase: bool = True, nfkd_normalize: bool = True): - text_tokens = self.encode(text, language_id=language_id, lowercase=lowercase, nfkd_normalize=nfkd_normalize) + def text_to_tokens( + self, + text: str, + language_id: str = None, + lowercase: bool = True, + nfkd_normalize: bool = True, + ): + text_tokens = self.encode( + text, + language_id=language_id, + lowercase=lowercase, + nfkd_normalize=nfkd_normalize, + ) text_tokens = torch.IntTensor(text_tokens).unsqueeze(0) return text_tokens - def encode(self, txt: str, language_id: str = None, lowercase: bool = True, nfkd_normalize: bool = True): - txt = self.preprocess_text(txt, language_id=language_id, lowercase=lowercase, nfkd_normalize=nfkd_normalize) - + def encode( + self, + txt: str, + language_id: str = None, + lowercase: bool = True, + nfkd_normalize: bool = True, + ): + txt = self.preprocess_text( + txt, + language_id=language_id, + lowercase=lowercase, + nfkd_normalize=nfkd_normalize, + ) + # Language-specific text processing - if language_id == 'zh': + if language_id == "zh": txt = self.cangjie_converter(txt) - elif language_id == 'ja': + elif language_id == "ja": txt = hiragana_normalize(txt) - elif language_id == 'he': + elif language_id == "he": txt = add_hebrew_diacritics(txt) - elif language_id == 'ko': + elif language_id == "ko": txt = korean_normalize(txt) - elif language_id == 'ru': + elif language_id == "ru": txt = add_russian_stress(txt) - + # Prepend language token if language_id: txt = f"[{language_id.lower()}]{txt}" - - txt = txt.replace(' ', SPACE) + + txt = txt.replace(" ", SPACE) return self.tokenizer.encode(txt).ids def decode(self, seq): @@ -309,5 +347,5 @@ class MTLTokenizer: seq = seq.cpu().numpy() txt = self.tokenizer.decode(seq, skip_special_tokens=False) - txt = txt.replace(' ', '').replace(SPACE, ' ').replace(EOT, '').replace(UNK, '') + txt = txt.replace(" ", "").replace(SPACE, " ").replace(EOT, "").replace(UNK, "") return txt