From 87815b7aae54edf9467cc809556b477926293fde Mon Sep 17 00:00:00 2001 From: Tung Nguyen Date: Sat, 16 Sep 2023 06:41:59 +0700 Subject: [PATCH] add gif preview & support pingping gif --- __init__.py | 4 +- animatediff/nodes.py | 10 ++- web/js/gif_preview.js | 150 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 162 insertions(+), 2 deletions(-) create mode 100644 web/js/gif_preview.js diff --git a/__init__.py b/__init__.py index 79afcc5..6d8dfa4 100644 --- a/__init__.py +++ b/__init__.py @@ -5,4 +5,6 @@ from .animatediff.model_utils import get_available_models if len(get_available_models()) == 0: logger.error("No models available. Please download one and put it in models folder") -__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 77b4d06..74f5a8d 100644 --- a/animatediff/nodes.py +++ b/animatediff/nodes.py @@ -334,6 +334,9 @@ class AnimateDiffCombine: "save_image": (["Enabled", "Disabled"],), "filename_prefix": ("STRING", {"default": "AnimateDiff"}), }, + "optional": { + "pingpong": (["Disabled", "Enabled"],), + }, "hidden": { "prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO", @@ -352,6 +355,7 @@ class AnimateDiffCombine: loop_count: int, save_image="Enabled", filename_prefix="AnimateDiff", + pingpong="Disabled", prompt=None, extra_pnginfo=None, ): @@ -392,6 +396,10 @@ class AnimateDiffCombine: compress_level=4, ) + # make revert gif + if pingpong == "Enabled": + pil_images = pil_images + pil_images[-2:0:-1] + # save gif file = f"{filename}_{counter:05}_.gif" file_path = os.path.join(full_output_folder, file) @@ -413,7 +421,7 @@ class AnimateDiffCombine: "type": "output" if save_image == "Enabled" else "temp", } ] - return {"ui": {"images": previews}} + return {"ui": {"gif": previews}} NODE_CLASS_MAPPINGS = { diff --git a/web/js/gif_preview.js b/web/js/gif_preview.js new file mode 100644 index 0000000..1f7f81d --- /dev/null +++ b/web/js/gif_preview.js @@ -0,0 +1,150 @@ +// From https://github.com/kijai/ComfyUI-AnimateDiff-Evolved + +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: `${(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 hasWidgets = (node) => { + if (!node.widgets || !node.widgets?.[Symbol.iterator]) { + return false + } + return true +} + +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: 'ad_gif_preview', + async beforeRegisterNodeDef(nodeType, nodeData, app) { + switch (nodeData.name) { + case '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 = () => { + cleanupNode(this) + return onRemoved?.() + } + } + return r + } + + break + } + } + } +} + +app.registerExtension(gif_preview) \ No newline at end of file