From 8f8ca6847e087308b2b9a5e91bc59d4a2e8ef61d Mon Sep 17 00:00:00 2001
From: Hawk Lee <l.y@live.cn>
Date: Sun, 28 Dec 2025 16:48:46 +0800
Subject: [PATCH] fix: manual model loading and device placement for Resemble
 Enhance (v1.4.49)

---
 aiia_audio_enhance.py | 50 +++++++++++++++++++++++++++++--------------
 pyproject.toml        |  2 +-
 2 files changed, 35 insertions(+), 17 deletions(-)

diff --git a/aiia_audio_enhance.py b/aiia_audio_enhance.py
index a23ce0c..68c5973 100644
--- a/aiia_audio_enhance.py
+++ b/aiia_audio_enhance.py
@@ -72,7 +72,10 @@ def _install_resemble_if_needed():
 
     # 2. Try import first
     try:
-        from resemble_enhance.enhancer.inference import enhance as eh, denoise as dn
+        from resemble_enhance.enhancer.inference import enhance as eh, denoise as dn, load_enhancer, inference
+        # Inject into global namespace so we can use them
+        globals()["load_enhancer"] = load_enhancer
+        globals()["inference"] = inference
         enhance = eh
         denoise = dn
         return
@@ -81,7 +84,9 @@ def _install_resemble_if_needed():
         
     try:
         # Fallback import
-        from resemble_enhance.inference import enhance as eh, denoise as dn
+        from resemble_enhance.inference import enhance as eh, denoise as dn, load_enhancer, inference
+        globals()["load_enhancer"] = load_enhancer
+        globals()["inference"] = inference
         enhance = eh
         denoise = dn
         return
@@ -103,12 +108,16 @@ def _install_resemble_if_needed():
         
         # 5. Retry import
         try:
-            from resemble_enhance.enhancer.inference import enhance as eh, denoise as dn
+            from resemble_enhance.enhancer.inference import enhance as eh, denoise as dn, load_enhancer, inference
+            globals()["load_enhancer"] = load_enhancer
+            globals()["inference"] = inference
             enhance = eh
             denoise = dn
         except ImportError:
              try:
-                from resemble_enhance.inference import enhance as eh, denoise as dn
+                from resemble_enhance.inference import enhance as eh, denoise as dn, load_enhancer, inference
+                globals()["load_enhancer"] = load_enhancer
+                globals()["inference"] = inference
                 enhance = eh
                 denoise = dn
              except ImportError as e:
@@ -245,22 +254,31 @@ class AIIA_Audio_Enhance:
                 # It returns (waveform, new_sr)
                 processed_wav, new_sr = denoise(wav_tensor, sample_rate, device)
             else:
+                # Manual Inference to ensure correct Device placement
+                # Because DeepSpeed mock might confuse the library's device logic
+                global _cached_enhancer
+                if "_cached_enhancer" not in globals() or _cached_enhancer is None:
+                     print(f"[AIIA] Loading Enhancer model...")
+                     # load_enhancer(run_dir=None, device=device)
+                     _cached_enhancer = load_enhancer(None, device)
+                
+                # Force move to device (crucial fix)
+                _cached_enhancer.to(device)
+                _cached_enhancer.eval()
+
+                # Call inference primitives
                 # enhance(waveform, sample_rate, device, nfe=64, solver="Midpoint", lambd=0.5, tau=0.5)
-                # Note: 'denoising' param is implicitly handled if we use 'enhance' wrapper?
-                # Actually enhance() calls denoise() internally if we look at source, 
-                # but let's see if we can control it. 
-                # The python API `enhance` func signature:
-                # def enhance(waveform, sample_rate, device, nfe=64, solver="Midpoint", lambd=0.9, tau=0.5):
-                # We map our inputs.
-                processed_wav, new_sr = enhance(
-                    wav_tensor, 
-                    sample_rate, 
-                    device, 
+                # Signature: inference(model, dwav, sr, device, nfe=64, solver='midpoint', lambd=0.5, tau=0.5)
+                
+                processed_wav, new_sr = inference(
+                    model=_cached_enhancer,
+                    dwav=wav_tensor, 
+                    sr=sample_rate, 
+                    device=device, 
                     nfe=nfe, 
                     solver=solver.lower(), 
                     tau=tau,
-                    lambd=0.9 if denoising else 0.1 # Approximate control, usually high lambd = more denoising focus?
-                    # Actually standard usage just calls clean, enhanced = enhance(...)
+                    lambd=0.9 if denoising else 0.1
                 )
         except Exception as e:
             print(f"Error in resemble-enhance: {e}")
diff --git a/pyproject.toml b/pyproject.toml
index 517f825..4404034 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
 name = "aiia"
 description = "A suite of powerful, intuitive, and highly customizable nodes for ComfyUI, featuring a high-performance Media Browser and OOM-safe video tools."
-version = "1.4.48"
+version = "1.4.49"
 license = {file = "LICENSE"}
 readme = "README.md"
 authors = [