diff --git a/README.md b/README.md index f645071..8bf025b 100644 --- a/README.md +++ b/README.md @@ -144,6 +144,7 @@ list in ComfyUI's settings to turn the toasts off. | **Offline** | Packs are fetched by a first-run download and a local pack is never re-downloaded afterwards. | | **VRAM** | Weights load under ComfyUI's manager and offload between runs; ~2.5 GB per pack on disk as safetensors. | | **Layout** | Node code lives in `nodes/`, the vendored Chatterbox core in `src/chatterbox/`. | +| **Emotion tags & chunking** | Turbo and Nano packs speak the paralinguistic tags natively: write `[laugh]`, `[chuckle]`, `[sigh]`, `[clear throat]` and the rest into the text box and the model says them. Classic, multilingual and Flash packs have no such vocabulary, so the same words are read out as plain text and the run says so in the log. ``-style v2 tokens are re-spelled for Turbo and left alone elsewhere. Long scripts are cut into chunks, spoken one chunk at a time and joined again -- `auto` picks the join by script length (clean under 500 characters, crossfaded above it, padded with silence above 1000), and the cut never lands inside a tag. | | **Upgrading from 1.x** | 2.0.0 moved packs from `models/chatterbox_tts/` to `models/tts/chatterbox/`. No compatibility shim is provided, so delete the old directory by hand. Widget order did not change and node ids did not change, so saved workflows still load. | | **Changelog** | [CHANGELOG.md](./CHANGELOG.md) | diff --git a/example_workflows/ChatterboxTTS-workflow.json b/example_workflows/ChatterboxTTS-workflow.json index 5fbe5c2..65a033a 100644 --- a/example_workflows/ChatterboxTTS-workflow.json +++ b/example_workflows/ChatterboxTTS-workflow.json @@ -313,7 +313,11 @@ 1000, true, 10, - 1 + 1, + true, + 400, + "auto", + 100 ], "widgets_values_named": { "model_pack_name": "chatterbox", @@ -335,7 +339,11 @@ "top_k": 1000, "normalize_loudness": true, "num_steps": 10, - "cfg_scale": 1 + "cfg_scale": 1, + "enable_chunking": true, + "max_chars_per_chunk": 400, + "chunk_combination_method": "auto", + "silence_between_chunks_ms": 100 }, "color": "#232", "bgcolor": "#353" diff --git a/modules/generation.py b/modules/generation.py index 8852a93..6abf66e 100644 --- a/modules/generation.py +++ b/modules/generation.py @@ -486,70 +486,86 @@ def force_offload_model( # --------------------------------------------------------------------------- -def generate_audio( - model: Any, - text: str, - audio_prompt: Optional[dict] = None, - exaggeration: float = 0.5, - temperature: float = 0.8, - cfg_weight: float = 0.5, - repetition_penalty: float = 1.2, - min_p: float = 0.05, - top_p: float = 1.0, - seed: int = 0, - max_new_tokens: int = 1000, - flow_cfg_scale: float = 0.7, - use_watermark: bool = False, - top_k: int = 1000, - norm_loudness: bool = True, - num_steps: int = 10, - cfg_scale: float = 1.0, -) -> tuple[torch.Tensor, int]: - """Synthesise ``text`` with a loaded Chatterbox TTS model. +def _tts_family(model: Any) -> str: + """Which TTS family ``model`` belongs to: ``"flash"``, ``"turbo"`` or + ``"classic"``. - Turbo-family packs (``ChatterboxTurboTTS`` -- Turbo and Nano) are - dispatched to :func:`generate_turbo_audio`: their sampler is a top-k / - top-p pair with no CFG, so ``cfg_weight`` and ``min_p`` are ignored for - them, and ``top_k``/``norm_loudness`` only mean anything there. The - dispatch is by model class name -- the vendored package is never imported - by this module. - - Args: - model: The ``ChatterboxTTS`` instance from :func:`load_chatterbox_model`. - text: The text to speak. Must be non-blank. - audio_prompt: Optional ComfyUI ``AUDIO`` dict to clone the voice from. - Downmixed, scrubbed and resampled to 24 kHz in memory, then handed - to the model as a waveform. - exaggeration: Expressiveness. - temperature: Sampling temperature. - cfg_weight: Classifier-free guidance weight. - repetition_penalty: Token repetition penalty. - min_p: Nucleus-sampling floor. - top_p: Nucleus-sampling ceiling. - seed: Seed; ``0`` means random. - max_new_tokens: Token budget. 25 tokens is roughly one second. - flow_cfg_scale: Flow-matching CFG scale for the mel decoder. - use_watermark: Whether to apply the PERTH watermark. Off by default - because ``perth`` is an optional dependency; the vendored models - take ``watermark=False`` rather than needing a stub object. - - Returns: - ``(waveform[1, T] float32 on CPU, sample_rate)``. - - Raises: - ValueError: ``text`` is blank, or the model returned no audio. + The dispatch key is the class name, exactly as it has always been: the + vendored package is never imported by this module, and the family decides + the sampler, the chunk semantics and whether the square-bracket tags are + the model's own vocabulary. """ - if not text or not text.strip(): - raise ValueError( - "Chatterbox needs some text to speak. The text input was empty." - ) + name = type(model).__name__ + if name == "ChatterboxFlashTTS": + return "flash" + if name == "ChatterboxTurboTTS": + return "turbo" + return "classic" - # Flash dispatch, before the classic progress bar: a block-diffusion - # decoder with no token budget to count. - if type(model).__name__ == "ChatterboxFlashTTS": + +def _chunk_seed(seed: int, index: int) -> int: + """The seed for chunk ``index`` of a run. + + ``0`` is the widget's "random" convention and stays ``0``, so every chunk + is sampled from fresh noise. Any other seed is re-derived per chunk + (``seed + index``): one noise sample replayed for every chunk gives every + chunk the same prosody -- the audible artefact chunking is supposed to + avoid -- while the run as a whole stays reproducible, because the same + script and the same seed still produce the same chunk waveforms. + """ + if not seed: + return 0 + return (int(seed) + int(index)) % (2 ** 64) + + +def _log_plan_notices(notices: tuple[str, ...]) -> None: + """Say once per run what the engine could not honour. + + A notice is never an error and never changes the audio: it is the engine + admitting that a tag in the script is not going to do what it looks like + it does -- a Turbo tag on a pack without the Turbo vocabulary, a pause + that is parsed but never rendered, a language that is parsed but never + switched. Dropping them on the floor is what made the feature look + broken, so each one is logged as a warning, once per run. + """ + for notice in notices: + logging.warning("[Chatterbox TTS] %s", notice) + + +def _generate_one_chunk( + model: Any, + chunk: str, + *, + family: str, + audio_prompt: Optional[dict], + exaggeration: float, + temperature: float, + cfg_weight: float, + repetition_penalty: float, + min_p: float, + top_p: float, + seed: int, + max_new_tokens: int, + flow_cfg_scale: float, + use_watermark: bool, + top_k: int, + norm_loudness: bool, + num_steps: int, + cfg_scale: float, +) -> tuple[torch.Tensor, int]: + """Synthesise one chunk with the sampler its family uses. + + The family dispatch that used to sit at the top of :func:`generate_audio` + lives here now, so the chunk loop above it can call the same code once + per chunk. A Flash pack is a block-diffusion decoder with no token budget + to count; a Turbo pack has no ``max_new_tokens`` either, takes the top-k + sampler and keeps its square-bracket tags in the text for the tokenizer + to speak; everything else is the classic sampler. + """ + if family == "flash": return generate_flash_audio( model, - text, + chunk, audio_prompt=audio_prompt, exaggeration=exaggeration, temperature=temperature, @@ -558,19 +574,19 @@ def generate_audio( cfg_scale=cfg_scale, ) - # Turbo-family dispatch, before the classic progress bar: Turbo has no - # max_new_tokens to count, and its square-bracket tags need resolving. - if type(model).__name__ == "ChatterboxTurboTTS": - plan = plan_text(text, turbo_tags=True) + if family == "turbo": return generate_turbo_audio( model, - plan.text, + chunk, audio_prompt=audio_prompt, temperature=temperature, top_k=top_k, top_p=top_p, repetition_penalty=repetition_penalty, - exaggeration=exaggeration, + # Turbo ignores exaggeration for sampling -- ``hp.emotion_adv`` is + # False on this family (``src/chatterbox/tts_turbo.py:259``) -- and + # upstream warns about every non-zero value, once per chunk. + exaggeration=0.0, norm_loudness=norm_loudness, seed=seed, use_watermark=use_watermark, @@ -592,7 +608,7 @@ def generate_audio( try: with torch.no_grad(): waveform = model.generate( - text, + chunk, audio_prompt_path=prompt, exaggeration=exaggeration, temperature=temperature, @@ -617,7 +633,7 @@ def generate_audio( except Exception: logging.exception( "[Chatterbox TTS] Generation failed for a %d-character text.", - len(text), + len(chunk), ) raise finally: @@ -632,13 +648,167 @@ def generate_audio( result = waveform.detach().to(device="cpu", dtype=torch.float32) if result.dim() == 1: result = result.unsqueeze(0) + return result, S3GEN_SAMPLE_RATE + + +def generate_audio( + model: Any, + text: str, + audio_prompt: Optional[dict] = None, + exaggeration: float = 0.5, + temperature: float = 0.8, + cfg_weight: float = 0.5, + repetition_penalty: float = 1.2, + min_p: float = 0.05, + top_p: float = 1.0, + seed: int = 0, + max_new_tokens: int = 1000, + flow_cfg_scale: float = 0.7, + use_watermark: bool = False, + top_k: int = 1000, + norm_loudness: bool = True, + num_steps: int = 10, + cfg_scale: float = 1.0, + *, + enable_chunking: bool = True, + max_chars_per_chunk: int = 400, + chunk_combination_method: str = "auto", + silence_between_chunks_ms: int = 100, + crash_protection_template: str = "", +) -> tuple[torch.Tensor, int]: + """Synthesise ``text`` with a loaded Chatterbox TTS model. + + Turbo-family packs (``ChatterboxTurboTTS`` -- Turbo and Nano) are + dispatched to :func:`generate_turbo_audio`: their sampler is a top-k / + top-p pair with no CFG, so ``cfg_weight`` and ``min_p`` are ignored for + them, and ``top_k``/``norm_loudness`` only mean anything there. The + dispatch is by model class name -- the vendored package is never imported + by this module. + + **Every family goes through the same text engine and the same chunk + loop.** The script is parsed once by :func:`modules.text_processing.plan_text` + -- which is what resolves ``[Alice]``, keeps Turbo's ``[laugh]`` tags for + the tokenizer and cuts the script into chunks -- and each chunk is then + synthesised with its own model call and joined with + :func:`combine_audio_chunks`. The chunks are what make a long script + speakable: one call per chunk means one token budget per chunk, so a + script of any length is spoken in full instead of running out of tokens + half way through. + + Args: + model: The ``ChatterboxTTS`` instance from :func:`load_chatterbox_model`. + text: The text to speak. Must be non-blank. + audio_prompt: Optional ComfyUI ``AUDIO`` dict to clone the voice from. + Downmixed, scrubbed and resampled to 24 kHz in memory, then handed + to the model as a waveform. + exaggeration: Expressiveness. + temperature: Sampling temperature. + cfg_weight: Classifier-free guidance weight. + repetition_penalty: Token repetition penalty. + min_p: Nucleus-sampling floor. + top_p: Nucleus-sampling ceiling. + seed: Seed; ``0`` means random. A non-zero seed is re-derived per + chunk (``seed + chunk_index``), so the run stays reproducible + while the chunks do not all share one noise sample. + max_new_tokens: Token budget **per chunk**. 25 tokens is roughly one + second. + flow_cfg_scale: Flow-matching CFG scale for the mel decoder. + use_watermark: Whether to apply the PERTH watermark. Off by default + because ``perth`` is an optional dependency; the vendored models + take ``watermark=False`` rather than needing a stub object. + enable_chunking: Whether to cut the script into chunks at all. Off + means one model call for the whole script. + max_chars_per_chunk: Soft ceiling on one chunk, in characters. + chunk_combination_method: ``auto`` / ``concatenate`` / + ``silence_padding`` / ``crossfade``. ``auto`` resolves by the + length of the whole script. + silence_between_chunks_ms: The gap between chunks when the method is + ``silence_padding``. + crash_protection_template: Padding template for short segments. + Empty by default: the TTS node has no such widget, and turning it + on here would change every existing user's audio. + + Returns: + ``(waveform[1, T] float32 on CPU, sample_rate)``. + + Raises: + ValueError: ``text`` is blank, the script held nothing but tags, or + the model returned no audio. + """ + if not text or not text.strip(): + raise ValueError( + "Chatterbox needs some text to speak. The text input was empty." + ) + + family = _tts_family(model) + + # One plan for every family. On a Turbo pack the square-bracket tags are + # kept verbatim for the tokenizer; on every other pack they stay literal + # text and are reported as a notice, so nothing is swallowed silently. + plan = plan_text( + text, + turbo_tags=(family == "turbo"), + crash_protection_template=crash_protection_template, + enable_chunking=enable_chunking, + max_chars_per_chunk=max_chars_per_chunk, + chunk_combination_method=chunk_combination_method, + silence_between_chunks_ms=silence_between_chunks_ms, + ) + _log_plan_notices(plan.notices) + + if not plan.text.strip(): + raise ValueError( + "Chatterbox needs some text to speak. The script held nothing but " + "tags, so there is nothing left to say." + ) + + chunks = plan.chunks or (plan.text,) + logging.info( + "[Chatterbox TTS] Synthesising %d chunk(s) of %d characters, joined " + "with '%s'.", + len(chunks), + plan.char_count, + plan.chunk_method, + ) + + waveforms: list[torch.Tensor] = [] + sample_rate = S3GEN_SAMPLE_RATE + for index, chunk in enumerate(chunks): + waveform, sample_rate = _generate_one_chunk( + model, + chunk, + family=family, + audio_prompt=audio_prompt, + exaggeration=exaggeration, + temperature=temperature, + cfg_weight=cfg_weight, + repetition_penalty=repetition_penalty, + min_p=min_p, + top_p=top_p, + seed=_chunk_seed(seed, index), + max_new_tokens=max_new_tokens, + flow_cfg_scale=flow_cfg_scale, + use_watermark=use_watermark, + top_k=top_k, + norm_loudness=norm_loudness, + num_steps=num_steps, + cfg_scale=cfg_scale, + ) + waveforms.append(waveform) + + joined = combine_audio_chunks( + waveforms, + sample_rate, + method=plan.chunk_method, + silence_between_chunks_ms=silence_between_chunks_ms, + ) logging.info( "[Chatterbox TTS] Generated %.2f s of audio at %d Hz.", - result.shape[-1] / float(S3GEN_SAMPLE_RATE), - S3GEN_SAMPLE_RATE, + joined.shape[-1] / float(sample_rate), + sample_rate, ) - return result, S3GEN_SAMPLE_RATE + return joined, sample_rate def convert_voice_audio( diff --git a/modules/text_processing.py b/modules/text_processing.py index 47e5ff5..dc08f26 100644 --- a/modules/text_processing.py +++ b/modules/text_processing.py @@ -13,7 +13,8 @@ corpus specifies: 2. language tags -- ``[de:Alice] Hallo.`` 3. pause tags -- ``[pause:1.5]`` / ``[wait:2s]`` / ``[stop:500ms]`` 4. per-segment parameter overrides -- ``[Alice|seed:42|temp:0.5]`` -5. the v2 angle-bracket tokens -- ```` +5. the v2 angle-bracket tokens -- ```` (re-spelled ``[giggle]`` on a + Turbo pack, whose tokenizer knows the square spelling and not the angle one) Stages 1-4 classify one ``[...]`` token, so they are tried **in that order** against each token in turn: a token that is not a character is then tested as a @@ -22,14 +23,22 @@ because a ```` is only rewritten once the character names around it have resolved -- rewriting it first is what makes ``[Alice] `` lose its character. +The nine Turbo tag names (``[laugh]``, ``[clear throat]``, ...) are claimed +**before** stage 1, on every path. A Turbo pack needs them kept verbatim for +the tokenizer, and a pack that has no Turbo vocabulary still has to say the +word rather than swallow it as a voice name -- so the tag is never a +character, anywhere. + Three rules that follow from the syntax and are worth stating up front: * **Square brackets are the Turbo tag set's home, and angle brackets are the v2 - token set's.** They are deliberately not interchangeable. ``[giggle]`` is a - *character name* -- an upstream footgun this engine reproduces faithfully and - documents rather than silently second-guessing -- while ```` is a v2 - special token. Turbo tags (``[laugh]``, ``[clear throat]``, ...) are only - recognised when the caller says it is speaking to a Turbo model. + token set's.** They are deliberately not interchangeable. ```` is a + v2 special token; ``[laugh]`` is a Turbo paralinguistic tag, and it is only + *recognised* as one when the caller says it is speaking to a Turbo model. + On any other pack the nine Turbo names are kept as **literal text** and + reported as a notice -- swallowing ``[laugh]`` as a voice name would take the + word out of the speech without saying so. Every other bare bracket stays a + *character name*, which is what upstream does. * **Unknown parameters are ignored, not fatal.** A script carried over from another tool is full of knobs this node has no socket for; refusing the whole run over ``[Alice|batch:4]`` would be hostile. @@ -55,6 +64,7 @@ __all__ = [ "AUTO_CHUNK_THRESHOLDS", "PARAM_ALIASES", "TURBO_TAGS", + "TURBO_TAG_V2_SPELLING", "V2_SPECIAL_TOKENS", "UPPERCASE_V2_TOKENS", "TextSegment", @@ -140,6 +150,14 @@ TURBO_TAGS: tuple[str, ...] = ( "clear throat", ) +#: The v2 tokens whose Turbo spelling differs from their own name. Only one +#: does: the v2 set writes ```` with an underscore, Turbo's +#: tokenizer was trained on ``[clear throat]`` with a space. Every other token +#: keeps its own spelling on both sides of the translation. +TURBO_TAG_V2_SPELLING: dict[str, str] = { + "clear_throat": "clear throat", +} + #: The 30 named v2 special tokens. The tokenizer upstream actually contains 118 #: special tokens; these are the 30 the reference documentation names, and they #: are undocumented upstream, experimental, and tracked in upstream issue #186. @@ -193,6 +211,15 @@ _TURBO_TAG_LOOKUP = frozenset(TURBO_TAGS) # because they are never tags -- only v2 special tokens. _BRACKET_RE = re.compile(r"\[([^\[\]]*)\]") _ANGLE_RE = re.compile(r"<([^<>]*)>") +#: A sentence end: the mark that ends a sentence and the space that proves it. +_SENTENCE_END_RE = re.compile(r"[.!?] ") +#: The bracket pairs a chunk boundary must never land inside. +_BRACKET_PAIRS = (("[", "]"), ("<", ">")) +#: A Turbo tag as it appears in *text*: the square spelling, in any case. +_TURBO_TAG_IN_TEXT_RE = re.compile( + r"\[\s*(" + "|".join(re.escape(tag) for tag in TURBO_TAGS) + r")\s*\]", + re.IGNORECASE, +) _PAUSE_RE = re.compile(r"^(pause|wait|stop)\s*:\s*(?P.+)$", re.IGNORECASE) _OVERRIDE_SPLIT_RE = re.compile(r"\|") _DURATION_RE = re.compile(r"^\s*(?P\d+(?:\.\d+)?|\.\d+)\s*(?P[a-z]*)\s*$", re.IGNORECASE) @@ -230,7 +257,12 @@ class TextSegment: pause_after: Seconds of silence to insert **after** this segment, from the ``[pause:…]`` / ``[wait:…]`` / ``[stop:…]`` tag that follows it. A pause with nothing before it in the script is dropped. - turbo_tags: Turbo paralinguistic tags found inside this segment. + turbo_tags: The Turbo paralinguistic tags found inside this segment. + Recorded on a Turbo pack, where they are kept verbatim in ``text`` + and handed to the model, which is where they have an effect. On any + other pack a Turbo tag is kept as literal text -- that pack has no + such vocabulary -- and the plan reports a notice about it instead + of recording it here. index: Position in the original script, preserved through chunking so a caller can still map a chunk back to where it came from. """ @@ -251,13 +283,29 @@ class TextSegment: @dataclass class TextPlan: - """A parsed script: its segments and how to synthesise them.""" + """A parsed script: its segments and how to synthesise them. + + Attributes: + segments: The script, split and resolved. + text: Everything that will be spoken, as one string. + chunks: ``text`` cut into the pieces that are synthesised separately. + chunk_method: The concrete join method -- never ``"auto"``. + char_count: Length of ``text``. + notices: Things the engine could not honour, said out loud rather + than dropped. A notice is never an error and never changes the + audio: it is the engine admitting that a tag in the script is not + going to do what it looks like it does -- a Turbo tag on a pack + without the Turbo vocabulary, a pause that is parsed but never + rendered as silence, a language that is parsed but never switched. + The caller decides whether to log them. + """ segments: tuple[TextSegment, ...] text: str chunks: tuple[str, ...] chunk_method: str char_count: int + notices: tuple[str, ...] = () @property def characters(self) -> tuple[str, ...]: @@ -421,7 +469,11 @@ def _normalise_v2_token(inner: str) -> Optional[str]: return candidate.upper() if candidate.upper() in UPPERCASE_V2_TOKENS else candidate -def resolve_v2_tokens(text: str) -> tuple[str, tuple[str, ...]]: +def resolve_v2_tokens( + text: str, + *, + turbo_tags: bool = False, +) -> tuple[str, tuple[str, ...]]: """Normalise the v2 ``<...>`` special tokens in ``text``. Runs **after** the script has been split into segments, which is what makes @@ -437,9 +489,20 @@ def resolve_v2_tokens(text: str) -> tuple[str, tuple[str, ...]]: Args: text: One segment's text. + turbo_tags: Whether the pack speaks Turbo. A Turbo checkpoint's + tokenizer knows the **square** spelling and has never seen + ````, so on a Turbo pack a recognised token is re-spelled + into that vocabulary: ```` becomes ``[giggle]`` and + ```` becomes ``[clear throat]``, the spelling Turbo + was actually trained on. On every other pack the angle spelling is + left exactly as it is -- those packs have no tag vocabulary at all, + and rewriting a token into a spelling that means nothing there + would be a change for its own sake. Returns: - ``(resolved_text, tokens_seen)``. + ``(resolved_text, tokens_seen)``. The second element lists the + canonical token *names* (``"giggle"``, ``"UH"``, ``"clear_throat"``) + whatever spelling was emitted for them. """ seen: list[str] = [] @@ -448,7 +511,10 @@ def resolve_v2_tokens(text: str) -> tuple[str, tuple[str, ...]]: if canonical is None: return match.group(0) seen.append(canonical) - return f"<{canonical}>" + if not turbo_tags: + return f"<{canonical}>" + spelling = TURBO_TAG_V2_SPELLING.get(canonical, canonical) + return f"[{spelling}]" return _ANGLE_RE.sub(_replace, text), tuple(seen) @@ -476,10 +542,21 @@ def split_segments( Args: text: The script. - turbo_tags: Recognise the square-bracket Turbo tag set. ``False`` (the - default) means ``[laugh]`` is a character named ``"laugh"``, which - is exactly what upstream does and what the reference corpus warns - about. + turbo_tags: Recognise the square-bracket Turbo tag set. + + On a Turbo pack the tag is **kept exactly as written** and recorded + in ``segment.turbo_tags``: upstream hands the tag to the tokenizer + verbatim (``src/chatterbox/tts_turbo.py:401-403``) and never parses + it, so a Turbo checkpoint only speaks the tag if the text still + carries it. Deleting it here is what made the feature silent. + + On any other pack the nine Turbo names have no vocabulary to land + in, so they are kept as literal text too -- recorded as text, not + as a tag -- and reported, by :func:`plan_text`, as a notice. They + are deliberately **not** treated as character names: ``[laugh]`` + would be swallowed as a voice called "laugh", and the word would + leave the script without anyone being told. Every other bare + bracket keeps today's character behaviour. Returns: The segments, in order. Empty input yields no segments. @@ -494,8 +571,6 @@ def split_segments( # means 0.5 s of silence after "Hi.". It is held here rather than written # straight onto ``current`` because the tag closes the segment it lands on. pending_pause = 0.0 - # Set when a consumed tag left trailing space that must not double up. - pending_lstrip = False def _flush() -> None: nonlocal current, pending_pause @@ -508,21 +583,26 @@ def split_segments( for match in _BRACKET_RE.finditer(text): chunk = text[cursor : match.start()] - current.text += chunk.lstrip() if pending_lstrip else chunk - pending_lstrip = False + current.text += chunk cursor = match.end() body = match.group(1) + tag_name = body.strip().lower() - if turbo_tags and body.strip().lower() in _TURBO_TAG_LOOKUP: - current.turbo_tags = current.turbo_tags + (body.strip().lower(),) - # Removing a tag must not leave the double space that surrounded it: - # ``"a test. [laugh] Isn't"`` would otherwise become ``"a test. Is"`` - # and reach the model with a gap the user never typed. Trailing - # space before the tag is dropped; leading space after it is kept, - # so the single separator survives. A tag at the very start of a - # segment has nothing before it to keep the separator. - current.text = current.text.rstrip() - pending_lstrip = not current.text + if ( + tag_name in _TURBO_TAG_LOOKUP + and ":" not in body + and "|" not in body + ): + # A Turbo paralinguistic tag. The text is kept exactly as written + # on both paths: upstream hands the tag to the tokenizer verbatim + # and never parses it, so a Turbo checkpoint only speaks the tag + # if the text still carries it -- and a pack that cannot speak it + # says the word rather than swallowing it. What differs is + # whether the tag is *recorded* as one, and that is what the + # caller's ``turbo_tags`` flag means. + if turbo_tags: + current.turbo_tags = current.turbo_tags + (tag_name,) + current.text += match.group(0) continue kind, payload = _classify_bracket(body) @@ -568,7 +648,7 @@ def split_segments( ) tail = text[cursor:] - current.text += tail.lstrip() if pending_lstrip else tail + current.text += tail _flush() return segments @@ -637,6 +717,61 @@ def select_chunk_method(method: str, char_count: int) -> str: return "concatenate" +def _last_sentence_end(text: str, limit: int) -> int: + """The index of the last sentence end inside the window, or ``-1``. + + A sentence end is ``.`` / ``!`` / ``?`` followed by a space -- the three + marks that end a sentence and the space that proves it. The window is + ``text[:limit]`` plus one character, because a sentence end whose space + falls just past the window is still a sentence end, and cutting there + fills the chunk exactly. + """ + found = -1 + for match in _SENTENCE_END_RE.finditer(text[: limit + 1]): + if match.start() < limit: + found = match.start() + return found + + +def _bracket_run_at(text: str, position: int) -> Optional[tuple[int, int]]: + """The ``[...]`` / ``<...>`` run that a cut at ``position`` lands inside. + + Returns ``(open_index, close_index)``, or ``None`` when the cut is not + inside a run. The run is the one opened by the last opening bracket + before the cut whose closing bracket is at or after it. + """ + for opener, closer in _BRACKET_PAIRS: + start = text.rfind(opener, 0, position) + if start < 0: + continue + end = text.find(closer, start + 1) + if end < 0: + continue + if end >= position: + return (start, end) + return None + + +def _snap_out_of_bracket_run(text: str, take: int, limit: int) -> int: + """Move a cut that landed inside a bracket run out of it. + + Forward to the closing bracket when that still fits inside the limit, + otherwise back to before the opening bracket. A run that is longer than + the limit on its own is never split: the chunk grows past the limit + instead, which is the one documented exception to the ceiling. + """ + run = _bracket_run_at(text, take) + if run is None: + return take + start, end = run + forward = end + 1 + if forward <= limit: + return forward + if start > 0: + return start + return forward + + def chunk_text( text: str, max_chars_per_chunk: int = 400, @@ -650,6 +785,15 @@ def chunk_text( because it is the one that sounds wrong, and it only happens on a single unbroken run longer than the limit (a URL, a long number). + A cut never lands **inside** a tag. Half a ``[chuckle]`` in one chunk and + half in the next is not a tag at all: the model sees two pieces of + punctuation and says neither. So a cut that falls inside a ``[...]`` or + ``<...>`` run is moved out of it -- forward to the closing bracket when + that still fits inside the limit, otherwise back to before the opening + one. The single documented exception is a tag run that is longer than the + limit by itself: it stays whole and the chunk grows past the ceiling, + because splitting it would be worse than exceeding it. + Args: text: The script to cut. max_chars_per_chunk: The soft ceiling. Values below 1 disable cutting. @@ -664,9 +808,10 @@ def chunk_text( chunks: list[str] = [] remaining = text while len(remaining) > max_chars_per_chunk: - window = remaining[:max_chars_per_chunk] + limit = max_chars_per_chunk + window = remaining[:limit] # Sentence end, then line break, then space: in that order of preference. - cut = window.rfind(". ") + cut = _last_sentence_end(remaining, limit) if cut < 0: cut = window.rfind("\n") if cut < 0: @@ -674,10 +819,12 @@ def chunk_text( if cut <= 0: # No boundary inside the window -- one unbreakable run. Consume the # whole window, or the loop never advances. - take = max_chars_per_chunk + take = limit else: # The boundary character belongs to the chunk that ends on it. take = cut + 1 + # The cut may have landed inside a tag; move it out rather than split. + take = _snap_out_of_bracket_run(remaining, take, limit) chunks.append(remaining[:take].strip()) remaining = remaining[take:] @@ -687,11 +834,77 @@ def chunk_text( return [chunk for chunk in chunks if chunk] +# --------------------------------------------------------------------------- +# Notices (D8) -- what the engine could not honour +# --------------------------------------------------------------------------- + + +def _turbo_tag_notice(tag: str) -> str: + """The notice a Turbo tag earns on a pack that has no Turbo vocabulary.""" + return ( + f"[{tag}] is a Turbo/Nano paralinguistic tag. This pack has no Turbo " + f"tag vocabulary, so the tag is spoken as literal text -- use a Turbo " + f"or Nano pack for it to have an effect." + ) + + +#: A pause is parsed onto its segment and never rendered: nothing in this +#: repository turns ``pause_after`` into silence. +_PAUSE_NOTICE = ( + "A pause tag ([pause:1.5]) is parsed, but nothing renders it as silence: " + "the pause is not audible." +) + +#: A language is parsed onto its segment and never switched: no pack here +#: changes voice or language halfway through a script. +_LANGUAGE_NOTICE = ( + "A language tag ([de:Alice]) is parsed, but no Chatterbox pack switches " + "language per segment: the whole run speaks one language." +) + + # --------------------------------------------------------------------------- # The whole pipeline # --------------------------------------------------------------------------- +def _collect_notices( + segments: list[TextSegment], + spoken: str, + *, + turbo_tags: bool, +) -> tuple[str, ...]: + """What the engine could not honour, said out loud rather than dropped. + + Three things can be parsed and then not mean anything, and all three are + announced instead of swallowed: + + * a Turbo tag on a pack that has no Turbo vocabulary -- the word is + spoken, but it will not do what it looks like it does; + * a pause, which lands on ``pause_after`` where nothing renders it; + * a language, which lands on ``segment.language`` where no pack switches. + + A notice is never an error and never changes the audio. + """ + notices: list[str] = [] + + if not turbo_tags: + # The tags are still in the spoken text -- they were kept, not eaten -- + # so they are found there rather than remembered on the way through. + for match in _TURBO_TAG_IN_TEXT_RE.finditer(spoken): + notice = _turbo_tag_notice(match.group(1).lower()) + if notice not in notices: + notices.append(notice) + + if any(segment.pause_after for segment in segments): + notices.append(_PAUSE_NOTICE) + + if any(segment.language for segment in segments): + notices.append(_LANGUAGE_NOTICE) + + return tuple(notices) + + def plan_text( text: str, *, @@ -726,13 +939,16 @@ def plan_text( Returns: A :class:`TextPlan`. Blank input yields an empty plan with one empty - chunk, so a caller never has to special-case it. + chunk, so a caller never has to special-case it. The plan's + ``notices`` say what the engine could not honour -- a Turbo tag on a + pack without the Turbo vocabulary, a pause that is never rendered, a + language that is never switched. """ segments = split_segments(text or "", turbo_tags=turbo_tags) resolved: list[TextSegment] = [] for segment in segments: - body, _tokens = resolve_v2_tokens(segment.text) + body, _tokens = resolve_v2_tokens(segment.text, turbo_tags=turbo_tags) segment.text = apply_crash_protection(body.strip(), crash_protection_template) resolved.append(segment) @@ -748,4 +964,5 @@ def plan_text( chunks=tuple(chunks), chunk_method=method, char_count=len(spoken), + notices=_collect_notices(resolved, spoken, turbo_tags=turbo_tags), ) diff --git a/nodes/tts_node.py b/nodes/tts_node.py index cc7495c..e21227f 100644 --- a/nodes/tts_node.py +++ b/nodes/tts_node.py @@ -44,6 +44,7 @@ from ..modules.audio_utils import interrupt_exception_class, to_comfy_audio from ..modules.dtype_utils import DTYPE_AUTO from ..modules.generation import generate_audio, load_chatterbox_model from ..modules.model_discovery import get_chatterbox_model_pack_names +from ..modules.text_processing import CHUNK_COMBINATION_METHODS from .common import default_pack_name, runtime_widgets, silent_audio, validate_common_inputs __all__ = ["ChatterboxTTSNode"] @@ -266,6 +267,56 @@ class ChatterboxTTSNode(io.ComfyNode): "only; every other family ignores this." ), ), + # Long-script controls. Appended last, so every widget index + # an existing saved workflow holds is untouched. + io.Boolean.Input( + "enable_chunking", + default=True, + tooltip=( + "Split a long script into chunks and synthesise them " + "one at a time, then join them. This is what lets a " + "script of any length be spoken in full: the token " + "budget applies per chunk, not to the whole script. " + "Turn it off to send the script to the model in one " + "piece." + ), + ), + io.Int.Input( + "max_chars_per_chunk", + default=400, + min=100, + max=1000, + step=50, + tooltip=( + "How many characters go into one chunk. The cut " + "prefers a sentence end, then a line break, then a " + "space, and never lands inside a tag. Smaller chunks " + "follow a long script more closely but cost more " + "model calls." + ), + ), + io.Combo.Input( + "chunk_combination_method", + options=list(CHUNK_COMBINATION_METHODS), + default="auto", + tooltip=( + "How the chunks are joined. 'auto' resolves by the " + "length of the script: up to 500 characters joins " + "cleanly, above 500 crossfades, above 1000 pads with " + "silence." + ), + ), + io.Int.Input( + "silence_between_chunks_ms", + default=100, + min=0, + max=500, + step=25, + tooltip=( + "The gap between two chunks when the join method is " + "'silence_padding'. Ignored by the other methods." + ), + ), io.Audio.Input( "audio_prompt", optional=True, @@ -322,6 +373,10 @@ class ChatterboxTTSNode(io.ComfyNode): normalize_loudness: bool = True, num_steps: int = 10, cfg_scale: float = 1.0, + enable_chunking: bool = True, + max_chars_per_chunk: int = 400, + chunk_combination_method: str = "auto", + silence_between_chunks_ms: int = 100, audio_prompt: Optional[dict] = None, client_id: Optional[str] = None, ) -> io.NodeOutput: @@ -383,6 +438,10 @@ class ChatterboxTTSNode(io.ComfyNode): norm_loudness=normalize_loudness, num_steps=num_steps, cfg_scale=cfg_scale, + enable_chunking=enable_chunking, + max_chars_per_chunk=max_chars_per_chunk, + chunk_combination_method=chunk_combination_method, + silence_between_chunks_ms=silence_between_chunks_ms, ) except interrupt_error: # InterruptProcessingException derives from BaseException, so this diff --git a/tests/test_chunking.py b/tests/test_chunking.py new file mode 100644 index 0000000..533ef32 --- /dev/null +++ b/tests/test_chunking.py @@ -0,0 +1,181 @@ +"""Chunk boundaries -- sentence first, and never inside a tag. + +A chunk boundary is a place where the speech stops and starts again, so it has +to be somewhere a speaker would also stop: after a sentence, after a line, at +a space. And it has to be somewhere the *tags* are not, because half a +``[chuckle]`` is not a tag at all -- it is punctuation the model reads out +loud, in two pieces, neither of which means anything. + +The rule under test is therefore twofold: prefer the best boundary available, +and move a boundary that landed inside a tag out of it. The one documented +exception is a tag run that is longer than the chunk limit by itself: it stays +whole and the chunk grows past the ceiling, because splitting it would be +worse than exceeding it. +""" + +from __future__ import annotations + +import pytest + +from ComfyUI_Chatterbox.modules.text_processing import chunk_text + + +def _assert_no_split_tag(chunks: list[str]) -> None: + """No chunk may contain an opening bracket without its closing one. + + That is the whole property: a tag that is split across two chunks shows up + here as a chunk that opens a bracket and never closes it. + """ + for chunk in chunks: + opened = chunk.count("[") - chunk.count("]") + assert opened <= 0, f"a square tag was split: {chunk!r}" + opened = chunk.count("<") - chunk.count(">") + assert opened <= 0, f"an angle tag was split: {chunk!r}" + + +class TestBoundaries: + def test_a_sentence_end_is_preferred_over_a_space(self) -> None: + """The cut lands after the sentence, not at the last word before it.""" + text = "Alpha beta gamma delta. Epsilon zeta eta theta. Iota kappa." + chunks = chunk_text(text, 45) + assert chunks[0] == "Alpha beta gamma delta.", ( + "the cut took the last space instead of the sentence end" + ) + + def test_an_exclamation_mark_ends_a_sentence_too(self) -> None: + chunks = chunk_text("Alpha beta gamma delta! Epsilon zeta eta theta!", 45) + assert chunks[0] == "Alpha beta gamma delta!" + + def test_a_question_mark_ends_a_sentence_too(self) -> None: + chunks = chunk_text("Alpha beta gamma delta? Epsilon zeta eta theta?", 45) + assert chunks[0] == "Alpha beta gamma delta?" + + def test_a_space_is_used_when_no_sentence_end_exists(self) -> None: + chunks = chunk_text("alpha beta gamma delta epsilon zeta", 20) + assert chunks[0] == "alpha beta gamma" + + def test_a_line_break_is_used_when_no_sentence_end_exists(self) -> None: + chunks = chunk_text("alpha beta gamma\ndelta epsilon zeta", 20) + assert chunks[0] == "alpha beta gamma" + + def test_a_hard_cut_only_happens_for_an_unbreakable_run(self) -> None: + """A URL has no spaces to cut at, so the limit is the only boundary.""" + url = "https://example.com/" + "a" * 480 + chunks = chunk_text(url, 100) + assert len(chunks) == 5 + assert all(len(chunk) == 100 for chunk in chunks) + + def test_every_chunk_is_at_most_the_limit(self) -> None: + text = "A sentence here. " * 100 + chunks = chunk_text(text, 200) + assert all(len(chunk) <= 200 for chunk in chunks) + + +class TestTagSafeBoundaries: + def test_a_chunk_boundary_never_splits_a_square_tag(self) -> None: + """The tag sits where a naive cut would land inside it.""" + text = "x" * 398 + "[chuckle]" + "y" * 500 + chunks = chunk_text(text, 400) + _assert_no_split_tag(chunks) + assert "[chuckle]" in "".join(chunks) + + def test_a_chunk_boundary_never_splits_an_angle_tag(self) -> None: + text = "x" * 398 + "" + "y" * 500 + chunks = chunk_text(text, 400) + _assert_no_split_tag(chunks) + assert "" in "".join(chunks) + + def test_a_tag_with_a_space_in_it_is_never_split(self) -> None: + """``[clear throat]`` is the one Turbo tag with an internal space, + which is exactly where a space-based cut would land.""" + text = "[clear throat] and the rest of the script goes on here" + chunks = chunk_text(text, 12) + _assert_no_split_tag(chunks) + assert "[clear throat]" in chunks[0] + + def test_a_tag_longer_than_the_limit_is_not_split(self) -> None: + """The one documented exception: the chunk grows past the ceiling.""" + tag = "[" + "a" * 498 + "]" + chunks = chunk_text(tag, 100) + assert chunks == [tag] + assert len(chunks[0]) > 100, "the exception is supposed to exceed the limit" + + def test_a_tag_at_the_very_start_of_a_script_is_not_split(self) -> None: + text = "[chuckle]" + "y" * 500 + chunks = chunk_text(text, 400) + _assert_no_split_tag(chunks) + assert chunks[0].startswith("[chuckle]") + + def test_a_tag_is_never_split_in_a_real_script(self) -> None: + """The whole thing end to end, with tags scattered through it.""" + script = ( + "First sentence here. [laugh] Second sentence here. " + " Third sentence here. [clear throat] Fourth sentence. " + ) * 12 + chunks = chunk_text(script, 120) + _assert_no_split_tag(chunks) + assert "[laugh]" in "".join(chunks) + assert "[clear throat]" in "".join(chunks) + + +class TestNothingIsLost: + def test_the_concatenation_of_the_chunks_loses_no_words(self) -> None: + text = "A sentence here. " * 100 + chunks = chunk_text(text, 200) + assert text.split() == " ".join(chunks).split() + + def test_no_words_are_lost_when_tags_are_in_the_script(self) -> None: + script = ("One two three. [laugh] Four five six. Seven. ") * 20 + chunks = chunk_text(script, 100) + assert script.split() == " ".join(chunks).split() + + def test_chunking_can_still_be_disabled(self) -> None: + text = "A sentence here. " * 100 + assert chunk_text(text, 200, enabled=False) == [text] + + @pytest.mark.parametrize("limit", [1, 7, 40, 400]) + def test_the_chunker_always_makes_progress(self, limit: int) -> None: + """A limit of one is the pathological case: it must still terminate.""" + chunks = chunk_text("a bb ccc [d] eeeee", limit) + assert "".join(chunks).replace(" ", "") == "abbccc[d]eeeee" + + +class TestThePlanCutsTheScript: + """The engine's own view of a long script. + + The model-level half of this -- a stub that says more when the script is + longer -- lives in ``tests/test_generation.py``, because the module that + drives the model imports ComfyUI and this file deliberately does not. + """ + + def test_a_long_script_is_cut_into_several_chunks(self) -> None: + from ComfyUI_Chatterbox.modules.text_processing import plan_text + + plan = plan_text( + "A sentence here. " * 100, + crash_protection_template="", + max_chars_per_chunk=200, + ) + assert len(plan.chunks) > 1 + + def test_the_cut_loses_no_words(self) -> None: + from ComfyUI_Chatterbox.modules.text_processing import plan_text + + script = "A sentence here. " * 100 + plan = plan_text( + script, + crash_protection_template="", + max_chars_per_chunk=200, + ) + assert " ".join(plan.chunks).split() == script.split() + + def test_the_cut_can_be_switched_off(self) -> None: + from ComfyUI_Chatterbox.modules.text_processing import plan_text + + plan = plan_text( + "A sentence here. " * 100, + crash_protection_template="", + enable_chunking=False, + max_chars_per_chunk=200, + ) + assert len(plan.chunks) == 1 diff --git a/tests/test_e2e_smoke.py b/tests/test_e2e_smoke.py index dfe4d67..5a2ffae 100644 --- a/tests/test_e2e_smoke.py +++ b/tests/test_e2e_smoke.py @@ -255,4 +255,53 @@ class TestVoiceConversion: assert wav.numel() > 0 assert bool(wav.isfinite().all()) assert rms(wav) > RMS_FLOOR - assert duration_seconds(wav) > MIN_SECONDS \ No newline at end of file + assert duration_seconds(wav) > MIN_SECONDS + +#: Long enough to be cut into more than one chunk at the default limit. +CHUNKED_SCRIPT = " ".join([SHORT_SCRIPT] * 12) + + +class TestChunkedRealGeneration: + """The chunked path, against the real model. + + The same property as the unchunked test -- a longer script has to produce + longer audio -- but through the chunk loop and the join, so it also proves + that real waveforms survive being cut and put back together. + """ + + def test_a_long_script_grows_with_the_chunked_path(self, model) -> None: + from ComfyUI_Chatterbox.modules.generation import ( # noqa: PLC0415 + generate_audio, + ) + + short, rate = generate_audio( + model, SHORT_SCRIPT, seed=1, max_chars_per_chunk=400 + ) + long, _ = generate_audio( + model, CHUNKED_SCRIPT, seed=1, max_chars_per_chunk=400 + ) + + short_seconds = duration_seconds(short, rate) + long_seconds = duration_seconds(long, rate) + ratio = long_seconds / short_seconds + + assert math.isfinite(ratio) + assert ratio > MIN_LENGTH_RATIO, ( + f"12x the script produced only {ratio:.2f}x the audio " + f"({short_seconds:.2f}s -> {long_seconds:.2f}s) through the " + "chunked path" + ) + + def test_the_chunked_audio_is_not_silent(self, model) -> None: + from ComfyUI_Chatterbox.modules.generation import ( # noqa: PLC0415 + generate_audio, + ) + + wave, _ = generate_audio( + model, CHUNKED_SCRIPT, seed=1, max_chars_per_chunk=400 + ) + level = rms(wave) + assert level > RMS_FLOOR, ( + f"RMS {level:.5f} is at or below the silence floor {RMS_FLOOR}; " + "the chunked path produced audio of the right length but no signal" + ) diff --git a/tests/test_generation.py b/tests/test_generation.py index 8235016..cd430ec 100644 --- a/tests/test_generation.py +++ b/tests/test_generation.py @@ -651,16 +651,22 @@ class TestInterruptHandling: assert returned is None def test_the_interrupt_clause_precedes_the_generic_one(self) -> None: - """The source-level contract, asserted structurally rather than by eye.""" + """The source-level contract, asserted structurally rather than by eye. + + The guarded try/except lives in ``_generate_one_chunk`` since the + chunk loop was added: that is the function that calls + ``model.generate``, and it is the one that has to let a cancel + through. + """ tree = ast.parse(read_source(generation_module), filename=generation_module.__file__) function = next( node for node in ast.walk(tree) - if isinstance(node, ast.FunctionDef) and node.name == "generate_audio" + if isinstance(node, ast.FunctionDef) and node.name == "_generate_one_chunk" ) guarded = [node for node in ast.walk(function) if isinstance(node, ast.Try) and len(node.handlers) >= 2] - assert guarded, "generate_audio has no multi-clause try/except to inspect" + assert guarded, "_generate_one_chunk has no multi-clause try/except to inspect" for node in guarded: kinds = [] @@ -925,4 +931,179 @@ class TestModuleSurface: the same function is one more thing that can drift. """ with pytest.raises(ImportError): - importlib.import_module("modules.utils") \ No newline at end of file + importlib.import_module("modules.utils") + +# --------------------------------------------------------------------------- +# Long scripts: one model call per chunk, joined with the resolved method +# --------------------------------------------------------------------------- + + +class SeededStub(StubPack): + """A stub whose output depends on the seed that was set before the call. + + That is what makes determinism observable: the same seed gives the same + waveform, a different seed gives a different one. + """ + + def generate(self, *args, **kwargs) -> torch.Tensor: + self.calls.append({"args": list(args), **kwargs}) + level = float(torch.rand(1).item()) + return torch.full((1, self.samples), level) + + +class LengthStub: + """A stub whose audio is as long as the text it was handed. + + The real model speaks longer scripts for longer, so a stub that does the + same is enough to show that a long script is spoken *in full* rather than + truncated. + """ + + def __init__(self) -> None: + self.calls: list[dict] = [] + + def generate(self, *args, **kwargs) -> torch.Tensor: + self.calls.append({"args": list(args), **kwargs}) + text = args[0] if args else kwargs.get("text", "") + return torch.full((1, max(1, len(text) * 10)), 0.1) + + +#: Long enough to be cut into several chunks at the default limit. +LONG_SCRIPT = "This is one sentence of the long script. " * 40 + + +class TestChunkedSynthesis: + def test_a_short_script_is_one_model_call(self) -> None: + model = StubPack() + generate_audio(model, "hello there", seed=1) + assert len(model.calls) == 1 + + def test_a_long_script_calls_the_model_once_per_chunk(self) -> None: + model = StubPack() + generate_audio(model, LONG_SCRIPT, seed=1, max_chars_per_chunk=400) + assert len(model.calls) >= 2 + + def test_the_chunks_are_the_whole_script(self) -> None: + """No word is dropped between the chunks.""" + model = StubPack() + generate_audio(model, LONG_SCRIPT, seed=1, max_chars_per_chunk=400) + spoken = " ".join(call["args"][0] for call in model.calls) + assert spoken.split() == LONG_SCRIPT.split() + + def test_the_chunks_are_joined_with_the_resolved_method(self) -> None: + """``auto`` resolves a long script to ``silence_padding``, so the + joined audio is the parts plus one gap between each pair.""" + model = StubPack(samples=24000) + waveform, rate = generate_audio( + model, LONG_SCRIPT, seed=1, max_chars_per_chunk=400, silence_between_chunks_ms=100 + ) + calls = len(model.calls) + gap = int(rate * 100 / 1000) + assert waveform.shape[-1] == calls * 24000 + (calls - 1) * gap + + def test_concatenate_is_used_for_a_500_char_script(self) -> None: + """Exactly 500 characters is still ``concatenate``: the comparison is + strictly greater than.""" + script = "word " * 100 + assert len(script.strip()) == 499 + model = StubPack(samples=24000) + waveform, _rate = generate_audio( + model, script, seed=1, max_chars_per_chunk=400 + ) + assert len(model.calls) == 2 + assert waveform.shape[-1] == 2 * 24000 + + def test_crossfade_is_used_for_a_501_char_script(self) -> None: + """One character more and the join crossfades: the total is the sum + of the parts minus the overlap.""" + script = "word " * 100 + "x" + assert len(script.strip()) == 501 + model = StubPack(samples=24000) + waveform, rate = generate_audio( + model, script, seed=1, max_chars_per_chunk=400 + ) + assert len(model.calls) == 2 + overlap = int(rate * 10 / 1000) + assert waveform.shape[-1] == 2 * 24000 - overlap + + def test_chunking_can_be_disabled(self) -> None: + model = StubPack() + generate_audio(model, LONG_SCRIPT, seed=1, enable_chunking=False) + assert len(model.calls) == 1 + + def test_the_chunk_count_is_logged(self, caplog) -> None: + with caplog.at_level(logging.INFO, logger=""): + generate_audio(StubPack(), LONG_SCRIPT, seed=1, max_chars_per_chunk=400) + messages = [record.getMessage() for record in caplog.records] + assert any( + "chunk" in message and "silence_padding" in message + for message in messages + ), f"the chunk count and the join method were not logged: {messages!r}" + + +class TestChunkedDeterminism: + """Same script, same seed, same audio -- chunk by chunk.""" + + def test_each_chunk_gets_its_own_seed(self, monkeypatch) -> None: + seen: list[int] = [] + monkeypatch.setattr(generation_module, "set_seed", seen.append) + generate_audio(StubPack(), LONG_SCRIPT, seed=42, max_chars_per_chunk=400) + assert len(seen) >= 2 + assert seen == [42 + index for index in range(len(seen))] + + def test_a_zero_seed_stays_random(self, monkeypatch) -> None: + """Zero is the widget's "random" convention: every chunk keeps it.""" + seen: list[int] = [] + monkeypatch.setattr(generation_module, "set_seed", seen.append) + generate_audio(StubPack(), LONG_SCRIPT, seed=0, max_chars_per_chunk=400) + assert seen and all(value == 0 for value in seen) + + def test_the_same_script_and_seed_produce_identical_audio(self) -> None: + first, _ = generate_audio( + SeededStub(), LONG_SCRIPT, seed=7, max_chars_per_chunk=400 + ) + second, _ = generate_audio( + SeededStub(), LONG_SCRIPT, seed=7, max_chars_per_chunk=400 + ) + assert torch.equal(first, second) + + def test_a_different_seed_produces_different_audio(self) -> None: + first, _ = generate_audio( + SeededStub(), LONG_SCRIPT, seed=7, max_chars_per_chunk=400 + ) + second, _ = generate_audio( + SeededStub(), LONG_SCRIPT, seed=8, max_chars_per_chunk=400 + ) + assert not torch.equal(first, second) + + +class TestLongScriptSynthesis: + """A long script is spoken in full: more text, more audio, no silence.""" + + def test_a_long_script_produces_more_audio_than_a_short_one(self) -> None: + short = "A short line of dialogue." + long = "A short line of dialogue. " * 120 + short_wave, rate = generate_audio(LengthStub(), short, seed=1) + long_wave, _ = generate_audio( + LengthStub(), long, seed=1, max_chars_per_chunk=400 + ) + assert long_wave.shape[-1] > short_wave.shape[-1] + assert rate == 24000 + + def test_the_joined_audio_is_above_the_rms_floor(self) -> None: + """The same floor the real-model smoke test uses (0.01).""" + wave, _ = generate_audio( + LengthStub(), + "A short line of dialogue. " * 120, + seed=1, + max_chars_per_chunk=400, + ) + level = float(torch.sqrt(torch.mean(wave.float() ** 2)).item()) + assert level > 0.01, f"RMS {level:.5f} is at or below the silence floor" + + def test_the_long_script_is_synthesised_in_several_chunks(self) -> None: + stub = LengthStub() + generate_audio( + stub, "A short line of dialogue. " * 120, seed=1, max_chars_per_chunk=400 + ) + assert len(stub.calls) >= 2 diff --git a/tests/test_node_compat.py b/tests/test_node_compat.py index 29529fb..70552b7 100644 --- a/tests/test_node_compat.py +++ b/tests/test_node_compat.py @@ -74,6 +74,11 @@ APPENDED_WIDGET_ORDER: dict[str, tuple[str, ...]] = { # Flash controls (Single Language Pack round, 2026-10-07). "num_steps", "cfg_scale", + # Long-script controls (emotion-tag and chunking round, 2026-10-22). + "enable_chunking", + "max_chars_per_chunk", + "chunk_combination_method", + "silence_between_chunks_ms", ), VC_NODE_ID: ( "device", diff --git a/tests/test_node_schema.py b/tests/test_node_schema.py index b90eaed..97f3a55 100644 --- a/tests/test_node_schema.py +++ b/tests/test_node_schema.py @@ -60,6 +60,12 @@ APPENDED_WIDGETS: dict[str, tuple[str, ...]] = { # (2026-10-07). Flash packs ignore the Turbo ones and vice versa. "num_steps", "cfg_scale", + # Long-script controls, appended with the emotion-tag and chunking + # round (2026-10-22). Every family uses them. + "enable_chunking", + "max_chars_per_chunk", + "chunk_combination_method", + "silence_between_chunks_ms", ), VC_NODE_ID: ( "device", diff --git a/tests/test_text_processing.py b/tests/test_text_processing.py index b8dc9f8..26f9e88 100644 --- a/tests/test_text_processing.py +++ b/tests/test_text_processing.py @@ -26,6 +26,7 @@ from ComfyUI_Chatterbox.modules.text_processing import ( MIN_SAFE_SEGMENT_CHARS, PARAM_ALIASES, TURBO_TAGS, + TURBO_TAG_V2_SPELLING, V2_SPECIAL_TOKENS, apply_crash_protection, chunk_text, @@ -245,12 +246,85 @@ class TestV2Tokens: assert "" in plan.text -class TestTurboTagsAreADifferentVocabulary: - """``[laugh]`` is a character name here, and a Turbo tag there.""" +class TestV2TokensOnTheTurboPath: + """A Turbo checkpoint knows the square spelling and not the angle one. - def test_square_brackets_are_a_character_on_the_classic_path(self) -> None: - segments = split_segments("[laugh] Hello.") - assert segments[0].character == "laugh" + The v2 set writes ```` with an underscore and Turbo's + tokenizer was trained on ``[clear throat]`` with a space, so the one + spelling that differs is published rather than buried. + """ + + def test_v2_tokens_are_untouched_on_the_classic_path(self) -> None: + assert resolve_v2_tokens("")[0] == "" + + def test_v2_tokens_become_square_brackets_on_the_turbo_path(self) -> None: + assert resolve_v2_tokens("", turbo_tags=True)[0] == "[giggle]" + + def test_clear_throat_uses_the_turbo_spelling(self) -> None: + assert resolve_v2_tokens("", turbo_tags=True)[0] == ( + "[clear throat]" + ) + + def test_uh_keeps_its_casing_in_both_spellings(self) -> None: + assert resolve_v2_tokens("")[0] == "" + assert resolve_v2_tokens("", turbo_tags=True)[0] == "[UH]" + + def test_an_unknown_angle_token_is_left_alone_on_the_turbo_path(self) -> None: + resolved, seen = resolve_v2_tokens("", turbo_tags=True) + assert seen == () + assert resolved == "" + + def test_a_v2_token_after_a_character_keeps_the_character(self) -> None: + """The conversion happens after the character is resolved, so the + character survives it.""" + segments = split_segments("[Alice] Hello.", turbo_tags=True) + assert segments[0].character == "Alice" + assert "" in segments[0].text + + plan = plan_text("[Alice] Hello.", turbo_tags=True) + assert plan.segments[0].character == "Alice" + assert "[giggle]" in plan.text + + def test_the_turbo_plan_text_carries_the_square_spelling(self) -> None: + # Crash protection is off here: it pads any segment shorter than + # twenty-one characters, and a two-word script is shorter than that. + plan = plan_text(" Hi.", turbo_tags=True, crash_protection_template="") + assert plan.text == "[giggle] Hi." + + def test_the_only_differing_spelling_is_published(self) -> None: + assert TURBO_TAG_V2_SPELLING == {"clear_throat": "clear throat"} + + +class TestTurboTagsAreADifferentVocabulary: + """``[laugh]`` is a Turbo tag, and a Turbo tag is never eaten.""" + + def test_a_turbo_tag_is_preserved_in_the_segment_text(self) -> None: + """The tag reaches the model exactly as the user typed it. + + Upstream does not parse these tags; it hands the text to the + tokenizer, which knows the square spelling. A tag that is deleted + before the model sees it is a feature that looks implemented and is + not. + """ + segment = split_segments("[laugh] Hi.", turbo_tags=True)[0] + assert segment.text == "[laugh] Hi." + + def test_a_turbo_tag_is_still_recorded(self) -> None: + """Recognition is kept; only the deletion is gone.""" + segment = split_segments("[laugh] Hi.", turbo_tags=True)[0] + assert segment.turbo_tags == ("laugh",) + assert segment.character is None + + def test_a_turbo_tag_at_the_start_of_a_segment_is_preserved(self) -> None: + segment = split_segments("[Alice] [chuckle] Hello.", turbo_tags=True) + assert segment[0].text == "[chuckle] Hello." + assert segment[0].turbo_tags == ("chuckle",) + assert segment[0].character == "Alice" + + def test_two_turbo_tags_are_both_preserved_and_recorded(self) -> None: + segment = split_segments("[laugh] and [chuckle] that", turbo_tags=True)[0] + assert segment.text == "[laugh] and [chuckle] that" + assert segment.turbo_tags == ("laugh", "chuckle") def test_the_same_text_is_a_tag_on_the_turbo_path(self) -> None: segments = split_segments("[laugh] Hello.", turbo_tags=True) @@ -282,6 +356,49 @@ class TestTurboTagsAreADifferentVocabulary: assert "" in segments[0].text +class TestATurboTagIsNotAVoiceName: + """On a pack that cannot speak them, Turbo tags are text -- not a name. + + Treating ``[laugh]`` as a voice called "laugh" swallows the word: the + script says nothing where the user wrote a direction, and the voice never + changes either, because this repository ships no voice library. + """ + + def test_a_turbo_tag_name_on_the_classic_path_is_not_a_character(self) -> None: + segment = split_segments("[laugh] Hello.")[0] + assert segment.character is None + assert "[laugh]" in segment.text + + def test_an_unknown_bracket_on_the_classic_path_is_still_a_character(self) -> None: + """The unchanged behaviour: every other bare bracket is a name.""" + segment = split_segments("[Alice] Hello.")[0] + assert segment.character == "Alice" + assert segment.text == "Hello." + + def test_a_turbo_tag_name_on_the_classic_path_produces_a_notice(self) -> None: + notices = plan_text("[laugh] Hi.").notices + assert notices, "the tag was swallowed without a word" + assert any("Turbo" in notice and "laugh" in notice for notice in notices) + + def test_the_classic_path_does_not_record_the_tag_as_a_turbo_tag(self) -> None: + """It is not a Turbo tag here -- it is the word the user typed.""" + assert split_segments("[laugh] Hello.")[0].turbo_tags == () + + @pytest.mark.parametrize("body", ["de:Alice", "Alice|seed:42", "Alice|temp:0.5"]) + def test_a_bracket_with_a_colon_or_pipe_is_never_treated_as_a_turbo_tag( + self, body: str + ) -> None: + """Only the nine bare names are Turbo tags; the other forms keep + their own meaning.""" + segment = split_segments(f"[{body}] Hello.")[0] + assert segment.turbo_tags == () + + def test_a_turbo_tag_is_not_noticed_on_the_turbo_path(self) -> None: + """It is native there, so there is nothing to announce.""" + plan = plan_text("[laugh] Hi.", turbo_tags=True) + assert not [notice for notice in plan.notices if "Turbo" in notice] + + # --------------------------------------------------------------------------- # Chunking # --------------------------------------------------------------------------- @@ -342,6 +459,8 @@ class TestChunking: assert len(chunks) > 1 assert all(len(chunk) <= 200 for chunk in chunks) assert "".join(chunks).replace(" ", "") == text.replace(" ", "") + # No word is lost or split in half on the way through. + assert text.split() == " ".join(chunks).split() def test_cutting_is_disabled_by_the_flag(self) -> None: text = "A sentence here. " * 100 @@ -377,3 +496,60 @@ class TestPlanText: assert CRASH_PROTECTION_TEMPLATE == "hmm ,, {seg} hmm ,," assert MIN_SAFE_SEGMENT_CHARS == 21 assert apply_crash_protection("Short.").startswith("hmm ,,") + + def test_auto_is_unchanged_by_the_new_boundary_rule(self) -> None: + """The tag-safe cut changes *where* a script is cut, not *how long* it + is: ``auto`` still resolves on the script's length alone.""" + assert plan_text("x" * 500).chunk_method == "concatenate" + assert plan_text("x" * 501).chunk_method == "crossfade" + assert plan_text("x" * 1001).chunk_method == "silence_padding" + + def test_a_script_with_tags_resolves_auto_on_its_length(self) -> None: + script = "[laugh] " + ("x" * 600) + assert plan_text(script).chunk_method == "crossfade" + + +# --------------------------------------------------------------------------- +# Notices -- what the engine could not honour +# --------------------------------------------------------------------------- + + +class TestNotices: + """A tag that is parsed and then means nothing is announced, not dropped. + + The alternative is a script that quietly loses a word, and a user who has + no idea why ``[laugh]`` did nothing. + """ + + def test_a_plain_script_has_no_notices(self) -> None: + assert plan_text("Hello there, this is a plain script.").notices == () + + def test_a_turbo_tag_on_a_non_turbo_pack_is_noticed(self) -> None: + notices = plan_text("[laugh] Hi.").notices + assert len(notices) == 1 + assert "[laugh]" in notices[0] + assert "Turbo" in notices[0] + + def test_the_same_tag_is_not_noticed_twice(self) -> None: + plan = plan_text("[laugh] Hi. [laugh] Bye.") + assert len([notice for notice in plan.notices if "Turbo" in notice]) == 1 + + def test_a_pause_is_noticed_rather_than_silently_dropped(self) -> None: + notices = plan_text("Hi there. [pause:1.5] And more.").notices + assert any("pause" in notice.lower() for notice in notices) + + def test_a_language_is_noticed_rather_than_silently_dropped(self) -> None: + notices = plan_text("[de:Alice] Hallo.").notices + assert any("language" in notice.lower() for notice in notices) + + def test_a_notice_never_changes_the_audio(self) -> None: + """A notice is a note about the text, not a change to it: the same + words are spoken on either path.""" + classic = plan_text("[laugh] Hi there, this is long enough to speak.") + turbo = plan_text( + "[laugh] Hi there, this is long enough to speak.", turbo_tags=True + ) + assert classic.text == turbo.text + assert classic.chunks == turbo.chunks + assert classic.notices != () + assert turbo.notices == () diff --git a/tests/test_tts_routing.py b/tests/test_tts_routing.py index b6ab3a5..08a45e4 100644 --- a/tests/test_tts_routing.py +++ b/tests/test_tts_routing.py @@ -12,6 +12,8 @@ whole family. The pieces that make that safe: from __future__ import annotations +import logging + import numpy as np import pytest import torch @@ -272,15 +274,30 @@ class TestGenerateAudioDispatch: assert "max_new_tokens" not in call assert "flow_cfg_scale" not in call - def test_turbo_tags_are_recognised_before_the_model_call(self) -> None: - """The parser consumes the nine square-bracket tags; the model never - sees one. Recognition is the point, not pass-through.""" + def test_a_turbo_tag_reaches_the_model_verbatim(self) -> None: + """The tag is the model's own vocabulary, so it is handed over as it + was written. + + Upstream never parses the nine square-bracket tags: it injects them + into the prompt and lets the tokenizer speak them + (``src/chatterbox/tts_turbo.py:401-403``). A Turbo checkpoint only + says ``[chuckle]`` if the text still carries it, so stripping it here + made the feature silent while looking implemented. + """ model = ChatterboxTurboTTS() - generate_audio(model, "Hello [laugh] there") + generate_audio(model, "Hello [chuckle] there") spoken = model.calls[0]["args"][0] - assert "[laugh]" not in spoken + assert "[chuckle]" in spoken assert "Hello" in spoken and "there" in spoken + def test_a_turbo_model_receives_exaggeration_zero(self) -> None: + """Turbo ignores exaggeration for sampling (``hp.emotion_adv`` is + False on that family) and upstream warns about every non-zero + value -- once per chunk.""" + model = ChatterboxTurboTTS() + generate_audio(model, "Hello there", exaggeration=0.5) + assert model.calls[0]["exaggeration"] == 0.0 + def test_a_classic_model_keeps_its_own_sampler(self) -> None: class PlainChatterboxTTS: def __init__(self) -> None: @@ -299,6 +316,60 @@ class TestGenerateAudioDispatch: assert "norm_loudness" not in call +class TestClassicPathTextSemantics: + """The classic path goes through the text engine too. + + That is what makes chunking possible at all -- there is nowhere else to + cut a long script -- and it is what makes ``[Alice]`` open a character + segment instead of being read out loud as a name. + """ + + class _Plain: + def __init__(self) -> None: + self.calls: list[dict] = [] + + def generate(self, *args, **kwargs) -> torch.Tensor: + self.calls.append({"args": list(args), **kwargs}) + return torch.zeros((1, 24000)) + + def test_the_classic_path_consumes_a_character_tag(self) -> None: + model = self._Plain() + generate_audio(model, "[Alice] Hi.") + assert model.calls[0]["args"][0] == "Hi." + + def test_the_classic_path_keeps_a_turbo_tag_name_as_text(self) -> None: + """A Turbo tag on a pack with no Turbo vocabulary is *said*, not + swallowed as a voice name.""" + model = self._Plain() + generate_audio(model, "[laugh] Hi.") + assert model.calls[0]["args"][0] == "[laugh] Hi." + + def test_the_classic_path_reports_a_notice(self, caplog) -> None: + with caplog.at_level(logging.WARNING, logger=""): + generate_audio(self._Plain(), "[laugh] Hi.") + messages = [record.getMessage() for record in caplog.records] + assert any("Turbo" in message for message in messages), ( + "a Turbo tag on a classic pack must be announced, not dropped: " + f"{messages!r}" + ) + + def test_a_pause_tag_is_reported_rather_than_silently_dropped( + self, caplog + ) -> None: + with caplog.at_level(logging.WARNING, logger=""): + generate_audio(self._Plain(), "Hi. [pause:1.5] There.") + messages = [record.getMessage() for record in caplog.records] + assert any("pause" in message.lower() for message in messages), ( + f"the parsed-but-unrendered pause was dropped on the floor: {messages!r}" + ) + + def test_a_clean_script_reports_nothing(self, caplog) -> None: + with caplog.at_level(logging.WARNING, logger=""): + generate_audio(self._Plain(), "Hello there.") + messages = [record.getMessage() for record in caplog.records] + assert messages == [] + + # --------------------------------------------------------------------------- # The main node's schema carries the Turbo controls # --------------------------------------------------------------------------- @@ -357,6 +428,81 @@ class TestTTSNodeCarriesTheTurboControls: assert seen["top_k"] == 64 assert seen["norm_loudness"] is False + def test_the_chunking_widgets_reach_generate_audio(self, monkeypatch) -> None: + """The four long-script widgets are not decoration: the values the + user picked have to arrive at the chunk loop.""" + from ComfyUI_Chatterbox.nodes import tts_node + + seen: dict = {} + monkeypatch.setattr( + tts_node, + "load_chatterbox_model", + lambda *a, **k: (object(), object(), object()), + ) + + def _fake_generate_audio(model, text, **kwargs): + seen.update(kwargs) + return torch.zeros((1, 24000)), 24000 + + monkeypatch.setattr(tts_node, "generate_audio", _fake_generate_audio) + tts_node.ChatterboxTTSNode.execute( + model_pack_name="chatterbox", + text="Hello.", + max_new_tokens=1000, + flow_cfg_scale=0.7, + exaggeration=0.5, + temperature=0.8, + cfg_weight=0.5, + repetition_penalty=1.2, + min_p=0.05, + top_p=1.0, + seed=1, + use_watermark=False, + enable_chunking=False, + max_chars_per_chunk=250, + chunk_combination_method="crossfade", + silence_between_chunks_ms=25, + ) + assert seen["enable_chunking"] is False + assert seen["max_chars_per_chunk"] == 250 + assert seen["chunk_combination_method"] == "crossfade" + assert seen["silence_between_chunks_ms"] == 25 + + def test_the_node_defaults_match_the_schema_defaults(self, monkeypatch) -> None: + """A run that says nothing about chunking gets the schema's defaults.""" + from ComfyUI_Chatterbox.nodes import tts_node + + seen: dict = {} + monkeypatch.setattr( + tts_node, + "load_chatterbox_model", + lambda *a, **k: (object(), object(), object()), + ) + + def _fake_generate_audio(model, text, **kwargs): + seen.update(kwargs) + return torch.zeros((1, 24000)), 24000 + + monkeypatch.setattr(tts_node, "generate_audio", _fake_generate_audio) + tts_node.ChatterboxTTSNode.execute( + model_pack_name="chatterbox", + text="Hello.", + max_new_tokens=1000, + flow_cfg_scale=0.7, + exaggeration=0.5, + temperature=0.8, + cfg_weight=0.5, + repetition_penalty=1.2, + min_p=0.05, + top_p=1.0, + seed=1, + use_watermark=False, + ) + assert seen["enable_chunking"] is True + assert seen["max_chars_per_chunk"] == 400 + assert seen["chunk_combination_method"] == "auto" + assert seen["silence_between_chunks_ms"] == 100 + class TestSharedFamilyFolderRouting: """One folder, two members: the dropdown entry picks the weights.""" diff --git a/tests/test_workflow_compat.py b/tests/test_workflow_compat.py index 3c65860..fe3109d 100644 --- a/tests/test_workflow_compat.py +++ b/tests/test_workflow_compat.py @@ -138,15 +138,15 @@ class TestShippedWorkflowsWereFound: class TestTTSWorkflowBindsByIndex: - """13 stored values, ``randomize`` at 11, ``use_watermark`` at 12.""" + """24 stored values, ``randomize`` at 11, ``use_watermark`` at 12.""" @pytest.mark.parametrize("path", WORKFLOW_FILES, ids=WORKFLOW_IDS) - def test_widget_values_are_13_long(self, path: str) -> None: + def test_widget_values_are_24_long(self, path: str) -> None: tts_nodes = _nodes_by_type(_load(path)).get("ChatterboxTTS", []) for node in tts_nodes: - assert len(node["widgets_values"]) == 13, ( + assert len(node["widgets_values"]) == 24, ( f"{path}: ChatterboxTTS stores {len(node['widgets_values'])} " - f"values, expected 13 (12 required inputs + one seed companion)" + f"values, expected 24 (23 widget inputs + one seed companion)" ) @pytest.mark.parametrize("path", WORKFLOW_FILES, ids=WORKFLOW_IDS) @@ -253,7 +253,8 @@ class TestShippedWidgetValuesAgainstTheLiveSchema: # The S6.3 runtime widgets, then the Turbo-family controls appended # when the Turbo node merged into this one (2026-10-06), then the # Flash controls appended with the Single Language Pack round - # (2026-10-07). + # (2026-10-07), then the long-script controls appended with the + # emotion-tag and chunking round (2026-10-22). assert widgets[len(frozen) :] == ( "device", "dtype", @@ -262,8 +263,12 @@ class TestShippedWidgetValuesAgainstTheLiveSchema: "normalize_loudness", "num_steps", "cfg_scale", + "enable_chunking", + "max_chars_per_chunk", + "chunk_combination_method", + "silence_between_chunks_ms", ) - assert len(widgets) == len(frozen) + 7 + assert len(widgets) == len(frozen) + 11 def test_the_tts_node_validates_the_shipped_pack(self) -> None: path = os.path.join(ROOT_DIR, "example_workflows", "ChatterboxTTS-workflow.json")