nodes/DawnW0lf-ComfyUI-DCW-Diffusion-Color-Wavelets-Node
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update DCW_ModelPatch.py

Browse Source
This commit is contained in:
Dawn Wolf
2026-04-21 19:07:44 +02:00
committed by GitHub
parent e79cbc53ac
commit fc550c5bef
1 changed files with 33 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
DCW_ModelPatch.py
+33 -17
View File
@@ -1,4 +1,5 @@
import torch
import torch.nn.functional as F
from pytorch_wavelets import DWTForward, DWTInverse
class DCW_ModelPatch:
@@ -35,19 +36,17 @@ class DCW_ModelPatch:
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.
# 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)))
# If 5D (Video), we squeeze the 'Frames' dim to do Wavelet math, then unsqueeze
# 2. 5D to 4D Conversion for Wavelets
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)
@@ -57,10 +56,24 @@ class DCW_ModelPatch:
work_output = output
work_s = s
# 1. Calculate x0
x0 = work_input - (work_s * work_output)
# 3. Handle Padding
orig_h, orig_w = work_input.shape[-2:]
pad_h = orig_h % 2
pad_w = orig_w % 2
# 2. Wavelet Logic
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)
@@ -70,21 +83,24 @@ class DCW_ModelPatch:
yl_corrected = yl_0 + (strength * (yl_t - yl_0) * t_factor)
x0_corrected = idwt((yl_corrected, yh_0))
# 3. Final Prediction
# 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)
# 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}")
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"[DCW ERROR] Still failing: {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()
@@ -92,4 +108,4 @@ class DCW_ModelPatch:
return (m,)
NODE_CLASS_MAPPINGS = {"DCW_ModelPatch": DCW_ModelPatch}
NODE_DISPLAY_NAME_MAPPINGS = {"DCW_ModelPatch": "Apply DCW (Wavelet Patch)"}
NODE_DISPLAY_NAME_MAPPINGS = {"DCW_ModelPatch": "Apply DCW (Wavelet Patch)"}