wildminder-ComfyUI-VibeVoice/modules/external_loader.py

"""External model loading for VibeVoice.

Loads a VibeVoice model from a standalone weight file (safetensors / .bin /
.gguf) placed in ComfyUI's ``diffusion_models`` folder, binding the required
config / tokenizer / preprocessor JSONs via sidecar files or packaged defaults.

This module bypasses ComfyUI's ``model_detection.detect_unet_config()`` entirely
(VibeVoice's state-dict keys are not recognized by it) and instead reuses the
existing :class:`~modules.loader.VibeVoiceLoader` internals for config,
tokenizer, processor, and model instantiation.

Sidecar file convention (place next to the weight file):
    - ``<weight>.config.json``          → architecture config (preferred)
    - ``config.json``                   → architecture config (same directory)
    - ``<weight>.preprocessor.json``    → audio preprocessor config (preferred)
    - ``preprocessor_config.json``      → audio preprocessor config (same dir)
    - ``tokenizer.json``                → Qwen2.5 text tokenizer (same dir)
"""

import os
import json
import logging
import contextlib

import torch

import comfy.utils
import comfy.model_management as model_management

from transformers import BitsAndBytesConfig

from ..src.vibevoice.modular.configuration_vibevoice import VibeVoiceConfig, VibeVoiceASRConfig
from ..src.vibevoice.modular.configuration_vibevoice_streaming import VibeVoiceStreamingConfig
from ..src.vibevoice.modular.modeling_vibevoice_asr import VibeVoiceASRForConditionalGeneration
from ..src.vibevoice.modular.modular_vibevoice_text_tokenizer import VibeVoiceASRTextTokenizerFast
from ..src.vibevoice.processor.vibevoice_asr_processor import VibeVoiceASRProcessor
from ..src.vibevoice.processor.vibevoice_tokenizer_processor import VibeVoiceTokenizerProcessor

from .loader import VibeVoiceLoader
from .attention_utils import (
    SAGE_ATTENTION_AVAILABLE,
    resolve_attention_mode,
    get_attn_implementation_for_load,
    check_sage_attention_compatible,
)
from .dtype_utils import resolve_dtype, cast_model_to_dtype_if_needed, set_config_dtype
from .convrot_quant import UnsupportedQuantFormat
from .quant_common import validate_weight_plan

if SAGE_ATTENTION_AVAILABLE:
    from ..src.vibevoice.modular.sage_attention_patch import set_sage_attention

logger = logging.getLogger(__name__)

# Packaged default config filenames keyed by config_name.
# Only models with a packaged default are listed here; all other config_name
# values require a sidecar config file next to the weight file.
_PACKAGED_CONFIG_FILES = {
    "VibeVoice-1.5B": "default_VibeVoice-1.5B_config.json",
    "VibeVoice-7B": "default_VibeVoice-Large_config.json",
}

# Sentinel config_name that resolves the architecture from the weight file's
# embedding fingerprint at load time (plan 2026-08-27, D6).
AUTO_CONFIG_NAME = "Auto-detect"

# All config_name values accepted by the loader node dropdown.
EXTERNAL_CONFIG_OPTIONS = [
    AUTO_CONFIG_NAME,
    "VibeVoice-1.5B",
    "VibeVoice-7B",
    "VibeVoice-Realtime-0.5B",
    "VibeVoice-ASR",
]

# Legacy config_name values removed from the dropdown but still honored so
# saved workflows keep loading (plan 2026-08-27, D1/D2). Keys are compared
# case-insensitively.
_LEGACY_CONFIG_ALIASES = {
    "vibevoice-large": "VibeVoice-7B",
}


def normalize_config_name(config_name: str) -> str:
    """Map legacy/alias config_name values onto their canonical option.

    Exact options pass through unchanged; case-insensitive alias hits return
    the canonical name; anything unknown is returned as-is (callers decide
    validity).

    Args:
        config_name: Raw config_name value (may come from a saved workflow).

    Returns:
        Canonical config_name when an alias applies, else the input value.
    """
    if not config_name:
        return config_name
    if config_name in EXTERNAL_CONFIG_OPTIONS:
        return config_name
    return _LEGACY_CONFIG_ALIASES.get(config_name.lower(), config_name)


def resolve_auto_config_name(weight_path: str, gguf_reader=None, weights_fp=None) -> str:
    """Resolve ``AUTO_CONFIG_NAME`` to a concrete family from the weights.

    Plan 2026-08-27 (D6): Auto-detect reads the weight file's architecture
    fingerprint (header-only). A conclusive fingerprint becomes the selected
    config_name; an inconclusive one raises an actionable ``ValueError``
    BEFORE any heavy load work. Shared by the node (which must resolve the
    name before computing the cache identity) and the loader (defense in
    depth for direct callers) so both fail fast with the same message.

    Args:
        weight_path: Absolute path to the external weight file.
        gguf_reader: Optional open ``gguf.GGUFReader`` for ``.gguf`` files
            (avoids re-opening when the loader already holds one).
        weights_fp: Optional precomputed ``WeightsFingerprint``; when omitted
            the fingerprint is computed here (header-only).

    Returns:
        The detected config_name (e.g. ``"VibeVoice-7B"``).

    Raises:
        ValueError: When the architecture cannot be determined.
    """
    if weights_fp is None:
        from .config_detect import fingerprint_weights

        weights_fp = fingerprint_weights(weight_path, gguf_reader=gguf_reader)

    detected = weights_fp.config_name if weights_fp is not None else None
    if not detected:
        explicit = [o for o in EXTERNAL_CONFIG_OPTIONS if o != AUTO_CONFIG_NAME]
        raise ValueError(
            f"config_name '{AUTO_CONFIG_NAME}': could not determine the "
            f"VibeVoice architecture from "
            f"'{os.path.basename(weight_path)}'. Auto-detect reads the "
            f"checkpoint's embedding fingerprint, which is unavailable for "
            f"this file (.bin/.pt checkpoints and unknown architectures "
            f"cannot be auto-detected). Select config_name explicitly "
            f"(one of: {', '.join(explicit)}) or place a sidecar config.json "
            f"next to the weight file."
        )
    logger.info(
        f"Auto-detected architecture '{detected}' from "
        f"'{os.path.basename(weight_path)}'"
    )
    return detected


def reconcile_config(selected_name: str, resolved_config, weights_fp):
    """Reconcile the selected config against the weights' fingerprint.

    Plan 2026-08-27 (D5): the checkpoint's embedding shape is a perfect
    architecture fingerprint. When it CONTRADICTS the resolved config family,
    the fingerprint wins and the detected family's config is substituted —
    a mismatching config (dropdown or sidecar) is a guaranteed size-mismatch
    crash otherwise. When the fingerprint is inconclusive (``None``) or
    agrees, the selection stands.

    Args:
        selected_name: The (normalized) config_name in effect.
        resolved_config: The config object loaded for ``selected_name``.
        weights_fp: ``WeightsFingerprint`` from the weight file, or ``None``.

    Returns:
        ``(effective_name, changed)`` — the config_name to use and whether
        it differs from ``selected_name``.
    """
    from .config_detect import config_fingerprint

    if weights_fp is None:
        return selected_name, False

    cfg_fp = config_fingerprint(resolved_config)
    if cfg_fp is None or cfg_fp == (weights_fp.hidden_size, weights_fp.vocab_size):
        return selected_name, False

    detected = weights_fp.config_name
    if not detected or detected == selected_name:
        return selected_name, False

    logger.warning(
        f"Config mismatch: weights contain '{weights_fp.source_key}' with "
        f"shape (vocab={weights_fp.vocab_size}, hidden={weights_fp.hidden_size}) "
        f"-> {detected}, but config '{selected_name}' "
        f"(hidden={cfg_fp[0]}, vocab={cfg_fp[1]}) was selected. "
        f"Using '{detected}'."
    )
    return detected, True


# ====================================================================
# Path resolution helpers
# ====================================================================

def _packaged_configs_dir() -> str:
    """Return the absolute path to the packaged configs directory."""
    return os.path.join(
        os.path.dirname(__file__), "..", "src", "vibevoice", "configs"
    )


def _get_packaged_config_path(config_name: str) -> str:
    """Return the packaged default config path for ``config_name``.

    Args:
        config_name: One of the keys in ``_PACKAGED_CONFIG_FILES``.

    Returns:
        Absolute path to the packaged config file, or empty string if no
        packaged default exists for this config_name.
    """
    filename = _PACKAGED_CONFIG_FILES.get(config_name)
    if filename is None:
        return ""
    return os.path.normpath(os.path.join(_packaged_configs_dir(), filename))


def resolve_sidecar_config(weight_path: str, config_name: str) -> str:
    """Resolve the architecture config JSON path for an external weight file.

    Resolution priority:
        1. ``<weight_path>.config.json`` (sidecar, preferred)
        2. ``config.json`` in the same directory as the weight file
        3. Packaged default selected by ``config_name``

    Args:
        weight_path: Absolute path to the external weight file.
        config_name: Config selector used for the packaged-default fallback.

    Returns:
        Absolute path to the config JSON to use.

    Raises:
        FileNotFoundError: If no config can be resolved.
    """
    # 1. Sidecar: <weight_path>.config.json
    sidecar_path = weight_path + ".config.json"
    if os.path.exists(sidecar_path):
        logger.debug(f"Using sidecar config: {sidecar_path}")
        return sidecar_path

    # 2. Sidecar: config.json in the same directory
    dir_sidecar = os.path.join(os.path.dirname(weight_path), "config.json")
    if os.path.exists(dir_sidecar):
        logger.debug(f"Using directory sidecar config: {dir_sidecar}")
        return dir_sidecar

    # 3. Packaged default based on config_name
    packaged = _get_packaged_config_path(config_name)
    if packaged and os.path.exists(packaged):
        logger.debug(f"Using packaged default config for '{config_name}': {packaged}")
        return packaged

    raise FileNotFoundError(
        f"No architecture config found for external model "
        f"'{os.path.basename(weight_path)}'.\n"
        f"Place a sidecar config next to the weight file:\n"
        f"  - '{sidecar_path}'  (preferred), or\n"
        f"  - 'config.json' in '{os.path.dirname(weight_path)}'\n"
        f"Alternatively select a config_name with a packaged default "
        f"({', '.join(_PACKAGED_CONFIG_FILES.keys())})."
    )


def resolve_sidecar_preprocessor(weight_path: str) -> str:
    """Resolve the audio preprocessor config JSON path for an external weight file.

    Resolution priority:
        1. ``<weight_path>.preprocessor.json`` (sidecar, preferred)
        2. ``preprocessor_config.json`` in the same directory as the weight file

    Args:
        weight_path: Absolute path to the external weight file.

    Returns:
        Absolute path to the preprocessor config, or empty string if not found
        (the processor will fall back to built-in defaults).
    """
    sidecar_path = weight_path + ".preprocessor.json"
    if os.path.exists(sidecar_path):
        logger.debug(f"Using sidecar preprocessor config: {sidecar_path}")
        return sidecar_path

    dir_sidecar = os.path.join(os.path.dirname(weight_path), "preprocessor_config.json")
    if os.path.exists(dir_sidecar):
        logger.debug(f"Using directory sidecar preprocessor config: {dir_sidecar}")
        return dir_sidecar

    return ""


def resolve_sidecar_tokenizer_dir(weight_path: str) -> str:
    """Return the directory to search for ``tokenizer.json``.

    The tokenizer is expected to live in the same directory as the weight file.
    If absent, :meth:`VibeVoiceLoader._load_tokenizer` falls back to the
    packaged tokenizer or a HuggingFace download.

    Args:
        weight_path: Absolute path to the external weight file.

    Returns:
        The directory containing the weight file.
    """
    return os.path.dirname(weight_path)


# ====================================================================
# ASR-specific helpers
# ====================================================================

# config_name values that select the ASR loading branch.
ASR_CONFIG_NAMES = {"VibeVoice-ASR"}


def is_asr_config_name(config_name: str) -> bool:
    """Return True if ``config_name`` selects the ASR loading branch."""
    return config_name in ASR_CONFIG_NAMES


def _load_asr_config(config_path: str) -> "VibeVoiceASRConfig":
    """Load a :class:`VibeVoiceASRConfig` from a sidecar config JSON.

    ASR configs share ``model_type == "vibevoice"`` with TTS configs, so the
    caller must select this loader explicitly (via the ``config_name`` dropdown).

    Args:
        config_path: Path to the ASR config.json.

    Returns:
        VibeVoiceASRConfig instance.

    Raises:
        FileNotFoundError: If the config file does not exist.
    """
    if not os.path.exists(config_path):
        raise FileNotFoundError(f"ASR config not found: {config_path}")
    return VibeVoiceASRConfig.from_pretrained(config_path)


def _load_asr_tokenizer(tokenizer_dir: str) -> "VibeVoiceASRTextTokenizerFast":
    """Load the ASR text tokenizer from ``tokenizer.json``.

    Resolution priority:
        1. ``tokenizer.json`` in ``tokenizer_dir`` (next to the weight file)
        2. The packaged ``tokenizer.json`` fallback

    Args:
        tokenizer_dir: Directory to search for ``tokenizer.json``.

    Returns:
        VibeVoiceASRTextTokenizerFast instance.

    Raises:
        FileNotFoundError: If no tokenizer.json can be found.
    """
    tokenizer_file_path = os.path.join(tokenizer_dir, "tokenizer.json")

    if not os.path.exists(tokenizer_file_path):
        packaged_tokenizer_path = os.path.join(
            _packaged_configs_dir(), "tokenizer.json"
        )
        if os.path.exists(packaged_tokenizer_path):
            logger.debug(
                f"Using packaged tokenizer.json fallback for ASR: {packaged_tokenizer_path}"
            )
            tokenizer_file_path = packaged_tokenizer_path
        else:
            raise FileNotFoundError(
                f"No 'tokenizer.json' found for external ASR model. "
                f"Place a 'tokenizer.json' next to the weight file in "
                f"'{tokenizer_dir}'."
            )

    return VibeVoiceASRTextTokenizerFast(tokenizer_file=tokenizer_file_path)


def _load_asr_processor(
    tokenizer,
    preprocessor_config_path: str,
) -> "VibeVoiceASRProcessor":
    """Build a :class:`VibeVoiceASRProcessor` from a tokenizer + preprocessor config.

    Mirrors :meth:`VibeVoiceASRProcessor.from_pretrained` but accepts
    already-resolved local paths instead of a model directory / HF repo.

    Args:
        tokenizer: VibeVoiceASRTextTokenizerFast instance.
        preprocessor_config_path: Path to preprocessor_config.json (may be empty).

    Returns:
        VibeVoiceASRProcessor instance.
    """
    config = {}
    if preprocessor_config_path and os.path.exists(preprocessor_config_path):
        with open(preprocessor_config_path, "r", encoding="utf-8") as f:
            config = json.load(f)

    speech_tok_compress_ratio = config.get("speech_tok_compress_ratio", 3200)
    target_sample_rate = config.get("target_sample_rate", 24000)
    normalize_audio = config.get("normalize_audio", True)

    audio_processor = VibeVoiceTokenizerProcessor(
        sampling_rate=target_sample_rate,
        normalize_audio=normalize_audio,
        target_dB_FS=config.get("target_dB_FS", -25),
        eps=config.get("eps", 1e-6),
    )

    return VibeVoiceASRProcessor(
        tokenizer=tokenizer,
        audio_processor=audio_processor,
        speech_tok_compress_ratio=speech_tok_compress_ratio,
        target_sample_rate=target_sample_rate,
        normalize_audio=normalize_audio,
    )


def _instantiate_asr_model(
    config,
    attn_implementation: str,
    final_load_dtype: torch.dtype,
    use_meta: bool = True,
):
    """Instantiate a :class:`VibeVoiceASRForConditionalGeneration` directly.

    ASR counterpart of :meth:`VibeVoiceLoader._instantiate_model`. Instantiates
    the model class directly (bypassing ``from_pretrained``) so the state dict
    can be loaded in-memory afterwards. By default construction runs under a
    ``torch.device("meta")`` context (plan 2026-08-18, D1) — zero RAM, zero
    random init; weights are bound afterwards by ``_apply_state_dict``.

    Args:
        config: VibeVoiceASRConfig instance.
        attn_implementation: Attention implementation string.
        final_load_dtype: torch.dtype for the model.
        use_meta: Construct under a meta device context (default True).

    Returns:
        Model instance (weights not yet loaded).
    """
    # Set attention implementation on the decoder config
    if hasattr(config, "decoder_config"):
        config.decoder_config._attn_implementation = attn_implementation

    # Set dtype on config (version-safe: transformers v5 deprecated torch_dtype)
    set_config_dtype(config, final_load_dtype)
    if hasattr(config, "decoder_config"):
        set_config_dtype(config.decoder_config, final_load_dtype)

    ctx = torch.device("meta") if use_meta else contextlib.nullcontext()
    with ctx:
        return VibeVoiceASRForConditionalGeneration(config)


# ====================================================================
# Weight file state-dict loading (safetensors / bin / gguf dispatch)
# ====================================================================

def _load_gguf_state_dict(weight_path: str, device=None) -> dict:
    """Load a state dict from a ``.gguf`` file via the ``gguf`` package.

    ComfyUI's :func:`comfy.utils.load_torch_file` does not understand the GGUF
    container (it would route ``.gguf`` to ``torch.load`` and fail), so GGUF
    weights are parsed here with :class:`gguf.GGUFReader` and dequantized
    tensor-by-tensor.

    Args:
        weight_path: Absolute path to the ``.gguf`` file.
        device: Target torch device for the tensors (defaults to CPU). The
            patcher owns the single host-to-device transfer, so CPU is the
            normal choice.

    Returns:
        dict mapping tensor name -> torch.Tensor (dequantized, on ``device``).

    Raises:
        RuntimeError: If the ``gguf`` package is not installed.
    """
    if device is None:
        device = torch.device("cpu")

    try:
        import gguf
    except ImportError as e:
        raise RuntimeError(
            "Loading .gguf weights requires the 'gguf' Python package. "
            "Install it with: pip install gguf"
        ) from e

    logger.debug(f"Loading GGUF state dict from: {weight_path}")
    reader = gguf.GGUFReader(weight_path)

    state_dict = {}
    for tensor in reader.tensors:
        # gguf.dequantize returns a correctly-shaped numpy array for all
        # quantization types (F32/F16 pass through unchanged). The array may be
        # read-only (mmap-backed), so copy it to make it writable for torch.
        dequantized = gguf.dequantize(tensor.data, tensor.tensor_type)
        state_dict[tensor.name] = torch.from_numpy(dequantized.copy()).to(device)

    logger.debug(f"Loaded {len(state_dict)} tensors from GGUF file")
    return state_dict


def _load_weight_state_dict(weight_path: str, device) -> dict:
    """Load a state dict from a weight file, dispatching GGUF to the gguf parser.

    Args:
        weight_path: Absolute path to the weight file
            (``.safetensors`` / ``.bin`` / ``.pt`` / ``.gguf``).
        device: Target torch device for the tensors.

    Returns:
        dict mapping tensor name -> torch.Tensor.
    """
    if weight_path.lower().endswith(".gguf"):
        return _load_gguf_state_dict(weight_path, device=device)
    return comfy.utils.load_torch_file(weight_path, device=device)


# ====================================================================
# Quant-resident loading (GGUF raw-block + ConvRot INT8)
# ====================================================================

def _open_gguf_reader(weight_path: str):
    """Open a ``.gguf`` file with :class:`gguf.GGUFReader`.

    Raises:
        RuntimeError: If the ``gguf`` package is not installed.
    """
    try:
        import gguf
    except ImportError as e:
        raise RuntimeError(
            "Loading .gguf weights requires the 'gguf' Python package. "
            "Install it with: pip install gguf"
        ) from e
    return gguf.GGUFReader(weight_path)


def _gguf_kquant_present(reader) -> bool:
    """True when any tensor uses a K-quant format (Q4_K/Q5_K/Q6_K)."""
    from .gguf_quant import _T

    kquants = {_T.Q4_K, _T.Q5_K, _T.Q6_K}
    return any(t.tensor_type in kquants for t in reader.tensors)


# Quant-storage dtypes that must NEVER reach the dense float loader.
_QUANT_STORAGE_DTYPES = frozenset({
    torch.int8, torch.uint8, torch.float8_e4m3fn, torch.float8_e5m2,
})


def _prepare_quantized_safetensors_load(state_dict: dict, quant_map: dict):
    """Split a scanned comfy_quant map into execution strategies (in place).

    - ``convrot=True`` layers  -> module-replacement plan (int8 resident +
      comfy-kitchen kernels); their int8/scale tensors stay in the state dict
      for assign.
    - ``convrot=False`` layers -> dequant-at-load: ``weight = q * per-row
      scale`` materialized back to the declared orig dtype inside
      ``state_dict``; scale + metadata keys removed.

    Returns:
        ``(layer_plan, n_rowwise)``.

    Raises:
        QuantTargetMismatch: On missing/malformed rowwise tensor pairs or an
            undeclared orig dtype.
    """
    from .convrot_quant import (
        QUANT_META_SUFFIX,
        make_convrot_linear,
        resolve_orig_dtype,
    )
    from .quant_common import QuantTargetMismatch

    layer_plan = {}
    n_rowwise = 0
    for prefix, info in quant_map.items():
        if info.convrot:
            layer_plan[prefix] = make_convrot_linear(info)
            continue

        w_key = f"{prefix}.weight"
        s_key = f"{prefix}.weight_scale"
        w = state_dict.get(w_key)
        s = state_dict.get(s_key)
        if w is None or s is None:
            raise QuantTargetMismatch(
                f"Rowwise-quantized layer '{prefix}' is missing its "
                f"'{w_key}' / '{s_key}' tensors in the checkpoint"
            )
        storage_dtype = info.rowwise_dtype or torch.int8
        if w.dtype != storage_dtype and w.dtype != torch.uint8:
            raise QuantTargetMismatch(
                f"Rowwise layer '{prefix}': expected {storage_dtype} storage, "
                f"checkpoint has {w.dtype}"
            )
        if w.dim() != 2 or w.shape[0] != info.out_features \
                or (info.in_features and w.shape[1] != info.in_features):
            raise QuantTargetMismatch(
                f"Rowwise layer '{prefix}': weight shape {tuple(w.shape)} "
                f"disagrees with metadata [out={info.out_features}, "
                f"in={info.in_features}]"
            )
        orig_dtype = resolve_orig_dtype(info.orig_dtype)

        if info.group_size:
            # BLOCKWISE scales: shape [out / gs, in / gs], one scale per
            # gs x gs block (int8-blockwise checkpoints).
            gs = int(info.group_size)
            if w.shape[0] % gs or w.shape[1] % gs:
                raise QuantTargetMismatch(
                    f"Blockwise layer '{prefix}': weight {tuple(w.shape)} "
                    f"not divisible by group_size {gs} on both dims"
                )
            og, ig = w.shape[0] // gs, w.shape[1] // gs
            if tuple(s.shape) != (og, ig):
                raise QuantTargetMismatch(
                    f"Blockwise layer '{prefix}': scale shape "
                    f"{tuple(s.shape)} disagrees with weight "
                    f"{tuple(w.shape)} at group_size {gs} (expected "
                    f"({og}, {ig}))"
                )
            state_dict[w_key] = (
                (
                    w.to(torch.float32).view(og, gs, ig, gs)
                    * s.to(torch.float32).view(og, 1, ig, 1)
                ).reshape(w.shape).to(orig_dtype)
            )
        else:
            # SCALAR (per-tensor, fp8) or PER-ROW [out, 1] scales; both
            # broadcast.
            if not (s.dim() == 0 or (s.dim() == 2 and s.shape[1] == 1
                                     and s.shape[0] in (w.shape[0], 1))):
                raise QuantTargetMismatch(
                    f"Rowwise layer '{prefix}': scale shape {tuple(s.shape)} "
                    f"is neither a scalar nor [out,1] (weight "
                    f"{tuple(w.shape)})"
                )
            # fp32 intermediate: exact rescale before the final cast.
            state_dict[w_key] = (
                (w.to(torch.float32) * s.to(torch.float32)).to(orig_dtype)
            )
        state_dict.pop(s_key, None)
        n_rowwise += 1

    for prefix in quant_map:
        state_dict.pop(f"{prefix}.{QUANT_META_SUFFIX}", None)
    return layer_plan, n_rowwise


def _assert_dense_loadable(state_dict: dict) -> None:
    """Hard-fail when unplanned quantized weights reach the dense loader.

    Assigning int8/uint8/fp8 storages into float parameters either crashes
    cryptically or (fp8) silently misloads with scales ignored. Both are
    worse than a clear error at load time.
    """
    bad = [
        (k, str(v.dtype))
        for k, v in state_dict.items()
        if v.dtype in _QUANT_STORAGE_DTYPES
    ]
    if bad:
        sample = ", ".join(f"{k} [{d}]" for k, d in bad[:5])
        raise ValueError(
            "Checkpoint contains quantized-weight tensors but carries no "
            f"executable quantization metadata ({len(bad)} tensors, e.g. "
            f"{sample}). Loading them as floats would corrupt the model. "
            "Re-export it with *.comfy_quant metadata (comfy-model-tools) "
            "or use the dense/BF16 checkpoint."
        )


def _install_gguf_weights(model, reader) -> dict:
    """Install GGUF weights onto a meta-initialized model WITHOUT float
    materialization (plan 2026-08-24, D1 — the RAM-spike kill).

    Pipeline:
    1. Map every reader tensor onto the model tree (HF pass-through or
       llamacpp naming).
    2. Quantized tensors whose target is an ``nn.Linear`` become RESIDENTS:
       the Linear is swapped for :class:`~modules.gguf_quant.GGUFLinear` and
       the RAW BLOCK BYTES are installed as its uint8 parameter (the one
       unavoidable copy — it IS the residency).
    3. Float tensors (F32/F16/BF16) take zero-copy views at their NATIVE
       dtype into a filtered dense state dict applied through
       :meth:`VibeVoiceLoader._apply_state_dict` (assign semantics, re-tie,
       sentinel/RoPE fixes preserved).

    Peak RAM ≈ raw file size instead of ~2x full-float size.

    Returns:
        Stats dict ``{"n_resident_layers", "raw_bytes", "weight_family"}``.
    """
    import numpy as np

    from .gguf_quant import (
        FLOAT_GGML_TYPES,
        GGUFTensor,
        SUPPORTED_GGML_TYPES,
        UnsupportedGGMLType,
        gguf_linear_factory,
        map_keys,
    )
    from .quant_common import (
        QuantTargetMismatch,
        replace_linears_for_quant,
        resolve_module,
    )

    tensors = list(reader.tensors)
    mapping = map_keys([t.name for t in tensors])
    modules_by_name = dict(model.named_modules())

    resident_plan = {}
    resident_tensors = {}
    dense_state = {}
    unsupported = []

    for t in tensors:
        target_key = mapping[t.name]
        logical_shape = tuple(int(s) for s in reversed(t.shape))
        tt = t.tensor_type

        if tt in FLOAT_GGML_TYPES:
            data = t.data
            if data.dtype == np.float32:
                tensor = torch.from_numpy(data)
            elif data.dtype == np.float16:
                tensor = torch.from_numpy(data)
            else:  # BF16 arrives as raw uint8 bytes
                tensor = torch.from_numpy(data).view(torch.bfloat16)
            dense_state[target_key] = tensor.reshape(logical_shape)
            continue

        if tt not in SUPPORTED_GGML_TYPES:
            unsupported.append((t.name, tt))
            continue

        if not target_key.endswith(".weight"):
            raise QuantTargetMismatch(
                f"Quantized GGUF tensor '{t.name}' maps to '{target_key}', "
                f"which is not a Linear weight; this checkpoint cannot be "
                f"executed with quant-resident linears."
            )
        module_path = target_key[: -len(".weight")]
        target_module = modules_by_name.get(module_path)
        if target_module is None or not isinstance(target_module, torch.nn.Linear):
            kind = type(target_module).__name__ if target_module is not None else "missing"
            raise QuantTargetMismatch(
                f"GGUF-quantized tensor '{t.name}' maps to '{module_path}' "
                f"({kind}), expected an nn.Linear. The sidecar config may not "
                f"match this checkpoint."
            )
        expected_shape = tuple(target_module.weight.shape)
        if logical_shape != expected_shape:
            raise QuantTargetMismatch(
                f"GGUF tensor '{t.name}' shape {logical_shape} disagrees with "
                f"model {type(target_module).__name__} shape {expected_shape}"
            )
        # Copy the raw block bytes off the mmap now (this copy IS the final
        # residency); float materialization never happens.
        resident_tensors[module_path] = GGUFTensor.from_reader_tensor(t)
        resident_plan[module_path] = gguf_linear_factory(tt)

    if unsupported:
        name, tt = unsupported[0]
        raise UnsupportedGGMLType(tt, f"{name}" + (
            f" (+{len(unsupported) - 1} more)" if len(unsupported) > 1 else ""
        ))

    replaced = replace_linears_for_quant(model, resident_plan)
    total_raw = 0
    for module_path in replaced:
        module = resolve_module(model, module_path)
        gtensor = resident_tensors[module_path]
        module.set_raw_weight(gtensor.raw)
        total_raw += gtensor.raw.numel()
        resident_tensors[module_path] = None

    known_missing = {f"{p}.weight" for p in replaced}
    VibeVoiceLoader._apply_state_dict(model, dense_state, known_missing=known_missing)

    return {
        "n_resident_layers": len(replaced),
        "raw_bytes": total_raw,
        "weight_family": "gguf_block",
    }


# ====================================================================
# Low-bit / naive quantization detection (defensive warning)
# ====================================================================

# safetensors dtypes that indicate integer (quantized) storage.
_SAFETENSORS_INT_DTYPES = {
    "I8", "U8", "I16", "U16", "I32", "U32", "I64", "U64", "I4", "U4",
}

# Tensor-name substrings that indicate proper dequantization metadata is present
# (a real quantized checkpoint stores per-channel scales / zero-points).
_QUANT_SCALE_NAME_HINTS = ("scale", "zero_point", "qzero", "absmax")

# Fraction of integer/low-bit tensors above which we consider the file
# "mostly quantized" (as opposed to a couple of incidental index tensors).
_LOWBIT_FRACTION_THRESHOLD = 0.10


def _inspect_safetensors_quantization(weight_path: str):
    """Parse only the safetensors header (no tensor data) to detect naive int casts.

    A *proper* quantized checkpoint stores dequantization metadata (per-channel
    scale / zero-point tensors). A *naive* ``int8`` cast stores raw integer values
    with no scales, which is unusable (loads as garbage weights).

    Args:
        weight_path: Absolute path to a ``.safetensors`` file.

    Returns:
        dict ``{"total", "int_count", "has_scale_meta"}`` or ``None`` if the
        header could not be parsed.
    """
    import struct

    try:
        with open(weight_path, "rb") as f:
            (header_len,) = struct.unpack("<Q", f.read(8))
            header = json.loads(f.read(header_len))
    except Exception:
        return None

    header.pop("__metadata__", None)
    total = len(header)
    int_count = 0
    has_scale_meta = False
    for name, info in header.items():
        if info.get("dtype", "") in _SAFETENSORS_INT_DTYPES:
            int_count += 1
        low = name.lower()
        if any(hint in low for hint in _QUANT_SCALE_NAME_HINTS):
            has_scale_meta = True

    return {"total": total, "int_count": int_count, "has_scale_meta": has_scale_meta}


def _inspect_gguf_quantization(weight_path: str):
    """Inspect the GGUF tensor table (no dequantization) to count sub-4-bit I-quants.

    Args:
        weight_path: Absolute path to a ``.gguf`` file.

    Returns:
        dict ``{"total", "low_bit_count"}`` or ``None`` if the file could not be
        read or the ``gguf`` package is unavailable.
    """
    try:
        import gguf
        from gguf.constants import GGMLQuantizationType
    except ImportError:
        return None

    try:
        reader = gguf.GGUFReader(weight_path)
    except Exception:
        return None

    # Aggressive sub-4-bit importance-matrix quants. These are faithful but very
    # lossy; on a small (~1.5B) TTS LM they commonly degrade below the threshold
    # needed for text-conditioned generation.
    low_bit_types = {
        GGMLQuantizationType.IQ1_S,
        GGMLQuantizationType.IQ1_M,
        GGMLQuantizationType.IQ2_XXS,
        GGMLQuantizationType.IQ2_XS,
        GGMLQuantizationType.IQ2_S,
        GGMLQuantizationType.IQ3_XXS,
        GGMLQuantizationType.IQ3_S,
    }
    total = len(reader.tensors)
    low_bit_count = sum(1 for t in reader.tensors if t.tensor_type in low_bit_types)
    return {"total": total, "low_bit_count": low_bit_count}


def warn_if_lowbit_quantization(weight_path: str) -> None:
    """Log a warning if ``weight_path`` looks over-quantized or naively cast.

    This is a defensive, load-time heuristic. It does not block loading; it only
    surfaces a clear, actionable warning so the user knows a garbled / silent
    result is likely caused by the checkpoint rather than the node.

    Two failure modes are detected:
      * safetensors with many raw integer tensors but NO scale/zero-point metadata
        -> naive int cast (unusable).
      * GGUF with many sub-4-bit I-quant tensors -> aggressive quant that may
        degrade TTS quality (garbled output / reference echo).

    Args:
        weight_path: Absolute path to the external weight file.
    """
    lower = weight_path.lower()

    if lower.endswith(".safetensors"):
        info = _inspect_safetensors_quantization(weight_path)
        if info is None or info["total"] == 0:
            return
        frac = info["int_count"] / info["total"]
        if frac >= _LOWBIT_FRACTION_THRESHOLD and not info["has_scale_meta"]:
            logger.warning(
                f"[low-bit check] '{os.path.basename(weight_path)}' stores "
                f"{info['int_count']}/{info['total']} tensors as raw integers with "
                f"NO dequantization scale/zero-point metadata. This looks like a "
                f"naive int cast, which loads as garbage weights and will produce "
                f"silent or broken audio. Use a proper quantized or full-precision "
                f"(BF16/FP16) checkpoint instead."
            )
        return

    if lower.endswith(".gguf"):
        info = _inspect_gguf_quantization(weight_path)
        if info is None or info["total"] == 0:
            return
        frac = info["low_bit_count"] / info["total"]
        if frac >= _LOWBIT_FRACTION_THRESHOLD:
            logger.warning(
                f"[low-bit check] '{os.path.basename(weight_path)}' contains "
                f"{info['low_bit_count']}/{info['total']} tensors at sub-4-bit "
                f"I-quant precision. This is very aggressive for a small TTS LM and "
                f"may degrade below the threshold for text-conditioned generation "
                f"(symptoms: garbled syllables, reference-audio echo, never "
                f"terminating). Consider a higher-quality quant (Q4_K_M / Q5_K_M / "
                f"Q8_0 / BF16)."
            )
        return


# ====================================================================
# In-memory state dict loading
# ====================================================================

def _load_state_dict_into_model_from_memory(model, state_dict: dict):
    """Load an in-memory state dict into an already-instantiated model.

    Mirrors :meth:`VibeVoiceLoader._load_state_dict_into_model` but accepts a
    state dict that is already in memory (loaded via
    ``comfy.utils.load_torch_file``) instead of resolving a checkpoint path.

    Delegates to :meth:`VibeVoiceLoader._apply_state_dict` (plan 2026-08-18,
    D2/D3): assign semantics (no copy pass), tied-weight re-tie, and
    zero-materialization of meta stragglers.

    Args:
        model: Instantiated model (weights not yet loaded).
        state_dict: State dict mapping (on CPU).

    Returns:
        The model with the state dict loaded (still on CPU).
    """
    VibeVoiceLoader._apply_state_dict(model, state_dict)
    return model


# ====================================================================
# Core external loading function
# ====================================================================

def load_external_vibevoice_model(
    weight_path: str,
    config_name: str,
    attention_mode: str = "eager",
    use_llm_4bit: bool = False,
    dtype_str: str = "auto",
    device=None,
) -> dict:
    """Load a VibeVoice model from an external weight file.

    The model is built entirely on CPU. The single host-to-device transfer is
    owned by :class:`~modules.patcher.VibeVoicePatcher` after ComfyUI's VRAM
    arbitration (same contract as the standard loader path).

    Args:
        weight_path: Absolute path to the external weight file
            (safetensors / .bin / .gguf).
        config_name: Architecture config selector. Used for the packaged-default
            fallback when no sidecar config is present. One of
            ``EXTERNAL_CONFIG_OPTIONS``; ``AUTO_CONFIG_NAME`` resolves the
            family from the weights' embedding fingerprint.
        attention_mode: Attention implementation
            ("eager", "sdpa", "flash_attention_2", "sage").
        use_llm_4bit: Whether to quantize the LLM to 4-bit NF4.
        dtype_str: Dtype string ("auto", "bf16", "fp16", "fp32").
        device: Optional target device hint for dtype-auto resolution.

    Returns:
        The ``VIBEVOICE_MODEL`` bundle dict::

            {
                "state_dict": dict[str, torch.Tensor],
                "config": VibeVoiceConfig | VibeVoiceStreamingConfig,
                "processor": VibeVoiceProcessor | VibeVoiceStreamingProcessor,
                "model": torch.nn.Module,
                "model_name": str,
                "source_path": str,
                "is_streaming": bool,
            }

    Raises:
        FileNotFoundError: If the weight file or config cannot be found.
        ValueError: If ``AUTO_CONFIG_NAME`` is selected but the architecture
            cannot be determined from the weight file.
        RuntimeError: If model instantiation or loading fails.
    """
    # Legacy alias normalization (plan 2026-08-27, D1): saved workflows may
    # still carry removed option values (e.g. "VibeVoice-Large").
    config_name = normalize_config_name(config_name)

    if not os.path.isfile(weight_path):
        raise FileNotFoundError(f"External weight file not found: {weight_path}")

    # File-identity stamp for the cache key (plan 2026-08-20, B1): captured
    # once at entry so the bundle records exactly which bytes were loaded.
    try:
        _stat = os.stat(weight_path)
        source_mtime_ns = _stat.st_mtime_ns
        source_size = _stat.st_size
    except OSError:
        source_mtime_ns = 0
        source_size = 0

    # Defensive: warn early if the checkpoint looks over-quantized / naively cast.
    warn_if_lowbit_quantization(weight_path)

    # ASR models use a distinct loading branch (different config / tokenizer /
    # processor / model classes, and no 4-bit quantization).
    if is_asr_config_name(config_name):
        return load_external_vibevoice_asr_model(
            weight_path=weight_path,
            config_name=config_name,
            attention_mode=attention_mode,
            dtype_str=dtype_str,
            device=device,
        )

    # Resolve attention mode with fallback logic (same as standard loader)
    attention_mode = resolve_attention_mode(attention_mode, use_llm_4bit)

    # Weight-plan validation + lazy source opening (plan 2026-08-24, D1/D2/D4).
    # GGUF keeps an mmap READER (no float materialization); plain safetensors
    # are scanned for ConvRot *.comfy_quant metadata BEFORE any load.
    lower_path = weight_path.lower()
    is_gguf_file = lower_path.endswith(".gguf")
    gguf_reader = _open_gguf_reader(weight_path) if is_gguf_file else None
    convrot_quant_map = {}
    if gguf_reader is None and lower_path.endswith(".safetensors"):
        # Scan-only pass. UNREADABLE headers degrade to "no quant metadata"
        # (a genuinely corrupt file still fails at the actual load), but an
        # UNSUPPORTED QUANT FORMAT must propagate: silently falling through
        # to the dense loader misloads int8/fp8 weights as floats.
        try:
            from .convrot_quant import scan_checkpoint_quantization

            convrot_quant_map = scan_checkpoint_quantization(weight_path)
        except UnsupportedQuantFormat:
            raise
        except Exception as e:
            logger.debug(f"ConvRot scan skipped for {weight_path}: {e}")
            convrot_quant_map = {}
    validate_weight_plan(
        is_gguf_file=is_gguf_file,
        convrot_quant_map=convrot_quant_map,
        use_llm_4bit=use_llm_4bit,
        attention_mode=attention_mode,
        gguf_kquant_present=_gguf_kquant_present(gguf_reader) if gguf_reader else False,
    )

    # Architecture fingerprint (plan 2026-08-27, D3): header-only for
    # safetensors, reuses the open GGUF reader. Computed BEFORE the heavy
    # state-dict load so Auto-detect can fail fast (D6) and reconciliation
    # (D5) reuses it below without re-reading the file.
    from .config_detect import fingerprint_weights

    weights_fp = fingerprint_weights(weight_path, gguf_reader=gguf_reader)

    # Auto-detect (D6): resolve the architecture from the fingerprint BEFORE
    # any heavy work. Conclusive -> adopt the detected family; inconclusive
    # -> actionable ValueError (fail-fast, no partial load). Explicit
    # selections skip this and self-heal via reconciliation below (D5).
    if config_name == AUTO_CONFIG_NAME:
        config_name = resolve_auto_config_name(weight_path, weights_fp=weights_fp)

    cpu_device = torch.device("cpu")
    state_dict = None
    if gguf_reader is not None:
        logger.info(
            f"Opening external VibeVoice GGUF weights (raw-block residency): "
            f"{weight_path}"
        )
    else:
        logger.info(f"Loading external VibeVoice weights from: {weight_path}")
        # Step 1: Load the state dict onto CPU (safetensors/bin via ComfyUI's
        # loader). Always CPU — the patcher owns the single H2D transfer.
        state_dict = _load_weight_state_dict(weight_path, cpu_device)

    # Step 2: Resolve and load the architecture config, reconciled against
    # the weights' architecture fingerprint (plan 2026-08-27, D5). When the
    # fingerprint contradicts the selected config family, the detected
    # family's packaged config is substituted with a WARNING.
    config_path = resolve_sidecar_config(weight_path, config_name)
    config = VibeVoiceLoader._load_config(config_path, config_name)
    effective_name, config_changed = reconcile_config(config_name, config, weights_fp)
    if config_changed:
        config_name = effective_name
        config_path = resolve_sidecar_config(weight_path, config_name)
        config = VibeVoiceLoader._load_config(config_path, config_name)

    # Step 3: Detect streaming from the loaded config.
    is_streaming = isinstance(config, VibeVoiceStreamingConfig)
    if is_streaming:
        logger.debug(f"External model '{config_name}' detected as streaming model")

    # Step 4: Resolve and load the tokenizer.
    tokenizer_dir = resolve_sidecar_tokenizer_dir(weight_path)
    vibevoice_tokenizer = VibeVoiceLoader._load_tokenizer(tokenizer_dir, config_name)

    # Step 5: Resolve and load the processor.
    preprocessor_path = resolve_sidecar_preprocessor(weight_path)
    processor = VibeVoiceLoader._load_processor(
        vibevoice_tokenizer, preprocessor_path, is_streaming=is_streaming
    )

    # Step 6: Resolve dtype + attention implementation.
    load_device = (
        model_management.get_torch_device()
        if not isinstance(device, torch.device)
        else device
    )
    model_dtype = resolve_dtype(dtype_str, load_device)

    quant_config = None
    final_load_dtype = model_dtype
    if use_llm_4bit:
        bnb_compute_dtype = model_dtype
        if attention_mode == "sage":
            bnb_compute_dtype, final_load_dtype = torch.float32, torch.float32
        quant_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_use_double_quant=True,
            bnb_4bit_compute_dtype=bnb_compute_dtype,
        )

    attn_implementation_for_load = get_attn_implementation_for_load(attention_mode)

    try:
        logger.debug(
            f"Instantiating external VibeVoice model '{config_name}' with "
            f"dtype={final_load_dtype}, attention='{attn_implementation_for_load}'"
        )

        # Step 7: Instantiate the model class directly (no meta device init).
        model = VibeVoiceLoader._instantiate_model(
            config=config,
            is_streaming=is_streaming,
            attn_implementation=attn_implementation_for_load,
            final_load_dtype=final_load_dtype,
        )

        # Step 8: Bind the weights — dispatch by weight family:
        # - gguf_block: raw-block residency install (no float materialization)
        # - convrot_int8: linears already swapped; int8/scale assign directly
        # - dense: legacy in-memory assign of the full float state dict
        weight_family = "dense"
        quant_stats = {}
        if gguf_reader is not None:
            quant_stats = _install_gguf_weights(model, gguf_reader)
            weight_family = "gguf_block"
            del gguf_reader
            gguf_reader = None
        elif convrot_quant_map:
            from .quant_common import replace_linears_for_quant

            layer_plan, n_rowwise = _prepare_quantized_safetensors_load(
                state_dict, convrot_quant_map
            )
            replaced = replace_linears_for_quant(model, layer_plan)
            model = _load_state_dict_into_model_from_memory(model, state_dict)
            del state_dict
            state_dict = None
            weight_family = "convrot_int8"
            quant_stats = {
                "n_resident_layers": len(replaced),
                "n_rowwise_layers": n_rowwise,
            }
        else:
            # Defensive net: unplanned int8/uint8/fp8 weights must never
            # reach the float loader (crash or silent corruption).
            _assert_dense_loadable(state_dict)
            model = _load_state_dict_into_model_from_memory(model, state_dict)

            # Free the state dict immediately — the model now owns the tensors
            # (assign semantics), so the dict is dead weight (plan D7/RC-4).
            del state_dict
            state_dict = None

        # Step 9: Apply the final dtype ON CPU (single H2D owned by patcher) —
        # only where needed (plan 2026-08-18, D4/RC-3).
        cast_model_to_dtype_if_needed(model, final_load_dtype)

        # Step 10: Apply 4-bit quantization if requested (post-load).
        if quant_config is not None:
            try:
                from transformers.integrations.bitsandbytes import replace_with_bnb_linear
            except ImportError:
                from transformers.utils.bitsandbytes import replace_with_bnb_linear
            replace_with_bnb_linear(
                model,
                quantization_config=quant_config,
                modules_to_not_convert=None,
            )

        # Step 11: Apply SageAttention post-load if requested.
        if attention_mode == "sage":
            if check_sage_attention_compatible():
                set_sage_attention(model)
            else:
                raise RuntimeError("Incompatible hardware/setup for SageAttention.")

        # Step 12: Eval mode.
        model.eval()
        setattr(model, "_llm_4bit", bool(quant_config))

        logger.info(
            f"Successfully loaded external VibeVoice model '{config_name}' "
            f"from {weight_path}"
        )

        # Step 13: Assemble the bundle.
        # NOTE (plan 2026-08-18, D7/RC-4): the state dict is intentionally NOT
        # included — the model owns the tensors after assign-loading, and
        # retaining a second ~model-size copy in the node output is dead RAM.
        #
        # Plan 2026-08-20 (B1/P1): the bundle also records the full build
        # identity — file stat, RESOLVED attention mode, 4-bit flag, and the
        # requested dtype string — so consumers can derive a cache key that is
        # byte-for-byte equal to the loader-node's pre-build request identity.
        return {
            "config": config,
            "processor": processor,
            "model": model,
            "model_name": config_name,
            "source_path": weight_path,
            "source_mtime_ns": source_mtime_ns,
            "source_size": source_size,
            "attention_mode": attention_mode,
            "use_llm_4bit": bool(use_llm_4bit),
            "dtype_str": dtype_str,
            "is_streaming": is_streaming,
            "is_asr": False,
            "weight_family": weight_family,
            "quant_stats": quant_stats,
        }

    except Exception as e:
        logger.error(
            f"Failed to load external VibeVoice model '{config_name}' "
            f"from {weight_path}: {e}"
        )
        raise RuntimeError(
            f"Failed to load external VibeVoice model '{config_name}': {e}"
        )


# ====================================================================
# ASR external loading branch
# ====================================================================

def load_external_vibevoice_asr_model(
    weight_path: str,
    config_name: str,
    attention_mode: str = "sdpa",
    dtype_str: str = "auto",
    device=None,
) -> dict:
    """Load a VibeVoice ASR model from an external weight file.

    ASR counterpart of :func:`load_external_vibevoice_model`. ASR models use a
    distinct config (:class:`VibeVoiceASRConfig`), tokenizer
    (:class:`VibeVoiceASRTextTokenizerFast`), processor
    (:class:`VibeVoiceASRProcessor`), and model class
    (:class:`VibeVoiceASRForConditionalGeneration`), and do not support 4-bit
    quantization.

    The model is built entirely on CPU. The single host-to-device transfer is
    owned by :class:`~modules.patcher.VibeVoiceASRPatcher` after ComfyUI's VRAM
    arbitration (same contract as the standard ASR loader path).

    Args:
        weight_path: Absolute path to the external weight file
            (safetensors / .bin / .gguf).
        config_name: Architecture config selector (e.g. ``"VibeVoice-ASR"``).
        attention_mode: Attention implementation
            ("eager", "sdpa", "flash_attention_2", "sage").
        dtype_str: Dtype string ("auto", "bf16", "fp16", "fp32").
        device: Optional target device hint for dtype-auto resolution.

    Returns:
        The ``VIBEVOICE_MODEL`` bundle dict::

            {
                "state_dict": dict[str, torch.Tensor],
                "config": VibeVoiceASRConfig,
                "processor": VibeVoiceASRProcessor,
                "model": VibeVoiceASRForConditionalGeneration,
                "model_name": str,
                "source_path": str,
                "is_streaming": False,
                "is_asr": True,
            }

    Raises:
        FileNotFoundError: If the weight file, config, or tokenizer cannot be found.
        RuntimeError: If model instantiation or loading fails.
    """
    # Legacy alias normalization (plan 2026-08-27, D1) — mirrors the TTS branch.
    config_name = normalize_config_name(config_name)

    if not os.path.isfile(weight_path):
        raise FileNotFoundError(f"External weight file not found: {weight_path}")

    # File-identity stamp (plan 2026-08-20, B1) — mirrors the TTS branch.
    try:
        _stat = os.stat(weight_path)
        source_mtime_ns = _stat.st_mtime_ns
        source_size = _stat.st_size
    except OSError:
        source_mtime_ns = 0
        source_size = 0

    # Defensive: warn early if the checkpoint looks over-quantized / naively cast.
    warn_if_lowbit_quantization(weight_path)

    # Resolve attention mode with fallback logic (no 4-bit for ASR).
    attention_mode = resolve_attention_mode(attention_mode, quantize_4bit=False)

    # Weight-plan validation + lazy source opening (mirrors the TTS branch).
    lower_path = weight_path.lower()
    is_gguf_file = lower_path.endswith(".gguf")
    gguf_reader = _open_gguf_reader(weight_path) if is_gguf_file else None
    convrot_quant_map = {}
    if gguf_reader is None and lower_path.endswith(".safetensors"):
        # Scan-only pass. UNREADABLE headers degrade to "no quant metadata"
        # (a genuinely corrupt file still fails at the actual load), but an
        # UNSUPPORTED QUANT FORMAT must propagate: silently falling through
        # to the dense loader misloads int8/fp8 weights as floats.
        try:
            from .convrot_quant import scan_checkpoint_quantization

            convrot_quant_map = scan_checkpoint_quantization(weight_path)
        except UnsupportedQuantFormat:
            raise
        except Exception as e:
            logger.debug(f"ConvRot scan skipped for {weight_path}: {e}")
            convrot_quant_map = {}
    validate_weight_plan(
        is_gguf_file=is_gguf_file,
        convrot_quant_map=convrot_quant_map,
        use_llm_4bit=False,
        attention_mode=attention_mode,
        gguf_kquant_present=_gguf_kquant_present(gguf_reader) if gguf_reader else False,
    )

    # Step 1: Load the state dict onto CPU (safetensors/bin via ComfyUI's
    # loader). GGUF keeps a lazy mmap READER instead (raw-block residency).
    cpu_device = torch.device("cpu")
    state_dict = None
    if gguf_reader is not None:
        logger.info(
            f"Opening external VibeVoice ASR GGUF weights (raw-block "
            f"residency): {weight_path}"
        )
    else:
        logger.info(f"Loading external VibeVoice ASR weights from: {weight_path}")
        state_dict = _load_weight_state_dict(weight_path, cpu_device)

    # Step 2: Resolve and load the ASR architecture config.
    config_path = resolve_sidecar_config(weight_path, config_name)
    config = _load_asr_config(config_path)

    # Step 3: Resolve and load the ASR tokenizer.
    tokenizer_dir = resolve_sidecar_tokenizer_dir(weight_path)
    asr_tokenizer = _load_asr_tokenizer(tokenizer_dir)

    # Step 4: Resolve and load the ASR processor.
    preprocessor_path = resolve_sidecar_preprocessor(weight_path)
    processor = _load_asr_processor(asr_tokenizer, preprocessor_path)

    # Step 5: Resolve dtype + attention implementation.
    load_device = (
        model_management.get_torch_device()
        if not isinstance(device, torch.device)
        else device
    )
    model_dtype = resolve_dtype(dtype_str, load_device)
    attn_implementation_for_load = get_attn_implementation_for_load(attention_mode)

    try:
        logger.debug(
            f"Instantiating external VibeVoice ASR model '{config_name}' with "
            f"dtype={model_dtype}, attention='{attn_implementation_for_load}'"
        )

        # Step 6: Instantiate the ASR model class directly (no meta device init).
        model = _instantiate_asr_model(
            config=config,
            attn_implementation=attn_implementation_for_load,
            final_load_dtype=model_dtype,
        )

        # Step 7: Bind the weights (same dispatch as the TTS branch).
        weight_family = "dense"
        quant_stats = {}
        if gguf_reader is not None:
            quant_stats = _install_gguf_weights(model, gguf_reader)
            weight_family = "gguf_block"
            del gguf_reader
            gguf_reader = None
        elif convrot_quant_map:
            from .quant_common import replace_linears_for_quant

            layer_plan, n_rowwise = _prepare_quantized_safetensors_load(
                state_dict, convrot_quant_map
            )
            replaced = replace_linears_for_quant(model, layer_plan)
            model = _load_state_dict_into_model_from_memory(model, state_dict)
            del state_dict
            state_dict = None
            weight_family = "convrot_int8"
            quant_stats = {
                "n_resident_layers": len(replaced),
                "n_rowwise_layers": n_rowwise,
            }
        else:
            # Defensive net: unplanned int8/uint8/fp8 weights must never
            # reach the float loader (crash or silent corruption).
            _assert_dense_loadable(state_dict)
            model = _load_state_dict_into_model_from_memory(model, state_dict)

            # Free the state dict immediately — the model now owns the tensors
            # (assign semantics), so the dict is dead weight (plan D7/RC-4).
            del state_dict
            state_dict = None

        # Step 8: Apply the final dtype ON CPU (single H2D owned by patcher) —
        # only where needed (plan 2026-08-18, D4/RC-3).
        cast_model_to_dtype_if_needed(model, model_dtype)

        # Step 9: Apply SageAttention post-load if requested.
        if attention_mode == "sage":
            if check_sage_attention_compatible():
                set_sage_attention(model)
            else:
                raise RuntimeError("Incompatible hardware/setup for SageAttention.")

        # Step 10: Eval mode.
        model.eval()

        logger.info(
            f"Successfully loaded external VibeVoice ASR model '{config_name}' "
            f"from {weight_path}"
        )

        # Step 11: Assemble the bundle (identity fields mirror the TTS bundle,
        # plan 2026-08-20 B1/P1).
        # NOTE (plan 2026-08-18, D7/RC-4): the state dict is intentionally NOT
        # included — the model owns the tensors after assign-loading, and
        # retaining a second ~model-size copy in the node output is dead RAM.
        return {
            "config": config,
            "processor": processor,
            "model": model,
            "model_name": config_name,
            "source_path": weight_path,
            "source_mtime_ns": source_mtime_ns,
            "source_size": source_size,
            "attention_mode": attention_mode,
            "use_llm_4bit": False,
            "dtype_str": dtype_str,
            "is_streaming": False,
            "is_asr": True,
            "weight_family": weight_family,
            "quant_stats": quant_stats,
        }

    except Exception as e:
        logger.error(
            f"Failed to load external VibeVoice ASR model '{config_name}' "
            f"from {weight_path}: {e}"
        )
        raise RuntimeError(
            f"Failed to load external VibeVoice ASR model '{config_name}': {e}"
        )