DemonGatanjieuandClaude Opus 5.5
f65f12152b
Checked against the headers of 315 models in ComfyUI-Manager's list: a Hunyuan-DiT checkpoint and LBM were read as VAEs, SUPIR as a ControlNet, XLabs' Flux LoRAs and IP-Adapter as diffusion models, a whole CLIP as a text encoder only; also PhotoMaker, LoRAs of weight differences or OFT, and checkpoints without a VAE. Files two loaders read from different folders (SUPIR, a whole CLIP) are not sure. A lone drawing model in checkpoints loads there (the checkpoint loader reads it as a diffusion model). Newly told: diffusers UNets, Stable Cascade, SDXL T2I-Adapters, RCAN/SPAN/compact upscalers. 301 of 315 now agree. "Looks like no drawing model" said what it could not tell: now "Kind not told". Kept header answers carry the rules version and are read again when it changes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
354 lines
17 KiB
Python
354 lines
17 KiB
Python
"""What a model file is, from its header alone: which models folder it belongs in (LoRA,
|
|
checkpoint, VAE…) and its base model. Reads a .safetensors header (JSON) or a .gguf header
|
|
(key/value metadata and tensor names); never loads tensors and never opens formats that can
|
|
carry code (.ckpt/.pt/.pth/.bin). Shared by model import (api/model_import.py) and the
|
|
scanner's offline base model (model_identity.infer_base_model_from_header). No imports from
|
|
ComfyUI, so it is tested on its own.
|
|
|
|
`classify(fmt, header)` -> {"kind": a ComfyUI folder name or "", "base": base model or "",
|
|
"sure": True when the header names both outright, False when it is a reasoned guess, "form":
|
|
"lllite" for an Anima LLLite model patch (a node pack reads those from controlnet), else ""}.
|
|
`diffusion_params(path)` -> how many parameters the file's drawing model has (the tensor shapes
|
|
the header lists), the yardstick for how well a model draws hard things such as hands.
|
|
"""
|
|
|
|
import json
|
|
import struct
|
|
|
|
MAX_HEADER = 100 * 1024 * 1024
|
|
GGUF_READ = 16 * 1024 * 1024
|
|
# Raised whenever a rule changes: answers kept on disk (api/model_heads.py) are then read again.
|
|
RULES = 2
|
|
|
|
# Cross-attention width (the text encoder's output) -> base model, for the SD family.
|
|
CONTEXT_WIDTH = {768: "SD 1.5", 1024: "SD 2.1", 2048: "SDXL"}
|
|
# ComfyUI-GGUF's architecture names.
|
|
GGUF_DIFFUSION = {
|
|
"flux": "Flux.1 D", "sd1": "SD 1.5", "sdxl": "SDXL", "sd3": "SD3", "aura": "AuraFlow",
|
|
"hidream": "HiDream", "cosmos": "Cosmos", "ltxv": "LTXV", "hyvid": "Hunyuan Video",
|
|
"wan": "Wan Video", "lumina2": "Lumina", "qwen_image": "Qwen", "chroma": "Chroma",
|
|
}
|
|
GGUF_TEXT = ("t5", "t5encoder", "umt5", "llama", "qwen2", "qwen2vl", "qwen3", "gemma", "clip", "mistral")
|
|
EMBEDDING_KEYS = {"emb_params", "clip_l", "clip_g", "string_to_param", "string_to_token", "name", "step",
|
|
"sd_checkpoint", "sd_checkpoint_name"}
|
|
LORA_MARKS = ("lora_up.", "lora_down.", "lora_a.", "lora_b.", "lora.up.", "lora.down.", "hada_w1_a",
|
|
"lokr_w1", ".lora_linear_layer.", "dora_scale", "oft_blocks", "oft_diag", "boft_blocks",
|
|
".proj_lora", ".qkv_lora") # the last two: XLabs' Flux LoRAs
|
|
# kohya's names (any LoRA kind, GLoRA and weight differences too); ".diff"/".diff_b": full-weight differences.
|
|
LORA_PREFIXES = ("lora_unet_", "lora_te_", "lora_te1_", "lora_te2_", "lora_te3_")
|
|
LORA_ENDS = (".diff", ".diff_b")
|
|
CONTROL_MARKS = ("control_model_", "input_hint_block", "controlnet_cond_embedding", "controlnet_blocks",
|
|
"controlnet_x_embedder", "zero_convs", "controlnet_mid_block")
|
|
# Tensors only a model patch has (ComfyUI's "Load Model Patch" node, models/model_patches): Anima
|
|
# LLLite, the Qwen-Image block-wise and Z-Image Fun ControlNets, SigLIP features, Wan's Uni3C. Some
|
|
# are called ControlNets, but the ControlNet loader cannot read them.
|
|
PATCH_KEYS = {"lllite_conditioning1.conv1.weight", "controlnet_blocks.0.y_rms.weight", "feature_embedder.mid_layer_norm.bias",
|
|
"control_all_x_embedder.2-1.weight", "controlnet_patch_embedding.weight"}
|
|
# Image-prompt layers in a model's blocks: the Anima IP-Adapter (beside LoRA layers of its own),
|
|
# XLabs' Flux IP-Adapter.
|
|
IPADAPTER_MARKS = (".ip_k_proj.", ".adaln_ip.", "ip_adapter_")
|
|
# The parts of a whole model besides its drawing part (a checkpoint's VAE and text encoders).
|
|
VAE_PARTS = ("first_stage_model.", "vae.")
|
|
TEXT_PARTS = ("cond_stage_model.", "conditioner.", "text_encoders.")
|
|
DIFFUSION_PREFIXES = ("model.diffusion_model.", "diffusion_model.", "model.model.")
|
|
DIFFUSION_STARTS = ("double_blocks.", "single_blocks.", "joint_blocks.", "transformer_blocks.",
|
|
"single_transformer_blocks.", "input_blocks.", "blocks.0.self_attn", "net.blocks.", "net.x_embedder.",
|
|
"noise_refiner.", "context_refiner.", "double_stream_blocks.", # Lumina 2 / Z-Image, HiDream
|
|
"down_downscalers.") # Stable Cascade's stages B and C
|
|
|
|
|
|
def parse_safetensors(data):
|
|
"""The header dict from a file's first bytes (8-byte length + JSON); None if not one."""
|
|
if len(data) < 8:
|
|
return None
|
|
size = struct.unpack("<Q", data[:8])[0]
|
|
if not 2 <= size <= MAX_HEADER or len(data) < 8 + size:
|
|
return None
|
|
try:
|
|
header = json.loads(data[8:8 + size].decode("utf-8"))
|
|
except (UnicodeDecodeError, ValueError):
|
|
return None
|
|
return header if isinstance(header, dict) else None
|
|
|
|
|
|
class _Reader:
|
|
def __init__(self, data):
|
|
self.data, self.at = data, 0
|
|
|
|
def take(self, count):
|
|
if self.at + count > len(self.data):
|
|
raise EOFError
|
|
part = self.data[self.at:self.at + count]
|
|
self.at += count
|
|
return part
|
|
|
|
def number(self, fmt):
|
|
return struct.unpack("<" + fmt, self.take(struct.calcsize(fmt)))[0]
|
|
|
|
def text(self):
|
|
return self.take(self.number("Q")).decode("utf-8", "replace")
|
|
|
|
|
|
_GGUF_SCALARS = {0: "B", 1: "b", 2: "H", 3: "h", 4: "I", 5: "i", 6: "f", 7: "?", 10: "Q", 11: "q", 12: "d"}
|
|
|
|
|
|
def _gguf_value(reader, kind):
|
|
if kind in _GGUF_SCALARS:
|
|
return reader.number(_GGUF_SCALARS[kind])
|
|
if kind == 8:
|
|
return reader.text()
|
|
if kind == 9: # an array: skipped (tokenizer word lists are long)
|
|
item, count = reader.number("I"), reader.number("Q")
|
|
for _ in range(count):
|
|
_gguf_value(reader, item)
|
|
return None
|
|
raise ValueError("unknown gguf value type")
|
|
|
|
|
|
def parse_gguf(data):
|
|
"""{"meta": general.* values, "keys": tensor names, "params": their parameters, 0 when not all
|
|
were read} from a .gguf's first bytes; None if not one. What did not fit in `data` is left out."""
|
|
if data[:4] != b"GGUF":
|
|
return None
|
|
reader = _Reader(data)
|
|
reader.take(4)
|
|
meta, keys, params = {}, [], 0
|
|
try:
|
|
if reader.number("I") < 2:
|
|
return None
|
|
tensors, values = reader.number("Q"), reader.number("Q")
|
|
for _ in range(values):
|
|
key = reader.text()
|
|
value = _gguf_value(reader, reader.number("I"))
|
|
if key.startswith("general.") and isinstance(value, (str, int, float)):
|
|
meta[key] = value
|
|
for _ in range(tensors):
|
|
keys.append(reader.text())
|
|
size = 1
|
|
for _ in range(reader.number("I")):
|
|
size *= reader.number("Q")
|
|
params += size
|
|
reader.take(4 + 8) # its type and offset
|
|
except (EOFError, ValueError, struct.error):
|
|
params = 0
|
|
return {"meta": meta, "keys": keys, "params": params}
|
|
|
|
|
|
def read_head(path):
|
|
"""(fmt, header) of a model file: ("safetensors", dict), ("gguf", dict), or (None, None)
|
|
for other formats and unreadable files."""
|
|
lower = path.lower()
|
|
try:
|
|
with open(path, "rb") as source:
|
|
if lower.endswith((".safetensors", ".sft")):
|
|
start = source.read(8)
|
|
if len(start) < 8:
|
|
return None, None
|
|
size = struct.unpack("<Q", start)[0]
|
|
if size > MAX_HEADER:
|
|
return None, None
|
|
return "safetensors", parse_safetensors(start + source.read(size))
|
|
if lower.endswith(".gguf"):
|
|
return "gguf", parse_gguf(source.read(GGUF_READ))
|
|
except OSError:
|
|
pass
|
|
return None, None
|
|
|
|
|
|
def _context_width(header):
|
|
"""The cross-attention input width a SD-family file shows (768 / 1024 / 2048), else 0."""
|
|
for key, value in header.items():
|
|
low = key.lower().replace(".", "_")
|
|
if "attn2_to_k" not in low or not isinstance(value, dict):
|
|
continue
|
|
shape = value.get("shape") or []
|
|
if len(shape) == 2 and ("lora_down" in low or "lora_a" in low or low.endswith("to_k_weight")):
|
|
return shape[1]
|
|
return 0
|
|
|
|
|
|
def _family(meta):
|
|
"""An SDXL file's family from the name of the model it was trained on (kohya's metadata)."""
|
|
trained_on = " ".join(str(meta.get(key, "")) for key in ("ss_sd_model_name", "ss_base_model", "modelspec.title")).lower()
|
|
for mark, name in (("pony", "Pony"), ("illustrious", "Illustrious"), ("noob", "NoobAI")):
|
|
if mark in trained_on:
|
|
return name
|
|
return ""
|
|
|
|
|
|
def _base_from_meta(meta):
|
|
arch = str(meta.get("modelspec.architecture", "")).lower()
|
|
trained = str(meta.get("ss_base_model_version", "")).lower()
|
|
pairs = (("flux", "Flux.1 D"), ("chroma", "Chroma"), ("stable-diffusion-xl", "SDXL"), ("sdxl", "SDXL"),
|
|
("stable-diffusion-v1", "SD 1.5"), ("sd_v1", "SD 1.5"), ("stable-diffusion-v2", "SD 2.1"),
|
|
("sd_v2", "SD 2.1"), ("stable-diffusion-3", "SD3"), ("sd3", "SD3"), ("wan", "Wan Video"),
|
|
("qwen", "Qwen"), ("hunyuan", "Hunyuan Video"))
|
|
for mark, name in pairs:
|
|
if (arch and mark in arch) or (trained and trained.startswith(mark)):
|
|
return name
|
|
return ""
|
|
|
|
|
|
def _base_from_keys(header, text):
|
|
"""(base, sure) from tensor names; ("", False) when they do not tell."""
|
|
if "individual_token_refiner" in text:
|
|
return "Hunyuan Video", True
|
|
# Cosmos-Predict2 (with or without "net."); Anima adds an adapter for its Qwen3 text
|
|
# encoder. Before Wan: their blocks' attention names look alike.
|
|
if "llm_adapter" in text:
|
|
return "Anima", True
|
|
if "adaln_modulation_self_attn" in text:
|
|
return "Cosmos", True
|
|
if "distilled_guidance_layer" in text:
|
|
return "Chroma", True
|
|
if "img_mod_1" in text or "txt_norm" in text or "img_mlp_net" in text:
|
|
return "Qwen", True
|
|
if "double_blocks_0_img_attn" in text or "single_blocks_0_linear1" in text:
|
|
return "Flux.1 D", True
|
|
if "single_transformer_blocks_0_attn" in text or ("single_transformer_blocks" in text and "x_embedder" in text):
|
|
return "Flux.1 D", False
|
|
if "joint_blocks_0_x_block" in text:
|
|
return "SD3", True
|
|
if "blocks_0_self_attn_q" in text and "blocks_0_cross_attn_k" in text:
|
|
return "Wan Video", True
|
|
if any(mark in text for mark in ("conditioner_embedders_1", "label_emb_0_0", "lora_te1_", "lora_te2_",
|
|
"lora_unet_input_blocks", "lora_unet_output_blocks")):
|
|
return "SDXL", True
|
|
width = _context_width(header)
|
|
if width in CONTEXT_WIDTH:
|
|
return CONTEXT_WIDTH[width], True
|
|
if "cond_stage_model_transformer_text_model" in text:
|
|
return "SD 1.5", True
|
|
if "input_blocks_1_1_transformer_blocks" in text or "lora_unet_down_blocks" in text:
|
|
return "SD 1.5", False
|
|
return "", False
|
|
|
|
|
|
def _embedding_base(header):
|
|
if "clip_g" in header:
|
|
return "SDXL"
|
|
shape = (header.get("emb_params") or {}).get("shape") or []
|
|
return CONTEXT_WIDTH.get(shape[-1], "") if shape else ""
|
|
|
|
|
|
def _kind(keys, raw, text):
|
|
"""The models folder from tensor names: (kind, sure)."""
|
|
names = set(keys)
|
|
if names and names <= EMBEDDING_KEYS and names & {"emb_params", "clip_l", "clip_g", "string_to_param"}:
|
|
return "embeddings", True
|
|
if names & PATCH_KEYS:
|
|
return "model_patches", True
|
|
# SUPIR: SDXL with a control model. ComfyUI's "Load Model Patch" reads it from model_patches,
|
|
# the ComfyUI-SUPIR pack from checkpoints: no folder is sure.
|
|
if any(key.startswith("model.control_model.") for key in keys):
|
|
return "model_patches", False
|
|
if any(mark in text for mark in CONTROL_MARKS):
|
|
return "controlnet", True
|
|
diffusion = any(key.startswith(DIFFUSION_PREFIXES + DIFFUSION_STARTS) for key in keys)
|
|
vae = any(key.startswith(VAE_PARTS) for key in keys)
|
|
encoders = any(key.startswith(TEXT_PARTS) for key in keys)
|
|
# A whole model first: some merged checkpoints still carry a LoRA's keys. Its drawing part may
|
|
# have names not listed here (Hunyuan-DiT's "model.blocks"): a VAE and text encoders tell then.
|
|
if vae and (diffusion or encoders):
|
|
return "checkpoints", True
|
|
if any(key.startswith("id_encoder.") for key in keys): # PhotoMaker (its own folder), LoRA layers inside
|
|
return "", False
|
|
# Before LoRA: some IP-Adapters carry LoRA layers too.
|
|
if any(key.startswith(("image_proj.", "ip_adapter.")) for key in keys) or any(mark in raw for mark in IPADAPTER_MARKS):
|
|
return "ipadapter", True
|
|
if any(mark in raw for mark in LORA_MARKS) or any(key.startswith(LORA_PREFIXES) or key.endswith(LORA_ENDS) for key in keys):
|
|
return "loras", True
|
|
# A checkpoint saved without its VAE: the checkpoint loader reads it all the same.
|
|
if diffusion and encoders:
|
|
return "checkpoints", True
|
|
if diffusion:
|
|
return "diffusion_models", True
|
|
if any(key.startswith("adapter.body.") for key in keys): # a T2I-Adapter (diffusers)
|
|
return "controlnet", True
|
|
# A UNet saved by diffusers (IC-Light, Kolors, Lotus): its own blocks at the top, no VAE's
|
|
# "encoder."; AnimateDiff's motion modules have the same block names.
|
|
if "motion_modules" not in raw and any(key.startswith("mid_block.") for key in keys) and any(key.startswith(("time_embedding.", "class_embedding.")) for key in keys):
|
|
return "diffusion_models", True
|
|
# A VAE inside a model of another kind (LBM's relighting model): not a VAE file.
|
|
if vae and any(not key.startswith(VAE_PARTS) for key in keys):
|
|
return "", False
|
|
if any(mark in text for mark in ("encoder_down_0_block", "decoder_up_0_block", "encoder_down_blocks_0",
|
|
"encoder_downsamples", "first_stage_model_encoder")):
|
|
return "vae", True
|
|
# A whole CLIP (text and vision): the clip vision loader reads it, so does a text encoder loader.
|
|
if "vision_model_encoder_layers" in text:
|
|
return "clip_vision", "text_model" not in text
|
|
if "text_model_encoder_layers" in text or "encoder_block_0_layer_0_selfattention" in text \
|
|
or ("model_layers_0_self_attn" in text and "embed_tokens" in text):
|
|
return "text_encoders", True
|
|
if "conv_first_weight" in text or "body_0_rdb1" in text or ("model_0_weight" in text and "model_1_sub" in text) or "body_0_body_0_body_0" in text or "block_1_c1_r" in text or (keys and all(key.startswith("body.") and key.count(".") == 2 for key in keys)): # RCAN, SPAN, compact ESRGAN
|
|
return "upscale_models", False
|
|
return "", False
|
|
|
|
|
|
def classify(fmt, header):
|
|
"""What a parsed header says: {"kind", "base", "sure"} (see the module text)."""
|
|
if not header:
|
|
return {"kind": "", "base": "", "sure": False}
|
|
if fmt == "gguf":
|
|
meta, keys = header.get("meta") or {}, header.get("keys") or []
|
|
arch = str(meta.get("general.architecture", "")).lower()
|
|
if arch in GGUF_DIFFUSION:
|
|
return {"kind": "diffusion_models", "base": GGUF_DIFFUSION[arch], "sure": True}
|
|
if arch.startswith(GGUF_TEXT):
|
|
return {"kind": "text_encoders", "base": "", "sure": True}
|
|
header = {key: {} for key in keys}
|
|
meta = header.get("__metadata__") if isinstance(header.get("__metadata__"), dict) else {}
|
|
keys = [key for key in header if key != "__metadata__"]
|
|
raw = "\n".join(keys).lower()
|
|
text = raw.replace(".", "_")
|
|
kind, sure = _kind(keys, raw, text)
|
|
if kind in ("vae", "text_encoders", "clip_vision", "upscale_models", "ipadapter", "model_patches"):
|
|
base, base_sure = "", True # shared by several base models: the type's own folder
|
|
elif kind == "embeddings":
|
|
base, base_sure = _embedding_base(header), True
|
|
else:
|
|
base = _base_from_meta(meta)
|
|
base_sure = bool(base)
|
|
if not base:
|
|
base, base_sure = _base_from_keys(header, text)
|
|
if base == "SDXL":
|
|
base = _family(meta) or base
|
|
form = "lllite" if kind == "model_patches" and any(key.startswith("lllite_") for key in keys) else ""
|
|
return {"kind": kind, "base": base, "sure": bool(kind) and sure and (base_sure or not base), "form": form}
|
|
|
|
|
|
# Parts of a checkpoint that are not its drawing model: text encoders, VAE, a kept EMA copy.
|
|
NOT_DRAWING = ("first_stage_model.", "cond_stage_model.", "conditioner.", "text_encoders.", "vae.", "model_ema.")
|
|
|
|
|
|
def diffusion_params(path):
|
|
"""How many parameters a model file's drawing model has, from its header (a checkpoint's text
|
|
encoders and VAE left out); 0 when the header cannot tell."""
|
|
fmt, header = read_head(path)
|
|
if fmt == "gguf":
|
|
return (header or {}).get("params") or 0
|
|
if not header:
|
|
return 0
|
|
tensors = {key: value for key, value in header.items() if key != "__metadata__" and isinstance(value, dict)}
|
|
if any(key.startswith("model.diffusion_model.") for key in tensors):
|
|
tensors = {key: value for key, value in tensors.items() if key.startswith("model.diffusion_model.")}
|
|
total = 0
|
|
for key, value in tensors.items():
|
|
shape = value.get("shape")
|
|
if key.startswith(NOT_DRAWING) or not isinstance(shape, list):
|
|
continue
|
|
size = 1
|
|
for dim in shape:
|
|
size *= dim if isinstance(dim, int) and dim >= 0 else 0
|
|
total += size
|
|
return total
|
|
|
|
|
|
def inspect_file(path):
|
|
"""classify() of a file on disk, plus its "format" (None when its header cannot be read)."""
|
|
fmt, header = read_head(path)
|
|
found = classify(fmt, header)
|
|
found["format"] = fmt if header else None
|
|
return found
|