nodes/JTriggerFish-ComfyLatentTools
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fixes and improvements to guidance rescaling. Attention upsampling tests stubs in place

Browse Source
This commit is contained in:
JTriggerFish
2025-02-08 23:02:06 +00:00
parent 52a88503df
commit 2ec272a662
5 changed files with 406 additions and 137 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/guidance.py
+81 -102
View File
@@ -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(
core/latent_filters.py
+252 -16
View File
@@ -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)
core/utils.py
+5 -5
View File
@@ -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(
nodes/generic_guidance.py
+40 -4
View File
@@ -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,
nodes/spec_guidance.py
+28 -10
View File
@@ -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)