350 lines
10 KiB
TypeScript
350 lines
10 KiB
TypeScript
import { get, writable } from 'svelte/store';
|
|
import type { Readable, Writable } from 'svelte/store';
|
|
import { isComfyWidgetNode, type DragItemID, type IDragItem } from './layoutStates';
|
|
import { LiteGraph, type LGraphNode, type NodeID, type UUID } from '@litegraph-ts/core';
|
|
import type { SerializedAppState } from '$lib/components/ComfyApp';
|
|
import { getWorkflowRestoreParamsFromWorkflow, type RestoreParamSourceWorkflowNode, type RestoreParamTargets, type RestoreParamWorkflowNodeTargets } from '$lib/restoreParameters';
|
|
import { v4 as uuidv4 } from "uuid";
|
|
import deepEqual from "deep-equal";
|
|
import notify from '$lib/notify';
|
|
import type { ComfyBoxWorkflow } from './workflowState';
|
|
import type { ComfyNodeID, PromptID } from '$lib/api';
|
|
import type { SerializedPromptOutput } from '$lib/utils';
|
|
import type { QueueEntry } from './queueState';
|
|
|
|
export type JourneyNodeType = "root" | "patch";
|
|
|
|
export type JourneyNodeID = UUID;
|
|
|
|
export interface JourneyNode {
|
|
id: JourneyNodeID,
|
|
type: JourneyNodeType,
|
|
children: JourneyPatchNode[],
|
|
promptID?: PromptID,
|
|
images?: string[]
|
|
}
|
|
|
|
export interface JourneyRootNode extends JourneyNode {
|
|
type: "root"
|
|
|
|
/*
|
|
* This contains all the values of the workflow to set
|
|
*/
|
|
base: RestoreParamWorkflowNodeTargets
|
|
}
|
|
|
|
export interface JourneyPatchNode extends JourneyNode {
|
|
type: "patch"
|
|
|
|
parent: JourneyNode,
|
|
|
|
/*
|
|
* This contains only the subset of parameters that were changed from the
|
|
* parent
|
|
*/
|
|
patch: RestoreParamWorkflowNodeTargets
|
|
}
|
|
|
|
function isRoot(node: JourneyNode): node is JourneyRootNode {
|
|
return node.type === "root";
|
|
}
|
|
|
|
function isPatch(node: JourneyNode): node is JourneyPatchNode {
|
|
return node.type === "patch";
|
|
}
|
|
|
|
export function resolvePatch(node: JourneyNode, memoize?: Record<JourneyNodeID, RestoreParamWorkflowNodeTargets>): RestoreParamWorkflowNodeTargets {
|
|
if (node.type === "root") {
|
|
return { ...(node as JourneyRootNode).base }
|
|
}
|
|
|
|
if (memoize && memoize[node.id] != null)
|
|
return { ...memoize[node.id] }
|
|
|
|
const patchNode = (node as JourneyPatchNode);
|
|
const patch = { ...patchNode.patch };
|
|
const base = resolvePatch(patchNode.parent);
|
|
for (const [k, v] of Object.entries(patch)) {
|
|
base[k] = v;
|
|
}
|
|
|
|
if (memoize) {
|
|
memoize[node.id] = base;
|
|
}
|
|
|
|
return base;
|
|
}
|
|
|
|
export function diffParams(base: RestoreParamWorkflowNodeTargets, updated: RestoreParamWorkflowNodeTargets): RestoreParamWorkflowNodeTargets {
|
|
const result = {}
|
|
|
|
for (const [k, v] of Object.entries(updated)) {
|
|
if (!(k in base) || !deepEqual(base[k].finalValue, v.finalValue, { strict: true })) {
|
|
result[k] = v
|
|
}
|
|
}
|
|
|
|
return result;
|
|
}
|
|
|
|
export function calculateWorkflowParamsPatch(parent: JourneyNode, newParams: RestoreParamWorkflowNodeTargets): RestoreParamWorkflowNodeTargets {
|
|
const patch = resolvePatch(parent);
|
|
const diff = diffParams(patch, newParams)
|
|
return diff;
|
|
}
|
|
|
|
/*
|
|
* A "journey" is like browser history for prompts, except organized in a
|
|
* tree-like graph. It lets you save incremental changes to your workflow and
|
|
* jump between past and present sets of parameters.
|
|
*/
|
|
export type JourneyState = {
|
|
root: JourneyRootNode | null,
|
|
nodesByID: Record<JourneyNodeID, JourneyNode>,
|
|
activeNodeID: JourneyNodeID | null,
|
|
|
|
/*
|
|
* Incremented when graph structure is updated
|
|
*/
|
|
version: number
|
|
}
|
|
|
|
type JourneyStateOps = {
|
|
clear: () => void,
|
|
getActiveNode: () => JourneyNode | null,
|
|
// addNode: (params: RestoreParamWorkflowNodeTargets, parent?: JourneyNodeID | JourneyNode) => JourneyNode,
|
|
selectNode: (id?: JourneyNodeID | JourneyNode) => void,
|
|
iterateBreadthFirst: (id?: JourneyNodeID | null) => Iterable<JourneyNode>,
|
|
iterateLinearPath: (id: JourneyNodeID) => Iterable<JourneyNode>,
|
|
pushPatchOntoActive: (workflow: ComfyBoxWorkflow, activeNode?: JourneyNode, showNotification?: boolean) => JourneyNode | null
|
|
onExecuted: (promptID: PromptID, nodeID: ComfyNodeID, output: SerializedPromptOutput, queueEntry: QueueEntry) => void
|
|
}
|
|
|
|
export type WritableJourneyStateStore = Writable<JourneyState> & JourneyStateOps;
|
|
|
|
function create() {
|
|
const store: Writable<JourneyState> = writable(
|
|
{
|
|
root: null,
|
|
nodesByID: {},
|
|
activeNodeID: null,
|
|
version: 0
|
|
})
|
|
|
|
function clear() {
|
|
store.set({
|
|
root: null,
|
|
nodesByID: {},
|
|
activeNodeID: null,
|
|
version: 0
|
|
})
|
|
}
|
|
|
|
function getActiveNode(): JourneyNode | null {
|
|
const state = get(store)
|
|
if (state.activeNodeID === null)
|
|
return null;
|
|
const active = state.nodesByID[state.activeNodeID]
|
|
if (active == null) {
|
|
console.error("[journeyStates] Active node not found in graph!", state.activeNodeID);
|
|
}
|
|
return active;
|
|
}
|
|
|
|
/*
|
|
* params: full state or state patch of widgets in the UI
|
|
* parent: parent node to patch against
|
|
*/
|
|
function addNode(params: RestoreParamWorkflowNodeTargets, parent?: JourneyNodeID | JourneyNode): JourneyNode {
|
|
let _node: JourneyRootNode | JourneyPatchNode;
|
|
|
|
store.update(s => {
|
|
let parentNode: JourneyNode | null = null
|
|
if (parent != null) {
|
|
if (typeof parent === "object")
|
|
parent = parent.id;
|
|
parentNode = s.nodesByID[parent];
|
|
if (parentNode == null) {
|
|
throw new Error(`Could not find parent node ${parent} to insert into!`)
|
|
}
|
|
}
|
|
if (parentNode == null) {
|
|
_node = {
|
|
id: uuidv4(),
|
|
type: "root",
|
|
children: [],
|
|
base: { ...params }
|
|
}
|
|
s.root = _node
|
|
}
|
|
else {
|
|
_node = {
|
|
id: uuidv4(),
|
|
type: "patch",
|
|
parent: parentNode,
|
|
children: [],
|
|
patch: params,
|
|
}
|
|
parentNode.children.push(_node);
|
|
}
|
|
s.nodesByID[_node.id] = _node;
|
|
s.version += 1;
|
|
return s;
|
|
});
|
|
return _node;
|
|
}
|
|
|
|
function pushPatchOntoActive(workflow: ComfyBoxWorkflow, activeNode?: JourneyNode, showNotification: boolean = false): JourneyNode | null {
|
|
const workflowParams = getWorkflowRestoreParamsFromWorkflow(workflow)
|
|
|
|
let journeyNode
|
|
|
|
if (activeNode == null) {
|
|
// add root node
|
|
if (get(store).root != null) {
|
|
console.debug("[journeyStates] Root already exists")
|
|
return null;
|
|
}
|
|
journeyNode = addNode(workflowParams, null);
|
|
if (showNotification)
|
|
notify("Pushed a new base workflow state.", { type: "info" })
|
|
}
|
|
else {
|
|
// add patch node
|
|
const patch = calculateWorkflowParamsPatch(activeNode, workflowParams);
|
|
const patchedCount = Object.keys(patch).length;
|
|
if (patchedCount === 0) {
|
|
console.debug("[journeyStates] Patch had no diff")
|
|
if (showNotification)
|
|
notify("No changes were made to active parameters yet.", { type: "warning" })
|
|
return null;
|
|
}
|
|
journeyNode = addNode(patch, activeNode);
|
|
if (showNotification)
|
|
notify(`Pushed new state with ${patchedCount} changes.`, { type: "info" })
|
|
}
|
|
|
|
if (journeyNode != null) {
|
|
selectNode(journeyNode);
|
|
}
|
|
|
|
console.debug("[journeyStates] added node", journeyNode)
|
|
return journeyNode;
|
|
}
|
|
|
|
function selectNode(obj?: JourneyNodeID | JourneyNode) {
|
|
store.update(s => {
|
|
if (typeof obj === "string")
|
|
s.activeNodeID = obj;
|
|
else
|
|
s.activeNodeID = obj.id;
|
|
return s;
|
|
})
|
|
}
|
|
|
|
// function removeNode(id: JourneyNodeID) {
|
|
// store.update(s => {
|
|
// const node = s.nodesByID[id];
|
|
// if (node == null) {
|
|
// throw new Error(`Journey node not found: ${id}`)
|
|
// }
|
|
|
|
// if (node.type === "patch") {
|
|
|
|
// }
|
|
// else {
|
|
// s.root = null;
|
|
// }
|
|
|
|
// delete s.nodesByID[id];
|
|
// s.version += 1;
|
|
|
|
// return s;
|
|
// });
|
|
// }
|
|
|
|
function* iterateBreadthFirst(id?: JourneyNodeID | null): Iterable<JourneyNode> {
|
|
const state = get(store);
|
|
|
|
id ||= state.root?.id;
|
|
if (id == null)
|
|
return;
|
|
|
|
const queue = [state.nodesByID[id]];
|
|
while (queue.length > 0) {
|
|
const node = queue.shift();
|
|
yield node;
|
|
if (node.children) {
|
|
for (const child of node.children) {
|
|
queue.push(state.nodesByID[child.id]);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
function* iterateNodeParents(node: JourneyNode): Iterable<JourneyNode> {
|
|
while (isPatch(node)) {
|
|
yield node.parent;
|
|
node = node.parent;
|
|
}
|
|
}
|
|
|
|
function iterateLinearPath(id: JourneyNodeID): Iterable<JourneyNode> {
|
|
const state = get(store);
|
|
|
|
const node = state.nodesByID[id];
|
|
if (node == null) {
|
|
console.error("[journeyStates] Journey node not found!", id);
|
|
return
|
|
}
|
|
|
|
let path = Array.from(iterateNodeParents(node)).reverse()
|
|
path.push(node)
|
|
|
|
// pick first child for nodes downstream
|
|
let child = node.children[0]
|
|
while (child != null) {
|
|
path.push(child);
|
|
child = child.children[0];
|
|
}
|
|
|
|
return path;
|
|
}
|
|
|
|
function onExecuted(promptID: PromptID, nodeID: ComfyNodeID, output: SerializedPromptOutput, queueEntry: QueueEntry) {
|
|
const journeyNode = Array.from(iterateBreadthFirst()).find(j => j.promptID === promptID);
|
|
if (journeyNode == null)
|
|
return;
|
|
|
|
// TODO
|
|
store.update(s => {
|
|
s.version += 1;
|
|
s.activeNodeID = journeyNode.id;
|
|
return s;
|
|
})
|
|
}
|
|
|
|
return {
|
|
...store,
|
|
getActiveNode,
|
|
clear,
|
|
// addNode,
|
|
pushPatchOntoActive,
|
|
selectNode,
|
|
iterateBreadthFirst,
|
|
iterateLinearPath,
|
|
onExecuted
|
|
}
|
|
}
|
|
|
|
export type JourneyStateStaticOps = {
|
|
create: () => WritableJourneyStateStore
|
|
}
|
|
|
|
// These will be attached to workflows.
|
|
const ops: JourneyStateStaticOps = {
|
|
create
|
|
}
|
|
|
|
export default ops
|