nodes/Flow-two-ComfyUI-WanStartEndFramesNative
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

bug fixes

Browse Source
This commit is contained in:
Flow-two
2025-03-28 00:45:58 +09:00
parent 7ee27a3353
commit 08cf659cf3
3 changed files with 95 additions and 135 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+4 -2
View File
@@ -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",
}
nodes/nodes.py
+90 -132
View File
@@ -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, )
pyproject.toml
+1 -1
View File
@@ -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 = []