diff --git a/essentials.py b/essentials.py index 0602e61..4b20e8f 100644 --- a/essentials.py +++ b/essentials.py @@ -803,6 +803,41 @@ class ConsoleDebug: return (None,) +class DebugTensorShape: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "tensor": (any, {}), + }, + } + + RETURN_TYPES = () + FUNCTION = "execute" + CATEGORY = "essentials" + OUTPUT_NODE = True + + def execute(self, tensor): + shapes = [] + 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'): + shapes.append(list(tensor.shape)) + + tensorShape(tensor) + + print(f"\033[96mShapes found: {shapes}\033[0m") + + return (None,) + class BatchCount: @classmethod def INPUT_TYPES(s): @@ -903,6 +938,7 @@ NODE_CLASS_MAPPINGS = { "SimpleMath+": SimpleMath, "ConsoleDebug+": ConsoleDebug, + "DebugTensorShape+": DebugTensorShape, "ModelCompile+": ModelCompile, "BatchCount+": BatchCount, @@ -936,6 +972,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "SimpleMath+": "🔧 Simple Math", "ConsoleDebug+": "🔧 Console Debug", + "DebugTensorShape+": "🔧 Tensor Shape Debug", "ModelCompile+": "🔧 Compile Model", "BatchCount+": "🔧 Batch Count",