wildminder-ComfyUI-Chatterbox/modules/audio_utils.py

"""Audio conversion helpers for the Chatterbox nodes.

Everything that turns a ComfyUI ``AUDIO`` dict into something Chatterbox can
consume -- and back again -- lives here, in one place:

* :func:`extract_waveform` / :func:`to_comfy_audio` -- ``AUDIO`` <-> tensor.
* :func:`downmix_to_mono` -- the ``[B, C, T]`` mean-downmix Chatterbox's
  tokenizer needs.
* :func:`comfy_audio_to_numpy` / :func:`numpy_to_comfy_audio` -- ``AUDIO`` <->
  numpy.
* :func:`preprocess_comfy_audio` -- the whole chain, resampled and scrubbed,
  in one call. What it returns is handed straight to the vendored
  ``generate(audio_prompt_path=...)``: the vendor's loader accepts an
  in-memory waveform (``load_audio_at`` in ``src/chatterbox/models/dsp.py``),
  so no temporary ``.wav`` is ever written -- and the PCM_16 quantization a
  file round-trip imposed is gone with it.
* :func:`set_seed` -- the seeding routine that replaces the deleted
  ``set_chatterbox_seed``.
* :func:`check_for_interrupt` -- a cooperative cancellation poll.

``comfy.model_management`` is reached through
:func:`modules.device_utils._get_model_management`, never imported at module
scope, so this module is importable without ComfyUI booted. ``torchaudio`` is
imported inside the function that needs it: it is not a node-load-time cost,
and it is heavier than the rest of the module.
"""

from __future__ import annotations

import logging
import random
from typing import Optional, Union

import numpy as np
import torch

from .device_utils import _get_model_management

__all__ = [
    "S3GEN_SAMPLE_RATE",
    "S3_TOKENIZER_SAMPLE_RATE",
    "set_seed",
    "check_for_interrupt",
    "interrupt_exception_class",
    "resample_waveform",
    "downmix_to_mono",
    "extract_waveform",
    "to_comfy_audio",
    "comfy_audio_to_numpy",
    "numpy_to_comfy_audio",
    "preprocess_comfy_audio",
]


# Every Chatterbox model synthesises through S3Gen at 24 kHz
# (``src/chatterbox/models/s3gen/const.py::S3GEN_SR``). Reference audio is
# resampled to it so the tokenizer's token-rate and mel-rate invariants hold.
S3GEN_SAMPLE_RATE = 24000


# The S3 tokenizer consumes its clip at 16 kHz
# (``src/chatterbox/models/s3tokenizer/s3tokenizer.py::S3_SR``). Only voice
# conversion's *source* clip needs it -- everything the models condition on
# arrives at ``S3GEN_SAMPLE_RATE``. The value is redeclared here rather than
# imported: the vendored tree is never this module's dependency, and
# ``modules/generation.py`` says why that holds for the node layer too.
S3_TOKENIZER_SAMPLE_RATE = 16000


# ---------------------------------------------------------------------------
# Seeding and cancellation
# ---------------------------------------------------------------------------


def set_seed(seed: int) -> None:
    """Seed torch, numpy and ``random`` from one widget value.

    Replaces the deleted ``set_chatterbox_seed``. ``seed == 0`` means "random"
    (the node widget's convention), which is resolved here to a real seed so
    the run that actually happened can be logged.

    Args:
        seed: Non-negative integer. ``0`` picks a random seed.
    """
    if seed == 0:
        seed = random.randint(1, 0xFFFFFFFFFFFFFFFF)

    # numpy's legacy seeder only accepts a 32-bit value, and torch's accepts
    # the full range. Clamping numpy rather than the seed keeps the two RNGs
    # correlated for every value a 64-bit widget can produce.
    max_numpy_seed = 2**32 - 1

    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
    np.random.seed(seed % max_numpy_seed)
    random.seed(seed)


def interrupt_exception_class() -> type:
    """The exception class ComfyUI raises when a run is cancelled.

    ``comfy.model_management.InterruptProcessingException`` subclasses
    ``BaseException``, **not** ``Exception``. Anything written as
    ``except Exception`` therefore cannot catch it, and anything written as
    ``except BaseException`` catches it by accident. Generation code that wants
    one wants exactly this class.
    """
    candidate = getattr(_get_model_management(), "InterruptProcessingException", None)
    return candidate if isinstance(candidate, type) else BaseException


def check_for_interrupt() -> bool:
    """Poll for a user cancellation without ever raising.

    Only the interrupt class is caught: a ``KeyboardInterrupt`` arriving
    mid-generation must keep propagating rather than being reported as "the
    user pressed stop".

    Returns:
        ``True`` when the run was cancelled, ``False`` otherwise.
    """
    try:
        _get_model_management().throw_exception_if_processing_interrupted()
    except interrupt_exception_class():
        return True
    return False


# ---------------------------------------------------------------------------
# Resampling
# ---------------------------------------------------------------------------


def resample_waveform(
    waveform: torch.Tensor,
    orig_sr: int,
    target_sr: int,
) -> torch.Tensor:
    """Resample ``waveform`` from ``orig_sr`` to ``target_sr``.

    Uses ``torchaudio.functional.resample`` -- the ComfyUI-core idiom -- rather
    than reaching for librosa, which is a much heavier import for a sinc
    resampler. Resampling in tensor space also avoids a numpy round-trip.

    Args:
        waveform: A float tensor of any shape; the last axis is time.
        orig_sr: Source sample rate.
        target_sr: Destination sample rate.

    Returns:
        The resampled tensor. The input object itself when the rates match.

    Raises:
        ValueError: Either sample rate is not positive.
        RuntimeError: ``torchaudio`` is not importable.
    """
    orig_sr = int(orig_sr)
    target_sr = int(target_sr)
    if orig_sr <= 0 or target_sr <= 0:
        raise ValueError(
            f"Sample rates must be positive (got orig={orig_sr}, target={target_sr})"
        )
    if orig_sr == target_sr:
        return waveform

    try:
        import torchaudio.functional as audio_functional
    except Exception as exc:  # pragma: no cover - torchaudio ships with torch
        raise RuntimeError(
            "Resampling reference audio requires torchaudio, which could not be "
            f"imported: {exc}"
        ) from exc

    return audio_functional.resample(waveform, orig_sr, target_sr)


# ---------------------------------------------------------------------------
# Waveform plumbing
# ---------------------------------------------------------------------------


def extract_waveform(audio: Union[dict, torch.Tensor], name: str = "audio") -> torch.Tensor:
    """Pull the waveform tensor out of a ComfyUI ``AUDIO`` dict.

    The batch axis is stripped: ComfyUI hands over ``[B, C, T]`` and every
    Chatterbox entry point wants ``[C, T]``. Multi-channel audio is preserved
    here on purpose -- the downmix is a separate, explicit step
    (:func:`downmix_to_mono`), because whether to mix is the caller's
    decision, not a side effect of unpacking.

    Args:
        audio: A ComfyUI ``AUDIO`` dict, or a bare waveform tensor.
        name: Field name used in error messages.

    Returns:
        A ``[C, T]`` tensor, or ``[T]`` when the input was already
        single-channel and unbatched.

    Raises:
        ValueError: The dict is malformed or the waveform is empty.
    """
    if isinstance(audio, torch.Tensor):
        waveform = audio
    else:
        if not isinstance(audio, dict):
            raise ValueError(
                f"{name}: expected an AUDIO dict, got {type(audio).__name__}"
            )
        if "waveform" not in audio or "sample_rate" not in audio:
            raise ValueError(
                f"{name}: an AUDIO dict needs both 'waveform' and 'sample_rate'"
            )
        waveform = audio["waveform"]
        if waveform is None:
            raise ValueError(f"{name}: 'waveform' is None")

    if waveform is None:
        raise ValueError(f"{name}: 'waveform' is None")
    if waveform.numel() == 0:
        raise ValueError(f"{name}: the waveform is empty")
    if waveform.dim() == 3:
        waveform = waveform[0]
    if waveform.dim() == 1:
        waveform = waveform.unsqueeze(0)
    if waveform.dim() != 2:
        raise ValueError(
            f"{name}: expected [B, C, T] / [C, T] / [T], got {tuple(waveform.shape)}"
        )
    return waveform


def downmix_to_mono(waveform: torch.Tensor) -> torch.Tensor:
    """Mean-mix a ``[B, C, T]`` / ``[C, T]`` waveform down to one channel.

    Chatterbox's voice encoder and tokenizer are mono-only. Averaging (rather
    than taking channel 0) is what ComfyUI itself does and what keeps an
    anti-phase stereo pair from cancelling to silence.

    Args:
        waveform: ``[B, C, T]``, ``[C, T]`` or ``[T]``.

    Returns:
        ``[B, T]`` for a batched input, ``[T]`` otherwise.
    """
    if waveform.dim() == 1:
        return waveform
    if waveform.dim() == 2:
        if waveform.shape[0] > 1:
            return waveform.mean(dim=0)
        return waveform[0]
    if waveform.dim() == 3:
        return waveform.mean(dim=1)
    raise ValueError(
        f"Expected a waveform of rank 1-3, got shape {tuple(waveform.shape)}"
    )


def to_comfy_audio(waveform: torch.Tensor, sample_rate: int) -> dict:
    """Build a ComfyUI ``AUDIO`` dict from a waveform tensor.

    Normalises ``[T]`` / ``[C, T]`` / ``[B, C, T]`` to ComfyUI's ``[B, C, T]``,
    so a node can hand back whatever its model returned.

    Args:
        waveform: The audio samples.
        sample_rate: Sample rate in Hz.

    Returns:
        ``{"waveform": FloatTensor[B, C, T], "sample_rate": int}`` on CPU.
    """
    if not isinstance(waveform, torch.Tensor):
        waveform = torch.as_tensor(waveform)
    if waveform.dim() == 1:
        waveform = waveform.unsqueeze(0).unsqueeze(0)
    elif waveform.dim() == 2:
        waveform = waveform.unsqueeze(0)
    elif waveform.dim() != 3:
        raise ValueError(
            f"Cannot build an AUDIO dict from shape {tuple(waveform.shape)}"
        )
    return {
        "waveform": waveform.detach().to(device="cpu", dtype=torch.float32),
        "sample_rate": int(sample_rate),
    }


def comfy_audio_to_numpy(
    audio: Union[dict, torch.Tensor], name: str = "audio"
) -> np.ndarray:
    """Mono float32 numpy view of a ComfyUI ``AUDIO`` dict.

    Args:
        audio: A ComfyUI ``AUDIO`` dict, or a bare waveform tensor.
        name: Field name used in error messages.

    Returns:
        A 1-D ``float32`` array.
    """
    mono = downmix_to_mono(extract_waveform(audio, name=name))
    return mono.detach().to(device="cpu", dtype=torch.float32).numpy()


def numpy_to_comfy_audio(waveform: np.ndarray, sample_rate: int) -> dict:
    """Build a ComfyUI ``AUDIO`` dict from a numpy waveform.

    Args:
        waveform: ``[T]`` or ``[C, T]`` samples.
        sample_rate: Sample rate in Hz.

    Returns:
        ``{"waveform": FloatTensor[B, C, T], "sample_rate": int}``.
    """
    if not isinstance(waveform, np.ndarray):
        waveform = np.asarray(waveform)
    return to_comfy_audio(torch.from_numpy(np.ascontiguousarray(waveform)), sample_rate)


def preprocess_comfy_audio(
    audio: Optional[dict],
    target_sr: int = S3GEN_SAMPLE_RATE,
    name: str = "audio",
) -> Optional[np.ndarray]:
    """Turn a ComfyUI ``AUDIO`` dict into a mono, resampled float32 array.

    The full chain, in the order that has to happen: extract, downmix, scrub
    non-finite samples **before** resampling (one NaN would otherwise poison
    the whole sinc kernel), resample, scrub again.

    Args:
        audio: A ComfyUI ``AUDIO`` dict, or ``None``.
        target_sr: Sample rate the caller needs.
        name: Field name used in error messages.

    Returns:
        A mono ``float32`` numpy array, or ``None`` when ``audio`` is ``None``
        or carries no samples.

    Raises:
        RuntimeError: Resampling was required but no backend was available.
    """
    if audio is None:
        return None
    if isinstance(audio, dict):
        waveform = audio.get("waveform")
        if waveform is None or waveform.numel() == 0:
            return None
        original_sr = int(audio.get("sample_rate") or 0)
    else:
        waveform = audio
        original_sr = 0
        if waveform is None or waveform.numel() == 0:
            return None

    if original_sr <= 0:
        raise ValueError(
            f"{name}: an AUDIO dict needs a positive 'sample_rate' to resample; "
            f"got {original_sr!r}"
        )

    mono = downmix_to_mono(extract_waveform(audio, name=name))

    if not torch.isfinite(mono).all():
        logging.error(
            "[Chatterbox TTS] %s contains NaN or Inf samples; replacing them with "
            "zeros before resampling",
            name,
        )
        mono = torch.nan_to_num(mono, nan=0.0, posinf=0.0, neginf=0.0)

    if original_sr != int(target_sr):
        mono = resample_waveform(mono, original_sr, int(target_sr))

    samples = mono.detach().to(device="cpu", dtype=torch.float32).numpy()

    if not np.all(np.isfinite(samples)):
        logging.error(
            "[Chatterbox TTS] %s still contains NaN or Inf after resampling; "
            "replacing them with zeros",
            name,
        )
        samples = np.nan_to_num(samples, nan=0.0, posinf=0.0, neginf=0.0)

    peak = float(np.abs(samples).max()) if samples.size else 0.0
    if peak > 10.0:
        logging.warning(
            "[Chatterbox TTS] %s peak amplitude is %.2f, far outside the usual "
            "[-1, 1] range; scaling it down",
            name,
            peak,
        )
        samples = samples / peak

    return np.ascontiguousarray(samples, dtype=np.float32)