From e22dd206f49ca5343974790ae8d5d6923d5b1443 Mon Sep 17 00:00:00 2001 From: ExponentialML <59846140+ExponentialML@users.noreply.github.com> Date: Thu, 14 Mar 2024 19:21:01 -0700 Subject: [PATCH] Initial Commit. --- ComfyUI_VisualStylePrompting/__init__.py | 3 + ComfyUI_VisualStylePrompting/nodes.py | 57 ++++++++++++++++++ .../utils/attention_functions.py | 45 ++++++++++++++ .../utils/style_functions.py | 60 +++++++++++++++++++ 4 files changed, 165 insertions(+) create mode 100644 ComfyUI_VisualStylePrompting/__init__.py create mode 100644 ComfyUI_VisualStylePrompting/nodes.py create mode 100644 ComfyUI_VisualStylePrompting/utils/attention_functions.py create mode 100644 ComfyUI_VisualStylePrompting/utils/style_functions.py diff --git a/ComfyUI_VisualStylePrompting/__init__.py b/ComfyUI_VisualStylePrompting/__init__.py new file mode 100644 index 0000000..3d998ea --- /dev/null +++ b/ComfyUI_VisualStylePrompting/__init__.py @@ -0,0 +1,3 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/ComfyUI_VisualStylePrompting/nodes.py b/ComfyUI_VisualStylePrompting/nodes.py new file mode 100644 index 0000000..00622ec --- /dev/null +++ b/ComfyUI_VisualStylePrompting/nodes.py @@ -0,0 +1,57 @@ +import comfy +import torch +import nodes + +from .utils.attention_functions import VisualStyleProcessor + +class ApplyVisualStyle: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "vae": ("VAE", ), + "reference_image": ("IMAGE",), + "conditioning_prompt": ("CONDITIONING",), + "reference_image_prompt": ("CONDITIONING",), + "negative_prompt": ("CONDITIONING", ), + "enabled": ("BOOLEAN", {"default": True}) + } + } + + RETURN_TYPES = ("CONDITIONING","CONDITIONING", "LATENT") + RETURN_NAMES = ("positive", "negative", "latents") + + FUNCTION = "apply_visual_style_prompt" + + def apply_visual_style_prompt( + self, + model: comfy.model_patcher.ModelPatcher, + vae, + reference_image, + conditioning_prompt, + reference_image_prompt, + negative_prompt, + enabled + ): + self.model = model + reference_latent = vae.encode(reference_image[:,:,:,:3]) + + for n, m in model.model.diffusion_model.named_modules(): + if m.__class__.__name__ == "CrossAttention": + processor = VisualStyleProcessor(m, enabled=enabled) + setattr(m, 'forward', processor.visual_style_forward) + + conditioning_prompt[0][0][1:] = reference_image_prompt[0][0][1:] + latents = torch.zeros_like(reference_latent) + latents = torch.cat([latents] * 2) + + latents[0] = reference_latent + denoise_mask = torch.ones_like(latents)[:, :1, ...] + denoise_mask[0] = 0. + + return (conditioning_prompt, negative_prompt, {"samples": latents, "noise_mask": denoise_mask}) + +NODE_CLASS_MAPPINGS = {"ApplyVisualStyle": ApplyVisualStyle} + +NODE_DISPLAY_NAME_MAPPINGS = {"ApplyVisualStyle": "Apply Visual Style Prompting"} \ No newline at end of file diff --git a/ComfyUI_VisualStylePrompting/utils/attention_functions.py b/ComfyUI_VisualStylePrompting/utils/attention_functions.py new file mode 100644 index 0000000..29be745 --- /dev/null +++ b/ComfyUI_VisualStylePrompting/utils/attention_functions.py @@ -0,0 +1,45 @@ +from comfy.ldm.modules.attention import default, optimized_attention, optimized_attention_masked +from .style_functions import adain, concat_first + +class VisualStyleProcessor(object): + def __init__(self, + module_self, + keys_scale: float = 1.0, + enabled: bool = True, + adain_queries: bool = True, + adain_keys: bool = True, + adain_values: bool = False + ): + self.module_self = module_self + self.keys_scale = keys_scale + self.enabled = enabled + self.adain_queries = adain_queries + self.adain_keys = adain_keys + self.adain_values = adain_values + + def visual_style_forward(self, x, context, value, mask=None): + q = self.module_self.to_q(x) + context = default(context, x) + k = self.module_self.to_k(context) + if value is not None: + v = self.module_self.to_v(value) + del value + else: + v = self.module_self.to_v(context) + + if self.enabled: + if self.adain_queries: + q = adain(q) + if self.adain_keys: + k = adain(k) + if self.adain_values: + v = adain(v) + + k = concat_first(k, -2, self.keys_scale) + v = concat_first(v, -2) + + if mask is None: + out = optimized_attention(q, k, v, self.module_self.heads) + else: + out = optimized_attention_masked(q, k, v, self.module_self.heads, mask) + return self.module_self.to_out(out) \ No newline at end of file diff --git a/ComfyUI_VisualStylePrompting/utils/style_functions.py b/ComfyUI_VisualStylePrompting/utils/style_functions.py new file mode 100644 index 0000000..a41321e --- /dev/null +++ b/ComfyUI_VisualStylePrompting/utils/style_functions.py @@ -0,0 +1,60 @@ +import torch + +from einops import rearrange +from dataclasses import dataclass + +T = torch.Tensor + +@dataclass(frozen=True) +class StyleAlignedArgs: + share_group_norm: bool = True + share_layer_norm: bool = True, + share_attention: bool = True + adain_queries: bool = True + adain_keys: bool = True + adain_values: bool = False + full_attention_share: bool = False + keys_scale: float = 1. + only_self_level: float = 0. + +def expand_first(feat: T, scale=1., ) -> T: + b = feat.shape[0] + feat_style = torch.stack((feat[0], feat[b // 2])).unsqueeze(1) + if scale == 1: + feat_style = feat_style.expand(2, b // 2, *feat.shape[1:]) + else: + feat_style = feat_style.repeat(1, b // 2, 1, 1, 1) + feat_style = torch.cat([feat_style[:, :1], scale * feat_style[:, 1:]], dim=1) + return feat_style.reshape(*feat.shape) + + +def concat_first(feat: T, dim=2, scale=1.) -> T: + feat_style = expand_first(feat, scale=scale) + return torch.cat((feat, feat_style), dim=dim) + + +def calc_mean_std(feat, eps: float = 1e-5) -> tuple[T, T]: + feat_std = (feat.var(dim=-2, keepdims=True) + eps).sqrt() + feat_mean = feat.mean(dim=-2, keepdims=True) + return feat_mean, feat_std + + +def adain(feat: T) -> T: + feat_mean, feat_std = calc_mean_std(feat) + feat_style_mean = expand_first(feat_mean) + feat_style_std = expand_first(feat_std) + feat = (feat - feat_mean) / feat_std + feat = feat * feat_style_std + feat_style_mean + return feat + +def swapping_attention(key, value, chunk_size=2): + chunk_length = key.size()[0] // chunk_size # [text-condition, null-condition] + reference_image_index = [0] * chunk_length # [0 0 0 0 0] + key = rearrange(key, "(b f) d c -> b f d c", f=chunk_length) + key = key[:, reference_image_index] # ref to all + key = rearrange(key, "b f d c -> (b f) d c") + value = rearrange(value, "(b f) d c -> b f d c", f=chunk_length) + value = value[:, reference_image_index] # ref to all + value = rearrange(value, "b f d c -> (b f) d c") + + return key, value \ No newline at end of file