activation_dump_node/profiler_node.py

"""star_inspect_model_v2.py — dump full module tree for blocks.0"""
import torch

class StarInspectModelV2:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"model": ("MODEL",)}}

    RETURN_TYPES = ("MODEL",)
    FUNCTION = "inspect"
    CATEGORY = "⭐StarNodes/Model Tools"

    def inspect(self, model):
        real_model = model.model
        target = real_model.diffusion_model if hasattr(real_model, "diffusion_model") else real_model

        print("=" * 70)
        print("FULL TREE for blocks.0 (and top-level non-block modules):")
        print("=" * 70)

        for name, module in target.named_modules():
            cls = type(module).__name__
            if name.startswith("blocks.0.") or name == "blocks.0" or "blocks" not in name:
                params = [f"{pn}:{tuple(p.shape)}" for pn, p in module.named_parameters(recurse=False)]
                print(f"name={name!r}  class={cls}  params={params}")

        print("=" * 70)
        print("ALL PARAMETER NAMES under blocks.0 (raw state_dict style):")
        for pname, p in target.named_parameters():
            if pname.startswith("blocks.0."):
                print(f"  {pname}  shape={tuple(p.shape)}  dtype={p.dtype}")
        print("=" * 70)

        return (model,)


NODE_CLASS_MAPPINGS = {"StarInspectModelV2": StarInspectModelV2}
NODE_DISPLAY_NAME_MAPPINGS = {"StarInspectModelV2": "⭐ Star Inspect Model V2"}