From 2d18a14d5cbbeceecf7786f6c328fb185c904341 Mon Sep 17 00:00:00 2001
From: spacepxl <spacetug93@gmail.com>
Date: Tue, 23 Jul 2024 02:15:16 -0400
Subject: [PATCH] updated latent stats to support arbitrary channel counts,
 added ModelTest node for messing around with loaded models

---
 nodes.py | 29 +++++++++++++++++++++++++++--
 1 file changed, 27 insertions(+), 2 deletions(-)

diff --git a/nodes.py b/nodes.py
index 6506755..64a9e76 100644
--- a/nodes.py
+++ b/nodes.py
@@ -1236,19 +1236,22 @@ class LatentStats:
 
     def notify(self, latent):
         latents = latent["samples"]
+        channels = latents.size(1)
         width, height = latents.size(3), latents.size(2)
         
         text = ["",]
         text[0] = f"batch size: {latents.size(0)}"
+        text.append(f"channels: {channels}")
         text.append(f"width: {width} ({width * 8})")
         text.append(f"height: {height} ({height * 8})")
         
         cmean = [0,0,0,0]
-        for i in range(4):
+        for i in range(channels):
             minimum = torch.min(latents[:,i,:,:]).item()
             maximum = torch.max(latents[:,i,:,:]).item()
             std_dev, mean = torch.std_mean(latents[:,i,:,:], dim=None)
-            cmean[i] = mean
+            if i < 4:
+                cmean[i] = mean
             
             text.append(f"c{i} mean: {mean:.1f} std_dev: {std_dev:.1f} min: {minimum:.1f} max: {maximum:.1f}")
         
@@ -1953,6 +1956,26 @@ class GameOfLife:
         
         return (image, mask, off, on)
 
+modeltest_code_default = """d = model.model.model_config.unet_config
+for k in d.keys():
+    print(k, d[k])"""
+
+class ModelTest:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+            "model": ("MODEL",),
+            "code": ("STRING", {"multiline": True, "default": modeltest_code_default}),
+            }}
+    RETURN_TYPES = ()
+    FUNCTION = "test"
+    OUTPUT_NODE = True
+    CATEGORY = "utils"
+    
+    def test(self, model, code):
+        exec(code)
+        return ()
+
 NODE_CLASS_MAPPINGS = {
     "AdainFilterLatent": AdainFilterLatent,
     "AdainImage": AdainImage,
@@ -1989,6 +2012,7 @@ NODE_CLASS_MAPPINGS = {
     "LatentNormalizeShuffle": LatentNormalizeShuffle,
     "LatentStats": LatentStats,
     "MedianFilterImage": MedianFilterImage,
+    "ModelTest": ModelTest,
     "NormalMapSimple": NormalMapSimple,
     "OffsetLatentImage": OffsetLatentImage,
     "PrintSigmas": PrintSigmas,
@@ -2039,6 +2063,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "LatentNormalizeShuffle": "LatentNormalizeShuffle",
     "LatentStats": "Latent Stats",
     "MedianFilterImage": "Median Filter Image",
+    "ModelTest": "Model Test",
     "NormalMapSimple": "Normal Map (Simple)",
     "OffsetLatentImage": "Offset Latent Image",
     "PrintSigmas": "PrintSigmas",