Re-Commit Initial Files

__init__.py

@@ -0,0 +1,3 @@
+from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']

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"}

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)

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