Preview sampled images with TAESD

web/scripts/api.js

@@ -42,6 +42,7 @@ class ComfyApi extends EventTarget {
 		this.socket = new WebSocket(
 			`ws${window.location.protocol === "https:" ? "s" : ""}://${location.host}/ws${existingSession}`
 		);
+		this.socket.binaryType = "arraybuffer";
 
 		this.socket.addEventListener("open", () => {
 			opened = true;
@@ -70,39 +71,66 @@ class ComfyApi extends EventTarget {
 
 		this.socket.addEventListener("message", (event) => {
 			try {
-				const msg = JSON.parse(event.data);
-				switch (msg.type) {
-					case "status":
-						if (msg.data.sid) {
-							this.clientId = msg.data.sid;
-							window.name = this.clientId;
+				if (event.data instanceof ArrayBuffer) {
+					const view = new DataView(event.data);
+					const eventType = view.getUint32(0);
+					const buffer = event.data.slice(4);
+					console.error("BINARY", eventType);
+					switch (eventType) {
+					case 1:
+						const view2 = new DataView(event.data);
+						const imageType = view2.getUint32(0)
+						let imageMime
+						switch (imageType) {
+							case 1:
+							default:
+								imageMime = "image/jpeg";
+								break;
+							case 2:
+								imageMime = "image/png"
 						}
-						this.dispatchEvent(new CustomEvent("status", { detail: msg.data.status }));
-						break;
-					case "progress":
-						this.dispatchEvent(new CustomEvent("progress", { detail: msg.data }));
-						break;
-					case "executing":
-						this.dispatchEvent(new CustomEvent("executing", { detail: msg.data.node }));
-						break;
-					case "executed":
-						this.dispatchEvent(new CustomEvent("executed", { detail: msg.data }));
-						break;
-					case "execution_start":
-						this.dispatchEvent(new CustomEvent("execution_start", { detail: msg.data }));
-						break;
-					case "execution_error":
-						this.dispatchEvent(new CustomEvent("execution_error", { detail: msg.data }));
+						const jpegBlob = new Blob([buffer.slice(4)], { type: imageMime });
+						this.dispatchEvent(new CustomEvent("b_preview", { detail: jpegBlob }));
 						break;
 					default:
-						if (this.#registered.has(msg.type)) {
-							this.dispatchEvent(new CustomEvent(msg.type, { detail: msg.data }));
-						} else {
-							throw new Error("Unknown message type");
-						}
+						throw new Error(`Unknown binary websocket message of type ${eventType}`);
+					}
+				}
+				else {
+				    const msg = JSON.parse(event.data);
+				    switch (msg.type) {
+					    case "status":
+						    if (msg.data.sid) {
+							    this.clientId = msg.data.sid;
+							    window.name = this.clientId;
+						    }
+						    this.dispatchEvent(new CustomEvent("status", { detail: msg.data.status }));
+						    break;
+					    case "progress":
+						    this.dispatchEvent(new CustomEvent("progress", { detail: msg.data }));
+						    break;
+					    case "executing":
+						    this.dispatchEvent(new CustomEvent("executing", { detail: msg.data.node }));
+						    break;
+					    case "executed":
+						    this.dispatchEvent(new CustomEvent("executed", { detail: msg.data }));
+						    break;
+					    case "execution_start":
+						    this.dispatchEvent(new CustomEvent("execution_start", { detail: msg.data }));
+						    break;
+					    case "execution_error":
+						    this.dispatchEvent(new CustomEvent("execution_error", { detail: msg.data }));
+						    break;
+					    default:
+						    if (this.#registered.has(msg.type)) {
+							    this.dispatchEvent(new CustomEvent(msg.type, { detail: msg.data }));
+						    } else {
+							    throw new Error(`Unknown message type ${msg.type}`);
+						    }
+				    }
 				}
 			} catch (error) {
-				console.warn("Unhandled message:", event.data);
+				console.warn("Unhandled message:", event.data, error);
 			}
 		});
 	}

web/scripts/app.js

@@ -44,6 +44,12 @@ export class ComfyApp {
 		 */
 		this.nodeOutputs = {};
 
+		/**
+		 * Stores the preview image data for each node
+		 * @type {Record<string, Image>}
+		 */
+		this.nodePreviewImages = {};
+
 		/**
 		 * If the shift key on the keyboard is pressed
 		 * @type {boolean}
@@ -367,29 +373,52 @@ export class ComfyApp {
 
 		node.prototype.onDrawBackground = function (ctx) {
 			if (!this.flags.collapsed) {
+				let imgURLs = []
+				let imagesChanged = false
+
 				const output = app.nodeOutputs[this.id + ""];
 				if (output && output.images) {
 					if (this.images !== output.images) {
 						this.images = output.images;
-						this.imgs = null;
-						this.imageIndex = null;
+						imagesChanged = true;
+						imgURLs = imgURLs.concat(output.images.map(params => {
+							return "/view?" + new URLSearchParams(src).toString() + app.getPreviewFormatParam();
+						}))
+					}
+				}
+
+				const preview = app.nodePreviewImages[this.id + ""]
+				if (this.preview !== preview) {
+					this.preview = preview
+					imagesChanged = true;
+					if (preview != null) {
+						imgURLs.push(preview);
+					}
+				}
+
+				if (imagesChanged) {
+					this.imageIndex = null;
+					if (imgURLs.length > 0) {
 						Promise.all(
-							output.images.map((src) => {
+							imgURLs.map((src) => {
 								return new Promise((r) => {
 									const img = new Image();
 									img.onload = () => r(img);
 									img.onerror = () => r(null);
-									img.src = "/view?" + new URLSearchParams(src).toString() + app.getPreviewFormatParam();
+									img.src = src
 								});
 							})
 						).then((imgs) => {
-							if (this.images === output.images) {
+							if ((!output || this.images === output.images) && (!preview || this.preview === preview)) {
 								this.imgs = imgs.filter(Boolean);
 								this.setSizeForImage?.();
 								app.graph.setDirtyCanvas(true);
 							}
 						});
 					}
+					else {
+						this.imgs = null;
+					}
 				}
 
 				if (this.imgs && this.imgs.length) {
@@ -901,17 +930,20 @@ export class ComfyApp {
 			this.progress = null;
 			this.runningNodeId = detail;
 			this.graph.setDirtyCanvas(true, false);
+			delete this.nodePreviewImages[this.runningNodeId]
 		});
 
 		api.addEventListener("executed", ({ detail }) => {
 			this.nodeOutputs[detail.node] = detail.output;
 			const node = this.graph.getNodeById(detail.node);
-			if (node?.onExecuted) {
-				node.onExecuted(detail.output);
+			if (node) {
+				if (node.onExecuted)
+					node.onExecuted(detail.output);
 			}
 		});
 
 		api.addEventListener("execution_start", ({ detail }) => {
+			this.runningNodeId = null;
 			this.lastExecutionError = null
 		});
 
@@ -922,6 +954,16 @@ export class ComfyApp {
 			this.canvas.draw(true, true);
 		});
 
+		api.addEventListener("b_preview", ({ detail }) => {
+			const id = this.runningNodeId
+			if (id == null)
+				return;
+
+			const blob = detail
+			const blobUrl = URL.createObjectURL(blob)
+			this.nodePreviewImages[id] = [blobUrl]
+		});
+
 		api.init();
 	}
 
@@ -1465,8 +1507,10 @@ export class ComfyApp {
 	 */
 	clean() {
 		this.nodeOutputs = {};
+		this.nodePreviewImages = {}
 		this.lastPromptError = null;
 		this.lastExecutionError = null;
+		this.runningNodeId = null;
 	}
 }