Performance optimization on queueStore (#521)

* Eager calculate flat outputs

* Optimize flat tasks performance
This commit is contained in:
Chenlei Hu
2024-08-18 21:48:48 -04:00
committed by GitHub
parent 8576d3797b
commit 4a43dfe6b9

View File

@@ -9,7 +9,6 @@ import type {
TaskOutput,
ResultItem
} from '@/types/apiTypes'
import { validateTaskItem } from '@/types/apiTypes'
import type { NodeId } from '@/types/comfyWorkflow'
import { instanceToPlain, plainToClass } from 'class-transformer'
import _ from 'lodash'
@@ -51,18 +50,33 @@ export class ResultItemImpl {
}
export class TaskItemImpl {
taskType: TaskType
prompt: TaskPrompt
status?: TaskStatus
outputs: TaskOutput
readonly taskType: TaskType
readonly prompt: TaskPrompt
readonly status?: TaskStatus
readonly outputs: TaskOutput
readonly flatOutputs: ReadonlyArray<ResultItemImpl>
get flatOutputs(): ResultItemImpl[] {
constructor(
taskType: TaskType,
prompt: TaskPrompt,
status: TaskStatus | undefined,
outputs: TaskOutput,
flatOutputs?: ReadonlyArray<ResultItemImpl>
) {
this.taskType = taskType
this.prompt = prompt
this.status = status
this.outputs = outputs
this.flatOutputs = flatOutputs ?? this.calculateFlatOutputs()
}
private calculateFlatOutputs(): ReadonlyArray<ResultItemImpl> {
if (!this.outputs) {
return []
}
return Object.entries(this.outputs).flatMap(([nodeId, nodeOutputs]) =>
Object.entries(nodeOutputs).flatMap(([mediaType, items]) =>
(items as ResultItem[]).flatMap((item: ResultItem) =>
(items as ResultItem[]).map((item: ResultItem) =>
plainToClass(ResultItemImpl, {
...item,
nodeId,
@@ -230,22 +244,25 @@ export const useQueueStore = defineStore('queue', {
return [task]
}
return task.flatOutputs.map((output: ResultItemImpl, i: number) =>
plainToClass(TaskItemImpl, {
...instanceToPlain(task),
prompt: [
task.queueIndex,
`${task.promptId}-${i}`,
task.promptInputs,
task.extraData,
task.outputsToExecute
],
outputs: {
[output.nodeId]: {
[output.mediaType]: [instanceToPlain(output)]
}
}
})
return task.flatOutputs.map(
(output: ResultItemImpl, i: number) =>
new TaskItemImpl(
task.taskType,
[
task.queueIndex,
`${task.promptId}-${i}`,
task.promptInputs,
task.extraData,
task.outputsToExecute
],
task.status,
{
[output.nodeId]: {
[output.mediaType]: [output]
}
},
[output]
)
)
})
},
@@ -263,9 +280,15 @@ export const useQueueStore = defineStore('queue', {
const toClassAll = (tasks: TaskItem[]): TaskItemImpl[] =>
tasks
.map((task) => validateTaskItem(task))
.filter((result) => result.success)
.map((result) => plainToClass(TaskItemImpl, result.data))
.map(
(task: TaskItem) =>
new TaskItemImpl(
task.taskType,
task.prompt,
task['status'],
task['outputs'] || {}
)
)
// Desc order to show the latest tasks first
.sort((a, b) => b.queueIndex - a.queueIndex)