nodes/havvk-ComfyUI_AIIA
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: 删除 from_pretrained 保存/恢复 hack，修复 classmethod 描述符损坏

Browse Source 
保存 bound classmethod 再恢复会破坏描述符：所有子类调用 from_pretrained 时
cls 参数都变成 PreTrainedModel 而不是实际子类。这导致 VibeVoice 创建出空模型
（0个module、空 device_map）并触发 accelerate IndexError。

modelscope 导入已改为 from transformers，不再需要这个 hack。
This commit is contained in:
Hawk Lee
2026-02-19 13:37:11 +08:00
parent 133daf8eee
commit 29d64b64d3
2 changed files with 11 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files
aiia_indextts_nodes.py
+10 -22
View File
@@ -318,28 +318,16 @@ class AIIA_IndexTTS2_Loader:
# (infer_v2 → model_v2 → transformers_gpt2 → transformers_generation_utils) does
# top-level `from transformers.cache_utils import QuantizedCacheConfig` which needs
# the patch to exist first.
#
# IMPORTANT: indextts's infer_v2.py does `from modelscope import AutoModelForCausalLM`
# at module level, which globally monkey-patches transformers.PreTrainedModel.from_pretrained.
# This breaks NeMo's SortformerEncLabelModel.restore_from() (hangs indefinitely).
# We save and restore the original method to neutralize modelscope's side-effect.
import transformers
_orig_from_pretrained = transformers.PreTrainedModel.from_pretrained
try:
_apply_transformers_patches()
from indextts.infer_v2 import IndexTTS2
with _patch_indextts_loading(model_dir):
tts = IndexTTS2(
cfg_path=cfg_path,
model_dir=model_dir,
use_fp16=use_fp16,
use_cuda_kernel=use_cuda_kernel,
use_deepspeed=False,
)
finally:
# Always restore original from_pretrained, even if loading fails
transformers.PreTrainedModel.from_pretrained = _orig_from_pretrained
_apply_transformers_patches()
from indextts.infer_v2 import IndexTTS2
with _patch_indextts_loading(model_dir):
tts = IndexTTS2(
cfg_path=cfg_path,
model_dir=model_dir,
use_fp16=use_fp16,
use_cuda_kernel=use_cuda_kernel,
use_deepspeed=False,
)
_INDEXTTS_MODEL_CACHE[cache_key] = tts
print("[AIIA] IndexTTS-2 loaded successfully.")
aiia_vibevoice_nodes.py
+1 -38
View File
@@ -146,44 +146,7 @@ class AIIA_VibeVoice_Loader:
# 5. Load Model
print(f"[AIIA] Loading VibeVoice model variant: {VibeVoiceClass.__name__}")
config = VibeVoiceConfig.from_pretrained(load_path)
# --- Diagnostic: intercept dispatch_model to see the device_map ---
from accelerate.big_modeling import dispatch_model as _orig_dispatch
import accelerate.big_modeling
def _debug_dispatch(model, device_map, **kwargs):
print(f"[AIIA DEBUG] dispatch_model called with device_map={device_map}")
print(f"[AIIA DEBUG] device_map type={type(device_map)}, len={len(device_map)}")
print(f"[AIIA DEBUG] model type={type(model).__name__}")
named_mods = list(model.named_modules())
print(f"[AIIA DEBUG] model has {len(named_mods)} named modules")
if len(named_mods) <= 5:
print(f"[AIIA DEBUG] modules: {[n for n,_ in named_mods]}")
import torch
if torch.cuda.is_available():
free, total = torch.cuda.mem_get_info(0)
print(f"[AIIA DEBUG] GPU mem: free={free/1e9:.2f}GB, total={total/1e9:.2f}GB, used={(total-free)/1e9:.2f}GB")
if len(device_map) == 0:
print("[AIIA DEBUG] *** EMPTY device_map! Will crash in accelerate. ***")
# Print model._no_split_modules
nsm = getattr(model, '_no_split_modules', 'NOT SET')
print(f"[AIIA DEBUG] model._no_split_modules = {nsm}")
# Try to get a working device_map manually
from accelerate import infer_auto_device_map
try:
manual_map = infer_auto_device_map(model, dtype=dtype)
print(f"[AIIA DEBUG] manual infer_auto_device_map result: {manual_map}")
except Exception as e2:
print(f"[AIIA DEBUG] manual infer_auto_device_map failed: {e2}")
return _orig_dispatch(model, device_map=device_map, **kwargs)
accelerate.big_modeling.dispatch_model = _debug_dispatch
# Also patch the import reference in transformers
import transformers.modeling_utils as _tmu
_tmu.dispatch_model = _debug_dispatch
try:
model = VibeVoiceClass.from_pretrained(load_path, config=config, torch_dtype=dtype, device_map="auto", trust_remote_code=False)
finally:
accelerate.big_modeling.dispatch_model = _orig_dispatch
_tmu.dispatch_model = _orig_dispatch
model = VibeVoiceClass.from_pretrained(load_path, config=config, torch_dtype=dtype, device_map="auto", trust_remote_code=False)
# Load Generation Config
try: