173 lines
7.5 KiB
JavaScript
173 lines
7.5 KiB
JavaScript
/*
|
||
* Part of ComfyUI-AI2Go-Utils.
|
||
*
|
||
* Shared front-end for the four Batch Loader nodes. GPL-3.0, like the rest of the pack.
|
||
*
|
||
* Files dropped/picked are uploaded once to ComfyUI's input/ folder; the hidden `files_json`
|
||
* widget (a JSON array of {name, subfolder, type}) is the single source of truth for save,
|
||
* restore and execution — the Prompt Batch pattern. The Python schema declares MAX_FILES
|
||
* output groups; syncOutputs() trims node.outputs to `count` + the loaded groups and re-adds
|
||
* up to the ceiling when files come back. Trimming only ever cuts from the end — ComfyUI
|
||
* validates output types by slot position, so used slots must stay contiguous from slot 0.
|
||
* parseFiles mirrors parse_files in nodes/batch_loader_core.py — keep the two in sync.
|
||
*/
|
||
import { chainCallback } from "./utility.js";
|
||
const { app } = window.comfyAPI.app;
|
||
|
||
const MAX_FILES = 8;
|
||
|
||
// group: [prefix, TYPE] per output within one file's group, in schema order.
|
||
const NODES = {
|
||
AI2GoBatchImageLoader: { kind: "image", group: [["image_", "IMAGE"]] },
|
||
AI2GoBatchImageLoaderAdvanced: { kind: "image", group: [["image_", "IMAGE"], ["mask_", "MASK"], ["filename_", "STRING"]] },
|
||
AI2GoBatchAudioLoader: { kind: "audio", group: [["audio_", "AUDIO"]] },
|
||
AI2GoBatchAudioLoaderAdvanced: { kind: "audio", group: [["audio_", "AUDIO"], ["filename_", "STRING"]] },
|
||
};
|
||
|
||
// ── Mirror of parse_files in nodes/batch_loader_core.py ──
|
||
function parseFiles(raw) {
|
||
const text = (raw || "").trim();
|
||
if (!text) return { ok: true, files: [] }; // empty is fine in the UI; Python rejects at run time
|
||
let data;
|
||
try { data = JSON.parse(text); } catch (e) { return { ok: false, error: "Malformed files_json: " + e.message }; }
|
||
if (!Array.isArray(data)) return { ok: false, error: "Expected a JSON array of files." };
|
||
const files = [];
|
||
for (let i = 0; i < data.length && i < MAX_FILES; i++) {
|
||
const e = data[i];
|
||
if (!e || typeof e !== "object" || typeof e.name !== "string" || !e.name.trim()) {
|
||
return { ok: false, error: `File #${i + 1}: 'name' must be a non-empty string.` };
|
||
}
|
||
files.push({ name: e.name, subfolder: typeof e.subfolder === "string" ? e.subfolder : "", type: "input" });
|
||
}
|
||
return { ok: true, files };
|
||
}
|
||
|
||
const findWidget = (node, name) => node.widgets?.find((w) => w.name === name);
|
||
|
||
function hideWidget(w) {
|
||
if (!w) return;
|
||
w.hidden = true;
|
||
w.computeSize = () => [0, -4];
|
||
}
|
||
|
||
// Trim node.outputs to count + fileCount groups; re-add (in schema order) up to the ceiling.
|
||
// removeOutput disconnects any links on the removed slot — that is the intended behavior.
|
||
function syncOutputs(node, cfg, fileCount) {
|
||
const want = 1 + Math.min(fileCount, MAX_FILES) * cfg.group.length;
|
||
while (node.outputs.length > want) node.removeOutput(node.outputs.length - 1);
|
||
while (node.outputs.length < want) {
|
||
const slot = node.outputs.length; // next slot index to create
|
||
const gi = (slot - 1) % cfg.group.length; // position within the file's group
|
||
const fi = Math.floor((slot - 1) / cfg.group.length) + 1; // 1-based file number
|
||
const [prefix, type] = cfg.group[gi];
|
||
node.addOutput(prefix + fi, type);
|
||
}
|
||
node.setDirtyCanvas?.(true, true);
|
||
}
|
||
|
||
// Upload one File to ComfyUI's input/ via the same endpoint the stock upload widget uses.
|
||
// The endpoint is generic despite its name; the form field is "image" even for audio.
|
||
async function uploadFile(file) {
|
||
const form = new FormData();
|
||
form.append("image", file);
|
||
form.append("type", "input");
|
||
const api = window.comfyAPI?.api?.api;
|
||
const res = api?.fetchApi
|
||
? await api.fetchApi("/upload/image", { method: "POST", body: form })
|
||
: await fetch("/upload/image", { method: "POST", body: form });
|
||
if (res.status !== 200) throw new Error(`upload of ${file.name} failed (HTTP ${res.status})`);
|
||
const data = await res.json();
|
||
return { name: data.name, subfolder: data.subfolder || "", type: "input" };
|
||
}
|
||
|
||
app.registerExtension({
|
||
name: "AI2Go.BatchLoader",
|
||
|
||
async beforeRegisterNodeDef(nodeType, nodeData) {
|
||
const cfg = NODES[nodeData?.name];
|
||
if (!cfg) return;
|
||
|
||
chainCallback(nodeType.prototype, "onNodeCreated", function () {
|
||
const node = this;
|
||
const jsonW = findWidget(node, "files_json");
|
||
hideWidget(jsonW);
|
||
node._blRows = [];
|
||
|
||
// INT guard (image nodes): a number widget can serialize '' which fails INT validation.
|
||
const maxSizeW = findWidget(node, "max_size");
|
||
if (maxSizeW) {
|
||
maxSizeW.serializeValue = () => {
|
||
const v = parseInt(maxSizeW.value, 10);
|
||
return Number.isFinite(v) && v >= 8 ? v : 1200;
|
||
};
|
||
}
|
||
|
||
// Status line (read-only DOM widget).
|
||
const statusEl = document.createElement("div");
|
||
statusEl.style.cssText = "width:100%;box-sizing:border-box;padding:3px 6px;text-align:center;line-height:1.4;font:12px sans-serif;";
|
||
const setStatus = (text, color) => { statusEl.textContent = text; statusEl.style.color = color; node.setDirtyCanvas?.(true, true); };
|
||
node._blSetStatus = setStatus;
|
||
|
||
function syncJson() {
|
||
if (jsonW) jsonW.value = JSON.stringify(node._blRows);
|
||
}
|
||
node._blSyncJson = syncJson;
|
||
node._blSyncOutputs = () => syncOutputs(node, cfg, node._blRows.length);
|
||
|
||
async function addFiles(fileList) {
|
||
const files = [...fileList].filter((f) => f.type.startsWith(cfg.kind + "/"));
|
||
const rejected = fileList.length - files.length;
|
||
const free = MAX_FILES - node._blRows.length;
|
||
const taking = files.slice(0, free);
|
||
const skipped = files.length - taking.length;
|
||
for (const f of taking) {
|
||
try {
|
||
node._blRows.push(await uploadFile(f));
|
||
} catch (e) {
|
||
setStatus("❌ " + e.message, "#e0555a");
|
||
break;
|
||
}
|
||
}
|
||
syncJson();
|
||
node._blSyncOutputs();
|
||
node._blRender?.();
|
||
const parts = [`${node._blRows.length} file${node._blRows.length === 1 ? "" : "s"} loaded`];
|
||
if (skipped) parts.push(`only ${MAX_FILES} fit — ${skipped} skipped`);
|
||
if (rejected) parts.push(`${rejected} not ${cfg.kind} — ignored`);
|
||
setStatus((skipped || rejected ? "⚠ " : "✅ ") + parts.join("; "), skipped || rejected ? "#e0a03c" : "#46b4e6");
|
||
}
|
||
node._blAddFiles = addFiles;
|
||
|
||
const addBtn = node.addWidget("button", cfg.kind === "image" ? "＋ Add images" : "＋ Add audio", null, () => {
|
||
const picker = document.createElement("input");
|
||
picker.type = "file";
|
||
picker.multiple = true;
|
||
picker.accept = cfg.kind + "/*";
|
||
picker.onchange = () => picker.files?.length && addFiles(picker.files);
|
||
picker.click();
|
||
});
|
||
addBtn.serialize = false;
|
||
|
||
node.addDOMWidget("batch_loader_status", "info", statusEl, { serialize: false });
|
||
setStatus(`Drop ${cfg.kind} files here or press ＋ Add.`, "#8a8a8a");
|
||
|
||
// Fresh node: no files yet -> trim the declared ceiling down to just `count`.
|
||
node._blSyncOutputs();
|
||
});
|
||
|
||
// After a workflow loads: rebuild rows from the restored files_json, then re-trim.
|
||
// (Serialized nodes save their trimmed outputs, so this is normally a no-op — it heals
|
||
// hand-edited or older workflows.)
|
||
chainCallback(nodeType.prototype, "onConfigure", function () {
|
||
const node = this;
|
||
requestAnimationFrame(() => {
|
||
const res = parseFiles(findWidget(node, "files_json")?.value);
|
||
node._blRows = res.ok ? res.files : [];
|
||
node._blSyncJson?.();
|
||
node._blSyncOutputs?.();
|
||
node._blRender?.();
|
||
});
|
||
});
|
||
},
|
||
});
|