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

Fixed and cleanup in progress. Improved SEG

Browse Source
This commit is contained in:
JTriggerFish
2025-01-31 01:06:11 +00:00
parent 0316eb3b63
commit ce8b59720c
3 changed files with 289 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/guidance.py
+186
View File
@@ -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
core/unet.py
+75 -3
View File
@@ -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]
nodes/attention.py
+28 -62
View File
@@ -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()