fix(IndexTTS): preload kaldifst to prevent NeMo hang (revert NeMo changes)

This commit is contained in:
Hawk Lee
2026-02-19 15:52:33 +08:00
parent 8bdc5cb7d3
commit d8de95f7f3
2 changed files with 13 additions and 48 deletions
+2 -48
View File
@@ -240,57 +240,11 @@ class AIIA_GenerateSpeakerSegments:
try:
# --- DIAGNOSTIC START ---
import threading, time, traceback, sys, faulthandler, os
# Use stderr for immediate output (unbuffered)
def log(msg):
try: sys.stderr.write(f"{node_name_log} {msg}\n")
except: pass
log("Starting strict NeMo import watchdog (30s timeout)...")
# 1. Use faulthandler (C-level signal) to dump trace even if GIL is held
try:
faulthandler.dump_traceback_later(30, repeat=False, file=sys.stderr)
log("Enabled faulthandler (will dump to stderr after 30s)")
except Exception as e:
log(f"Failed to enable faulthandler: {e}")
# 2. Keep python thread for heartbeats
def _watchdog():
start_t = time.time()
while not getattr(sys.modules[__name__], "_nemo_imported", False):
elapsed = time.time() - start_t
if elapsed > 35: # Backup dump if faulthandler fails
log(f"[WATCHDOG] STILL IMPORTING after {elapsed:.1f}s! Force dumping stacks:")
try:
for thread_id, frame in sys._current_frames().items():
log(f"\n=== Thread {thread_id} ===")
traceback.print_stack(frame, file=sys.stderr)
log("[WATCHDOG] End of stack dump.\n")
except:
log("[WATCHDOG] Failed to dump stacks!")
break
if elapsed > 2 and int(elapsed) % 5 == 0:
log(f"[WATCHDOG] still importing... ({int(elapsed)}s)")
time.sleep(1)
sys.modules[__name__]._nemo_imported = False
_t = threading.Thread(target=_watchdog, daemon=True)
_t.start()
# --- DIAGNOSTIC END ---
try: from nemo.collections.asr.models.msdd_models import SortformerEncLabelModel
except ImportError: from nemo.collections.asr.models import SortformerEncLabelModel
from nemo.collections.asr.parts.mixins.diarization import DiarizeConfig
# Disable faulthandler timer
try: faulthandler.cancel_dump_traceback_later()
except: pass
sys.modules[__name__]._nemo_imported = True
log("成功导入 NeMo 类。")
# PostProcessingParams 和 asdict 在此流程中不再直接从 Python 导入和使用
print(f"{node_name_log} 成功导入 NeMo 类。")
except ImportError as e_import_model:
return self._create_error_output(f"导入 NeMo 类失败 ({e_import_model})")
+11
View File
@@ -12,6 +12,17 @@ import numpy as np
import contextlib
import shutil
# --- Preload kaldifst to prevent hang when used with NeMo Diarization ---
# IndexTTS-2 seems to put the process (OpenMP/MKL/dlopen lock) in a state
# that causes kaldifst initialization to hang if imported LATER.
# By importing it here at startup (before IndexTTS runs), we ensure it
# initializes safely.
try:
import kaldifst
except ImportError:
pass
# ---------------------------------------------------------------------------
# Lazy setup: add libs/index-tts to sys.path so `from indextts...` works
# ---------------------------------------------------------------------------