From f4f43d12e72c7034b981193c7d3360918ca2f81a Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Fri, 22 Nov 2024 10:23:50 +0200
Subject: [PATCH] Add StyleModelApplyAdvanced

---
 __init__.py    |  1 +
 nodes/nodes.py | 22 ++++++++++++++++++++++
 2 files changed, 23 insertions(+)

diff --git a/__init__.py b/__init__.py
index 02bd738..ba6ffd8 100644
--- a/__init__.py
+++ b/__init__.py
@@ -115,6 +115,7 @@ NODE_CONFIG = {
     "ModelPassThrough": {"class": ModelPassThrough, "name": "ModelPass"},
     "ModelSaveKJ": {"class": ModelSaveKJ, "name": "Model Save KJ"},
     "SetShakkerLabsUnionControlNetType": {"class": SetShakkerLabsUnionControlNetType, "name": "Set Shakker Labs Union ControlNet Type"},
+    "StyleModelApplyAdvanced": {"class": StyleModelApplyAdvanced, "name": "Style Model Apply Advanced"},
     #audioscheduler stuff
     "NormalizedAmplitudeToMask": {"class": NormalizedAmplitudeToMask},
     "NormalizedAmplitudeToFloatList": {"class": NormalizedAmplitudeToFloatList},
diff --git a/nodes/nodes.py b/nodes/nodes.py
index 0d51792..33e0177 100644
--- a/nodes/nodes.py
+++ b/nodes/nodes.py
@@ -2482,3 +2482,25 @@ class TorchCompileControlNet:
                 raise RuntimeError("Failed to compile model")
        
         return (controlnet, )
+
+class StyleModelApplyAdvanced:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {"conditioning": ("CONDITIONING", ),
+                             "style_model": ("STYLE_MODEL", ),
+                             "clip_vision_output": ("CLIP_VISION_OUTPUT", ),
+                             "strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
+                             }}
+    RETURN_TYPES = ("CONDITIONING",)
+    FUNCTION = "apply_stylemodel"
+    CATEGORY = "KJNodes/experimental"
+    DESCRIPTION = "StyleModelApply but with strength parameter"
+
+    def apply_stylemodel(self, clip_vision_output, style_model, conditioning, strength=1.0):
+        cond = style_model.get_cond(clip_vision_output).flatten(start_dim=0, end_dim=1).unsqueeze(dim=0)
+        cond = strength * cond
+        c = []
+        for t in conditioning:
+            n = [torch.cat((t[0], cond), dim=1), t[1].copy()]
+            c.append(n)
+        return (c, )