472 lines
20 KiB
Python
472 lines
20 KiB
Python
#code originally taken from: https://github.com/ChenyangSi/FreeU (under MIT License)
|
|
|
|
import torch
|
|
import torch as th
|
|
import torch.fft as fft
|
|
import math
|
|
|
|
def normalize(latent, target_min=None, target_max=None):
|
|
"""
|
|
Normalize a tensor `latent` between `target_min` and `target_max`.
|
|
|
|
Args:
|
|
latent (torch.Tensor): The input tensor to be normalized.
|
|
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.
|
|
- When `None` max will be tensor max range value.
|
|
|
|
Returns:
|
|
torch.Tensor: The normalized tensor
|
|
"""
|
|
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
|
|
|
|
def hslerp(a, b, t):
|
|
"""
|
|
Perform Hybrid Spherical Linear Interpolation (HSLERP) between two tensors.
|
|
|
|
This function combines two input tensors `a` and `b` using HSLERP, which is a specialized
|
|
interpolation method for smooth transitions between orientations or colors.
|
|
|
|
Args:
|
|
a (tensor): The first input tensor.
|
|
b (tensor): The second input tensor.
|
|
t (float): The blending factor, a value between 0 and 1 that controls the interpolation.
|
|
|
|
Returns:
|
|
tensor: The result of HSLERP interpolation between `a` and `b`.
|
|
|
|
Note:
|
|
HSLERP provides smooth transitions between orientations or colors, particularly useful
|
|
in applications like image processing and 3D graphics.
|
|
"""
|
|
if a.shape != b.shape:
|
|
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
|
|
|
|
result = (1 - t) * a + t * b
|
|
|
|
if t < 0.5:
|
|
result += (torch.norm(b - a, dim=1, keepdim=True) / 6) * interpolation_tensor
|
|
else:
|
|
result -= (torch.norm(b - a, dim=1, keepdim=True) / 6) * interpolation_tensor
|
|
|
|
return result
|
|
|
|
blending_modes = {
|
|
# Args:
|
|
# - a (tensor): Latent input 1
|
|
# - b (tensor): Latent input 2
|
|
# - t (float): Blending factor
|
|
|
|
# Interpolates between tensors a and b using normalized linear interpolation.
|
|
'bislerp': lambda a, b, t: normalize((1 - t) * a + t * b),
|
|
# Transfer the color from `b` to `a` by t` factor
|
|
'colorize': lambda a, b, t: a + (b - a) * t,
|
|
# Interpolates between tensors a and b using cosine interpolation.
|
|
'cosine interp': lambda a, b, t: (a + b - (a - b) * torch.cos(t * torch.tensor(math.pi))) / 2,
|
|
# Interpolates between tensors a and b using cubic interpolation.
|
|
'cuberp': lambda a, b, t: a + (b - a) * (3 * t ** 2 - 2 * t ** 3),
|
|
# Interpolates between tensors a and b using normalized linear interpolation,
|
|
# with a twist when t is greater than or equal to 0.5.
|
|
'hslerp': hslerp,
|
|
# Adds tensor b to tensor a, scaled by t.
|
|
'inject': lambda a, b, t: a + b * t,
|
|
# Interpolates between tensors a and b using linear interpolation.
|
|
'lerp': lambda a, b, t: (1 - t) * a + t * b,
|
|
# Simulates a brightening effect by adding tensor b to tensor a, scaled by t.
|
|
'linear dodge': lambda a, b, t: normalize(a + b * t),
|
|
}
|
|
|
|
mscales = {
|
|
"Default": None,
|
|
"Bandpass": [
|
|
(5, 0.0), # Low-pass filter
|
|
(15, 1.0), # Pass-through filter (allows mid-range frequencies)
|
|
(25, 0.0), # High-pass filter
|
|
],
|
|
"Low-Pass": [
|
|
(10, 1.0), # Allows low-frequency components, suppresses high-frequency components
|
|
],
|
|
"High-Pass": [
|
|
(10, 0.0), # Suppresses low-frequency components, allows high-frequency components
|
|
],
|
|
"Pass-Through": [
|
|
(10, 1.0), # Passes all frequencies unchanged, no filtering
|
|
],
|
|
"Gaussian-Blur": [
|
|
(10, 0.5), # Blurs the image by allowing a range of frequencies with a Gaussian shape
|
|
],
|
|
"Edge-Enhancement": [
|
|
(10, 2.0), # Enhances edges and high-frequency features while suppressing low-frequency details
|
|
],
|
|
"Sharpen": [
|
|
(10, 1.5), # Increases the sharpness of the image by emphasizing high-frequency components
|
|
],
|
|
"Multi-Bandpass": [
|
|
[(5, 0.0), (15, 1.0), (25, 0.0)], # Multi-scale bandpass filter
|
|
],
|
|
"Multi-Low-Pass": [
|
|
[(5, 1.0), (10, 0.5), (15, 0.2)], # Multi-scale low-pass filter
|
|
],
|
|
"Multi-High-Pass": [
|
|
[(5, 0.0), (10, 0.5), (15, 0.8)], # Multi-scale high-pass filter
|
|
],
|
|
"Multi-Pass-Through": [
|
|
[(5, 1.0), (10, 1.0), (15, 1.0)], # Pass-through at different scales
|
|
],
|
|
"Multi-Gaussian-Blur": [
|
|
[(5, 0.5), (10, 0.8), (15, 0.2)], # Multi-scale Gaussian blur
|
|
],
|
|
"Multi-Edge-Enhancement": [
|
|
[(5, 1.2), (10, 1.5), (15, 2.0)], # Multi-scale edge enhancement
|
|
],
|
|
"Multi-Sharpen": [
|
|
[(5, 1.5), (10, 2.0), (15, 2.5)], # Multi-scale sharpening
|
|
],
|
|
}
|
|
|
|
# forward function from comfy.ldm.modules.diuffusionmodules.openaimodel
|
|
# Hopefully temporary replacement
|
|
def __temp__forward(self, x, timesteps=None, context=None, y=None, control=None, transformer_options={}, **kwargs):
|
|
"""
|
|
Apply the model to an input batch.
|
|
:param x: an [N x C x ...] Tensor of inputs.
|
|
:param timesteps: a 1-D batch of timesteps.
|
|
:param context: conditioning plugged in via crossattn
|
|
:param y: an [N] Tensor of labels, if class-conditional.
|
|
:return: an [N x C x ...] Tensor of outputs.
|
|
"""
|
|
transformer_options["original_shape"] = list(x.shape)
|
|
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", getattr(self, "default_image_only_indicator", None))
|
|
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(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
|
|
for id, module in enumerate(self.input_blocks):
|
|
transformer_options["block"] = ("input", id)
|
|
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 = 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, 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 = p(h, transformer_options)
|
|
|
|
for id, module in enumerate(self.output_blocks):
|
|
transformer_options["block"] = ("output", id)
|
|
hsp = hs.pop()
|
|
hsp = apply_control(hsp, control, 'output')
|
|
|
|
if "output_block_patch" in transformer_patches:
|
|
patch = transformer_patches["output_block_patch"]
|
|
for p in patch:
|
|
h, hsp = p(h, hsp, transformer_options)
|
|
|
|
h = th.cat([h, hsp], dim=1)
|
|
del hsp
|
|
if len(hs) > 0:
|
|
output_shape = hs[-1].shape
|
|
else:
|
|
output_shape = None
|
|
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)
|
|
else:
|
|
return self.out(h)
|
|
|
|
print("Patching UNetModel.forward")
|
|
import comfy.ldm.modules.diffusionmodules.openaimodel
|
|
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:
|
|
print("UNetModel.forward has been successfully patched.")
|
|
else:
|
|
print("UNetModel.forward patching failed.")
|
|
|
|
def Fourier_filter(x, threshold, scale, scales=None, strength=1.0):
|
|
# FFT
|
|
if isinstance(x, list):
|
|
x = x[0]
|
|
if isinstance(x, torch.Tensor):
|
|
x_freq = fft.fftn(x.float(), dim=(-2, -1))
|
|
x_freq = fft.fftshift(x_freq, dim=(-2, -1))
|
|
|
|
B, C, H, W = x_freq.shape
|
|
mask = torch.ones((B, C, H, W), device=x.device)
|
|
|
|
crow, ccol = H // 2, W // 2
|
|
mask[..., crow - threshold:crow + threshold, ccol - threshold:ccol + threshold] = scale
|
|
|
|
if scales is not None:
|
|
if isinstance(scales[0], tuple):
|
|
# Single-scale mode
|
|
for scale_params in scales:
|
|
if len(scale_params) == 2:
|
|
scale_threshold, scale_value = scale_params
|
|
scaled_scale_value = scale_value * strength
|
|
scale_mask = torch.ones((B, C, H, W), device=x.device)
|
|
scale_mask[..., crow - scale_threshold:crow + scale_threshold, ccol - scale_threshold:ccol + scale_threshold] = scaled_scale_value
|
|
mask = mask + (scale_mask - mask) * strength
|
|
else:
|
|
# Multi-scale mode
|
|
for scale_params in scales:
|
|
if isinstance(scale_params, list):
|
|
for scale_tuple in scale_params:
|
|
if len(scale_tuple) == 2:
|
|
scale_threshold, scale_value = scale_tuple
|
|
scaled_scale_value = scale_value * strength
|
|
scale_mask = torch.ones((B, C, H, W), device=x.device)
|
|
scale_mask[..., crow - scale_threshold:crow + scale_threshold, ccol - scale_threshold:ccol + scale_threshold] = scaled_scale_value
|
|
mask = mask + (scale_mask - mask) * strength
|
|
|
|
x_freq = x_freq * mask
|
|
|
|
# IFFT
|
|
x_freq = fft.ifftshift(x_freq, dim=(-2, -1))
|
|
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": {
|
|
"model": ("MODEL",),
|
|
"target_block": (["output_block", "middle_block", "input_block", "all"],),
|
|
"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": {
|
|
"b1_mode": (list(blending_modes.keys()),),
|
|
"b1_blend": ("FLOAT", {"default": 1.0, "max": 100, "min": 0, "step": 0.001}),
|
|
"b2_mode": (list(blending_modes.keys()),),
|
|
"b2_blend": ("FLOAT", {"default": 1.0, "max": 100, "min": 0, "step": 0.001}),
|
|
"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, 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() != "":
|
|
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
|
|
|
|
print(f"FreeU Plate Portions: {slice_b1} over {slice_b2}")
|
|
print(f"FreeU Multi-Scales: {scales}")
|
|
|
|
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)
|
|
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_hsp)
|
|
if target_block == "all" or target_block == "input_block":
|
|
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)",
|
|
}
|