diff --git a/nodes/debug.py b/nodes/debug.py index 1aaf5f0..10fe10f 100644 --- a/nodes/debug.py +++ b/nodes/debug.py @@ -7,6 +7,64 @@ from typing import Optional from pathlib import Path +# region processors +def process_tensor(tensor): + log.debug(f"Tensor: {tensor.shape}") + + image = tensor2pil(tensor) + b64_imgs = [] + for im in image: + buffered = io.BytesIO() + im.save(buffered, format="PNG") + b64_imgs.append( + "data:image/png;base64," + + base64.b64encode(buffered.getvalue()).decode("utf-8") + ) + + return {"b64_images": b64_imgs} + + +def process_list(anything): + text = [] + if not anything: + return {"text": []} + + first_element = anything[0] + if ( + isinstance(first_element, list) + and first_element + and isinstance(first_element[0], torch.Tensor) + ): + text.append( + f"List of List of Tensors: {first_element[0].shape} (x{len(anything)})" + ) + + elif isinstance(first_element, torch.Tensor): + text.append(f"List of Tensors: {first_element.shape} (x{len(anything)})") + + return {"text": text} + + +def process_dict(anything): + text = [] + if "samples" in anything: + is_empty = "(empty)" if torch.count_nonzero(anything["samples"]) == 0 else "" + text.append(f"Latent Samples: {anything['samples'].shape} {is_empty}") + + return {"text": text} + + +def process_bool(anything): + return {"text": ["True" if anything else "False"]} + + +def process_text(anything): + return {"text": [str(anything)]} + + +# endregion + + class Debug: """Experimental node to debug any Comfy values, support for more types and widgets is planned""" @@ -26,33 +84,23 @@ class Debug: "ui": {"b64_images": [], "text": []}, "result": ("A"), } - for k, v in kwargs.items(): - anything = v - text = "" - if isinstance(anything, torch.Tensor): - log.debug(f"Tensor: {anything.shape}") - # write the images to temp + processors = { + torch.Tensor: process_tensor, + list: process_list, + dict: process_dict, + bool: process_bool, + } - image = tensor2pil(anything) - b64_imgs = [] - for im in image: - buffered = io.BytesIO() - im.save(buffered, format="PNG") - b64_imgs.append( - "data:image/png;base64," - + base64.b64encode(buffered.getvalue()).decode("utf-8") - ) + for anything in kwargs.values(): + processor = processors.get(type(anything), process_text) + processed_data = processor(anything) - output["ui"]["b64_images"] += b64_imgs - log.debug(f"Input {k} contains {len(b64_imgs)} images") - elif isinstance(anything, bool): - log.debug(f"Input {k} contains boolean: {anything}") - output["ui"]["text"] += ["True" if anything else "False"] - else: - text = str(anything) - log.debug(f"Input {k} contains text: {text}") - output["ui"]["text"] += [text] + for ui_key, ui_value in processed_data.items(): + output["ui"][ui_key].extend(ui_value) + # log.debug( + # f"Processed input {k}, found {len(processed_data.get('b64_images', []))} images and {len(processed_data.get('text', []))} text items." + # ) return output diff --git a/web/debug.js b/web/debug.js index fbdadc2..0f144c7 100644 --- a/web/debug.js +++ b/web/debug.js @@ -15,6 +15,14 @@ import { MtbWidgets } from './mtb_widgets.js' // TODO: respect inputs order... +function escapeHtml(unsafe) { + return unsafe + .replace(/&/g, '&') + .replace(//g, '>') + .replace(/"/g, '"') + .replace(/'/g, ''') +} app.registerExtension({ name: 'mtb.Debug', async beforeRegisterNodeDef(nodeType, nodeData, app) { @@ -66,7 +74,7 @@ app.registerExtension({ if (message.text) { for (const txt of message.text) { const w = this.addCustomWidget( - MtbWidgets.DEBUG_STRING(`${prefix}_${widgetI}`, txt) + MtbWidgets.DEBUG_STRING(`${prefix}_${widgetI}`, escapeHtml(txt)) ) w.parent = this widgetI++ @@ -82,9 +90,10 @@ app.registerExtension({ } // this.onResize?.(this.size); // this.resize?.(this.size) - this.setSize(this.computeSize()) } + this.setSize(this.computeSize()) + this.onRemoved = function () { // When removing this node we need to remove the input from the DOM for (let y in this.widgets) { diff --git a/web/mtb_widgets.js b/web/mtb_widgets.js index 968c50c..1ec93d4 100644 --- a/web/mtb_widgets.js +++ b/web/mtb_widgets.js @@ -16,6 +16,36 @@ import { log } from './comfy_shared.js' const newTypes = [, /*'BOOL'*/ 'COLOR', 'BBOX'] +const withFont = (ctx, font, cb) => { + const oldFont = ctx.font + ctx.font = font + cb() + ctx.font = oldFont +} + +const calculateTextDimensions = (ctx, value, width, fontSize = 16) => { + const words = value.split(' ') + const lines = [] + let currentLine = '' + for (const word of words) { + const testLine = currentLine.length === 0 ? word : `${currentLine} ${word}` + const testWidth = ctx.measureText(testLine).width + if (testWidth > width) { + lines.push(currentLine) + currentLine = word + } else { + currentLine = testLine + } + } + if (lines.length === 0) lines.push(value) + const textHeight = (lines.length + 1) * fontSize + const maxLineWidth = lines.reduce( + (maxWidth, line) => Math.max(maxWidth, ctx.measureText(line).width), + 0 + ) + return { textHeight, maxLineWidth } +} + export const MtbWidgets = { BBOX: (key, val) => { /** @type {import("./types/litegraph").IWidget} */ @@ -317,46 +347,22 @@ export const MtbWidgets = { // const [cw, ch] = this.computeSize(widgetWidth) shared.offsetDOMWidget(this, ctx, node, widgetWidth, widgetY, height) }, - computeSize: function (width) { - const value = this.inputEl.innerHTML - if (!value) { + computeSize(width) { + if (!this.value) { return [32, 32] } if (!width) { - log(`No width ${this.parent.size}`) + console.debug(`No width ${this.parent.size}`) } - - const oldFont = app.ctx.font - app.ctx.font = `${fontSize}px monospace` - - const words = value.split(' ') - const lines = [] - let currentLine = '' - for (const word of words) { - const testLine = - currentLine.length === 0 ? word : `${currentLine} ${word}` - - const testWidth = app.ctx.measureText(testLine).width - - if (testWidth > width) { - lines.push(currentLine) - currentLine = word - } else { - currentLine = testLine - } - } - app.ctx.font = oldFont - if (lines.length === 0) lines.push(currentLine) - - const textHeight = (lines.length + 1) * fontSize - - const maxLineWidth = lines.reduce( - (maxWidth, line) => - Math.max(maxWidth, app.ctx.measureText(line).width), - 0 + let dimensions + withFont(app.ctx, `${fontSize}px monospace`, () => { + dimensions = calculateTextDimensions(app.ctx, this.value, width) + }) + const widgetWidth = Math.max( + width || this.width || 32, + dimensions.maxLineWidth ) - const widgetWidth = Math.max(width || this.width || 32, maxLineWidth) - const widgetHeight = textHeight * 1.5 + const widgetHeight = dimensions.textHeight * 1.5 return [widgetWidth, widgetHeight] }, onRemoved: function () { @@ -364,25 +370,23 @@ export const MtbWidgets = { this.inputEl.remove() } }, - } - - Object.defineProperty(w, 'value', { - get() { + get value() { return this.inputEl.innerHTML }, - set(value) { - this.inputEl.innerHTML = value + set value(val) { + this.inputEl.innerHTML = val this.parent?.setSize?.(this.parent?.computeSize()) }, - }) + } w.inputEl = document.createElement('p') - w.inputEl.style.textAlign = 'center' - w.inputEl.style.fontSize = `${fontSize}px` - w.inputEl.style.color = 'var(--input-text)' - w.inputEl.style.lineHeight = 0 - - w.inputEl.style.fontFamily = 'monospace' + w.inputEl.style = ` + text-align: center; + font-size: ${fontSize}px; + color: var(--input-text); + line-height: 0; + font-family: monospace; + ` w.value = val document.body.appendChild(w.inputEl)