266 lines
9.6 KiB
Python
266 lines
9.6 KiB
Python
from comfy_extras.nodes_custom_sampler import Noise_RandomNoise, BasicScheduler, BasicGuider, SamplerCustomAdvanced
|
|
from comfy_extras.nodes_latent import LatentBatch
|
|
from comfy_extras.nodes_model_advanced import ModelSamplingFlux, ModelSamplingAuraFlow
|
|
from node_helpers import conditioning_set_values
|
|
import comfy.samplers
|
|
import re
|
|
import os
|
|
from pathlib import Path
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import torchvision.transforms.v2 as T
|
|
# import folder_paths
|
|
|
|
FONTS_DIR = os.path.join(os.path.dirname(os.path.realpath(__file__)), "fonts")
|
|
|
|
|
|
class FluxBlockPatcherSampler:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"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}),
|
|
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.05}),
|
|
|
|
"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"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT", "SAMPLER_PARAMS", "STRING",)
|
|
RETURN_NAMES = ("latent", "sampler_params", "patched_blocks",)
|
|
FUNCTION = "apply_style"
|
|
|
|
def apply_style(self, model, conditioning, latent_image, noise_seed, steps, sampler, scheduler, guidance, denoise, blocks):
|
|
# is_schnell = model.model.model_type == comfy.model_base.ModelType.FLOW
|
|
|
|
sd = model.model_state_dict()
|
|
|
|
blocks = blocks.split("\n")
|
|
blocks = [b.strip() for b in blocks if b.strip()]
|
|
|
|
patched_blocks = []
|
|
fbi_params = []
|
|
out_latent = None
|
|
|
|
noise = Noise_RandomNoise(noise_seed)
|
|
sigmas = BasicScheduler().get_sigmas(model, scheduler, steps, denoise)[0]
|
|
cond = conditioning_set_values(conditioning, {"guidance": guidance})
|
|
sca = SamplerCustomAdvanced()
|
|
latentbatch = LatentBatch()
|
|
samplerobject = comfy.samplers.sampler_object(sampler)
|
|
|
|
for b in blocks:
|
|
b = b.split("=")
|
|
block = b[0].strip()
|
|
value = float(b[1].strip())
|
|
m = model.clone()
|
|
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]
|
|
fbi_params.append(out)
|
|
|
|
if out_latent is None:
|
|
out_latent = latent
|
|
else:
|
|
out_latent = latentbatch.batch(out_latent, latent)[0]
|
|
|
|
# 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"}),
|
|
}}
|
|
|
|
RETURN_TYPES = ("IMAGE", )
|
|
FUNCTION = "execute"
|
|
CATEGORY = "essentials/sampling"
|
|
|
|
def execute(self, images, params, cols_num, add_params):
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
import math
|
|
# import textwrap
|
|
|
|
if images.shape[0] != len(params):
|
|
raise ValueError(
|
|
"Number of images and number of parameters do not match.")
|
|
|
|
_params = params.copy()
|
|
|
|
if cols_num == 0:
|
|
cols_num = int(math.sqrt(images.shape[0]))
|
|
cols_num = max(1, min(cols_num, 1024))
|
|
|
|
width = images.shape[2]
|
|
out_image = []
|
|
|
|
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
|
|
|
|
for (image, param) in zip(images, _params):
|
|
image = image.permute(2, 0, 1)
|
|
|
|
if add_params != "false":
|
|
text = f"{param['regex']}: {param['value']}"
|
|
|
|
lines = text.split("\n")
|
|
text_height = line_height * len(lines)
|
|
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))
|
|
|
|
text_image = T.ToTensor()(text_image).to(image.device)
|
|
image = torch.cat([image, text_image], 1)
|
|
|
|
# a little cleanup
|
|
image = torch.nan_to_num(image, nan=0.0).clamp(0.0, 1.0)
|
|
out_image.append(image)
|
|
|
|
out_image = torch.stack(out_image, 0).permute(0, 2, 3, 1)
|
|
|
|
# merge images
|
|
if cols_num > -1:
|
|
cols = min(cols_num, out_image.shape[0])
|
|
b, h, w, c = out_image.shape
|
|
rows = math.ceil(b / cols)
|
|
|
|
# Pad the tensor if necessary
|
|
if b % cols != 0:
|
|
padding = cols - (b % cols)
|
|
out_image = F.pad(out_image, (0, 0, 0, 0, 0, 0, 0, padding))
|
|
b = out_image.shape[0]
|
|
|
|
# Reshape and transpose
|
|
out_image = out_image.reshape(rows, cols, h, w, c)
|
|
out_image = out_image.permute(0, 2, 1, 3, 4)
|
|
out_image = out_image.reshape(rows * h, cols * w, c).unsqueeze(0)
|
|
|
|
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
|
|
|
|
from comfy.ldm.flux.layers import DoubleStreamBlock, SingleStreamBlock
|
|
|
|
for i in range(len(double_blocks)):
|
|
if i in share_double_blocks_layers:
|
|
setattr(double_blocks[i], "forward", MethodType(
|
|
gen_double_blocks_new_forward(i), double_blocks[i]))
|
|
else:
|
|
setattr(double_blocks[i], "forward", MethodType(
|
|
DoubleStreamBlock.forward, double_blocks[i]))
|
|
|
|
for i in range(len(single_blocks)):
|
|
if i in share_single_blocks_layers:
|
|
setattr(single_blocks[i], "forward", MethodType(
|
|
gen_single_block_new_forward(i), single_blocks[i]))
|
|
else:
|
|
setattr(single_blocks[i], "forward", MethodType(
|
|
SingleStreamBlock.forward, 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",
|
|
}
|