diff --git a/execution.py b/execution.py index 47d58b9d..7063bd95 100644 --- a/execution.py +++ b/execution.py @@ -34,10 +34,12 @@ def get_input_data(inputs, class_def, unique_id, outputs={}, prompt={}, extra_da for x in h: if h[x] == "PROMPT": input_data_all[x] = [prompt] - if h[x] == "EXTRA_PNGINFO": + elif h[x] == "EXTRA_PNGINFO": input_data_all[x] = [extra_data.get('extra_pnginfo', None)] - if h[x] == "UNIQUE_ID": + elif h[x] == "UNIQUE_ID": input_data_all[x] = [unique_id] + elif h[x] == "EXTRA_DATA": + input_data_all[x] = [extra_data.get(x, None)] return input_data_all def map_node_over_list(obj, input_data_all, func, allow_interrupt=False): diff --git a/tests-ui/tests/extensions.test.js b/tests-ui/tests/extensions.test.js index 159e5113..b854ead0 100644 --- a/tests-ui/tests/extensions.test.js +++ b/tests-ui/tests/extensions.test.js @@ -25,6 +25,7 @@ describe("extensions", () => { nodeCreated: jest.fn(), beforeConfigureGraph: jest.fn(), afterConfigureGraph: jest.fn(), + // provideExtraData not testable without prompt. }; const { app, ez, graph } = await start({ diff --git a/web/scripts/api.js b/web/scripts/api.js index 8c8155be..10c7bd24 100644 --- a/web/scripts/api.js +++ b/web/scripts/api.js @@ -194,12 +194,14 @@ class ComfyApi extends EventTarget { * * @param {number} number The index at which to queue the prompt, passing -1 will insert the prompt at the front of the queue * @param {object} prompt The prompt data to queue + * @param {object} extra_data extra data added to the prompt request */ - async queuePrompt(number, { output, workflow }) { + async queuePrompt(number, { output, workflow }, extra_data) { + extra_data = { extra_pnginfo: { workflow }, ...extra_data } const body = { client_id: this.clientId, prompt: output, - extra_data: { extra_pnginfo: { workflow } }, + extra_data, }; if (number === -1) { diff --git a/web/scripts/app.js b/web/scripts/app.js index a516be70..38e1e4cd 100644 --- a/web/scripts/app.js +++ b/web/scripts/app.js @@ -2119,7 +2119,9 @@ export class ComfyApp { const p = await this.graphToPrompt(); try { - const res = await api.queuePrompt(number, p); + const extra_datas = await this.#invokeExtensionsAsync("provideExtraData"); + const extra_data = Object.assign({}, ...extra_datas) + const res = await api.queuePrompt(number, p, extra_data); this.lastNodeErrors = res.node_errors; if (this.lastNodeErrors.length > 0) { this.canvas.draw(true, true); diff --git a/web/types/comfy.d.ts b/web/types/comfy.d.ts index f7129b55..86695e9c 100644 --- a/web/types/comfy.d.ts +++ b/web/types/comfy.d.ts @@ -58,6 +58,11 @@ export interface ComfyExtension { * @param app The ComfyUI app instance */ nodeCreated(node: LGraphNode, app: ComfyApp); + /** + * Allows the extension to add custom extra_data to the prompt queue + * @param app The ComfyUI app instance + */ + provideExtraData(app: ComfyApp): Promise; } export type ComfyObjectInfo = {