add FluxBlockShareKV
This commit is contained in:
wailovet
+101
-6
@@ -13,6 +13,7 @@ import torchvision.transforms.v2 as T
|
||||
|
||||
FONTS_DIR = os.path.join(os.path.dirname(os.path.realpath(__file__)), "fonts")
|
||||
|
||||
|
||||
class FluxBlockPatcherSampler:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
@@ -69,7 +70,8 @@ class FluxBlockPatcherSampler:
|
||||
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:
|
||||
@@ -84,6 +86,7 @@ class FluxBlockPatcherSampler:
|
||||
|
||||
return (out_latent, fbi_params, patched_blocks)
|
||||
|
||||
|
||||
class PlotBlockParams:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
@@ -104,7 +107,8 @@ class PlotBlockParams:
|
||||
# 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,9 +119,11 @@ 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
|
||||
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):
|
||||
@@ -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",
|
||||
}
|
||||
|
||||
+151
@@ -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
|
||||
Reference in New Issue
Block a user