diff --git a/web/js/batch_loader.js b/web/js/batch_loader.js new file mode 100644 index 0000000..dcee8b0 --- /dev/null +++ b/web/js/batch_loader.js @@ -0,0 +1,172 @@ +/* + * 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?.(); + }); + }); + }, +});