codemirror option for text widgets & danbooru tag autocomplete

src/lib/ComfyBoxTemplate.ts

@@ -10,6 +10,7 @@ import { download } from "./utils";
  * components they represent in the UI.
  */
 export type ComfyBoxTemplate = {
+    version: 1,
     nodes: LGraphNode[],
     links: LLink[],
     container?: DragItemEntry
@@ -20,6 +21,8 @@ export type ComfyBoxTemplate = {
  * components they represent in the UI.
  */
 export type SerializedComfyBoxTemplate = {
+    version: 1,
+
     /*
      * Serialized nodes
      */
@@ -270,6 +273,7 @@ export function serializeTemplate(canvas: ComfyGraphCanvas, template: ComfyBoxTe
     [nodes, links] = pruneDetachedLinks(nodes, links);
 
     let comfyBoxTemplate: SerializedComfyBoxTemplate = {
+        version: 1,
         nodes: nodes,
         links: links,
         layout: layout
@@ -320,6 +324,7 @@ export function createTemplate(nodes: LGraphNode[]): ComfyBoxTemplateResult {
         }
 
         return {
+            version: 1,
             nodes: nodes,
             links: links,
             container: container
@@ -328,6 +333,7 @@ export function createTemplate(nodes: LGraphNode[]): ComfyBoxTemplateResult {
     else {
         // No UI to serialize.
         return {
+            version: 1,
             nodes: nodes,
             links: links,
         }

src/lib/DanbooruTags.ts

@@ -0,0 +1,198 @@
+import { parse } from 'csv-parse/browser/esm/sync';
+import { timeExecutionMs } from './utils';
+import { insertCompletionText, type Completion, type CompletionContext, type CompletionResult, type CompletionSource, type CompletionConfig, autocompletion } from '@codemirror/autocomplete';
+import { syntaxTree } from '@codemirror/language';
+import type { Extension, TransactionSpec } from '@codemirror/state';
+import type { EditorView } from '@codemirror/view';
+import type { StyleSpec } from "style-mod"
+
+export enum DanbooruTagCategory {
+    General = 0,
+    Artist = 1,
+    Copyright = 3,
+    Character = 4,
+}
+
+export type DanbooruTagCategoryData = {
+    name: string,
+    color: string
+}
+
+const TAG_CATEGORY_DATA: Record<DanbooruTagCategory, DanbooruTagCategoryData> = {
+    [DanbooruTagCategory.General]: {
+        name: "general",
+        color: "lightblue"
+    },
+    [DanbooruTagCategory.Artist]: {
+        name: "artist",
+        color: "red",
+    },
+    [DanbooruTagCategory.Copyright]: {
+        name: "copyright",
+        color: "lightpurple"
+    },
+    [DanbooruTagCategory.Character]: {
+        name: "character",
+        color: "green"
+    }
+}
+
+export const TAG_CATEGORY_COLORS: StyleSpec = Object.values(TAG_CATEGORY_DATA)
+    .flatMap(d => {
+        return [
+            [`.cm-autocompletion-${d.name}`, { color: d.color + " !important" }],
+        ]
+    })
+    .reduce((dict, el) => (dict[el[0]] = el[1], dict), {})
+
+export type DanbooruTag = {
+    text: string,
+    category: DanbooruTagCategory,
+    count: number,
+    aliases: string[]
+}
+
+function escapeRegExp(s: string) {
+    return s.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
+}
+
+function formatPostCount(postCount: number): string {
+    if (!postCount || !isNaN(postCount))
+        return ""
+
+    let formatter: Intl.NumberFormat;
+
+    // Danbooru formats numbers with a padded fraction for 1M or 1k, but not for 10/100k
+    if (postCount >= 1000000 || (postCount >= 1000 && postCount < 10000))
+        formatter = Intl.NumberFormat("en", { notation: "compact", minimumFractionDigits: 1, maximumFractionDigits: 1 });
+    else
+        formatter = Intl.NumberFormat("en", { notation: "compact" });
+
+    return formatter.format(postCount);
+}
+
+export default class DanbooruTags {
+    private static _instance: DanbooruTags;
+
+    tags: DanbooruTag[] = [];
+    tagsByCategory: Record<string, DanbooruTag> = {}
+    aliases: Record<string, number> = {}
+
+    static get instance(): DanbooruTags {
+        if (!DanbooruTags._instance)
+            DanbooruTags._instance = new DanbooruTags()
+        return DanbooruTags._instance
+    }
+
+    async load(force: boolean = false) {
+        console.log("Parsing danbooru tags CSV...")
+
+        if (this.tags.length > 0 && !force) {
+            console.info("Danbooru tags already parsed")
+            return;
+        }
+
+        this.tags = []
+
+        const transformTag = (data: string[]): DanbooruTag => {
+            return {
+                text: data[0],
+                category: parseInt(data[1]),
+                count: parseInt(data[2]),
+                aliases: data[3].split(",")
+            }
+        }
+
+        const time = await timeExecutionMs(async () => {
+            const resp = await fetch("/extra/danbooru.csv");
+            const csv = await resp.text();
+            const raw: string[][] = parse(csv, {
+                delimiter: ','
+            });
+
+            const refined = raw.map(transformTag);
+
+            this.tags = refined
+        })
+
+        console.log(`Parsed ${this.tags.length} tags in ${time / 1000}ms.`)
+        console.error(this.tags[0])
+    }
+
+    autocomplete(context: CompletionContext): CompletionResult {
+        let nodeBefore = syntaxTree(context.state).resolveInner(context.pos, -1)
+        let textBefore = context.state.sliceDoc(nodeBefore.from, context.pos)
+
+        let weightBefore = /:[0-9.]+$/.exec(textBefore)
+        if (weightBefore) return null
+
+        let tagBefore = /\b[a-zA-Z0-9_()-]+$/.exec(textBefore)
+        if (!tagBefore) return null
+
+        let tagword = tagBefore[0]
+        console.warn(tagword)
+
+        let searchRegex: RegExp;
+        if (tagword.startsWith("*")) {
+            tagword = tagword.slice(1);
+            searchRegex = new RegExp(`${escapeRegExp(tagword)}`, 'i');
+        } else {
+            searchRegex = new RegExp(`(^|[^a-zA-Z])${escapeRegExp(tagword)}`, 'i');
+        }
+
+        const sanitize = (rawTag: string): string => {
+            let sanitized = rawTag.replaceAll("_", " ");
+
+            // TODO config
+            const escapeParentheses = true;
+            const isTagType = true;
+            if (escapeParentheses && isTagType) {
+                sanitized = sanitized
+                    .replaceAll("(", "\\(")
+                    .replaceAll(")", "\\)")
+                    .replaceAll("[", "\\[")
+                    .replaceAll("]", "\\]");
+            }
+
+            return sanitized
+        }
+
+        const apply = (view: EditorView, completion: Completion, from: number, to: number) => {
+            const sanitized = sanitize(completion.label);
+            view.dispatch(insertCompletionText(view.state, sanitized, from, to));
+        }
+
+        const filter = (x: DanbooruTag) => x.text.toLowerCase().search(searchRegex) > -1;
+
+        const options: Completion[] = this.tags.filter(filter).map(t => {
+            const categoryName = TAG_CATEGORY_DATA[t.category]?.name || "unknown";
+            return {
+                label: t.text,
+                apply,
+                detail: formatPostCount(t.count),
+                type: categoryName,
+                section: "Tags"
+            }
+        })
+
+        return {
+            from: tagBefore ? nodeBefore.from + tagBefore.index : context.pos,
+            options,
+            validFor: /^\b([\w_()-]+)?$/
+        }
+    }
+
+    static getCompletionExt(): Extension {
+        const source: CompletionSource = DanbooruTags.instance.autocomplete.bind(DanbooruTags.instance)
+
+        const optionClass = (completion: Completion): string => {
+            return `cm-autocompletion-${completion.type}`
+        }
+
+        return autocompletion({
+            override: [source],
+            interactionDelay: 250,
+            optionClass
+        })
+    }
+}

src/lib/ImageViewer.ts

@@ -6,7 +6,7 @@ export class ImageViewer {
     currentImages: string[] = []
     selectedIndex: number = -1;
     currentGallery: HTMLDivElement | null = null;
-    static _instance: ImageViewer;
+    private static _instance: ImageViewer;
 
     static get instance(): ImageViewer {
         if (!ImageViewer._instance)

src/lib/components/ComfyApp.ts

@@ -34,6 +34,7 @@ import { tick } from "svelte";
 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";
 
 export const COMFYBOX_SERIAL_VERSION = 1;
 
@@ -232,6 +233,8 @@ export default class ComfyApp {
 
         await this.updateHistoryAndQueue();
 
+        await this.initFrontendFeatures();
+
         // await this.#invokeExtensionsAsync("setup");
 
         // Ensure the canvas fills the window
@@ -586,6 +589,10 @@ export default class ComfyApp {
         });
     }
 
+    private async initFrontendFeatures() {
+        await DanbooruTags.instance.load();
+    }
+
     private async updateHistoryAndQueue() {
         const queue = await this.api.getQueue();
         const history = await this.api.getHistory();

src/lib/init.ts

@@ -3,8 +3,9 @@ import { LGraphCanvas, LiteGraph, Subgraph } from '@litegraph-ts/core';
 import layoutStates from './stores/layoutStates';
 import { get } from 'svelte/store';
 import workflowState from './stores/workflowState';
+import DanbooruTags from './DanbooruTags';
 
-export function configureLitegraph(isMobile: boolean = false) {
+function configureLitegraph(isMobile: boolean = false) {
     LiteGraph.catch_exceptions = false;
 
     // Must be enabled, otherwise subgraphs won't work (because of non-unique node/link IDs)
@@ -26,10 +27,18 @@ export function configureLitegraph(isMobile: boolean = false) {
     }
 
     Subgraph.default_lgraph_factory = () => new ComfyGraph;
-
-    (window as any).LiteGraph = LiteGraph;
-    (window as any).LGraphCanvas = LGraphCanvas;
-    (window as any).layoutStates = layoutStates;
-    (window as any).workflowState = workflowState;
-    (window as any).svelteGet = get;
+}
+
+function configureGlobals() {
+    const win = window as any
+    win.LiteGraph = LiteGraph;
+    win.LGraphCanvas = LGraphCanvas;
+    win.layoutStates = layoutStates;
+    win.workflowState = workflowState;
+    win.svelteGet = get;
+}
+
+export default async function init(isMobile: boolean = false) {
+    configureLitegraph(isMobile);
+    configureGlobals();
 }

src/lib/stores/layoutStates.ts

@@ -405,7 +405,7 @@ const ALL_ATTRIBUTES: AttributesSpecList = [
                 defaultValue: ""
             },
 
-            // Editor
+            // Image Editor
             {
                 name: "variant",
                 type: "enum",
@@ -430,6 +430,16 @@ const ALL_ATTRIBUTES: AttributesSpecList = [
             },
 
             // Text
+            {
+                name: "variant",
+                type: "enum",
+                location: "widget",
+                editable: true,
+                validNodeTypes: ["ui/text"],
+                values: ["text", "code"],
+                defaultValue: "text",
+                refreshPanelOnChange: true
+            },
             {
                 name: "multiline",
                 type: "boolean",

src/lib/utils.ts

@@ -52,6 +52,14 @@ export function* enumerate<T>(iterable: Iterable<T>): Iterable<[number, T]> {
     }
 }
 
+export async function timeExecutionMs(fn: (...any) => Promise<any>, ...args: any[]): Promise<number> {
+    const start = new Date().getTime();
+
+    await fn.apply(null, args)
+
+    return new Date().getTime() - start;
+}
+
 export function download(filename: string, text: string, type: string = "text/plain") {
     const blob = new Blob([text], { type: type });
     const url = URL.createObjectURL(blob);

src/lib/widgets/TextWidget.svelte

@@ -6,6 +6,7 @@
  import type { ComfyTextNode } from "$lib/nodes/widgets";
  export let widget: WidgetLayout | null = null;
  export let isMobile: boolean = false;
+ import TextWidgetCodeVariant from "./TextWidgetCodeVariant.svelte"
 
  let node: ComfyTextNode | null = null;
  let nodeValue: Writable<string> | null = null;
@@ -31,18 +32,22 @@
 
 <div class="wrapper gradio-textbox">
     {#if node !== null && nodeValue !== null}
-        <TextBox
-            bind:value={$nodeValue}
-            label={widget.attrs.title}
-            disabled={isDisabled(widget)}
-            lines={node.properties.multiline ? node.properties.lines : 1}
-            max_lines={node.properties.multiline ? node.properties.maxLines : 1}
-            show_label={widget.attrs.title !== ""}
-            on:change
-            on:submit
-            on:blur
-            on:select
-        />
+        {#if widget.attrs.variant === "code"}
+            <TextWidgetCodeVariant {widget} {node} {nodeValue} />
+        {:else}
+            <TextBox
+                bind:value={$nodeValue}
+                label={widget.attrs.title}
+                disabled={isDisabled(widget)}
+                lines={node.properties.multiline ? node.properties.lines : 1}
+                max_lines={node.properties.multiline ? node.properties.maxLines : 1}
+                show_label={widget.attrs.title !== ""}
+                on:change
+                on:submit
+                on:blur
+                on:select
+            />
+        {/if}
     {/if}
 </div>
 

src/lib/widgets/TextWidgetCodeVariant.svelte

@@ -0,0 +1,233 @@
+<script lang="ts">
+ import type { ComfyTextNode } from "$lib/nodes/widgets";
+ import { type WidgetLayout } from "$lib/stores/layoutStates";
+ import { writable, type Writable } from "svelte/store";
+
+ import type { ViewUpdate } from "@codemirror/view";
+ import { EditorView, keymap, placeholder as placeholderExt } from "@codemirror/view";
+ import { StateEffect, EditorState, type Extension } from "@codemirror/state";
+ import { basicDark } from "cm6-theme-basic-dark";
+ import { basicLight } from "cm6-theme-basic-light";
+
+ import { basicSetup } from "./TextWidgetCodeVariant";
+	import { createEventDispatcher, onMount } from "svelte";
+	import { TAG_CATEGORY_COLORS } from "$lib/DanbooruTags";
+
+ export let widget: WidgetLayout;
+ export let node: ComfyTextNode;
+ export let nodeValue: Writable<string> = writable("");
+ export let extraExtensions: Extension[] = [];
+ let lines = 5;
+ let classNames = ""
+
+ let element: HTMLDivElement;
+ let view: EditorView;
+
+ const dispatch = createEventDispatcher<{ change: string }>();
+
+ $: lines = node?.properties?.lines || 5;
+
+ let BaseTheme: Extension = EditorView.theme({
+     "&": {
+         width: "100%",
+         maxWidth: "100%",
+         height: "12rem",
+         fontSize: "var(--text-sm)",
+         backgroundColor: "var(--input-background-fill)"
+     },
+     ".cm-content": {
+         paddingTop: "5px",
+         paddingBottom: "5px",
+         color: "var(--body-text-color)",
+         fontFamily: "var(--font-mono)",
+         minHeight: "100%"
+     },
+     ".cm-gutters": {
+         marginRight: "1px",
+         borderRight: "1px solid var(--border-color-primary)",
+         backgroundColor: "transparent",
+         color: "var(--body-text-color-subdued)"
+     },
+     ".cm-focused": {
+         outline: "none"
+     },
+     ".cm-scroller": {
+         height: "auto"
+     },
+     ".cm-cursor": {
+         borderLeftColor: "var(--body-text-color)"
+     },
+     ".cm-selectionBackground": {
+         backgroundColor: "var(--secondary-600) !important",
+     },
+     ".cm-tooltip": {
+         backgroundColor: "var(--panel-background-fill) !important",
+         border: "1px solid var(--panel-border-color) !important",
+     },
+     ".cm-tooltip-autocomplete": {
+         color: "var(--body-text-color) !important",
+     },
+     ".cm-tooltip-autocomplete > ul > li[aria-selected]": {
+         color: "unset"
+     },
+     ...TAG_CATEGORY_COLORS
+ })
+
+ onMount(() => {
+     view = createEditorView();
+     return () => view?.destroy();
+ });
+
+ $: reconfigure()
+ $: setDoc($nodeValue);
+ $: updateLines(lines);
+
+ function reconfigure(): void {
+     view?.dispatch({
+         effects: StateEffect.reconfigure.of(getExtensions())
+     });
+ }
+
+ function setDoc(newDoc: string) {
+     if (view && newDoc !== view.state.doc.toString()) {
+         view.dispatch({
+             changes: {
+                 from: 0,
+                 to: view.state.doc.length,
+                 insert: newDoc
+             }
+         });
+     }
+ }
+
+ function updateLines(newLines: number) {
+     if (view) {
+         view.requestMeasure({ read: updateGutters });
+     }
+ }
+
+ function getGutterLineHeight(view: EditorView): string | null {
+     let elements = view.dom.querySelectorAll<HTMLElement>(".cm-gutterElement");
+     if (elements.length === 0) {
+         return null;
+     }
+     for (var i = 0; i < elements.length; i++) {
+         let node = elements[i];
+         let height = getComputedStyle(node)?.height ?? "0px";
+         if (height != "0px") {
+             return height;
+         }
+     }
+     return null;
+ }
+
+ function updateGutters(view: EditorView): any {
+     let gutters = view.dom.querySelectorAll<HTMLElement>(".cm-gutter");
+     let _lines = lines + 1;
+     let lineHeight = getGutterLineHeight(view);
+     if (!lineHeight) {
+         return null;
+     }
+     // for (var i = 0; i < gutters.length; i++) {
+     //     let node = gutters[i];
+     //     node.style.minHeight = `calc(${lineHeight} * ${_lines})`;
+     // }
+     return null;
+ }
+
+ function handleChange(vu: ViewUpdate): void {
+     if (vu.docChanged) {
+         const doc = vu.state.doc;
+         const text = doc.toString();
+         $nodeValue = text;
+         dispatch("change", text);
+     }
+     view.requestMeasure({ read: updateGutters });
+ }
+
+ function getBaseExtensions(readonly: boolean, placeholder: string | null): Extension[] {
+     const extensions: Extension[] = [
+         EditorView.editable.of(!readonly),
+         EditorState.readOnly.of(readonly)
+     ];
+
+     extensions.push(basicSetup);
+
+     if (placeholder) {
+         extensions.push(placeholderExt(placeholder));
+     }
+
+     extensions.push(EditorView.updateListener.of(handleChange));
+
+     return extensions
+ }
+
+ function getTheme(dark_mode: boolean): Extension[] {
+     const extensions: Extension[] = [];
+
+     if (dark_mode) {
+         extensions.push(basicDark);
+     } else {
+         extensions.push(basicLight);
+     }
+     return extensions;
+ }
+
+ function getExtensions(): Extension[] {
+     // TODO
+     const readonly = false;
+     const placeholder = "Placeholder..."
+     const dark_mode = true;
+
+     const stateExtensions: Extension[] = [
+         ...getBaseExtensions(
+             readonly,
+             placeholder,
+         ),
+         BaseTheme,
+         ...getTheme(dark_mode),
+         ...extraExtensions
+     ];
+     return stateExtensions;
+ }
+
+ function createEditorState(value: string | null | undefined): EditorState {
+     return EditorState.create({
+         doc: value ?? undefined,
+         extensions: getExtensions()
+     });
+ }
+
+ function createEditorView(): EditorView {
+     return new EditorView({
+         parent: element,
+         state: createEditorState($nodeValue)
+     });
+ }
+</script>
+
+<div class="wrap">
+    <div class="codemirror-wrapper {classNames}" bind:this={element} />
+</div>
+<!-- <CodeMirror bind:value={$nodeValue} {styles} /> -->
+
+<style lang="scss">
+ .wrap {
+     display: flex;
+     flex-direction: column;
+     flex-flow: column;
+     margin: 0;
+     padding: 0;
+     height: 100%;
+     border: 1px solid var(--input-border-color);
+
+     .codemirror-wrapper {
+         height: 100%;
+         overflow: auto;
+     }
+ }
+
+ :global(.cm-scroller) {
+     height: 100% !important;
+ }
+</style>

src/lib/widgets/TextWidgetCodeVariant.ts

@@ -0,0 +1,55 @@
+import type { Extension } from "@codemirror/state";
+import {
+    lineNumbers,
+    highlightSpecialChars,
+    drawSelection,
+    rectangularSelection,
+    crosshairCursor,
+    keymap
+} from "@codemirror/view";
+import { EditorView } from "@codemirror/view";
+import { EditorState } from "@codemirror/state";
+import {
+    foldGutter,
+    indentOnInput,
+    syntaxHighlighting,
+    defaultHighlightStyle,
+    foldKeymap
+} from "@codemirror/language";
+import { history, defaultKeymap, historyKeymap } from "@codemirror/commands";
+import {
+    closeBrackets,
+    closeBracketsKeymap,
+    completionKeymap
+} from "@codemirror/autocomplete";
+import { lintKeymap } from "@codemirror/lint";
+import {
+    type CompletionSource, autocompletion, CompletionContext, startCompletion,
+    currentCompletions, completionStatus, completeFromList, acceptCompletion
+} from "@codemirror/autocomplete"
+import DanbooruTags from "$lib/DanbooruTags";
+
+export const basicSetup: Extension = /*@__PURE__*/ (() => [
+    lineNumbers(),
+    highlightSpecialChars(),
+    history(),
+    foldGutter(),
+    drawSelection(),
+    EditorState.allowMultipleSelections.of(true),
+    indentOnInput(),
+    syntaxHighlighting(defaultHighlightStyle, { fallback: true }),
+    closeBrackets(),
+    rectangularSelection(),
+    crosshairCursor(),
+    EditorView.lineWrapping,
+    DanbooruTags.getCompletionExt(),
+
+    keymap.of([
+        ...closeBracketsKeymap,
+        ...defaultKeymap,
+        ...historyKeymap,
+        ...foldKeymap,
+        ...completionKeymap,
+        ...lintKeymap
+    ])
+])();