comfyui-mcp/dist/services/training-datasets.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),

// Dataset + job-config readers behind the training surfaces (panel Jobs/
// Datasets tabs, the mobile monitor). READ-ONLY, and bounded to the training
// root: dataset names are basename-checked, action:"file" paths are realpath-
// contained, and inlining is size-capped (a phone is the other end).
import { existsSync, readdirSync, readFileSync, realpathSync, statSync, writeFileSync } from "node:fs";
import { basename, extname, join, resolve, sep } from "node:path";
import { parse as parseYaml } from "yaml";
import { datasetsRoot, getJob, trainingRoot } from "./training-jobs.js";
const IMAGE_EXTS = new Set([".png", ".jpg", ".jpeg", ".webp"]);
/** train_prepare_dataset (action:"file") inlining cap — a sample/dataset photo, not a checkpoint. */
const MAX_FILE_BYTES = 2 * 1024 * 1024;
/** Containment on REAL paths — a symlinked dir must not smuggle reads across
 *  the root (same lesson as the ref resolver, #302). Missing paths resolve
 *  lexically and fail later with an honest not-found. */
function realOrLexical(p) {
    try {
        return realpathSync(p);
    }
    catch {
        return resolve(p);
    }
}
function contained(child, root) {
    const c = realOrLexical(child);
    const r = realOrLexical(root);
    const norm = (x) => (process.platform === "win32" ? x.toLowerCase() : x);
    return norm(c) === norm(r) || norm(c).startsWith(norm(r) + sep);
}
/** Resolve a dataset dir by NAME. Rules mirror prepareDataset's sanitizeDirName
 *  (dot-prefixed names like ".portrait" are legal dataset dirs; dots-only
 *  collapses and separators are not). Staging dirs from in-flight/crashed
 *  prepares are internal and never resolve here (codex finding). */
function datasetDir(name) {
    const clean = name.trim();
    if (!clean || clean !== basename(clean) || /^\.+$/.test(clean)) {
        throw new Error(`invalid dataset name "${name}"`);
    }
    if (clean.includes(".staging-")) {
        throw new Error(`"${clean}" is an internal staging dir, not a dataset`);
    }
    const dir = join(datasetsRoot(), clean);
    if (!contained(dir, datasetsRoot())) {
        throw new Error(`dataset "${name}" escapes the datasets root`);
    }
    if (!existsSync(dir))
        throw new Error(`no dataset "${clean}" (looked in ${datasetsRoot()})`);
    return dir;
}
function summarize(dir, name) {
    const files = readdirSync(dir);
    const images = files.filter((f) => IMAGE_EXTS.has(extname(f).toLowerCase()));
    const captioned = images.filter((f) => existsSync(join(dir, `${basename(f, extname(f))}.txt`)));
    const modified = Math.max(0, ...files.map((f) => {
        try {
            return statSync(join(dir, f)).mtimeMs;
        }
        catch {
            return 0;
        }
    }));
    return {
        name,
        imageCount: images.length,
        captionedCount: captioned.length,
        modified: new Date(modified || 0).toISOString(),
    };
}
/** Every staged dataset, newest first. Internal staging dirs from in-flight
 *  or crashed prepares (`<name>.staging-<pid>-<rand>`) are NOT datasets — an
 *  abandoned partial set must never be advertised as trainable (codex). */
export function listDatasets() {
    const root = datasetsRoot();
    if (!existsSync(root))
        return [];
    // Newest first; a deterministic name tiebreaker keeps the order STABLE when two
    // datasets share an mtime (coarse filesystem granularity) — otherwise the order
    // follows readdir and varies by platform.
    return readdirSync(root, { withFileTypes: true })
        .filter((d) => d.isDirectory() && !d.name.includes(".staging-"))
        .map((d) => summarize(join(root, d.name), d.name))
        .filter((s) => s.imageCount > 0)
        .sort((a, b) => b.modified.localeCompare(a.modified) || a.name.localeCompare(b.name));
}
/** One dataset's items with their captions (null when uncaptioned). */
export function getDataset(name) {
    const dir = datasetDir(name);
    const files = readdirSync(dir)
        .filter((f) => IMAGE_EXTS.has(extname(f).toLowerCase()))
        .sort();
    const items = files.map((f) => {
        const capPath = join(dir, `${basename(f, extname(f))}.txt`);
        let caption = null;
        try {
            caption = readFileSync(capPath, "utf-8").trim() || null;
        }
        catch {
            /* uncaptioned */
        }
        return { file: f, caption };
    });
    return { ...summarize(dir, name), datasetPath: dir, items };
}
// ── Dataset mutations (edit captions, delete images, delete datasets) ───────
import { rmSync, renameSync } from "node:fs";
import { listJobs } from "./training-jobs.js";
import { buildTrainingConfig } from "./training-config.js";
function samePath(a, b) {
    const norm = (x) => (process.platform === "win32" ? x.toLowerCase() : x);
    return norm(resolve(a)) === norm(resolve(b));
}
/** A dataset dir a RUNNING/QUEUED job trains from must not be edited or
 *  deleted (the dir is bind-mounted into the running container — same guard
 *  as prepareDataset's replace check, codex finding lineage). */
export async function assertDatasetEditable(dir, what) {
    const active = await listJobs();
    const inUse = active.find((j) => (j.status === "running" || j.status === "queued") && j.datasetPath && samePath(j.datasetPath, dir));
    if (inUse) {
        throw new Error(`dataset is in use by ${inUse.status} job ${inUse.id} — ${what} is blocked until it finishes (or cancel it first)`);
    }
}
/** Edit captions and/or delete images in a staged dataset. Caption writes are
 *  atomic per file (tmp+rename); unknown files are warnings, never silent. */
export async function updateDataset(name, opts) {
    const dir = datasetDir(name);
    await assertDatasetEditable(dir, "editing");
    const warnings = [];
    let captionsSet = 0;
    let imagesDeleted = 0;
    for (const [file, caption] of Object.entries(opts.setCaptions ?? {})) {
        if (file !== basename(file) || !IMAGE_EXTS.has(extname(file).toLowerCase())) {
            warnings.push(`${file}: skipped (not a supported image file)`);
            continue;
        }
        const img = join(dir, file);
        if (!existsSync(img) || !statSync(img).isFile()) {
            warnings.push(`${file}: no such image — caption not written`);
            continue;
        }
        const cap = join(dir, `${basename(file, extname(file))}.txt`);
        const tmp = `${cap}.tmp-${process.pid}`;
        writeFileSync(tmp, caption.trim());
        renameSync(tmp, cap);
        captionsSet++;
    }
    for (const file of opts.deleteImages ?? []) {
        if (file !== basename(file) || !IMAGE_EXTS.has(extname(file).toLowerCase())) {
            warnings.push(`${file}: skipped (not a supported image file)`);
            continue;
        }
        const img = join(dir, file);
        if (!existsSync(img) || !statSync(img).isFile()) {
            warnings.push(`${file}: no such image`);
            continue;
        }
        rmSync(img, { force: true });
        rmSync(join(dir, `${basename(file, extname(file))}.txt`), { force: true });
        imagesDeleted++;
    }
    return { captionsSet, imagesDeleted, warnings };
}
/** Delete a whole staged dataset (dir + captions). In-use guard as above. */
export async function deleteDataset(name) {
    const dir = datasetDir(name);
    await assertDatasetEditable(dir, "deleting");
    rmSync(dir, { recursive: true, force: true });
}
/** The raw YAML train_start would write for these settings — the ostris-UI
 *  "raw config" preview. NO side effects (nothing is written or started). */
export function previewConfig(input) {
    const built = buildTrainingConfig({
        name: input.name,
        flow: "character",
        model: "flux1-dev",
        datasetPath: input.datasetPath,
        outputDir: join(trainingRoot(), "jobs", "<job>", "output"),
        trigger: input.trigger,
        params: input.params,
    });
    return { jobName: built.jobName, yaml: built.yaml };
}
// ── action:"file": bounded inline reads under the training root ────────────
/** Read an image under the training root (dataset images, job samples) as
 *  base64 for the MCP image channel. Contained, image-only, size-capped. */
export function readTrainingFile(absPath) {
    if (!contained(absPath, trainingRoot())) {
        throw new Error("path escapes the training root");
    }
    const p = realOrLexical(absPath);
    const ext = extname(p).toLowerCase();
    if (!IMAGE_EXTS.has(ext)) {
        throw new Error(`only image files are served (${[...IMAGE_EXTS].join("/")}), not ${ext || "(none)"}`);
    }
    const st = statSync(p); // ENOENT surfaces honestly as "no such file"
    if (st.size > MAX_FILE_BYTES) {
        throw new Error(`file too large to inline (${(st.size / 1024 / 1024).toFixed(1)} MB > 2 MB) — fetch it on the rig instead`);
    }
    const mimeType = ext === ".png" ? "image/png" : ext === ".webp" ? "image/webp" : "image/jpeg";
    return { data: readFileSync(p).toString("base64"), mimeType };
}
/** Read a job's effective settings back from the ai-toolkit config.yml we
 *  wrote for it (the single source of truth run.py consumed). */
export async function getJobConfig(id) {
    const job = await getJob(id);
    if (!job)
        throw new Error(`no training job ${id}`);
    const params = {};
    let source = "unknown";
    let rawYaml = null;
    try {
        rawYaml = readFileSync(join(job.jobDir, "config.yml"), "utf-8");
        const doc = parseYaml(rawYaml);
        const proc = (doc?.config?.process?.[0] ?? {});
        const train = proc.train ?? {};
        const network = proc.network ?? {};
        const save = proc.save ?? {};
        const sample = proc.sample ?? {};
        const model = proc.model ?? {};
        const ds = Array.isArray(proc.datasets) ? (proc.datasets[0] ?? {}) : {};
        if (typeof train.steps === "number")
            params.steps = train.steps;
        if (typeof train.lr === "number")
            params.lr = train.lr;
        if (typeof network.linear === "number")
            params.rank = network.linear;
        if (Array.isArray(ds.resolution))
            params.resolution = ds.resolution;
        if (typeof train.batch_size === "number")
            params.batchSize = train.batch_size;
        if (typeof save.save_every === "number")
            params.saveEvery = save.save_every;
        if (typeof sample.sample_every === "number")
            params.sampleEvery = sample.sample_every;
        if (typeof model.quantize === "boolean")
            params.quantize = model.quantize;
        source = "config.yml";
    }
    catch {
        // config.yml missing/unreadable — return the record fields, empty params
    }
    return {
        id: job.id,
        name: job.name,
        flow: job.flow,
        model: job.model,
        trigger: job.trigger,
        datasetPath: job.datasetPath,
        rawYaml,
        params,
        source,
    };
}
//# sourceMappingURL=training-datasets.js.map