File size: 1,420 Bytes
c450fa4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
"""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"}