From 675fa919272ea7bd14907929249a09508fd4c81d Mon Sep 17 00:00:00 2001 From: Hawk Lee Date: Sun, 28 Dec 2025 17:35:18 +0800 Subject: [PATCH] fix: active download patch and cpu fallback for resemble enhance (v1.4.52) --- aiia_audio_enhance.py | 194 ++++++++---------------------------------- pyproject.toml | 2 +- 2 files changed, 36 insertions(+), 160 deletions(-) diff --git a/aiia_audio_enhance.py b/aiia_audio_enhance.py index 1891256..9f88f69 100644 --- a/aiia_audio_enhance.py +++ b/aiia_audio_enhance.py @@ -254,53 +254,28 @@ class AIIA_Audio_Enhance: # It returns (waveform, new_sr) processed_wav, new_sr = denoise(wav_tensor, sample_rate, device) else: - # FORCE STOP GIT PULL - # usage of load_enhancer calls 'download' internally which checks github every time. - # We monkey patch it to avoid the delay. + # FORCE STOP GIT PULL (Aggressive Patch) try: import resemble_enhance.enhancer.download as dl_module - # Save original just in case, but we overwrite it - # Logic: if path exists, just return path. def noop_download(*args, **kwargs): - # The original returns the path to the model dir - # We assume it's already set up by us or previous run + # Force return local path return Path(folder_paths.models_dir) / "resemble_enhance" / "model_repo" / "enhancer_stage2" - # Only patch if we are sure + # ACTUALLY APPLY THE PATCH if hasattr(dl_module, "download"): - # Check if files exist first? - # Actually better to just let it run ONCE, but since we are caching the model, - # subsequent calls won't hit this. - # The user compains "every time execute". - # Because process_audio is called, and if _cached_enhancer is not None, we skip load. - # Wait, if _cached_enhancer IS CACHED, why did it pull again? - # Ah, the user probably restarted ComfyUI or invalidation happened? - # Or maybe globals() cleared? - # If standard workflow, it should be cached. - # But regardless, let's patch it. - pass # Actually, let's just implement the path return in load_enhancer call if possible. - # load_enhancer(run_dir) -> if run_dir is None, it calls download. - # We can pass the local path! + dl_module.download = noop_download + if hasattr(dl_module, "get_dataset_file"): + # Some versions use this + dl_module.get_dataset_file = noop_download except ImportError: pass global _cached_enhancer if "_cached_enhancer" not in globals() or _cached_enhancer is None: print(f"[AIIA] Loading Enhancer model...") - - # 1. Calculate Local Model Path to avoid Git Pull - model_dir = Path(folder_paths.models_dir) / "resemble_enhance" / "model_repo" / "enhancer_stage2" - if model_dir.exists(): - # Pass local path to avoid download check - run_dir = model_dir - else: - run_dir = None # Will trigger download - - _cached_enhancer = load_enhancer(run_dir, device) - - # Force move to device - _cached_enhancer.to(device) - _cached_enhancer.eval() + # Pass local path to be safe + run_dir = Path(folder_paths.models_dir) / "resemble_enhance" / "model_repo" / "enhancer_stage2" + _cached_enhancer = load_enhancer(str(run_dir) if run_dir.exists() else None, device) # Configure Params if hasattr(_cached_enhancer, "config"): @@ -308,131 +283,32 @@ class AIIA_Audio_Enhance: _cached_enhancer.config.solver = solver.lower() _cached_enhancer.config.tau = tau - # CUSTOM INFERENCE LOOP (fixes Device mismatch & gives control) - # Re-implementing 'inference' and 'inference_chunk' logic locally - # Source: resemble_enhance/inference.py - - dt = 1.0 / sample_rate - # dwav is [T] - # Normalize - abs_max = wav_tensor.abs().max() + 1e-6 - wav_tensor_norm = wav_tensor / abs_max - - # Chunking - chunk_duration = 20.0 # seconds, default in lib - chunk_length = int(chunk_duration * sample_rate) - - # If short audio, simple run - if wav_tensor_norm.shape[0] <= chunk_length: - # Run directly - # inference_chunk logic: - t = torch.arange(wav_tensor_norm.shape[0], device=device).float() * dt - # Unconditional generation? No, it's conditional on dwav - # The model forward: model(dwav, t) check signature? - # Actually model is ConditionalFlowMatching. - # We use the wrapper's logic for "inference_chunk" equivalent? - # Wait, `inference_chunk` does specific CFM solving. - # We better call the internal `model.inference` or similar if available? - # No, `load_enhancer` returns a `Enhancer` which is likely a `ConditionalFlowMatching`. - # Let's trust the `inference` function BUT handle the device mismatch of `abs_max` manually. - - # The error was `hwav * abs_max`. - # If we do the normalization OURSELVES and pass the normalized tensor, - # maybe we can avoid the error? - # But `inference` does normalization internally? - # function inference(...) does: - # ... - # chunks.append(inference_chunk(...)) - # function inference_chunk(...) does: - # abs_max = dwav.abs().max() - # dwav = dwav / abs_max - # ... solve ... - # hwav = hwav * abs_max - - # If we implement the loop, we control this. - pass - - # RE-IMPLEMENTATION of inference() to fix device issues - # Logic: Divide into chunks, process, concat. - - model = _cached_enhancer - # Ensure abs_max is ensuring device match - # Convert to Python float to be safe for multiplication with any tensor - abs_max_val = float(wav_tensor.abs().max() + 1e-6) - - # Normalize input - dwav_norm = wav_tensor / abs_max_val - - chunks = [] - start = 0 - N = dwav_norm.shape[0] - - # Create a local inference_chunk helper that uses our model - def local_infer_chunk(chunk_tensor): - # chunk_tensor is [L] on device - # Need time encoding - L = chunk_tensor.shape[0] - t = torch.arange(L, device=device).float() * dt - - # For CFM, we usually call sample() or similar. - # Enhancer.forward? - # Let's inspect library source via memory... - # It seems `inference_chunk` calls `model.sample(dwav, t)`? - # IF we can't be sure, we risk breaking it. - - # ALTERNATIVE: Use the library's inference, but patch the tensor? - # The error `hwav * abs_max` involves a scalar `abs_max` calculated inside. - # If `dwav` is on CUDA, `dwav.abs().max()` is a CUDA scalar. - # `hwav` (output) is likely CUDA. - # CUDA * CUDA is valid. - # Why "cpu"? - # Maybe `t` (time) created inside is on CPU? - # `t = torch.linspace(0, 1, nfe + 1, device=device)` - - # Let's try to just WRAP the call in a torch.autocast? No. - # Let's try to use the library `inference` but pass `device` object instead of string? - # Someties strings cause issues. - pass + def run_inference_safe(active_device): + # Force move + _cached_enhancer.to(active_device) + _cached_enhancer.eval() + + # Run + return inference( + model=_cached_enhancer, + dwav=wav_tensor.to(active_device), + sr=sample_rate, + device=active_device, + ) - # Let's default to calling the library inference, but fixing the RE-DOWNLOAD issue first. - # And for the device error: - # We will force the INPUT to be on CPU, then GPU inside? No that's slow. - # We will try `model.cuda()` explicitly. - - # If I pass `run_dir` as local path, it skips download. That fixes issue #2. - # Issue #1 (Device Error): - # I suspect `abs_max` is the culprit. - # If I modify `inference.py` in memory? - # Or... - # What if passing `device=torch.device("cuda")` object works better? - - run_dir = Path(folder_paths.models_dir) / "resemble_enhance" / "model_repo" / "enhancer_stage2" - if not run_dir.exists(): run_dir = None - - if "_cached_enhancer" not in globals() or _cached_enhancer is None: - _cached_enhancer = load_enhancer(str(run_dir) if run_dir else None, device) - - # Force cast - _cached_enhancer.to(device) - - # Update config (v1.4.50 fix) - if hasattr(_cached_enhancer, "config"): - _cached_enhancer.config.nfe = nfe - _cached_enhancer.config.solver = solver.lower() - _cached_enhancer.config.tau = tau - - # RUN IT - # We catch the runtime error and retry on CPU if it fails? - # No, that's too slow. - # I'll try to monkeypatch torch.arange? No. - - # Let's try to pass `dwav` and `src` carefully. - processed_wav, new_sr = inference( - model=_cached_enhancer, - dwav=wav_tensor, - sr=sample_rate, - device=device, - ) + try: + # Try on requested device + processed_wav, new_sr = run_inference_safe(device) + except RuntimeError as e: + if "device" in str(e) and device != "cpu": + print(f"[AIIA] Device mismatch error on {device}. Retrying on CPU (fallback)...") + # Clear cache or just move? + try: + processed_wav, new_sr = run_inference_safe("cpu") + except Exception as e2: + raise e2 + else: + raise e except Exception as e: print(f"Error in resemble-enhance: {e}") raise e diff --git a/pyproject.toml b/pyproject.toml index ea44d99..11a4a9c 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.51" +version = "1.4.52" license = {file = "LICENSE"} readme = "README.md" authors = [