init

.gitignore

@@ -0,0 +1,3 @@
+**/__pycache__
+**/*.pyc
+.DS_Store

__init__.py

@@ -0,0 +1,19 @@
+from .nodes.apply_ssr_attention_node import ApplySSRNode
+from .nodes.encode_ssr_query_node import EncodeSSRQueryNode
+from .nodes.load_ssr_aligner_node import LoadSSRAlignerNode
+from .nodes.load_ssr_attention_node import LoadSSRAttentionNode
+
+
+NODE_CLASS_MAPPINGS = {
+    'ApplySSR': ApplySSRNode,
+    'EncodeSSRQuery': EncodeSSRQueryNode,
+    'LoadSSRAligner': LoadSSRAlignerNode,
+    'LoadSSRAttention': LoadSSRAttentionNode,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    'ApplySSR': 'SSR Apply Model',
+    'EncodeSSRQuery': 'SSR Encode Query',
+    'LoadSSRAligner': 'SSR Load Aligner',
+    'LoadSSRAttention': 'SSR Load Attentions',
+}

modules/clip_vision_override.py

@@ -0,0 +1,26 @@
+from comfy.clip_vision import clip_preprocess
+from comfy.ldm.modules.attention import optimized_attention_for_device
+import comfy.model_management
+
+
+def encoder(encoder, x, mask=None):
+    optimized_attention = optimized_attention_for_device(x.device, mask=mask is not None, small_input=True)
+    intermediates = []
+    for i, l in enumerate(encoder.layers):
+        x = l(x, mask, optimized_attention)
+        intermediates.append(x.clone())
+    return intermediates
+
+
+def vision_model(vision_model, pixel_values):
+    x = vision_model.embeddings(pixel_values)
+    x = vision_model.pre_layrnorm(x)
+    intermediates = encoder(vision_model.encoder, x)
+    return intermediates
+
+
+def encode_image(clip_vision, image):
+    comfy.model_management.load_model_gpu(clip_vision.patcher)
+    pixel_values = clip_preprocess(image.to(clip_vision.load_device), size=clip_vision.image_size).float()
+    intermediates = vision_model(clip_vision.model.vision_model, pixel_values)
+    return intermediates

modules/ssr_modules.py

@@ -0,0 +1,89 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+import comfy.ops
+ops = comfy.ops.disable_weight_init
+from comfy.ldm.modules.attention import optimized_attention
+
+
+class SSRAligner(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        query_dim = 768
+        inner_dim = 512
+        cross_dim = 1024
+        self.heads = 8
+
+        self.to_q = ops.Linear(query_dim, inner_dim, bias=False)
+        self.to_k = ops.Linear(cross_dim, inner_dim, bias=False)
+        
+        self.to_v = nn.ModuleList([])
+        self.to_out = nn.ModuleList([])
+        for _ in range(6):
+            self.to_v.append(ops.Linear(cross_dim, inner_dim, bias=False))
+            self.to_out.append(ops.Linear(inner_dim, query_dim, bias=False))
+        self.norm = ops.LayerNorm(query_dim)
+
+    def forward(self, x, context_list):
+        q = self.to_q(x)
+        k = self.to_k(context_list[-1])
+
+        outs = []
+        for i in range(6):
+            v = self.to_v[i](context_list[i])
+            attn_out = optimized_attention(q, k, v, self.heads)
+            outs.append(self.to_out[i](attn_out))
+
+        return self.norm(torch.cat(outs, dim=1)) # [1, 77*6, 768]
+    
+
+# me being lazy
+ssr_caa_dims = {
+    'input': [320, 320, 640, 640, 1280, 1280],
+    'middle': [1280],
+    'output': [1280, 1280, 1280, 640, 640, 640, 320, 320, 320]
+}
+
+
+class SSRCrossAttention(nn.Module):
+    def __init__(self, hidden_size) -> None:
+        super().__init__()
+        self.heads = 8
+
+        self.to_k = ops.Linear(768, hidden_size, bias=False)
+        self.to_v = ops.Linear(768, hidden_size, bias=False)
+
+        self.pos_contexts = [] # [{ 'scale': int, 'embed': tensor }, ...]
+        self.neg_contexts = []
+
+    def to(self, device):
+        super().to(device)
+        for i in range(len(self.pos_contexts)):
+            self.pos_contexts[i]['embed'] = self.pos_contexts[i]['embed'].to(device)
+
+        for i in range(len(self.neg_contexts)):
+            self.neg_contexts[i]['embed'] = self.neg_contexts[i]['embed'].to(device)
+
+        return self
+
+    def forward(self, q, k, v, extra_options):
+
+        hidden_states = optimized_attention(q, k, v, extra_options['n_heads'])
+        cond_or_uncond = extra_options['cond_or_uncond']
+        len_conds = len(cond_or_uncond)
+        batch_size = len(q)//len_conds
+
+        for cond_idx, cond in enumerate(cond_or_uncond):
+            contexts = self.pos_contexts if cond == 0 else self.neg_contexts
+            q_cond = q[cond_idx*batch_size:(cond_idx+1)*batch_size]
+            for context in contexts:
+                scale = context['scale']
+                embed = context['embed']
+                ssr_k = self.to_k(embed).repeat(batch_size, 1, 1)
+                ssr_v = self.to_v(embed).repeat(batch_size, 1, 1)
+                ssr_hidden_states = optimized_attention(q_cond, ssr_k, ssr_v, self.heads)
+
+                hidden_states[cond_idx*batch_size:(cond_idx+1)*batch_size] += scale * ssr_hidden_states
+
+        return hidden_states

nodes/apply_ssr_attention_node.py

@@ -0,0 +1,27 @@
+
+class ApplySSRNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "ssr_attentions": ("SSR_ATTENTIONS",), 
+                "ssr_embeds": ("SSR_EMBEDS",),
+            }
+        }
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "apply"
+
+    CATEGORY = "ssr"
+
+    def apply(self, model, ssr_attentions, ssr_embeds):
+        model = model.clone()
+
+        for key, ssr_attention in ssr_attentions.items():
+            ssr_attention.pos_contexts = ssr_embeds['positive']
+            ssr_attention.neg_contexts = ssr_embeds['negative']
+            block_name, block_number = key
+            model.set_model_attn2_replace(ssr_attention, block_name, block_number)
+
+        return (model,)
+

nodes/encode_ssr_query_node.py

@@ -0,0 +1,70 @@
+import comfy.model_management
+
+from ..modules.clip_vision_override import encode_image
+
+
+class EncodeSSRQueryNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "positive_scale": ("FLOAT", {"default": 0.4, "min": 0.0, "max": 10.0, "step": 0.01}),
+                "positive": ("STRING", {"multiline": True, "dynamicPrompts": True, "tooltip": "The positive image query"}), 
+                "negative_scale": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 10.0, "step": 0.01}),
+                "negative": ("STRING", {"multiline": True, "dynamicPrompts": True, "tooltip": "The negative image query"}), 
+                "clip": ("CLIP", {"tooltip": "The CLIP model used for encoding the text."}),
+                "clip_vision": ("CLIP_VISION",),
+                "ssr_aligner": ("SSR_ALIGNER",),
+                "image": ("IMAGE",),
+            },
+            "optional": {
+                "prev_ssr_embeds": ("SSR_EMBEDS",)
+            }
+        }
+    RETURN_TYPES = ("SSR_EMBEDS",)
+    FUNCTION = "encode"
+    CATEGORY = "ssr"
+
+    def _encode_text(self, text, clip):
+        tokens = clip.tokenize(text)
+        output = clip.encode_from_tokens(tokens, return_pooled=True, return_dict=True)
+        cond = output.pop("cond")
+        return cond
+
+    def encode(self, 
+               positive_scale,
+               positive, 
+               negative_scale, 
+               negative, 
+               clip,
+               clip_vision, 
+               ssr_aligner, 
+               image, 
+               prev_ssr_embeds=None):
+        if prev_ssr_embeds is None:
+            prev_ssr_embeds = { 'negative': [], 'positive': [] }
+        else:
+            prev_ssr_embeds = { 
+                'negative': [*prev_ssr_embeds['negative']], 
+                'positive': [*prev_ssr_embeds['positive']] 
+            }
+
+
+        intermediate_device = comfy.model_management.intermediate_device()
+        torch_device = comfy.model_management.get_torch_device()
+        image_embeds = encode_image(clip_vision, image)[3::4]
+
+        ssr_aligner.to(torch_device)
+        pos_embed = None
+        if positive and positive_scale > 0:
+            pos_embed = self._encode_text(positive, clip).to(torch_device)
+            pos_embed = ssr_aligner(pos_embed, image_embeds).to(intermediate_device)
+            prev_ssr_embeds['positive'].append({ 'scale': positive_scale, 'embed': pos_embed })
+        if negative and negative_scale > 0:
+            neg_embed = self._encode_text(negative, clip).to(torch_device)
+            neg_embed = ssr_aligner(neg_embed, image_embeds).to(intermediate_device)
+            prev_ssr_embeds['negative'].append({ 'scale': negative_scale, 'embed': neg_embed })
+        
+        ssr_aligner.to(intermediate_device)
+
+        return (prev_ssr_embeds, )

nodes/load_ssr_aligner_node.py

@@ -0,0 +1,30 @@
+import os
+
+from safetensors.torch import load_file
+
+import folder_paths
+from folder_paths import models_dir
+
+from ..modules.ssr_modules import SSRAligner
+
+
+SSR_PATH = os.path.join(models_dir, 'ssr')
+os.makedirs(SSR_PATH, exist_ok=True)
+
+
+class LoadSSRAlignerNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": { "checkpoint": (os.listdir(SSR_PATH), )}}
+
+    RETURN_TYPES = ("SSR_ALIGNER",)
+    FUNCTION = "load"
+
+    CATEGORY = "ssr"
+
+    def load(self, checkpoint):
+        checkpoint_path = os.path.join(SSR_PATH, checkpoint)
+        state_dict = load_file(checkpoint_path)
+        ssr_aligner = SSRAligner()
+        ssr_aligner.load_state_dict(state_dict)
+        return (ssr_aligner,)

nodes/load_ssr_attention_node.py

@@ -0,0 +1,57 @@
+import os
+
+from safetensors.torch import load_file
+
+from folder_paths import models_dir
+
+from ..modules.ssr_modules import SSRCrossAttention, ssr_caa_dims
+
+
+SSR_PATH = os.path.join(models_dir, 'ssr')
+os.makedirs(SSR_PATH, exist_ok=True)
+
+
+class LoadSSRAttentionNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": { "checkpoint": (os.listdir(SSR_PATH), )}}
+
+    RETURN_TYPES = ("SSR_ATTENTIONS",)
+    FUNCTION = "load"
+
+    CATEGORY = "ssr"
+
+    def load(self, checkpoint):
+        checkpoint_path = os.path.join(SSR_PATH, checkpoint)
+        state_dict = load_file(checkpoint_path)
+        ssr_idx = 1
+        ssr_attentions = {}
+        for idx, i in enumerate([1,2,3,4,5,6]):
+            ssr_aligner = SSRCrossAttention(ssr_caa_dims['input'][idx])
+            attn_dict = {
+                'to_k.weight': state_dict[f'{ssr_idx}.to_k_SSR.weight'],
+                'to_v.weight': state_dict[f'{ssr_idx}.to_v_SSR.weight'],
+            }
+            ssr_aligner.load_state_dict(attn_dict)
+            ssr_attentions[('input', i)] = ssr_aligner
+            ssr_idx += 2
+
+        for idx, i in enumerate([3,4,5,6,7,8,9,10,11]):
+            ssr_aligner = SSRCrossAttention(ssr_caa_dims['output'][idx])
+            attn_dict = {
+                'to_k.weight': state_dict[f'{ssr_idx}.to_k_SSR.weight'],
+                'to_v.weight': state_dict[f'{ssr_idx}.to_v_SSR.weight'],
+            }
+            ssr_aligner.load_state_dict(attn_dict)
+            ssr_attentions[('output', i)] = ssr_aligner
+            ssr_idx += 2
+            
+        ssr_aligner = SSRCrossAttention(ssr_caa_dims['middle'][0])
+        attn_dict = {
+            'to_k.weight': state_dict[f'{ssr_idx}.to_k_SSR.weight'],
+            'to_v.weight': state_dict[f'{ssr_idx}.to_v_SSR.weight'],
+        }
+        ssr_aligner.load_state_dict(attn_dict)
+        ssr_attentions[('middle', 0)] = ssr_aligner
+
+        return (ssr_attentions,)