Merge pull request #1 from liusida/dev

Two nodes for inspecting models

__init__.py

@@ -0,0 +1,162 @@
+import inspect
+import io
+
+def size_to_human_readable(size_in_bytes, in_bytes=True):
+    if size_in_bytes < 1024:
+        return f"{size_in_bytes} "
+    elif size_in_bytes < 1024**2:
+        return f"{size_in_bytes / 1024:.1f} K"
+    elif size_in_bytes < 1024**3:
+        return f"{size_in_bytes / 1024**2:.1f} M"
+    elif size_in_bytes < 1024**4:
+        if in_bytes:
+            return f"{size_in_bytes / 1024**3:.1f} G"
+        else:
+            return f"{size_in_bytes / 1024**3:.1f} B"
+    else:
+        return f"{size_in_bytes / 1024**4:.1f} T"
+
+class DebugInspectorNode:
+    def __init__(self):
+        pass
+    
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "string_field": ("STRING", {
+                    "multiline": True, #True if you want the field to look like the one on the ClipTextEncode node
+                    "default": "Connect any of the inputs to get more information."
+                }),
+            },
+            "optional": {
+                "model": ("MODEL",),
+                "clip": ("CLIP",),
+                "vae": ("VAE",),
+                "conditioning": ("CONDITIONING",),
+                "latent": ("LATENT",),
+            },
+        }
+
+    RETURN_TYPES = ()
+    FUNCTION = "watch"
+    OUTPUT_NODE = True
+    CATEGORY = "utils"
+
+    def watch(self, string_field, model=None, clip=None, vae=None, CONDITIONING=None, LATENT=None):
+        results = ""
+        if model:
+            results += "MODEL:\n"
+            results += "\tArch:\t" + model.model.latent_format.__class__.__name__ + "\n"
+            results += f"\tSize:\t{size_to_human_readable(model.size)}\n"
+            results += "\n"
+
+        if clip:
+            results += "CLIP:\n\n"
+            results += str(type(clip.tokenizer)) + "\n"
+            results += "\n"
+
+        if vae:
+            results += "VAE:\n\n"
+            results += str(type(vae.first_stage_model)) + "\n"
+            results += "\n"
+
+        return { "ui": { "string_field": results } }
+
+class DebugModelInspectorNode:
+    def __init__(self):
+        pass
+    
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "output": ("STRING", {
+                    "multiline": True, #True if you want the field to look like the one on the ClipTextEncode node
+                    "default": ""
+                }),
+            },
+        }
+
+    RETURN_TYPES = ()
+    FUNCTION = "watch"
+    OUTPUT_NODE = True
+    CATEGORY = "utils"
+
+    def watch(self, output, model=None):
+        results = ""
+        if model:
+            results += "Base:\t\t" + model.model.latent_format.__class__.__name__ + "\n"
+            results += f"Size:\t\t{size_to_human_readable(model.size, in_bytes=True)}B\n"
+            parameter_size = sum(p.numel() for p in model.model.parameters())
+            results += f"Parameters:\t{size_to_human_readable(parameter_size, in_bytes=False)}\n"
+
+            dtype_value = model.model.get_dtype()  # Example: <class 'torch.bfloat16'>
+            dtype_str = str(dtype_value).split('.')[-1]  # Split by '.' and take the last part
+            print(dtype_str)  # This will print 'bfloat16'
+            results += f"dtype:\t\t{dtype_str}\n"
+
+            results += f"Model Type:\t{model.model.model_type.name}\n"
+
+            results += "\n"
+
+        return { "ui": { "output": results } }
+
+class DebugModelPrintOutNode:
+    def __init__(self):
+        pass
+    
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "output": ("STRING", {
+                    "multiline": True, #True if you want the field to look like the one on the ClipTextEncode node
+                    "default": ""
+                }),
+            },
+        }
+
+    RETURN_TYPES = ()
+    FUNCTION = "watch"
+    OUTPUT_NODE = True
+    CATEGORY = "utils"
+
+    def watch(self, output, model=None):
+        results = ""
+        # Check if the model object itself exists and has the attribute 'model'
+        if model and hasattr(model, 'model'):
+            # Check if the 'model' attribute of the model has the attribute 'diffusion_model'
+            if hasattr(model.model, 'diffusion_model'):
+                # Now we can safely print the diffusion_model
+                buffer = io.StringIO()
+                print(model.model.diffusion_model, file=buffer)
+                results += buffer.getvalue()
+                buffer.close()
+            else:
+                results += "No diffusion_model attribute found.\n"
+        else:
+            results += "Model is not defined or does not have a 'model' attribute.\n"
+
+        return { "ui": { "output": results } }
+
+
+# Set the web directory, any .js file in that directory will be loaded by the frontend as a frontend extension
+WEB_DIRECTORY = "./js"
+
+# A dictionary that contains all nodes you want to export with their names
+# NOTE: names should be globally unique
+NODE_CLASS_MAPPINGS = {
+    # "DebugInspectorNode": DebugInspectorNode,
+    "DebugModelInspectorNode": DebugModelInspectorNode,
+    "DebugModelPrintOutNode": DebugModelPrintOutNode,
+}
+
+# A dictionary that contains the friendly/humanly readable titles for the nodes
+NODE_DISPLAY_NAME_MAPPINGS = {
+    # "DebugInspectorNode": "Debug: Inspector",
+    "DebugModelInspectorNode": "Debug: Model Metrics",
+    "DebugModelPrintOutNode": "Debug: Model Architecture"
+}

js/debug_node.js

@@ -0,0 +1,6 @@
+import {app} from "../../scripts/app.js";
+
+// // Node that add notes to your project
+// app.registerExtension({
+// 	name: "Comfy.Debug",
+// });

js/string_output.js

@@ -0,0 +1,43 @@
+import { api } from '../../scripts/api.js';
+import { app } from '../../scripts/app.js'; // Assuming `app` exports the current application instance
+
+// Function to update STRING widget's content based on server response
+function updateStringWidgetContent(nodeId, widgetName, newContent) {
+    const node = app.graph._nodes.find(n => n.id === nodeId);
+    if (!node) {
+        console.error('Node not found:', nodeId);
+        return;
+    }
+
+    const widget = node.widgets.find(w => w.name === widgetName && w.type === 'customtext');
+    if (!widget) {
+        console.error('STRING widget not found:', widgetName);
+        return;
+    }
+
+    widget.value = newContent;
+    // Optionally, if there's a callback function defined for the widget, call it
+    if (widget.callback) {
+        widget.callback(newContent);
+    }
+}
+
+// Listen for a specific event from the server indicating a script execution is completed
+api.addEventListener('executed', event => {
+    const detail = event.detail;
+
+    // Check if `output` exists and has a `output`
+    if (detail.output && Array.isArray(detail.output.output)) {
+        // Join the array elements into a string
+        const content = detail.output.output.join('');
+
+        // Example node ID and widget name, replace with actual ones
+        const nodeId = parseInt(app.runningNodeId); // Replace with actual node ID
+        const widgetName = 'output'; // Replace with actual widget name
+
+        updateStringWidgetContent(nodeId, widgetName, content);
+    } else {
+        console.error('Unexpected detail format:', detail);
+    }
+
+});