comfyui-mcp/dist/services/train-caption.js

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),

// LLM captioning for dataset images — one vision turn per image through the
// user's OWN Claude subscription (the same Agent SDK the panel agent uses),
// NOT a paid API. Deliberately a direct `query()` call, not the panel's
// ClaudeBackend: the panel persona/session machinery would bleed chat-flavored
// prose into what must be a bare caption string.
//
// Images come from readTrainingFile (contained to the training root, ≤2MB) so
// dataset files never touch the /view channel.
import { assertDatasetEditable, getDataset, updateDataset } from "./training-datasets.js";
import { buildAgentSpawnEnv } from "./panel-secrets.js";
import { trainingRoot } from "./training-jobs.js";
import { readFileSync, realpathSync, statSync } from "node:fs";
import { extname, resolve, sep } from "node:path";
/** Captioning reads staged TRAINING images, not phone-bound display thumbs —
 *  the 2MB action:"file" cap is a display budget, not a property of the data
 *  (codex finding: a valid 8MP training PNG must still caption). 10MB keeps
 *  the vision context sane without rejecting real sets. Contained + image-only
 *  like readTrainingFile. */
const MAX_CAPTION_BYTES = 10 * 1024 * 1024;
const CAPTION_IMAGE_EXTS = new Set([".png", ".jpg", ".jpeg", ".webp"]);
function readCaptionFile(absPath) {
    let p;
    try {
        p = realpathSync(absPath);
    }
    catch {
        p = resolve(absPath);
    }
    const root = realpathSync(trainingRoot());
    const norm = (x) => (process.platform === "win32" ? x.toLowerCase() : x);
    if (norm(p) !== norm(root) && !norm(p).startsWith(norm(root) + sep)) {
        throw new Error("path escapes the training root");
    }
    const ext = extname(p).toLowerCase();
    if (!CAPTION_IMAGE_EXTS.has(ext)) {
        throw new Error(`only image files can be captioned (${[...CAPTION_IMAGE_EXTS].join("/")})`);
    }
    const st = statSync(p);
    if (st.size > MAX_CAPTION_BYTES) {
        throw new Error(`image too large to caption (${(st.size / 1024 / 1024).toFixed(1)} MB > 10 MB)`);
    }
    const mimeType = ext === ".png" ? "image/png" : ext === ".webp" ? "image/webp" : "image/jpeg";
    return { data: readFileSync(p).toString("base64"), mimeType };
}
// Lazy SDK import (same pattern as claude-backend — the SDK is a heavy optional
// dependency that must not load at MCP boot).
let queryFn = null;
async function getQuery() {
    if (!queryFn) {
        const mod = await import("@anthropic-ai/claude-agent-sdk");
        queryFn = mod.query;
    }
    return queryFn;
}
/** Model used for captioning — cheap and fast by default (this is a bulk
 *  utility call on the user's subscription). Override via COMFYUI_MCP_CAPTION_MODEL. */
const CAPTION_MODEL = process.env.COMFYUI_MCP_CAPTION_MODEL?.trim() || "claude-haiku-4-5";
/** A PERSISTENT failure that dooms every remaining image — a missing/invalid
 *  Claude Code login, or a backend that can't run vision captioning at all.
 *  captionDataset bails the batch on the first one instead of re-hitting the
 *  same auth wall 32 times (issue #438: the panel backend was Codex, so every
 *  image returned "Not logged in — Please run /login"). */
export class CaptionAuthError extends Error {
    constructor(message) {
        super(message);
        this.name = "CaptionAuthError";
    }
}
/** Does this SDK error text mean "no usable Claude Code credentials", i.e. a
 *  failure that will repeat identically for every image? Matches the Agent
 *  SDK's not-logged-in / invalid-key / auth phrasings. */
function isAuthFailureText(text) {
    const t = text.toLowerCase();
    return (t.includes("not logged in") ||
        t.includes("/login") ||
        t.includes("please run /login") ||
        t.includes("invalid api key") ||
        t.includes("invalid x-api-key") ||
        t.includes("authentication_error") ||
        t.includes("authentication error") ||
        t.includes("oauth token has expired") ||
        t.includes("unauthorized") ||
        (t.includes("credit balance") && t.includes("too low")));
}
function captionPrompt(opts) {
    const lines = [
        "Write ONE caption for this image for LoRA training. Rules:",
        "- Describe what CHANGES between training images: pose, clothing, background, lighting, expression, framing — NOT the subject's identity (the model learns that from the set, not the captions).",
        "- One or two short comma-separated phrases. No prefix, no quotes, no commentary.",
    ];
    if (opts.trigger)
        lines.push(`- Start the caption EXACTLY with "${opts.trigger}, " (the set's trigger word).`);
    if (opts.guide)
        lines.push(`- User guidance: ${opts.guide}`);
    lines.push("Reply with ONLY the caption.");
    return lines.join("\n");
}
/** Caption ONE image (path under the training root). Returns the bare caption. */
export async function captionImage(opts) {
    const { data, mimeType } = readCaptionFile(opts.imagePath);
    const query = await getQuery();
    const userContent = [
        {
            type: "image",
            source: {
                type: "base64",
                media_type: mimeType,
                data,
            },
        },
        { type: "text", text: captionPrompt(opts) },
    ];
    async function* turns() {
        yield {
            type: "user",
            message: { role: "user", content: userContent },
            parent_tool_use_id: null,
        };
    }
    // No tools, single turn — a utility call, not an agent session. The spawn
    // env is SCRUBBED of tool-only secrets (RunPod/CivitAI/HF keys) — the
    // caption subprocess has no business seeing them (codex finding; same
    // invariant as ClaudeBackend/ai-proposer).
    const q = query({
        prompt: turns(),
        options: { model: CAPTION_MODEL, maxTurns: 1, tools: [], env: buildAgentSpawnEnv() },
    });
    let text = "";
    for await (const msg of q) {
        const m = msg;
        if (m.type === "assistant" && m.message?.content) {
            text += m.message.content
                .filter((c) => c.type === "text")
                .map((c) => c.text ?? "")
                .join("");
        }
        else if (m.type === "result" && (m.is_error || (m.subtype && m.subtype !== "success"))) {
            // The SDK reports a failed run (e.g. "Not logged in — Please run /login")
            // as an error `result` message, NOT a thrown exception — previously this
            // was silently dropped and every image failed with the generic "empty
            // caption" message. Pull the text from `errors[]` (the error variant's
            // field), and flag auth failures so the batch bails.
            const detail = (Array.isArray(m.errors) ? m.errors.join("; ") : "").trim() ||
                (m.result ?? "").trim() ||
                `caption run failed (${m.subtype ?? "error"})`;
            if (isAuthFailureText(detail)) {
                throw new CaptionAuthError(`Claude captioning is unavailable — the Claude Code session is not authenticated ` +
                    `(SDK said: ${detail}). Run \`claude /login\`, or set ANTHROPIC_API_KEY, then retry. ` +
                    `Captioning always runs through Claude regardless of the panel's active backend.`);
            }
            throw new Error(detail);
        }
    }
    const caption = text.trim().replace(/^["']|["']$/g, "");
    if (!caption)
        throw new Error("the model returned an empty caption");
    return caption;
}
/** Caption a whole staged dataset (or a subset) and WRITE the captions into
 *  its .txt files. Sequential on purpose — a subscription isn't a rate pool. */
export async function captionDataset(name, opts = {}) {
    const detail = getDataset(name);
    // In-use check BEFORE spending subscription turns: a dataset a running job
    // trains from would reject every write anyway (codex finding — the loop
    // used to spend one vision turn per image and write zero captions).
    await assertDatasetEditable(detail.datasetPath, "captioning");
    const only = opts.only?.length ? new Set(opts.only) : null;
    const targets = detail.items.filter((it) => !only || only.has(it.file));
    if (!targets.length)
        throw new Error(only ? "none of the requested files are in this dataset" : "dataset has no images");
    const failed = [];
    let captioned = 0;
    for (const it of targets) {
        try {
            const caption = await captionImage({
                imagePath: `${detail.datasetPath}/${it.file}`,
                guide: opts.guide,
                trigger: opts.trigger,
            });
            await updateDataset(name, { setCaptions: { [it.file]: caption } });
            captioned++;
        }
        catch (err) {
            // A PERSISTENT auth/credential failure dooms every remaining image — the
            // whole batch runs on one Claude session (issue #438). Fail fast with an
            // actionable error instead of iterating N files re-hitting the same wall.
            // Per-file transient errors (bad image, empty caption, write conflict)
            // still just get recorded and the loop continues.
            if (err instanceof CaptionAuthError ||
                (err instanceof Error && isAuthFailureText(err.message))) {
                const message = err instanceof CaptionAuthError ? err.message : err.message;
                throw new CaptionAuthError(message);
            }
            failed.push({ file: it.file, error: err instanceof Error ? err.message : String(err) });
        }
    }
    return { captioned, failed };
}
//# sourceMappingURL=train-caption.js.map