inspect_dav.py

"""Safely inspect a MiniMax Music 3 ``dav.pth`` checkpoint.

This module deliberately inspects checkpoint structure and tensor metadata only.
It never imports a model implementation and always maps tensors to CPU.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import pickle
import re
from collections import Counter
from collections.abc import Mapping, Sequence
from pathlib import Path
from typing import Any

import torch


SEARCH_TERMS = (
    "generator",
    "62000_generator",
    "encoder",
    "decoder",
    "quantizer",
    "vq",
    "rvq",
    "codebook",
    "codec",
    "tokenizer",
)
MAX_KEY_EXAMPLES = 12
MAX_TENSOR_EXAMPLES = 4


def safe_load_checkpoint(path: str | os.PathLike[str]) -> Any:
    """Load a checkpoint without permitting arbitrary pickled Python objects."""

    checkpoint_path = Path(path).expanduser()
    if not checkpoint_path.is_file():
        raise FileNotFoundError(f"DAV checkpoint does not exist: {checkpoint_path}")
    try:
        return torch.load(str(checkpoint_path), map_location="cpu", weights_only=True)
    except (pickle.UnpicklingError, EOFError, OSError, RuntimeError) as error:
        raise ValueError(f"unable to safely load DAV checkpoint {checkpoint_path}: {error}") from error


def _child_path(path: str, component: str) -> str:
    return f"{path}.{component}" if path != "$" else f"$.{component}"


def _walk(
    value: Any,
    path: str = "$",
    *,
    seen: set[int] | None = None,
):
    """Yield ``(path, value)`` recursively for safe checkpoint containers."""

    if seen is None:
        seen = set()
    yield path, value
    if isinstance(value, torch.Tensor):
        return
    if isinstance(value, Mapping):
        identity = id(value)
        if identity in seen:
            return
        seen.add(identity)
        for key, child in value.items():
            yield from _walk(child, _child_path(path, str(key)), seen=seen)
        return
    if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)):
        identity = id(value)
        if identity in seen:
            return
        seen.add(identity)
        for index, child in enumerate(value):
            yield from _walk(child, _child_path(path, str(index)), seen=seen)


def _state_dict_candidates(checkpoint: Any) -> list[tuple[str, Mapping[str, Any], int]]:
    candidates: list[tuple[str, Mapping[str, Any], int]] = []
    for path, value in _walk(checkpoint):
        if not isinstance(value, Mapping) or not all(isinstance(key, str) for key in value):
            continue
        direct_tensors = sum(isinstance(child, torch.Tensor) for child in value.values())
        if direct_tensors:
            candidates.append((path, value, direct_tensors))
    return candidates


def _select_state_dict(checkpoint: Any) -> tuple[str, Mapping[str, Any]]:
    candidates = _state_dict_candidates(checkpoint)
    if not candidates:
        raise ValueError("checkpoint contains no string-keyed tensor state dictionary")

    # The number of direct tensor entries is the strongest state-dict signal.
    # A named wrapper breaks ties in favor of the conventional payload path.
    wrapper_names = ("62000_generator", "state_dict", "generator", "model")

    def rank(candidate: tuple[str, Mapping[str, Any], int]) -> tuple[int, int, int]:
        path, _, tensor_count = candidate
        wrapper_score = sum(name in path.lower() for name in wrapper_names)
        depth = path.count(".")
        return tensor_count, wrapper_score, depth

    path, state, _ = max(candidates, key=rank)
    return path, state


def _summarize_keys(keys: list[str], limit: int = MAX_KEY_EXAMPLES) -> dict[str, Any]:
    return {
        "count": len(keys),
        "examples": keys[:limit],
        "truncated": len(keys) > limit,
    }


def _tensor_info(name: str, tensor: torch.Tensor) -> dict[str, Any]:
    return {
        "name": name,
        "shape": list(tensor.shape),
        "dtype": str(tensor.dtype).removeprefix("torch."),
        "numel": tensor.numel(),
    }


def _term_matches(term: str, path: str) -> bool:
    lowered = path.lower()
    if term == "vq":
        return re.search(r"(?:^|[./_])vq(?:$|[./_])", lowered) is not None
    if term in {"encoder", "decoder", "generator", "codec", "tokenizer"}:
        return re.search(rf"(?:^|[./_]){re.escape(term)}(?:$|[./_])", lowered) is not None
    return term in lowered


def _prefix_counts(tensors: Mapping[str, torch.Tensor]) -> Counter[str]:
    return Counter(name.split(".", 1)[0] for name in tensors)


def _has_module(keys: list[str], name: str) -> bool:
    pattern = re.compile(rf"(?:^|[./_]){re.escape(name)}(?:$|[./_])", re.IGNORECASE)
    return any(pattern.search(key) for key in keys)


def inspect_checkpoint(
    path: str | os.PathLike[str],
    *,
    include_sha256: bool = False,
) -> dict[str, Any]:
    """Return a JSON-serializable structural and capability report."""

    checkpoint_path = Path(path).expanduser().resolve()
    checkpoint = safe_load_checkpoint(checkpoint_path)
    state_path, state = _select_state_dict(checkpoint)
    tensors = {key: value for key, value in state.items() if isinstance(value, torch.Tensor)}
    tensor_keys = list(tensors)

    walked = list(_walk(checkpoint))
    all_paths = [item_path.removeprefix("$.") for item_path, _ in walked]
    non_tensor_paths = [
        item_path.removeprefix("$.")
        for item_path, value in walked
        if not isinstance(value, (torch.Tensor, Mapping, list, tuple))
    ]
    serialized_config = any(
        re.search(r"(?:^|[./_])(config|architecture|hparams|hyper_parameters|args)(?:$|[./_])", item, re.I)
        for item in non_tensor_paths
    )

    has_analysis_encoder = _has_module(tensor_keys, "encoder")
    has_mean = any(key.startswith("mean_proj.") or ".mean_proj." in key for key in tensor_keys)
    has_logs = any(key.startswith("logs_proj.") or ".logs_proj." in key for key in tensor_keys)
    has_quantizer = _has_module(tensor_keys, "quantizer") or _has_module(tensor_keys, "rvq") or _has_module(
        tensor_keys, "vq"
    )
    has_codebooks = _has_module(tensor_keys, "codebook")
    has_decoder = _has_module(tensor_keys, "decoder") and any("dec_in_proj" in key for key in tensor_keys)
    has_flow = _has_module(tensor_keys, "flow")
    native_weight_set = has_analysis_encoder and has_quantizer and has_codebooks
    continuous_analysis_evidence = has_analysis_encoder and has_mean and has_logs and has_flow
    # Names and tensor shapes can identify candidates, but they cannot prove the
    # exact token semantics or supply an executable compatible API.
    native_tokenizer_verified = False

    prefix_counts = _prefix_counts(tensors)
    prefix_groups: dict[str, Any] = {}
    for prefix, count in sorted(prefix_counts.items()):
        members = [(key, tensors[key]) for key in tensor_keys if key.split(".", 1)[0] == prefix]
        examples = members[:MAX_TENSOR_EXAMPLES]
        if len(members) > MAX_TENSOR_EXAMPLES:
            examples += members[-1:]
        prefix_groups[prefix] = {
            "tensor_count": count,
            "examples": [_tensor_info(key, tensor) for key, tensor in examples],
        }

    searches: dict[str, Any] = {}
    for term in SEARCH_TERMS:
        matching_paths = [candidate for candidate in all_paths if _term_matches(term, candidate)]
        matching_tensors = [key for key in tensor_keys if _term_matches(term, key)]
        searches[term] = {
            "present": bool(matching_paths),
            "path_count": len(matching_paths),
            "tensor_count": len(matching_tensors),
            "examples": matching_paths[:MAX_KEY_EXAMPLES],
            "truncated": len(matching_paths) > MAX_KEY_EXAMPLES,
        }

    total_numel = sum(tensor.numel() for tensor in tensors.values())
    total_bytes = sum(tensor.numel() * tensor.element_size() for tensor in tensors.values())
    report: dict[str, Any] = {
        "checkpoint": {
            "path": str(checkpoint_path),
            "size_bytes": checkpoint_path.stat().st_size,
            "safe_load": "torch.load(weights_only=True, map_location='cpu')",
            "root_type": type(checkpoint).__name__,
        },
        "top_level": _summarize_keys([str(key) for key in checkpoint] if isinstance(checkpoint, Mapping) else []),
        "state_dict_path": state_path,
        "state_dict": {
            "entry_count": len(state),
            "tensor_count": len(tensors),
            "non_tensor_entry_count": len(state) - len(tensors),
            "total_numel": total_numel,
            "tensor_bytes": total_bytes,
            "dtypes": dict(sorted(Counter(str(value.dtype).removeprefix("torch.") for value in tensors.values()).items())),
        },
        "prefix_groups": prefix_groups,
        "searches": searches,
        "capabilities": {
            "waveform_analysis_encoder_weights": has_analysis_encoder,
            "continuous_gaussian_posterior_heads": has_mean and has_logs,
            "rvq_or_vq_quantizer_weights": has_quantizer,
            "codebook_embedding_weights": has_codebooks,
            "waveform_decoder_weights": has_decoder,
            "continuous_flow_weights": has_flow,
            "serialized_architecture_config": serialized_config,
            "native_tokenizer_weight_set": native_weight_set,
            "compatible_executable_tokenizer_api_verified": native_tokenizer_verified,
            "native_discrete_tokenizer_complete": native_tokenizer_verified,
            "can_encode_native_music3_tokens": native_tokenizer_verified,
        },
        "verdict": _capability_verdict(
            continuous_analysis_evidence=continuous_analysis_evidence,
            has_quantizer=has_quantizer,
            has_codebooks=has_codebooks,
            native_weight_set=native_weight_set,
            serialized_config=serialized_config,
        ),
    }
    if include_sha256:
        digest = hashlib.sha256()
        with checkpoint_path.open("rb") as checkpoint_file:
            for chunk in iter(lambda: checkpoint_file.read(1024 * 1024), b""):
                digest.update(chunk)
        report["checkpoint"]["sha256"] = digest.hexdigest()
    return report


def _capability_verdict(
    *,
    continuous_analysis_evidence: bool,
    has_quantizer: bool,
    has_codebooks: bool,
    native_weight_set: bool,
    serialized_config: bool,
) -> str:
    if continuous_analysis_evidence and not has_quantizer and not has_codebooks:
        return (
            "BLOCKED: no complete RVQ/VQ quantizer and codebook weights were found; "
            "the encoder, Gaussian posterior heads, and flow weights support a continuous "
            "Flow-VAE analysis path, not Music 3 token IDs."
        )
    if native_weight_set and serialized_config:
        return (
            "BLOCKED: tokenizer-like candidate weights and config were found, but no exact "
            "compatible executable Music 3 tokenizer architecture/API has been implemented and verified."
        )
    missing = []
    if not native_weight_set:
        missing.append("a complete candidate encoder/quantizer/codebook weight set")
    if not serialized_config:
        missing.append("serialized tokenizer architecture/config")
    return "BLOCKED: native Music 3 token encoding is unavailable; missing " + ", ".join(missing) + "."


def format_human(report: Mapping[str, Any]) -> str:
    checkpoint = report["checkpoint"]
    state = report["state_dict"]
    lines = [
        f"Checkpoint: {checkpoint['path']}",
        f"Safe load: {checkpoint['safe_load']}",
        f"Root: {checkpoint['root_type']}; top-level entries: {report['top_level']['count']}",
        f"Selected state dict: {report['state_dict_path']}",
        f"Tensors: {state['tensor_count']} ({state['total_numel']:,} values, {state['tensor_bytes']:,} bytes)",
        "Prefix groups:",
    ]
    if "sha256" in checkpoint:
        lines.insert(2, f"SHA-256: {checkpoint['sha256']}")
    top_level = report["top_level"]
    rendered_keys = ", ".join(top_level["examples"])
    if top_level["truncated"]:
        rendered_keys += ", ..."
    lines.insert(-1, f"Top-level key examples: {rendered_keys or '(non-mapping root)'}")
    for name, group in report["prefix_groups"].items():
        examples = "; ".join(
            f"{item['name']} {item['shape']} {item['dtype']}" for item in group["examples"]
        )
        lines.append(f"  {name}: {group['tensor_count']} tensor(s); {examples}")
    lines.append("Requested-name search:")
    for name, search in report["searches"].items():
        lines.append(f"  {name}: {'present' if search['present'] else 'absent'} ({search['tensor_count']} tensors)")
    lines.append("Capabilities:")
    for name, available in report["capabilities"].items():
        lines.append(f"  {name}: {'YES' if available else 'NO'}")
    lines.append(f"Verdict: {report['verdict']}")
    return "\n".join(lines)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("checkpoint", help="Path to dav.pth")
    parser.add_argument("--json", action="store_true", help="Emit machine-readable JSON")
    parser.add_argument("--sha256", action="store_true", help="Also hash the checkpoint file")
    return parser


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    try:
        report = inspect_checkpoint(args.checkpoint, include_sha256=args.sha256)
    except (FileNotFoundError, RuntimeError, ValueError) as error:
        if args.json:
            print(json.dumps({"status": "ERROR", "error": str(error)}, indent=2))
        else:
            print(f"ERROR: {error}")
        return 2
    print(json.dumps(report, indent=2) if args.json else format_human(report))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())