nodes/WASasquatch-FreeU_Advanced
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #7 from pamparamm/main

Browse Source 
Fix SVD, Add FreeU V2, Refactoring
This commit is contained in:
WAS
2024-01-05 17:59:22 -08:00
committed by GitHub
parent 120c23a3f4 680d97a2dd
commit 93e2bc1842
1 changed files with 144 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes.py
+144 -48
View File
@@ -3,7 +3,6 @@
import torch
import torch as th
import torch.fft as fft
import torch.nn.functional as F
import math
def normalize(latent, target_min=None, target_max=None):
@@ -12,9 +11,9 @@ def normalize(latent, target_min=None, target_max=None):
Args:
latent (torch.Tensor): The input tensor to be normalized.
target_min (float, optional): The minimum value after normalization.
target_min (float, optional): The minimum value after normalization.
- When `None` min will be tensor min range value.
target_max (float, optional): The maximum value after normalization.
target_max (float, optional): The maximum value after normalization.
- When `None` max will be tensor max range value.
Returns:
@@ -22,12 +21,12 @@ def normalize(latent, target_min=None, target_max=None):
"""
min_val = latent.min()
max_val = latent.max()
if target_min is None:
target_min = min_val
if target_max is None:
target_max = max_val
normalized = (latent - min_val) / (max_val - min_val)
scaled = normalized * (target_max - target_min) + target_min
return scaled
@@ -55,7 +54,7 @@ def hslerp(a, b, t):
raise ValueError("Input tensors a and b must have the same shape.")
num_channels = a.size(1)
interpolation_tensor = torch.zeros(1, num_channels, 1, 1, device=a.device, dtype=a.dtype)
interpolation_tensor[0, 0, 0, 0] = 1.0
@@ -153,58 +152,53 @@ def __temp__forward(self, x, timesteps=None, context=None, y=None, control=None,
:return: an [N x C x ...] Tensor of outputs.
"""
transformer_options["original_shape"] = list(x.shape)
transformer_options["current_index"] = 0
transformer_options["transformer_index"] = 0
transformer_patches = transformer_options.get("patches", {})
num_video_frames = kwargs.get("num_video_frames", self.default_num_video_frames)
image_only_indicator = kwargs.get("image_only_indicator", self.default_image_only_indicator)
time_context = kwargs.get("time_context", None)
assert (y is not None) == (
self.num_classes is not None
), "must specify y if and only if the model is class-conditional"
hs = []
t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(self.dtype)
t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x.dtype)
emb = self.time_embed(t_emb)
if self.num_classes is not None:
assert y.shape[0] == x.shape[0]
emb = emb + self.label_emb(y)
h = x.type(self.dtype)
h = x
for id, module in enumerate(self.input_blocks):
transformer_options["block"] = ("input", id)
h = forward_timestep_embed(module, h, emb, context, transformer_options)
if control is not None and 'input' in control and len(control['input']) > 0:
ctrl = control['input'].pop()
if ctrl is not None:
h += ctrl
hs.append(h)
hsp = hs
h = forward_timestep_embed(module, h, emb, context, transformer_options, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator)
h = apply_control(h, control, 'input')
if "input_block_patch" in transformer_patches:
patch = transformer_patches["input_block_patch"]
for p in patch:
h, hsp = p(h, hsp, transformer_options)
del hsp
h = p(h, transformer_options)
hs.append(h)
if "input_block_patch_after_skip" in transformer_patches:
patch = transformer_patches["input_block_patch_after_skip"]
for p in patch:
h = p(h, transformer_options)
transformer_options["block"] = ("middle", 0)
h = forward_timestep_embed(self.middle_block, h, emb, context, transformer_options)
if control is not None and 'middle' in control and len(control['middle']) > 0:
ctrl = control['middle'].pop()
if ctrl is not None:
h += ctrl
hsp = [h]
h = forward_timestep_embed(self.middle_block, h, emb, context, transformer_options, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator)
h = apply_control(h, control, 'middle')
if "middle_block_patch" in transformer_patches:
patch = transformer_patches["middle_block_patch"]
for p in patch:
h, hsp = p(h, hsp, transformer_options)
del hsp
h = p(h, transformer_options)
for id, module in enumerate(self.output_blocks):
transformer_options["block"] = ("output", id)
hsp = hs.pop()
if control is not None and 'output' in control and len(control['output']) > 0:
ctrl = control['output'].pop()
if ctrl is not None:
hsp += ctrl
hsp = apply_control(hsp, control, 'output')
if "output_block_patch" in transformer_patches:
patch = transformer_patches["output_block_patch"]
@@ -217,7 +211,7 @@ def __temp__forward(self, x, timesteps=None, context=None, y=None, control=None,
output_shape = hs[-1].shape
else:
output_shape = None
h = forward_timestep_embed(module, h, emb, context, transformer_options, output_shape)
h = forward_timestep_embed(module, h, emb, context, transformer_options, output_shape, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator)
h = h.type(x.dtype)
if self.predict_codebook_ids:
return self.id_predictor(h)
@@ -226,7 +220,7 @@ def __temp__forward(self, x, timesteps=None, context=None, y=None, control=None,
print("Patching UNetModel.forward")
import comfy.ldm.modules.diffusionmodules.openaimodel
from comfy.ldm.modules.diffusionmodules.openaimodel import forward_timestep_embed
from comfy.ldm.modules.diffusionmodules.openaimodel import forward_timestep_embed, apply_control
from comfy.ldm.modules.diffusionmodules.util import timestep_embedding
comfy.ldm.modules.diffusionmodules.openaimodel.UNetModel.forward = __temp__forward
if comfy.ldm.modules.diffusionmodules.openaimodel.UNetModel.forward is __temp__forward:
@@ -277,13 +271,13 @@ def Fourier_filter(x, threshold, scale, scales=None, strength=1.0):
x_filtered = fft.ifftn(x_freq, dim=(-2, -1)).real
return x_filtered.to(x.dtype)
return x
class WAS_FreeU:
@classmethod
def INPUT_TYPES(s):
return {"required": {
return {"required": {
"model": ("MODEL",),
"target_block": (["output_block", "middle_block", "input_block", "all"],),
"multiscale_mode": (list(mscales.keys()),),
@@ -308,20 +302,20 @@ class WAS_FreeU:
# 10, 1.5''', "multiline": True}),
}
}
RETURN_TYPES = ("MODEL",)
FUNCTION = "patch"
CATEGORY = "_for_testing"
def patch(self, model, target_block, multiscale_mode, multiscale_strength, slice_b1, slice_b2, b1, b2, s1, s2, b1_mode="add", b1_blend=1.0, b2_mode="add", b2_blend=1.0, threshold=1.0, use_override_scales="false", override_scales=""):
min_slice = 64
max_slice_b1 = 1280
max_slice_b2 = 640
slice_b1 = max(min(max_slice_b1, slice_b1), min_slice)
slice_b2 = max(min(min(slice_b1, max_slice_b2), slice_b2), min_slice)
scales_list = []
if use_override_scales == "true":
if override_scales.strip() != "":
@@ -331,45 +325,147 @@ class WAS_FreeU:
scale_values = line.split(',')
if len(scale_values) == 2:
scales_list.append((int(scale_values[0]), float(scale_values[1])))
if use_override_scales == "true" and not scales_list:
print("No valid override scales found. Using default scale.")
scales_list = None
scales = mscales[multiscale_mode] if use_override_scales == "false" else scales_list
print(f"FreeU Plate Portions: {slice_b1} over {slice_b2}")
print(f"FreeU Multi-Scales: {scales}")
def block_patch(h, hsp, transformer_options):
def block_patch(h, transformer_options):
if h.shape[1] == 1280:
h_t = h[:,:slice_b1]
h_r = h_t * b1
h[:,:slice_b1] = blending_modes[b1_mode](h_t, h_r, b1_blend)
hsp = Fourier_filter(hsp, threshold=threshold, scale=s1, scales=scales, strength=multiscale_strength)
if h.shape[1] == 640:
h_t = h[:,:slice_b2]
h_r = h_t * b2
h[:,:slice_b2] = blending_modes[b2_mode](h_t, h_r, b2_blend)
return h
def block_patch_hsp(h, hsp, transformer_options):
if h.shape[1] == 1280:
h = block_patch(h, transformer_options)
hsp = Fourier_filter(hsp, threshold=threshold, scale=s1, scales=scales, strength=multiscale_strength)
if h.shape[1] == 640:
h = block_patch(h, transformer_options)
hsp = Fourier_filter(hsp, threshold=threshold, scale=s2, scales=scales, strength=multiscale_strength)
return h, hsp
print(f"Patching {target_block}")
m = model.clone()
if target_block == "all" or target_block == "output_block":
m.set_model_output_block_patch(block_patch)
m.set_model_output_block_patch(block_patch_hsp)
if target_block == "all" or target_block == "input_block":
m.set_model_patch(block_patch, "input_block_patch")
m.set_model_input_block_patch(block_patch)
if target_block == "all" or target_block == "middle_block":
m.set_model_patch(block_patch, "middle_block_patch")
return (m, )
class WAS_FreeU_V2:
@classmethod
def INPUT_TYPES(s):
return {"required": {
"model": ("MODEL",),
"input_block": ("BOOLEAN", {"default": False}),
"middle_block": ("BOOLEAN", {"default": False}),
"output_block": ("BOOLEAN", {"default": False}),
"multiscale_mode": (list(mscales.keys()),),
"multiscale_strength": ("FLOAT", {"default": 1.0, "max": 1.0, "min": 0, "step": 0.001}),
"slice_b1": ("INT", {"default": 640, "min": 64, "max": 1280, "step": 1}),
"slice_b2": ("INT", {"default": 320, "min": 64, "max": 640, "step": 1}),
"b1": ("FLOAT", {"default": 1.1, "min": 0.0, "max": 10.0, "step": 0.001}),
"b2": ("FLOAT", {"default": 1.2, "min": 0.0, "max": 10.0, "step": 0.001}),
"s1": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 10.0, "step": 0.001}),
"s2": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 10.0, "step": 0.001}),
},
"optional": {
"threshold": ("INT", {"default": 1.0, "max": 10, "min": 1, "step": 1}),
"use_override_scales": (["false", "true"],),
"override_scales": ("STRING", {"default": '''# OVERRIDE SCALES
# Sharpen
# 10, 1.5''', "multiline": True}),
}
}
RETURN_TYPES = ("MODEL",)
FUNCTION = "patch"
CATEGORY = "_for_testing"
def patch(self, model, input_block, middle_block, output_block, multiscale_mode, multiscale_strength, slice_b1, slice_b2, b1, b2, s1, s2, threshold=1.0, use_override_scales="false", override_scales=""):
min_slice = 64
max_slice_b1 = 1280
max_slice_b2 = 640
slice_b1 = max(min(max_slice_b1, slice_b1), min_slice)
slice_b2 = max(min(min(slice_b1, max_slice_b2), slice_b2), min_slice)
scales_list = []
if use_override_scales == "true":
if override_scales.strip() != "":
scales_str = override_scales.strip().splitlines()
for line in scales_str:
if not line.strip().startswith('#') and not line.strip().startswith('!') and not line.strip().startswith('//'):
scale_values = line.split(',')
if len(scale_values) == 2:
scales_list.append((int(scale_values[0]), float(scale_values[1])))
if use_override_scales == "true" and not scales_list:
print("No valid override scales found. Using default scale.")
scales_list = None
scales = mscales[multiscale_mode] if use_override_scales == "false" else scales_list
def _hidden_mean(h):
hidden_mean = h.mean(1).unsqueeze(1)
B = hidden_mean.shape[0]
hidden_max, _ = torch.max(hidden_mean.view(B, -1), dim=-1, keepdim=True)
hidden_min, _ = torch.min(hidden_mean.view(B, -1), dim=-1, keepdim=True)
hidden_mean = (hidden_mean - hidden_min.unsqueeze(2).unsqueeze(3)) / (hidden_max - hidden_min).unsqueeze(2).unsqueeze(3)
return hidden_mean
def block_patch(h, transformer_options):
if h.shape[1] == 1280:
hidden_mean = _hidden_mean(h)
h[:,:slice_b1] = h[:,:slice_b1] * ((b1 - 1 ) * hidden_mean + 1)
if h.shape[1] == 640:
hidden_mean = _hidden_mean(h)
h[:,:slice_b2] = h[:,:slice_b2] * ((b2 - 1 ) * hidden_mean + 1)
return h
def block_patch_hsp(h, hsp, transformer_options):
if h.shape[1] == 1280:
h = block_patch(h, transformer_options)
hsp = Fourier_filter(hsp, threshold=threshold, scale=s1, scales=scales, strength=multiscale_strength)
if h.shape[1] == 640:
h = block_patch(h, transformer_options)
hsp = Fourier_filter(hsp, threshold=threshold, scale=s2, scales=scales, strength=multiscale_strength)
return h, hsp
m = model.clone()
if output_block:
print("Patching output block")
m.set_model_output_block_patch(block_patch_hsp)
if input_block:
print("Patching input block")
m.set_model_input_block_patch(block_patch)
if middle_block:
print("Patching middle block")
m.set_model_patch(block_patch, "middle_block_patch")
return (m, )
NODE_CLASS_MAPPINGS = {
"FreeU (Advanced)": WAS_FreeU,
"FreeU_V2 (Advanced)": WAS_FreeU_V2,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"FreeU (Advanced)": "FreeU (Advanced Plus)",
"FreeU_V2 (Advanced)": "FreeU V2 (Advanced Plus)",
}