"""Check JoyAI-Echo GGUF files for the missing text_embedding_projection tensors.

A GGUF stores its tensor count as a little-endian uint64 at byte offset 8, so
this reads 16 bytes per file. No model load, no dependencies, no GPU.

    4448 = complete
    4444 = the four text_embedding_projection tensors are missing

Those four inject text conditioning into the audio and video branches. When they
are absent they get built random-initialised at load time, the model runs, and
nothing errors -- you just get a flat synthetic voice, the model reading your
prompt aloud, identity resetting between shots, and style instructions ignored.

Usage:
    python check_gguf_tensors.py                 # scan current directory
    python check_gguf_tensors.py  path/to/models # scan a directory (recursive)
    python check_gguf_tensors.py  a.gguf b.gguf  # check specific files
"""
import os
import sys

EXPECTED = 4448
BROKEN = 4444


def n_tensors(path):
    """Return the tensor count from a GGUF header, or None if not a GGUF."""
    with open(path, "rb") as f:
        head = f.read(16)
    if len(head) < 16 or head[:4] != b"GGUF":
        return None
    return int.from_bytes(head[8:16], "little")


def collect(args):
    if not args:
        args = ["."]
    files = []
    for a in args:
        if os.path.isdir(a):
            for root, _, names in os.walk(a):
                files += [os.path.join(root, n) for n in sorted(names)
                          if n.lower().endswith(".gguf")]
        elif os.path.isfile(a):
            files.append(a)
        else:
            print(f"  ! not found: {a}")
    return files


def main():
    files = collect(sys.argv[1:])
    if not files:
        print("No .gguf files found.")
        return 0

    worst = 0
    for p in files:
        try:
            n = n_tensors(p)
        except OSError as e:
            print(f"  ?  {os.path.basename(p)}  (unreadable: {e})")
            continue

        name = os.path.basename(p)
        # The 4448/4444 comparison is only meaningful for a JoyAI-Echo DiT.
        # Other models -- including base LTX-2.3 -- have their own tensor
        # counts, and 4444 may well be correct for them. Only judge files
        # that identify themselves as JoyAI-Echo.
        is_joy = "joyai" in name.lower() or "joy-echo" in name.lower()

        if n is None:
            print(f"  ?       {name}  (not a GGUF file)")
        elif not is_joy:
            print(f"  --      {n}  {name}  (not a JoyAI-Echo DiT - not judged)")
        elif n == EXPECTED:
            print(f"  OK      {n}  {name}")
        elif n == BROKEN:
            print(f"  BROKEN  {n}  {name}  <- 4 tensors missing")
            worst = 1
        else:
            print(f"  ?       {n}  {name}  (unexpected count for a JoyAI-Echo DiT)")

    print()
    if worst:
        print("At least one JoyAI-Echo DiT is missing the text_embedding_projection")
        print("tensors. Replace it with a build that reports 4448.")
    else:
        print("No JoyAI-Echo DiT with the 4444 defect found.")
    print()
    print("Note: files that are not JoyAI-Echo builds are listed but not judged.")
    print("Base LTX-2.3 and other architectures have their own tensor counts, and")
    print("4444 is not necessarily wrong for them.")
    return worst


if __name__ == "__main__":
    sys.exit(main())