From e05046af10b63a39c2a6ffddea1d952fff9caee1 Mon Sep 17 00:00:00 2001 From: Will Miao Date: Tue, 25 Aug 2026 08:30:17 +0800 Subject: [PATCH] fix(loaders): sanitize invalid control_after_generate values when loading old workflows Old workflows (saved before the control_after_generate feature) carry a shorter widgets_values array. The frontend's index-based widget restore then shifts the old weight_dtype value into the hidden control widget (leaving an invalid value like 'default') and silently resets weight_dtype to its default. On graph load, hand the shifted value back to weight_dtype when it still sits at its default, then reset the control mode to 'fixed' so old workflows keep loading deterministically. --- .../frontend/core/randomLoaderControl.test.js | 255 ++++++++++++++++++ web/comfyui/random_loader_control.js | 78 +++++- 2 files changed, 332 insertions(+), 1 deletion(-) create mode 100644 tests/frontend/core/randomLoaderControl.test.js diff --git a/tests/frontend/core/randomLoaderControl.test.js b/tests/frontend/core/randomLoaderControl.test.js new file mode 100644 index 00000000..1a13d47a --- /dev/null +++ b/tests/frontend/core/randomLoaderControl.test.js @@ -0,0 +1,255 @@ +import { describe, it, expect, beforeEach, vi } from "vitest"; + +const state = vi.hoisted(() => { + const APP_MODULE = new URL("../../../scripts/app.js", import.meta.url).pathname; + const API_MODULE = new URL("../../../scripts/api.js", import.meta.url).pathname; + return { + APP_MODULE, + API_MODULE, + graph: { onConfigure: null, nodes: [] }, + registerExtension: vi.fn(), + fetchApi: vi.fn(), + }; +}); + +vi.mock(state.APP_MODULE, () => ({ + app: { + registerExtension: state.registerExtension, + graph: state.graph, + }, +})); + +vi.mock(state.API_MODULE, () => ({ + api: { + fetchApi: state.fetchApi, + }, +})); + +const { sanitizeControlWidget } = await import( + "../../../web/comfyui/random_loader_control.js" +); + +const CONTROL_VALUES = [ + "fixed", + "increment", + "decrement", + "randomize", + "increment-wrap", +]; +const DTYPE_VALUES = ["default", "fp8_e4m3fn", "fp8_e4m3fn_fast", "fp8_e5m2"]; + +function makeUnetNode(overrides = {}) { + return { + comfyClass: "Unet Loader (LoraManager)", + widgets: [ + { + name: "unet_name", + type: "combo", + value: "model.safetensors", + options: { values: ["model.safetensors"] }, + }, + { + // Real ComfyUI names the control widget after the string option + // (e.g. 'fixed'), not 'control_after_generate'. + name: "fixed", + type: "combo", + value: "fixed", + options: { values: CONTROL_VALUES }, + }, + { name: "control_filter_list", type: "string", value: "", options: {} }, + { + name: "weight_dtype", + type: "combo", + value: "default", + options: { values: DTYPE_VALUES }, + }, + { name: "base_model", type: "combo", value: "Any", options: { values: ["Any"] } }, + ], + ...overrides, + }; +} + +function makeCheckpointNode(overrides = {}) { + return { + comfyClass: "Checkpoint Loader (LoraManager)", + widgets: [ + { + name: "ckpt_name", + type: "combo", + value: "model.safetensors", + options: { values: ["model.safetensors"] }, + }, + { + // Real ComfyUI names the control widget after the string option + // (e.g. 'fixed'), not 'control_after_generate'. + name: "fixed", + type: "combo", + value: "fixed", + options: { values: CONTROL_VALUES }, + }, + { name: "control_filter_list", type: "string", value: "", options: {} }, + { name: "base_model", type: "combo", value: "Any", options: { values: ["Any"] } }, + ], + ...overrides, + }; +} + +function controlWidget(node) { + return node.widgets.find( + (widget) => + widget.name === "control_after_generate" || + (widget.type === "combo" && + Array.isArray(widget.options?.values) && + widget.options.values.length > 0 && + widget.options.values.every((v) => CONTROL_VALUES.includes(v))) + ); +} + +function dtypeWidget(node) { + return node.widgets.find((widget) => widget.name === "weight_dtype"); +} + +beforeEach(() => { + state.graph.onConfigure = null; + state.graph.nodes = []; + state.graph.__loraManagerConfigureHooked = false; + state.fetchApi.mockReset(); + state.fetchApi.mockResolvedValue({ + ok: true, + json: async () => ({ + items: [{ name: "model.safetensors", base_model: "Any" }], + }), + }); +}); + +describe("sanitizeControlWidget", () => { + it("hands the shifted weight_dtype value back and resets control to fixed", () => { + const node = makeUnetNode(); + controlWidget(node).value = "fp8_e4m3fn"; + dtypeWidget(node).value = "default"; + + sanitizeControlWidget(node, { + modelWidget: "unet_name", + subType: "diffusion_model", + dtypeWidget: "weight_dtype", + }); + + expect(dtypeWidget(node).value).toBe("fp8_e4m3fn"); + expect(controlWidget(node).value).toBe("fixed"); + }); + + it("resets control when the shifted value is the dtype default itself", () => { + const node = makeUnetNode(); + controlWidget(node).value = "default"; + dtypeWidget(node).value = "default"; + + sanitizeControlWidget(node, { + modelWidget: "unet_name", + subType: "diffusion_model", + dtypeWidget: "weight_dtype", + }); + + expect(dtypeWidget(node).value).toBe("default"); + expect(controlWidget(node).value).toBe("fixed"); + }); + + it("leaves valid control modes untouched", () => { + const node = makeUnetNode(); + controlWidget(node).value = "randomize"; + dtypeWidget(node).value = "fp8_e4m3fn"; + + sanitizeControlWidget(node, { + modelWidget: "unet_name", + subType: "diffusion_model", + dtypeWidget: "weight_dtype", + }); + + expect(controlWidget(node).value).toBe("randomize"); + expect(dtypeWidget(node).value).toBe("fp8_e4m3fn"); + }); + + it("does not overwrite a weight_dtype that is not at its default", () => { + const node = makeUnetNode(); + controlWidget(node).value = "fp8_e4m3fn"; + dtypeWidget(node).value = "fp8_e5m2"; + + sanitizeControlWidget(node, { + modelWidget: "unet_name", + subType: "diffusion_model", + dtypeWidget: "weight_dtype", + }); + + expect(dtypeWidget(node).value).toBe("fp8_e5m2"); + expect(controlWidget(node).value).toBe("fixed"); + }); + + it("only resets the control mode for nodes without a dtype widget", () => { + const node = makeCheckpointNode(); + controlWidget(node).value = "default"; + + sanitizeControlWidget(node, { + modelWidget: "ckpt_name", + subType: "checkpoint", + }); + + expect(controlWidget(node).value).toBe("fixed"); + }); + + it("is a no-op when the node has no control widget", () => { + const node = makeUnetNode(); + node.widgets = node.widgets.filter( + (widget) => widget.name !== "control_after_generate" + ); + + expect(() => + sanitizeControlWidget(node, { + modelWidget: "unet_name", + subType: "diffusion_model", + dtypeWidget: "weight_dtype", + }) + ).not.toThrow(); + }); + + it("falls back to the known control modes when options.values is missing", () => { + const node = makeUnetNode(); + const control = controlWidget(node); + // Standard ComfyUI name, so the widget is found by name while its + // options.values is gone (exercises the CONTROL_MODES fallback). + control.name = "control_after_generate"; + control.value = "randomize"; + control.options = {}; + + sanitizeControlWidget(node, { + modelWidget: "unet_name", + subType: "diffusion_model", + dtypeWidget: "weight_dtype", + }); + + expect(controlWidget(node).value).toBe("randomize"); + }); +}); + +describe("extension graph configure hook", () => { + it("sanitizes loader nodes after graph configure", async () => { + const extension = state.registerExtension.mock.calls.map( + (call) => call[0] + )[0]; + await extension.setup(); + + const node = makeUnetNode(); + controlWidget(node).value = "fp8_e4m3fn"; + dtypeWidget(node).value = "default"; + state.graph.nodes = [node]; + + const nodeType = { comfyClass: "Unet Loader (LoraManager)", prototype: {} }; + extension.beforeRegisterNodeDef(nodeType, {}); + nodeType.prototype.onAdded.call({ graph: state.graph }); + + state.graph.onConfigure({}); + + await new Promise((resolve) => setTimeout(resolve, 0)); + + expect(controlWidget(node).value).toBe("fixed"); + expect(dtypeWidget(node).value).toBe("fp8_e4m3fn"); + }); +}); \ No newline at end of file diff --git a/web/comfyui/random_loader_control.js b/web/comfyui/random_loader_control.js index 1b0be8c6..91febd0b 100644 --- a/web/comfyui/random_loader_control.js +++ b/web/comfyui/random_loader_control.js @@ -9,9 +9,26 @@ const NODE_CONFIGS = { "Unet Loader (LoraManager)": { modelWidget: "unet_name", subType: "diffusion_model", + // Old workflows (saved before the control_after_generate feature) carry a + // shorter widgets_values array; the frontend's index-based restore then + // shifts the old weight_dtype value into the hidden control widget and + // silently resets weight_dtype to its default. Sanitization hands the + // shifted value back to this widget. + dtypeWidget: "weight_dtype", }, }; +// Fallback set of valid control modes, used only when the widget's +// options.values list is unavailable. Combo targets additionally get +// 'increment-wrap' appended by ComfyUI. +const CONTROL_MODES = new Set([ + "fixed", + "increment", + "decrement", + "randomize", + "increment-wrap", +]); + const poolCache = new Map(); async function fetchPool(subType) { @@ -66,10 +83,69 @@ function applyBaseModelFilter(node, config) { } } +function isControlWidget(widget) { + if (!widget || widget.type !== "combo") return false; + const values = widget.options?.values; + return ( + Array.isArray(values) && + values.length > 0 && + values.every((value) => CONTROL_MODES.has(value)) + ); +} + +/** + * Repair a control_after_generate widget that holds a value outside its option + * list after loading an old workflow. The invalid value is a side effect of + * the frontend's index-based widget restore: old workflows serialized fewer + * widget values (no control slot), so the value that followed the model combo + * (e.g. weight_dtype) shifted into the control widget while its real widget + * was silently reset to its default. Hand the shifted value back, then reset + * the control mode to 'fixed' (the node's declared default) so old workflows + * keep loading deterministically and the invalid value stops persisting. + */ +export function sanitizeControlWidget(node, config) { + const controlWidget = node.widgets?.find( + (widget) => + // ComfyUI names the widget after the input option string when it is a + // string (e.g. 'fixed'), so it cannot be located by name alone. + widget.name === "control_after_generate" || isControlWidget(widget) + ); + if (!controlWidget) return; + + const validModes = controlWidget.options?.values; + if (Array.isArray(validModes)) { + if (validModes.includes(controlWidget.value)) return; + } else if (CONTROL_MODES.has(controlWidget.value)) { + return; + } + + if (config.dtypeWidget) { + const dtypeWidget = node.widgets?.find( + (widget) => widget.name === config.dtypeWidget + ); + const dtypeOptions = dtypeWidget?.options?.values; + if ( + dtypeWidget && + Array.isArray(dtypeOptions) && + // Only hand the value back when the real widget still sits at its + // default; a non-default value means it was restored or edited + // correctly and the control value is just stale workflow data. + dtypeWidget.value === dtypeOptions[0] && + dtypeOptions.includes(controlWidget.value) + ) { + dtypeWidget.value = controlWidget.value; + } + } + + controlWidget.value = "fixed"; +} + function applyToAllNodes() { app.graph?.nodes?.forEach((node) => { const config = NODE_CONFIGS[node.comfyClass]; - if (config) applyBaseModelFilter(node, config); + if (!config) return; + sanitizeControlWidget(node, config); + applyBaseModelFilter(node, config); }); }