nodes/cubiq-Block_Patcher_ComfyUI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add FluxBlockShareKV

Browse Source
This commit is contained in:
wailovet
2024-09-22 02:00:19 +08:00
parent bb92c9248a
commit f59681cf64
2 changed files with 267 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
blockpatcher.py
+116 -21
View File
@@ -9,10 +9,11 @@ from pathlib import Path
import torch
import torch.nn.functional as F
import torchvision.transforms.v2 as T
#import folder_paths
# import folder_paths
FONTS_DIR = os.path.join(os.path.dirname(os.path.realpath(__file__)), "fonts")
class FluxBlockPatcherSampler:
@classmethod
def INPUT_TYPES(s):
@@ -21,14 +22,14 @@ class FluxBlockPatcherSampler:
"model": ("MODEL", ),
"conditioning": ("CONDITIONING", ),
"latent_image": ("LATENT", ),
"noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
"steps": ("INT", {"default": 24, "min": 1, "max": 10000}),
"sampler": (comfy.samplers.KSampler.SAMPLERS, ),
"scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
"guidance": ("FLOAT", {"default": 3.5, "min": -10.0, "max": 10.0, "step": 0.1}),
"blocks": ("STRING", { "multiline": True, "dynamicPrompts": True, "default": "double_blocks\.([0-9]+)\.(img|txt)_(mod|attn|mlp\.[02])\.(lin|qkv|proj)\.(weight|bias)=1.1\nsingle_blocks\.([0-9]+)\.(linear[12]|modulation\.lin)\.(weight|bias)=1.1" }),
"blocks": ("STRING", {"multiline": True, "dynamicPrompts": True, "default": "double_blocks\.([0-9]+)\.(img|txt)_(mod|attn|mlp\.[02])\.(lin|qkv|proj)\.(weight|bias)=1.1\nsingle_blocks\.([0-9]+)\.(linear[12]|modulation\.lin)\.(weight|bias)=1.1"}),
}
}
@@ -37,7 +38,7 @@ class FluxBlockPatcherSampler:
FUNCTION = "apply_style"
def apply_style(self, model, conditioning, latent_image, noise_seed, steps, sampler, scheduler, guidance, blocks):
#is_schnell = model.model.model_type == comfy.model_base.ModelType.FLOW
# is_schnell = model.model.model_type == comfy.model_base.ModelType.FLOW
sd = model.model_state_dict()
@@ -60,16 +61,17 @@ class FluxBlockPatcherSampler:
block = b[0].strip()
value = float(b[1].strip())
m = model.clone()
out = { "regex": block, "value": value, "blocks": [] }
out = {"regex": block, "value": value, "blocks": []}
for k in sd:
if re.search(block, k):
m.add_patches({k: (None,)}, 0.0, value)
patched_blocks.append(f"{k}: {value}")
out["blocks"].append(k)
guider = BasicGuider().get_guider(m, cond)[0]
latent = sca.sample(noise, guider, samplerobject, sigmas, latent_image)[1]
latent = sca.sample(noise, guider, samplerobject,
sigmas, latent_image)[1]
fbi_params.append(out)
if out_latent is None:
@@ -77,22 +79,23 @@ class FluxBlockPatcherSampler:
else:
out_latent = latentbatch.batch(out_latent, latent)[0]
#m = None
#del m
# m = None
# del m
patched_blocks = "\n".join(patched_blocks)
return (out_latent, fbi_params, patched_blocks)
class PlotBlockParams:
@classmethod
def INPUT_TYPES(s):
return {"required": {
"images": ("IMAGE", ),
"params": ("SAMPLER_PARAMS", ),
"cols_num": ("INT", {"default": -1, "min": -1, "max": 1024 }),
"add_params": (["false", "true"], {"default": "true"}),
}}
"images": ("IMAGE", ),
"params": ("SAMPLER_PARAMS", ),
"cols_num": ("INT", {"default": -1, "min": -1, "max": 1024}),
"add_params": (["false", "true"], {"default": "true"}),
}}
RETURN_TYPES = ("IMAGE", )
FUNCTION = "execute"
@@ -101,10 +104,11 @@ class PlotBlockParams:
def execute(self, images, params, cols_num, add_params):
from PIL import Image, ImageDraw, ImageFont
import math
#import textwrap
# import textwrap
if images.shape[0] != len(params):
raise ValueError("Number of images and number of parameters do not match.")
raise ValueError(
"Number of images and number of parameters do not match.")
_params = params.copy()
@@ -115,10 +119,12 @@ class PlotBlockParams:
width = images.shape[2]
out_image = []
font = ImageFont.truetype(os.path.join(FONTS_DIR, 'ShareTechMono-Regular.ttf'), min(32, int(20*(width/1024))))
font = ImageFont.truetype(os.path.join(
FONTS_DIR, 'ShareTechMono-Regular.ttf'), min(32, int(20 * (width / 1024))))
text_padding = 3
line_height = font.getmask('Q').getbbox()[3] + font.getmetrics()[1] + text_padding*2
#char_width = font.getbbox('M')[2]+1 # using monospace font
line_height = font.getmask('Q').getbbox(
)[3] + font.getmetrics()[1] + text_padding * 2
# char_width = font.getbbox('M')[2]+1 # using monospace font
for (image, param) in zip(images, _params):
image = image.permute(2, 0, 1)
@@ -128,11 +134,13 @@ class PlotBlockParams:
lines = text.split("\n")
text_height = line_height * len(lines)
text_image = Image.new('RGB', (width, text_height), color=(0, 0, 0))
text_image = Image.new(
'RGB', (width, text_height), color=(0, 0, 0))
for i, line in enumerate(lines):
draw = ImageDraw.Draw(text_image)
draw.text((text_padding, i * line_height + text_padding), line, font=font, fill=(255, 255, 255))
draw.text((text_padding, i * line_height + text_padding),
line, font=font, fill=(255, 255, 255))
text_image = T.ToTensor()(text_image).to(image.device)
image = torch.cat([image, text_image], 1)
@@ -163,12 +171,99 @@ class PlotBlockParams:
return (out_image, )
class FluxBlockShareKV:
@classmethod
def INPUT_TYPES(s):
double_blocks_in_def = []
single_block_in_def = []
for i in range(18):
double_blocks_in_def.append(f"{i}")
for i in range(38):
single_block_in_def.append(f"{i}")
return {"required": {
"model": ("MODEL", ),
"blocks": ("STRING", {"multiline": True, "dynamicPrompts": True, "default": "double_blocks\.([0-9]+)\nsingle_blocks\.([0-9]+)"}),
}}
RETURN_TYPES = ("MODEL",)
FUNCTION = "apply_style"
def apply_style(self, model, blocks):
import importlib
from . import flux_hook
diffusion_model = model.model.diffusion_model
def gen_double_blocks_new_forward(i):
def double_blocks_new_forward(self, img: torch.Tensor, txt: torch.Tensor, vec: torch.Tensor, pe: torch.Tensor):
importlib.reload(flux_hook)
return flux_hook.double_blocks_forward(self, img, txt, vec, pe, i)
return double_blocks_new_forward
def gen_single_block_new_forward(i):
def single_block_new_forward(self, x: torch.Tensor, vec: torch.Tensor, pe: torch.Tensor):
importlib.reload(flux_hook)
return flux_hook.single_block_forward(self, x, vec, pe, i, diffusion_model)
return single_block_new_forward
double_blocks = diffusion_model.double_blocks
share_double_blocks_layers = []
for i in range(len(double_blocks)):
for b in blocks.split("\n"):
b = b.strip()
block = b
if re.match(block, f"double_blocks.{i}"):
share_double_blocks_layers.append(i)
break
single_blocks = diffusion_model.single_blocks
share_single_blocks_layers = []
for i in range(len(single_blocks)):
for b in blocks.split("\n"):
b = b.strip()
block = b
if re.match(block, f"single_blocks.{i}"):
share_single_blocks_layers.append(i)
break
from types import MethodType
for i in range(len(double_blocks)):
if hasattr(double_blocks[i], "sharekv_original_forward"):
setattr(double_blocks[i], "forward",
MethodType(double_blocks[i].sharekv_original_forward, double_blocks[i]))
for i in range(len(single_blocks)):
if hasattr(single_blocks[i], "sharekv_original_forward"):
setattr(single_blocks[i], "forward",
MethodType(single_blocks[i].sharekv_original_forward, single_blocks[i]))
for i in share_double_blocks_layers:
setattr(double_blocks[i], "sharekv_original_forward", MethodType(
double_blocks[i].forward, double_blocks[i]))
setattr(double_blocks[i], "forward", MethodType(
gen_double_blocks_new_forward(i), double_blocks[i]))
for i in share_single_blocks_layers:
setattr(single_blocks[i], "sharekv_original_forward", MethodType(
single_blocks[i].forward, single_blocks[i]))
setattr(single_blocks[i], "forward", MethodType(
gen_single_block_new_forward(i), single_blocks[i]))
return (model,)
NODE_CLASS_MAPPINGS = {
"FluxBlockPatcherSampler": FluxBlockPatcherSampler,
"FluxBlockShareKV": FluxBlockShareKV,
"PlotBlockParams": PlotBlockParams,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"FluxBlockPatcherSampler": "Flux Block Patcher Sampler",
"FluxBlockShareKV": "Flux Block Share KV",
"PlotBlockParams": "Plot Block Params",
}
flux_hook.py
+151
View File
@@ -0,0 +1,151 @@
import comfy
import torch
from einops import rearrange
from torch import Tensor
from comfy.ldm.modules.attention import optimized_attention
import comfy.model_management
def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor) -> Tensor:
q, k = apply_rope(q, k, pe)
heads = q.shape[1]
x = optimized_attention(q, k, v, heads, skip_reshape=True)
return x
def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
assert dim % 2 == 0
if comfy.model_management.is_device_mps(pos.device) or comfy.model_management.is_intel_xpu():
device = torch.device("cpu")
else:
device = pos.device
scale = torch.linspace(0, (dim - 2) / dim, steps=dim //
2, dtype=torch.float64, device=device)
omega = 1.0 / (theta**scale)
out = torch.einsum(
"...n,d->...nd", pos.to(dtype=torch.float32, device=device), omega)
out = torch.stack([torch.cos(out), -torch.sin(out),
torch.sin(out), torch.cos(out)], dim=-1)
out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)
return out.to(dtype=torch.float32, device=pos.device)
def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor):
xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
# print("freqs_cis: ", freqs_cis.shape)
freqs_cis_q = freqs_cis[:, :, :xq_.shape[2], :, :]
xq_out = freqs_cis_q[..., 0] * xq_[..., 0] + \
freqs_cis_q[..., 1] * xq_[..., 1]
xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)
def single_block_forward(self, x: Tensor, vec: Tensor, pe: Tensor, layer_id: int, ctx) -> Tensor:
txt_size = 512
bs = x.shape[0]
mod, _ = self.modulation(vec)
x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift
qkv, mlp = torch.split(self.linear1(
x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1)
q, k, v = qkv.view(qkv.shape[0], qkv.shape[1],
3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
if True:
txt_k = k[:, :, :txt_size, :]
img_k = k[:, :, txt_size:, :]
img_k = img_k.permute(0, 2, 1, 3).reshape(
1, bs * img_k.shape[2], img_k.shape[1], img_k.shape[3]).permute(0, 2, 1, 3).repeat(bs, 1, 1, 1)
k = torch.cat((txt_k, img_k), dim=2)
txt_v = v[:, :, :txt_size, :]
img_v = v[:, :, txt_size:, :]
img_v = img_v.permute(0, 2, 1, 3).reshape(
1, bs * img_v.shape[2], img_v.shape[1], img_v.shape[3]).permute(0, 2, 1, 3).repeat(bs, 1, 1, 1)
v = torch.cat((txt_v, img_v), dim=2)
txt_pe = pe[:, :, :txt_size, :, :, :]
img_pe = pe[:, :, txt_size:, :, :, :]
img_pe = img_pe.permute(0, 2, 1, 3, 4, 5).reshape(
1, bs * img_pe.shape[2], img_pe.shape[1], img_pe.shape[3], img_pe.shape[4], img_pe.shape[5]).permute(0, 2, 1, 3, 4, 5).repeat(bs, 1, 1, 1, 1, 1)
pe = torch.cat((txt_pe, img_pe), dim=2)
q, k = self.norm(q, k, v)
attn = attention(q, k, v, pe=pe)
# compute activation in mlp stream, cat again and run second linear layer
output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
x += mod.gate * output
if x.dtype == torch.float16:
x = torch.nan_to_num(x, nan=0.0, posinf=65504, neginf=-65504)
return x
def double_blocks_forward(self, img: torch.Tensor, txt: torch.Tensor, vec: torch.Tensor, pe: torch.Tensor, layer_id: int):
# print("input txt shape: ", txt.shape)
bs = img.shape[0]
img_mod1, img_mod2 = self.img_mod(vec)
txt_mod1, txt_mod2 = self.txt_mod(vec)
# prepare image for attention
img_modulated = self.img_norm1(img)
img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift
img_qkv = self.img_attn.qkv(img_modulated)
img_q, img_k, img_v = img_qkv.view(
img_qkv.shape[0], img_qkv.shape[1], 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
img_q, img_k = self.img_attn.norm(img_q, img_k, img_v)
# prepare txt for attention
txt_modulated = self.txt_norm1(txt)
txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift
txt_qkv = self.txt_attn.qkv(txt_modulated)
txt_q, txt_k, txt_v = txt_qkv.view(
txt_qkv.shape[0], txt_qkv.shape[1], 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v)
if True:
img_k = img_k.permute(0, 2, 1, 3).reshape(
1, bs * img_k.shape[2], img_k.shape[1], img_k.shape[3]).permute(0, 2, 1, 3).repeat(bs, 1, 1, 1)
img_v = img_v.permute(0, 2, 1, 3).reshape(
1, bs * img_v.shape[2], img_v.shape[1], img_v.shape[3]).permute(0, 2, 1, 3).repeat(bs, 1, 1, 1)
txt_pe = pe[:, :, :txt_k.shape[2], :, :, :]
img_pe = pe[:, :, txt_k.shape[2]:, :, :, :]
img_pe = img_pe.permute(0, 2, 1, 3, 4, 5).reshape(
1, bs * img_pe.shape[2], img_pe.shape[1], img_pe.shape[3], img_pe.shape[4], img_pe.shape[5]).permute(0, 2, 1, 3, 4, 5).repeat(bs, 1, 1, 1, 1, 1)
pe = torch.cat((txt_pe, img_pe), dim=2)
# run actual attention
attn = attention(torch.cat((txt_q, img_q), dim=2),
torch.cat((txt_k, img_k), dim=2),
torch.cat((txt_v, img_v), dim=2), pe=pe)
txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1]:]
# calculate the img bloks
img = img + img_mod1.gate * self.img_attn.proj(img_attn)
img = img + img_mod2.gate * \
self.img_mlp((1 + img_mod2.scale) *
self.img_norm2(img) + img_mod2.shift)
# calculate the txt bloks
txt += txt_mod1.gate * self.txt_attn.proj(txt_attn)
txt += txt_mod2.gate * \
self.txt_mlp((1 + txt_mod2.scale) *
self.txt_norm2(txt) + txt_mod2.shift)
if txt.dtype == torch.float16:
txt = torch.nan_to_num(txt, nan=0.0, posinf=65504, neginf=-65504)
return img, txt