From 22d0d82e4ad9f3fa2ec40836585eecbd5cb4762d Mon Sep 17 00:00:00 2001 From: JTriggerFish Date: Wed, 29 Jan 2025 15:11:40 +0000 Subject: [PATCH] SEG initial code --- core/latent_filters.py | 25 +++++ core/{diffusion.py => unet.py} | 39 ++++++- nodes/attention.py | 184 ++++++++++++++++++++++++++++++--- 3 files changed, 232 insertions(+), 16 deletions(-) rename core/{diffusion.py => unet.py} (62%) diff --git a/core/latent_filters.py b/core/latent_filters.py index 635fa34..d92ba65 100644 --- a/core/latent_filters.py +++ b/core/latent_filters.py @@ -5,6 +5,8 @@ import torchvision.transforms as transforms import torch import torch.nn.functional as F +from torch import Tensor +from torch.nn import functional as F def center_tensor( @@ -180,3 +182,26 @@ def latent_upscale( # Upscale each channel separately res = upscale(latent) return res + + +def gaussian_blur_2d(img: Tensor, kernel_size: int, sigma: float) -> Tensor: + height = img.shape[-1] + kernel_size = min(kernel_size, height - (height % 2 - 1)) + ksize_half = (kernel_size - 1) * 0.5 + + x = torch.linspace(-ksize_half, ksize_half, steps=kernel_size) + + pdf = torch.exp(-0.5 * (x / sigma).pow(2)) + + x_kernel = pdf / pdf.sum() + x_kernel = x_kernel.to(device=img.device, dtype=img.dtype) + + kernel2d = torch.mm(x_kernel[:, None], x_kernel[None, :]) + kernel2d = kernel2d.expand(img.shape[-3], 1, kernel2d.shape[0], kernel2d.shape[1]) + + padding = [kernel_size // 2, kernel_size // 2, kernel_size // 2, kernel_size // 2] + + img = F.pad(img, padding, mode="reflect") + img = F.conv2d(img, kernel2d, groups=img.shape[-3]) + + return img diff --git a/core/diffusion.py b/core/unet.py similarity index 62% rename from core/diffusion.py rename to core/unet.py index fa7f420..67836bf 100644 --- a/core/diffusion.py +++ b/core/unet.py @@ -1,16 +1,49 @@ import math from typing import Callable -import torch from torch import Tensor -import torch.nn.functional as F from itertools import groupby +from latent_filters import gaussian_blur_2d -def pag_perturbed_attention(q: Tensor, k: Tensor, v: Tensor, extra_options, mask=None): + +def pag_perturbed_attention( + q: Tensor, k: Tensor, v: Tensor, extra_options, mask=None +) -> Tensor: """Perturbed self-attention corresponding to an identity matrix replacing the attention matrix.""" return v +def seg_attention_wrapper(attention: Callable, blur_sigma: float = 10.0): + def seg_perturbed_attention( + q: Tensor, k: Tensor, v: Tensor, extra_options, mask=None + ) -> Tensor: + """Smoothed Energy Guidance self-attention""" + heads = extra_options["n_heads"] + bs, area, inner_dim = q.shape + + height_orig, width_orig = extra_options["original_shape"][2:4] + aspect_ratio = width_orig / height_orig + + if aspect_ratio >= 1.0: + height = round((area / aspect_ratio) ** 0.5) + q = q.permute(0, 2, 1).reshape(bs, inner_dim, height, -1) + else: + width = round((area * aspect_ratio) ** 0.5) + q = q.permute(0, 2, 1).reshape(bs, inner_dim, -1, width) + + if blur_sigma >= 0: + kernel_size = math.ceil(6 * blur_sigma) + 1 - math.ceil(6 * blur_sigma) % 2 + q = gaussian_blur_2d(q, kernel_size, blur_sigma) + else: + q[:] = q.mean(dim=(-2, -1), keepdim=True) + + q = q.reshape(bs, inner_dim, -1).permute(0, 2, 1) + + return attention(q, k, v, heads=heads) + + return seg_perturbed_attention + + def parse_unet_blocks(model, unet_block_list: str): """ Copied from https://github.com/pamparamm/sd-perturbed-attention/blob/master/pag_utils.py#L9 diff --git a/nodes/attention.py b/nodes/attention.py index 17d0aec..988affd 100644 --- a/nodes/attention.py +++ b/nodes/attention.py @@ -1,9 +1,8 @@ import torch -import comfy from comfy.model_patcher import ModelPatcher, set_model_options_patch_replace from comfy.samplers import calc_cond_batch from comfy.ldm.modules.attention import optimized_attention -from diffusion import parse_unet_blocks, pag_perturbed_attention +from unet import parse_unet_blocks, pag_perturbed_attention, seg_attention_wrapper # import torch @@ -16,7 +15,7 @@ class RescaledPAG: return { "required": { "model": ("MODEL",), - "scale": ( + "pag_weight": ( "FLOAT", { "default": 1.5, @@ -26,7 +25,7 @@ class RescaledPAG: "round": 0.01, }, ), - "post_rescale": ( + "rescaling_fraction": ( "FLOAT", { "default": 0.9, @@ -76,8 +75,8 @@ class RescaledPAG: def patch( self, model: ModelPatcher, - scale: float = 3.0, - post_rescale: float = 0.7, + pag_weight: float = 3.0, + rescaling_fraction: float = 0.7, unet_block: str = "middle", unet_block_id: int = 0, noise_fraction_start: float = 1.0, @@ -86,8 +85,8 @@ class RescaledPAG: ): """ :param model: - :param scale: - :param post_rescale: + :param pag_weight: + :param rescaling_fraction: :param unet_block: :param unet_block_id: :param noise_fraction_start: @@ -147,13 +146,15 @@ class RescaledPAG: v_out_unnorm = ( uncond_v + cond_scale * (cond_v - uncond_v) - - scale * (pag_cond_v - cond_v) + - pag_weight * (pag_cond_v - cond_v) ) s0 = torch.std(cond_v, dim=(1, 2, 3), keepdim=True) s1 = torch.std(v_out_unnorm, dim=(1, 2, 3), keepdim=True) v_out_normed = v_out_unnorm * s0 / s1 - v_final = v_out_normed * post_rescale + v_out_unnorm * (1 - post_rescale) + v_final = v_out_normed * rescaling_fraction + v_out_unnorm * ( + 1 - rescaling_fraction + ) pred_final = alpha * x - v_final * sigma return pred_final @@ -163,9 +164,166 @@ class RescaledPAG: return (m,) -NODE_CLASS_MAPPINGS = { - "RescaledPAG": RescaledPAG, -} +class RescaledSEG: + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "seg_sigma": ( + "FLOAT", + { + "default": 30.0, + "min": 0.0, + "max": 9999.0, + "step": 5.0, + "round": 0.01, + }, + ), + "seg_weight": ( + "FLOAT", + { + "default": 1.5, + "min": 0.0, + "max": 100.0, + "step": 0.1, + "round": 0.01, + }, + ), + "rescaling_fraction": ( + "FLOAT", + { + "default": 0.9, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "round": 0.0001, + }, + ), + "unet_block": (["input", "middle", "output"], {"default": "middle"}), + "unet_block_id": ("INT", {"default": 0}), + "noise_fraction_start": ( + "FLOAT", + { + "default": 1.0, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "round": False, + }, + ), + "noise_fraction_end": ( + "FLOAT", + { + "default": 0.0, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "round": False, + }, + ), + }, + "optional": { + "unet_block_list": ("STRING", {"default": ""}), + }, + } + + RETURN_TYPES = ("MODEL",) + FUNCTION = "patch" + CATEGORY = "model_patches/unet" + OUTPUT_NODE = False + + @classmethod + def help(cls): + return "Implementation of Smoothed Energy Guidance ( SEG ) that also incorporates the Rescale CFG approach" + + def patch( + self, + model: ModelPatcher, + seg_sigma: float = 30.0, + seg_weight: float = 3.0, + rescaling_fraction: float = 0.7, + unet_block: str = "middle", + unet_block_id: int = 0, + noise_fraction_start: float = 1.0, + noise_fraction_end: float = 0.0, + unet_block_list: str = "", + ): + + def cfg_function(args): + """Rescaled CFG+PAG""" + model = args["model"] + cond_pred = args["cond_denoised"] # x_0 + uncond_pred = args["uncond_denoised"] + cond_scale = args["cond_scale"] + cond = args["cond"] + sigma_ = args["sigma"] + model_options = args["model_options"].copy() + sigma = sigma_.view(sigma_.shape[:1] + (1,) * (cond_pred.ndim - 1)) + x = args["input"] + + noise_fracs = model_options["transformer_options"]["sample_sigmas"] ** 2 + current_frac = sigma_**2 / noise_fracs[0] + + alpha = 1 / (sigma**2 + 1.0) ** 0.5 + sigma = sigma / (sigma**2 + 1.0) ** 0.5 + sigma_inv = 1.0 / sigma + + cond_v = (alpha * x - cond_pred) * sigma_inv + uncond_v = (alpha * x - uncond_pred) * sigma_inv + + if (noise_fraction_end <= current_frac) and ( + current_frac <= noise_fraction_start + ): + if unet_block_list: + blocks = parse_unet_blocks(model, unet_block_list) + else: + blocks = [(unet_block, unet_block_id, None)] + + for block in blocks: + layer, number, index = block + model_options = set_model_options_patch_replace( + model_options, + seg_attention_wrapper( + optimized_attention, blur_sigma=seg_sigma + ), + "attn1", + layer, + number, + index, + ) + + (pag_cond_pred,) = calc_cond_batch( + model, [cond], x, sigma_, model_options + ) + pag_cond_v = (alpha * x - pag_cond_pred) * sigma_inv + else: + pag_cond_v = cond_v + + v_out_unnorm = ( + uncond_v + + cond_scale * (cond_v - uncond_v) + - seg_weight * (pag_cond_v - cond_v) + ) + s0 = torch.std(cond_v, dim=(1, 2, 3), keepdim=True) + s1 = torch.std(v_out_unnorm, dim=(1, 2, 3), keepdim=True) + + v_out_normed = v_out_unnorm * s0 / s1 + v_final = v_out_normed * rescaling_fraction + v_out_unnorm * ( + 1 - rescaling_fraction + ) + + pred_final = alpha * x - v_final * sigma + return pred_final + + m = model.clone() + m.set_model_sampler_post_cfg_function(cfg_function) + return (m,) + + +NODE_CLASS_MAPPINGS = {"RescaledPAG": RescaledPAG, "RescaledSEG": RescaledSEG} NODE_DISPLAY_NAME_MAPPINGS = { "RescaledPAG": "Rescaled PAG", + "RescaledSEG": "Rescaled SEG", }