fix: active download patch and cpu fallback for resemble enhance (v1.4.52)
This commit is contained in:
Hawk Lee
+35
-159
@@ -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
|
||||
|
||||
+1
-1
@@ -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 = [
|
||||
|
||||
Reference in New Issue
Block a user