#!/usr/bin/env python3
"""
Inspect SDXL checkpoint structure.
Shows positional embeddings, encoder structure, and embedding keys.
"""

import sys
from safetensors import safe_open


def inspect(path):
    print(f"Inspecting: {path}\n")

    with safe_open(path, framework="pt", device="cpu") as f:
        keys = sorted(f.keys())

        print("=" * 70)
        print("POSITION-RELATED KEYS:")
        print("=" * 70)
        for key in keys:
            if "position" in key.lower():
                tensor = f.get_tensor(key)
                print(f"  {key}")
                print(f"    shape={tensor.shape} dtype={tensor.dtype}")
                print()

        print("=" * 70)
        print("ENCODER STRUCTURE:")
        print("=" * 70)
        prefixes = set()
        for key in keys:
            if any(x in key for x in ["conditioner.", "cond_stage", "text_encoder"]):
                prefix = ".".join(key.split(".")[:3])
                if prefix not in prefixes:
                    prefixes.add(prefix)

        for prefix in sorted(prefixes):
            count = sum(1 for k in keys if k.startswith(prefix))
            print(f"  {prefix}.*  ({count} tensors)")

        print()
        print("=" * 70)
        print("ALL EMBEDDING KEYS:")
        print("=" * 70)
        for key in keys:
            if "embedding" in key.lower():
                tensor = f.get_tensor(key)
                print(f"  {key}")
                print(f"    shape={tensor.shape}")

        print()
        print(f"TOTAL TENSORS: {len(keys)}")


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python inspect_checkpoint.py <path.safetensors>")
        sys.exit(1)
    inspect(sys.argv[1])