diff --git a/core/guidance.py b/core/guidance.py index 4e60951..8b6da39 100644 --- a/core/guidance.py +++ b/core/guidance.py @@ -15,6 +15,9 @@ try: gaussian_blur_2d, gaussian_kernel_size_for_img, moment_match, + tokens_to_spatial, + spatial_to_tokens, + normalize_tensor, ) from . import utils except ImportError: @@ -34,11 +37,11 @@ class GuidanceType(str, enum.Enum): class GuidanceScalingMethod(str, enum.Enum): NONE = "None" - MOMENT_MATCH = "MomentMatch" PRED_SPACE_RESCALE = "PredSpaceRescale" V_SPACE_RESCALE = "VSpaceRescale" SNF = "SNF" SOFTMAX = "Softmax" + NORMALIZE = "Normalize" def plain_guidance_combine( @@ -108,6 +111,45 @@ def softmax_guidance_combine( return base + utils.softmax_weighted_combination(cfg, guidance_adj, temperature) +def guidance_normalize( + cond_pred: Tensor, + uncond_pred: Tensor, + cfg: float, + alternate_pred: Tensor | None, + alternate_pred_ref: Tensor | None, + alternate_guidance_weight: float, + apply_rescaling_to_guidance: bool, + rescaling_fraction: float = 1.0, + target_std: float = 1.0, + per_channel: bool = True, +) -> Tensor: + + base_guidance = uncond_pred + cfg * (cond_pred - uncond_pred) + + if alternate_pred is None: + guidance_adj = 0 + else: + if alternate_pred_ref is None: + alternate_pred_ref = 0 + guidance_adj = alternate_guidance_weight * (alternate_pred_ref - alternate_pred) + + if apply_rescaling_to_guidance: + guidance = base_guidance + guidance_adj + guidance = ( + 1 - rescaling_fraction + ) * guidance + rescaling_fraction * normalize_tensor( + guidance, target_std=target_std, per_channel=per_channel + ) + return guidance + else: + base_guidance = ( + 1 - rescaling_fraction + ) * base_guidance + rescaling_fraction * normalize_tensor( + base_guidance, target_std=target_std, per_channel=per_channel + ) + return base_guidance + guidance_adj + + def pred_rescaled_guidance_combine( cond_pred: Tensor, uncond_pred: Tensor, @@ -130,7 +172,10 @@ def pred_rescaled_guidance_combine( if apply_rescaling_to_guidance: guidance = base_guidance + guidance_adj guidance = utils.partial_rescaling( - ref_tensor=cond_pred, z=guidance, rescaled_fraction=rescaling_fraction + ref_tensor=cond_pred, + z=guidance, + rescaled_fraction=rescaling_fraction, + match_mean=False, ) return guidance else: @@ -140,40 +185,6 @@ def pred_rescaled_guidance_combine( return base_guidance + guidance_adj -def moment_match_guidance_combine( - cond_pred: Tensor, - uncond_pred: Tensor, - cfg: float, - alternate_pred: Tensor | None, - alternate_pred_ref: Tensor | None, - alternate_guidance_weight: float, - apply_rescaling_to_guidance: bool, - rescaling_fraction: float, -) -> Tensor: - base_guidance = uncond_pred + cfg * (cond_pred - uncond_pred) - - if alternate_pred is None: - guidance_adj = 0 - else: - if alternate_pred_ref is None: - alternate_pred_ref = 0 - guidance_adj = alternate_guidance_weight * (alternate_pred_ref - alternate_pred) - - if apply_rescaling_to_guidance: - guidance = base_guidance + guidance_adj - guidance_rescaled = moment_match(cond_pred, guidance) - return ( - 1 - rescaling_fraction - ) * guidance + rescaling_fraction * guidance_rescaled - else: - base_guidance_rescaled = moment_match(cond_pred, base_guidance) - return ( - (1 - rescaling_fraction) * base_guidance - + rescaling_fraction * base_guidance_rescaled - + guidance_adj - ) - - def v_space_rescaled_guidance_combine( x: Tensor, sigma: Tensor, @@ -210,12 +221,18 @@ def v_space_rescaled_guidance_combine( if apply_rescaling_to_guidance: guidance = base_guidance + guidance_adj guidance = utils.partial_rescaling( - ref_tensor=cond, z=guidance, rescaled_fraction=rescaling_fraction + ref_tensor=cond, + z=guidance, + rescaled_fraction=rescaling_fraction, + match_mean=False, ) final_pred = utils.v_to_pred(x, sigma, [guidance])[0] else: base_guidance = utils.partial_rescaling( - ref_tensor=cond, z=base_guidance, rescaled_fraction=rescaling_fraction + ref_tensor=cond, + z=base_guidance, + rescaled_fraction=rescaling_fraction, + match_mean=False, ) final_pred = utils.v_to_pred(x, sigma, [base_guidance])[0] final_pred += guidance_adj @@ -290,8 +307,8 @@ def guidance_combine_and_scale( alternate_guidance_weight, temperature=rescaling_fraction, ) - case GuidanceScalingMethod.MOMENT_MATCH: - return moment_match_guidance_combine( + case GuidanceScalingMethod.NORMALIZE: + return guidance_normalize( cond_pred, uncond_pred, cfg, @@ -366,13 +383,7 @@ def seg_attention_wrapper(scaled_blur_sigma: float = 0.1) -> callable: 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) + q = tokens_to_spatial(q, aspect_ratio) blur_sigma = scaled_blur_sigma * math.sqrt(height_orig * width_orig * 64) / 10 @@ -385,8 +396,7 @@ def seg_attention_wrapper(scaled_blur_sigma: float = 0.1) -> callable: # 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) + q = spatial_to_tokens(q) return optimized_attention(q, k, v, heads=heads) @@ -549,68 +559,37 @@ def fuzzy_attention_wrapper( return fuzzy_attention -def downsampled_attention_wrapper( - attention: callable, - downsample_factor: float, +def upscale_and_transfer_previous_attention_wrapper( + downsample_factor: float = 2.0, ) -> callable: - def downsampled_attention( + prev_k = None + prev_v = None + prev_q = None + + def transer_attention( q: Tensor, k: Tensor, v: Tensor, - extra_options, + extra_options: dict, 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) - k = k.permute(0, 2, 1).reshape(bs, inner_dim, height, -1) - # v = v.permute(0, 2, 1).reshape(bs, inner_dim, height, -1) + nonlocal prev_k, prev_v, prev_q + if prev_k is None: + prev_k = k + prev_v = v + prev_q = q + return optimized_attention( + q, k, v, heads=extra_options["n_heads"], mask=mask + ) else: - width = round((area * aspect_ratio) ** 0.5) - # q = q.permute(0, 2, 1).reshape(bs, inner_dim, -1, width) - k = k.permute(0, 2, 1).reshape(bs, inner_dim, -1, width) - # v = v.permute(0, 2, 1).reshape(bs, inner_dim, -1, width) + # TODO + pass - # q = kornia.geometry.resize( - # q, (new_height, new_width), interpolation="area", antialias=True - # ) - k = kornia.geometry.rescale( - k, 1 / downsample_factor, interpolation="area", antialias=True - ) - k = kornia.geometry.rescale( - k, float(downsample_factor), interpolation="nearest", antialias=True - ) - # v = kornia.geometry.resize( - # v, (new_height, new_width), interpolation="area", antialias=True - # ) - - # Reshape back to (B, area, dim) for the attention function - # q = q.reshape(bs, inner_dim, -1).permute(0, 2, 1) - k = k.reshape(bs, inner_dim, -1).permute(0, 2, 1) - # v = v.reshape(bs, inner_dim, -1).permute(0, 2, 1) - - res = attention(q, k, v, heads=heads) - - # Re upsample the attention map - # res = res.permute(0, 2, 1).reshape(bs, inner_dim, new_height, new_width) - # res = kornia.geometry.resize( - # res, (height, width), interpolation="area", antialias=True - # ) - # res = res.reshape(bs, inner_dim, -1).permute(0, 2, 1) - return res - - return downsampled_attention + return transer_attention -def copy_attention_keys_queries_wrapper() -> callable: +def batch_copy_attention_keys_queries_wrapper() -> callable: def copy_attention_keys_queries( q: Tensor, @@ -684,7 +663,7 @@ def random_rotation_wrapper( S *= 0.5 return S - def partial_random_orthonormal_attention( + def random_orthonormal_attention( q: Tensor, k: Tensor, v: Tensor, @@ -722,7 +701,7 @@ def random_rotation_wrapper( return optimized_attention(rotated_q, k, v, heads=heads, mask=mask) - return partial_random_orthonormal_attention + return random_orthonormal_attention def sliding_window_guidance_wrapper( diff --git a/core/latent_filters.py b/core/latent_filters.py index 866d74c..2e617c6 100644 --- a/core/latent_filters.py +++ b/core/latent_filters.py @@ -6,6 +6,7 @@ import torch from kornia.filters import GaussianBlur2d from torch import Tensor from torch.nn import functional as F +import matplotlib as plt def center_tensor( @@ -86,30 +87,59 @@ def huberize_quantile( return y +def normalize_tensor( + tensor: torch.Tensor, + per_channel: bool = True, + normalize_mean: bool = False, + eps: float = 1e-8, + target_mean: float | Tensor = 0.0, + target_std: float | Tensor = 1.0, +) -> torch.Tensor: + """ + Normalize a tensor to have zero mean and unit variance. + :param tensor: + :param per_channel: + :param eps: + :return: + """ + if not per_channel: + mean = tensor.mean() + std = tensor.std() + eps + else: + mean = tensor.mean(dim=(-2, -1), keepdim=True) + std = tensor.std(dim=(-2, -1), keepdim=True) + eps + if normalize_mean: + return (tensor - mean) * target_std / std + target_mean + else: + return (tensor - mean) * target_std / std + mean + + def moment_match( - from_tensor: torch.Tensor, to_tensor: torch.Tensor, per_channel: bool = True + from_tensor: torch.Tensor, + to_tensor: torch.Tensor, + per_channel: bool = True, + match_mean: bool = True, ) -> torch.Tensor: """ Match the mean and variance of a tensor to another tensor. :param from_tensor: The tensor to match. - :param to_tensor: The tensor to match to. - :return: The transformed tensor. + :param to_tensor: + :param per_channel: + :param match_mean: """ if not per_channel: - from_mean = from_tensor.mean() - from_std = from_tensor.std() - to_mean = to_tensor.mean() - to_std = to_tensor.std() - return (to_tensor + to_mean - from_mean) * to_std / from_std + target_mean = from_tensor.mean() + target_std = from_tensor.std() else: - ret = to_tensor.clone() - ret += from_tensor.mean(dim=(2, 3), keepdim=True) - ret.mean( - dim=(2, 3), keepdim=True - ) - ret *= from_tensor.std(dim=(2, 3), keepdim=True) / ret.std( - dim=(2, 3), keepdim=True - ) - return ret + target_mean = from_tensor.mean(dim=(-2, -1), keepdim=True) + target_std = from_tensor.std(dim=(-2, -1), keepdim=True) + return normalize_tensor( + to_tensor, + per_channel, + target_mean=target_mean, + target_std=target_std, + normalize_mean=match_mean, + ) def add_correlated_gaussian_noise( @@ -280,3 +310,209 @@ def _gaussian_blur_2d_manual(img: Tensor, kernel_size: int, sigma: float) -> Ten img = F.conv2d(img, kernel2d, groups=img.shape[-3]) return img + + +def compute_cosine_distance( + tensorA: Tensor, tensorB: Tensor, eps=1e-8, return_similarity: bool = False +) -> Tensor: + """ + Compute element-wise cosine distance between two tensors of shape [B, C, H, W]. + Returns a heatmap of shape [B, H, W] where each pixel is: + 1 - cos_sim = 1 - (a·b / (||a||*||b||)). + """ + # Ensure same shape + if tensorA.shape != tensorB.shape: + raise ValueError("tensorA and tensorB must have the same shape.") + + # [B, C, H, W] -> (a·b) along C dimension + dot_product = (tensorA * tensorB).sum(dim=1) # => [B, H, W] + normA = tensorA.norm(dim=1) # => [B, H, W] + normB = tensorB.norm(dim=1) # => [B, H, W] + + cos_sim = dot_product / (normA * normB + eps) + if return_similarity: + return cos_sim + cos_dist = 1.0 - cos_sim + return cos_dist + + +def tokens_to_spatial(x: Tensor, aspect_ratio: float) -> Tensor: + """ + Beshape (B, area, dim) -> (B, dim, H, W) + :param x: + :param aspect_ratio: + :return: + """ + bs, area, inner_dim = x.shape + height_orig, width_orig = int(math.sqrt(area / aspect_ratio)), int( + math.sqrt(area * aspect_ratio) + ) + if aspect_ratio >= 1.0: + height = round((area / aspect_ratio) ** 0.5) + x = x.permute(0, 2, 1).reshape(bs, inner_dim, height, -1) + else: + width = round((area * aspect_ratio) ** 0.5) + x = x.permute(0, 2, 1).reshape(bs, inner_dim, -1, width) + return x + + +def spatial_to_tokens(x: Tensor) -> Tensor: + """ + Reshape (B, dim, H, W) -> (B, area, dim) + :param x: + :return: + """ + bs, inner_dim, height, width = x.shape + return x.reshape(bs, inner_dim, -1).permute(0, 2, 1) + + +def compare_kqv_resolutions( + k_hi: Tensor, + q_hi: Tensor, + v_hi: Tensor, # [B, C, L_hi] each (high resolution) + k_lo: Tensor, + q_lo: Tensor, + v_lo: Tensor, # [B, C, L_lo] each (low resolution) + aspect_ratio: float, + upsample_mode="bicubic", + attention_hi=None, # [B, n_heads, L_hi, L_hi] (optional) + attention_lo=None, # [B, n_heads, L_lo, L_lo] (optional) +): + """ + 1. Reshape hi-res and lo-res (K,Q,V) to [B, C, H, W] using the given aspect ratio. + 2. Upsample lo-res to hi-res shape. + 3. Compute and plot cosine-distance heatmaps for K, Q, V. + 4. Optionally compare attention maps by upsampling lo-res attention to hi-res + and plotting difference (or ratio). + + Parameters: + k_hi, q_hi, v_hi: [B, C, L_hi] + k_lo, q_lo, v_lo: [B, C, L_lo] + aspect_ratio: float, used to factor L into H,W + upsample_mode: 'nearest' or 'bicubic', etc. + attention_hi: optional [B, n_heads, L_hi, L_hi] + attention_lo: optional [B, n_heads, L_lo, L_lo] + """ + # 1. Reshape to spatial + k_hi_2d = tokens_to_spatial(k_hi, aspect_ratio) # [B, C, H_hi, W_hi] + q_hi_2d = tokens_to_spatial(q_hi, aspect_ratio) + v_hi_2d = tokens_to_spatial(v_hi, aspect_ratio) + + k_lo_2d = tokens_to_spatial(k_lo, aspect_ratio) # [B, C, H_lo, W_lo] + q_lo_2d = tokens_to_spatial(q_lo, aspect_ratio) + v_lo_2d = tokens_to_spatial(v_lo, aspect_ratio) + + B, C, H_hi, W_hi = k_hi_2d.shape + _, _, H_lo, W_lo = k_lo_2d.shape + + # 2. Upsample lo-res -> hi-res + scale_factor_h = H_hi / H_lo + scale_factor_w = W_hi / W_lo + + k_lo_up = F.interpolate( + k_lo_2d, size=(H_hi, W_hi), mode=upsample_mode + ) # [B, C, H_hi, W_hi] + q_lo_up = F.interpolate(q_lo_2d, size=(H_hi, W_hi), mode=upsample_mode) + v_lo_up = F.interpolate(v_lo_2d, size=(H_hi, W_hi), mode=upsample_mode) + + # 3. Compute cosine distance heatmaps for K, Q, V + k_dist = compute_cosine_distance(k_hi_2d, k_lo_up) # [B, H_hi, W_hi] + q_dist = compute_cosine_distance(q_hi_2d, q_lo_up) + v_dist = compute_cosine_distance(v_hi_2d, v_lo_up) + + # Example: Plot these heatmaps for the first item in the batch. + # Real usage might store them or handle multiple batch elements. + fig, axes = plt.subplots(1, 3, figsize=(12, 4)) + for ax, dist_map, title in zip( + axes, [k_dist, q_dist, v_dist], ["K Dist", "Q Dist", "V Dist"] + ): + ax.imshow( + dist_map[0].cpu().detach().numpy(), cmap="magma", interpolation="nearest" + ) + ax.set_title(title) + ax.axis("off") + plt.tight_layout() + plt.show() + + # 4. Optionally compare attention maps + # attention_hi: [B, n_heads, L_hi, L_hi] + # attention_lo: [B, n_heads, L_lo, L_lo] + if attention_hi is not None and attention_lo is not None: + # We'll upsample attention_lo to [B, n_heads, L_hi, L_hi] by 2D interpolation. + # But we must treat each head [L_lo, L_lo] like an image [1, 1, H_lo, W_lo]. + # So let's reshape to [B*n_heads, 1, L_lo_h, L_lo_w] if we want to do the same + # aspect-based approach. However, typically L_lo_h = L_lo_w = sqrt(L_lo), + # or you need a separate approach if the attention map isn't truly a 2D "image". + + # For simplicity, assume L_lo = H_lo * W_lo from above. + # We'll just do naive approach: reshape attention to [B*n_heads, 1, H_lo, W_lo] + # and upsample to [B*n_heads, 1, H_hi, W_hi], then flatten back to [B, n_heads, L_hi, L_hi]. + + # 4a) Reshape hi attention to [B*n_heads, 1, H_hi, W_hi] + # and lo attention to [B*n_heads, 1, H_lo, W_lo] + # We have to factor L_hi and L_lo the same way we did for K,Q,V. + # If the aspect ratio is the same, we can reuse tokens_to_spatial or do it manually. + + # For demonstration, let's do the quick version: + + B2, n_heads, L_hi2, _ = attention_hi.shape + B3, n_heads2, L_lo2, _ = attention_lo.shape + if not (B2 == B3 == B) or n_heads != n_heads2: + raise ValueError( + "Attention shapes and batch size do not match the given Q,K,V shapes." + ) + + # Factor L_hi into (H_hi, W_hi) the same way: + # But note that L_hi2 might be different from L_hi above if your n_heads differ or if aspect ratio differs. + # We'll assume we used the same approach for L_hi => H_hi, W_hi. + # Let's do it carefully: + H_hi_attn, W_hi_attn = find_hw_for_aspect_ratio(L_hi2, aspect_ratio) + if H_hi_attn * W_hi_attn != L_hi2: + raise ValueError( + "Cannot reshape attention_hi tokens to the same H_hi,W_hi with this aspect ratio." + ) + + H_lo_attn, W_lo_attn = find_hw_for_aspect_ratio(L_lo2, aspect_ratio) + if H_lo_attn * W_lo_attn != L_lo2: + raise ValueError( + "Cannot reshape attention_lo tokens to the same H_lo,W_lo with this aspect ratio." + ) + + attn_hi_2d = attention_hi.reshape(B * n_heads, L_hi2, L_hi2) + attn_lo_2d = attention_lo.reshape(B * n_heads, L_lo2, L_lo2) + + # We need a 4D shape for interpolation: [B*n_heads, 1, H_lo_attn, W_lo_attn]. + # But note that the attention is 2D in each dimension (query vs. key). + # A simpler approach is to just do a 2D interpolation for each row, but that’s more complex. + # We'll do a naive approach: treat each row as 1D and upsample. + # For a full 2D approach, we'd interpret (query index -> y-axis, key index -> x-axis) + # and reshape to [B*n_heads, 1, H_lo_attn, W_lo_attn], then upsample to [B*n_heads, 1, H_hi_attn, W_hi_attn]. + + # Quick demonstration: we can just do a normal 2D interpolation: + attn_hi_4d = attn_hi_2d.reshape(B * n_heads, 1, H_hi_attn, W_hi_attn) + attn_lo_4d = attn_lo_2d.reshape(B * n_heads, 1, H_lo_attn, W_lo_attn) + + attn_lo_up = F.interpolate( + attn_lo_4d, size=(H_hi_attn, W_hi_attn), mode=upsample_mode + ) + # Now compare attn_hi_4d vs. attn_lo_up + # For example, compute difference: + attn_diff = attn_hi_4d - attn_lo_up # [B*n_heads, 1, H_hi_attn, W_hi_attn] + + # Visualize for one example (and maybe one head): + fig2, axes2 = plt.subplots(1, 3, figsize=(15, 4)) + idx = 0 # pick the first from B*n_heads + axes2[0].imshow(attn_hi_4d[idx, 0].cpu().detach().numpy(), cmap="viridis") + axes2[0].set_title("Attention Hi") + axes2[1].imshow(attn_lo_up[idx, 0].cpu().detach().numpy(), cmap="viridis") + axes2[1].set_title("Attention Lo Up") + diff_map = attn_diff[idx, 0].cpu().detach().numpy() + im = axes2[2].imshow(diff_map, cmap="bwr") + axes2[2].set_title("Attn Difference (Hi - LoUp)") + plt.colorbar(im, ax=axes2[2], fraction=0.046, pad=0.04) + for ax in axes2: + ax.axis("off") + plt.tight_layout() + plt.show() + + return (k_dist, q_dist, v_dist) diff --git a/core/utils.py b/core/utils.py index 74217f9..01c72d4 100644 --- a/core/utils.py +++ b/core/utils.py @@ -2,7 +2,7 @@ from itertools import groupby import torch.nn.functional as F import torch -from .latent_filters import gaussian_blur_2d +from .latent_filters import gaussian_blur_2d, moment_match def pred_to_v( @@ -48,10 +48,10 @@ def partial_rescaling( ref_tensor: torch.Tensor, z: torch.Tensor, rescaled_fraction: float, -): - target_std = torch.std(ref_tensor, dim=(-3, -2, -1), keepdim=True) - std = torch.std(z, dim=(-3, -2, -1), keepdim=True) - return z * (1 - rescaled_fraction) + target_std * z / std * rescaled_fraction + match_mean: bool = False, +) -> torch.Tensor: + rescaled = moment_match(ref_tensor, z, per_channel=True, match_mean=match_mean) + return rescaled_fraction * rescaled + (1 - rescaled_fraction) * z def saliency_tensor_combination( diff --git a/nodes/generic_guidance.py b/nodes/generic_guidance.py index 0cab11f..b7f7efe 100644 --- a/nodes/generic_guidance.py +++ b/nodes/generic_guidance.py @@ -2,6 +2,8 @@ from comfy.model_patcher import ModelPatcher from comfy.samplers import calc_cond_batch from ..core import guidance as guidance from ..core import utils as utils +import torch +import numpy as np class GenericAttentionGuidance: @@ -97,6 +99,10 @@ class GenericAttentionGuidance: "round": False, }, ), + "apply_cosine_schedule_to_guidance": ( + "BOOLEAN", + {"default": False}, + ), }, "optional": { "unet_block_list": ("STRING", {"default": ""}), @@ -115,8 +121,8 @@ class GenericAttentionGuidance: def patch( self, model: ModelPatcher, - guidance_type: str = "SEG", - guidance_weight: float = 1.5, + guidance_type: str = "AAT", + guidance_weight: float = 2.0, param1: float = 0.0, param2: float = 0.0, param3: float = 0.0, @@ -127,8 +133,29 @@ class GenericAttentionGuidance: unet_block_id: int = 0, noise_fraction_start: float = 1.0, noise_fraction_end: float = 0.0, + apply_cosine_schedule_to_guidance: bool = False, unet_block_list: str = "", ): + """ + + :param model: + :param guidance_type: + :param guidance_weight: Recommeded around 2.0 + :param param1: 1.0 for most + :param param2: -100 for AAT, 0 to 1 for RandomRotation + :param param3: + :param apply_rescaling_to_alternate_guidance: + :param rescaling_method: Recommended : PredSpaceRescale + :param rescaling_fraction: Recommended : 0.7 for VSpaceRescale, 0.3-0.7 for PredSpaceRescale, + 1.0 for Softmax ( there it is the temperature) + :param unet_block: + :param unet_block_id: + :param noise_fraction_start: + :param noise_fraction_end: + :param apply_cosine_schedule_to_guidance: + :param unet_block_list: + :return: + """ if unet_block_list: blocks = utils.parse_unet_blocks(model, unet_block_list) else: @@ -150,12 +177,21 @@ class GenericAttentionGuidance: x = args["input"] sigmas = model_options["transformer_options"]["sample_sigmas"] + current_frac = sigma**2 / sigmas[0] ** 2 + current_step = torch.where(sigmas == sigma)[0].item() + total_steps = len(sigmas) + + if apply_cosine_schedule_to_guidance: + cosine_schedule = 0.5 * (1 + np.cos(np.pi * current_step / total_steps)) + guidance_w = guidance_weight * cosine_schedule + else: + guidance_w = guidance_weight if ( (current_frac > noise_fraction_start) or (current_frac <= noise_fraction_end) - or guidance_weight == 0 + or abs(guidance_w) < 1e-6 ): # Skip alternate_cond_pred = cond_pred @@ -193,7 +229,7 @@ class GenericAttentionGuidance: alternate_cond_pred, # alternate_uncond_pred, cond_pred, - guidance_weight, + guidance_w, scaling_method=rescaling_method, apply_rescaling_to_guidance=apply_rescaling_to_alternate_guidance, rescaling_fraction=rescaling_fraction, diff --git a/nodes/spec_guidance.py b/nodes/spec_guidance.py index b10beb6..b0c8959 100644 --- a/nodes/spec_guidance.py +++ b/nodes/spec_guidance.py @@ -152,10 +152,15 @@ class DownsampledLatentGuidance: # x, lf.gaussian_kernel_size_for_img(sigma_filter, x), sigma_filter # ) - ref_latent = reference_latent.to(device=cond_pred.device) + ref_latent = ( + reference_latent.to(device=cond_pred.device) + * model.latent_format.scale_factor + ) + noise = sigma.to(device=ref_latent.device) * torch.randn_like( ref_latent, device=ref_latent.device ) + # Renoise the reference latent to the current noise level x_reference = ref_latent + noise # x_downscaled = kornia.geometry.transform.rescale( @@ -178,12 +183,18 @@ class DownsampledLatentGuidance: interpolation="nearest", antialias=False, ) - attention_fn = guidance.random_rotation_wrapper(1.0, 1.0) + copy_attention = ( + guidance.upscale_and_transfer_previous_attention_wrapper( + downsample_factor + ) + ) + + perturbed_attention_fn = guidance.random_rotation_wrapper(1.0, 1.0) # attention_fn = guidance.fuzzy_attention_wrapper(1.0, 0.5) # attention_fn = guidance.affine_attention_transform_wrapper(1.0, -100, 0) model_options_perturbed = guidance.patch_attention_in_model_blocks( - model_options, attention_fn, blocks + model_options, perturbed_attention_fn, blocks ) (cond_down_peturbed,) = calc_cond_batch( @@ -212,16 +223,23 @@ class DownsampledLatentGuidance: # cond_pred + 1 * (cond_down - cond_down_peturbed), # guidance_weight, # ) + # x_reference_upscaled = kornia.geometry.transform.rescale( + # x_reference, + # float(downsample_factor), + # interpolation="nearest", + # antialias=False, + # ) final_pred = ( - uncond_pred - + cond_scale * (cond_pred - uncond_pred) + cond_pred + + (cond_scale - 1) * (cond_pred - uncond_pred) + guidance_weight * (cond_down - cond_down_peturbed) ) - moment_matched = lf.moment_match(cond_pred, final_pred) - return ( - final_pred * (1 - rescaling_fraction) - + moment_matched * rescaling_fraction - ) + return final_pred + # moment_matched = lf.moment_match(cond_pred, final_pred) + # return ( + # final_pred * (1 - rescaling_fraction) + # + moment_matched * rescaling_fraction + # ) m = model.clone() m.set_model_sampler_post_cfg_function(cfg_function)