From e21cb5987b31d62d803d7d7a77bbe281eefdafed Mon Sep 17 00:00:00 2001 From: kijai Date: Fri, 15 Sep 2023 21:19:05 +0300 Subject: [PATCH] First working gif preview javascript code is pretty much directly from mtb node's "save gif" widget extension code --- __init__.py | 3 +- animatediff/nodes.py | 22 ++++--- web/js/gif_preview.js | 144 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 160 insertions(+), 9 deletions(-) create mode 100644 web/js/gif_preview.js diff --git a/__init__.py b/__init__.py index c5afdf6..7dd7250 100644 --- a/__init__.py +++ b/__init__.py @@ -5,4 +5,5 @@ from .animatediff.model_utils import get_available_models, get_folder_path, Fold if len(get_available_models()) == 0: logger.error(f"No motion models found. Please download one and place in: {get_folder_path(Folders.MODELS)}") -__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] +WEB_DIRECTORY = "./web" +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] diff --git a/animatediff/nodes.py b/animatediff/nodes.py index 8cf1a56..72241eb 100644 --- a/animatediff/nodes.py +++ b/animatediff/nodes.py @@ -436,8 +436,9 @@ class AnimateDiffCombine: {"default": 8, "min": 1, "max": 24, "step": 1}, ), "loop_count": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}), - "save_image": (["Enabled", "Disabled"],), + "save_image": ("BOOLEAN", {"default": True}), "filename_prefix": ("STRING", {"default": "AnimateDiff"}), + "pingpong": ("BOOLEAN", {"default": False}), }, "hidden": { "prompt": "PROMPT", @@ -455,8 +456,9 @@ class AnimateDiffCombine: images, frame_rate: int, loop_count: int, - save_image="Enabled", + save_image=True, filename_prefix="AnimateDiff", + pingpong=False, prompt=None, extra_pnginfo=None, ): @@ -466,11 +468,13 @@ class AnimateDiffCombine: img = 255.0 * image.cpu().numpy() img = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8)) pil_images.append(img) - + if pingpong: + pil_images += pil_images[::-1] + # save image output_dir = ( folder_paths.get_output_directory() - if save_image == "Enabled" + if save_image else folder_paths.get_temp_directory() ) ( @@ -496,7 +500,8 @@ class AnimateDiffCombine: pnginfo=metadata, compress_level=4, ) - + + # save gif file = f"{filename}_{counter:05}_.gif" file_path = os.path.join(full_output_folder, file) @@ -514,11 +519,12 @@ class AnimateDiffCombine: previews = [ { "filename": file, - "subfolder": subfolder, - "type": "output" if save_image == "Enabled" else "temp", + "subfolder": "", + "type": "output" if save_image else "temp", } ] - return {"ui": {"images": previews}} + print(previews) + return {"ui": {"gif": previews}} class CheckpointLoaderSimpleWithNoiseSelect: diff --git a/web/js/gif_preview.js b/web/js/gif_preview.js new file mode 100644 index 0000000..225196a --- /dev/null +++ b/web/js/gif_preview.js @@ -0,0 +1,144 @@ +import { app } from '../../../scripts/app.js' +import { api } from '../../../scripts/api.js' + +function offsetDOMWidget( + widget, + ctx, + node, + widgetWidth, + widgetY, + height + ) { + const margin = 10 + const elRect = ctx.canvas.getBoundingClientRect() + const transform = new DOMMatrix() + .scaleSelf( + elRect.width / ctx.canvas.width, + elRect.height / ctx.canvas.height + ) + .multiplySelf(ctx.getTransform()) + .translateSelf(margin, margin + widgetY) + + const scale = new DOMMatrix().scaleSelf(transform.a, transform.d) + Object.assign(widget.inputEl.style, { + transformOrigin: '0 0', + transform: scale, + left: `${transform.a + transform.e}px`, + top: `${transform.d + transform.f}px`, + width: `${widgetWidth - margin * 2}px`, + // height: `${(widget.parent?.inputHeight || 32) - (margin * 2)}px`, + height: `${(height || widget.parent?.inputHeight || 32) - margin * 2}px`, + + position: 'absolute', + background: !node.color ? '' : node.color, + color: !node.color ? '' : 'white', + zIndex: 5, //app.graph._nodes.indexOf(node), + }) + } + + export const cleanupNode = (node) => { + if (!hasWidgets(node)) { + return + } + + for (const w of node.widgets) { + if (w.canvas) { + w.canvas.remove() + } + if (w.inputEl) { + w.inputEl.remove() + } + // calls the widget remove callback + w.onRemoved?.() + } + } + +const DEBUG_IMG = (name, val) => { + const w = { + name, + type: 'image', + value: val, + draw: function (ctx, node, widgetWidth, widgetY, height) { + const [cw, ch] = this.computeSize(widgetWidth) + offsetDOMWidget(this, ctx, node, widgetWidth, widgetY, ch) + }, + computeSize: function (width) { + const ratio = this.inputRatio || 1 + if (width) { + return [width, width / ratio + 4] + } + return [128, 128] + }, + onRemoved: function () { + if (this.inputEl) { + this.inputEl.remove() + } + }, + } + + w.inputEl = document.createElement('img') + w.inputEl.src = w.value + w.inputEl.onload = function () { + w.inputRatio = w.inputEl.naturalWidth / w.inputEl.naturalHeight + } + document.body.appendChild(w.inputEl) + return w + } + +const gif_preview = { + name: 'gif.preview', + async beforeRegisterNodeDef(nodeType, nodeData, app) { + switch (nodeData.name) { + case 'ADE_AnimateDiffCombine':{ + const onExecuted = nodeType.prototype.onExecuted + nodeType.prototype.onExecuted = function (message) { + const prefix = 'anything_' + const r = onExecuted ? onExecuted.apply(this, message) : undefined + + if (this.widgets) { + const pos = this.widgets.findIndex((w) => w.name === `${prefix}_0`) + if (pos !== -1) { + for (let i = pos; i < this.widgets.length; i++) { + this.widgets[i].onRemoved?.() + } + this.widgets.length = pos + } + + let imgURLs = [] + if (message) { + if (message.gif) { + console.log("found gif") + imgURLs = imgURLs.concat( + message.gif.map((params) => { + return api.apiURL( + '/view?' + new URLSearchParams(params).toString() + ) + }) + ) + } + let i = 0 + for (const img of imgURLs) { + const w = this.addCustomWidget( + DEBUG_IMG(`${prefix}_${i}`, img) + ) + w.parent = this + i++ + } + } + const onRemoved = this.onRemoved + this.onRemoved = () => { + shared.cleanupNode(this) + return onRemoved?.() + } + } + //this.setSize?.(this.computeSize()) # this seems to reset the node size on each generation + return r + } + + break + } + } + } +} + +app.registerExtension(gif_preview)