diff --git a/core/guidance.py b/core/guidance.py new file mode 100644 index 0000000..ae9af5e --- /dev/null +++ b/core/guidance.py @@ -0,0 +1,186 @@ +import torch +import enum +import dataclasses +from .latent_filters import gaussian_blur_2d + + +class GuidanceCombinationMethod(str, enum.Enum): + PLAIN = "Plain" + PRED_SPACE_RESCALE = "PredSpaceRescale" + V_SPACE_RESCALE = "VSpaceRescale" + SNF = "SNF" + + +class AlternateConditionalComparator(str, enum.Enum): + BASE_COND = "BaseCond" + ALTERNATE_UNCOND = "AlternateUncond" + + +def plain_combination( + cond_x, + uncond_x, + cfg: float, + alternate_cond_x, + alternate_uncond_x, + alternate_weight: float, + alternate_cond_comparator: AlternateConditionalComparator, +) -> torch.Tensor: + if ( + alternate_cond_comparator == AlternateConditionalComparator.ALTERNATE_UNCOND + and alternate_uncond_x is None + ): + raise ValueError( + "Alternate uncond x is required for alternate uncond comparator" + ) + if alternate_cond_comparator == AlternateConditionalComparator.ALTERNATE_UNCOND: + alternate_guidance = alternate_uncond_x - alternate_cond_x + else: + alternate_guidance = alternate_cond_x - cond_x + return uncond_x + cfg * (cond_x - uncond_x) - alternate_weight * alternate_guidance + + +def v_space_normalised_combination( + x, + sigma_, + cond_x, + uncond_x, + cfg: float, + alternate_cond_x, + alternate_uncond_x, + alternate_weight: float, + rescaling_weight: float, + apply_rescaling_to_alternate: bool, + alternate_cond_comparator: AlternateConditionalComparator, +) -> torch.Tensor: + """ + Convert from x_0 prediction space to v prediction space and combine + with partial rescaling + + """ + # Reshape to match the dimensions of the input tensor + sigma = sigma_.view(sigma_.shape[:1] + (1,) * (cond_x.ndim - 1)) + + # Convert to variance preserving space + alpha = 1 / (sigma**2 + 1.0) ** 0.5 + sigma = sigma * alpha + sigma_inv = 1.0 / sigma + x = x * alpha + + cond_v = (alpha * x - cond_x) * sigma_inv + uncond_v = (alpha * x - uncond_x) * sigma_inv + + s0 = torch.std(cond_v, dim=(1, 2, 3), keepdim=True) + if apply_rescaling_to_alternate: + alternate_cond_v = (alpha * x - alternate_cond_x) * sigma_inv + if alternate_cond_comparator == AlternateConditionalComparator.ALTERNATE_UNCOND: + alternate_uncond_v = (alpha * x - alternate_uncond_x) * sigma_inv + alternate_guidance = alternate_cond_v - alternate_uncond_v + else: + alternate_guidance = alternate_cond_v - cond_v + v_out_unnorm = ( + uncond_v + cfg * (cond_v - uncond_v) - alternate_weight * alternate_guidance + ) + 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_weight + v_out_unnorm * ( + 1 - rescaling_weight + ) + # Convert back to prediction space and rescale + pred_final = x * alpha - v_final * sigma + else: + # Apply standard CFG rescaling and combine with alternate + v_cfg_unnorm = uncond_v + cfg * (cond_v - uncond_v) + s1 = torch.std(v_cfg_unnorm, dim=(1, 2, 3), keepdim=True) + v_cfg_normed = v_cfg_unnorm * s0 / s1 + v_cfg = v_cfg_normed * rescaling_weight + v_cfg_unnorm * (1 - rescaling_weight) + pred_cfg = x * alpha - v_cfg * sigma + if alternate_cond_comparator == AlternateConditionalComparator.ALTERNATE_UNCOND: + alternate_guidance = alternate_uncond_x - alternate_cond_x + else: + alternate_guidance = alternate_cond_x - cond_x + pred_final = pred_cfg - alternate_weight * alternate_guidance + + return pred_final + + +def pred_space_normalised_combination( + cond_x, + uncond_x, + cfg: float, + alternate_cond_x, + alternate_uncond_x, + alternate_weight: float, + rescaling_weight: float, + apply_rescaling_to_alternate: bool, + alternate_cond_comparator: AlternateConditionalComparator, +) -> torch.Tensor: + + s0 = torch.std(cond_x, dim=(1, 2, 3), keepdim=True) + if apply_rescaling_to_alternate: + if alternate_cond_comparator == AlternateConditionalComparator.ALTERNATE_UNCOND: + alternate_guidance = alternate_cond_x - alternate_uncond_x + else: + alternate_guidance = alternate_cond_x - cond_x + x_out_unnorm = ( + uncond_x + cfg * (cond_x - uncond_x) - alternate_weight * alternate_guidance + ) + s1 = torch.std(x_out_unnorm, dim=(1, 2, 3), keepdim=True) + x_out_normed = x_out_unnorm * s0 / s1 + x_final = x_out_normed * rescaling_weight + x_out_unnorm * ( + 1 - rescaling_weight + ) + else: + # Apply standard CFG rescaling and combine with alternate + x_cfg_unnorm = uncond_x + cfg * (cond_x - uncond_x) + s1 = torch.std(x_cfg_unnorm, dim=(1, 2, 3), keepdim=True) + x_cfg_normed = x_cfg_unnorm * s0 / s1 + x_cfg = x_cfg_normed * rescaling_weight + x_cfg_unnorm * (1 - rescaling_weight) + if alternate_cond_comparator == AlternateConditionalComparator.ALTERNATE_UNCOND: + alternate_guidance = alternate_uncond_x - alternate_cond_x + else: + alternate_guidance = alternate_cond_x - cond_x + x_final = x_cfg - alternate_weight * alternate_guidance + + return x_final + + +def snf_combination( + cond_a: torch.Tensor, + uncond_a: torch.Tensor, + cond_b: torch.Tensor, + uncond_b: torch.Tensor, + weight_a: float, + weight_b: float, + alternate_cond_comparator: AlternateConditionalComparator, +): + """ + From Saliency-adaptive Noise Fusion based on High-fidelity Person-centric Subject-to-Image Synthesis (Wang et al.) + :param cond_a: + :param uncond_a: + :param cond_b: + :param uncond_b: + :param weight_a: + :param weight_b: + :param alternate_cond_comparator: + :return: + """ + x_a = weight_a * (cond_a - uncond_a) + if alternate_cond_comparator == AlternateConditionalComparator.ALTERNATE_UNCOND: + x_b = weight_b * (cond_b - uncond_b) + else: + x_b = weight_b * (cond_b - cond_a) + + b, c, h, w = x_a.shape + + a_bar = gaussian_blur_2d(torch.abs(x_a), 3, 1.0) + b_bar = gaussian_blur_2d(torch.abs(x_b), 3, 1.0) + a_softmax = torch.softmax(a_bar.reshape(b * c, h * w), dim=1).reshape(b, c, h, w) + b_softmax = torch.softmax(b_bar.reshape(b * c, h * w), dim=1).reshape(b, c, h, w) + guidance_stacked = torch.stack([x_a, x_b], dim=0) + + ab_softmax = torch.stack([a_softmax, b_softmax], dim=0) + argeps = torch.argmax(ab_softmax, dim=0, keepdim=True) + + # TODO : should do softmax instead of argmax + snf = torch.gather(guidance_stacked, dim=0, index=argeps).squeeze(0) + return snf diff --git a/core/unet.py b/core/unet.py index 2f28200..9c05b18 100644 --- a/core/unet.py +++ b/core/unet.py @@ -1,7 +1,7 @@ import math -from typing import Callable from torch import Tensor from itertools import groupby +from kornia.filters import GaussianBlur2d from .latent_filters import gaussian_blur_2d @@ -13,7 +13,79 @@ def pag_perturbed_attention( return v -def seg_attention_wrapper(attention: Callable, blur_sigma: float = 10.0): +def seg_attention_wrapper( + attention: callable, blur_sigma: float = 10.0, border_mode: str = "reflect" +): + """ + Wraps an attention function to apply a Gaussian blur (via Kornia) on q before computing attention. + + Args: + attention: The attention function to wrap (must accept q, k, v, ...) + blur_sigma: If >= 0, apply Gaussian blur with this sigma. If < 0, replace q by the global mean. + border_mode: Passed to Kornia's GaussianBlur2d for handling edges + ('reflect', 'replicate', 'constant', 'circular'). + """ + + def seg_perturbed_attention( + q: Tensor, k: Tensor, v: Tensor, extra_options, mask=None + ) -> Tensor: + 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 + + # Reshape (B, area, dim) -> (B, dim, H, W) + 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: + # Compute kernel size from sigma + kernel_size = math.ceil(6 * blur_sigma) + if kernel_size % 2 == 0: + kernel_size += 1 + + # Cap kernel if using reflection (or any border mode that can't handle huge pads) + # The largest reflection pad we can do is half the dimension minus 1 + # so the kernel size can't exceed 2*dim - 1 in either dimension. + max_k = 2 * min(q.shape[-2], q.shape[-1]) - 1 + kernel_size = min(kernel_size, max_k) + + # 3) If the kernel is still >= 3, apply Kornia blur. Else, skip + if kernel_size >= 3: + blur = GaussianBlur2d( + kernel_size=(kernel_size, kernel_size), + sigma=(blur_sigma, blur_sigma), + border_type=border_mode, + ) + q = blur(q) + else: + pass # Can't blur with kernel_size < 3 + + else: + # Negative blur_sigma => set entire q to the mean + q[:] = q.mean(dim=(-2, -1), keepdim=True) + + # Reshape back to (B, area, dim) for the attention function + q = q.reshape(bs, inner_dim, -1).permute(0, 2, 1) + + return attention(q, k, v, heads=heads) + + return seg_perturbed_attention + + +def seg_attention_wrapper_old(attention: callable, blur_sigma: float = 10.0): + """ + Faster but not as clean version of seg_attention_wrapper. + :param attention: + :param blur_sigma: + :return: + """ + def seg_perturbed_attention( q: Tensor, k: Tensor, v: Tensor, extra_options, mask=None ) -> Tensor: @@ -55,7 +127,7 @@ def parse_unet_blocks(model, unet_block_list: str): # Get all Self-attention blocks input_blocks, middle_blocks, output_blocks = [], [], [] - for name, module in model.model.diffusion_model.named_modules(): + for name, module in model.diffusion_model.named_modules(): if module.__class__.__name__ == "CrossAttention" and name.endswith("attn1"): parts = name.split(".") block_name = parts[0] diff --git a/nodes/attention.py b/nodes/attention.py index b348370..88d0ee6 100644 --- a/nodes/attention.py +++ b/nodes/attention.py @@ -3,65 +3,12 @@ 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 ..core import unet as unet - - -# import torch -def compute_v_space_normalised_combination( - x: torch.Tensor, - sigma_: float, - cond_x: torch.Tensor, - uncond_x: torch.Tensor, - cond_scale: float, - alternate_cond_x: torch.Tensor, - alternate_weight: float, - normalised_fraction: float, - apply_rescaling_to_alternate: bool, -) -> torch.Tensor: - """ - Convert from x_0 prediction space to v prediction space and combine - with partial rescaling - - """ - # Reshape to match the dimensions of the input tensor - sigma = sigma_.view(sigma_.shape[:1] + (1,) * (cond_x.ndim - 1)) - - # Convert to variance preserving space - alpha = 1 / (sigma**2 + 1.0) ** 0.5 - sigma = sigma * alpha - sigma_inv = 1.0 / sigma - x = x * alpha - - cond_v = (alpha * x - cond_x) * sigma_inv - uncond_v = (alpha * x - uncond_x) * sigma_inv - - s0: torch.Tensor = torch.std(cond_v, dim=(1, 2, 3), keepdim=True) - if apply_rescaling_to_alternate: - alternate_cond_v = (alpha * x - alternate_cond_x) * sigma_inv - v_out_unnorm: torch.Tensor = ( - uncond_v - + cond_scale * (cond_v - uncond_v) - - alternate_weight * (alternate_cond_v - cond_v) - ) - s1: torch.Tensor = torch.std(v_out_unnorm, dim=(1, 2, 3), keepdim=True) - v_out_normed: torch.Tensor = v_out_unnorm * s0 / s1 - v_final: torch.Tensor = v_out_normed * normalised_fraction + v_out_unnorm * ( - 1 - normalised_fraction - ) - # Convert back to prediction space and rescale - # pred_final = (alpha * x - v_final * sigma) / alpha - pred_final = x * alpha - v_final * sigma - else: - # Apply standard CFG rescaling and combine with alternate - v_cfg_unnorm: torch.Tensor = uncond_v + cond_scale * (cond_v - uncond_v) - s1: torch.Tensor = torch.std(v_cfg_unnorm, dim=(1, 2, 3), keepdim=True) - v_cfg_normed: torch.Tensor = v_cfg_unnorm * s0 / s1 - v_cfg: torch.Tensor = v_cfg_normed * normalised_fraction + v_cfg_unnorm * ( - 1 - normalised_fraction - ) - pred_cfg = x * alpha - v_cfg * sigma - pred_final = pred_cfg - alternate_weight * (alternate_cond_x - cond_x) - - return pred_final +from ..core.guidance import ( + v_space_normalised_combination, + pred_space_normalised_combination, + snf_combination, + AlternateConditionalComparator, +) class RescaledPAG: @@ -181,16 +128,18 @@ class RescaledPAG: else: pag_cond_pred = cond_pred - pred_final = compute_v_space_normalised_combination( + pred_final = v_space_normalised_combination( x, sigma, cond_pred, uncond_pred, cond_scale, pag_cond_pred, + None, pag_weight, rescaling_fraction, apply_rescaling_to_pag, + AlternateConditionalComparator.BASE_COND, ) return pred_final @@ -310,13 +259,19 @@ class RescaledSEG: blocks = unet.parse_unet_blocks(model, unet_block_list) else: blocks = [(unet_block, unet_block_id, None)] + # blocks = ( + # # [("input", i, None) for i in range(4)] + # # + [("middle", 0, None)] + # # + [("output", i, None) for i in range(4)] + # ) for block in blocks: layer, number, index = block model_options = set_model_options_patch_replace( model_options, unet.seg_attention_wrapper( - optimized_attention, blur_sigma=seg_sigma + optimized_attention, + blur_sigma=seg_sigma, ), "attn1", layer, @@ -330,17 +285,28 @@ class RescaledSEG: else: seg_cond_pred = cond_pred - pred_final = compute_v_space_normalised_combination( + pred_final = v_space_normalised_combination( x, sigma, cond_pred, uncond_pred, cond_scale, seg_cond_pred, + None, seg_weight, rescaling_fraction, apply_rescaling_to_seg, + AlternateConditionalComparator.BASE_COND, ) + # pred_final = uncond_pred + snf_combination( + # cond_pred, + # uncond_pred, + # seg_cond_pred, + # None, + # cond_scale, + # seg_weight, + # AlternateConditionalComparator.BASE_COND, + # ) return pred_final m = model.clone()