nodes/pentarab-ComfyUI-For-ChatterBox
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: use ComfyUI models folder for model downloads instead of HF cache

Browse Source 
- 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
This commit is contained in:
james-song
2025-12-28 19:04:09 +09:00
parent b625c94cfd
commit 735e28ed82
3 changed files with 113 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files
node.py
+7 -2
View File
@@ -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]
pyproject.toml
+1 -1
View File
@@ -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"
src/chatterbox/models/tokenizers/tokenizer.py
+105 -67
View File
@@ -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