diff --git a/comfy/ldm/flux/math.py b/comfy/ldm/flux/math.py index 084dbec0f..b9e21f9cd 100644 --- a/comfy/ldm/flux/math.py +++ b/comfy/ldm/flux/math.py @@ -1,6 +1,8 @@ import torch from einops import rearrange from torch import Tensor +from torch.nn.functional import interpolate + from comfy.ldm.modules.attention import optimized_attention import comfy.model_management @@ -33,3 +35,4 @@ def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor): xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 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) + diff --git a/comfy/ldm/flux/model.py b/comfy/ldm/flux/model.py index 97ad8ffea..02be06845 100644 --- a/comfy/ldm/flux/model.py +++ b/comfy/ldm/flux/model.py @@ -4,6 +4,8 @@ from dataclasses import dataclass import torch from torch import Tensor, nn +from einops import rearrange, repeat +import comfy.ldm.common_dit from .layers import ( DoubleStreamBlock, @@ -14,9 +16,6 @@ from .layers import ( timestep_embedding, ) -from einops import rearrange, repeat -import comfy.ldm.common_dit - @dataclass class FluxParams: in_channels: int @@ -98,8 +97,9 @@ class Flux(nn.Module): timesteps: Tensor, y: Tensor, guidance: Tensor = None, - control=None, + control = None, transformer_options={}, + attn_mask: Tensor = None, ) -> Tensor: patches_replace = transformer_options.get("patches_replace", {}) if img.ndim != 3 or txt.ndim != 3: @@ -124,14 +124,27 @@ class Flux(nn.Module): if ("double_block", i) in blocks_replace: def block_wrap(args): out = {} - out["img"], out["txt"] = block(img=args["img"], txt=args["txt"], vec=args["vec"], pe=args["pe"]) + out["img"], out["txt"] = block(img=args["img"], + txt=args["txt"], + vec=args["vec"], + pe=args["pe"], + attn_mask=args.get("attn_mask")) return out - out = blocks_replace[("double_block", i)]({"img": img, "txt": txt, "vec": vec, "pe": pe}, {"original_block": block_wrap}) + out = blocks_replace[("double_block", i)]({"img": img, + "txt": txt, + "vec": vec, + "pe": pe, + "attn_mask": attn_mask}, + {"original_block": block_wrap}) txt = out["txt"] img = out["img"] else: - img, txt = block(img=img, txt=txt, vec=vec, pe=pe) + img, txt = block(img=img, + txt=txt, + vec=vec, + pe=pe, + attn_mask=attn_mask) if control is not None: # Controlnet control_i = control.get("input") @@ -146,13 +159,20 @@ class Flux(nn.Module): if ("single_block", i) in blocks_replace: def block_wrap(args): out = {} - out["img"] = block(args["img"], vec=args["vec"], pe=args["pe"]) + out["img"] = block(args["img"], + vec=args["vec"], + pe=args["pe"], + attn_mask=args.get("attn_mask")) return out - out = blocks_replace[("single_block", i)]({"img": img, "vec": vec, "pe": pe}, {"original_block": block_wrap}) + out = blocks_replace[("single_block", i)]({"img": img, + "vec": vec, + "pe": pe, + "attn_mask": attn_mask}, + {"original_block": block_wrap}) img = out["img"] else: - img = block(img, vec=vec, pe=pe) + img = block(img, vec=vec, pe=pe, attn_mask=attn_mask) if control is not None: # Controlnet control_o = control.get("output") @@ -181,5 +201,5 @@ class Flux(nn.Module): img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs) txt_ids = torch.zeros((bs, context.shape[1], 3), device=x.device, dtype=x.dtype) - out = self.forward_orig(img, img_ids, context, txt_ids, timestep, y, guidance, control, transformer_options) + out = self.forward_orig(img, img_ids, context, txt_ids, timestep, y, guidance, control, transformer_options, attn_mask=kwargs.get("attention_mask", None)) return rearrange(out, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=h_len, w=w_len, ph=2, pw=2)[:,:,:h,:w] diff --git a/comfy/model_base.py b/comfy/model_base.py index 8f37af660..0fadef1f0 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -687,6 +687,7 @@ class StableAudio1(BaseModel): sd["{}{}".format(k, l)] = s[l] return sd + class HunyuanDiT(BaseModel): def __init__(self, model_config, model_type=ModelType.V_PREDICTION, device=None): super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.hydit.models.HunYuanDiT) @@ -769,6 +770,16 @@ class Flux(BaseModel): cross_attn = kwargs.get("cross_attn", None) if cross_attn is not None: out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn) + # upscale the attention mask, since now we + attention_mask = kwargs.get("attention_mask", None) + if attention_mask is not None: + shape = kwargs["noise"].shape + mask_ref_size = kwargs["attention_mask_img_shape"] + # the model will pad to the patch size, and then divide + # essentially dividing and rounding up + (h_tok, w_tok) = (math.ceil(shape[2] / self.diffusion_model.patch_size), math.ceil(shape[3] / self.diffusion_model.patch_size)) + attention_mask = utils.upscale_dit_mask(attention_mask, mask_ref_size, (h_tok, w_tok)) + out['attention_mask'] = comfy.conds.CONDRegular(attention_mask) out['guidance'] = comfy.conds.CONDRegular(torch.FloatTensor([kwargs.get("guidance", 3.5)])) return out diff --git a/comfy/utils.py b/comfy/utils.py index 985cd9a1b..f9fb4e7f3 100644 --- a/comfy/utils.py +++ b/comfy/utils.py @@ -26,6 +26,8 @@ import numpy as np from PIL import Image import logging import itertools +from torch.nn.functional import interpolate +from einops import rearrange def load_torch_file(ckpt, safe_load=False, device=None): if device is None: @@ -867,5 +869,46 @@ def reshape_mask(input_mask, output_shape): mask = torch.nn.functional.interpolate(input_mask, size=output_shape[2:], mode=scale_mode) if mask.shape[1] < output_shape[1]: mask = mask.repeat((1, output_shape[1]) + (1,) * dims)[:,:output_shape[1]] - mask = comfy.utils.repeat_to_batch_size(mask, output_shape[0]) + mask = repeat_to_batch_size(mask, output_shape[0]) return mask + +def upscale_dit_mask(mask: torch.Tensor, img_size_in, img_size_out): + hi, wi = img_size_in + ho, wo = img_size_out + # if it's already the correct size, no need to do anything + if (hi, wi) == (ho, wo): + return mask + if mask.ndim == 2: + mask = mask.unsqueeze(0) + if mask.ndim != 3: + raise ValueError(f"Got a mask of shape {list(mask.shape)}, expected [b, q, k] or [q, k]") + txt_tokens = mask.shape[1] - (hi * wi) + # quadrants of the mask + txt_to_txt = mask[:, :txt_tokens, :txt_tokens] + txt_to_img = mask[:, :txt_tokens, txt_tokens:] + img_to_img = mask[:, txt_tokens:, txt_tokens:] + img_to_txt = mask[:, txt_tokens:, :txt_tokens] + + # convert to 1d x 2d, interpolate, then back to 1d x 1d + txt_to_img = rearrange (txt_to_img, "b t (h w) -> b t h w", h=hi, w=wi) + txt_to_img = interpolate(txt_to_img, size=img_size_out, mode="bilinear") + txt_to_img = rearrange (txt_to_img, "b t h w -> b t (h w)") + # this one is hard because we have to do it twice + # convert to 1d x 2d, interpolate, then to 2d x 1d, interpolate, then 1d x 1d + img_to_img = rearrange (img_to_img, "b hw (h w) -> b hw h w", h=hi, w=wi) + img_to_img = interpolate(img_to_img, size=img_size_out, mode="bilinear") + img_to_img = rearrange (img_to_img, "b (hk wk) hq wq -> b (hq wq) hk wk", hk=hi, wk=wi) + img_to_img = interpolate(img_to_img, size=img_size_out, mode="bilinear") + img_to_img = rearrange (img_to_img, "b (hq wq) hk wk -> b (hk wk) (hq wq)", hq=ho, wq=wo) + # convert to 2d x 1d, interpolate, then back to 1d x 1d + img_to_txt = rearrange (img_to_txt, "b (h w) t -> b t h w", h=hi, w=wi) + img_to_txt = interpolate(img_to_txt, size=img_size_out, mode="bilinear") + img_to_txt = rearrange (img_to_txt, "b t h w -> b (h w) t") + + # reassemble the mask from blocks + out = torch.cat([ + torch.cat([txt_to_txt, txt_to_img], dim=2), + torch.cat([img_to_txt, img_to_img], dim=2)], + dim=1 + ) + return out diff --git a/comfy_extras/nodes_flux.py b/comfy_extras/nodes_flux.py index 35d932512..623ddefd2 100644 --- a/comfy_extras/nodes_flux.py +++ b/comfy_extras/nodes_flux.py @@ -41,99 +41,8 @@ class FluxGuidance: c = node_helpers.conditioning_set_values(conditioning, {"guidance": guidance}) return (c, ) -class _ReduxAttnWrapper: - def __init__(self, previous, token_counts, bias=0.0, is_first=False): - self.previous = previous - self.token_counts = token_counts - self.bias = bias - self.is_first = is_first - - def __call__(self, args, extra_args): - # args: {"img": img, <"txt": txt>, "vec": vec, "pe": pe} - if self.is_first: - self.token_counts["img"] = args["img"].shape[1] - - # determine the total number of tokens in the mask, depending on whether we're wrapping a single block or a double one - total_tokens = args["img"].shape[1] - if "txt" in args: - total_tokens += args["txt"].shape[1] - # create the mask (or bias map) - mask = extra_args.get("attn_mask", torch.zeros((total_tokens, total_tokens), device=args["img"].device, dtype=args["img"].dtype)) - # if this wrapper was called by another ReduxAttnWrapper, compute the range of tokens that correspond to our image - redux_end = extra_args.get("redux_end", -self.token_counts["img"]) - redux_start = redux_end - self.token_counts["redux"] - # modify the mask - # first 256 tokens are the text prompt - mask[:256, redux_start:redux_end] = self.bias - # last 'img' tokens are the image being generated - mask[-self.token_counts["img"]:, redux_start:redux_end] = self.bias - # nice case for a match statement - if isinstance(self.previous, DoubleStreamBlock): - x, c = self.previous(img=args["img"], txt=args["txt"],vec=args["vec"], pe=args["pe"], attn_mask=mask) - return {"img": x, "txt": c} - elif isinstance(self.previous, SingleStreamBlock): - x = self.previous(args["img"], vec=args["vec"], pe=args["pe"], attn_mask=mask) - return {"img": x} - elif isinstance(self.previous, _ReduxAttnWrapper): - # pass along the mask, and tell the next redux what its part of the mask is - extra_args["attn_mask"] = mask - extra_args["redux_end"] = redux_start - return self.previous(args, extra_args) - else: - print(f"Can't wrap {repr(self.previous)} with mask.") - return self.previous(args, extra_args) - -class ReduxApplyWithAttnMask: - @classmethod - def INPUT_TYPES(s): - return {"required": { - "model": ("MODEL", ), - "conditioning": ("CONDITIONING", ), - "style_model": ("STYLE_MODEL", ), - "clip_vision_output": ("CLIP_VISION_OUTPUT", ), - "attn_bias": ("FLOAT", {"default": 0.0, "min": -10.0, "max": 10.0, "step": 0.01}), - }} - RETURN_TYPES = ("MODEL", "CONDITIONING") - FUNCTION = "apply_stylemodel" - - CATEGORY = "conditioning/style_model" - - def apply_stylemodel(self, model: ModelPatcher, clip_vision_output, style_model, conditioning, attn_bias): - cond = style_model.get_cond(clip_vision_output).flatten(start_dim=0, end_dim=1).unsqueeze(dim=0) - - c = [] - for t in conditioning: - n = [torch.cat((t[0], cond), dim=1), t[1].copy()] - c.append(n) - - if attn_bias != 0.0: - token_counts = { - "redux": cond.shape[1], - "img": None - } - - m = model.clone() - # patch the model - previous_patches = m.model_options["transformer_options"].get("patches_replace", {}).get("dit", {}) - - for i, block in enumerate(m.model.diffusion_model.double_blocks): - # is there already a patch there? - # if so, the attnwrapper can chain off it - previous = previous_patches.get(("double_block", i), block) - wrapper = _ReduxAttnWrapper(previous, token_counts, bias=attn_bias, is_first=i==0) - # I think this properly clones things? - m.set_model_patch_replace(wrapper, "dit", "double_block", i) - - for i, block in enumerate(m.model.diffusion_model.single_blocks): - previous = previous_patches.get(("single_block", i), block) - wrapper = _ReduxAttnWrapper(previous, token_counts, bias=attn_bias) - m.set_model_patch_replace(wrapper, "dit", "single_block", i) - else: - m = model - return (m, c) NODE_CLASS_MAPPINGS = { "CLIPTextEncodeFlux": CLIPTextEncodeFlux, "FluxGuidance": FluxGuidance, - "ReduxWithAttnMask": ReduxApplyWithAttnMask } diff --git a/nodes.py b/nodes.py index 1cb4b5a5a..a56e4bd66 100644 --- a/nodes.py +++ b/nodes.py @@ -1010,23 +1010,58 @@ class StyleModelApply: "style_model": ("STYLE_MODEL", ), "clip_vision_output": ("CLIP_VISION_OUTPUT", ), "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}), - "strength_type": (["multiply"], ), + "strength_type": (["multiply", "attn_bias"], ), }} RETURN_TYPES = ("CONDITIONING",) FUNCTION = "apply_stylemodel" CATEGORY = "conditioning/style_model" - def apply_stylemodel(self, clip_vision_output, style_model, conditioning, strength, strength_type): + def apply_stylemodel(self, conditioning, style_model, clip_vision_output, strength, strength_type): cond = style_model.get_cond(clip_vision_output).flatten(start_dim=0, end_dim=1).unsqueeze(dim=0) if strength_type == "multiply": cond *= strength - c = [] + n = cond.shape[1] + c_out = [] for t in conditioning: - n = [torch.cat((t[0], cond), dim=1), t[1].copy()] - c.append(n) - return (c, ) + (txt, keys) = t + keys = keys.copy() + if strength_type == "attn_bias" and strength != 1.0: + # math.log raises an error if the argument is zero + # torch.log returns -inf, which is what we want + attn_bias = torch.log(torch.Tensor([strength])) + # get the size of the mask image + mask_ref_size = keys.get("attention_mask_img_shape", (1, 1)) + n_ref = mask_ref_size[0] * mask_ref_size[1] + n_txt = txt.shape[1] + # grab the existing mask + mask = keys.get("attention_mask", None) + # create a default mask if it doesn't exist + if mask is None: + mask = torch.zeros((txt.shape[0], n_txt + n_ref, n_txt + n_ref), dtype=torch.float16) + # convert the mask dtype, because it might be boolean + # we want it to be interpreted as a bias + if mask.dtype == torch.bool: + # log(True) = log(1) = 0 + # log(False) = log(0) = -inf + mask = torch.log(mask.to(dtype=torch.float16)) + # now we make the mask bigger to add space for our new tokens + new_mask = torch.zeros((txt.shape[0], n_txt + n + n_ref, n_txt + n + n_ref), dtype=torch.float16) + # copy over the old mask, in quandrants + new_mask[:, :n_txt, :n_txt] = mask[:, :n_txt, :n_txt] + new_mask[:, :n_txt, n_txt+n:] = mask[:, :n_txt, n_txt:] + new_mask[:, n_txt+n:, :n_txt] = mask[:, n_txt:, :n_txt] + new_mask[:, n_txt+n:, n_txt+n:] = mask[:, n_txt:, n_txt:] + # now fill in the attention bias to our redux tokens + new_mask[:, :n_txt, n_txt:n_txt+n] = attn_bias + new_mask[:, n_txt+n:, n_txt:n_txt+n] = attn_bias + keys["attention_mask"] = new_mask.to(txt.device) + keys["attention_mask_img_shape"] = mask_ref_size + + c_out.append([torch.cat((txt, cond), dim=1), keys]) + + return (c_out,) class unCLIPConditioning: @classmethod