From be421dba5807faca417817b86b83c7eaf804c197 Mon Sep 17 00:00:00 2001
From: kaidu <git@k-du.de>
Date: Fri, 22 Nov 2024 22:39:51 +0100
Subject: [PATCH] add arbitrary downsampling factors and methods

---
 nodes.py | 32 +++++++++++++++++++++++++++++++-
 1 file changed, 31 insertions(+), 1 deletion(-)

diff --git a/nodes.py b/nodes.py
index 11df9f7..2fcc4ba 100644
--- a/nodes.py
+++ b/nodes.py
@@ -67,6 +67,34 @@ class StyleModelApplySimple:
         return (c, )
 
 
+class StyleModelApplyInterpolation:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {"conditioning": ("CONDITIONING", ),
+                             "style_model": ("STYLE_MODEL", ),
+                             "clip_vision_output": ("CLIP_VISION_OUTPUT", ),
+                             "downsampling_factor": ("INT", {"default": 3, "min": 1, "max": 9}),
+                             "mode": (["nearest", "bicubic", "bilinear","area", "nearest-exact" ], {"default": "area"})
+                             }}
+    RETURN_TYPES = ("CONDITIONING",)
+    FUNCTION = "apply_stylemodel"
+
+    CATEGORY = "conditioning/style_model"
+
+    def apply_stylemodel(self, clip_vision_output, style_model, conditioning, downsampling_factor, mode):
+        cond = style_model.get_cond(clip_vision_output).flatten(start_dim=0, end_dim=1).unsqueeze(dim=0)
+        if downsampling_factor>1:
+            (b,t,h)=cond.shape
+            m = int(np.sqrt(t))
+            cond=torch.nn.functional.interpolate(cond.view(b, m, m, h).transpose(1,-1), size=(m//downsampling_factor, m//downsampling_factor), mode=mode)#
+            cond=cond.transpose(1,-1).reshape(b,-1,h)
+        c = []
+        for t in conditioning:
+            n = [torch.cat((t[0], cond), dim=1), t[1].copy()]
+            c.append(n)
+        return (c, )
+
+
 class StyleModelApplyAdvanced:
     @classmethod
     def INPUT_TYPES(s):
@@ -106,10 +134,12 @@ class StyleModelApplyAdvanced:
 NODE_CLASS_MAPPINGS = {
     "StyleModelApplySimple": StyleModelApplySimple,
     "StyleModelApplyAdvanced": StyleModelApplyAdvanced,
+    "StyleModelApplyInterpolation": StyleModelApplyInterpolation
 }
 
 # A dictionary that contains the friendly/humanly readable titles for the nodes
 NODE_DISPLAY_NAME_MAPPINGS = {
     "StyleModelApplySimple": "Apply Style model (simple)",
-    "StyleModelApplyAdvanced": "Apply Style model (advanced)"
+    "StyleModelApplyAdvanced": "Apply Style model (advanced)",
+    "StyleModelApplyInterpolation": "Apply style model (interpolation)"
 }
\ No newline at end of file