From 08cf659cf30fefb93cd71efb5729dd0424673692 Mon Sep 17 00:00:00 2001 From: Flow-two Date: Fri, 28 Mar 2025 00:45:58 +0900 Subject: [PATCH] bug fixes --- __init__.py | 6 +- nodes/nodes.py | 222 ++++++++++++++++++++----------------------------- pyproject.toml | 2 +- 3 files changed, 95 insertions(+), 135 deletions(-) diff --git a/__init__.py b/__init__.py index 49fc799..e199e60 100644 --- a/__init__.py +++ b/__init__.py @@ -6,10 +6,12 @@ WEB_DIRECTORY = "./web" NODE_CLASS_MAPPINGS = { "WanImageToVideo_F2": WanImageToVideo_F2, - "GetImagesFromBatchRanged_F2": GetImagesFromBatchRanged_F2 + "GetImagesFromBatchRanged_F2": GetImagesFromBatchRanged_F2, + "WanSkipEndFrameImages_F2": WanSkipEndFrameImages_F2, } NODE_DISPLAY_NAME_MAPPINGS = { "WanImageToVideo_F2": "WanImageToVideo (Flow2)", - "GetImagesFromBatchRanged_F2": "Get Images From Batch Ranged" + "GetImagesFromBatchRanged_F2": "Get Images From Batch Ranged", + "WanSkipEndFrameImages_F2": "Wan Skip End Frame Images", } \ No newline at end of file diff --git a/nodes/nodes.py b/nodes/nodes.py index 463a738..d73d5db 100644 --- a/nodes/nodes.py +++ b/nodes/nodes.py @@ -1,129 +1,66 @@ from comfy.utils import common_upscale import torch import comfy.model_management as mm +import comfy.ldm.wan.vae import node_helpers import nodes +original_encode = comfy.ldm.wan.vae.WanVAE.encode +original_decode = comfy.ldm.wan.vae.WanVAE.decode + # reference: https://github.com/raindrop313/ComfyUI-WanVideoStartEndFrames # reference: https://github.com/kijai/ComfyUI-WanVideoWrapper/ +def encode(self, x): + self.clear_cache() -class WanVaeWrapper_F2: - def __init__(self, wv, end_image): - self.wv = wv - self.end_image = end_image + t = x.shape[2] + iter_ = 2 + (t - 2) // 4 - def __getattr__(self, name): - return getattr(self.wv, name) - - def encode(self, x): - if self.end_image is not None: - mu = self.encode_2(x) + for i in range(iter_): + self._enc_conv_idx = [0] + if i == 0: + out = self.encoder(x[:, :, :1, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx) + elif i== iter_-1: + out_ = self.encoder(x[:, :, -1:, :, :], + feat_cache=[None] * self._enc_conv_num, + feat_idx=self._enc_conv_idx) + out = torch.cat([out, out_], 2) else: - mu = self.encode_1(x) - return mu + out_ = self.encoder(x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx) + out = torch.cat([out, out_], 2) + out_head = out[:, :, :iter_ - 1, :, :] + out_tail = out[:, :, -1, :, :].unsqueeze(2) + mu, log_var = torch.cat([self.conv1(out_head), self.conv1(out_tail)], dim=2).chunk(2, dim=1) + return mu + +def decode(self, z): + self.clear_cache() - def decode(self, x): - if self.end_image is not None: - mu = self.decode_2(x) + iter_ = z.shape[2] + z_head=z[:,:,:-1,:,:] + z_tail=z[:,:,-1,:,:].unsqueeze(2) + x = torch.cat([self.conv2(z_head), self.conv2(z_tail)], dim=2) + for i in range(iter_): + self._conv_idx = [0] + if i == 0: + out = self.decoder(x[:, :, i:i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx) + elif i==iter_-1: + out_ = self.decoder(x[:, :, -1, :, :].unsqueeze(2), + feat_cache=None, + feat_idx=self._conv_idx) + out = torch.cat([out, out_], 2) else: - mu = self.decode_1(x) - return mu - - def encode_1(self, x): - self.clear_cache() - - t = x.shape[2] - iter_ = 1 + (t - 1) // 4 - - for i in range(iter_): - self._enc_conv_idx = [0] - if i == 0: - out = self.encoder( - x[:, :, :1, :, :], - feat_cache=self._enc_feat_map, - feat_idx=self._enc_conv_idx) - else: - out_ = self.encoder( - x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], - feat_cache=self._enc_feat_map, - feat_idx=self._enc_conv_idx) - out = torch.cat([out, out_], 2) - mu, log_var = self.conv1(out).chunk(2, dim=1) - self.clear_cache() - return mu - - def encode_2(self, x): - self.clear_cache() - - t = x.shape[2] - iter_ = 2 + (t - 2) // 4 - - for i in range(iter_): - self._enc_conv_idx = [0] - if i == 0: - out = self.encoder(x[:, :, :1, :, :], - feat_cache=self._enc_feat_map, - feat_idx=self._enc_conv_idx) - elif i== iter_-1: - out_ = self.encoder(x[:, :, -1:, :, :], - feat_cache=[None] * self._enc_conv_num, - feat_idx=self._enc_conv_idx) - out = torch.cat([out, out_], 2) - else: - out_ = self.encoder(x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], - feat_cache=self._enc_feat_map, - feat_idx=self._enc_conv_idx) - out = torch.cat([out, out_], 2) - out_head = out[:, :, :iter_ - 1, :, :] - out_tail = out[:, :, -1, :, :].unsqueeze(2) - mu, log_var = torch.cat([self.conv1(out_head), self.conv1(out_tail)], dim=2).chunk(2, dim=1) - return mu - - def decode_1(self, z): - self.clear_cache() - - iter_ = z.shape[2] - x = self.conv2(z) - for i in range(iter_): - self._conv_idx = [0] - if i == 0: - out = self.decoder( - x[:, :, i:i + 1, :, :], - feat_cache=self._feat_map, - feat_idx=self._conv_idx) - else: - out_ = self.decoder( - x[:, :, i:i + 1, :, :], - feat_cache=self._feat_map, - feat_idx=self._conv_idx) - out = torch.cat([out, out_], 2) - self.clear_cache() - return out - - def decode_2(self, z): - self.clear_cache() - - iter_ = z.shape[2] - z_head=z[:,:,:-1,:,:] - z_tail=z[:,:,-1,:,:].unsqueeze(2) - x = torch.cat([self.conv2(z_head), self.conv2(z_tail)], dim=2) - for i in range(iter_): - self._conv_idx = [0] - if i == 0: - out = self.decoder(x[:, :, i:i + 1, :, :], - feat_cache=self._feat_map, - feat_idx=self._conv_idx) - elif i==iter_-1: - out_ = self.decoder(x[:, :, -1, :, :].unsqueeze(2), - feat_cache=None, - feat_idx=self._conv_idx) - out = torch.cat([out, out_], 2) - else: - out_ = self.decoder(x[:, :, i:i + 1, :, :], - feat_cache=self._feat_map, - feat_idx=self._conv_idx) - out = torch.cat([out, out_], 2) - return out + out_ = self.decoder(x[:, :, i:i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx) + out = torch.cat([out, out_], 2) + return out class WanImageToVideo_F2: @classmethod @@ -131,9 +68,9 @@ class WanImageToVideo_F2: return {"required": {"positive": ("CONDITIONING", ), "negative": ("CONDITIONING", ), "vae": ("VAE", ), - "width": ("INT", {"default": 832, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 16}), - "height": ("INT", {"default": 480, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 16}), - "length": ("INT", {"default": 81, "min": 1, "max": nodes.MAX_RESOLUTION, "step": 4}), + "width": ("INT", {"default": 832, "min": 320, "max": nodes.MAX_RESOLUTION, "step": 16}), + "height": ("INT", {"default": 480, "min": 320, "max": nodes.MAX_RESOLUTION, "step": 16}), + "length": ("INT", {"default": 81, "min": 21, "max": nodes.MAX_RESOLUTION, "step": 4}), "start_image": ("IMAGE", ), }, "optional": {"end_image": ("IMAGE", ), @@ -142,39 +79,42 @@ class WanImageToVideo_F2: RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT") RETURN_NAMES = ("positive", "negative", "latent") - FUNCTION = "encode" + FUNCTION = "initialize" CATEGORY = "Flow2/Wan 2.1" - def encode(self, positive, negative, vae, width, height, length, start_image=None, end_image=None, clip_vision_output=None): - offload_device = mm.unet_offload_device() - device = offload_device - + def initialize(self, positive, negative, vae, width, height, length, start_image=None, end_image=None, clip_vision_output=None): valid_end_image = end_image is not None - resized_start_image = common_upscale(start_image.movedim(-1, 1), width, height, "lanczos", "center").movedim(1, -1) if valid_end_image: - resized_end_image = common_upscale(end_image.movedim(-1, 1), width, height, "lanczos", "center").movedim(1, -1) + comfy.ldm.wan.vae.WanVAE.encode = encode + comfy.ldm.wan.vae.WanVAE.decode = decode + else: + comfy.ldm.wan.vae.WanVAE.encode = original_encode + comfy.ldm.wan.vae.WanVAE.decode = original_decode - empty_frames = torch.ones(length - 1, height, width, 3, device=device) * 0.5 + offload_device = mm.unet_offload_device() + device = offload_device + + resized_start_image = common_upscale(start_image.movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1) + if valid_end_image: + resized_end_image = common_upscale(end_image.movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1) + + latent = torch.zeros([1, 16, ((length - 1) // 4) + 1, height // 8, width // 8], device=device) + + empty_frames = torch.ones(length - (1 + bool(valid_end_image)), height, width, 3, device=device) * 0.5 if valid_end_image: concatenated = torch.cat([resized_start_image.to(device), empty_frames, resized_end_image.to(device)], dim=0) else: concatenated = torch.cat([resized_start_image.to(device), empty_frames], dim=0) - if valid_end_image: # Replacing comfyui method can cause problems. - vae.first_stage_model.encode = WanVaeWrapper_F2(vae.first_stage_model, end_image).encode - la = vae.encode(concatenated) - latent = torch.zeros([1, 16, ((length - 1) // 4) + (1 + bool(valid_end_image)), height // 8, width // 8], device=device) - mask = torch.ones((1, 1, latent.shape[2], la.shape[-2], la.shape[-1]), device=device) - - mask[:, :, :1] = 0.0 + mask[:, :, 0] = 0.0 if valid_end_image: - mask[:, :, -1:] = 0.0 + mask[:, :, -1] = 0.0 positive = node_helpers.conditioning_set_values(positive, {"concat_latent_image": la, "concat_mask": mask}) negative = node_helpers.conditioning_set_values(negative, {"concat_latent_image": la, "concat_mask": mask}) @@ -207,4 +147,22 @@ class GetImagesFromBatchRanged_F2: end_index = round(count * end_percent) images = images[start_index:end_index] + return (images, ) + +class WanSkipEndFrameImages_F2: + @classmethod + def INPUT_TYPES(s): + return {"required": {"images": ("IMAGE", ), + "skip_end_frames": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}), + }} + + RETURN_TYPES = ("IMAGE", ) + FUNCTION = "skip" + + CATEGORY = "Flow2/Wan 2.1" + + def skip(self, images, skip_end_frames): + images = (images - images.min()) / (images.max() - images.min()) + images = images[0:-skip_end_frames] + images = torch.clamp(images, 0.0, 1.0) return (images, ) \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 941ac5d..4714867 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "ComfyUI-WanStartEndFramesNative" description = "Start and end frames video generation node that supports native ComfyUI." -version = "1.0.2" +version = "1.0.3" license = {file = "LICENSE"} dependencies = []