"""Dtype utilities for Chatterbox TTS nodes. Dtype resolution, config bookkeeping and casting helpers built on ComfyUI's ``model_management`` so precision selection matches the rest of the UI. Two things here are load-bearing rather than cosmetic: * **transformers v5 renamed the config attribute.** ``config.torch_dtype`` is a deprecated property in transformers >= 5.0 that warns on every get *and* set; ``config.dtype`` is the replacement. Reading or writing the old name on the installed 5.3.0 floods the log. :func:`set_config_dtype` / :func:`get_config_dtype` probe for the deprecated property and use whichever name the installed transformers actually defines. * **An explicit dtype is validated, not obeyed.** ``resolve_dtype("bf16")`` on hardware that cannot run bf16 raises instead of handing back a ``torch.bfloat16`` the caller will discover three minutes into a CUDA kernel. Only ``"auto"`` is a request that ComfyUI gets to answer. * **Some frontends are fp32-resident.** The speaker encoder, voice encoder and speech tokenizer are excluded from every cast: their own code hard-casts their inputs to float32, and half-precision there either crashes outright (``Kaldi.fbank``'s FFT has no half-dtype kernel) or mismatches the fp32 tensors they are fed. """ import logging from typing import List, Optional import torch from .device_utils import _get_model_management, should_use_bf16, should_use_fp16 __all__ = [ "DTYPE_AUTO", "DTYPE_BF16", "DTYPE_FP16", "DTYPE_FP32", "get_dtype_options", "set_config_dtype", "get_config_dtype", "resolve_dtype", "get_dtype_str", "cast_model_to_dtype", "cast_model_to_dtype_if_needed", "representative_dtype", ] # Dtype string constants DTYPE_AUTO = "auto" DTYPE_BF16 = "bf16" DTYPE_FP16 = "fp16" DTYPE_FP32 = "fp32" # Mapping from string to torch.dtype _DTYPE_MAP = { DTYPE_BF16: torch.bfloat16, DTYPE_FP16: torch.float16, DTYPE_FP32: torch.float32, } # Reverse mapping _DTYPE_STR_MAP = { torch.bfloat16: DTYPE_BF16, torch.float16: DTYPE_FP16, torch.float32: DTYPE_FP32, } # Class names of feature-extraction frontends that stay float32 no matter what # the cast dtype is. Matched by *name* on purpose: importing the vendor tree # here would pull librosa/transformers into every cast call and break the # mocked test environment. # # - ``CAMPPlus`` (s3gen's speaker encoder): its ``inference()`` hard-casts the # fbank features to float32 before the forward, so upstream intends fp32 # weights; and the fbank frontend's ``torch.fft.rfft`` has no half-dtype # kernel at all ("Unsupported dtype BFloat16" was the field failure). # - ``VoiceEncoder`` (t3's speaker embedding): a cuDNN LSTM over fp32 librosa # mels. # - ``S3Tokenizer``: the vendored wrapper over the external ``s3tokenizer`` # package, whose mel path and encoder cannot be patched from here. # - ``HiFTGenerator`` (the s3gen vocoder): SineGen's harmonic math promotes to # fp32 (``torch.arange`` defaults) and ``torch.istft`` consumes the fp32 # window; also the quality-sensitive end of the pipeline, where half # precision buys nothing worth the risk. _FP32_RESIDENT_CLASS_NAMES = frozenset( {"CAMPPlus", "VoiceEncoder", "S3Tokenizer", "HiFTGenerator"} ) # Explicit dtypes validated against the target device. "fp32" is absent on # purpose: it is the universally safe fallback and must always resolve. _VALIDATED_DTYPES = frozenset({DTYPE_BF16}) def get_dtype_options() -> List[str]: """Dtype values for the node dropdown. ``"auto"`` leads because it is the default and the widget's index 0; the rest follow in the widget's own fixed order, which no existing workflow can be bound to yet (the V1 nodes had no dtype widget). Returns: List of dtype option strings. """ return [DTYPE_AUTO, DTYPE_BF16, DTYPE_FP16, DTYPE_FP32] def _torch_dtype_is_deprecated_property(config) -> bool: """True when ``type(config).torch_dtype`` is the transformers >= 5.0 deprecated property alias (which warns on every get and set).""" return isinstance(getattr(type(config), "torch_dtype", None), property) def set_config_dtype(config, dtype) -> None: """Record the load dtype on a transformers config, without the v5 warning. Writes ``config.dtype`` when the installed transformers defines ``torch_dtype`` as the deprecated property, and ``config.torch_dtype`` otherwise (transformers 4.x, and any duck-typed config that stores a plain attribute). """ if _torch_dtype_is_deprecated_property(config): config.dtype = dtype else: config.torch_dtype = dtype def get_config_dtype(config): """Return the torch dtype recorded on a transformers config, or ``None``.""" if _torch_dtype_is_deprecated_property(config): dtype = getattr(config, "dtype", None) else: dtype = getattr(config, "torch_dtype", None) if dtype is None: dtype = getattr(config, "dtype", None) if isinstance(dtype, str): dtype = getattr(torch, dtype, None) return dtype def _validate_dtype_on_device(dtype_str: str, device: torch.device) -> None: """Raise when an explicit half dtype cannot run on ``device``. CPU is exempt: torch implements bf16 there, just slowly, and a user who asks for bf16 on a CPU-only box should get what they asked for. Raises: ValueError: when the dtype is unsupported on the target device. """ if dtype_str not in _VALIDATED_DTYPES: return if device.type == "cpu": return supported = should_use_bf16(device) if dtype_str == DTYPE_BF16 else should_use_fp16(device) if supported: return raise ValueError( f"dtype '{dtype_str}' was requested explicitly but device '{device}' cannot " f"run it (ComfyUI's precision check says this hardware has no " f"{dtype_str} support). Use dtype 'auto' to let ComfyUI choose, or pick " f"'{DTYPE_FP32}' / '{DTYPE_FP16}'." ) def resolve_dtype(dtype_str: str, device: Optional[torch.device] = None) -> torch.dtype: """Resolve a dtype string to a ``torch.dtype``. Args: dtype_str: Dtype string (``"auto"``, ``"bf16"``, ``"fp16"``, ``"fp32"``). device: Device the dtype has to run on. ``None`` means ComfyUI's default compute device; it is only consulted for validation of an explicit choice. Returns: The resolved ``torch.dtype``. Raises: ValueError: If ``dtype_str`` is not recognised, or is an explicit dtype the target device cannot run. """ if dtype_str == DTYPE_AUTO or dtype_str is None: mm = _get_model_management() if device is None: device = mm.get_torch_device() if should_use_bf16(device): return torch.bfloat16 if should_use_fp16(device): return torch.float16 return torch.float32 if dtype_str in _DTYPE_MAP: if device is None: device = _get_model_management().get_torch_device() _validate_dtype_on_device(dtype_str, device) return _DTYPE_MAP[dtype_str] raise ValueError( f"Unknown dtype: {dtype_str}. Valid options: {get_dtype_options()}" ) def get_dtype_str(dtype: torch.dtype) -> str: """Convert a ``torch.dtype`` back to its widget string. Raises: ValueError: If the dtype is not one this module offers. """ if dtype in _DTYPE_STR_MAP: return _DTYPE_STR_MAP[dtype] raise ValueError( f"Unknown torch.dtype: {dtype}. Supported: {sorted(_DTYPE_STR_MAP.values())}" ) def _iter_named_modules(model): """``(name, module)`` pairs for an ``nn.Module`` tree *or* a plain container. The vendored pack classes (``ChatterboxTTS``, ``ChatterboxVC``, ``ChatterboxMultilingualTTS``, ``ChatterboxTurboTTS``) are plain objects whose weights live in ``nn.Module`` attributes (``t3``, ``s3gen``, ``ve``). ``modules.loader`` registers those submodules on the handler for ComfyUI's own machinery, but the dtype walkers below are handed the *pack objects* -- so they must not assume ``named_modules()`` exists. """ if isinstance(model, torch.nn.Module): yield from model.named_modules() return seen: set = set() for attr, value in list(getattr(model, "__dict__", {}).items()): if not isinstance(value, torch.nn.Module): continue if id(value) in seen: continue seen.add(id(value)) # Delegate to the child's own walker: an nn.Module registers its # children in ``_modules``, which plain ``__dict__`` iteration does # not descend into. for name, mod in value.named_modules(): yield (f"{attr}.{name}" if name else attr), mod def _iter_named_parameters(model): """``(name, parameter)`` pairs, as tolerant as :func:`_iter_named_modules`.""" if isinstance(model, torch.nn.Module): yield from model.named_parameters() return seen: set = set() for attr, value in list(getattr(model, "__dict__", {}).items()): if isinstance(value, torch.nn.Parameter): if id(value) not in seen: seen.add(id(value)) yield attr, value for name, mod in _iter_named_modules(model): for pname, param in mod.named_parameters(recurse=False): if id(param) in seen: continue seen.add(id(param)) yield f"{name}.{pname}", param def _quant_protected_names(model) -> set: """Parameter names excluded from bulk dtype casts. - Parameters of modules marked ``_quant_resident`` (GGUF / int8 / FP8 linears), whose storage is raw packed bytes and is not castable at all. - Any parameter whose leaf name is ``weight_scale``, which is kept in the quantizer's own precision. """ protected = set() for name, mod in _iter_named_modules(model): if getattr(mod, "_quant_resident", False): prefix = f"{name}." if name else "" protected.add(f"{prefix}weight") protected.add(f"{prefix}weight_scale") return protected def _quant_resident_module_names(model) -> set: """Names of modules whose parameters hold packed quant data.""" return { name for name, mod in _iter_named_modules(model) if getattr(mod, "_quant_resident", False) } def _fp32_resident_prefixes(model) -> tuple: """``"name."`` prefixes of the fp32-resident frontend modules. A resident frontend (see ``_FP32_RESIDENT_CLASS_NAMES``) owns a whole subtree -- ``CAMPPlus`` contains dozens of nested layers -- so protection is by prefix, not by direct-owner match. """ return tuple( f"{name}." for name, mod in _iter_named_modules(model) if type(mod).__name__ in _FP32_RESIDENT_CLASS_NAMES ) def _fp32_protected_names(model) -> set: """Parameter names under fp32-resident frontends (whole subtree).""" prefixes = _fp32_resident_prefixes(model) if not prefixes: return set() return { name for name, _param in _iter_named_parameters(model) if name.startswith(prefixes) } def representative_dtype(model): """The model's floating-point parameter dtype, ignoring quant and fp32 resident storage. Returns ``None`` when the tree holds no floating parameter at all and the model carries no ``dtype`` attribute either. """ try: resident_modules = _quant_resident_module_names(model) fp32_prefixes = _fp32_resident_prefixes(model) for name, param in _iter_named_parameters(model): if not param.dtype.is_floating_point: continue owner = name.rpartition(".")[0] if owner in resident_modules: continue if fp32_prefixes and name.startswith(fp32_prefixes): continue return param.dtype except Exception: pass return getattr(model, "dtype", None) def _cast_mismatched_params(model, dtype: torch.dtype) -> None: """Cast only floating, unprotected, mismatched parameters, then re-tie. Unprotected means: not quant-resident packed storage, and not part of an fp32-resident frontend subtree (see ``_FP32_RESIDENT_CLASS_NAMES``). """ protected = _quant_protected_names(model) | _fp32_protected_names(model) for name, param in _iter_named_parameters(model): if name in protected: continue if not param.dtype.is_floating_point: continue if param.dtype != dtype: param.data = param.data.to(dtype) config = getattr(model, "config", None) if config is not None and hasattr(model, "tie_weights"): decoder_config = getattr(config, "decoder_config", None) tied = bool(getattr(decoder_config, "tie_word_embeddings", False)) or bool( getattr(config, "tie_word_embeddings", False) ) if tied: model.tie_weights() def cast_model_to_dtype(model, dtype: torch.dtype) -> None: """Cast a model's floating parameters to ``dtype``. Quant-resident storage (raw packed GGUF blocks, int8 ConvRot weights, the quantizer's own fp32 ``weight_scale``) is **never** touched: those bytes are a payload, not numbers, and casting them corrupts the checkpoint. """ if dtype is None: return _cast_mismatched_params(model, dtype) logging.debug( "[Chatterbox TTS] Model cast to dtype (filtered): %s", dtype ) def cast_model_to_dtype_if_needed(model, dtype: torch.dtype) -> None: """Cast to ``dtype`` only when a castable parameter actually mismatches.""" if dtype is None: return protected = _quant_protected_names(model) | _fp32_protected_names(model) mismatched = [] for name, param in _iter_named_parameters(model): if name in protected or not param.dtype.is_floating_point: continue if param.dtype != dtype: mismatched.append((name, param.dtype)) if not mismatched: return _cast_mismatched_params(model, dtype) sources = ", ".join(sorted({str(src) for _, src in mismatched})) logging.debug( "[Chatterbox TTS] Model cast %s -> %s (%d mismatched params)", sources, dtype, len(mismatched), )