diff --git a/blockpatcher.py b/blockpatcher.py index 77f617f..67e062b 100644 --- a/blockpatcher.py +++ b/blockpatcher.py @@ -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", } diff --git a/flux_hook.py b/flux_hook.py new file mode 100644 index 0000000..267a7b7 --- /dev/null +++ b/flux_hook.py @@ -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