Grid viewer

web/extensions/core/showGrid.js

@@ -0,0 +1,247 @@
+import { app } from "/scripts/app.js";
+
+// Show grids from combinatorial outputs
+
+app.registerExtension({
+	name: "Comfy.ShowGrid",
+	async beforeRegisterNodeDef(nodeType, nodeData, app) {
+		if (!(nodeData.name === "SaveImage" || nodeData.name === "PreviewImage")) {
+			return
+		}
+
+		const onNodeCreated = nodeType.prototype.onNodeCreated;
+		nodeType.prototype.onNodeCreated = function () {
+			const r = onNodeCreated ? onNodeCreated.apply(this, arguments) : undefined;
+
+			this.addWidget("button", "Show Grid", "Show Grid", () => {
+				const grid = app.nodeGrids[this.id];
+				if (grid == null) {
+					console.warn("No grid to show!");
+					return;
+				}
+
+				const graphCanvas = LiteGraph.LGraphCanvas.active_canvas
+				if (graphCanvas == null)
+					return;
+
+				if (this._gridPanel != null)
+					return
+
+				this._gridPanel = graphCanvas.createPanel("Grid", { closable: true });
+				this._gridPanel.onClose = () => {
+					this._gridPanel = null;
+				}
+				this._gridPanel.node = this;
+				this._gridPanel.classList.add("grid_dialog");
+
+				const rootHtml = `
+<div class="axis-selectors">
+</div>
+<table class="image-table">
+</table>
+`;
+				const rootElem = this._gridPanel.addHTML(rootHtml, "grid-root");
+				const axisSelectors = rootElem.querySelector(".axis-selectors");
+				const imageTable = rootElem.querySelector(".image-table");
+
+				const footerHtml = `
+<label for="image-size">Image size</label>
+<input class="image-size" id="image-size" type="range" min="64" max="1024" step="1" value="512">
+</input>
+`
+				const footerElem = this._gridPanel.addHTML(footerHtml, "grid-footer", true);
+				const imageSizeInput = footerElem.querySelector(".image-size");
+
+				const frozenCoords = Array.from({length: grid.axes.length}, (v, i) => 0)
+
+				const getAxisData = (index) => {
+					let data = grid.axes[index];
+					if (data == null) {
+						data = {
+							nodeID: null,
+							id: "none",
+							label: "(Nothing)",
+							values: ["(None)"]
+						}
+					}
+					return data;
+				}
+
+				const selectAxis = (isY, axisID, change) => {
+					const axisName = isY ? "y" : "x";
+					const group = axisSelectors.querySelector(`.${axisName}-axis-selector`);
+
+					for (const input of group.querySelectorAll(`input#${axisName}-${axisID}`)) {
+						input.checked = true;
+						if (change) {
+							input.dispatchEvent(new Event('change'));
+						}
+					}
+				}
+
+				const getImagesAt = (x, y) => {
+					return grid.images.filter(image => {
+						for (let i = 0; i < grid.axes.length; i++) {
+							if (i === this.xAxis) {
+								if (image.coords[this.xAxis] !== x)
+									return false;
+							}
+							else if (i === this.yAxis) {
+								if (image.coords[this.yAxis] !== y)
+									return false;
+							}
+							else {
+								if (image.coords[i] !== frozenCoords[i])
+									return false;
+							}
+						}
+						return true;
+					});
+				}
+
+				const refreshGrid = (xAxis, yAxis) => {
+					this.xAxis = xAxis;
+					this.yAxis = yAxis;
+					this.xAxisData = getAxisData(this.xAxis);
+					this.yAxisData = getAxisData(this.yAxis);
+
+					selectAxis(false, this.xAxisData.id)
+					selectAxis(true, this.yAxisData.id)
+
+					if (xAxis === yAxis) {
+						this.yAxisData = getAxisData(-1);
+					}
+
+					imageTable.innerHTML = "";
+
+					const thead = document.createElement("thead")
+
+					const trXAxisLabel = document.createElement("tr");
+					const thXAxisLabel = document.createElement("th");
+					thXAxisLabel.setAttribute("colspan", String(this.xAxisData.values.length + 2))
+					thXAxisLabel.classList.add("axis", "x-axis")
+					thXAxisLabel.innerHTML = "<span>" + this.xAxisData.label + "</span>";
+					trXAxisLabel.appendChild(thXAxisLabel);
+					thead.appendChild(trXAxisLabel);
+
+					const trLabel = document.createElement("tr");
+					trLabel.appendChild(document.createElement("th")) // blank
+					trLabel.appendChild(document.createElement("th")) // blank
+					for (const xValue of this.xAxisData.values) {
+						const th = document.createElement("th");
+						th.classList.add("label", "x-label");
+						th.innerHTML = "<span>" + String(xValue) + "</span>";
+						trLabel.appendChild(th);
+					}
+					thead.appendChild(trLabel)
+
+					imageTable.appendChild(thead)
+
+					const tableBody = document.createElement("tbody");
+					imageTable.appendChild(tableBody);
+
+					const trYAxisLabel = document.createElement("tr");
+					const thYAxisLabel = document.createElement("th");
+					thYAxisLabel.setAttribute("rowspan", String(this.yAxisData.values.length + 1))
+					thYAxisLabel.classList.add("axis", "y-axis")
+					thYAxisLabel.style.textAlign = "center"
+					thYAxisLabel.style.textAlign = "center"
+					thYAxisLabel.innerHTML = "<span>" + this.yAxisData.label + "</span>";
+					trYAxisLabel.appendChild(thYAxisLabel);
+					tableBody.appendChild(trYAxisLabel);
+
+					for (const [y, yValue] of this.yAxisData.values.entries()) {
+						const tr = document.createElement("tr");
+
+						const tdLabel = document.createElement("td");
+						tdLabel.innerHTML = "<span>" + String(yValue) + "</span>";
+						tdLabel.classList.add("label", "y-label")
+						tr.append(tdLabel);
+
+						for (const [x, xValue] of this.xAxisData.values.entries()) {
+							const td = document.createElement("td");
+
+							const img = document.createElement("img");
+							img.style.width = `${this.imageSize}px`
+							img.style.height = `${this.imageSize}px`
+							const gridImages = getImagesAt(x, y);
+							if (gridImages.length > 0) {
+								img.src = "/view?" + new URLSearchParams(gridImages[0].image).toString() + app.getPreviewFormatParam();
+							}
+							td.append(img);
+
+							tr.append(td);
+						}
+						tableBody.appendChild(tr);
+					}
+				}
+
+				for (let i = 0; i < 2; i++) {
+					const axisName = i === 0 ? "x" : "y";
+					const isY = i === 1;
+
+					const group = document.createElement("div")
+					group.setAttribute("role", "group");
+					group.classList.add("axis-selector", `${axisName}-axis-selector`)
+
+					group.innerHTML = `${axisName.toUpperCase()} Axis:&nbsp `;
+
+					const addAxis = (index, axis) => {
+						const axisID = `${axisName}-${axis.id}`;
+
+						const input = document.createElement("input")
+						input.setAttribute("type", "radio")
+						input.setAttribute("name", `${axisName}-axis-selector`)
+						input.setAttribute("id", axisID)
+						input.classList.add("axis-radio")
+						input.addEventListener("change", () => {
+							if (input.checked) {
+								if (isY)
+									this.yAxis = index;
+								else
+									this.xAxis = index;
+							}
+
+							refreshGrid(this.xAxis, this.yAxis);
+						})
+
+						const label = document.createElement("label")
+						label.setAttribute("for", axisID)
+						label.classList.add("axis-label")
+						label.innerHTML = String(axis.label);
+						label.addEventListener("click", () => {
+							console.warn("SETAXIS", axis);
+							selectAxis(isY, axis.id, true);
+						})
+
+						group.appendChild(input)
+						group.appendChild(label)
+					}
+
+					// Add "None" entry
+					addAxis(-1, getAxisData(-1));
+
+					for (const [index, axis] of grid.axes.entries()) {
+						addAxis(index, axis);
+					}
+
+					axisSelectors.appendChild(group);
+				}
+
+				this.imageSize = 256;
+
+				imageSizeInput.addEventListener("input", () => {
+					this.imageSize = parseInt(imageSizeInput.value);
+					for (const img of imageTable.querySelectorAll("img")) {
+						img.style.width = `${this.imageSize}px`
+						img.style.height = `${this.imageSize}px`
+					}
+				})
+
+				refreshGrid(1, 2);
+
+				document.body.appendChild(this._gridPanel);
+			})
+		}
+	}
+})

web/scripts/app.js

@@ -1031,6 +1031,7 @@ export class ComfyApp {
 		while (queue.length > 0) {
 			const nodeID = queue.pop();
 			const promptInput = runningPrompt.output[nodeID];
+			const nodeClass = promptInput.class_type
 
 			// Ensure input keys are sorted alphanumerically
 			// This is important for the plot to have the same order as
@@ -1047,6 +1048,9 @@ export class ComfyApp {
 				if (typeof input === "object" && "__inputType__" in input) {
 					axes.push({
 						nodeID,
+						nodeClass,
+						id: `${nodeID}-${inputName}`.replace(" ", "-"),
+						label: `${nodeClass}: ${inputName}`,
 						inputName,
 						values: input.values
 					})
@@ -1075,19 +1079,23 @@ export class ComfyApp {
 			coords: []
 		}})
 
+		// TODO i don't know if this can generalize across arbitrary batch sizes
 		let factor = 1
+		let maxFactor = axes.map(a => a.values.length).reduce((s, l) => s * l, 1)
+		let batchFactor = images.length / maxFactor;
 
 		for (const axis of axes) {
 			factor *= axis.values.length;
 			for (const [index, image] of images.entries()) {
-				image.coords.push(Math.floor((index / factor) * axis.values.length) % axis.values.length);
+				const coord = Math.floor((index / factor / batchFactor) * axis.values.length) % axis.values.length;
+				image.coords.push(coord);
 			}
 		}
 
 		const grid = { axes, images };
 		console.error("GRID", grid);
 
-		return null;
+		return grid;
 	}
 
 	#addKeyboardHandler() {

web/style.css

@@ -356,3 +356,123 @@ button.comfy-queue-btn {
 	color: var(--input-text);
 	filter: brightness(50%);
 }
+
+.litegraph .dialog.grid_dialog {
+    display: inline-block;
+    text-align: right;
+	color: #AAA;
+    top: 0;
+    left: 0;
+    width: 100%;
+    height: 100%;
+    max-width: initial;
+    min-width: 200px;
+	max-height: initial;
+    min-height: 20px;
+    padding: 4px;
+    margin: auto;
+	overflow: hidden;
+	border-radius: 3px;
+    z-index: 200;
+}
+
+.litegraph .dialog.grid_dialog .grid-root {
+    margin: auto;
+    width: fit-content;
+}
+
+.litegraph .dialog.grid_dialog .grid-root .axis-selectors {
+    text-align: right;
+    margin: auto;
+    width: fit-content;
+    display: flex;
+    flex-direction: column;
+}
+
+.litegraph .dialog.grid_dialog .grid-root .axis-selector {
+    position: relative;
+    display: inline-flex;
+    vertical-align: middle;
+
+}
+
+.litegraph .dialog.grid_dialog .grid-root .axis-label {
+    position: relative;
+    flex: 1 1 auto;
+    display: inline-block;
+    padding: 0.5rem;
+    font-size: 14pt;
+    color: #ccc;
+    text-align: center;
+    text-decoration: none;
+    vertical-align: middle;
+    cursor: pointer;
+}
+
+.litegraph .dialog.grid_dialog .grid-root .axis-radio:checked + .axis-label {
+    color: black;
+    background-color: #aaa;
+    border-color: black;
+}
+
+.litegraph .dialog.grid_dialog .grid-root .axis-radio {
+    position: absolute;
+    clip: rect(0, 0, 0, 0);
+    pointer-events: none;
+}
+
+.litegraph .dialog.grid_dialog .grid-root .axis-label {
+}
+
+.litegraph .dialog.grid_dialog .grid-root img {
+    width: 512px;
+    height: 512px;
+    display: block;
+    margin: auto;
+    object-fit: contain;
+}
+
+.litegraph .dialog.grid_dialog .grid-root table {
+    border-spacing: 0;
+    border-collapse: collapse;
+}
+
+.litegraph .dialog.grid_dialog .grid-root .axis {
+    text-align: center;
+    font-size: 20pt;
+    padding: 1rem;
+}
+
+.litegraph .dialog.grid_dialog .grid-root .label {
+    text-align: center;
+    font-weight: bold;
+    font-size: 16pt;
+    padding: 0.5rem;
+}
+
+.litegraph .dialog.grid_dialog .grid-root .y-label {
+    vertical-align: middle;
+}
+
+.litegraph .dialog.grid_dialog .grid-root .y-axis {
+    writing-mode: tb-rl;
+    text-orientation: mixed;
+    transform: rotate(180deg);
+}
+
+.litegraph .dialog.grid_dialog table {
+    margin: auto;
+}
+
+.litegraph .dialog.grid_dialog .grid-footer {
+    display: flex;
+}
+
+.litegraph .dialog.grid_dialog .image-size {
+    width: 400px;
+}
+
+.litegraph .dialog.grid_dialog table tr td, .litegraph .dialog.grid_dialog table tr th {
+    border: 1px solid #ccc;
+    vertical-align: top;
+}