| import os
|
| import torch
|
| from collections import Counter, defaultdict
|
|
|
| ROOT = r"E:\Models SD XL\UNET\Minimax_H3"
|
|
|
| H3_FILE = os.path.join(
|
| ROOT,
|
| "h3_bridge_hidden_480.pt"
|
| )
|
|
|
| SN_FILE = os.path.join(
|
| ROOT,
|
| "sensenova_bridge_hidden_480.pt"
|
| )
|
|
|
| BRIDGE_FILE = (
|
| r"D:\ComfyUI_Krea2\ComfyUI\models\bridge"
|
| r"\SN_L32_to_H3_L49_rank128.safetensors"
|
| )
|
|
|
| REPORT_FILE = os.path.join(
|
| ROOT,
|
| "distillation_source_inspection.txt"
|
| )
|
|
|
|
|
|
|
|
|
|
|
|
|
| def tensor_info(t):
|
| return (
|
| f"shape={tuple(t.shape)} "
|
| f"dtype={t.dtype} "
|
| f"device={t.device}"
|
| )
|
|
|
|
|
| def walk(obj, path="root", out=None, depth=0, max_depth=8):
|
| if out is None:
|
| out = []
|
|
|
| if depth > max_depth:
|
| return out
|
|
|
| if torch.is_tensor(obj):
|
| out.append({
|
| "path": path,
|
| "shape": tuple(obj.shape),
|
| "dtype": str(obj.dtype),
|
| "tensor": obj,
|
| })
|
| return out
|
|
|
| if isinstance(obj, dict):
|
| for key, value in obj.items():
|
| walk(
|
| value,
|
| f"{path}.{key}",
|
| out,
|
| depth + 1,
|
| max_depth,
|
| )
|
|
|
| elif isinstance(obj, (list, tuple)):
|
| for i, value in enumerate(obj):
|
| walk(
|
| value,
|
| f"{path}[{i}]",
|
| out,
|
| depth + 1,
|
| max_depth,
|
| )
|
|
|
| return out
|
|
|
|
|
| def print_structure(obj, name, lines, depth=0, max_depth=4, max_items=4):
|
| indent = " " * depth
|
|
|
| if depth > max_depth:
|
| lines.append(f"{indent}...")
|
| return
|
|
|
| if torch.is_tensor(obj):
|
| lines.append(
|
| f"{indent}{tensor_info(obj)}"
|
| )
|
| return
|
|
|
| if isinstance(obj, dict):
|
| lines.append(
|
| f"{indent}dict keys={len(obj)}"
|
| )
|
|
|
| for i, (key, value) in enumerate(obj.items()):
|
| if i >= max_items:
|
| lines.append(
|
| f"{indent} ... ({len(obj) - max_items} more)"
|
| )
|
| break
|
|
|
| lines.append(
|
| f"{indent} KEY: {repr(key)}"
|
| )
|
|
|
| print_structure(
|
| value,
|
| name,
|
| lines,
|
| depth + 2,
|
| max_depth,
|
| max_items,
|
| )
|
|
|
| elif isinstance(obj, (list, tuple)):
|
| lines.append(
|
| f"{indent}{type(obj).__name__} len={len(obj)}"
|
| )
|
|
|
| for i, value in enumerate(obj[:max_items]):
|
| lines.append(
|
| f"{indent} ITEM [{i}]"
|
| )
|
|
|
| print_structure(
|
| value,
|
| name,
|
| lines,
|
| depth + 2,
|
| max_depth,
|
| max_items,
|
| )
|
|
|
| if len(obj) > max_items:
|
| lines.append(
|
| f"{indent} ... ({len(obj) - max_items} more)"
|
| )
|
|
|
| else:
|
| value_repr = repr(obj)
|
|
|
| if len(value_repr) > 180:
|
| value_repr = value_repr[:180] + "..."
|
|
|
| lines.append(
|
| f"{indent}{type(obj).__name__}: {value_repr}"
|
| )
|
|
|
|
|
| def analyze_file(label, path, lines):
|
| lines.append("")
|
| lines.append("=" * 100)
|
| lines.append(label)
|
| lines.append("=" * 100)
|
| lines.append(f"Path: {path}")
|
| lines.append(
|
| f"Size: {os.path.getsize(path) / (1024 ** 2):.2f} MiB"
|
| )
|
|
|
| print(f"\nLoading {label}...")
|
| data = torch.load(
|
| path,
|
| map_location="cpu",
|
| weights_only=False,
|
| )
|
|
|
| print(f"{label} loaded.")
|
|
|
| lines.append("")
|
| lines.append("TOP-LEVEL STRUCTURE")
|
| lines.append("-" * 100)
|
|
|
| print_structure(
|
| data,
|
| label,
|
| lines,
|
| )
|
|
|
| tensors = walk(data)
|
|
|
| lines.append("")
|
| lines.append(
|
| f"Total tensors found recursively: {len(tensors)}"
|
| )
|
|
|
| shape_counter = Counter(
|
| item["shape"]
|
| for item in tensors
|
| )
|
|
|
| lines.append("")
|
| lines.append("MOST COMMON TENSOR SHAPES")
|
| lines.append("-" * 100)
|
|
|
| for shape, count in shape_counter.most_common(20):
|
| lines.append(
|
| f"{count:6d} x {shape}"
|
| )
|
|
|
|
|
|
|
|
|
|
|
| candidates_4096 = []
|
| candidates_5120 = []
|
|
|
| for item in tensors:
|
| shape = item["shape"]
|
|
|
| if len(shape) >= 1:
|
| if shape[-1] == 4096:
|
| candidates_4096.append(item)
|
|
|
| if shape[-1] == 5120:
|
| candidates_5120.append(item)
|
|
|
| lines.append("")
|
| lines.append("4096-DIM CANDIDATES")
|
| lines.append("-" * 100)
|
| lines.append(
|
| f"Count: {len(candidates_4096)}"
|
| )
|
|
|
| for item in candidates_4096[:30]:
|
| lines.append(
|
| f"{item['path']} | "
|
| f"{item['shape']} | "
|
| f"{item['dtype']}"
|
| )
|
|
|
| if len(candidates_4096) > 30:
|
| lines.append(
|
| f"... {len(candidates_4096) - 30} more"
|
| )
|
|
|
| lines.append("")
|
| lines.append("5120-DIM CANDIDATES")
|
| lines.append("-" * 100)
|
| lines.append(
|
| f"Count: {len(candidates_5120)}"
|
| )
|
|
|
| for item in candidates_5120[:30]:
|
| lines.append(
|
| f"{item['path']} | "
|
| f"{item['shape']} | "
|
| f"{item['dtype']}"
|
| )
|
|
|
| if len(candidates_5120) > 30:
|
| lines.append(
|
| f"... {len(candidates_5120) - 30} more"
|
| )
|
|
|
|
|
|
|
|
|
|
|
| lines.append("")
|
| lines.append("PATHS CONTAINING LAYER 32 / L32")
|
| lines.append("-" * 100)
|
|
|
| layer32 = [
|
| item
|
| for item in tensors
|
| if (
|
| "32" in item["path"].lower()
|
| or "l32" in item["path"].lower()
|
| or "layer_32" in item["path"].lower()
|
| or "layer32" in item["path"].lower()
|
| )
|
| ]
|
|
|
| lines.append(
|
| f"Count: {len(layer32)}"
|
| )
|
|
|
| for item in layer32[:40]:
|
| lines.append(
|
| f"{item['path']} | "
|
| f"{item['shape']}"
|
| )
|
|
|
| lines.append("")
|
| lines.append("PATHS CONTAINING LAYER 49 / L49")
|
| lines.append("-" * 100)
|
|
|
| layer49 = [
|
| item
|
| for item in tensors
|
| if (
|
| "49" in item["path"].lower()
|
| or "l49" in item["path"].lower()
|
| or "layer_49" in item["path"].lower()
|
| or "layer49" in item["path"].lower()
|
| )
|
| ]
|
|
|
| lines.append(
|
| f"Count: {len(layer49)}"
|
| )
|
|
|
| for item in layer49[:40]:
|
| lines.append(
|
| f"{item['path']} | "
|
| f"{item['shape']}"
|
| )
|
|
|
| return data, tensors
|
|
|
|
|
|
|
|
|
|
|
|
|
| print("=" * 100)
|
| print("DISTILLATION SOURCE INSPECTION")
|
| print("=" * 100)
|
|
|
| for path in [
|
| H3_FILE,
|
| SN_FILE,
|
| BRIDGE_FILE,
|
| ]:
|
| if not os.path.exists(path):
|
| raise FileNotFoundError(
|
| f"Missing required file:\n{path}"
|
| )
|
|
|
| print("H3:")
|
| print(H3_FILE)
|
|
|
| print("\nSenseNova:")
|
| print(SN_FILE)
|
|
|
| print("\nBridge:")
|
| print(BRIDGE_FILE)
|
|
|
| lines = []
|
|
|
| lines.append(
|
| "SenseNova -> MiniMax H3 Distillation Source Inspection"
|
| )
|
|
|
| lines.append(
|
| "=" * 100
|
| )
|
|
|
| h3_data, h3_tensors = analyze_file(
|
| "MINIMAX H3 DATA",
|
| H3_FILE,
|
| lines,
|
| )
|
|
|
| sn_data, sn_tensors = analyze_file(
|
| "SENSENOVA DATA",
|
| SN_FILE,
|
| lines,
|
| )
|
|
|
|
|
|
|
|
|
|
|
|
|
| lines.append("")
|
| lines.append("=" * 100)
|
| lines.append("BRIDGE FILE")
|
| lines.append("=" * 100)
|
| lines.append(f"Path: {BRIDGE_FILE}")
|
|
|
| from safetensors.torch import load_file
|
|
|
| bridge = load_file(
|
| BRIDGE_FILE,
|
| device="cpu",
|
| )
|
|
|
| for key, tensor in bridge.items():
|
| lines.append(
|
| f"{key:40s} "
|
| f"{tuple(tensor.shape)} "
|
| f"{tensor.dtype}"
|
| )
|
|
|
|
|
|
|
|
|
|
|
|
|
| h3_5120 = [
|
| x for x in h3_tensors
|
| if (
|
| len(x["shape"]) > 0
|
| and x["shape"][-1] == 5120
|
| )
|
| ]
|
|
|
| sn_4096 = [
|
| x for x in sn_tensors
|
| if (
|
| len(x["shape"]) > 0
|
| and x["shape"][-1] == 4096
|
| )
|
| ]
|
|
|
| lines.append("")
|
| lines.append("=" * 100)
|
| lines.append("FINAL SUMMARY")
|
| lines.append("=" * 100)
|
|
|
| lines.append(
|
| f"H3 5120-dim tensors found: {len(h3_5120)}"
|
| )
|
|
|
| lines.append(
|
| f"SenseNova 4096-dim tensors found: {len(sn_4096)}"
|
| )
|
|
|
| lines.append("")
|
| lines.append(
|
| "The next script will use this report to identify the exact "
|
| "H3 input representation and SenseNova teacher target without "
|
| "guessing the .pt layout."
|
| )
|
|
|
|
|
|
|
|
|
|
|
|
|
| with open(
|
| REPORT_FILE,
|
| "w",
|
| encoding="utf-8",
|
| ) as f:
|
| f.write(
|
| "\n".join(lines)
|
| )
|
|
|
| print()
|
| print("=" * 100)
|
| print("SUCCESS")
|
| print("=" * 100)
|
|
|
| print("Report saved:")
|
| print(REPORT_FILE)
|
|
|
| print()
|
| print(
|
| "H3 5120-dim tensors:",
|
| len(h3_5120)
|
| )
|
|
|
| print(
|
| "SenseNova 4096-dim tensors:",
|
| len(sn_4096)
|
| )
|
|
|
| print()
|
| print(
|
| "Send me the contents of "
|
| "distillation_source_inspection.txt"
|
| ) |
