Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Refactor] Move Widget.beforeQueued invocation from graphToPrompt to queuePrompt #2667

Merged
merged 1 commit into from
Feb 21, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 12 additions & 16 deletions src/scripts/app.ts
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ import {
import { ExtensionManager } from '@/types/extensionTypes'
import { ColorAdjustOptions, adjustColor } from '@/utils/colorUtil'
import { graphToPrompt } from '@/utils/executionUtil'
import { isImageNode } from '@/utils/litegraphUtil'
import { executeWidgetsCallback, isImageNode } from '@/utils/litegraphUtil'
import { deserialiseAndCreate } from '@/utils/vintageClipboard'

import { type ComfyApi, api } from './api'
Expand Down Expand Up @@ -1260,7 +1260,7 @@ export class ComfyApp {
return '(unknown error)'
}

async queuePrompt(number, batchCount = 1) {
async queuePrompt(number: number, batchCount: number = 1): Promise<boolean> {
this.#queueItems.push({ number, batchCount })

// Only have one action process the items so each one gets a unique seed correctly
Expand All @@ -1276,8 +1276,11 @@ export class ComfyApp {
;({ number, batchCount } = this.#queueItems.pop())

for (let i = 0; i < batchCount; i++) {
const p = await this.graphToPrompt()
// Allow widgets to run callbacks before a prompt has been queued
// e.g. random seed before every gen
executeWidgetsCallback(this.graph.nodes, 'beforeQueued')

const p = await this.graphToPrompt()
try {
const res = await api.queuePrompt(number, p)
this.lastNodeErrors = res.node_errors
Expand All @@ -1303,19 +1306,12 @@ export class ComfyApp {
break
}

for (const n of p.workflow.nodes) {
const node = this.graph.getNodeById(n.id)
if (node.widgets) {
for (const widget of node.widgets) {
// Allow widgets to run callbacks after a prompt has been queued
// e.g. random seed after every gen
if (widget.afterQueued) {
widget.afterQueued()
}
}
}
}

// Allow widgets to run callbacks after a prompt has been queued
// e.g. random seed after every gen
executeWidgetsCallback(
p.workflow.nodes.map((n) => this.graph.getNodeById(n.id)),
'afterQueued'
)
this.canvas.draw(true, true)
await this.ui.queue.update()
}
Expand Down
8 changes: 0 additions & 8 deletions src/utils/executionUtil.ts
Original file line number Diff line number Diff line change
Expand Up @@ -15,14 +15,6 @@ export const graphToPrompt = async (
const { sortNodes = false } = options

for (const outerNode of graph.computeExecutionOrder(false)) {
if (outerNode.widgets) {
for (const widget of outerNode.widgets) {
// Allow widgets to run callbacks before a prompt has been queued
// e.g. random seed before every gen
widget.beforeQueued?.()
}
}

const innerNodes = outerNode.getInnerNodes
? outerNode.getInnerNodes()
: [outerNode]
Expand Down
11 changes: 11 additions & 0 deletions src/utils/litegraphUtil.ts
Original file line number Diff line number Diff line change
Expand Up @@ -42,3 +42,14 @@ export const getItemsColorOption = (items: unknown[]): ColorOption | null => {
? _.head(colorOptions)!
: null
}

export function executeWidgetsCallback(
nodes: LGraphNode[],
callbackName: 'onRemove' | 'beforeQueued' | 'afterQueued'
) {
for (const node of nodes) {
for (const widget of node.widgets ?? []) {
widget[callbackName]?.()
}
}
}