From edb7caf31873eab3a91a21d118e50fbb813b0132 Mon Sep 17 00:00:00 2001
From: cubiq <matteo@elf.io>
Date: Sat, 7 Sep 2024 15:11:31 +0200
Subject: [PATCH] add very experimental flux blocks buster

---
 conditioning.py | 20 +++++++++++++++++---
 misc.py         | 18 +++++++++++++++++-
 2 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/conditioning.py b/conditioning.py
index 8ac9796..65da2bf 100644
--- a/conditioning.py
+++ b/conditioning.py
@@ -210,7 +210,7 @@ class FluxBlocksBuster:
     def INPUT_TYPES(s):
         return {"required": {
             "model": ("MODEL",),
-            "blocks": ("STRING", {"default": "double_blocks\.0\.(img|txt)_(mod|attn|mlp)\.(lin|qkv|proj|0|2)\.(weight|bias)=1.5\nsingle_blocks\.0\.(linear[12]|modulation\.lin)\.(weight|bias)=1.5", "multiline": True, "dynamicPrompts": True}),
+            "blocks": ("STRING", {"default": "## 0 = 1.0\n## 1 = 1.0\n## 2 = 1.0\n## 3 = 1.0\n## 4 = 1.0\n## 5 = 1.0\n## 6 = 1.0\n## 7 = 1.0\n## 8 = 1.0\n## 9 = 1.0\n## 10 = 1.0\n## 11 = 1.0\n## 12 = 1.0\n## 13 = 1.0\n## 14 = 1.0\n## 15 = 1.0\n## 16 = 1.0\n## 17 = 1.0\n## 18 = 1.0\n# 0 = 1.0\n# 1 = 1.0\n# 2 = 1.0\n# 3 = 1.0\n# 4 = 1.0\n# 5 = 1.0\n# 6 = 1.0\n# 7 = 1.0\n# 8 = 1.0\n# 9 = 1.0\n# 10 = 1.0\n# 11 = 1.0\n# 12 = 1.0\n# 13 = 1.0\n# 14 = 1.0\n# 15 = 1.0\n# 16 = 1.0\n# 17 = 1.0\n# 18 = 1.0\n# 19 = 1.0\n# 20 = 1.0\n# 21 = 1.0\n# 22 = 1.0\n# 23 = 1.0\n# 24 = 1.0\n# 25 = 1.0\n# 26 = 1.0\n# 27 = 1.0\n# 28 = 1.0\n# 29 = 1.0\n# 30 = 1.0\n# 31 = 1.0\n# 32 = 1.0\n# 33 = 1.0\n# 34 = 1.0\n# 35 = 1.0\n# 36 = 1.0\n# 37 = 1.0", "multiline": True, "dynamicPrompts": True}),
             #**{f"double_block_{s}": ("FLOAT", { "display": "slider", "default": 1.0, "min": 0, "max": 5, "step": 0.05 }) for s in range(19)},
             #**{f"single_block_{s}": ("FLOAT", { "display": "slider", "default": 1.0, "min": 0, "max": 5, "step": 0.05 }) for s in range(38)},
         }}
@@ -223,18 +223,32 @@ class FluxBlocksBuster:
         m = model.clone()
         sd = model.model_state_dict()
 
-        # blocks is a regex string
+        """
+        Also compatible with the following format:
+
+        double_blocks\.0\.(img|txt)_(mod|attn|mlp)\.(lin|qkv|proj|0|2)\.(weight|bias)=1.1
+        single_blocks\.0\.(linear[12]|modulation\.lin)\.(weight|bias)=1.1
+
+        The regex is used to match the block names
+        """
+
         blocks = blocks.split("\n")
         blocks = [b.strip() for b in blocks if b.strip()]
 
+        print("Patched blocks:")
         for k in sd:
             for block in blocks:
                 block = block.split("=")
                 value = float(block[1].strip()) if len(block) > 1 else 1.0
                 block = block[0].strip()
+                if block.startswith("##"):
+                    block = r"double_blocks\." + block[2:].strip() + r"\.(img|txt)_(mod|attn|mlp)\.(lin|qkv|proj|0|2)\.(weight|bias)"
+                elif block.startswith("#"):
+                    block = r"single_blocks\." + block[1:].strip() + r"\.(linear[12]|modulation\.lin)\.(weight|bias)"
+
                 if value != 1.0 and re.search(block, k):
-                    print(k, block, value)
                     m.add_patches({k: (None,)}, 0.0, value)
+                    print(f"{k}: {value}")
 
         return (m, )
 
diff --git a/misc.py b/misc.py
index c54ecdc..237a0a7 100644
--- a/misc.py
+++ b/misc.py
@@ -496,6 +496,7 @@ class DisplayAny:
         return {
             "required": {
                 "input": (("*",{})),
+                "mode": (["raw value", "tensor shape"],),
             },
         }
 
@@ -509,7 +510,22 @@ class DisplayAny:
 
     CATEGORY = "essentials/utilities"
 
-    def execute(self, input):
+    def execute(self, input, mode):
+        if mode == "tensor shape":
+            text = []
+            def tensorShape(tensor):
+                if isinstance(tensor, dict):
+                    for k in tensor:
+                        tensorShape(tensor[k])
+                elif isinstance(tensor, list):
+                    for i in range(len(tensor)):
+                        tensorShape(tensor[i])
+                elif hasattr(tensor, 'shape'):
+                    text.append(list(tensor.shape))
+
+            tensorShape(input)
+            input = text
+
         text = str(input)
 
         return {"ui": {"text": text}, "result": (text,)}