WildAi
125e06594b
Every entry point that conditions on a reference clip now takes the waveform in memory: the vendored loaders gained load_audio_at (path, ndarray or tensor -> float32 mono at the target rate, soundfile + torchaudio, no librosa), and the resamples ride the new dsp module. The node hands over preprocess_comfy_audio output at the right rate per site (S3Gen 24k for conditioning, the tokenizer's 16k for the VC source), so temp_wav_path and its temp-file machinery are gone -- which also removes the 16-bit PCM quantization the file round-trip imposed. Resampling is the one deliberate numerical difference: torchaudio's sinc kernel instead of soxr, touching only clips supplied at call time. The built-in-voice shared-assets E2E stays byte-identical (md5 370acf3f9469d55759c3d6096b2628ff); the reference-clip E2E changes measurably in the last digits, as expected and reported.
1356 lines
50 KiB
Python
1356 lines
50 KiB
Python
"""Model acquisition and synthesis for the Chatterbox nodes.
|
|
|
|
The node layer (B9) calls into this module and nowhere else: give it a pack
|
|
name and it returns ``(patcher, tts_model, vc_model)``; give that a text and it
|
|
returns ``(tensor, sample_rate)``.
|
|
|
|
**Ordering is load-bearing.** :func:`load_chatterbox_model` performs, in this
|
|
order and no other:
|
|
|
|
1. resolve the attention mode against this machine's capabilities;
|
|
2. resolve the dtype and the devices;
|
|
3. compute the cache key from the *resolved* values;
|
|
4. **evict the previous model, before anything is allocated**;
|
|
5. build the handler;
|
|
6. build the patcher;
|
|
7. hand it to ``load_to_device``.
|
|
|
|
Steps 1-3 before step 4 are what make step 4 mean anything: a key built from
|
|
the *requested* attention mode would differ from the key the same pack already
|
|
holds under its *resolved* one, so a sage -> sdpa fallback would evict a
|
|
perfectly good model and reload it (or evict nothing and hold two copies of the
|
|
same weights). Step 4 before step 5 is plain eviction discipline -- on a
|
|
machine that can hold either of two packs but never both, freeing before
|
|
allocating is the difference between working and a CUDA OOM.
|
|
|
|
**``client_id`` is not part of the cache key.** It names the websocket progress
|
|
events go to, not the model. It is stamped onto the patcher per call, so a
|
|
second browser tab asking for a loaded pack reuses it *and* gets its own id for
|
|
progress. VibeVoice's ``generation.py:78`` stomped the id onto a cached
|
|
patcher at construction time and lost the second client's id entirely.
|
|
|
|
**Interrupt handling.** ``InterruptProcessingException`` subclasses
|
|
``BaseException``, not ``Exception``. :func:`generate_audio` names it *before*
|
|
the broad ``except Exception`` that logs and re-raises, so a cancelled run
|
|
always cancels. The reverse order would let a cancellation be reported as a
|
|
generation failure -- and, in the V1 nodes, as one second of silence.
|
|
|
|
Borrowed from VibeVoice's ``modules/generation.py``; see the plan's §3.1.8 for
|
|
the ``client_id`` correction.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import gc
|
|
import logging
|
|
import os
|
|
from typing import Any, Optional
|
|
|
|
import torch
|
|
|
|
from . import model_info
|
|
from .attention_utils import get_attn_implementation_for_load, resolve_attention_mode
|
|
from .audio_utils import (
|
|
S3GEN_SAMPLE_RATE,
|
|
S3_TOKENIZER_SAMPLE_RATE,
|
|
downmix_to_mono,
|
|
interrupt_exception_class,
|
|
preprocess_comfy_audio,
|
|
set_seed,
|
|
)
|
|
from .device_utils import (
|
|
DEVICE_CPU,
|
|
_get_model_management,
|
|
get_offload_device,
|
|
get_torch_device,
|
|
)
|
|
from .dtype_utils import DTYPE_AUTO, get_dtype_str, resolve_dtype
|
|
from .loader import ChatterboxLoader, ChatterboxModelHandler
|
|
from .model_registry import FAMILY_TTS, evict_if_changed
|
|
from .text_processing import plan_text
|
|
from .patcher import (
|
|
CHATTERBOX_PATCHER_CACHE as _SHARED_PATCHER_CACHE,
|
|
)
|
|
from .patcher import (
|
|
ChatterboxPatcher,
|
|
load_to_device,
|
|
select_patcher_class,
|
|
)
|
|
|
|
__all__ = [
|
|
"CHATTERBOX_PATCHER_CACHE",
|
|
"FAMILY_TTS",
|
|
"FAMILY_MTL",
|
|
"FAMILY_TURBO",
|
|
"REFINEMENT_CACHE_DEPTH",
|
|
"InterruptProgressBar",
|
|
"resolve_generation_family",
|
|
"resolve_attn_implementation",
|
|
"build_cache_key",
|
|
"variant_cache_stem",
|
|
"load_chatterbox_model",
|
|
"load_chatterbox_entry_point",
|
|
"generate_audio",
|
|
"generate_mtl_audio",
|
|
"generate_turbo_audio",
|
|
"convert_voice_audio",
|
|
"iterative_convert_voice",
|
|
"RefinementCache",
|
|
"split_audio_for_conversion",
|
|
"combine_audio_chunks",
|
|
"force_offload_model",
|
|
]
|
|
|
|
|
|
# The one cache, re-exported rather than re-created: two dicts means every
|
|
# eviction sweeps one of them and leaves the other holding the weights it just
|
|
# claimed to free.
|
|
CHATTERBOX_PATCHER_CACHE: dict[str, ChatterboxPatcher] = _SHARED_PATCHER_CACHE
|
|
|
|
#: Model-registry families. One per eviction class: a multilingual pack and a
|
|
#: Turbo pack are as mutually exclusive in VRAM as two classic packs are.
|
|
FAMILY_MTL = "mtl"
|
|
FAMILY_TURBO = "turbo"
|
|
FAMILY_FLASH = "flash"
|
|
|
|
#: How many intermediate iterations :class:`RefinementCache` keeps. The reference
|
|
#: implementation caps at 5 and its own tooltip says "recommended max 5 passes --
|
|
#: more causes distortions", so holding more would be memory nobody benefits from.
|
|
REFINEMENT_CACHE_DEPTH = 5
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Progress
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class InterruptProgressBar:
|
|
"""A ComfyUI progress bar that polls for cancellation on every tick.
|
|
|
|
Chatterbox's sampling loops call ``pbar.update(1)`` -- thousands of times
|
|
over a CPU run -- and the vendored code has no interrupt check of its own.
|
|
This is therefore the hook point for "the user pressed stop": the poll
|
|
raises ComfyUI's interrupt, which :func:`generate_audio` lets through.
|
|
|
|
Args:
|
|
total: The widget's token (or timestep) budget.
|
|
inner: A ``comfy.utils.ProgressBar`` to drive, or ``None`` to run
|
|
without one.
|
|
"""
|
|
|
|
def __init__(self, total: int, inner: Optional[Any] = None) -> None:
|
|
self.total = int(total)
|
|
self._inner = inner
|
|
self._current = 0
|
|
self._model_management = _get_model_management()
|
|
|
|
def update(self, value: int = 1) -> None:
|
|
"""Advance by ``value`` steps. Raises on a cancelled run."""
|
|
self._current += int(value)
|
|
self._model_management.throw_exception_if_processing_interrupted()
|
|
if self._inner is not None:
|
|
self._inner.update(value)
|
|
|
|
def update_absolute(self, value: int, total: Optional[int] = None) -> None:
|
|
"""Jump to ``value``. Raises on a cancelled run."""
|
|
if total is not None:
|
|
self.total = int(total)
|
|
self._current = int(value)
|
|
self._model_management.throw_exception_if_processing_interrupted()
|
|
if self._inner is not None:
|
|
self._inner.update_absolute(value, total=total)
|
|
|
|
def finish(self) -> None:
|
|
"""Close the bar out. Never polls, so it is safe from a ``finally``.
|
|
|
|
Polling here would re-raise the interrupt that the ``try`` block just
|
|
propagated, replacing a clean cancellation with one raised from a
|
|
cleanup path.
|
|
"""
|
|
if self._inner is not None:
|
|
try:
|
|
self._inner.update_absolute(self.total, total=self.total)
|
|
except Exception as exc: # pragma: no cover - cosmetic
|
|
logging.debug("[Chatterbox TTS] Could not close the progress bar: %s", exc)
|
|
|
|
|
|
def _build_progress_bar(total: int) -> InterruptProgressBar:
|
|
"""An :class:`InterruptProgressBar`, with a core ``ProgressBar`` if available."""
|
|
inner = None
|
|
try:
|
|
from comfy.utils import ProgressBar
|
|
|
|
inner = ProgressBar(total)
|
|
except Exception as exc: # pragma: no cover - core is present under ComfyUI
|
|
logging.debug("[Chatterbox TTS] ProgressBar unavailable: %s", exc)
|
|
return InterruptProgressBar(total, inner)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Routing helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def resolve_attn_implementation(attention_mode: str) -> str:
|
|
"""The ``attn_implementation`` string a pack would be built with.
|
|
|
|
``sage`` loads as ``sdpa``: transformers has no sage implementation to ask
|
|
for, and the kernel is patched in afterwards (see
|
|
:func:`modules.attention_utils.get_attn_implementation_for_load`).
|
|
"""
|
|
return get_attn_implementation_for_load(attention_mode)
|
|
|
|
|
|
def resolve_generation_family(
|
|
model_name: str,
|
|
external_model: Optional[dict] = None,
|
|
) -> str:
|
|
"""Which generation family a pack belongs to.
|
|
|
|
Args:
|
|
model_name: Pack name from the dropdown.
|
|
external_model: An externally-loaded pack descriptor, when the caller
|
|
has one. Wins over ``model_name``.
|
|
|
|
Returns:
|
|
``"tts"`` for a synthesis pack, ``"vc"`` for a conversion-only one.
|
|
|
|
Raises:
|
|
ValueError: No model was named, or the external descriptor declares a
|
|
family this node pack does not implement.
|
|
"""
|
|
if external_model is not None:
|
|
family = external_model.get("family") or external_model.get("type")
|
|
if family not in ("tts", "vc"):
|
|
raise ValueError(
|
|
f"An external Chatterbox model declares family '{family}'; this "
|
|
"node pack implements 'tts' and 'vc' only."
|
|
)
|
|
return str(family)
|
|
|
|
if not model_name:
|
|
raise ValueError(
|
|
"No Chatterbox model was selected. Pick a voice pack from the model "
|
|
"dropdown."
|
|
)
|
|
return "tts"
|
|
|
|
|
|
def build_cache_key(
|
|
model_name: str,
|
|
attention_mode: str,
|
|
dtype_str: str,
|
|
device: Optional[torch.device],
|
|
) -> str:
|
|
"""The process-wide identity of one loaded pack.
|
|
|
|
Built from the **resolved** attention mode and dtype, never from the
|
|
requested ones, so a workflow asking for ``sage`` on a machine that falls
|
|
back to ``sdpa`` shares the cache entry of a workflow that asked for
|
|
``sdpa`` directly instead of loading the same weights twice under two keys.
|
|
|
|
``client_id`` is deliberately absent (see the module docstring).
|
|
"""
|
|
device_key = str(device) if device is not None else "auto"
|
|
return f"{model_name}_attn_{attention_mode}_dtype_{dtype_str}_dev_{device_key}"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Loading
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _resolve_turbo_variant(model_name: str) -> Optional[str]:
|
|
"""``"turbo"``/``"nano"`` when ``model_name`` names a Turbo-family pack.
|
|
|
|
The answer comes from the registry entry when one exists -- an official
|
|
Turbo-family entry, or a local one the scan registered per variant -- and
|
|
from the pack directory's t3 checkpoint otherwise, for community folders
|
|
the scan has not seen. Registry first is what makes the shared-family
|
|
folder work: a directory holding both t3 checkpoints cannot answer "which
|
|
member" from its contents, but each dropdown entry names its own.
|
|
"""
|
|
try:
|
|
pack_dir = ChatterboxLoader.resolve_pack_dir(model_name)
|
|
except ValueError:
|
|
return None
|
|
|
|
entry = model_info.AVAILABLE_CHATTERBOX_MODELS.get(model_name) or {}
|
|
architecture = entry.get("architecture")
|
|
if architecture not in ("chatterbox_turbo", "chatterbox_nano") and os.path.isdir(pack_dir):
|
|
architecture = model_info.infer_pack_architecture(pack_dir)
|
|
if architecture == "chatterbox_nano":
|
|
return "nano"
|
|
if architecture == "chatterbox_turbo":
|
|
return "turbo"
|
|
return None
|
|
|
|
|
|
def _resolve_flash_pack(model_name: str) -> bool:
|
|
"""True when ``model_name`` names a Chatterbox Flash pack.
|
|
|
|
Same precedence as the Turbo helper: the registry entry's architecture
|
|
answers first (official entry, or the local one the scan registered),
|
|
the pack directory's t3 checkpoint second -- a flash dir is only ever its
|
|
own pack, so contents alone are unambiguous.
|
|
"""
|
|
try:
|
|
pack_dir = ChatterboxLoader.resolve_pack_dir(model_name)
|
|
except ValueError:
|
|
return False
|
|
|
|
entry = model_info.AVAILABLE_CHATTERBOX_MODELS.get(model_name) or {}
|
|
if entry.get("architecture") == "chatterbox_flash":
|
|
return True
|
|
return os.path.isdir(pack_dir) and (
|
|
model_info.infer_pack_architecture(pack_dir) == "chatterbox_flash"
|
|
)
|
|
|
|
|
|
def load_chatterbox_model(
|
|
model_name: str,
|
|
device: str = "auto",
|
|
dtype: str = DTYPE_AUTO,
|
|
attention_mode: str = "sdpa",
|
|
force_reload: bool = False,
|
|
client_id: Optional[str] = None,
|
|
) -> tuple[ChatterboxPatcher, Any, Optional[Any]]:
|
|
"""Load, or reuse, the pack ``model_name``.
|
|
|
|
Args:
|
|
model_name: Pack name from the model dropdown.
|
|
device: Device option value (see
|
|
:func:`modules.device_utils.get_device_options`).
|
|
dtype: Dtype widget value; ``"auto"`` lets ComfyUI decide.
|
|
attention_mode: Attention widget value, resolved against this machine.
|
|
force_reload: Build a fresh pack even when one is cached.
|
|
client_id: Websocket id for progress. Stamped onto the patcher per call;
|
|
**not** part of the cache key.
|
|
|
|
Returns:
|
|
``(patcher, tts_model, vc_model)``.
|
|
|
|
Raises:
|
|
RuntimeError: The pack could not be loaded.
|
|
"""
|
|
# 1. Resolve the attention mode -- before the cache key. A key built from
|
|
# the unresolved request would not match the key this same pack already
|
|
# holds under its resolved mode.
|
|
actual_attention_mode = resolve_attention_mode(attention_mode, False)
|
|
|
|
# 2. Resolve devices and dtype.
|
|
if device == DEVICE_CPU:
|
|
load_device = torch.device(DEVICE_CPU)
|
|
offload_device = torch.device(DEVICE_CPU)
|
|
else:
|
|
load_device = get_torch_device(device)
|
|
offload_device = get_offload_device()
|
|
target_dtype = resolve_dtype(dtype, load_device)
|
|
|
|
# 2b. Turbo-family packs (Turbo and its Nano member) load through their
|
|
# own entry point: no VC half, no classic checkpoint checks, and a
|
|
# cache entry of their own. The TTS node routes here so one node
|
|
# serves the whole family. The entry point's own eviction sweep
|
|
# clears the classic pack this node may still hold, the same way the
|
|
# classic path below clears a Turbo one.
|
|
if _resolve_flash_pack(model_name):
|
|
patcher, tts_model = load_chatterbox_entry_point(
|
|
model_name,
|
|
entry_point="flash",
|
|
variant="flash",
|
|
family=FAMILY_FLASH,
|
|
device=device,
|
|
dtype=dtype,
|
|
attention_mode=attention_mode,
|
|
force_reload=force_reload,
|
|
client_id=client_id,
|
|
)
|
|
return patcher, tts_model, None
|
|
|
|
turbo_variant = _resolve_turbo_variant(model_name)
|
|
if turbo_variant is not None:
|
|
patcher, tts_model = load_chatterbox_entry_point(
|
|
model_name,
|
|
entry_point="turbo",
|
|
variant=turbo_variant,
|
|
family=FAMILY_TURBO,
|
|
device=device,
|
|
dtype=dtype,
|
|
attention_mode=attention_mode,
|
|
force_reload=force_reload,
|
|
client_id=client_id,
|
|
)
|
|
return patcher, tts_model, None
|
|
|
|
# 3. The cache key.
|
|
cache_key = build_cache_key(
|
|
model_name, actual_attention_mode, get_dtype_str(target_dtype), load_device
|
|
)
|
|
|
|
# 4. Evict the previous pack -- BEFORE any allocation below. On a machine
|
|
# that holds either pack but not both, this ordering is the difference
|
|
# between a working load and a CUDA OOM.
|
|
evict_if_changed(FAMILY_TTS, cache_key, (CHATTERBOX_PATCHER_CACHE,))
|
|
|
|
# 5-6. Build the handler and the patcher.
|
|
if cache_key not in CHATTERBOX_PATCHER_CACHE or force_reload:
|
|
pack_dir = ChatterboxLoader.resolve_pack_dir(model_name)
|
|
model_handler = ChatterboxModelHandler(
|
|
model_name,
|
|
attention_mode=actual_attention_mode,
|
|
dtype_str=dtype,
|
|
model_path=pack_dir,
|
|
)
|
|
model_handler.cache_key = cache_key
|
|
|
|
patcher_cls = select_patcher_class(
|
|
None, load_device, legacy_cls=ChatterboxPatcher
|
|
)
|
|
patcher = patcher_cls(
|
|
model_handler,
|
|
attention_mode=actual_attention_mode,
|
|
load_device=load_device,
|
|
offload_device=offload_device,
|
|
size=model_handler.size,
|
|
dtype=target_dtype,
|
|
)
|
|
CHATTERBOX_PATCHER_CACHE[cache_key] = patcher
|
|
logging.debug(
|
|
"[Chatterbox TTS] Created a patcher for '%s' (attn=%s, dtype=%s, "
|
|
"estimated %d bytes)",
|
|
model_name,
|
|
actual_attention_mode,
|
|
get_dtype_str(target_dtype),
|
|
model_handler.size,
|
|
)
|
|
|
|
patcher = CHATTERBOX_PATCHER_CACHE[cache_key]
|
|
|
|
# Per-call, never at construction: a cached patcher must carry the *current*
|
|
# caller's id, not the first one's.
|
|
patcher.client_id = client_id
|
|
|
|
# 7. Place it.
|
|
load_to_device(patcher)
|
|
|
|
tts_model = getattr(patcher.model, "tts_model", None)
|
|
vc_model = getattr(patcher.model, "vc_model", None)
|
|
if tts_model is None and vc_model is None:
|
|
raise RuntimeError(
|
|
f"The Chatterbox voice pack '{model_name}' could not be loaded. "
|
|
"Check the log above for the download or build failure."
|
|
)
|
|
|
|
return patcher, tts_model, vc_model
|
|
|
|
|
|
def force_offload_model(
|
|
patcher: ChatterboxPatcher,
|
|
model_name: str,
|
|
warm: bool = False,
|
|
) -> None:
|
|
"""Take a pack out of VRAM on demand.
|
|
|
|
Args:
|
|
patcher: The live patcher.
|
|
model_name: Pack name, for the log line.
|
|
warm: Keep the tensors on the intermediate device for a fast reload
|
|
instead of dropping them.
|
|
"""
|
|
logging.info(
|
|
"[Chatterbox TTS] Force offloading voice pack '%s'%s...",
|
|
model_name,
|
|
" (warm)" if warm else "",
|
|
)
|
|
if patcher.is_loaded:
|
|
if warm:
|
|
patcher.unpatch_model(unpatch_weights=True, warm=True)
|
|
else:
|
|
patcher.unpatch_model(unpatch_weights=True, destroy=True)
|
|
|
|
model_management = _get_model_management()
|
|
try:
|
|
model_management.unload_all_models()
|
|
except Exception as exc:
|
|
logging.debug("[Chatterbox TTS] unload_all_models failed: %s", exc)
|
|
gc.collect()
|
|
try:
|
|
model_management.soft_empty_cache()
|
|
except Exception as exc:
|
|
logging.debug("[Chatterbox TTS] soft_empty_cache failed: %s", exc)
|
|
logging.info("[Chatterbox TTS] Force offload of '%s' complete.", model_name)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Synthesis
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
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.
|
|
|
|
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.
|
|
"""
|
|
if not text or not text.strip():
|
|
raise ValueError(
|
|
"Chatterbox needs some text to speak. The text input was empty."
|
|
)
|
|
|
|
# Flash dispatch, before the classic progress bar: a block-diffusion
|
|
# decoder with no token budget to count.
|
|
if type(model).__name__ == "ChatterboxFlashTTS":
|
|
return generate_flash_audio(
|
|
model,
|
|
text,
|
|
audio_prompt=audio_prompt,
|
|
exaggeration=exaggeration,
|
|
temperature=temperature,
|
|
seed=seed,
|
|
num_steps=num_steps,
|
|
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)
|
|
return generate_turbo_audio(
|
|
model,
|
|
plan.text,
|
|
audio_prompt=audio_prompt,
|
|
temperature=temperature,
|
|
top_k=top_k,
|
|
top_p=top_p,
|
|
repetition_penalty=repetition_penalty,
|
|
exaggeration=exaggeration,
|
|
norm_loudness=norm_loudness,
|
|
seed=seed,
|
|
use_watermark=use_watermark,
|
|
)
|
|
|
|
interrupt_error = interrupt_exception_class()
|
|
set_seed(seed)
|
|
progress_bar = _build_progress_bar(max_new_tokens)
|
|
|
|
try:
|
|
# The voice prompt is handed over in memory: preprocess_comfy_audio
|
|
# downmixes, scrubs and resamples it to the 24 kHz the model
|
|
# conditions on, and the vendored loader takes a waveform as-is
|
|
# (``load_audio_at``, ``src/chatterbox/models/dsp.py``). ``None`` when
|
|
# no clip is connected, which the vendor answers with the pack's
|
|
# built-in voice. Inside the ``try`` so a malformed clip still closes
|
|
# the progress bar, the way the temp-file write it replaces did.
|
|
prompt = preprocess_comfy_audio(audio_prompt, target_sr=S3GEN_SAMPLE_RATE)
|
|
try:
|
|
with torch.no_grad():
|
|
waveform = model.generate(
|
|
text,
|
|
audio_prompt_path=prompt,
|
|
exaggeration=exaggeration,
|
|
temperature=temperature,
|
|
cfg_weight=cfg_weight,
|
|
repetition_penalty=repetition_penalty,
|
|
min_p=min_p,
|
|
top_p=top_p,
|
|
pbar=progress_bar,
|
|
max_new_tokens=max_new_tokens,
|
|
flow_cfg_scale=flow_cfg_scale,
|
|
watermark=bool(use_watermark),
|
|
)
|
|
except interrupt_error:
|
|
# FIRST, and before `except Exception`. The interrupt derives
|
|
# from BaseException, so this clause is unreachable from the
|
|
# generic handler -- and a cancelled run is not a failed run:
|
|
# re-raising here is what actually cancels it.
|
|
logging.info(
|
|
"[Chatterbox TTS] Generation was interrupted; cancelling the run."
|
|
)
|
|
raise
|
|
except Exception:
|
|
logging.exception(
|
|
"[Chatterbox TTS] Generation failed for a %d-character text.",
|
|
len(text),
|
|
)
|
|
raise
|
|
finally:
|
|
progress_bar.finish()
|
|
|
|
if waveform is None or waveform.numel() == 0:
|
|
raise ValueError(
|
|
"Chatterbox produced no audio for this text. Shorten it, raise "
|
|
"'max_new_tokens', or pick a different voice pack."
|
|
)
|
|
|
|
result = waveform.detach().to(device="cpu", dtype=torch.float32)
|
|
if result.dim() == 1:
|
|
result = result.unsqueeze(0)
|
|
|
|
logging.info(
|
|
"[Chatterbox TTS] Generated %.2f s of audio at %d Hz.",
|
|
result.shape[-1] / float(S3GEN_SAMPLE_RATE),
|
|
S3GEN_SAMPLE_RATE,
|
|
)
|
|
return result, S3GEN_SAMPLE_RATE
|
|
|
|
|
|
def convert_voice_audio(
|
|
model: Any,
|
|
audio: dict,
|
|
target_voice_audio: Optional[dict] = None,
|
|
n_timesteps: int = 10,
|
|
temperature: float = 1.0,
|
|
flow_cfg_scale: float = 0.7,
|
|
use_watermark: bool = False,
|
|
) -> tuple[torch.Tensor, int]:
|
|
"""Convert ``audio`` into the voice described by ``target_voice_audio``.
|
|
|
|
Args:
|
|
model: The ``ChatterboxVC`` instance from :func:`load_chatterbox_model`.
|
|
audio: The source speech, as a ComfyUI ``AUDIO`` dict.
|
|
target_voice_audio: The target voice's timbre. ``None`` uses the pack's
|
|
built-in reference (``conds.pt``).
|
|
n_timesteps: Diffusion steps.
|
|
temperature: Temperature of the initial latent noise.
|
|
flow_cfg_scale: Flow-matching CFG scale.
|
|
use_watermark: Whether to apply the PERTH watermark.
|
|
|
|
Returns:
|
|
``(waveform[1, T] float32 on CPU, sample_rate)``.
|
|
|
|
Raises:
|
|
ValueError: The source audio is empty, or the model returned nothing.
|
|
"""
|
|
if not audio or audio.get("waveform") is None or audio["waveform"].numel() == 0:
|
|
raise ValueError(
|
|
"Chatterbox voice conversion needs source audio. Connect an AUDIO "
|
|
"input carrying the speech to convert."
|
|
)
|
|
|
|
interrupt_error = interrupt_exception_class()
|
|
progress_bar = _build_progress_bar(n_timesteps)
|
|
|
|
try:
|
|
# Two rates, on purpose. The vendored ``generate`` loads the *source*
|
|
# at the tokenizer's 16 kHz (``load_audio_at(audio, S3_SR)``,
|
|
# ``vc.py``), and the *target* voice through ``set_target_voice`` at
|
|
# S3Gen's 24 kHz -- the only dual-rate call site in the tree.
|
|
# ``target`` is ``None`` when no target clip is connected, which the
|
|
# vendor answers with the pack's own ``conds.pt``. Both are prepped
|
|
# inside the ``try`` so the progress bar is closed on a malformed
|
|
# clip, the way the temp-file writes they replace were.
|
|
source = preprocess_comfy_audio(
|
|
audio, target_sr=S3_TOKENIZER_SAMPLE_RATE, name="source audio"
|
|
)
|
|
target = preprocess_comfy_audio(
|
|
target_voice_audio, target_sr=S3GEN_SAMPLE_RATE, name="target voice"
|
|
)
|
|
try:
|
|
with torch.no_grad():
|
|
waveform = model.generate(
|
|
source,
|
|
target_voice_path=target,
|
|
n_timesteps=n_timesteps,
|
|
pbar=progress_bar,
|
|
temperature=temperature,
|
|
flow_cfg_scale=flow_cfg_scale,
|
|
watermark=bool(use_watermark),
|
|
)
|
|
except interrupt_error:
|
|
logging.info(
|
|
"[Chatterbox TTS] Voice conversion was interrupted; "
|
|
"cancelling the run."
|
|
)
|
|
raise
|
|
except Exception:
|
|
logging.exception("[Chatterbox TTS] Voice conversion failed.")
|
|
raise
|
|
finally:
|
|
progress_bar.finish()
|
|
|
|
if waveform is None or waveform.numel() == 0:
|
|
raise ValueError(
|
|
"Chatterbox voice conversion produced no audio. Check that the "
|
|
"source clip actually contains speech."
|
|
)
|
|
|
|
result = waveform.detach().to(device="cpu", dtype=torch.float32)
|
|
if result.dim() == 1:
|
|
result = result.unsqueeze(0)
|
|
return result, S3GEN_SAMPLE_RATE
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Other entry points: multilingual and Turbo (S8.1 / S8.2)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def variant_cache_stem(model_name: str, variant: str) -> str:
|
|
"""The cache-key stem for one *variant* of an entry point.
|
|
|
|
Two collisions this exists to prevent:
|
|
|
|
* **Turbo against Nano** -- the same module, the same directory layout, two
|
|
disjoint weight sets. Without the variant in the key the second load is
|
|
served the first one's weights, and the user gets the wrong model with no
|
|
error anywhere.
|
|
* **multilingual v2 against v3** -- one directory, three T3 checkpoints.
|
|
|
|
``build_cache_key`` keeps its four-argument signature on purpose (a test
|
|
pins it), so the variant is folded into the *name* rather than passed as a
|
|
fifth key input.
|
|
"""
|
|
if not variant:
|
|
return model_name
|
|
return f"{model_name}__{variant}"
|
|
|
|
|
|
def load_chatterbox_entry_point(
|
|
model_name: str,
|
|
*,
|
|
entry_point: str,
|
|
variant: str,
|
|
family: str,
|
|
device: str = "auto",
|
|
dtype: str = DTYPE_AUTO,
|
|
attention_mode: str = "sdpa",
|
|
force_reload: bool = False,
|
|
client_id: Optional[str] = None,
|
|
) -> tuple[ChatterboxPatcher, Any]:
|
|
"""Load, or reuse, one multilingual / Turbo pack.
|
|
|
|
:func:`load_chatterbox_model` states its seven ordering steps in prose and
|
|
then performs them. They are not restated here: the three entry points differ
|
|
only in which loader runs at step 5 and which family is evicted at step 4,
|
|
so this function performs the same seven steps against the other two.
|
|
|
|
Args:
|
|
model_name: Pack name from the dropdown.
|
|
entry_point: ``"multilingual"`` or ``"turbo"``.
|
|
variant: The resolved T3 checkpoint filename, or ``"nano"`` / ``"turbo"``.
|
|
family: Registry family to evict against (:data:`FAMILY_MTL` /
|
|
:data:`FAMILY_TURBO`).
|
|
device: Device option value.
|
|
dtype: Dtype widget value.
|
|
attention_mode: Attention widget value, resolved against this machine.
|
|
force_reload: Build a fresh pack even when one is cached.
|
|
client_id: Websocket id for progress; stamped per call, never keyed.
|
|
|
|
Returns:
|
|
``(patcher, model)``.
|
|
|
|
Raises:
|
|
RuntimeError: The pack could not be loaded.
|
|
"""
|
|
actual_attention_mode = resolve_attention_mode(attention_mode, False)
|
|
|
|
if device == DEVICE_CPU:
|
|
load_device = torch.device(DEVICE_CPU)
|
|
offload_device = torch.device(DEVICE_CPU)
|
|
else:
|
|
load_device = get_torch_device(device)
|
|
offload_device = get_offload_device()
|
|
target_dtype = resolve_dtype(dtype, load_device)
|
|
|
|
stem = variant_cache_stem(model_name, variant)
|
|
cache_key = build_cache_key(
|
|
stem, actual_attention_mode, get_dtype_str(target_dtype), load_device
|
|
)
|
|
evict_if_changed(family, cache_key, (CHATTERBOX_PATCHER_CACHE,))
|
|
|
|
if cache_key not in CHATTERBOX_PATCHER_CACHE or force_reload:
|
|
model_handler = ChatterboxModelHandler(
|
|
model_name,
|
|
attention_mode=actual_attention_mode,
|
|
dtype_str=dtype,
|
|
# The pack directory is the user's selection; the variant-folded
|
|
# stem is a cache key and must never name a folder. Deriving the
|
|
# path from the stem sent a Turbo download into 'X__nano' next to
|
|
# the pack the user actually picked.
|
|
model_path=ChatterboxLoader.resolve_pack_dir(model_name),
|
|
load_vc=False,
|
|
entry_point=entry_point,
|
|
variant=variant,
|
|
)
|
|
model_handler.cache_key = cache_key
|
|
|
|
patcher_cls = select_patcher_class(
|
|
None, load_device, legacy_cls=ChatterboxPatcher
|
|
)
|
|
patcher = patcher_cls(
|
|
model_handler,
|
|
attention_mode=actual_attention_mode,
|
|
load_device=load_device,
|
|
offload_device=offload_device,
|
|
size=model_handler.size,
|
|
dtype=target_dtype,
|
|
)
|
|
CHATTERBOX_PATCHER_CACHE[cache_key] = patcher
|
|
|
|
patcher = CHATTERBOX_PATCHER_CACHE[cache_key]
|
|
patcher.client_id = client_id
|
|
load_to_device(patcher)
|
|
|
|
model = getattr(patcher.model, "tts_model", None)
|
|
if model is None:
|
|
raise RuntimeError(
|
|
f"The Chatterbox '{entry_point}' pack '{model_name}' could not be "
|
|
"loaded. Check the log above for the download or build failure."
|
|
)
|
|
return patcher, model
|
|
|
|
|
|
def _normalise_waveform(waveform: Any) -> torch.Tensor:
|
|
"""``[1, T]`` float32 on CPU, whatever shape the model returned.
|
|
|
|
Raises:
|
|
ValueError: The model returned nothing, or an empty tensor. Both mean
|
|
the same thing to a user -- the text was not spoken -- and both are
|
|
worth a message rather than a downstream shape error.
|
|
"""
|
|
if waveform is None:
|
|
raise ValueError(
|
|
"Chatterbox produced no audio for this text. Shorten it, raise "
|
|
"'max_new_tokens', or pick a different voice pack."
|
|
)
|
|
result = waveform.detach().to(device="cpu", dtype=torch.float32)
|
|
if result.dim() == 1:
|
|
result = result.unsqueeze(0)
|
|
if result.numel() == 0:
|
|
raise ValueError(
|
|
"Chatterbox produced no audio for this text. Shorten it, raise "
|
|
"'max_new_tokens', or pick a different voice pack."
|
|
)
|
|
return result
|
|
|
|
|
|
def generate_mtl_audio(
|
|
model: Any,
|
|
text: str,
|
|
language_id: 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,
|
|
use_watermark: bool = False,
|
|
) -> tuple[torch.Tensor, int]:
|
|
"""Synthesise ``text`` with a loaded ``ChatterboxMultilingualTTS``.
|
|
|
|
**The language id is passed as a keyword, never spliced into ``text``.** The
|
|
vendored tokenizer prepends it itself (``MTLTokenizer.text_to_tokens`` does
|
|
``clean_text + lang_id`` before encoding), so a node that also inserted it
|
|
would produce ``<de><de>Hallo`` -- a doubled prefix that is not a valid token
|
|
sequence, and which fails deep inside the T3 with a message about tokens
|
|
rather than about the text box.
|
|
|
|
Args:
|
|
model: The ``ChatterboxMultilingualTTS`` instance.
|
|
text: The text to speak, already crash-protected and tag-resolved.
|
|
language_id: One of the 23 ISO codes in
|
|
:data:`modules.constants.MTL_LANGUAGES`.
|
|
audio_prompt: Optional ComfyUI ``AUDIO`` dict to clone the voice from.
|
|
exaggeration: Expressiveness. Inert on v3 weights -- see the node tooltip.
|
|
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.
|
|
use_watermark: Whether to apply the PERTH watermark.
|
|
|
|
Returns:
|
|
``(waveform[1, T] float32 on CPU, sample_rate)``.
|
|
"""
|
|
if not text or not text.strip():
|
|
raise ValueError(
|
|
"Chatterbox needs some text to speak. The text input was empty."
|
|
)
|
|
# Upstream would hit `assert self.conds is not None` deep inside its
|
|
# generate() -- a bare AssertionError that names no fix. Five of the six
|
|
# Single Language Packs ship no conds.pt, so this is the common case for
|
|
# them, not an edge case.
|
|
if audio_prompt is None and getattr(model, "conds", None) is None:
|
|
raise ValueError(
|
|
"This Chatterbox pack ships no built-in voice -- connect a "
|
|
"reference audio clip to the node's audio prompt input."
|
|
)
|
|
|
|
set_seed(seed)
|
|
prompt = preprocess_comfy_audio(audio_prompt, target_sr=S3GEN_SAMPLE_RATE)
|
|
with torch.no_grad():
|
|
waveform = model.generate(
|
|
text,
|
|
language_id=language_id,
|
|
audio_prompt_path=prompt,
|
|
exaggeration=exaggeration,
|
|
temperature=temperature,
|
|
cfg_weight=cfg_weight,
|
|
repetition_penalty=repetition_penalty,
|
|
min_p=min_p,
|
|
top_p=top_p,
|
|
watermark=bool(use_watermark),
|
|
)
|
|
return _normalise_waveform(waveform), S3GEN_SAMPLE_RATE
|
|
|
|
|
|
def generate_turbo_audio(
|
|
model: Any,
|
|
text: str,
|
|
audio_prompt: Optional[dict] = None,
|
|
temperature: float = 0.8,
|
|
top_k: int = 1000,
|
|
top_p: float = 0.95,
|
|
repetition_penalty: float = 1.2,
|
|
exaggeration: float = 0.0,
|
|
norm_loudness: bool = True,
|
|
seed: int = 0,
|
|
use_watermark: bool = False,
|
|
) -> tuple[torch.Tensor, int]:
|
|
"""Synthesise ``text`` with a loaded ``ChatterboxTurboTTS`` (Turbo or Nano).
|
|
|
|
Turbo's sampler is a top-k / top-p pair with no CFG: ``cfg_weight`` and
|
|
``min_p`` do not exist on this family, and upstream logs a warning whenever
|
|
either is non-zero. This function therefore does not accept them at all --
|
|
offering a widget the model ignores is how a user spends an afternoon on a
|
|
knob with no effect.
|
|
|
|
Args:
|
|
model: The ``ChatterboxTurboTTS`` instance.
|
|
text: The text to speak, with Turbo's square-bracket tags left in place.
|
|
audio_prompt: Optional ComfyUI ``AUDIO`` dict to clone the voice from.
|
|
temperature: Sampling temperature (0.05 - 2.0 on this family).
|
|
top_k: Candidate cut-off (1 - 5000).
|
|
top_p: Nucleus-sampling ceiling.
|
|
repetition_penalty: Token repetition penalty.
|
|
exaggeration: Reference-conditioning intensity. Turbo ignores it for
|
|
sampling and upstream warns when it is non-zero.
|
|
norm_loudness: Normalise the reference clip. Skipped with the vendored
|
|
warning when ``pyloudnorm`` is absent -- see ``src/PATCHES.md`` P-026.
|
|
seed: Seed; ``0`` means random.
|
|
use_watermark: Whether to apply the PERTH watermark.
|
|
|
|
Returns:
|
|
``(waveform[1, T] float32 on CPU, sample_rate)``.
|
|
"""
|
|
if not text or not text.strip():
|
|
raise ValueError(
|
|
"Chatterbox needs some text to speak. The text input was empty."
|
|
)
|
|
|
|
set_seed(seed)
|
|
prompt = preprocess_comfy_audio(audio_prompt, target_sr=S3GEN_SAMPLE_RATE)
|
|
with torch.no_grad():
|
|
waveform = model.generate(
|
|
text,
|
|
audio_prompt_path=prompt,
|
|
temperature=temperature,
|
|
top_k=int(top_k),
|
|
top_p=top_p,
|
|
repetition_penalty=repetition_penalty,
|
|
exaggeration=exaggeration,
|
|
norm_loudness=bool(norm_loudness),
|
|
watermark=bool(use_watermark),
|
|
)
|
|
return _normalise_waveform(waveform), S3GEN_SAMPLE_RATE
|
|
|
|
|
|
def generate_flash_audio(
|
|
model: Any,
|
|
text: str,
|
|
audio_prompt: Optional[dict] = None,
|
|
exaggeration: float = 0.5,
|
|
temperature: float = 0.8,
|
|
seed: int = 0,
|
|
num_steps: int = 10,
|
|
cfg_scale: float = 1.0,
|
|
) -> tuple[torch.Tensor, int]:
|
|
"""Synthesise ``text`` with a loaded ``ChatterboxFlashTTS``.
|
|
|
|
Flash is a block-diffusion decoder: there is no top-k / min-p surface and
|
|
no ``conds.pt`` ships in its repo -- the reference clip is the only voice
|
|
source, so a prompt is mandatory and conditioning is rebuilt per call
|
|
exactly like the classic entry point does. ``num_steps`` trades speed for
|
|
quality (upstream default 10); ``cfg_scale`` is the block-diffusion
|
|
guidance weight (upstream default 1.0, i.e. no guidance). The remaining
|
|
upstream kwargs (``backend``, ``time_shift_tau``, ...) stay at their
|
|
paper defaults; the engine backend can still be forced with the
|
|
``CHATTERBOX_FLASH_ENGINE`` environment variable.
|
|
"""
|
|
if audio_prompt is None:
|
|
raise ValueError(
|
|
"Chatterbox Flash ships no built-in voice -- connect a reference "
|
|
"audio clip to the node's audio prompt input."
|
|
)
|
|
|
|
set_seed(seed)
|
|
prompt = preprocess_comfy_audio(audio_prompt, target_sr=S3GEN_SAMPLE_RATE)
|
|
model.prepare_conditionals(prompt, exaggeration=exaggeration)
|
|
waveform = model.generate(
|
|
text,
|
|
exaggeration=exaggeration,
|
|
temperature=temperature,
|
|
num_steps=num_steps,
|
|
cfg_scale=cfg_scale,
|
|
)
|
|
return _normalise_waveform(waveform), S3GEN_SAMPLE_RATE
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Iterative and chunked voice conversion (S8.6)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class RefinementCache:
|
|
"""The intermediate iterations of one conversion, capped at five.
|
|
|
|
Iterative conversion is a fold: pass *n* is pass *n-1* run again. Asking for
|
|
six passes when five are already computed should cost one conversion, not
|
|
six; asking for three again should cost nothing.
|
|
|
|
The cap is not a memory optimisation, it is the documented behaviour: the
|
|
reference implementation keeps five and its own tooltip warns that more
|
|
passes distort the result, so retaining a sixth would be holding audio nobody
|
|
should be able to select.
|
|
|
|
Attributes:
|
|
depth: How many iterations are kept per job.
|
|
"""
|
|
|
|
def __init__(self, depth: int = REFINEMENT_CACHE_DEPTH) -> None:
|
|
self.depth = int(depth)
|
|
self._iterations: dict[str, list[dict]] = {}
|
|
|
|
def get(self, key: str) -> list[dict]:
|
|
"""The stored iterations for ``key`` (a copy, oldest first)."""
|
|
return list(self._iterations.get(key, ()))
|
|
|
|
def put(self, key: str, iterations: list[dict]) -> None:
|
|
"""Store ``iterations``, keeping only the newest :attr:`depth`."""
|
|
self._iterations[key] = list(iterations[-self.depth :])
|
|
|
|
def clear(self) -> None:
|
|
"""Forget everything. Exposed for tests and for an explicit reset."""
|
|
self._iterations.clear()
|
|
|
|
def iteration_counts(self) -> dict[str, int]:
|
|
"""How many iterations are held, per job.
|
|
|
|
The composed cache key is an implementation detail callers do not hold --
|
|
it folds in the audio shape and the sampling settings -- so this is how a
|
|
caller (or a test) checks that the cap is actually being applied rather
|
|
than assuming the key it passed in is the key used.
|
|
"""
|
|
return {key: len(value) for key, value in self._iterations.items()}
|
|
|
|
def __len__(self) -> int:
|
|
return len(self._iterations)
|
|
|
|
|
|
def _refinement_cache_key(
|
|
cache_key: Optional[str],
|
|
audio: Optional[dict],
|
|
n_timesteps: int,
|
|
temperature: float,
|
|
flow_cfg_scale: float,
|
|
) -> str:
|
|
"""A stable identity for one conversion job.
|
|
|
|
The source audio's *shape* stands in for its content. Hashing megabytes of
|
|
waveform on every run would cost more than the conversion the cache saves,
|
|
and the guarantee that matters -- within one run, a repeated request returns
|
|
the cached iteration -- does not depend on the bytes being identical.
|
|
"""
|
|
shape = "none"
|
|
if isinstance(audio, dict) and audio.get("waveform") is not None:
|
|
shape = str(tuple(audio["waveform"].shape))
|
|
return (
|
|
f"{cache_key or 'adhoc'}|{shape}|{int(n_timesteps)}|{temperature}|"
|
|
f"{flow_cfg_scale}"
|
|
)
|
|
|
|
|
|
def iterative_convert_voice(
|
|
model: Any,
|
|
audio: dict,
|
|
*,
|
|
target_voice_audio: Optional[dict] = None,
|
|
n_timesteps: int = 10,
|
|
temperature: float = 1.0,
|
|
flow_cfg_scale: float = 0.7,
|
|
use_watermark: bool = False,
|
|
refinement_passes: int = 1,
|
|
cache: Optional[RefinementCache] = None,
|
|
cache_key: Optional[str] = None,
|
|
single_pass: Optional[Any] = None,
|
|
) -> tuple[torch.Tensor, int]:
|
|
"""Convert ``audio`` into the target voice, optionally refining it N times.
|
|
|
|
Pass 1 is the ordinary single conversion. Pass *n* feeds pass *n-1*'s output
|
|
back in, so each pass moves the audio a little further toward the target
|
|
timbre -- and a little further from the source, which is why the widget stops
|
|
at 30 while its own tooltip recommends 5.
|
|
|
|
The passes are **deterministic**: nothing is re-seeded between them, each one
|
|
depends only on the previous one, and the cache key carries every input that
|
|
could change the result. That is what makes returning a cached iteration
|
|
safe rather than merely convenient.
|
|
|
|
The watermark is applied on the final pass only. Watermarking every pass and
|
|
then re-converting the watermarked audio is how a long run degrades into an
|
|
unintelligible buzz.
|
|
|
|
Args:
|
|
model: The ``ChatterboxVC`` instance.
|
|
audio: The source speech, as a ComfyUI ``AUDIO`` dict.
|
|
target_voice_audio: The target voice's timbre. ``None`` uses the pack's
|
|
built-in reference.
|
|
n_timesteps: Diffusion steps, applied to **every** pass.
|
|
temperature: Temperature of the initial latent noise.
|
|
flow_cfg_scale: Flow-matching CFG scale.
|
|
use_watermark: Whether to apply the PERTH watermark.
|
|
refinement_passes: How many passes to run. Clamped to at least 1.
|
|
cache: An optional :class:`RefinementCache` to reuse iterations in.
|
|
cache_key: Identity of this job within ``cache``.
|
|
single_pass: The one-pass conversion to call each time.
|
|
:func:`convert_voice_audio` unless a caller injects its own -- the
|
|
node passes the name it imported, so a test that substitutes one
|
|
conversion in the node module reaches the refinement path too
|
|
instead of silently bypassing it.
|
|
|
|
Returns:
|
|
``(waveform[1, T] float32 on CPU, sample_rate)``.
|
|
|
|
Raises:
|
|
ValueError: The source audio is empty, or every pass produced nothing.
|
|
"""
|
|
passes = max(1, int(refinement_passes))
|
|
convert = single_pass or convert_voice_audio
|
|
key = _refinement_cache_key(
|
|
cache_key, audio, n_timesteps, temperature, flow_cfg_scale
|
|
)
|
|
stored = cache.get(key) if cache is not None else []
|
|
|
|
current_audio = audio
|
|
result: Optional[torch.Tensor] = None
|
|
completed = 0
|
|
|
|
for iteration in stored:
|
|
result = iteration["audio"]
|
|
completed += 1
|
|
if completed >= passes:
|
|
break
|
|
current_audio = iteration["source"]
|
|
|
|
while completed < passes:
|
|
waveform, sample_rate = convert(
|
|
model,
|
|
current_audio,
|
|
target_voice_audio=target_voice_audio,
|
|
n_timesteps=n_timesteps,
|
|
temperature=temperature,
|
|
flow_cfg_scale=flow_cfg_scale,
|
|
use_watermark=use_watermark,
|
|
)
|
|
result = waveform
|
|
current_audio = {
|
|
"waveform": waveform.unsqueeze(0).to(dtype=torch.float32),
|
|
"sample_rate": sample_rate,
|
|
}
|
|
completed += 1
|
|
if cache is not None:
|
|
stored.append({"audio": result, "source": current_audio})
|
|
cache.put(key, stored)
|
|
|
|
if result is None:
|
|
raise ValueError(
|
|
"Chatterbox voice conversion produced no audio. Check that the "
|
|
"source clip actually contains speech."
|
|
)
|
|
return result, S3GEN_SAMPLE_RATE
|
|
|
|
|
|
def split_audio_for_conversion(
|
|
waveform: torch.Tensor,
|
|
sample_rate: int,
|
|
max_chunk_duration: float,
|
|
chunk_method: str = "smart",
|
|
search_window: float = 0.25,
|
|
frame_ms: int = 20,
|
|
) -> list[torch.Tensor]:
|
|
"""Cut ``waveform`` into chunks of at most ``max_chunk_duration`` seconds.
|
|
|
|
OOM protection: a long recording converts as one diffusion problem, and the
|
|
activation memory grows with the mel length. Cutting the source first bounds
|
|
it. ``0`` disables chunking and the clip comes back whole.
|
|
|
|
Args:
|
|
waveform: ``[C, T]`` or ``[B, C, T]`` audio.
|
|
sample_rate: Hz.
|
|
max_chunk_duration: The ceiling, in seconds. ``0`` disables.
|
|
chunk_method: ``"smart"`` cuts as close to silence as it can near the
|
|
ideal boundary -- a boundary in a pause is inaudible, a boundary
|
|
through a vowel is a click. ``"fixed"`` cuts at exact intervals,
|
|
which is cheaper and which the reference implementation documents as
|
|
making the joins more audible.
|
|
search_window: How far from the ideal boundary ``"smart"`` may look, as
|
|
a fraction of the chunk length.
|
|
frame_ms: Analysis frame size for the silence search.
|
|
|
|
Returns:
|
|
The chunks, in order. Always at least one; never empty.
|
|
"""
|
|
mono = downmix_to_mono(waveform)
|
|
total = int(mono.shape[-1])
|
|
|
|
if not max_chunk_duration or float(max_chunk_duration) <= 0:
|
|
return [mono]
|
|
limit = int(float(max_chunk_duration) * float(sample_rate))
|
|
if limit < 1 or total <= limit:
|
|
return [mono]
|
|
|
|
frame = max(1, int(sample_rate) * int(frame_ms) // 1000)
|
|
usable = (total // frame) * frame
|
|
flat = mono.reshape(-1)[:usable]
|
|
# Per-frame RMS, computed by hand: the mel spectrogram's own machinery is an
|
|
# order of magnitude more than choosing a cut location needs.
|
|
energy = (
|
|
flat.float().reshape(-1, frame).pow(2).mean(dim=1).sqrt()
|
|
if flat.numel()
|
|
else torch.zeros(1)
|
|
)
|
|
|
|
chunks: list[torch.Tensor] = []
|
|
start = 0
|
|
while total - start > limit:
|
|
cut = start + limit
|
|
if chunk_method == "smart" and energy.numel() > 1:
|
|
radius = max(1, int(limit * float(search_window) // frame))
|
|
first = max(0, int(cut // frame) - radius)
|
|
last = min(int(energy.shape[-1]), int(cut // frame) + radius + 1)
|
|
if last > first:
|
|
cut = int((first + int(torch.argmin(energy[first:last]))) * frame)
|
|
cut = max(start + 1, min(cut, total - 1))
|
|
chunks.append(mono[..., start:cut])
|
|
start = cut
|
|
chunks.append(mono[..., start:])
|
|
return [chunk for chunk in chunks if chunk.numel() > 0] or [mono]
|
|
|
|
|
|
def combine_audio_chunks(
|
|
chunks: list[torch.Tensor],
|
|
sample_rate: int,
|
|
method: str = "concatenate",
|
|
silence_between_chunks_ms: int = 100,
|
|
crossfade_ms: int = 10,
|
|
) -> torch.Tensor:
|
|
"""Join synthesised chunks with the mode ``chunk_combination_method`` names.
|
|
|
|
Args:
|
|
chunks: The chunk waveforms, in order.
|
|
sample_rate: Hz.
|
|
method: ``concatenate`` (hard join), ``silence_padding`` (a real gap, so
|
|
the boundary reads as a pause rather than a cut), or ``crossfade``
|
|
(overlap the tail and head with a linear ramp, so the join is
|
|
continuous).
|
|
silence_between_chunks_ms: The gap for ``silence_padding``.
|
|
crossfade_ms: The overlap for ``crossfade``.
|
|
|
|
Returns:
|
|
One ``[C, T]`` waveform. A single chunk is returned untouched.
|
|
|
|
Raises:
|
|
ValueError: ``chunks`` is empty.
|
|
"""
|
|
if not chunks:
|
|
raise ValueError("There are no audio chunks to join.")
|
|
if len(chunks) == 1:
|
|
return chunks[0]
|
|
|
|
if method == "silence_padding":
|
|
gap_samples = max(0, int(sample_rate * int(silence_between_chunks_ms) / 1000))
|
|
gap = torch.zeros(
|
|
chunks[0].shape[:-1] + (gap_samples,),
|
|
dtype=chunks[0].dtype,
|
|
device=chunks[0].device,
|
|
)
|
|
parts: list[torch.Tensor] = []
|
|
for position, chunk in enumerate(chunks):
|
|
parts.append(chunk)
|
|
if position != len(chunks) - 1:
|
|
parts.append(gap)
|
|
return torch.cat(parts, dim=-1)
|
|
|
|
if method == "crossfade":
|
|
overlap = max(
|
|
1,
|
|
min(
|
|
int(sample_rate * int(crossfade_ms) / 1000),
|
|
*(int(c.shape[-1]) for c in chunks),
|
|
),
|
|
)
|
|
joined = chunks[0]
|
|
for chunk in chunks[1:]:
|
|
tail = joined[..., -overlap:]
|
|
head = chunk[..., :overlap]
|
|
ramp = torch.linspace(
|
|
0.0, 1.0, overlap, dtype=joined.dtype, device=joined.device
|
|
)
|
|
blended = tail * (1.0 - ramp) + head * ramp
|
|
joined = torch.cat(
|
|
[joined[..., :-overlap], blended, chunk[..., overlap:]], dim=-1
|
|
)
|
|
return joined
|
|
|
|
return torch.cat(chunks, dim=-1)
|