diff --git a/animatediff/motion_module.py b/animatediff/motion_module.py index 17a732f..7e57ea2 100644 --- a/animatediff/motion_module.py +++ b/animatediff/motion_module.py @@ -61,7 +61,7 @@ class MotionWrapper(nn.Module): is_v2 = has_mid_block(mm_state_dict) mm = cls(mm_type, encoding_max_len=encoding_max_len, is_v2=is_v2) - mm.load_state_dict(mm_state_dict) + mm.load_state_dict(mm_state_dict, strict=False) return mm def set_video_length(self, video_length: int): diff --git a/animatediff/nodes.py b/animatediff/nodes.py index 74f5a8d..ffe10b6 100644 --- a/animatediff/nodes.py +++ b/animatediff/nodes.py @@ -14,7 +14,6 @@ import comfy.ldm.modules.diffusionmodules.openaimodel as openaimodel import comfy.model_management as model_management from comfy.model_base import BaseModel from comfy.ldm.modules.attention import SpatialTransformer -from comfy.cli_args import args as cli_args from comfy.utils import load_torch_file, calculate_parameters from nodes import KSampler @@ -75,6 +74,8 @@ def load_motion_module(model_name: str): if model_management.should_use_fp16(model_params=params): logger.info(f"Converting motion module to fp16.") motion_module.half() + offload_device = model_management.unet_offload_device() + motion_module = motion_module.to(offload_device) motion_modules[model_hash] = motion_module @@ -219,7 +220,7 @@ class AnimateDiffSampler(KSampler): def override_beta_schedule(self, model: BaseModel): logger.info(f"Override beta schedule.") - self.prev_beta = model.get_buffer("betas") + self.prev_beta = model.get_buffer("betas").cpu().clone() self.prev_linear_start = model.linear_start self.prev_linear_end = model.linear_end model.register_schedule( @@ -331,11 +332,10 @@ 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"],), - "filename_prefix": ("STRING", {"default": "AnimateDiff"}), - }, - "optional": { - "pingpong": (["Disabled", "Enabled"],), + "save_image": ([True, False],), + "filename_prefix": ("STRING", {"default": "animate_diff"}), + "format": (["image/gif", "image/webp", "video/webm"],), + "pingpong": ([False, True],), }, "hidden": { "prompt": "PROMPT", @@ -343,7 +343,7 @@ class AnimateDiffCombine: }, } - RETURN_TYPES = () + RETURN_TYPES = ("GIF",) OUTPUT_NODE = True CATEGORY = "Animate Diff" FUNCTION = "generate_gif" @@ -353,23 +353,24 @@ class AnimateDiffCombine: images, frame_rate: int, loop_count: int, - save_image="Enabled", + save_image=True, filename_prefix="AnimateDiff", - pingpong="Disabled", + format="image/gif", + pingpong=False, prompt=None, extra_pnginfo=None, ): # convert images to numpy - pil_images: List[Image.Image] = [] + frames: List[Image.Image] = [] for image in images: img = 255.0 * image.cpu().numpy() img = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8)) - pil_images.append(img) + frames.append(img) # save image output_dir = ( folder_paths.get_output_directory() - if save_image == "Enabled" + if save_image else folder_paths.get_temp_directory() ) ( @@ -390,38 +391,54 @@ class AnimateDiffCombine: # save first frame as png to keep metadata file = f"{filename}_{counter:05}_.png" file_path = os.path.join(full_output_folder, file) - pil_images[0].save( + frames[0].save( file_path, pnginfo=metadata, compress_level=4, ) - - # make revert gif - if pingpong == "Enabled": - pil_images = pil_images + pil_images[-2:0:-1] + if pingpong: + frames = frames + frames[-2:0:-1] # save gif - file = f"{filename}_{counter:05}_.gif" + format_type, format_ext = format.split("/") + file = f"{filename}_{counter:05}_.{format_ext}" file_path = os.path.join(full_output_folder, file) - pil_images[0].save( - file_path, - save_all=True, - append_images=pil_images[1:], - duration=round(1000 / frame_rate), - loop=loop_count, - compress_level=4, - ) - print("Saved gif to", file_path, os.path.exists(file_path)) + if format_type == "image": + frames[0].save( + file_path, + format=format_ext.upper(), + save_all=True, + append_images=frames[1:], + duration=round(1000 / frame_rate), + loop=loop_count, + compress_level=4, + ) + else: + # save webm + import shutil + import subprocess + + ffmpeg_path = shutil.which("ffmpeg") + if ffmpeg_path is None: + raise ProcessLookupError("Could not find ffmpeg") + dimensions = f"{frames[0].width}x{frames[0].height}" + args = [ffmpeg_path, "-v", "panic", "-n", "-f", "rawvideo", "-pix_fmt", "rgb24", "-s", + dimensions, "-r", str(frame_rate), "-i", "-", "-pix_fmt", "yuv420p", file_path] + + with subprocess.Popen(args, stdin=subprocess.PIPE) as proc: + for frame in frames: + proc.stdin.write(frame.tobytes()) previews = [ { "filename": file, "subfolder": subfolder, - "type": "output" if save_image == "Enabled" else "temp", + "type": "output" if save_image else "temp", + "format": format, } ] - return {"ui": {"gif": previews}} + return {"ui": {"gifs": previews}} NODE_CLASS_MAPPINGS = { diff --git a/web/js/gif_preview.js b/web/js/gif_preview.js index 1f7f81d..f67b870 100644 --- a/web/js/gif_preview.js +++ b/web/js/gif_preview.js @@ -1,5 +1,3 @@ -// From https://github.com/kijai/ComfyUI-AnimateDiff-Evolved - import { app } from '../../../scripts/app.js' import { api } from '../../../scripts/api.js' @@ -19,16 +17,16 @@ function offsetDOMWidget( elRect.height / ctx.canvas.height ) .multiplySelf(ctx.getTransform()) - .translateSelf(margin, margin + widgetY) + .translateSelf(0, widgetY + margin) 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`, + left: `${transform.e}px`, top: `${transform.d + transform.f}px`, - width: `${widgetWidth - margin * 2}px`, - height: `${(height || widget.parent?.inputHeight || 32) - margin * 2}px`, + width: `${widgetWidth}px`, + height: `${(height || widget.parent?.inputHeight || 32) - margin}px`, position: 'absolute', background: !node.color ? '' : node.color, color: !node.color ? '' : 'white', @@ -60,21 +58,21 @@ export const cleanupNode = (node) => { } } -const DEBUG_IMG = (name, val) => { +const CreatePreviewElement = (name, val, format) => { + const [type] = format.split('/') + const w = { name, - type: 'image', + type, 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) { + computeSize: function (_) { const ratio = this.inputRatio || 1 - if (width) { - return [width, width / ratio + 4] - } - return [128, 128] + const width = Math.max(220, this.parent.size[0]) + return [width, (width / ratio + 10)] }, onRemoved: function () { if (this.inputEl) { @@ -83,8 +81,14 @@ const DEBUG_IMG = (name, val) => { }, } - w.inputEl = document.createElement('img') + w.inputEl = document.createElement(type === 'video' ? 'video' : 'img') w.inputEl.src = w.value + if (type === 'video') { + w.inputEl.setAttribute('type', 'video/webm'); + w.inputEl.autoplay = true + w.inputEl.loop = true + w.inputEl.controls = false; + } w.inputEl.onload = function () { w.inputRatio = w.inputEl.naturalWidth / w.inputEl.naturalHeight } @@ -93,13 +97,13 @@ const DEBUG_IMG = (name, val) => { } const gif_preview = { - name: 'ad_gif_preview', + name: 'AnimateDiff.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 prefix = 'ad_gif_preview_' const r = onExecuted ? onExecuted.apply(this, message) : undefined if (this.widgets) { @@ -110,37 +114,29 @@ const gif_preview = { } 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() - ) - }) + if (message?.gifs) { + message.gifs.forEach((params, i) => { + const previewUrl = api.apiURL( + '/view?' + new URLSearchParams(params).toString() ) - } - let i = 0 - for (const img of imgURLs) { const w = this.addCustomWidget( - DEBUG_IMG(`${prefix}_${i}`, img) + CreatePreviewElement(`${prefix}_${i}`, previewUrl, params.format || 'image/gif') ) w.parent = this - i++ - } + }) } + const onRemoved = this.onRemoved this.onRemoved = () => { cleanupNode(this) return onRemoved?.() } } + + // keep width and update height + this.setSize([this.size[0], this.computeSize([this.size[0], this.size[1]])[1]]) return r } - break } }