Audio node widgets (#54)

* Add audio widgets

* Change rendering order

* Proper output node

* Change indent

* Fix param

* Hide audio widget when output is empty

* Load default value

* populate audio widget from history

* Prevent init if no audio selection

public/style.css

@@ -557,3 +557,7 @@ dialog::backdrop {
 		border-top: none;
 	}
 }
+
+audio.comfy-audio.empty-audio-widget {
+	display: none;
+}

src/extensions/core/index.js

@@ -21,3 +21,4 @@ import "./undoRedo";
 import "./uploadImage";
 import "./webcamCapture";
 import "./widgetInputs";
+import "./uploadAudio";

src/extensions/core/uploadAudio.ts

@@ -0,0 +1,164 @@
+import { app } from "../../scripts/app";
+import { api } from "../../scripts/api";
+import type { IWidget } from "/types/litegraph";
+import type { DOMWidget } from "/scripts/domWidget";
+
+type FolderType = "input" | "output" | "temp";
+
+function splitFilePath(path: string): [string, string] {
+  const folder_separator = path.lastIndexOf("/");
+  if (folder_separator === -1) {
+    return ["", path];
+  }
+  return [path.substring(0, folder_separator), path.substring(folder_separator + 1)];
+}
+
+function getResourceURL(subfolder: string, filename: string, type: FolderType = "input"): string {
+  const params = [
+    "filename=" + encodeURIComponent(filename),
+    "type=" + type,
+    "subfolder=" + subfolder,
+    app.getPreviewFormatParam().substring(1),
+    app.getRandParam().substring(1),
+  ].join("&");
+
+  return `/view?${params}`;
+}
+
+async function uploadFile(
+  audioWidget: IWidget,
+  audioUIWidget: DOMWidget<HTMLAudioElement>,
+  file: File,
+  updateNode: boolean,
+  pasted: boolean = false
+) {
+  try {
+    // Wrap file in formdata so it includes filename
+    const body = new FormData();
+    body.append("image", file);
+    if (pasted) body.append("subfolder", "pasted");
+    const resp = await api.fetchApi("/upload/image", {
+      method: "POST",
+      body,
+    });
+
+    if (resp.status === 200) {
+      const data = await resp.json();
+      // Add the file to the dropdown list and update the widget value
+      let path = data.name;
+      if (data.subfolder) path = data.subfolder + "/" + path;
+
+      if (!audioWidget.options.values.includes(path)) {
+        audioWidget.options.values.push(path);
+      }
+
+      if (updateNode) {
+        audioUIWidget.element.src = api.apiURL(getResourceURL(...splitFilePath(path)));
+        audioWidget.value = path;
+      }
+    } else {
+      alert(resp.status + " - " + resp.statusText);
+    }
+  } catch (error) {
+    alert(error);
+  }
+}
+
+// AudioWidget MUST be registered first, as AUDIOUPLOAD depends on AUDIO_UI to be
+// present.
+app.registerExtension({
+  name: "Comfy.AudioWidget",
+  async beforeRegisterNodeDef(nodeType, nodeData) {
+    if (["LoadAudio", "SaveAudio"].includes(nodeType.comfyClass)) {
+      nodeData.input.required.audioUI = ["AUDIO_UI"];
+    }
+  },
+  getCustomWidgets() {
+    return {
+      AUDIO_UI(node, inputName: string) {
+        const audio = document.createElement("audio");
+        audio.controls = true;
+        audio.classList.add("comfy-audio");
+        audio.setAttribute("name", "media");
+
+        const audioUIWidget: DOMWidget<HTMLAudioElement> = node.addDOMWidget(inputName, /* name=*/ "audioUI", audio);
+        // @ts-ignore
+        // TODO: Sort out the DOMWidget type.
+        audioUIWidget.serialize = false;
+
+        const isOutputNode = node.constructor.nodeData.output_node;
+        if (isOutputNode) {
+          // Hide the audio widget when there is no audio initially.
+          audioUIWidget.element.classList.add("empty-audio-widget");
+          // Populate the audio widget UI on node execution.
+          const onExecuted = node.onExecuted;
+          node.onExecuted = function (message) {
+            onExecuted?.apply(this, arguments);
+            const audios = message.audio;
+            if (!audios) return;
+            const audio = audios[0];
+            audioUIWidget.element.src = api.apiURL(getResourceURL(audio.subfolder, audio.filename, "output"));
+            audioUIWidget.element.classList.remove("empty-audio-widget");
+          }
+        }
+        return { widget: audioUIWidget };
+      }
+    }
+  },
+  onNodeOutputsUpdated(nodeOutputs: Record<number, any>) {
+    for (const [nodeId, output] of Object.entries(nodeOutputs)) {
+      const node = app.graph.getNodeById(Number.parseInt(nodeId));
+      if ("audio" in output) {
+        const audioUIWidget = node.widgets.find((w) => w.name === "audioUI") as unknown as DOMWidget<HTMLAudioElement>;
+        const audio = output.audio[0];
+        audioUIWidget.element.src = api.apiURL(getResourceURL(audio.subfolder, audio.filename, "output"));
+        audioUIWidget.element.classList.remove("empty-audio-widget");
+      }
+    }
+  },
+});
+
+app.registerExtension({
+  name: "Comfy.UploadAudio",
+  async beforeRegisterNodeDef(nodeType, nodeData) {
+    if (nodeData?.input?.required?.audio?.[1]?.audio_upload === true) {
+      nodeData.input.required.upload = ["AUDIOUPLOAD"];
+    }
+  },
+  getCustomWidgets() {
+    return {
+      AUDIOUPLOAD(node, inputName: string) {
+        // The widget that allows user to select file.
+        const audioWidget: IWidget = node.widgets.find((w: IWidget) => w.name === "audio");
+        const audioUIWidget: DOMWidget<HTMLAudioElement> = node.widgets.find((w: IWidget) => w.name === "audioUI");
+
+        const onAudioWidgetUpdate = () => {
+          audioUIWidget.element.src = api.apiURL(getResourceURL(...splitFilePath(audioWidget.value)));
+        };
+        // Initially load default audio file to audioUIWidget.
+        if (audioWidget.value) {
+          onAudioWidgetUpdate();
+        }
+        audioWidget.callback = onAudioWidgetUpdate;
+
+        const fileInput = document.createElement("input");
+        fileInput.type = "file";
+        fileInput.accept = "audio/*";
+        fileInput.style.display = "none";
+        fileInput.onchange = () => {
+          if (fileInput.files.length) {
+            uploadFile(audioWidget, audioUIWidget, fileInput.files[0], true);
+          }
+        };
+        // The widget to pop up the upload dialog.
+        const uploadWidget = node.addWidget("button", inputName, /* value=*/"", () => {
+          fileInput.click();
+        });
+        uploadWidget.label = "choose file to upload";
+        uploadWidget.serialize = false;
+
+        return { widget: uploadWidget };
+      }
+    }
+  },
+});

src/scripts/app.ts

@@ -65,7 +65,7 @@ export class ComfyApp {
 	ui: ComfyUI;
 	logging: ComfyLogging;
 	extensions: ComfyExtension[];
-	nodeOutputs: Record<string, any>;
+	_nodeOutputs: Record<string, any>;
 	nodePreviewImages: Record<string, typeof Image>;
 	shiftDown: boolean;
 	graph: LGraph;
@@ -116,6 +116,15 @@ export class ComfyApp {
 		this.shiftDown = false;
 	}
 
+	get nodeOutputs() {
+		return this._nodeOutputs;
+	}
+
+	set nodeOutputs(value) {
+		this._nodeOutputs = value;
+		this.#invokeExtensions("onNodeOutputsUpdated", value);
+	}
+
 	getPreviewFormatParam() {
 		let preview_format = this.ui.settings.getSettingValue("Comfy.PreviewFormat");
 		if(preview_format)

src/scripts/domWidget.ts

@@ -12,11 +12,11 @@ interface Rect {
 	y: number;
 }
 
-interface DOMWidget {
+export interface DOMWidget<T = HTMLElement> {
 	type: string;
 	name: string;
 	computedHeight?: number;
-	element?: HTMLElement;
+	element?: T;
 	options: any;
 	value?: any;
 	y?: number;

src/scripts/widgets.ts

@@ -400,7 +400,7 @@ export const ComfyWidgets: Record<string, ComfyWidgetConstructor> = {
 		}
 		return res;
 	},
-	IMAGEUPLOAD(node: LGraphNode, inputName, inputData, app) {
+	IMAGEUPLOAD(node: LGraphNode, inputName: string, inputData, app) {
 		// TODO make image upload handle a custom node type?
 		// @ts-ignore
 		const imageWidget = node.widgets.find((w) => w.name === (inputData[1]?.widget ?? "image"));