comfyui-mcp
Version:
Local-first, agent-native control plane for ComfyUI — MCP server + autonomous sidebar agent that drives your live graph in natural language on ANY LLM: Claude/ChatGPT/Gemini on your subscription (no API key), free local models via Ollama (fully offline),
217 lines • 8.66 kB
JavaScript
/**
* txt2img capability probing for checkpoint auto-select (issue #892).
*
* `generate_image`'s auto-select used to take the alphabetically-first entry
* of the checkpoints listing — which fails 100% of the time when that entry
* is a video model (LTX, Wan, HunyuanVideo), because the txt2img graph wires
* CheckpointLoaderSimple's CLIP output into CLIPTextEncode and a video
* checkpoint contains no text encoder weights (clip arrives as None).
*
* ComfyUI exposes no "does this checkpoint contain a CLIP" endpoint, so the
* probe reads the safetensors header (first bytes only, never the multi-GB
* tensor payload). A header WITH known text-encoder tensor prefixes is
* "capable"; a header with a recognized diffusion/VAE layout but no text
* encoder is "incapable" (an observed absence, e.g. the video models this
* fix targets); anything else — remote servers, non-safetensors formats,
* unreadable files, unrecognized layouts — is "unknown", never folded into
* "incapable". Selection prefers known-capable, falls back to unknown, and
* refuses only when EVERY candidate is known incapable (where enqueueing
* would be a guaranteed CLIPTextEncode failure).
*/
import { open, stat } from "node:fs/promises";
import { isAbsolute } from "node:path";
import { isRemoteMode } from "../config.js";
import { resolveExistingModelFile, } from "./model-resolver.js";
import { logger } from "../utils/logger.js";
/** Same cap as model-explorer's header reader: far beyond any real checkpoint
* header (KBs–low MBs); guards against a corrupt length prefix making us
* allocate the file's whole size. */
const SAFETENSORS_HEADER_MAX_BYTES = 64 * 1024 * 1024;
/**
* Tensor-key prefixes that mark embedded text-encoder weights. Covers the
* single-file image checkpoint families CheckpointLoaderSimple can serve CLIP
* from: SD1.x/SD2.x ("cond_stage_model."), SDXL ("conditioner."), SD3/Flux
* ("text_encoders."), PixArt/Hunyuan-DiT ("text_encoder."). Video checkpoints
* (LTX, Wan, HunyuanVideo) carry transformer+VAE only and match none of these.
*/
const TEXT_ENCODER_KEY_PREFIXES = [
"cond_stage_model.",
"conditioner.",
"text_encoders.",
"text_encoder.",
];
/**
* Tensor-key prefixes that positively identify a recognized diffusion/VAE
* checkpoint layout. "incapable" is only reported when the header shows one
* of these layouts AND no text encoder — an observed absence inside a
* recognized layout, not an unrecognized layout. A header matching NEITHER
* list is "unknown" (could not determine), never "incapable".
*/
const DIFFUSION_LAYOUT_KEY_PREFIXES = [
"model.diffusion_model.",
"diffusion_model.",
"double_blocks.",
"single_blocks.",
"first_stage_model.",
"vae.",
];
export function headerKeysHaveTextEncoder(keys) {
for (const key of keys) {
for (const prefix of TEXT_ENCODER_KEY_PREFIXES) {
if (key.startsWith(prefix))
return true;
}
}
return false;
}
/** True when the header keys show a recognized diffusion/VAE checkpoint layout. */
export function headerKeysHaveDiffusionLayout(keys) {
for (const key of keys) {
for (const prefix of DIFFUSION_LAYOUT_KEY_PREFIXES) {
if (key.startsWith(prefix))
return true;
}
}
return false;
}
/**
* Read just a safetensors file's header (8-byte little-endian length + that
* many bytes of JSON) and return its top-level keys. Returns null for
* anything that isn't a well-formed, capped header — the probe is
* best-effort and never throws.
*/
async function readSafetensorsHeaderKeys(filePath) {
let fh;
try {
fh = await open(filePath, "r");
const lenBuf = Buffer.alloc(8);
if ((await fh.read(lenBuf, 0, 8, 0)).bytesRead < 8)
return null;
const headerLen = Number(lenBuf.readBigUInt64LE(0));
if (!Number.isSafeInteger(headerLen) ||
headerLen <= 0 ||
headerLen > SAFETENSORS_HEADER_MAX_BYTES) {
return null;
}
const headerBuf = Buffer.alloc(headerLen);
if ((await fh.read(headerBuf, 0, headerLen, 8)).bytesRead < headerLen) {
return null;
}
const header = JSON.parse(headerBuf.toString("utf8"));
if (!header || typeof header !== "object" || Array.isArray(header)) {
return null;
}
return Object.keys(header);
}
catch {
return null;
}
finally {
await fh?.close().catch(() => undefined);
}
}
/** Header reads are cheap but not free; invalidated by mtime+size change. */
const probeCache = new Map();
/** Test hook: drop all cached probe results. */
export function resetCheckpointCapabilityCache() {
probeCache.clear();
}
/**
* Probe one listed checkpoint for txt2img capability. Never throws: any
* observation failure resolves to "unknown", never to a definite negative.
*/
async function probeCheckpoint(model) {
const lower = model.name.toLowerCase();
// CheckpointLoaderSimple cannot load GGUF at all (needs UnetLoaderGGUF).
if (lower.endsWith(".gguf"))
return "incapable";
// .ckpt/.pt/.bin have no cheap metadata header to inspect.
if (!lower.endsWith(".safetensors"))
return "unknown";
// The header lives on the server's filesystem; in remote mode that is not
// ours to read (config.ts warns when COMFYUI_PATH points elsewhere).
if (isRemoteMode())
return "unknown";
let filePath;
let mtimeMs;
let size;
try {
if (isAbsolute(model.path)) {
const info = await stat(model.path);
if (!info.isFile())
return "unknown";
filePath = model.path;
mtimeMs = info.mtimeMs;
size = info.size;
}
else {
// HTTP listing: path is "<dir>/<name>" relative to models/. Resolve via
// the shared resolver so extra_model_paths.yaml roots are honoured.
const resolved = await resolveExistingModelFile(`checkpoints/${model.name}`);
if (!resolved.info.isFile())
return "unknown";
filePath = resolved.path;
mtimeMs = resolved.info.mtimeMs;
size = resolved.info.size;
}
}
catch {
return "unknown";
}
const cached = probeCache.get(filePath);
if (cached && cached.mtimeMs === mtimeMs && cached.size === size) {
return cached.capability;
}
const keys = await readSafetensorsHeaderKeys(filePath);
const capability = keys === null
? "unknown"
: headerKeysHaveTextEncoder(keys)
? "capable"
: // No text encoder — but only a RECOGNIZED diffusion/VAE layout makes
// that an observed absence; an unrecognized layout is "unknown".
headerKeysHaveDiffusionLayout(keys)
? "incapable"
: "unknown";
probeCache.set(filePath, { mtimeMs, size, capability });
return capability;
}
/**
* Pick a checkpoint that can drive the txt2img graph: the first known-capable
* candidate in listing order, else the first whose capability could not be
* determined (status-quo behaviour — the graph may still work). Returns null
* only when every candidate is known incapable, where enqueueing would be a
* guaranteed failure.
*/
export async function selectTxt2ImgCheckpoint(models) {
// Sequential with early exit on the first known-capable candidate: the
// common case (first checkpoint is a normal image model) costs one header
// read. An unknown candidate is only a fallback — a LATER known-capable
// entry still beats an earlier unknown one.
const incapable = [];
let firstUnknown;
for (const model of models) {
const capability = await probeCheckpoint(model);
if (capability === "capable") {
if (incapable.length > 0) {
logger.debug("checkpoint auto-select skipped non-txt2img checkpoints", {
picked: model.name,
skipped: incapable,
});
}
return { name: model.name, capability, skippedIncapable: incapable };
}
if (capability === "incapable")
incapable.push(model.name);
else if (!firstUnknown)
firstUnknown = model;
}
if (firstUnknown) {
return {
name: firstUnknown.name,
capability: "unknown",
skippedIncapable: incapable,
};
}
return null;
}
//# sourceMappingURL=checkpoint-capability.js.map