DumasNodes/js/load_images_folder/index.js

import { app } from "/scripts/app.js";
import { applyAdaptiveCanvasOnly } from "../shared/nodes2.mjs";
import { listFolder, pickNativeFolder } from "./api.mjs";
import {
  COMFY_CLASS,
  HIDDEN_INPUT_NAME,
  readState,
  writeState,
} from "./state.mjs";
import {
  buildRoot,
  injectCSS,
  openBrowsePopup,
  openPickGallery,
} from "./ui.mjs";

const MIN_W = 300;
const DEFAULT_W = 370;

function hideJsonWidget(widgets, name) {
  const widget = (widgets || []).find((entry) => entry?.name === name);
  if (!widget) return;
  widget.type = "hidden";
  widget.computeSize = () => [0, 0];
}

function normalizePath(value) {
  if (!value) return "";
  let normalized = String(value).trim().replace(/\\/g, "/").replace(/\/+$/, "");
  if (/^[A-Za-z]:$/.test(normalized)) normalized += "/";
  return normalized;
}

function stripInputs(node) {
  if (!node?.inputs?.length) return;
  for (let index = node.inputs.length - 1; index >= 0; index -= 1) {
    if (node.inputs[index]?.link != null) {
      try {
        node.disconnectInput(index);
      } catch (_error) {
        // Ignore failed disconnects and keep stripping.
      }
    }
    node.removeInput(index);
  }
}

function selectionSummary(state, total) {
  const mode = state.selection_mode || "selected";
  if (mode === "all") return `All ${total} / ${total}`;
  if (mode === "first_n") return `First ${Math.min(Math.max(parseInt(state.first_n, 10) || 0, 0), total)} / ${total}`;
  if (mode === "random") return `Random ${total ? 1 : 0} / ${total}`;
  return `${(state.selected || []).length} / ${total}`;
}

function renderUI(node) {
  const ui = node._dlfUI;
  if (!ui) return;
  const state = readState(node);
  if (document.activeElement !== ui.folderInput) {
    ui.folderInput.value = state.folder || "";
  }
  const total = (node._dlfFiles || []).length;
  ui.pickBtn.textContent = `Pick images · ${selectionSummary(state, total)}`;
  ui.pickBtn.classList.toggle("empty", total === 0);
  ui.msgEl.textContent = node._dlfListError || "";
  node.setDirtyCanvas?.(true, true);
}

async function refreshListing(node) {
  const requestId = (node._dlfListReq = (node._dlfListReq || 0) + 1);
  const state = readState(node);
  if (!state.folder) {
    node._dlfFiles = [];
    node._dlfListError = "";
    renderUI(node);
    return;
  }

  const response = await listFolder(state.folder, state.recursive);
  if (node._dlfListReq !== requestId || !node._dlfUI) return;
  if (response?.ok) {
    node._dlfFiles = response.files || [];
    node._dlfListError = node._dlfFiles.length ? "" : "No images found in this folder.";
  } else {
    node._dlfFiles = [];
    node._dlfListError = response?.message || "Folder not found.";
  }

  const present = new Set((node._dlfFiles || []).map((file) => file.file));
  const nextState = readState(node);
  const kept = (nextState.selected || []).filter((file) => present.has(file));
  if (kept.length !== (nextState.selected || []).length) {
    nextState.selected = kept;
    writeState(node, nextState);
  }
  renderUI(node);
}

async function setFolder(node, folder) {
  const normalized = normalizePath(folder);
  const state = readState(node);
  const changed = (state.folder || "") !== normalized;
  state.folder = normalized;
  if (changed) state.selected = [];
  writeState(node, state);
  await refreshListing(node);
}

function setupNode(node) {
  injectCSS();
  hideJsonWidget(node.widgets, HIDDEN_INPUT_NAME);
  stripInputs(node);

  const ui = buildRoot();
  node._dlfUI = ui;
  const widget = node.addDOMWidget("dumas_load_images_folder", "custom", ui.root, {
    getValue: () => null,
    setValue: () => {},
    serialize: false,
    getMinHeight: () => 96,
  });
  widget.computeLayoutSize = () => ({ minHeight: 96, minWidth: 1 });
  applyAdaptiveCanvasOnly(widget);

  ui.folderInput.addEventListener("keydown", (event) => {
    event.stopImmediatePropagation();
    if (event.key === "Enter") {
      event.preventDefault();
      ui.folderInput.blur();
    }
  });
  ui.folderInput.addEventListener("change", () => setFolder(node, ui.folderInput.value.trim()));
  ui.folderInput.addEventListener("paste", () => {
    setTimeout(() => {
      const value = ui.folderInput.value.trim();
      if (normalizePath(value) !== (readState(node).folder || "")) setFolder(node, value);
    }, 0);
  });
  ui.browseBtn.addEventListener("click", async () => {
    const start = readState(node).folder || "";
    const previous = ui.browseLbl.textContent;
    ui.browseBtn.disabled = true;
    ui.browseLbl.textContent = "Opening…";
    let response;
    try {
      response = await pickNativeFolder(start);
    } catch (_error) {
      response = { ok: false };
    }
    ui.browseBtn.disabled = false;
    ui.browseLbl.textContent = previous || "Browse";
    if (response?.ok && response.path) {
      await setFolder(node, response.path);
      return;
    }
    if (response?.cancelled) return;
    openBrowsePopup(node, ui.browseBtn, {
      startPath: start,
      onPick: (folder) => setFolder(node, folder),
    });
  });
  ui.pickBtn.addEventListener("click", async () => {
    const typed = ui.folderInput.value.trim();
    if (typed !== (readState(node).folder || "")) await setFolder(node, typed);
    const state = readState(node);
    if (!state.folder) {
      ui.folderInput.focus();
      node._dlfListError = "Set a folder first.";
      renderUI(node);
      return;
    }
    ui.pickBtn.disabled = true;
    try {
      await refreshListing(node);
    } finally {
      ui.pickBtn.disabled = false;
    }
    openPickGallery(node, ui.pickBtn, {
      onChange: renderUI,
      refreshListing,
    });
  });

  node.size[0] = Math.max(node.size[0] || 0, DEFAULT_W);
  node.size[1] = Math.max(node.size[1] || 0, 110);
  queueMicrotask(() => refreshListing(node));
}

function collectNodes(graph, out) {
  if (!graph) return;
  const nodes = graph._nodes || graph.nodes || [];
  for (const node of nodes) {
    if (node?.comfyClass === COMFY_CLASS) out.push(node);
    const inner = node?.subgraph || node?.graph || node?._graph;
    if (inner && inner !== graph) collectNodes(inner, out);
  }
}

function matchNode(nodes, promptId) {
  let node = nodes.find((entry) => String(entry.id) === String(promptId));
  if (node) return node;
  const tail = String(promptId).split(":").pop();
  node = nodes.find((entry) => String(entry.id) === tail);
  return node || null;
}

function injectState(result) {
  const output = result?.output;
  if (!output) return;
  const nodes = [];
  collectNodes(app.graph, nodes);
  for (const id in output) {
    const entry = output[id];
    if (!entry || entry.class_type !== COMFY_CLASS) continue;
    const node = matchNode(nodes, id);
    if (!node) continue;
    if (!entry.inputs) entry.inputs = {};
    entry.inputs[HIDDEN_INPUT_NAME] = JSON.stringify(readState(node));
  }
}

function installGraphToPromptHook() {
  if (app._dlfGraphPatched) return;
  app._dlfGraphPatched = true;
  const original = app.graphToPrompt.bind(app);
  app.graphToPrompt = async function graphToPromptPatched(...args) {
    const result = await original(...args);
    injectState(result);
    return result;
  };
}

app.registerExtension({
  name: "Dumas.LoadImagesFolder",
  setup() {
    installGraphToPromptHook();
  },
  async beforeRegisterNodeDef(nodeType, nodeData) {
    if (nodeData?.name !== COMFY_CLASS) return;

    const originalNodeCreated = nodeType.prototype.onNodeCreated;
    nodeType.prototype.onNodeCreated = function onNodeCreated() {
      const result = originalNodeCreated?.apply(this, arguments);
      setupNode(this);
      return result;
    };

    const originalConfigure = nodeType.prototype.onConfigure;
    nodeType.prototype.onConfigure = function onConfigure() {
      const result = originalConfigure?.apply(this, arguments);
      stripInputs(this);
      queueMicrotask(() => refreshListing(this));
      return result;
    };

    const originalDrawForeground = nodeType.prototype.onDrawForeground;
    nodeType.prototype.onDrawForeground = function onDrawForeground(ctx) {
      const result = originalDrawForeground?.call(this, ctx);
      this.size[0] = Math.max(this.size[0], MIN_W);
      return result;
    };

    nodeType.prototype.onConnectInput = function onConnectInput() {
      return false;
    };

    const originalRemoved = nodeType.prototype.onRemoved;
    nodeType.prototype.onRemoved = function onRemoved() {
      this._dlfGallery?._dlfClose?.();
      this._dlfBrowsePopup?._dlfClose?.();
      document.querySelectorAll(".dlf-menu").forEach((menu) => menu._dlfClose?.());
      return originalRemoved?.apply(this, arguments);
    };
  },
});