95 lines
3.6 KiB
Python
95 lines
3.6 KiB
Python
import torch
|
|
from pytorch_wavelets import DWTForward, DWTInverse
|
|
|
|
class DCW_ModelPatch:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 4.0, "step": 0.05}),
|
|
"wavelet": (["haar", "db2", "db3"], {"default": "haar"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "apply_dcw"
|
|
CATEGORY = "model_patches/DCW"
|
|
|
|
def apply_dcw(self, model, strength, wavelet):
|
|
dwt = DWTForward(J=1, wave=wavelet, mode='zero')
|
|
idwt = DWTInverse(wave=wavelet, mode='zero')
|
|
|
|
def dcw_wrapper(model_function, params):
|
|
input_x = params.get("input")
|
|
timestep = params.get("timestep")
|
|
c = params.get("c")
|
|
|
|
output = model_function(input_x, timestep, **c)
|
|
|
|
if strength == 0:
|
|
return output
|
|
|
|
try:
|
|
device, dtype = output.device, output.dtype
|
|
dwt.to(device=device, dtype=dtype)
|
|
idwt.to(device=device, dtype=dtype)
|
|
|
|
# --- DYNAMIC BROADCASTING FIX ---
|
|
# We count how many dimensions the latent has (4 for image, 5 for video)
|
|
# and expand the scale 's' to match.
|
|
dims = len(input_x.shape)
|
|
s = timestep.view(-1, *([1] * (dims - 1)))
|
|
|
|
if s.shape[0] != output.shape[0]:
|
|
s = s[0].repeat(output.shape[0], *([1] * (dims - 1)))
|
|
|
|
# If 5D (Video), we squeeze the 'Frames' dim to do Wavelet math, then unsqueeze
|
|
is_5d = (dims == 5)
|
|
if is_5d:
|
|
# [B, C, F, H, W] -> [B*F, C, H, W]
|
|
b, c_feat, f, h, w = input_x.shape
|
|
work_input = input_x.transpose(1, 2).reshape(-1, c_feat, h, w)
|
|
work_output = output.transpose(1, 2).reshape(-1, c_feat, h, w)
|
|
work_s = s.transpose(1, 2).reshape(-1, 1, 1, 1)
|
|
else:
|
|
work_input = input_x
|
|
work_output = output
|
|
work_s = s
|
|
|
|
# 1. Calculate x0
|
|
x0 = work_input - (work_s * work_output)
|
|
|
|
# 2. Wavelet Logic
|
|
yl_t, yh_t = dwt(work_input)
|
|
yl_0, yh_0 = dwt(x0)
|
|
|
|
t_val = timestep[0].item()
|
|
t_factor = min(t_val / 1000.0, 1.0) if t_val > 1.0 else t_val
|
|
|
|
yl_corrected = yl_0 + (strength * (yl_t - yl_0) * t_factor)
|
|
x0_corrected = idwt((yl_corrected, yh_0))
|
|
|
|
# 3. Final Prediction
|
|
res_output = (work_input - x0_corrected) / torch.clamp(work_s, min=1e-7)
|
|
|
|
# 4. Restore 5D shape if necessary
|
|
if is_5d:
|
|
# [B*F, C, H, W] -> [B, C, F, H, W]
|
|
res_output = res_output.view(b, f, c_feat, h, w).transpose(1, 2)
|
|
|
|
if t_val % 10 == 0:
|
|
print(f"[DCW] Success | Mode: {'5D' if is_5d else '4D'} | T: {t_val:.2f}")
|
|
|
|
return res_output
|
|
|
|
except Exception as e:
|
|
print(f"[DCW ERROR] Still failing: {e}")
|
|
return output
|
|
|
|
m = model.clone()
|
|
m.model_options["model_function_wrapper"] = dcw_wrapper
|
|
return (m,)
|
|
|
|
NODE_CLASS_MAPPINGS = {"DCW_ModelPatch": DCW_ModelPatch}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"DCW_ModelPatch": "Apply DCW (Wavelet Patch)"} |