From 39e634ec726b8523200566b10a13266f1af15305 Mon Sep 17 00:00:00 2001 From: ExponentialML <59846140+ExponentialML@users.noreply.github.com> Date: Thu, 14 Mar 2024 19:23:39 -0700 Subject: [PATCH] Delete ComfyUI_VisualStylePrompting directory Moving to primary directory. --- ComfyUI_VisualStylePrompting/__init__.py | 3 - ComfyUI_VisualStylePrompting/nodes.py | 57 ------------------ .../utils/attention_functions.py | 45 -------------- .../utils/style_functions.py | 60 ------------------- 4 files changed, 165 deletions(-) delete mode 100644 ComfyUI_VisualStylePrompting/__init__.py delete mode 100644 ComfyUI_VisualStylePrompting/nodes.py delete mode 100644 ComfyUI_VisualStylePrompting/utils/attention_functions.py delete mode 100644 ComfyUI_VisualStylePrompting/utils/style_functions.py diff --git a/ComfyUI_VisualStylePrompting/__init__.py b/ComfyUI_VisualStylePrompting/__init__.py deleted file mode 100644 index 3d998ea..0000000 --- a/ComfyUI_VisualStylePrompting/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -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 deleted file mode 100644 index 00622ec..0000000 --- a/ComfyUI_VisualStylePrompting/nodes.py +++ /dev/null @@ -1,57 +0,0 @@ -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 deleted file mode 100644 index 29be745..0000000 --- a/ComfyUI_VisualStylePrompting/utils/attention_functions.py +++ /dev/null @@ -1,45 +0,0 @@ -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 deleted file mode 100644 index a41321e..0000000 --- a/ComfyUI_VisualStylePrompting/utils/style_functions.py +++ /dev/null @@ -1,60 +0,0 @@ -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