diff --git a/nodes.py b/nodes.py index b91cd39..d5b81e3 100644 --- a/nodes.py +++ b/nodes.py @@ -909,6 +909,62 @@ class WanVideoAddStoryMemLatents: updated["story_mem_latents"] = story_mem_latents["samples"].squeeze(2).permute(1, 0, 2, 3) # [C, T, H, W] return (updated,) + +class WanVideoSVIProEmbeds: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "anchor_samples": ("LATENT", {"tooltip": "Initial start image encoded"}), + "num_frames": ("INT", {"default": 81, "min": 1, "max": 10000, "step": 4, "tooltip": "Number of frames to encode"}), + }, + "optional": { + "prev_samples": ("LATENT", {"tooltip": "Last latent from previous generation"}), + "motion_latent_count": ("INT", {"default": 1, "min": 0, "max": 100, "step": 1, "tooltip": "Number of latents used to continue"}), + } + } + + RETURN_TYPES = ("WANVIDIMAGE_EMBEDS",) + RETURN_NAMES = ("image_embeds",) + FUNCTION = "add" + CATEGORY = "WanVideoWrapper" + + def add(self, anchor_samples, num_frames, prev_samples=None, motion_latent_count=1): + + anchor_latent = anchor_samples["samples"][0].clone() + + C, T, H, W = anchor_latent.shape + + total_latents = (num_frames - 1) // 4 + 1 + device = anchor_latent.device + dtype = anchor_latent.dtype + + if prev_samples is None or motion_latent_count == 0: + padding_size = total_latents - anchor_latent.shape[1] + padding = torch.zeros(C, padding_size, H, W, dtype=dtype, device=device) + y = torch.concat([anchor_latent, padding], dim=1) + else: + prev_latent = prev_samples["samples"][0].clone() + motion_latent = prev_latent[:, -motion_latent_count:] + padding_size = total_latents - anchor_latent.shape[1] - motion_latent.shape[1] + padding = torch.zeros(C, padding_size, H, W, dtype=dtype, device=device) + y = torch.concat([anchor_latent, motion_latent, padding], dim=1) + + msk = torch.ones(1, num_frames, H, W, device=device, dtype=dtype) + msk[:, 1:] = 0 + msk = torch.concat([torch.repeat_interleave(msk[:, 0:1], repeats=4, dim=1), msk[:, 1:]], dim=1) + msk = msk.view(1, msk.shape[1] // 4, 4, H, W) + msk = msk.transpose(1, 2)[0] + + image_embeds = { + "image_embeds": y, + "num_frames": num_frames, + "lat_h": H, + "lat_w": W, + "mask": msk + } + + return (image_embeds,) + #region I2V encode class WanVideoImageToVideoEncode: @classmethod @@ -2258,6 +2314,7 @@ NODE_CLASS_MAPPINGS = { "WanVideoUniLumosEmbeds": WanVideoUniLumosEmbeds, "WanVideoAddTTMLatents": WanVideoAddTTMLatents, "WanVideoAddStoryMemLatents": WanVideoAddStoryMemLatents, + "WanVideoSVIProEmbeds": WanVideoSVIProEmbeds, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -2299,4 +2356,5 @@ NODE_DISPLAY_NAME_MAPPINGS = { "WanVideoUniLumosEmbeds": "WanVideo UniLumos Embeds", "WanVideoAddTTMLatents": "WanVideo Add TTMLatents", "WanVideoAddStoryMemLatents": "WanVideo Add StoryMem Latents", + "WanVideoSVIProEmbeds": "WanVideo SVIPro Embeds", } diff --git a/nodes_sampler.py b/nodes_sampler.py index fc1f6c3..5eabb25 100644 --- a/nodes_sampler.py +++ b/nodes_sampler.py @@ -2784,8 +2784,13 @@ class WanVideoScheduler: handles, labels = ax.get_legend_handles_labels() if labels: ax.legend() - if start_idx < end_idx and 0 <= start_idx < len(sigmas_np) and 0 < end_idx < len(sigmas_np): - ax.axvspan(start_idx, end_idx, color='lightblue', alpha=0.1, label='Sampled Range') + # Draw shaded range + range_start_idx = start_idx if start_idx > 0 else 0 + range_end_idx = end_idx if end_idx > 0 and end_idx < len(sigmas_np) else len(sigmas_np) - 1 + if range_start_idx < range_end_idx: + ax.axvspan(range_start_idx, range_end_idx, color='lightblue', alpha=0.1, label='Sampled Range') + + plt.tight_layout() plt.savefig(buf, format='png') plt.close(fig)