112 lines
4.3 KiB
Python
112 lines
4.3 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
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)
|
|
|
|
# 1. Determine Dimensions
|
|
dims = len(input_x.shape)
|
|
# Correctly broadcast S to [B, 1, 1, 1] or [B, 1, 1, 1, 1]
|
|
s = timestep.view(-1, *([1] * (dims - 1)))
|
|
|
|
if s.shape[0] != output.shape[0]:
|
|
s = s[0].repeat(output.shape[0], *([1] * (dims - 1)))
|
|
|
|
# 2. 5D to 4D Conversion for Wavelets
|
|
is_5d = (dims == 5)
|
|
if is_5d:
|
|
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
|
|
|
|
# 3. Handle Padding
|
|
orig_h, orig_w = work_input.shape[-2:]
|
|
pad_h = orig_h % 2
|
|
pad_w = orig_w % 2
|
|
|
|
if pad_h > 0 or pad_w > 0:
|
|
work_input = F.pad(work_input, (0, pad_w, 0, pad_h))
|
|
work_output = F.pad(work_output, (0, pad_w, 0, pad_h))
|
|
# Ensure work_s matches the padded spatial dimensions
|
|
work_s_pad = work_s.expand(-1, -1, work_input.shape[-2], work_input.shape[-1])
|
|
else:
|
|
work_s_pad = work_s
|
|
|
|
# 4. Calculate x0
|
|
# If it fails here, the debug prints above will show why
|
|
x0 = work_input - (work_s_pad * work_output)
|
|
|
|
# 5. Wavelet Correction
|
|
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))
|
|
|
|
# 6. Crop and Return
|
|
if pad_h > 0 or pad_w > 0:
|
|
x0_corrected = x0_corrected[:, :, :orig_h, :orig_w]
|
|
work_input = work_input[:, :, :orig_h, :orig_w]
|
|
|
|
res_output = (work_input - x0_corrected) / torch.clamp(work_s, min=1e-7)
|
|
|
|
if is_5d:
|
|
res_output = res_output.view(b, f, c_feat, orig_h, orig_w).transpose(1, 2)
|
|
|
|
return res_output
|
|
|
|
except Exception as e:
|
|
print(f"\n[DCW FATAL ERROR]: {e}")
|
|
try:
|
|
print(f"Context Shapes -> Input: {work_input.shape}, Output: {work_output.shape}, S: {work_s_pad.shape}")
|
|
except:
|
|
pass
|
|
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)"}
|