485 lines
22 KiB
Python
485 lines
22 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 torch.nn as nn
|
|
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
|
|
|
|
def stable_slerp(a, b, t: float, eps: float = 1e-6):
|
|
"""
|
|
Numerically stable spherical linear interpolation over the channel dimension.
|
|
|
|
Treat each BCHW location's C-vector as a point on a hypersphere and SLERP from a->b.
|
|
Falls back to LERP when the angle is very small or vectors are near-zero.
|
|
"""
|
|
if a.shape != b.shape:
|
|
raise ValueError("Input tensors a and b must have the same shape.")
|
|
|
|
# Norms across channel dimension
|
|
a_norm = torch.linalg.norm(a, dim=1, keepdim=True).clamp_min(eps)
|
|
b_norm = torch.linalg.norm(b, dim=1, keepdim=True).clamp_min(eps)
|
|
a_n = a / a_norm
|
|
b_n = b / b_norm
|
|
|
|
# Cosine of angle between vectors
|
|
dot = (a_n * b_n).sum(dim=1, keepdim=True).clamp(-1.0 + eps, 1.0 - eps)
|
|
theta = torch.acos(dot)
|
|
sin_theta = torch.sin(theta).clamp_min(eps)
|
|
|
|
# Scalar t is expected; keep broadcast-friendly
|
|
s0 = torch.sin((1.0 - t) * theta) / sin_theta
|
|
s1 = torch.sin(t * theta) / sin_theta
|
|
|
|
slerp_out = s0 * a + s1 * b
|
|
lerp_out = (1.0 - t) * a + t * b
|
|
|
|
# When angle is too small, prefer LERP to avoid instabilities
|
|
use_lerp = (theta < 1e-3).squeeze(1)
|
|
out = torch.where(use_lerp.unsqueeze(1), lerp_out, slerp_out)
|
|
return out
|
|
|
|
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,
|
|
# Numerically stable SLERP over channel vectors
|
|
'stable_slerp': stable_slerp,
|
|
# 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,
|
|
"Low-Pass": [
|
|
(10, 1.0), # Allows low-frequency components, suppresses 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
|
|
],
|
|
}
|
|
|
|
|
|
|
|
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"], {"tooltip": "Which UNet block(s) to patch."}),
|
|
"multiscale_mode": (list(mscales.keys()), {"tooltip": "Frequency shaping preset used by Fourier_filter."}),
|
|
"multiscale_strength": ("FLOAT", {"default": 1.0, "max": 1.0, "min": 0, "step": 0.001, "tooltip": "Intensity of multi-scale shaping [0-1]."}),
|
|
"slice_b1": ("INT", {"default": 640, "min": 64, "max": 1280, "step": 1, "tooltip": "Slice width (channels) affected in 1280-channel features."}),
|
|
"slice_b2": ("INT", {"default": 320, "min": 64, "max": 640, "step": 1, "tooltip": "Slice width (channels) affected in 640-channel features."}),
|
|
"b1": ("FLOAT", {"default": 1.1, "min": 0.0, "max": 10.0, "step": 0.001, "tooltip": "Gain multiplier applied to the b1 slice (1280-ch)."}),
|
|
"b2": ("FLOAT", {"default": 1.2, "min": 0.0, "max": 10.0, "step": 0.001, "tooltip": "Gain multiplier applied to the b2 slice (640-ch)."}),
|
|
"s1": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 10.0, "step": 0.001, "tooltip": "Fourier scale at threshold for 1280-ch features."}),
|
|
"s2": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 10.0, "step": 0.001, "tooltip": "Fourier scale at threshold for 640-ch features."}),
|
|
},
|
|
"optional": {
|
|
"b1_mode": (list(blending_modes.keys()), {"tooltip": "Blending mode for b1 path."}),
|
|
"b1_blend": ("FLOAT", {"default": 1.0, "max": 100, "min": 0, "step": 0.001, "tooltip": "Blend strength for b1 path."}),
|
|
"b2_mode": (list(blending_modes.keys()), {"tooltip": "Blending mode for b2 path."}),
|
|
"b2_blend": ("FLOAT", {"default": 1.0, "max": 100, "min": 0, "step": 0.001, "tooltip": "Blend strength for b2 path."}),
|
|
"threshold": ("INT", {"default": 1.0, "max": 10, "min": 1, "step": 1, "tooltip": "Base radius for the Fourier mask."}),
|
|
"use_override_scales": (["false", "true"], {"tooltip": "Enable manual override of scale presets."}),
|
|
"override_scales": ("STRING", {"default": '''# OVERRIDE SCALES
|
|
|
|
# Sharpen
|
|
# 10, 1.5''', "multiline": True, "tooltip": "Custom scale lines: '<radius>, <scale>'. Comments with #,//,!"}),
|
|
}
|
|
}
|
|
|
|
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_PostCFGShift:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"model": ("MODEL",),
|
|
"steps": ("INT", {"default": 20, "min": 1, "max": 1000, "step": 1, "tooltip": "Number of steps to apply SHIFT."}),
|
|
"mode": (list(blending_modes.keys()), {"tooltip": "Blend strategy for denoised vs denoised*b (e.g., inject, stable_slerp)."}),
|
|
"blend": ("FLOAT", {"default": 1.0, "max": 100.0, "min": 0.0, "step": 0.001, "tooltip": "Blend amount between base and scaled tensors."}),
|
|
"b": ("FLOAT", {"default": 1.1, "min": 0.0, "max": 10.0, "step": 0.001, "tooltip": "Scale factor for the injected path (higher = stronger)."}),
|
|
"apply_fourier": ("BOOLEAN", {"default": False, "tooltip": "Apply frequency-domain shaping (Fourier_filter)."}),
|
|
"multiscale_mode": (list(mscales.keys()), {"tooltip": "Preset shaping curves for Fourier_filter."}),
|
|
"multiscale_strength": ("FLOAT", {"default": 1.0, "max": 1.0, "min": 0.0, "step": 0.001, "tooltip": "Intensity of multi-scale shaping [0-1]."}),
|
|
"threshold": ("INT", {"default": 1, "min": 1, "max": 10, "step": 1, "tooltip": "Base radius for frequency mask."}),
|
|
"s": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 10.0, "step": 0.001, "tooltip": "Base scale value applied at threshold radius."}),
|
|
"force_gain": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01, "tooltip": "Final multiplier to boost or attenuate effect."}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "patch"
|
|
|
|
CATEGORY = "_for_testing"
|
|
|
|
def patch(self, model, steps, mode, blend, b, apply_fourier, multiscale_mode, multiscale_strength, threshold, s, force_gain):
|
|
|
|
scales = mscales[multiscale_mode]
|
|
steps = max(1, min(1000, steps))
|
|
current_step = 0
|
|
|
|
print(
|
|
"[FluxU] inputs:",
|
|
f"mode={mode}", f"blend={blend}", f"b={b}",
|
|
f"apply_fourier={apply_fourier}", f"multiscale_mode={multiscale_mode}", f"multiscale_strength={multiscale_strength}",
|
|
f"threshold={threshold}", f"s={s}", f"force_gain={force_gain}"
|
|
)
|
|
|
|
m = model.clone()
|
|
|
|
def post_cfg_function(args):
|
|
|
|
nonlocal current_step
|
|
current_step += 1
|
|
if current_step > steps:
|
|
return args.get("denoised")
|
|
|
|
denoised = args.get("denoised")
|
|
eff_blend = float(blend)
|
|
t_scaled = denoised * b
|
|
y = blending_modes[mode](denoised, t_scaled, eff_blend)
|
|
|
|
if apply_fourier:
|
|
y = Fourier_filter(y, threshold=threshold, scale=s, scales=scales, strength=multiscale_strength)
|
|
|
|
if force_gain != 1.0:
|
|
y = y * float(force_gain)
|
|
return y
|
|
|
|
|
|
try:
|
|
m.set_model_sampler_post_cfg_function(post_cfg_function)
|
|
print("[FluxU] set_model_sampler_post_cfg_function registered")
|
|
except Exception as e:
|
|
print(f"[FluxU] set_model_sampler_post_cfg_function failed: {e}")
|
|
return (m, )
|
|
|
|
class WAS_FreeU_V2:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"model": ("MODEL",),
|
|
"input_block": ("BOOLEAN", {"default": False, "tooltip": "Enable patching on the input block."}),
|
|
"middle_block": ("BOOLEAN", {"default": False, "tooltip": "Enable patching on the middle block."}),
|
|
"output_block": ("BOOLEAN", {"default": False, "tooltip": "Enable patching on the output block."}),
|
|
"multiscale_mode": (list(mscales.keys()), {"tooltip": "Frequency shaping preset used by Fourier_filter."}),
|
|
"multiscale_strength": ("FLOAT", {"default": 1.0, "max": 1.0, "min": 0, "step": 0.001, "tooltip": "Intensity of multi-scale shaping [0-1]."}),
|
|
"slice_b1": ("INT", {"default": 640, "min": 64, "max": 1280, "step": 1, "tooltip": "Slice width (channels) affected in 1280-channel features."}),
|
|
"slice_b2": ("INT", {"default": 320, "min": 64, "max": 640, "step": 1, "tooltip": "Slice width (channels) affected in 640-channel features."}),
|
|
"b1": ("FLOAT", {"default": 1.1, "min": 0.0, "max": 10.0, "step": 0.001, "tooltip": "Gain multiplier for 1280-channel slice."}),
|
|
"b2": ("FLOAT", {"default": 1.2, "min": 0.0, "max": 10.0, "step": 0.001, "tooltip": "Gain multiplier for 640-channel slice."}),
|
|
"s1": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 10.0, "step": 0.001, "tooltip": "Fourier scale at threshold for 1280-ch features."}),
|
|
"s2": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 10.0, "step": 0.001, "tooltip": "Fourier scale at threshold for 640-ch features."}),
|
|
},
|
|
"optional": {
|
|
"threshold": ("INT", {"default": 1.0, "max": 10, "min": 1, "step": 1, "tooltip": "Base radius for the Fourier mask."}),
|
|
"use_override_scales": (["false", "true"], {"tooltip": "Enable manual override of scale presets."}),
|
|
"override_scales": ("STRING", {"default": '''# OVERRIDE SCALES
|
|
|
|
# Sharpen
|
|
# 10, 1.5''', "multiline": True, "tooltip": "Custom scale lines: '<radius>, <scale>'. Comments with #,//,!"}),
|
|
}
|
|
}
|
|
|
|
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 = {
|
|
"WAS_FreeU": WAS_FreeU,
|
|
"WAS_FreeU_V2": WAS_FreeU_V2,
|
|
"WAS_PostCFGShift": WAS_PostCFGShift,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"WAS_FreeU": "FreeU (Advanced Plus)",
|
|
"WAS_FreeU_V2": "FreeU V2 (Advanced Plus)",
|
|
"WAS_PostCFGShift": "Post-CFG SHIFT",
|
|
}
|
|
|