From d8de95f7f3af890b39b1e7dfce3c3af4959a00ee Mon Sep 17 00:00:00 2001
From: Hawk Lee <l.y@live.cn>
Date: Thu, 19 Feb 2026 15:52:33 +0800
Subject: [PATCH] fix(IndexTTS): preload kaldifst to prevent NeMo hang (revert
 NeMo changes)

---
 aiia_generate_segments.py | 50 ++-------------------------------------
 aiia_indextts_nodes.py    | 11 +++++++++
 2 files changed, 13 insertions(+), 48 deletions(-)

diff --git a/aiia_generate_segments.py b/aiia_generate_segments.py
index 63aa204..c2cd04c 100755
--- a/aiia_generate_segments.py
+++ b/aiia_generate_segments.py
@@ -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})")
 
diff --git a/aiia_indextts_nodes.py b/aiia_indextts_nodes.py
index 8f6b113..52d3eea 100644
--- a/aiia_indextts_nodes.py
+++ b/aiia_indextts_nodes.py
@@ -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
 # ---------------------------------------------------------------------------