olivv-cs-ComfyUI-FunPack/core/probe.py

"""Which model family a checkpoint file is, without loading it.

Core reads the file -- a safetensors header is a length-prefixed JSON object
whose keys are tensor names, readable in kilobytes regardless of whether the
file itself is 2GB or 40GB -- and that is ALL it knows how to do. Whether a
set of tensor names belongs to MiniMax H3, or anything else, is a model
module's own business: asking one to identify itself is the same "detect"
capability every model module can already offer, the same shape as the
`empty_latent` and `decode` capabilities modules/models/minimax_h3 provides.
Core aggregates whichever module claims a file and says plainly when none
does -- it is never the one deciding what a family IS called.

Only .safetensors today: no loader in this pack can read a GGUF file yet, so
detecting one accurately would be groundwork for a format nothing here can
use.
"""

from __future__ import annotations

import json
import struct
from pathlib import Path
from typing import List, Optional

from . import registry as registry_mod

#: A header past this size is not a safetensors file -- reading a wild length
#: prefix as a real one would try to allocate however many gigabytes a
#: corrupt or unrelated file's first eight bytes happen to spell out.
MAX_HEADER_BYTES = 256 * 1024 * 1024


def read_safetensors_keys(path: Path) -> Optional[List[str]]:
    """Tensor names out of a safetensors header, or None when this is not one.

    Layout: an 8-byte little-endian length, then that many bytes of JSON whose
    top-level keys are the tensor names (plus `__metadata__`, dropped here).
    Nothing past the header is read.
    """
    try:
        with open(path, "rb") as f:
            raw = f.read(8)
            if len(raw) < 8:
                return None
            (length,) = struct.unpack("<Q", raw)
            if length <= 0 or length > MAX_HEADER_BYTES:
                return None
            header = json.loads(f.read(length).decode("utf-8", errors="replace"))
    except (OSError, ValueError, json.JSONDecodeError, struct.error):
        return None
    if not isinstance(header, dict):
        return None
    return [k for k in header if k != "__metadata__"]


def resolve_diffusion_model(filename: str) -> Optional[Path]:
    """Where `filename` lives, across the folders a FunUI loader can
    actually read a diffusion model from -- `modules/loaders/diffusion_model`
    offers files from "diffusion_models", `modules/loaders/checkpoint` from
    "checkpoints"; nothing in this app reads a model from a "unet" folder,
    so probing one would risk answering for a file no loader here would ever
    actually load under that name. None when it is nowhere findable -- not
    this function's job to say why.
    """
    try:
        import folder_paths
    except ImportError:                                # not inside ComfyUI
        return None
    for folder in ("diffusion_models", "checkpoints"):
        try:
            found = folder_paths.get_full_path(folder, filename)
        except Exception:                              # noqa: BLE001
            found = None
        if found:
            return Path(found)
    return None


def detect(path: Path, registry=None) -> dict:
    """Inspect one checkpoint file.

    Returns ``{"module", "title", "detected", "reason", "traits"}``. `module`
    is the id of whichever model module claimed it, or None -- and None is a
    real answer, not a failure to try hard enough: it means this install has
    no module that recognises the file (or, said in `reason`, that a module's
    detector raised and could not look), and the caller should say so rather
    than guess a family. `reason` is plain text for a person, always
    populated. `traits` is only ever populated when `detected` is true --
    an undetected file says nothing about what it is, so it must not narrow
    what the UI shows via `/api/modules?traits=`; the caller should treat
    `detected: false` as "do not filter", not as "filter to nothing". The
    same goes for `traits: None`: the module that claimed the file names no
    traits (or failed to), which is unknown, not "has none".
    """
    p = Path(path)
    if not p.is_file():
        return {"module": None, "title": None, "detected": False,
                "reason": f"{p.name}: file not found", "traits": []}
    if p.suffix.lower() != ".safetensors":
        return {"module": None, "title": None, "detected": False,
                "reason": f"{p.name}: only .safetensors can be inspected without loading it",
                "traits": []}
    keys = read_safetensors_keys(p)
    if keys is None:
        return {"module": None, "title": None, "detected": False,
                "reason": f"{p.name}: not a readable safetensors header", "traits": []}

    keyset = set(keys)
    reg = registry or registry_mod.current()
    broken = []
    for spec, claims in reg.providers("detect"):
        try:
            matched = claims(keyset)
        except Exception as exc:                       # noqa: BLE001
            broken.append(spec.id)
            from . import log
            log.once(f"provider:{spec.id}.detect:{type(exc).__name__}", log.WARNING, f"{spec.id}.detect", f"failed while answering -- {type(exc).__name__}: {exc}")
            continue
        if matched:
            traits = None                              # unknown, not empty: a module that never says hides nothing
            probe_traits = spec.provides.get("probe_traits")
            if probe_traits is not None:
                try:
                    traits = list(probe_traits(keyset))
                except Exception as exc:                # noqa: BLE001
                    from . import log
                    log.once(f"provider:{spec.id}.probe_traits:{type(exc).__name__}", log.WARNING, f"{spec.id}.probe_traits", f"failed while answering -- {type(exc).__name__}: {exc}")
            blocks = None
            probe_blocks = spec.provides.get("probe_blocks")
            if probe_blocks is not None:
                try:
                    blocks = probe_blocks(keyset)
                except Exception as exc:                # noqa: BLE001
                    from . import log
                    log.once(f"provider:{spec.id}.probe_blocks:{type(exc).__name__}", log.WARNING, f"{spec.id}.probe_blocks", f"failed while answering -- {type(exc).__name__}: {exc}")
            return {"module": spec.id, "title": spec.title, "detected": True,
                    "reason": f"{p.name}: {spec.title}", "traits": traits, "blocks": blocks}
    return {"module": None, "title": None, "detected": False,
            "reason": (f"{p.name}: could not be checked, {', '.join(broken)} failed while looking at it (see the Log)" if broken
                       else f"{p.name}: no installed model module recognises this file"),
            "traits": []}