From 9054d257afd1299c1b8c60227d2dbb0b99ee296e Mon Sep 17 00:00:00 2001
From: space-nuko <24979496+space-nuko@users.noreply.github.com>
Date: Sat, 3 Jun 2023 17:25:24 -0500
Subject: [PATCH] Linear history mode

---
 src/lib/components/ComfyJourneyView.svelte |  11 +-
 src/lib/components/JourneyRenderer.svelte  | 143 +++++++++++++++------
 src/lib/components/PromptDisplay.svelte    |   4 +-
 src/lib/components/graph/GraphStyles.ts    |   5 +-
 src/lib/restoreParameters.ts               |  72 ++++++++++-
 src/lib/stores/journeyStates.ts            |  51 +++++++-
 6 files changed, 228 insertions(+), 58 deletions(-)

diff --git a/src/lib/components/ComfyJourneyView.svelte b/src/lib/components/ComfyJourneyView.svelte
index bbc227a..e4492eb 100644
--- a/src/lib/components/ComfyJourneyView.svelte
+++ b/src/lib/components/ComfyJourneyView.svelte
@@ -12,10 +12,10 @@
  import type { ComfyBoxWorkflow } from '$lib/stores/workflowState';
  import workflowState from '$lib/stores/workflowState';
  import uiState from '$lib/stores/uiState';
- import { calculateWorkflowParamsPatch, resolvePatch, type JourneyPatchNode, type WritableJourneyStateStore, diffParams, JourneyNode } from '$lib/stores/journeyStates';
+ import { resolvePatch, type JourneyPatchNode, type WritableJourneyStateStore, diffParams, type JourneyNode } from '$lib/stores/journeyStates';
  import JourneyRenderer, { type JourneyNodeEvent } from './JourneyRenderer.svelte';
  import { Trash, ClockHistory, Diagram3  } from "svelte-bootstrap-icons";
- import { getWorkflowRestoreParams, getWorkflowRestoreParamsFromWorkflow } from '$lib/restoreParameters';
+ import { getWorkflowRestoreParamsFromWorkflow } from '$lib/restoreParameters';
  import notify from '$lib/notify';
  import selectionState from '$lib/stores/selectionState';
  import { Checkbox } from '@gradio/form';
@@ -41,11 +41,6 @@
  $: workflow = $workflowState.activeWorkflow
  $: {
      journey = workflow?.journey
-     activeNode = null;
-     updateActiveNode();
- }
-
- function updateActiveNode() {
      activeNode = journey?.getActiveNode()
  }
 
@@ -151,7 +146,7 @@
         </button>
     </div>
     {#key $journey.version}
-        <JourneyRenderer {workflow} {journey} {mode} {activeNode}
+        <JourneyRenderer {workflow} {journey} {mode}
                          on:select_node={onSelectNode}
                          on:right_click_node={onRightClickNode}
                          on:hover_node={onHoverNode}
diff --git a/src/lib/components/JourneyRenderer.svelte b/src/lib/components/JourneyRenderer.svelte
index f50e418..a6f440d 100644
--- a/src/lib/components/JourneyRenderer.svelte
+++ b/src/lib/components/JourneyRenderer.svelte
@@ -6,7 +6,7 @@
 </script>
 
 <script lang="ts">
- import type { JourneyNode, JourneyPatchNode, WritableJourneyStateStore } from '$lib/stores/journeyStates';
+ import { resolvePatch, type JourneyNode, type JourneyPatchNode, type WritableJourneyStateStore } from '$lib/stores/journeyStates';
  import { ComfyBoxWorkflow } from '$lib/stores/workflowState';
  import { get } from 'svelte/store';
  import Graph from './graph/Graph.svelte'
@@ -18,11 +18,11 @@
  import { convertComfyOutputToComfyURL, countNewLines } from '$lib/utils';
 	import type { ElementDefinition } from 'cytoscape';
 	import type { JourneyMode } from './ComfyJourneyView.svelte';
+	import type { RestoreParamWorkflowNodeTargets } from '$lib/restoreParameters';
 
  export let workflow: ComfyBoxWorkflow | null = null
  export let journey: WritableJourneyStateStore | null = null
  export let mode: JourneyMode = "linear";
- export let activeNode: JourneyNode | null = null;
 
  const dispatch = createEventDispatcher<{
      select_node: JourneyNodeEvent;
@@ -31,7 +31,6 @@
      hover_node_out: JourneyNodeEvent;
  }>();
 
- let lastSelected = null;
  let lastMode = null;
  let lastVersion = -1;
 
@@ -43,6 +42,34 @@
      lastMode = mode;
  }
 
+ function makePatchText(patch: RestoreParamWorkflowNodeTargets, prev: RestoreParamWorkflowNodeTargets): string {
+     const lines = []
+
+     let sorted = Array.from(Object.entries(patch))
+     sorted.sort((a, b) => {
+         return a[1].name > b[1].name ? 1 : -1
+     })
+
+     for (const [nodeID, source] of sorted) {
+         let line = ""
+         switch (source.nodeType) {
+             case "ui/text":
+                 line = `${source.name} (changed)`
+                 break;
+             default:
+                 const prevValue = prev[nodeID];
+                 let prevValueStr = "???"
+                 if (prevValue)
+                     prevValueStr = prevValue.finalValue
+                 line = `${source.name}: ${prevValueStr} -> ${source.finalValue}`
+                 break;
+         }
+         lines.push(line)
+     }
+
+     return lines.join("\n")
+ }
+
  /*
   * Converts the journey tree into the renderable graph format Cytoscape expects
   */
@@ -51,19 +78,31 @@
          return [[], []]
      }
 
-     const journeyState = get(journey);
+     let activeNode = journey.getActiveNode()
 
      const nodes: ElementDefinition[] = []
      const edges: ElementDefinition[] = []
 
-     for (const node of journey.iterateBreadthFirst()) {
+     let iter: Iterable<JourneyNode> = [];
+     if (mode === "linear") {
+         if (activeNode != null)
+             iter = journey.iterateLinearPath(activeNode.id);
+     }
+     else {
+         iter = journey.iterateBreadthFirst();
+     }
+
+     const showPatches = mode === "linear";
+
+     const memoize = {}
+
+     for (const node of iter) {
          if (node.type === "root") {
              nodes.push({
                  data: {
                      id: node.id,
                      label: "Start",
                  },
-                 selected: node.id === activeNode?.id,
                  classes: "historyNode"
              })
              continue;
@@ -75,44 +114,62 @@
                      id: patchNode.id,
                      label: "P",
                  },
-                 selected: node.id === activeNode?.id,
                  classes: "historyNode"
              })
 
              // Display a small node between with the patch details
              const midNodeID = `${patchNode.id}_patch`;
 
-             const patchText = "cfg: 8 -> 10\nsteps: 20->30\na";
-             const patchNodeHeight = countNewLines(patchText) * 11 + 22;
+             console.debug("get", patchNode);
 
-             nodes.push({
-                 data: {
-                     id: midNodeID,
-                     label: patchText,
-                     patchNodeHeight
-                 },
-                 selectable: false,
-                 classes: "patchNode"
-             })
+             if (showPatches) {
+                 // show a node with the changes between gens
+                 const prev = resolvePatch(patchNode.parent, memoize);
+                 const patchText = makePatchText(patchNode.patch, prev);
+                 const patchNodeHeight = countNewLines(patchText) * 11 + 22;
 
-             edges.push({
-                 data: {
-                     id: `${patchNode.parent.id}_${midNodeID}`,
-                     source: patchNode.parent.id,
-                     target: midNodeID,
-                 },
-                 selectable: false,
-             })
+                 console.debug("[JourneyRenderer] Patch text", prev, patchText);
 
-             edges.push({
-                 data: {
-                     id: `${midNodeID}_${patchNode.id}`,
-                     source: midNodeID,
-                     target: patchNode.id,
-                 },
-                 selectable: false,
-                 locked: true
-             })
+                 nodes.push({
+                     data: {
+                         id: midNodeID,
+                         label: patchText,
+                         patchNodeHeight
+                     },
+                     selectable: false,
+                     classes: "patchNode"
+                 })
+
+                 edges.push({
+                     data: {
+                         id: `${patchNode.parent.id}_${midNodeID}`,
+                         source: patchNode.parent.id,
+                         target: midNodeID,
+                     },
+                     selectable: false,
+                 })
+
+                 edges.push({
+                     data: {
+                         id: `${midNodeID}_${patchNode.id}`,
+                         source: midNodeID,
+                         target: patchNode.id,
+                     },
+                     selectable: false,
+                     locked: true
+                 })
+             }
+             else {
+                 edges.push({
+                     data: {
+                         id: `${patchNode.parent.id}_${patchNode.id}`,
+                         source: patchNode.parent.id,
+                         target: patchNode.id,
+                     },
+                     selectable: false,
+                     locked: true
+                 })
+             }
          }
      }
 
@@ -120,8 +177,9 @@
  }
 
  function onNodeSelected(e: cytoscape.InputEventObject) {
-     console.warn("SELECT", e)
+     console.warn("[JourneyNode] onNodeSelected", e)
      const node = e.target as cytoscape.NodeSingular;
+
      journey.selectNode(node.id());
 
      e.cy.animate({
@@ -154,9 +212,12 @@
  function onRebuilt(e: CustomEvent<{cyto: cytoscape.Core}>) {
      const { cyto } = e.detail;
 
-     for (const node of cyto.nodes().components()) {
+     const activeNode = journey.getActiveNode();
+
+     for (const node of cyto.nodes(".historyNode").components()) {
          const nodeID = node.id()
-         if (nodeID === lastSelected) {
+         if (nodeID === activeNode?.id) {
+             node.select();
              cyto.zoom(1.25);
              cyto.center(node)
          }
@@ -171,7 +232,6 @@
                      if (outputs) {
                          node.data("bgImage", outputs[0]);
                      }
-                     console.warn("node.classes", node.classes())
                  }
              }
          }
@@ -179,8 +239,9 @@
 
      $selectionState.currentPatchHoveredNodes = new Set()
 
-     cyto.nodes()
-         .lock()
+     cyto.nodes().lock()
+
+     cyto.nodes(".historyNode")
          .on("select", onNodeSelected)
          .on("cxttapend ", onNodeRightClicked)
          .on("mouseout", onNodeHoveredOut)
diff --git a/src/lib/components/PromptDisplay.svelte b/src/lib/components/PromptDisplay.svelte
index ab5d3f7..d435fbb 100644
--- a/src/lib/components/PromptDisplay.svelte
+++ b/src/lib/components/PromptDisplay.svelte
@@ -19,7 +19,7 @@
  import ComfyBoxStdPromptSerializer from "$lib/ComfyBoxStdPromptSerializer";
  import JsonView from "./JsonView.svelte";
  import type { ZodError } from "zod";
- import { concatRestoreParams, getWorkflowRestoreParams, type RestoreParamTargets, type RestoreParamWorkflowNodeTargets } from "$lib/restoreParameters";
+ import { concatRestoreParams, getWorkflowRestoreParams, getWorkflowRestoreParamsUsingLayout, type RestoreParamTargets, type RestoreParamWorkflowNodeTargets } from "$lib/restoreParameters";
 
  const splitLength = 50;
 
@@ -40,7 +40,7 @@
 
      // TODO other sources than serialized workflow
      if (workflow != null) {
-         const workflowParams = getWorkflowRestoreParams(workflow.workflow)
+         const workflowParams = getWorkflowRestoreParamsUsingLayout(workflow.workflow, workflow.layout)
          restoreParams = concatRestoreParams(restoreParams, workflowParams);
      }
 
diff --git a/src/lib/components/graph/GraphStyles.ts b/src/lib/components/graph/GraphStyles.ts
index eacd87e..325ca4b 100644
--- a/src/lib/components/graph/GraphStyles.ts
+++ b/src/lib/components/graph/GraphStyles.ts
@@ -57,9 +57,10 @@ const styles: Stylesheet[] = [
     {
         selector: ".patchNode",
         style: {
-            "width": "100",
-            "height": "data(patchNodeHeight)",
+            "width": "label",
+            "height": "label",
             "shape": "round-rectangle",
+            "padding": "20",
             "font-family": "Arial",
             "font-size": "11",
             "font-weight": "normal",
diff --git a/src/lib/restoreParameters.ts b/src/lib/restoreParameters.ts
index c350077..a29be11 100644
--- a/src/lib/restoreParameters.ts
+++ b/src/lib/restoreParameters.ts
@@ -2,7 +2,7 @@ import type { INodeInputSlot, NodeID, SerializedLGraph } from "@litegraph-ts/cor
 import type { SerializedPrompt } from "./components/ComfyApp";
 import type { ComfyWidgetNode } from "./nodes/widgets";
 import type { SerializedComfyWidgetNode } from "./nodes/widgets/ComfyWidgetNode";
-import { isComfyWidgetNode } from "./stores/layoutStates";
+import { isComfyWidgetNode, type SerializedLayoutState } from "./stores/layoutStates";
 import type { ComfyBoxWorkflow } from "./stores/workflowState";
 import { isSerializedPromptInputLink } from "./utils";
 import ComfyBoxStdPromptSerializer from "./ComfyBoxStdPromptSerializer";
@@ -15,6 +15,16 @@ export type RestoreParamType = "workflow" | "backend" | "stdPrompt";
 export interface RestoreParamSource<T extends RestoreParamType = any> {
     type: T,
 
+    /*
+     * A human-readable name for this parameter
+     */
+    name?: string,
+
+    /*
+     * LiteGraph type of the widget node
+     */
+    nodeType: string,
+
     /*
      * The actual value to copy to the widget after all conversions have been
      * applied.
@@ -172,10 +182,22 @@ export function getWorkflowRestoreParamsFromWorkflow(workflow: ComfyBoxWorkflow,
         if (!noExclude && node.properties.excludeFromJourney)
             continue;
 
+        let name = null;
+        const realNode = workflow.graph.getNodeByIdRecursive(node.id);
+        if (realNode != null && isComfyWidgetNode(realNode)) {
+            name = realNode.title || name;
+            const widget = realNode.dragItem;
+            if (widget != null) {
+                name = widget.attrs.title || name;
+            }
+        }
+
         const finalValue = node.getValue();
         if (finalValue != null) {
             const source: RestoreParamSourceWorkflowNode = {
                 type: "workflow",
+                nodeType: node.type,
+                name,
                 finalValue,
             }
             result[node.id] = source;
@@ -185,7 +207,7 @@ export function getWorkflowRestoreParamsFromWorkflow(workflow: ComfyBoxWorkflow,
     return result
 }
 
-export function getWorkflowRestoreParams(serGraph: SerializedLGraph, noExclude: boolean = false): RestoreParamWorkflowNodeTargets {
+export function getWorkflowRestoreParams(serGraph: SerializedLGraph, workflow?: ComfyBoxWorkflow, noExclude: boolean = false): RestoreParamWorkflowNodeTargets {
     const result = {}
 
     for (const node of serGraph.nodes) {
@@ -195,10 +217,22 @@ export function getWorkflowRestoreParams(serGraph: SerializedLGraph, noExclude:
         if (!noExclude && node.properties.excludeFromJourney)
             continue;
 
+        let name = null;
+        const realNode = workflow.graph.getNodeByIdRecursive(node.id);
+        if (realNode != null && isComfyWidgetNode(realNode)) {
+            name = realNode.title || name;
+            const widget = realNode.dragItem;
+            if (widget != null) {
+                name = widget.attrs.title || name;
+            }
+        }
+
         const finalValue = node.comfyValue
         if (finalValue != null) {
             const source: RestoreParamSourceWorkflowNode = {
                 type: "workflow",
+                nodeType: node.type,
+                name,
                 finalValue,
             }
             result[node.id] = source;
@@ -208,6 +242,37 @@ export function getWorkflowRestoreParams(serGraph: SerializedLGraph, noExclude:
     return result
 }
 
+export function getWorkflowRestoreParamsUsingLayout(serGraph: SerializedLGraph, layout?: SerializedLayoutState, noExclude: boolean = false): RestoreParamWorkflowNodeTargets {
+    const result = {}
+
+    for (const serNode of serGraph.nodes) {
+        if (!isSerializedComfyWidgetNode(serNode))
+            continue;
+
+        if (!noExclude && serNode.properties.excludeFromJourney)
+            continue;
+
+        let name = null;
+        const serWidget = Array.from(Object.values(layout?.allItems || {})).find(di => di.dragItem.type === "widget" && di.dragItem.nodeId === serNode.id)
+        if (serWidget) {
+            name = serWidget.dragItem.attrs.title;
+        }
+
+        const finalValue = serNode.comfyValue
+        if (finalValue != null) {
+            const source: RestoreParamSourceWorkflowNode = {
+                type: "workflow",
+                nodeType: serNode.type,
+                name,
+                finalValue,
+            }
+            result[serNode.id] = source;
+        }
+    }
+
+    return result
+}
+
 export function getBackendRestoreParams(workflow: ComfyBoxWorkflow, prompt: SerializedPrompt): Record<NodeID, RestoreParamSourceBackendNodeInput[]> {
     const result = {}
 
@@ -236,6 +301,7 @@ export function getBackendRestoreParams(workflow: ComfyBoxWorkflow, prompt: Seri
                 if (isComfyWidgetNode(foundNode) && foundNode.type === serNode.type) {
                     const source: RestoreParamSourceBackendNodeInput = {
                         type: "backend",
+                        nodeType: foundNode.type,
                         finalValue: inputValue,
                         backendNode: serNode,
                         isDirectAttachment
@@ -252,7 +318,7 @@ export function getBackendRestoreParams(workflow: ComfyBoxWorkflow, prompt: Seri
 export default function getRestoreParameters(workflow: ComfyBoxWorkflow, prompt: SerializedPrompt): RestoreParamTargets {
     const result = {}
 
-    const workflowParams = getWorkflowRestoreParams(prompt.workflow);
+    const workflowParams = getWorkflowRestoreParams(prompt.workflow, workflow);
     concatRestoreParams(result, workflowParams);
 
     const backendParams = getBackendRestoreParams(workflow, prompt);
diff --git a/src/lib/stores/journeyStates.ts b/src/lib/stores/journeyStates.ts
index 0b7ec2e..31fdddf 100644
--- a/src/lib/stores/journeyStates.ts
+++ b/src/lib/stores/journeyStates.ts
@@ -45,17 +45,33 @@ export interface JourneyPatchNode extends JourneyNode {
     patch: RestoreParamWorkflowNodeTargets
 }
 
-export function resolvePatch(node: JourneyNode): 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;
 }
 
@@ -99,6 +115,7 @@ type JourneyStateOps = {
     // 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
 }
@@ -265,9 +282,38 @@ function create() {
         }
     }
 
+    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)
+        if (journeyNode == null)
             return;
 
         // TODO
@@ -286,6 +332,7 @@ function create() {
         pushPatchOntoActive,
         selectNode,
         iterateBreadthFirst,
+        iterateLinearPath,
         onExecuted
     }
 }