initial repo

__init__.py

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

nodes.py

@@ -0,0 +1,68 @@
+import numpy as np
+import torch
+import comfy
+import folder_paths
+import nodes
+import os
+import math
+import re
+import safetensors
+import glob
+from collections import namedtuple
+
+@torch.no_grad()
+def automerge(tensor, threshold):
+    (batchsize, slices, dim) = tensor.shape
+    newTensor=[]
+    for batch in range(batchsize):
+        tokens = []
+        lastEmbed = tensor[batch,0,:]
+        merge=[lastEmbed]
+        tokens.append(lastEmbed)
+        for i in range(1,slices):
+            tok = tensor[batch,i,:]
+            cosine = torch.dot(tok,lastEmbed)/torch.sqrt(torch.dot(tok,tok)*torch.dot(lastEmbed,lastEmbed))
+            if cosine >= threshold:
+                merge.append(tok)
+                lastEmbed = torch.stack(merge).mean(dim=0)
+            else:
+                tokens.append(lastEmbed)
+                merge=[]
+                lastEmbed=tok
+        newTensor.append(torch.stack(tokens))
+    return torch.stack(newTensor)
+
+class StyleModelApplyAdvanced:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {"conditioning": ("CONDITIONING", ),
+                             "style_model": ("STYLE_MODEL", ),
+                             "clip_vision_output": ("CLIP_VISION_OUTPUT", ),
+                             "merge_strength": ("FLOAT", {"default": 0.8, "min": 0.0, "max": 1.0, "step": 0.01}),
+                             "clipWeight": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01})
+                             }}
+    RETURN_TYPES = ("CONDITIONING",)
+    FUNCTION = "apply_stylemodel"
+
+    CATEGORY = "conditioning/style_model"
+
+    def apply_stylemodel(self, clip_vision_output, style_model, conditioning, merge_strength, clipWeight):
+        cond = style_model.get_cond(clip_vision_output).flatten(start_dim=0, end_dim=1).unsqueeze(dim=0)
+        cond = automerge(cond, merge_strength)*clipWeight
+        c = []
+        for t in conditioning:
+            n = [torch.cat((t[0], cond), dim=1), t[1].copy()]
+            c.append(n)
+        return (c, )
+
+
+# A dictionary that contains all nodes you want to export with their names
+# NOTE: names should be globally unique
+NODE_CLASS_MAPPINGS = {
+    "StyleModelApplyAdvanced": StyleModelApplyAdvanced,
+}
+
+# A dictionary that contains the friendly/humanly readable titles for the nodes
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "StyleModelApplyAdvanced": "Apply Style model (advanced)"
+}