Moser9815-ComfyUI-RMAutomation/js/rm_styles_full.js

/**
 * RM Styles Full - Frontend Widget Controller
 * Handles mode switching, increment/decrement logic.
 * Random mode is handled server-side (Python) with anti-repeat history.
 */

import { app } from "../../scripts/app.js";
import { api } from "../../scripts/api.js";

console.log("[RMStyles] Script loading...");

const RM_STYLES_NODES = ["RMStylesFull", "RMStylesFullDisplay", "RMStylesPipe"];

// The client-side counter that used to live here has been removed. Increment/Decrement
// are now computed in Python (nodes/styles_full.py::_step_sequence) and reported back
// via the `rm_styles_executed` message. Advancing it here silently froze or skipped the
// sequence whenever the browser wasn't driving execution — batch-queued prompts all
// snapshot the same next_prompt, and API-queued runs never load this file at all.

/**
 * Set up widget callbacks for a node
 */
function setupNodeWidgets(node) {
  if (!node.widgets) return;

  // Check if already set up
  if (node._rmSetupDone) return;
  node._rmSetupDone = true;

  // Find our widgets
  const modeWidget = node.widgets.find((w) => w.name === "mode");
  const prevWidget = node.widgets.find((w) => w.name === "previous_prompt");
  const nextWidget = node.widgets.find((w) => w.name === "next_prompt");
  const minWidget = node.widgets.find((w) => w.name === "minimum");
  const maxWidget = node.widgets.find((w) => w.name === "maximum");

  if (!modeWidget || !prevWidget || !nextWidget || !minWidget || !maxWidget) {
    return;
  }

  // Store widget references on node for easy access
  node._rmWidgets = { modeWidget, prevWidget, nextWidget, minWidget, maxWidget };

  // Flag to prevent recursive callbacks
  node._rmUpdating = false;

  // Hook mode widget callback
  const origModeCallback = modeWidget.callback;
  modeWidget.callback = function (value) {
    if (origModeCallback) {
      origModeCallback.call(this, value);
    }

    if (node._rmUpdating) return;
    // Deliberately does NOT touch next_prompt any more. Nudging the widget on a mode
    // change would skip a style, and Python (which now owns the counter) would read the
    // altered widget as a deliberate edit and resync to it. next_prompt is the value
    // that WILL be used; switching mode should not consume one.
    console.log(`[RMStyles] Mode changed to ${value}; next_prompt left at ${nextWidget.value}`);
  };

  // Hook next_prompt widget callback to detect manual edits
  const origNextCallback = nextWidget.callback;
  nextWidget.callback = function (value) {
    if (origNextCallback) {
      origNextCallback.call(this, value);
    }

    if (node._rmUpdating) return;

    // User manually edited - switch to Manual mode
    if (modeWidget.value !== "Manual") {
      node._rmUpdating = true;
      modeWidget.value = "Manual";
      node._rmUpdating = false;
    }
  };
}

/**
 * Handle node execution complete (Increment/Decrement only)
 * Random mode is handled by the rm_styles_executed message from Python.
 */
// NOTE: the counter is no longer advanced here. Every mode is now driven by Python via
// the `rm_styles_executed` message. Advancing client-side meant the sequence froze or
// skipped whenever the browser wasn't the thing driving execution — batch-queued
// prompts all snapshot the same next_prompt, and API-queued runs (Cloud Run controller,
// auto-retry) never execute this file at all. Python is the single source of truth.
function handleNodeExecuted(node) {
  // Intentionally a no-op; kept so the executing/execution_success wiring below is
  // unchanged and there is one obvious place to look for the old behavior.
}

// Register the extension
app.registerExtension({
  name: "RMAutomation.StylesFull",

  async beforeRegisterNodeDef(nodeType, nodeData, app) {
    if (!RM_STYLES_NODES.includes(nodeData.name)) return;

    console.log(`[RMStyles] Registering handlers for ${nodeData.name}`);

    // Hook into node creation
    const origOnNodeCreated = nodeType.prototype.onNodeCreated;
    nodeType.prototype.onNodeCreated = function () {
      const result = origOnNodeCreated?.apply(this, arguments);

      const node = this;
      setTimeout(() => {
        setupNodeWidgets(node);
      }, 200);

      return result;
    };

    // Also hook onConfigure for loaded workflows
    const origOnConfigure = nodeType.prototype.onConfigure;
    nodeType.prototype.onConfigure = function (data) {
      const result = origOnConfigure?.apply(this, arguments);

      const node = this;
      setTimeout(() => {
        setupNodeWidgets(node);
      }, 200);

      return result;
    };
  },

  async setup() {
    console.log("[RMStyles] Setting up API event listeners...");

    // Track nodes that are part of the current execution
    const pendingNodes = new Set();

    // Helper to process pending nodes
    const processPendingNodes = () => {
      for (const node of pendingNodes) {
        handleNodeExecuted(node);
      }
      pendingNodes.clear();
    };

    // Python reports what it used and what it will use next. Both widgets are display
    // only — this is the single place either of them is updated after execution.
    api.addEventListener("rm_styles_executed", (e) => {
      const { node_id, prompt_number, next_prompt } = e.detail;
      const node = app.graph?.getNodeById(Number(node_id));
      if (!node || !node._rmWidgets) return;

      node._rmUpdating = true;
      node._rmWidgets.prevWidget.value = prompt_number;
      if (typeof next_prompt === "number") {
        node._rmWidgets.nextWidget.value = next_prompt;
      }
      node._rmUpdating = false;

      console.log(`[RMStyles] Python used #${prompt_number}, next=${next_prompt ?? "(n/a)"} (node ${node_id})`);
    });

    api.addEventListener("execution_start", (e) => {
      pendingNodes.clear();
    });

    api.addEventListener("executing", (e) => {
      // Handle both old and new ComfyUI event formats
      const nodeId =
        typeof e.detail === "object" && e.detail !== null
          ? e.detail.node
          : e.detail;

      if (nodeId) {
        const node = app.graph?.getNodeById(Number(nodeId));
        if (node && RM_STYLES_NODES.includes(node.type)) {
          pendingNodes.add(node);
        }
      } else {
        // null detail = execution finished - update all our nodes that ran
        processPendingNodes();
      }
    });

    // Handle cached executions
    api.addEventListener("execution_cached", (e) => {
      const cachedNodes = e.detail?.nodes ?? [];
      for (const nodeId of cachedNodes) {
        const node = app.graph?.getNodeById(Number(nodeId));
        if (node && RM_STYLES_NODES.includes(node.type)) {
          pendingNodes.add(node);
        }
      }
    });

    // Handle execution interrupts - still update our nodes
    api.addEventListener("execution_interrupted", (e) => {
      console.log("[RMStyles] Execution interrupted, updating pending nodes...");
      processPendingNodes();
    });

    // Handle execution errors - still update our nodes
    api.addEventListener("execution_error", (e) => {
      console.log("[RMStyles] Execution error, updating pending nodes...");
      processPendingNodes();
    });

    console.log("[RMStyles] Extension setup complete!");
  },
});

console.log("[RMStyles] Script loaded");