diff --git a/nodes.py b/nodes.py index 11bcc9f..4ff1247 100644 --- a/nodes.py +++ b/nodes.py @@ -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)", }