nodes/wildminder-ComfyUI-VibeVoice
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
de6ee6bf81ba426a8a20e362fd26f2d58ea68950
wildminder-ComfyUI-VibeVoice/modules/loader.py
T
WildAi fddc0110e4 fix(loader): match the rooted encoder namespace 
Real keys hang off model.acoustic_tokenizer.encoder.*, so the
bare prefix matched nothing and every realtime load warned about
the 276 keys the rule exists to silence.
2026-10-06 01:49:23 +03:00

797 lines
33 KiB
Python
Raw Blame History

"""Model loading and caching for VibeVoice models.
Handles:
- Model path resolution (official, local_dir, standalone)
- Config loading (with fallback to packaged defaults)
- Tokenizer loading (with download fallback)
- Processor setup
- Model instantiation with dtype/attention/quantization settings
- Model caching by cache key
"""
import os
import json
import gc
import logging
import contextlib
import torch
import comfy.utils
import folder_paths
import comfy.model_management as model_management
from transformers import BitsAndBytesConfig
from ..src.vibevoice.modular.configuration_vibevoice import VibeVoiceConfig
from ..src.vibevoice.modular.configuration_vibevoice_streaming import VibeVoiceStreamingConfig
from ..src.vibevoice.modular.modeling_vibevoice import VibeVoiceForConditionalGeneration
from ..src.vibevoice.modular.modeling_vibevoice_streaming_inference import VibeVoiceStreamingForConditionalGenerationInference
from ..src.vibevoice.processor.vibevoice_processor import VibeVoiceProcessor
from ..src.vibevoice.processor.vibevoice_streaming_processor import VibeVoiceStreamingProcessor
from ..src.vibevoice.processor.vibevoice_tokenizer_processor import VibeVoiceTokenizerProcessor
from ..src.vibevoice.modular.modular_vibevoice_text_tokenizer import VibeVoiceTextTokenizerFast
from .model_info import AVAILABLE_VIBEVOICE_MODELS, MODEL_CONFIGS
from .base_loader import BaseVibeVoiceLoader, place_tensor_on_device
from .attention_utils import (
SAGE_ATTENTION_AVAILABLE,
ATTENTION_MODES,
resolve_attention_mode,
get_attn_implementation_for_load,
check_sage_attention_compatible,
)
from .dtype_utils import resolve_dtype, get_dtype_str, cast_model_to_dtype_if_needed, set_config_dtype
if SAGE_ATTENTION_AVAILABLE:
from ..src.vibevoice.modular.sage_attention_patch import set_sage_attention
from huggingface_hub import hf_hub_download
# Sentinel buffer values that must be restored to their intended initial state
# after meta-init materialization (regression fix, v2.3.1). The VibeVoice model
# registers ``speech_scaling_factor`` / ``speech_bias_factor`` as
# ``torch.tensor(float('nan'))`` and computes them at inference time
# (modeling_vibevoice.py:324-338). The diffusion-inversion gate at
# modeling_vibevoice.py:623 is ``if not torch.isnan(sf) and not torch.isnan(bf)``
# — zeroing these buffers makes the gate TRUE and applies ``speech / 0 - 0``,
# producing silent/garbage output. ``fix_std`` (persistent=False) is likewise
# restored to its config value (default 0.5 for the acoustic tokenizer).
#
# Keys are matched by the buffer's FINAL dotted component (e.g.
# ``speech_scaling_factor`` matches ``model.speech_scaling_factor``), so the
# lookup is robust to submodule nesting.
_SENTINEL_BUFFER_VALUES = {
"speech_scaling_factor": float("nan"),
"speech_bias_factor": float("nan"),
"fix_std": 0.5,
}
# Subtrees a released checkpoint may omit without indicating a broken load.
#
# VibeVoice-Realtime-0.5B ships an acoustic-tokenizer DECODER only (605 keys,
# 276 of them the decoder). Its encoder is not part of the released TTS weights
# and the generate path never calls it — it exists so the tokenizer can encode
# audio for voice-cloning / training pipelines. Vanilla
# `from_pretrained` reports the same 276 keys as MISSING for that checkpoint,
# which is the reference behaviour; we match it instead of printing an
# alarming warning on every realtime load, where it could mask a genuinely
# missing key.
#
# Prefixes are matched against the checkpoint's own key names, which are
# rooted at the module the tensors hang off — for every released tree that is
# ``model.`` (``VibeVoiceStreamingForConditionalGenerationInference.model
# .acoustic_tokenizer.encoder.*``). The unrooted spelling is kept so a
# tokenizer exported at the top level stays optional too, but on the released
# checkpoints it matches nothing.
#
# A prefix is only treated as optional when the checkpoint supplied NO key
# under it. If any encoder key is present, the rest are reported normally.
OPTIONAL_ABSENT_PREFIXES = (
"model.acoustic_tokenizer.encoder.",
"acoustic_tokenizer.encoder.",
)
# Quant-storage dtypes that must NEVER reach a dense float loader.
#
# Single source of truth for both dense routes: the batch check
# (``external_loader._assert_dense_loadable``) and the per-tensor streaming
# assign (``VibeVoiceLoader._stream_apply_dense``). ``external_loader``
# imports this name from here (it already imports ``VibeVoiceLoader``), so the
# dependency runs one way only — ``loader`` must never import from
# ``external_loader``, which would be a cycle.
#
# 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.
QUANT_STORAGE_DTYPES = frozenset({
torch.int8, torch.uint8, torch.float8_e4m3fn, torch.float8_e5m2,
})
def mark_optional_absent(
missing_keys: list[str],
assigned: set[str],
known_missing: set[str] | None = None,
) -> set[str]:
"""Add deliberately-omitted keys to ``known_missing``.
A prefix is only treated as optional when the checkpoint supplied NO key
under it. If even one key arrived, the remaining gaps under that prefix are
a real finding and stay reported.
"""
known = set(known_missing or ())
seen = {
prefix
for prefix in OPTIONAL_ABSENT_PREFIXES
if any(key.startswith(prefix) for key in assigned)
}
for key in missing_keys:
for prefix in OPTIONAL_ABSENT_PREFIXES:
if key.startswith(prefix) and prefix not in seen:
known.add(key)
break
return known
def _recompute_rope_buffers(model) -> int:
"""Recompute RoPE ``inv_freq`` buffers destroyed by meta-init (v2.3.2 fix).
Rotary embeddings compute ``inv_freq`` / ``original_inv_freq`` from the
config inside ``__init__`` — they are **non-persistent buffers that never
appear in the checkpoint**. Under ``torch.device("meta")`` instantiation
they are created on the meta device, and the generic zero-materialization
in :meth:`VibeVoiceLoader._apply_state_dict` turns them into all-zeros.
With ``inv_freq == 0`` RoPE yields ``cos(0)=1`` / ``sin(0)=0``, i.e. **no
positional encoding**: the language model cannot order tokens, so it emits
gibberish syllables and hits the speech-end token almost immediately
(observed as a ~2 s garbled clip). A deterministic eager-vs-meta diff on the
real VibeVoice-1.5B checkpoint showed these two buffers were the ONLY
tensors that differed (all 1204 checkpoint params matched exactly).
This helper re-runs each rotary module's own config-based computation on
CPU (mirroring the eager ``__init__``) and re-registers the buffers.
Args:
model: The loaded model (post assign-load + buffer materialization).
Returns:
Number of rotary modules whose ``inv_freq`` was recomputed.
"""
recomputed = 0
for module in model.modules():
buffers = dict(module.named_buffers(recurse=False))
if "inv_freq" not in buffers:
continue
config = getattr(module, "config", None)
if config is None:
continue
try:
rope_type = getattr(module, "rope_type", "default")
if rope_type == "default" and hasattr(module, "compute_default_rope_parameters"):
inv_freq, _ = module.compute_default_rope_parameters(config, device="cpu")
else:
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
rope_init_fn = ROPE_INIT_FUNCTIONS.get(rope_type)
if rope_init_fn is None:
logging.warning(
f"[VibeVoice TTS] No rope init function for type '{rope_type}' on "
f"{module.__class__.__name__}; leaving inv_freq as-is."
)
continue
inv_freq, _ = rope_init_fn(config, device="cpu")
inv_freq = inv_freq.to(device="cpu")
module.register_buffer("inv_freq", inv_freq, persistent=False)
if "original_inv_freq" in buffers:
module.register_buffer(
"original_inv_freq", inv_freq.clone(), persistent=False
)
recomputed += 1
except Exception as e:
logging.warning(
f"[VibeVoice TTS] Could not recompute RoPE inv_freq for "
f"{module.__class__.__name__}: {e}"
)
if recomputed:
logging.debug(f"[VibeVoice TTS] Recomputed RoPE inv_freq for {recomputed} rotary module(s).")
return recomputed
def _assert_shapes_compatible(model, state_dict, max_reported: int = 5) -> None:
"""Raise a friendly error when checkpoint shapes contradict the model."""
model_tensors = dict(model.named_parameters(remove_duplicate=False))
model_tensors.update(dict(model.named_buffers()))
mismatches = []
for key, tensor in state_dict.items():
model_tensor = model_tensors.get(key)
if model_tensor is None or not hasattr(tensor, "shape"):
continue
if tuple(tensor.shape) != tuple(model_tensor.shape):
mismatches.append(
(key, tuple(tensor.shape), tuple(model_tensor.shape))
)
if not mismatches:
return
raise _shape_mismatch_error(model, mismatches, max_reported)
def _shape_mismatch_error(model, mismatches, max_reported: int = 5):
"""Build the friendly config/weights-mismatch ValueError."""
lines = [
f" {key}: checkpoint shape {ckpt} vs model shape {mdl}"
for key, ckpt, mdl in mismatches[:max_reported]
]
if len(mismatches) > max_reported:
lines.append(f" ... and {len(mismatches) - max_reported} more")
return ValueError(
f"Checkpoint weight shapes do not match the "
f"'{type(model).__name__}' model for {len(mismatches)} shared "
f"key(s):\n" + "\n".join(lines) +
"\nThis almost always means the selected config does not match the "
"weight file. Verify config_name (or the sidecar config.json) "
"matches the checkpoint's architecture, or use 'Auto-detect'."
)
# Cache for loaded (model, processor) tuples, keyed by cache_key
LOADED_MODELS_CACHE = {}
def cleanup_old_models(keep_cache_key: str = None) -> None:
"""Remove all cached models except the one matching keep_cache_key."""
from .utils import VIBEVOICE_PATCHER_CACHE
from .model_registry import evict_patcher
keys_to_remove = []
for key in list(LOADED_MODELS_CACHE.keys()):
if key != keep_cache_key:
keys_to_remove.append(key)
del LOADED_MODELS_CACHE[key]
for key in list(VIBEVOICE_PATCHER_CACHE.keys()):
if key != keep_cache_key:
patcher = VIBEVOICE_PATCHER_CACHE.get(key)
try:
evict_patcher(patcher, VIBEVOICE_PATCHER_CACHE, key)
except Exception as e:
logging.warning(f"[VibeVoice TTS] Error cleaning up patcher {key}: {e}")
if keys_to_remove:
logging.debug(f"[VibeVoice TTS] Cleaned up cached models: {keys_to_remove}")
gc.collect()
model_management.soft_empty_cache()
class VibeVoiceModelHandler(torch.nn.Module):
"""A lightweight handler for a VibeVoice model.
Acts as a container that ComfyUI's ModelPatcher can manage, while the
actual heavy model is loaded on demand.
"""
def __init__(
self,
model_pack_name: str,
attention_mode: str = "eager",
use_llm_4bit: bool = False,
dtype_str: str = "auto",
):
super().__init__()
self.model_pack_name = model_pack_name
self.attention_mode = attention_mode
self.use_llm_4bit = use_llm_4bit
self.dtype_str = dtype_str
self.cache_key = f"{self.model_pack_name}_attn_{attention_mode}_q4_{int(use_llm_4bit)}"
self.model = None
self.processor = None
self.device = None
info = AVAILABLE_VIBEVOICE_MODELS.get(model_pack_name, {})
size_gb = MODEL_CONFIGS.get(model_pack_name, {}).get("size_gb", 4.0)
self.size = int(size_gb * (1024**3))
def load_model(self, device, attention_mode: str = "eager"):
"""Load the model and processor into memory."""
self.model, self.processor = VibeVoiceLoader.load_model(
self.model_pack_name,
device,
attention_mode,
use_llm_4bit=self.use_llm_4bit,
dtype_str=self.dtype_str,
)
self._refine_size()
def _refine_size(self) -> None:
"""Refine ``self.size`` from the real parameters after a load."""
try:
total = sum(
p.numel() * p.element_size() for p in self.model.parameters()
)
if total > 0:
self.size = total
except Exception:
pass
class VibeVoiceLoader(BaseVibeVoiceLoader):
"""Static loader class for VibeVoice models."""
@staticmethod
def _resolve_model_paths(model_name: str) -> tuple:
"""Resolve model paths based on the model type."""
model_info = AVAILABLE_VIBEVOICE_MODELS[model_name]
model_type = model_info["type"]
model_path = None
config_path = None
preprocessor_config_path = None
tokenizer_dir = None
if model_type == "official":
model_path = BaseVibeVoiceLoader._resolve_official_model_dir(model_name)
BaseVibeVoiceLoader._ensure_downloaded(
repo_id=model_info.get("repo_id", ""),
local_dir=model_path,
model_name=model_name,
)
config_path = os.path.join(model_path, "config.json")
preprocessor_config_path = os.path.join(model_path, "preprocessor_config.json")
tokenizer_dir = model_path
elif model_type == "local_dir":
model_path = model_info["path"]
config_path = os.path.join(model_path, "config.json")
preprocessor_config_path = os.path.join(model_path, "preprocessor_config.json")
tokenizer_dir = model_path
elif model_type == "standalone":
model_path = None
config_path = os.path.splitext(model_info["path"])[0] + ".config.json"
preprocessor_config_path = os.path.splitext(model_info["path"])[0] + ".preprocessor.json"
tokenizer_dir = os.path.dirname(model_info["path"])
return model_path, config_path, preprocessor_config_path, tokenizer_dir
@staticmethod
def _load_config(config_path: str, model_name: str):
"""Load VibeVoice config from path, with fallback to packaged default."""
if os.path.exists(config_path):
with open(config_path, 'r', encoding='utf-8') as f:
config_data = json.load(f)
model_type = config_data.get("model_type", "")
if model_type == "vibevoice_streaming":
logging.debug(f"[VibeVoice TTS] Detected streaming config for '{model_name}'")
return VibeVoiceStreamingConfig.from_pretrained(config_path)
else:
return VibeVoiceConfig.from_pretrained(config_path)
fallback_name = (
"default_VibeVoice-Large_config.json"
if "large" in model_name.lower()
else "default_VibeVoice-1.5B_config.json"
)
fallback_path = os.path.join(
os.path.dirname(__file__), "..", "src", "vibevoice", "configs", fallback_name
)
logging.warning(f"[VibeVoice TTS] Config not found for '{model_name}'. Using fallback: {fallback_name}")
return VibeVoiceConfig.from_pretrained(fallback_path)
@staticmethod
def _load_tokenizer(tokenizer_dir: str, model_name: str) -> "VibeVoiceTextTokenizerFast":
"""Load the VibeVoice text tokenizer."""
tokenizer_file_path = os.path.join(tokenizer_dir, "tokenizer.json")
if not os.path.exists(tokenizer_file_path):
logging.debug(f"[VibeVoice TTS] 'tokenizer.json' not found in model directory: {tokenizer_dir}")
packaged_configs_dir = os.path.join(
os.path.dirname(__file__), "..", "src", "vibevoice", "configs"
)
packaged_tokenizer_path = os.path.join(packaged_configs_dir, "tokenizer.json")
if os.path.exists(packaged_tokenizer_path):
logging.debug("[VibeVoice TTS] Using pre-packaged tokenizer directly from the node folder...")
return VibeVoiceTextTokenizerFast(tokenizer_file=packaged_tokenizer_path)
repos_to_try = ["Qwen/Qwen2.5-1.5B", "Qwen/Qwen2.5-7B"]
download_successful = False
last_error = None
for repo_id in repos_to_try:
logging.debug(f"[VibeVoice TTS] Attempting to download 'tokenizer.json' from Hugging Face repo '{repo_id}'...")
try:
hf_hub_download(
repo_id=repo_id,
filename="tokenizer.json",
local_dir=tokenizer_dir,
)
download_successful = True
logging.debug("[VibeVoice TTS] Download successful.")
break
except Exception as e:
logging.warning(f"[VibeVoice TTS] Failed to download from '{repo_id}': {e}")
last_error = e
if not download_successful:
error_message = (
f"FATAL: Could not get 'tokenizer.json'. All download attempts failed.\n"
f"Last error: {last_error}\n\n"
f"ACTION REQUIRED:\n"
f"1. Manually download 'tokenizer.json' from "
f"https://huggingface.co/{repos_to_try[0]}/blob/main/tokenizer.json\n"
f"2. Place the downloaded file in the following directory:\n '{tokenizer_dir}'"
)
raise RuntimeError(error_message)
return VibeVoiceTextTokenizerFast(tokenizer_file=tokenizer_file_path)
@staticmethod
def _load_processor(
tokenizer,
preprocessor_config_path: str,
is_streaming: bool = False,
):
"""Load the VibeVoice processor with tokenizer and audio processor."""
processor_config_data = {}
if os.path.exists(preprocessor_config_path):
with open(preprocessor_config_path, 'r', encoding='utf-8') as f:
processor_config_data = json.load(f)
audio_processor = VibeVoiceTokenizerProcessor()
if is_streaming:
processor = VibeVoiceStreamingProcessor(
tokenizer=tokenizer,
audio_processor=audio_processor,
speech_tok_compress_ratio=processor_config_data.get("speech_tok_compress_ratio", 3200),
db_normalize=processor_config_data.get("db_normalize", True),
)
else:
processor = VibeVoiceProcessor(
tokenizer=tokenizer,
audio_processor=audio_processor,
speech_tok_compress_ratio=processor_config_data.get("speech_tok_compress_ratio", 3200),
db_normalize=processor_config_data.get("db_normalize", True),
)
return processor
@staticmethod
def _instantiate_model(
config,
is_streaming: bool,
attn_implementation: str,
final_load_dtype: torch.dtype,
use_meta: bool = True,
):
"""Instantiate the model class directly, bypassing from_pretrained()."""
if hasattr(config, 'decoder_config'):
config.decoder_config._attn_implementation = attn_implementation
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:
if is_streaming:
model = VibeVoiceStreamingForConditionalGenerationInference(config)
else:
model = VibeVoiceForConditionalGeneration(config)
return model
@staticmethod
def _resolve_checkpoint_path(model_path, model_type, model_info):
"""Resolve the checkpoint file path(s) from a model directory or standalone path."""
if model_type == "standalone":
ckpt_path = model_info["path"]
if not os.path.isfile(ckpt_path):
raise FileNotFoundError(f"Standalone checkpoint not found: {ckpt_path}")
return ckpt_path, False
if model_path is None or not os.path.isdir(model_path):
raise FileNotFoundError(f"Model directory not found: {model_path}")
single_safetensors = os.path.join(model_path, "model.safetensors")
if os.path.isfile(single_safetensors):
logging.debug(f"[VibeVoice TTS] Found single safetensors checkpoint: {single_safetensors}")
return single_safetensors, False
sharded_safetensors_index = os.path.join(model_path, "model.safetensors.index.json")
if os.path.isfile(sharded_safetensors_index):
logging.debug(f"[VibeVoice TTS] Found sharded safetensors checkpoint: {sharded_safetensors_index}")
return sharded_safetensors_index, True
single_bin = os.path.join(model_path, "pytorch_model.bin")
if os.path.isfile(single_bin):
logging.debug(f"[VibeVoice TTS] Found single PyTorch checkpoint: {single_bin}")
return single_bin, False
sharded_bin_index = os.path.join(model_path, "pytorch_model.bin.index.json")
if os.path.isfile(sharded_bin_index):
logging.debug(f"[VibeVoice TTS] Found sharded PyTorch checkpoint: {sharded_bin_index}")
return sharded_bin_index, True
raise FileNotFoundError(
f"No checkpoint file found in model directory: {model_path}. "
f"Expected one of: model.safetensors, model.safetensors.index.json, "
f"pytorch_model.bin, pytorch_model.bin.index.json"
)
@staticmethod
def _load_sharded_state_dict(index_path, model_dir, device):
"""Load and merge a sharded checkpoint by delegating to the shared base."""
return BaseVibeVoiceLoader.load_state_dict_sharded(model_dir, device)
@staticmethod
def _apply_state_dict(model, state_dict, known_missing=None):
"""Load a state dict into the model with assign semantics."""
known_missing = known_missing or set()
_assert_shapes_compatible(model, state_dict)
missing_keys, unexpected_keys = model.load_state_dict(
state_dict, strict=False, assign=True
)
return VibeVoiceLoader._post_assign_fixups(
model, missing_keys, unexpected_keys, known_missing
)
@staticmethod
def _post_assign_fixups(model, missing_keys, unexpected_keys,
known_missing=None):
"""Post-assign model fixups shared by the batch and streaming loads."""
known_missing = known_missing or set()
config = getattr(model, "config", None)
tied = False
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()
for name, param in list(model.named_parameters()):
if param.is_meta:
param_new = torch.zeros(param.shape, dtype=param.dtype, device="cpu")
comfy.utils.set_attr_param(model, name, param_new)
for name, buf in list(model.named_buffers()):
if buf.is_meta:
leaf = name.split(".")[-1]
if leaf in _SENTINEL_BUFFER_VALUES:
val = _SENTINEL_BUFFER_VALUES[leaf]
buf_new = torch.tensor(
val, dtype=buf.dtype, device="cpu"
).reshape(buf.shape)
else:
buf_new = torch.zeros(buf.shape, dtype=buf.dtype, device="cpu")
comfy.utils.set_attr(model, name, buf_new)
_recompute_rope_buffers(model)
if tied and "lm_head.weight" in missing_keys:
logging.debug(
"[VibeVoice TTS] lm_head.weight absent from checkpoint (tied to input "
"embeddings via tie_word_embeddings) — expected."
)
reported_missing = [
k for k in missing_keys
if k not in known_missing and not (tied and k == "lm_head.weight")
]
if reported_missing:
logging.warning(f"[VibeVoice TTS] Missing keys when loading state dict: {len(reported_missing)} keys")
if len(reported_missing) < 20:
logging.warning(f"[VibeVoice TTS] Missing keys: {reported_missing}")
else:
logging.warning(f"[VibeVoice TTS] First 10 missing keys: {reported_missing[:10]}")
if unexpected_keys:
logging.warning(f"[VibeVoice TTS] Unexpected keys when loading state dict: {len(unexpected_keys)} keys")
if len(unexpected_keys) < 20:
logging.warning(f"[VibeVoice TTS] Unexpected keys: {unexpected_keys}")
else:
logging.warning(f"[VibeVoice TTS] First 10 unexpected keys: {unexpected_keys[:10]}")
return missing_keys, unexpected_keys
@staticmethod
def _stream_apply_dense(
model, tensor_pairs, known_missing=None, preserve_file_views: bool = True,
target_device=None,
):
"""Assign a dense checkpoint per-tensor directly from zero-copy file views.
``target_device`` moves each tensor onto the accelerator as it is
assigned, so the checkpoint is never materialised as a whole CPU model
first. This is the one dense assign used by every load route.
"""
known_missing = known_missing or set()
params = dict(model.named_parameters())
buffers = dict(model.named_buffers())
assigned = set()
unexpected = []
to_device = (
target_device is not None and getattr(target_device, "type", None) == "cuda"
)
def _placed(tensor):
if to_device:
return place_tensor_on_device(tensor, target_device)
return tensor if preserve_file_views else tensor.clone()
for key, tensor in tensor_pairs:
if tensor.dtype in QUANT_STORAGE_DTYPES:
raise ValueError(
"Checkpoint contains quantized-weight tensors but carries "
"no executable quantization metadata (first seen at "
f"'{key}' [{tensor.dtype}]). Loading them as floats would "
"corrupt the model. Re-export it with *.comfy_quant "
"metadata (comfy-model-tools) or use the dense/BF16 "
"checkpoint."
)
target = params.get(key)
if target is not None:
if tuple(tensor.shape) != tuple(target.shape):
raise _shape_mismatch_error(
model, [(key, tuple(tensor.shape), tuple(target.shape))]
)
comfy.utils.set_attr_param(model, key, _placed(tensor))
assigned.add(key)
continue
target_buf = buffers.get(key)
if target_buf is not None:
if tuple(tensor.shape) != tuple(target_buf.shape):
raise _shape_mismatch_error(
model, [(key, tuple(tensor.shape), tuple(target_buf.shape))]
)
comfy.utils.set_attr_buffer(model, key, _placed(tensor))
assigned.add(key)
continue
unexpected.append(key)
expected = set(model.state_dict().keys())
missing_keys = [k for k in expected if k not in assigned]
known_missing = mark_optional_absent(missing_keys, assigned, known_missing)
return VibeVoiceLoader._post_assign_fixups(
model, missing_keys, unexpected, known_missing
)
@staticmethod
def _load_state_dict_into_model(
model,
model_path: str,
model_type: str,
model_info: dict,
device,
):
"""Stream a dense checkpoint into the model using zero-copy file views."""
ckpt_path, is_sharded = VibeVoiceLoader._resolve_checkpoint_path(
model_path, model_type, model_info
)
if is_sharded:
model_dir = model_path if model_type != "standalone" else os.path.dirname(ckpt_path)
tensor_pairs = BaseVibeVoiceLoader.iter_sharded_tensors(model_dir)
else:
logging.debug(f"[VibeVoice TTS] Loading state dict from: {ckpt_path}")
tensor_pairs = BaseVibeVoiceLoader.iter_checkpoint_tensors(ckpt_path)
VibeVoiceLoader._stream_apply_dense(
model,
tensor_pairs,
preserve_file_views=True,
target_device=device,
)
return model
@staticmethod
def load_model(
model_name: str,
device,
attention_mode: str = "eager",
use_llm_4bit: bool = False,
dtype_str: str = "auto",
):
"""Load a VibeVoice model, downloading if necessary. Caches the loaded model."""
if model_name not in AVAILABLE_VIBEVOICE_MODELS:
raise ValueError(
f"Unknown VibeVoice model: {model_name}. "
f"Available models: {list(AVAILABLE_VIBEVOICE_MODELS.keys())}"
)
attention_mode = resolve_attention_mode(attention_mode, use_llm_4bit)
cache_key = f"{model_name}_attn_{attention_mode}_q4_{int(use_llm_4bit)}"
if cache_key in LOADED_MODELS_CACHE:
logging.debug(f"[VibeVoice TTS] Using cached model with {attention_mode} attention and q4={use_llm_4bit}")
return LOADED_MODELS_CACHE[cache_key]
model_info = AVAILABLE_VIBEVOICE_MODELS[model_name]
model_type = model_info["type"]
model_path, config_path, preprocessor_config_path, tokenizer_dir = \
VibeVoiceLoader._resolve_model_paths(model_name)
config = VibeVoiceLoader._load_config(config_path, model_name)
is_streaming = isinstance(config, VibeVoiceStreamingConfig)
if is_streaming:
logging.debug(f"[VibeVoice TTS] Model '{model_name}' detected as streaming model")
vibevoice_tokenizer = VibeVoiceLoader._load_tokenizer(tokenizer_dir, model_name)
processor = VibeVoiceLoader._load_processor(
vibevoice_tokenizer, preprocessor_config_path, is_streaming=is_streaming
)
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:
logging.debug(
f"[VibeVoice TTS] Loading model '{model_name}' with dtype: {final_load_dtype} "
f"and attention: '{attn_implementation_for_load}'"
)
model = VibeVoiceLoader._instantiate_model(
config=config,
is_streaming=is_streaming,
attn_implementation=attn_implementation_for_load,
final_load_dtype=final_load_dtype,
)
model = VibeVoiceLoader._load_state_dict_into_model(
model=model,
model_path=model_path,
model_type=model_type,
model_info=model_info,
device=load_device,
)
cast_model_to_dtype_if_needed(model, final_load_dtype)
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,
)
if attention_mode == "sage":
if check_sage_attention_compatible():
set_sage_attention(model)
else:
raise RuntimeError("Incompatible hardware/setup for SageAttention.")
model.eval()
setattr(model, "_llm_4bit", bool(quant_config))
LOADED_MODELS_CACHE[cache_key] = (model, processor)
logging.debug(f"[VibeVoice TTS] Successfully configured model '{model_name}' with {attention_mode} attention")
return model, processor
except Exception as e:
logging.error(f"[VibeVoice TTS] Failed to load model '{model_name}' with {attention_mode} attention: {e}")
raise RuntimeError(f"Failed to load model even with eager attention: {e}")
Reference in New Issue View Git Blame Copy Permalink