diff --git a/comfy_extras/nodes_freelunch.py b/comfy_extras/nodes_freelunch.py new file mode 100644 index 00000000..535eece3 --- /dev/null +++ b/comfy_extras/nodes_freelunch.py @@ -0,0 +1,56 @@ +#code originally taken from: https://github.com/ChenyangSi/FreeU (under MIT License) + +import torch + + +def Fourier_filter(x, threshold, scale): + # FFT + x_freq = torch.fft.fftn(x.float(), dim=(-2, -1)) + x_freq = torch.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 + x_freq = x_freq * mask + + # IFFT + x_freq = torch.fft.ifftshift(x_freq, dim=(-2, -1)) + x_filtered = torch.fft.ifftn(x_freq, dim=(-2, -1)).real + + return x_filtered.to(x.dtype) + + +class FreeU: + @classmethod + def INPUT_TYPES(s): + return {"required": { "model": ("MODEL",), + "b1": ("FLOAT", {"default": 1.1, "min": 0.0, "max": 10.0, "step": 0.01}), + "b2": ("FLOAT", {"default": 1.2, "min": 0.0, "max": 10.0, "step": 0.01}), + "s1": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 10.0, "step": 0.01}), + "s2": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 10.0, "step": 0.01}), + }} + RETURN_TYPES = ("MODEL",) + FUNCTION = "patch" + + CATEGORY = "_for_testing" + + def patch(self, model, b1, b2, s1, s2): + def output_block_patch(h, hsp, transformer_options): + if h.shape[1] == 1280: + h[:,:640] = h[:,:640] * b1 + hsp = Fourier_filter(hsp, threshold=1, scale=s1) + if h.shape[1] == 640: + h[:,:320] = h[:,:320] * b2 + hsp = Fourier_filter(hsp, threshold=1, scale=s2) + return h, hsp + + m = model.clone() + m.set_model_output_block_patch(output_block_patch) + return (m, ) + + +NODE_CLASS_MAPPINGS = { + "FreeU": FreeU, +} diff --git a/nodes.py b/nodes.py index 6e0d4374..11586260 100644 --- a/nodes.py +++ b/nodes.py @@ -1782,4 +1782,5 @@ def init_custom_nodes(): load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_tomesd.py")) load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_clip_sdxl.py")) load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_canny.py")) + load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_freelunch.py")) load_custom_nodes()