From fecd06d3f4cd09ef52433a0a66538b799050aebe Mon Sep 17 00:00:00 2001
From: space-nuko <24979496+space-nuko@users.noreply.github.com>
Date: Thu, 25 May 2023 15:46:04 -0500
Subject: [PATCH] Refactor for built-in templates

---
 litegraph                                |  2 +-
 src/lib/ComfyBoxTemplate.ts              |  1 +
 src/lib/ComfyGraphCanvas.ts              |  2 +-
 src/lib/components/ComfyApp.ts           | 12 ++-
 src/lib/components/ComfyTemplates.svelte | 46 +++++++-----
 src/lib/stores/templateState.ts          | 93 ++++++++++++++++--------
 src/tests/ComfyGraphTests.ts             | 35 ++++++---
 7 files changed, 126 insertions(+), 65 deletions(-)

diff --git a/litegraph b/litegraph
index 8ca6cca..d335948 160000
--- a/litegraph
+++ b/litegraph
@@ -1 +1 @@
-Subproject commit 8ca6cca77709af3644bb9474b28654f0459ec507
+Subproject commit d335948703b8705e7bc2b3a042142442a5d45853
diff --git a/src/lib/ComfyBoxTemplate.ts b/src/lib/ComfyBoxTemplate.ts
index d3ec966..c59b040 100644
--- a/src/lib/ComfyBoxTemplate.ts
+++ b/src/lib/ComfyBoxTemplate.ts
@@ -42,6 +42,7 @@ export type SerializedComfyBoxTemplate = {
     version: 1,
     id: UUID,
     commitHash: string,
+    isBuiltIn?: boolean,
 
     /*
      * Serialized metadata
diff --git a/src/lib/ComfyGraphCanvas.ts b/src/lib/ComfyGraphCanvas.ts
index 8dd4b11..c3b792e 100644
--- a/src/lib/ComfyGraphCanvas.ts
+++ b/src/lib/ComfyGraphCanvas.ts
@@ -489,7 +489,7 @@ export default class ComfyGraphCanvas extends LGraphCanvas {
         const serialized = serializeTemplate(this, template);
 
         try {
-            if (templateState.add(serialized)) {
+            if (templateState.addTemplate(serialized)) {
                 notify("Template saved!", { type: "success" })
             }
             else {
diff --git a/src/lib/components/ComfyApp.ts b/src/lib/components/ComfyApp.ts
index b08c4a5..96e1cda 100644
--- a/src/lib/components/ComfyApp.ts
+++ b/src/lib/components/ComfyApp.ts
@@ -36,7 +36,7 @@ import { type SvelteComponentDev } from "svelte/internal";
 import { get, writable, type Writable } from "svelte/store";
 import ComfyPromptSerializer, { isActiveBackendNode, UpstreamNodeLocator } from "./ComfyPromptSerializer";
 import DanbooruTags from "$lib/DanbooruTags";
-import { deserializeTemplateFromSVG } from "$lib/ComfyBoxTemplate";
+import { deserializeTemplateFromSVG, type SerializedComfyBoxTemplate } from "$lib/ComfyBoxTemplate";
 import templateState from "$lib/stores/templateState";
 
 export const COMFYBOX_SERIAL_VERSION = 1;
@@ -240,7 +240,9 @@ export default class ComfyApp {
         this.addKeyboardHandler();
 
         await this.updateHistoryAndQueue();
-        templateState.load();
+
+        const builtInTemplates = await this.loadBuiltInTemplates();
+        templateState.load(builtInTemplates);
 
         await this.initFrontendFeatures();
 
@@ -271,6 +273,10 @@ export default class ComfyApp {
         }
     }
 
+    async loadBuiltInTemplates(): Promise<SerializedComfyBoxTemplate[]> {
+        return []
+    }
+
     resizeCanvas() {
         if (!this.canvasEl)
             return;
@@ -1062,7 +1068,7 @@ export default class ComfyApp {
 
             const importTemplate = () => {
                 try {
-                    if (templateState.add(templateAndSvg)) {
+                    if (templateState.addTemplate(templateAndSvg)) {
                         notify("Template imported successfully!", { type: "success" })
                     }
                     else {
diff --git a/src/lib/components/ComfyTemplates.svelte b/src/lib/components/ComfyTemplates.svelte
index 3746a65..aae70a4 100644
--- a/src/lib/components/ComfyTemplates.svelte
+++ b/src/lib/components/ComfyTemplates.svelte
@@ -1,6 +1,7 @@
 <script lang="ts">
  import { embedTemplateInSvg, type SerializedComfyBoxTemplate } from "$lib/ComfyBoxTemplate";
- import templateState from "$lib/stores/templateState";
+ import templateState, { type TemplateState, type WritableTemplateStateStore } from "$lib/stores/templateState";
+ import { BoxSeam } from "svelte-bootstrap-icons";
  import uiState from "$lib/stores/uiState";
  import { download, truncateString } from "$lib/utils";
  import type ComfyApp from "./ComfyApp";
@@ -28,10 +29,10 @@
 
  let _sorted: TemplateLayout[] = []
 
- $: rebuildTemplates($templateState.templates);
+ $: $templateState && rebuildTemplates(templateState.getAllTemplates());
 
- function rebuildTemplates(templates: SerializedComfyBoxTemplate[]) {
-     _sorted = Array.from(templates).map(t => {
+ function rebuildTemplates(allTemplates: SerializedComfyBoxTemplate[]): void {
+     _sorted = Array.from(allTemplates).map(t => {
          return {
              type: "template", id: uuidv4(), template: t, attrs: {...defaultWidgetAttributes}, attrsChanged: writable(0)
          }
@@ -88,7 +89,7 @@
      const saveTemplate = (modal: ModalData) => {
          updateTemplate(modal);
          try {
-             templateState.update(layout.template);
+             templateState.updateTemplate(layout.template);
              notify("Saved template!", { type: "success" })
          }
          catch (error) {
@@ -101,7 +102,7 @@
              return false;
 
          try {
-             if (templateState.remove(layout.template.id)) {
+             if (templateState.removeTemplate(layout.template.id)) {
                  notify("Template deleted!", { type: "success" })
              }
              else {
@@ -165,17 +166,17 @@
                 </div>
                 <div class="template-entries-wrapper"
                      use:dndzone={{
-                                 type: "layout",
-                                 items: _sorted,
-                                 flipDurationMs,
-                                 dragDisabled: !draggable,
-                                 dropFromOthersDisabled: true
+                                    type: "layout",
+                                    items: _sorted,
+                                    flipDurationMs,
+                                    dragDisabled: !draggable,
+                                    dropFromOthersDisabled: true
                                  }}
                      on:consider={handleDndConsider}
                      on:finalize={handleDndFinalize}>
                     {#each _sorted.filter(i => i.id !== SHADOW_PLACEHOLDER_ITEM_ID) as item(item.id)}
                         <!-- svelte-ignore a11y-click-events-have-key-events -->
-                        <div class="template-entry" class:draggable on:click={() => handleClick(item)}>
+                        <div class="template-entry" class:built-in={item.template.isBuiltIn} class:draggable on:click={() => handleClick(item)}>
                             <div class="template-name">{item.template.metadata.title}</div>
                             <div class="template-desc">{item.template.metadata.description}</div>
                         </div>
@@ -186,8 +187,13 @@
                 </div>
             </div>
         {:else}
-            <div class="no-templates">
-                <span>(No templates)</span>
+            <div class="no-templates-container">
+                <div class="no-templates-icon">
+                    <BoxSeam width="100%" height="8rem" />
+                </div>
+                <div class="no-templates-message">
+                    (No templates)
+                </div>
             </div>
         {/if}
     </div>
@@ -253,15 +259,19 @@
      }
  }
 
- .no-templates {
+ .no-templates-container {
      display: flex;
      color: var(--comfy-accent-soft);
-     flex-direction: row;
+     flex-direction: column;
      margin: auto;
-     height: 100%;
      color: var(--comfy-accent-soft);
 
-     span {
+     .no-templates-icon {
+         margin: auto;
+         margin-bottom: 1rem;
+     }
+
+     .no-templates-message {
          margin: auto;
          font-size: 32px;
          font-weight: bolder;
diff --git a/src/lib/stores/templateState.ts b/src/lib/stores/templateState.ts
index d2633de..7b7cb81 100644
--- a/src/lib/stores/templateState.ts
+++ b/src/lib/stores/templateState.ts
@@ -5,33 +5,48 @@ import type { Readable, Writable } from 'svelte/store';
 import { v4 as uuidv4 } from "uuid";
 
 export type TemplateState = {
-    templates: SerializedComfyBoxTemplate[]
+    builtInTemplates: SerializedComfyBoxTemplate[]
+    userTemplates: SerializedComfyBoxTemplate[]
     templatesByID: Record<UUID, SerializedComfyBoxTemplate>
 }
 
 type TemplateStateOps = {
+    getAllTemplates: () => SerializedComfyBoxTemplate[],
+    addTemplate: (template: SerializedComfyBoxTemplate) => boolean,
+    updateTemplate: (template: SerializedComfyBoxTemplate) => boolean,
+    removeTemplate: (templateID: UUID) => boolean,
     save: () => void,
-    load: () => void,
-    add: (template: SerializedComfyBoxTemplate) => boolean,
-    update: (template: SerializedComfyBoxTemplate) => boolean,
-    remove: (templateID: UUID) => boolean,
+    load: (builtInTemplates: SerializedComfyBoxTemplate[]) => void,
 }
 
 export type WritableTemplateStateStore = Writable<TemplateState> & TemplateStateOps;
 const store: Writable<TemplateState> = writable(
     {
-        templates: [],
+        builtInTemplates: [],
+        userTemplates: [],
         templatesByID: {}
     })
 
-function add(template: SerializedComfyBoxTemplate): boolean {
+function getTemplateList(template: SerializedComfyBoxTemplate, state: TemplateState): SerializedComfyBoxTemplate[] {
+    if (template.isBuiltIn)
+        return state.builtInTemplates
+    return state.userTemplates;
+}
+
+function getAllTemplates(): SerializedComfyBoxTemplate[] {
+    const state = get(store);
+    return state.builtInTemplates.concat(state.userTemplates);
+}
+
+function addTemplate(template: SerializedComfyBoxTemplate): boolean {
     const state = get(store);
     if (state.templatesByID[template.id]) {
         return false;
     }
 
     store.update(s => {
-        s.templates.push(template);
+        const templateList = getTemplateList(template, s)
+        templateList.push(template)
         s.templatesByID[template.id] = template;
         return s;
     })
@@ -41,15 +56,16 @@ function add(template: SerializedComfyBoxTemplate): boolean {
     return true;
 }
 
-function remove(templateID: UUID): boolean {
+function removeTemplate(templateID: UUID): boolean {
     const state = get(store);
     if (!state.templatesByID[templateID]) {
         return false;
     }
 
     store.update(s => {
-        const index = s.templates.findIndex(t => t.id === templateID)
-        s.templates.splice(index, 1);
+        const templateList = getTemplateList(s.templatesByID[templateID], s)
+        const index = templateList.findIndex(t => t.id === templateID)
+        templateList.splice(index, 1);
         delete s.templatesByID[templateID];
         return s;
     })
@@ -59,7 +75,7 @@ function remove(templateID: UUID): boolean {
     return true;
 }
 
-function update(template: SerializedComfyBoxTemplate): boolean {
+function updateTemplate(template: SerializedComfyBoxTemplate): boolean {
     const state = get(store);
     if (!state.templatesByID[template.id]) {
         return false;
@@ -67,13 +83,14 @@ function update(template: SerializedComfyBoxTemplate): boolean {
 
     store.update(s => {
         const oldId = template.id
-        const index = s.templates.findIndex(t => t.id === oldId)
-        s.templates.splice(index, 1);
+        const templateList = getTemplateList(template, s)
+        const index = templateList.findIndex(t => t.id === oldId)
+        templateList.splice(index, 1);
         delete s.templatesByID[oldId];
 
         template.id = uuidv4();
 
-        s.templates.push(template);
+        templateList.push(template);
         s.templatesByID[template.id] = template;
         return s;
     })
@@ -84,12 +101,26 @@ function update(template: SerializedComfyBoxTemplate): boolean {
 }
 
 function save() {
-    const json = JSON.stringify(get(store).templates)
+    const json = JSON.stringify(get(store).userTemplates)
     localStorage.setItem("templates", json)
     store.set(get(store))
 }
 
-function load() {
+function load(builtInTemplates: SerializedComfyBoxTemplate[]) {
+    store.update(s => {
+        s.userTemplates = []
+        s.templatesByID = {}
+
+        for (const t of builtInTemplates) {
+            t.isBuiltIn = true;
+            s.templatesByID[t.id] = t;
+        }
+
+        s.builtInTemplates = builtInTemplates;
+
+        return s
+    })
+
     const json = localStorage.getItem("templates")
     if (!json) {
         console.info("No templates in local storage, creating store")
@@ -100,18 +131,15 @@ function load() {
     const data = JSON.parse(json) as SerializedComfyBoxTemplate[];
     if (Array.isArray(data)) {
         const templatesByID: Record<UUID, SerializedComfyBoxTemplate> =
-            data.map(d => [d.id, d])
-                .reduce((dict, el: [UUID, SerializedComfyBoxTemplate]) => (dict[el[0]] = el[1], dict), {})
+            data.map(t => {
+                t.isBuiltIn = false;
+                return [t.id, t]
+            }).reduce((dict, el: [UUID, SerializedComfyBoxTemplate]) => (dict[el[0]] = el[1], dict), {})
 
-        store.set({
-            templates: data,
-            templatesByID
-        })
-    }
-    else {
-        store.set({
-            templates: [],
-            templatesByID: {}
+        store.update(s => {
+            s.userTemplates = data
+            s.templatesByID = { ...s.templatesByID, ...templatesByID }
+            return s;
         })
     }
 }
@@ -119,10 +147,11 @@ function load() {
 const templateStateStore: WritableTemplateStateStore =
 {
     ...store,
-    add,
-    remove,
-    update,
+    getAllTemplates,
+    addTemplate,
+    removeTemplate,
+    updateTemplate,
     save,
-    load,
+    load
 }
 export default templateStateStore;
diff --git a/src/tests/ComfyGraphTests.ts b/src/tests/ComfyGraphTests.ts
index 2c741b4..4ce0fbb 100644
--- a/src/tests/ComfyGraphTests.ts
+++ b/src/tests/ComfyGraphTests.ts
@@ -1,16 +1,11 @@
-import { LGraph, LiteGraph, Subgraph, type SlotLayout } from "@litegraph-ts/core"
-import { Watch } from "@litegraph-ts/nodes-basic"
-import { expect } from 'vitest'
-import UnitTest from "./UnitTest"
 import ComfyGraph from "$lib/ComfyGraph";
-import ComfyPromptSerializer from "$lib/components/ComfyPromptSerializer";
-import { ComfyBackendNode } from "$lib/nodes/ComfyBackendNode";
-import ComfyGraphNode from "$lib/nodes/ComfyGraphNode";
-import { graphToGraphVis } from "$lib/utils";
 import { ComfyNumberNode } from "$lib/nodes/widgets";
-import { get } from "svelte/store";
-import layoutStates from "$lib/stores/layoutStates";
 import { ComfyBoxWorkflow } from "$lib/stores/workflowState";
+import { LiteGraph, Subgraph } from "@litegraph-ts/core";
+import { get } from "svelte/store";
+import { expect } from 'vitest';
+import UnitTest from "./UnitTest";
+import { Watch } from "@litegraph-ts/nodes-basic";
 
 export default class ComfyGraphTests extends UnitTest {
     test__onNodeAdded__updatesLayoutState() {
@@ -92,4 +87,24 @@ export default class ComfyGraphTests extends UnitTest {
         expect(Object.keys(state.allItems)).toHaveLength(3)
         expect(Object.keys(state.allItemsByNode)).toHaveLength(0)
     }
+
+    test__serialize__stripsLinkData() {
+        const [{ graph }, layoutState] = ComfyBoxWorkflow.create()
+        layoutState.initDefaultLayout()
+
+        const widget = LiteGraph.createNode(ComfyNumberNode);
+        const watch = LiteGraph.createNode(Watch);
+        graph.add(widget)
+        graph.add(watch)
+
+        widget.connect(0, watch, 0)
+        const link = widget.getOutputLinks(0)[0]
+        widget.setOutputData(0, 42);
+
+        const result = graph.serialize();
+
+        const serNode = result.nodes.find(n => n.id === widget.id);
+
+        expect(serNode.outputs[0]._data).toBeUndefined()
+    }
 }