diff --git a/nodes/nodes.py b/nodes/nodes.py index dbfb644..c606b24 100644 --- a/nodes/nodes.py +++ b/nodes/nodes.py @@ -110,11 +110,19 @@ class WanImageToVideo_F2: la = vae.encode(concatenated) latent = torch.zeros([1, 16, ((length - 1) // 4) + 1 + int(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[:, :, 0] = 0.0 + + mask = torch.ones((1, length + int(valid_end_image), la.shape[-2], la.shape[-1]), device=device) + mask[:, 0] = 0.0 if valid_end_image: - mask[:, :, -1] = 0.0 + mask[:, -1] = 0.0 + + start_mask_repeated = torch.repeat_interleave(mask[:, 0:1], repeats=4, dim=1) + end_mask_repeated = torch.repeat_interleave(mask[:, -1:], repeats=4, dim=1) + + mask = torch.cat([start_mask_repeated, mask[:, 1:-1], end_mask_repeated], dim=1) + mask = mask.view(1, mask.shape[1] // 4, 4, la.shape[-2], la.shape[-1]) + mask = mask.movedim(1, 2) + mask = mask.repeat(1, 4, 1, 1, 1) 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}) diff --git a/pyproject.toml b/pyproject.toml index 112e14f..76b2476 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.4" +version = "1.0.5" license = {file = "LICENSE"} dependencies = []