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

SEG initial code

Browse Source
This commit is contained in:
JTriggerFish
2025-01-29 15:11:40 +00:00
parent 8325707aa1
commit 22d0d82e4a
3 changed files with 232 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/latent_filters.py
+25
View File
@@ -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
core/diffusion.py → core/unet.py
+36 -3
View File
@@ -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
nodes/attention.py
+171 -13
View File
@@ -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",
}