diff --git a/nodes.py b/nodes.py index f0b0e84d..2865032c 100644 --- a/nodes.py +++ b/nodes.py @@ -1,3 +1,4 @@ +# Licensed under the Apache License, Version 2.0 import os, gc, math import torch import torch.nn.functional as F @@ -1208,7 +1209,10 @@ def INPUT_TYPES(s): "bg_images": ("IMAGE", {"tooltip": "background images"}), "mask": ("MASK", {"tooltip": "mask"}), "start_ref_image": ("IMAGE", {"tooltip": "start ref image"}), + "transition_video": ("IMAGE", {"default": None, "tooltip": "Transition video frames (32 images, encoded to 8 latent frames). Acts as hard conditioning guide for seamless connection."}), + "prefix_frames": ("IMAGE", {"default": None, "tooltip": "3 reference images. Expands canvas by 17 pixel frames, encoded together with bg frames. Image 0 ×5, image 1 ×4, image 2 ×4, image 0 ×4. Shifts pose/face by 17 frames."}), "tiled_vae": ("BOOLEAN", {"default": False, "tooltip": "Use tiled VAE encoding for reduced memory use"}), + "Prefix & Transition Video by wuwukasi(bilibili)": ("BOOLEAN", {"default": True, "label_on": "ON", "label_off": "ON"}), } } @@ -1218,7 +1222,8 @@ def INPUT_TYPES(s): CATEGORY = "WanVideoWrapper" def process(self, vae, width, height, num_frames, force_offload, frame_window_size, colormatch, pose_strength, face_strength, - ref_images=None, pose_images=None, face_images=None, clip_embeds=None, tiled_vae=False, bg_images=None, mask=None, start_ref_image=None): + ref_images=None, pose_images=None, face_images=None, clip_embeds=None, tiled_vae=False, bg_images=None, mask=None, start_ref_image=None, + transition_video=None, prefix_frames=None, **kwargs): W = (width // 16) * 16 H = (height // 16) * 16 @@ -1229,7 +1234,71 @@ def process(self, vae, width, height, num_frames, force_offload, frame_window_si num_refs = ref_images.shape[0] if ref_images is not None else 0 num_frames = ((num_frames - 1) // 4) * 4 + 1 - looping = num_frames > frame_window_size or start_ref_image is not None + if transition_video is not None and prefix_frames is None: + + # --- [Core Mod] Reserve space for insertion logic and shift subsequent actions --- + # 1. Expand canvas: Add space for 21 pixel frames (corresponding to 6 Latent frames). + num_frames += 21 + trim = (num_frames - 1) % 4 + num_frames -= trim + + # 2. Shift control signals by 21 pixel frames with sampled+reversed padding + if pose_images is not None: + sampled = pose_images[0:42:2] # 21 frames from indices 0,2,...,40 + sampled = torch.flip(sampled, [0]) # reverse order + pose_images = torch.cat([sampled, pose_images], dim=0) + if face_images is not None: + sampled = face_images[0:42:2] + sampled = torch.flip(sampled, [0]) + face_images = torch.cat([sampled, face_images], dim=0) + if bg_images is not None: + sampled = bg_images[0:42:2] + sampled = torch.flip(sampled, [0]) + bg_images = torch.cat([sampled, bg_images], dim=0) + if mask is not None: + sampled = mask[0:42:2] + sampled = torch.flip(sampled, [0]) + mask = torch.cat([sampled, mask], dim=0) + # ---------------------------------------------------- + + if start_ref_image is not None: + log.warning("Both transition_video and start_ref_image provided. Using transition_video only (loop disabled).") + # ============ Prefix frames: expand canvas and shift control signals ============ + if prefix_frames is not None: + # Expand canvas: always 37 = 17 prefix + 12 transition + 8 reserve + extra = 37 + num_frames += extra + # Trim 1-3 frames from end to keep num_frames % 4 == 1 (required by repeat_interleave + view) + trim = (num_frames - 1) % 4 + num_frames -= trim + + # Pad beginning with sampled+reversed frames for pose/face (sparse temporal context), + # repeat frame 0 for bg/mask + if pose_images is not None: + sampled = pose_images[0:extra*2:2] # indices 0,2,4,... (extra frames) + sampled = torch.flip(sampled, [0]) # reverse order + pose_images = torch.cat([sampled, pose_images], dim=0) + if face_images is not None: + sampled = face_images[0:extra*2:2] + sampled = torch.flip(sampled, [0]) + face_images = torch.cat([sampled, face_images], dim=0) + if bg_images is not None: + sampled = bg_images[0:extra*2:2] + sampled = torch.flip(sampled, [0]) + bg_images = torch.cat([sampled, bg_images], dim=0) + if mask is not None: + sampled = mask[0:extra*2:2] + sampled = torch.flip(sampled, [0]) + mask = torch.cat([sampled, mask], dim=0) + # ----------------------------------------------------------------- + + if prefix_frames is not None: + effective_frames = num_frames - 37 + elif transition_video is not None: + effective_frames = num_frames - 21 + else: + effective_frames = num_frames + looping = effective_frames > frame_window_size or start_ref_image is not None if num_frames < frame_window_size: frame_window_size = num_frames @@ -1239,6 +1308,9 @@ def process(self, vae, width, height, num_frames, force_offload, frame_window_si if not looping: num_frames = num_frames + num_refs * 4 + # latent_window_size must cover the full bg latent range (including prefix/transition expansion), + # otherwise context windows that reach past the original frame count will clamp pose indices + latent_window_size = target_shape[1] - num_refs else: latent_window_size = latent_window_size + 1 @@ -1275,13 +1347,82 @@ def process(self, vae, width, height, num_frames, force_offload, frame_window_si resized_bg_images = bg_images.permute(3, 0, 1, 2) # C, T, H, W resized_bg_images = (resized_bg_images[:3] * 2 - 1) + actual_prefix_px = 0 + prefix_pixel_data = None # holds [C, actual_prefix_px, H, W] normalized pixel data for reuse + if prefix_frames is not None: + pf = prefix_frames + b_pf, h_pf, w_pf, c_pf = pf.shape + log.info(f"Prefix frames input: {b_pf} frames, {h_pf}x{w_pf}") + if b_pf > 5: + log.warning(f"Prefix has {b_pf} images, max 5. Truncating.") + pf = pf[:5] + b_pf = 5 + pf_frames = pf[0:1] + for i in range(1, b_pf): + pf_frames = torch.cat([pf_frames, pf[i:i+1].repeat(4, 1, 1, 1)], dim=0) + actual_prefix_px = pf_frames.shape[0] + log.info(f"Prefix: {b_pf} images -> {actual_prefix_px} pixel frames") + if h_pf != H or w_pf != W: + pf_frames = common_upscale(pf_frames.movedim(-1, 1), W, H, "lanczos", "disabled").movedim(1, -1) + prefix_pixel_data = pf_frames.permute(3, 0, 1, 2)[:3] * 2 - 1 # [C, actual_prefix_px, H, W] + del pf, pf_frames + if not looping: if bg_images is None: resized_bg_images = torch.zeros(3, num_frames - num_refs, H, W, device=device, dtype=vae.dtype) + + # ============ Prefix: replace first N pixel frames of canvas ============ + if prefix_pixel_data is not None: + resized_bg_images[:, :actual_prefix_px] = prefix_pixel_data.to(device, dtype=resized_bg_images.dtype) + log.info(f"Prefix: replaced first {actual_prefix_px} pixel frames of black canvas") + # If transition_video also present, embed last 12 frames into canvas positions 25-37 + if transition_video is not None: + tv = transition_video # [B, H, W, C] + b_tv = tv.shape[0] + if b_tv >= 12: + tv = tv[-12:] + else: + tv = torch.cat([tv[0:1].repeat(12 - b_tv, 1, 1, 1), tv], dim=0) + if tv.shape[1] != H or tv.shape[2] != W: + tv = common_upscale(tv.movedim(-1, 1), W, H, "lanczos", "disabled").movedim(1, -1) + tv = tv.permute(3, 0, 1, 2)[:3] * 2 - 1 # [C, 12, H, W] + resized_bg_images[:, 25:37] = tv.to(device, dtype=resized_bg_images.dtype) + log.info("Prefix+Transition: embedded last 12 transition frames into canvas positions 25-37") + # ========================================================================== + + # ============ Transition (no prefix): embed into canvas first 21 frames ============ + if transition_video is not None and prefix_frames is None: + tv = transition_video # [B, H, W, C] + b_tv = tv.shape[0] + if b_tv >= 21: + tv = tv[-21:] + else: + tv = torch.cat([tv[0:1].repeat(21 - b_tv, 1, 1, 1), tv], dim=0) + if tv.shape[1] != H or tv.shape[2] != W: + tv = common_upscale(tv.movedim(-1, 1), W, H, "lanczos", "disabled").movedim(1, -1) + tv = tv.permute(3, 0, 1, 2)[:3] * 2 - 1 # [C, 21, H, W] + resized_bg_images[:, :21] = tv.to(device, dtype=resized_bg_images.dtype) + log.info("Transition: embedded first 21 pixel frames of black canvas") + # ========================================================================== + bg_latents = vae.encode([resized_bg_images.to(device, vae.dtype)], device,tiled=tiled_vae)[0].to(offload_device) del resized_bg_images elif bg_images is not None: resized_bg_images = resized_bg_images.to(offload_device, dtype=vae.dtype) + elif transition_video is not None or prefix_frames is not None: + # Looping mode: create canvas (transition and/or prefix handled separately via prefix_ctx) + resized_bg_images = torch.zeros(3, num_frames - num_refs, H, W, device=offload_device, dtype=vae.dtype) + if transition_video is not None: + tv = transition_video # [B, H, W, C] + b_tv = tv.shape[0] + if b_tv >= 21: + tv = tv[-21:] + else: + tv = torch.cat([tv[0:1].repeat(21 - b_tv, 1, 1, 1), tv], dim=0) + tv = tv.permute(3, 0, 1, 2)[:3] * 2 - 1 # [C, 21, H, W] + resized_bg_images[:, :21] = tv.to(offload_device, dtype=resized_bg_images.dtype) + log.info("Transition (loop): embedded first 21 pixel frames of canvas") + # Prefix NOT embedded in canvas for looping — handled via prefix_ctx prepend later if ref_images is not None: if ref_images.shape[1] != H or ref_images.shape[2] != W: @@ -1295,6 +1436,21 @@ def process(self, vae, width, height, num_frames, force_offload, frame_window_si msk[:, :num_refs] = 1 ref_latent_masked = torch.cat([msk, ref_latent], dim=0).to(offload_device) # 4+C 1 H W + # ============ Prefix: VAE encode for looping (prepended to each chunk like ref) ============ + prefix_ctx = None + prefix_T = 0 + if prefix_frames is not None and looping: + vae.to(device) + prefix_latent = vae.encode([prefix_pixel_data.to(device, vae.dtype)], device, tiled=tiled_vae)[0] + prefix_T = prefix_latent.shape[1] + prefix_msk = torch.ones(4, prefix_T, lat_h, lat_w, device=offload_device, dtype=vae.dtype) + prefix_latent_masked = torch.cat([prefix_msk, prefix_latent.to(offload_device)], dim=0) # [20, prefix_T, ...] + prefix_ctx = torch.cat([ref_latent_masked, prefix_latent_masked], dim=1) # [20, 1+prefix_T, ...] + log.info(f"Prefix (loop): encoded {actual_prefix_px}px -> {prefix_T} latent, prefix_ctx: {prefix_ctx.shape}") + if force_offload: + vae.to(offload_device) + # =========================================================================================== + if mask is None: bg_mask = torch.zeros(1, num_frames, lat_h, lat_w, device=offload_device, dtype=vae.dtype) else: @@ -1303,7 +1459,17 @@ def process(self, vae, width, height, num_frames, force_offload, frame_window_si bg_mask = torch.cat([bg_mask, bg_mask[-1:].repeat(num_frames - bg_mask.shape[0], 1, 1)], dim=0) bg_mask = common_upscale(bg_mask.unsqueeze(1), lat_w, lat_h, "nearest", "disabled").squeeze(1) bg_mask = bg_mask.unsqueeze(-1).permute(3, 0, 1, 2).to(offload_device, vae.dtype) # C, T, H, W - + + # ============ Prefix: set mask=1 for actual prefix frames and optionally transition ============ + if prefix_frames is not None: + bg_mask[:, :actual_prefix_px] = 1.0 # only actual prefix pixel frames + if transition_video is not None: + bg_mask[:, 25:37] = 1.0 + # ======= Transition (no prefix): set mask=1 for first 21 pixel frames ======= + elif transition_video is not None: + bg_mask[:, :21] = 1.0 + # ====================================================================================== + if bg_images is None and looping: bg_mask[:, :num_refs] = 1 bg_mask_mask_repeated = torch.repeat_interleave(bg_mask[:, 0:1], repeats=4, dim=1) # T, C, H, W @@ -1334,6 +1500,65 @@ def process(self, vae, width, height, num_frames, force_offload, frame_window_si resized_start_ref_image = start_ref_image.permute(3, 0, 1, 2) # C, T, H, W resized_start_ref_image = resized_start_ref_image[:3] * 2 - 1 + # ============ Transition video processing ============ + transition_latent = None + transition_mask_values = None + + if False: # Transition now embedded in canvas (non-looping) or bg_images (looping), no independent VAE encode needed + # transition_video input: 32 images [B, H, W, C] + # Expecting B=32, which encodes to 8 latent frames + b, h, w, c = transition_video.shape + log.info(f"Transition video input: {b} frames, {h}x{w}") + + # Verify frame count to ensure it is exactly 32 frames + expected_input_frames = 32 + if b != expected_input_frames: + log.warning(f"Transition video has {b} frames, expected {expected_input_frames}. Resizing time dimension.") + if b > expected_input_frames: + # Downsample to 32 frames + indices = torch.linspace(0, b-1, expected_input_frames).long() + transition_video = transition_video[indices] + else: + # Repeat frames to reach 32 frames + repeat_factor = math.ceil(expected_input_frames / b) + transition_video = transition_video.repeat(repeat_factor, 1, 1, 1)[:expected_input_frames] + + b, h, w, c = transition_video.shape # It should be 32 now + + # Adjust spatial dimensions to target WxH. + # Keep the same semantic flow as the matched-size path: + # BHWC -> (optional resize in BCHW) -> BHWC -> CTHW -> normalize -> encode + if h != H or w != W: + transition_video = common_upscale( + transition_video.movedim(-1, 1), W, H, "lanczos", "disabled" + ).movedim(1, -1) + + # Normalize to [-1, 1] + transition_video = transition_video.permute(3, 0, 1, 2) # [C, T, H, W] + transition_video = transition_video[:3] * 2 - 1 # Keep only RGB channels + + + # VAE Encoding (32 pixel frames -> 8 latent frames) + vae.to(device) + transition_latent = vae.encode([transition_video.to(device, vae.dtype)], device, tiled=tiled_vae)[0] + log.info(f"Transition latent encoded: {transition_latent.shape[1] if len(transition_latent.shape) > 1 else transition_latent.shape[0]} frames, shape {transition_latent.shape}") + transition_len = transition_latent.shape[1] # It should be 8 + log.info(f"Transition latent encoded: {transition_len} frames, shape {transition_latent.shape}") + + # ============ Generate Mask values ============ + # Force mask to all 1s, making it act purely as a hard conditioning guide. + # The model will strictly follow these frames without altering them. + transition_mask_values = torch.ones(transition_len) + + log.info("Transition mask: forced to all 1s for hard conditioning.") + log.info(f"Mask values: {transition_mask_values.tolist()}") + # ========================================== + + if force_offload: + transition_latent = transition_latent.to(offload_device) + transition_mask_values = transition_mask_values.to(offload_device) + # ================================================ + seq_len = math.ceil((target_shape[2] * target_shape[3]) / 4 * target_shape[1]) if force_offload: @@ -1347,12 +1572,18 @@ def process(self, vae, width, height, num_frames, force_offload, frame_window_si "max_seq_len": seq_len, "pose_latents": pose_latents, "pose_images": resized_pose_images if pose_images is not None and looping else None, - "bg_images": resized_bg_images if bg_images is not None and looping else None, - "ref_masks": bg_mask if mask is not None and looping else None, + "bg_images": resized_bg_images if (bg_images is not None or transition_video is not None or prefix_frames is not None) and looping else None, + "ref_masks": bg_mask if (mask is not None or prefix_frames is not None) and looping else None, "is_masked": mask is not None, "ref_latent": ref_latent, "ref_image": resized_ref_images if ref_images is not None else None, "start_ref_image": resized_start_ref_image if start_ref_image is not None else None, + "transition_latent": transition_latent, + "transition_mask_values": transition_mask_values, + "has_prefix": prefix_frames is not None, + "canvas_expansion_px": 37 if prefix_frames is not None else (21 if transition_video is not None else 0), + "prefix_ctx": prefix_ctx, + "prefix_T": prefix_T, "face_pixels": resized_face_images if face_images is not None else None, "num_frames": num_frames, "target_shape": target_shape, @@ -2124,6 +2355,8 @@ def decode(self, vae, samples, enable_vae_tiling, tile_x, tile_y, tile_stride_x, latents = samples["samples"].clone() end_image = samples.get("end_image", None) has_ref = samples.get("has_ref", False) + has_prefix = samples.get("has_prefix", False) + canvas_expansion_px = samples.get("canvas_expansion_px", 0) drop_last = samples.get("drop_last", False) is_looped = samples.get("looped", False) @@ -2165,9 +2398,12 @@ def decode(self, vae, samples, enable_vae_tiling, tile_x, tile_y, tile_stride_x, temp_images = (temp_images - temp_images.min()) / (temp_images.max() - temp_images.min()) images = torch.cat([temp_images[:, 9:].to(images), images[:, 5:]], dim=1) - if end_image is not None: + if end_image is not None: images = images[:, 0:-1] + if canvas_expansion_px and not is_looped: + images = images[:, canvas_expansion_px:] + vae.to(offload_device) mm.soft_empty_cache() diff --git a/nodes_sampler.py b/nodes_sampler.py index 32b51c6a..a6517aff 100644 --- a/nodes_sampler.py +++ b/nodes_sampler.py @@ -476,6 +476,22 @@ def process(self, model, image_embeds, shift, steps, cfg, seed, scheduler, rifle # region WanAnim inputs frame_window_size = image_embeds.get("frame_window_size", 77) wananimate_loop = image_embeds.get("looping", False) + # ============ Global transition_video read ============ + transition_latent = image_embeds.get("transition_latent", None) + transition_mask_values = image_embeds.get("transition_mask_values", None) + transition_len = transition_latent.shape[1] if transition_latent is not None else 0 + has_transition = transition_latent is not None + + + if has_transition: + log.info(f"Transition video: {transition_len} latent frames") + if transition_mask_values is not None: + log.info(f"Linear decay mask values: {transition_mask_values.tolist()}") + # ================================================================ + has_prefix = image_embeds.get("has_prefix", False) + canvas_expansion_px = image_embeds.get("canvas_expansion_px", 0) + if has_prefix: + log.info("Prefix frames: detected, will prepend indices 0-5 to each window (except first)") if wananimate_loop and context_options is not None: raise Exception("context_options are not compatible or necessary with WanAnim looping, since it creates the video in a loop.") wananim_pose_latents = image_embeds.get("pose_latents", None) @@ -630,6 +646,7 @@ def process(self, model, image_embeds, shift, steps, cfg, seed, scheduler, rifle seq_len = (context_frames * 2 + 1 + mocha_num_refs) * (noise.shape[2] * noise.shape[3] // 4) else: seq_len = math.ceil((noise.shape[2] * noise.shape[3]) / 4 * context_frames) + base_patches_per_frame = math.ceil((noise.shape[2] * noise.shape[3]) / 4) log.info(f"context window seq len: {seq_len}") if context_options["freenoise"]: @@ -844,7 +861,13 @@ def process(self, model, image_embeds, shift, steps, cfg, seed, scheduler, rifle render_latent = uni3c_embeds["render_latent"].to(device) uni3c_data = uni3c_embeds.copy() if render_latent.shape != noise.shape: - render_latent = torch.nn.functional.interpolate(render_latent, size=(noise.shape[1], noise.shape[2], noise.shape[3]), mode='trilinear', align_corners=False) + # If temporal is shorter (prefix/transition expansion), pad with first frame at beginning + if has_prefix: + pad_len = 10 # 37 pixel frames → 10 latent frames + first_frame = render_latent[:, :, :1].repeat(1, 1, pad_len, 1, 1) + render_latent = torch.cat([first_frame, render_latent], dim=2) + if render_latent.shape != noise.shape: + render_latent = torch.nn.functional.interpolate(render_latent, size=(noise.shape[1], noise.shape[2], noise.shape[3]), mode='trilinear', align_corners=False) uni3c_data["render_latent"] = render_latent # Enhance-a-video (feta) @@ -1857,8 +1880,14 @@ def predict_with_cfg(z, cfg_scale, positive_embeds, negative_embeds, timestep, i enhance_enabled = True #region context windowing if context_options is not None: + # ============ Get latent spatial dimensions ============ + lat_h = latent.shape[-2] + lat_w = latent.shape[-1] + # ============================================ counter = torch.zeros_like(latent_model_input, device=device) noise_pred = torch.zeros_like(latent_model_input, device=device) + if image_cond is not None and image_cond.ndim > 1: + latent_video_length = min(latent_video_length, image_cond.shape[1]) context_queue = list(context(idx, steps, latent_video_length, context_frames, context_stride, context_overlap)) fraction_per_context = 1.0 / len(context_queue) context_pbar = ProgressBar(steps) @@ -1891,22 +1920,55 @@ def predict_with_cfg(z, cfg_scale, positive_embeds, negative_embeds, timestep, i partial_img_emb = partial_control_latents = None if image_cond is not None: partial_img_emb = image_cond[:, c].to(device) - if c[0] != 0 and context_reference_latent is not None: - if context_reference_latent.shape[0] == 1: #only single extra init latent - new_init_image = context_reference_latent[0, :, 0].to(device) - # Concatenate the first 4 channels of partial_img_emb with new_init_image to match the required shape - partial_img_emb[:, 0] = torch.cat([image_cond[:4, 0].to(device), new_init_image], dim=0) - elif context_reference_latent.shape[0] > 1: - num_extra_inits = context_reference_latent.shape[0] - section_size = (latent_video_length / num_extra_inits) - extra_init_index = min(int(max(c) / section_size), num_extra_inits - 1) - if context_options["verbose"]: - log.info(f"extra init image index: {extra_init_index}") - new_init_image = context_reference_latent[extra_init_index, :, 0].to(device) - partial_img_emb[:, 0] = torch.cat([image_cond[:4, 0].to(device), new_init_image], dim=0) - else: - new_init_image = image_cond[:, 0].to(device) - partial_img_emb[:, 0] = new_init_image + # ============ Build msk channels ============ + window_frames = len(c) + window_msk = partial_img_emb[:4].clone() + # ======================================== + + # ============ Transition replacement (1st window only) ============ + if has_transition and c[0] == 0 and not has_prefix: + partial_img_emb[4:20, 1:1+transition_len] = transition_latent.to(device, dtype=transition_latent.dtype) + log.info(f"Replaced first {transition_len} latent frames with transition_video (context_options mode)") + # ============================================================ + + # ============ Set transition mask ============ + if has_transition and c[0] == 0 and not has_prefix: + window_trans_start = 1 + window_trans_end = 1 + transition_len + if window_trans_start < window_trans_end: + mask_slice = transition_mask_values[0:transition_len] + window_msk[:, window_trans_start:window_trans_end] = mask_slice.view(1, -1, 1, 1) + # ================================================================ + + # ============ Prefix: prepend indices 0-5 to non-first windows ============ + if has_prefix and c[0] != 0: + prefix_ctx = image_cond[:, :6].to(device, dtype) + partial_img_emb = torch.cat([prefix_ctx, partial_img_emb], dim=1) + prefix_msk = image_cond[:4, :6].to(device, dtype) + window_msk = torch.cat([prefix_msk, window_msk], dim=1) + # ====================================================================== + + # ============ ref_latent replacement (skip when prefix active) ============ + if not has_prefix: + if c[0] != 0 and context_reference_latent is not None: + if context_reference_latent.shape[0] == 1: + new_init_image = context_reference_latent[0, :, 0].to(device) + partial_img_emb[:, 0] = torch.cat([image_cond[:4, 0].to(device), new_init_image], dim=0) + window_msk[:, 0] = image_cond[:4, 0].to(device) + elif context_reference_latent.shape[0] > 1: + num_extra_inits = context_reference_latent.shape[0] + section_size = (latent_video_length / num_extra_inits) + extra_init_index = min(int(max(c) / section_size), num_extra_inits - 1) + if context_options["verbose"]: + log.info(f"extra init image index: {extra_init_index}") + new_init_image = context_reference_latent[extra_init_index, :, 0].to(device) + partial_img_emb[:, 0] = torch.cat([image_cond[:4, 0].to(device), new_init_image], dim=0) + window_msk[:, 0] = image_cond[:4, 0].to(device) + else: + new_init_image = image_cond[:, 0].to(device) + partial_img_emb[:, 0] = new_init_image + window_msk[:, 0] = image_cond[:4, 0].to(device) + # ========================================================================== if control_latents is not None: partial_control_latents = control_latents[:, c] @@ -1957,6 +2019,12 @@ def predict_with_cfg(z, cfg_scale, positive_embeds, negative_embeds, timestep, i if latents_to_insert is not None and c[0] != 0: partial_latent_model_input[:, :1] = latents_to_insert + # ============ Prefix: prepend noise for indices 0-5 ============ + if has_prefix and c[0] != 0: + prefix_noise = latent_model_input[:, :6].to(device) + partial_latent_model_input = torch.cat([prefix_noise, partial_latent_model_input], dim=1) + # ============================================================= + partial_unianim_data = None if unianim_data is not None: partial_dwpose = dwpose_data[:, :, c] @@ -2005,7 +2073,18 @@ def predict_with_cfg(z, cfg_scale, positive_embeds, negative_embeds, timestep, i end = c[-1] center_indices = torch.arange(start, end, 1) center_indices = torch.clamp(center_indices, min=0, max=wananim_pose_latents.shape[2] - 1) - partial_wananim_pose_latents = wananim_pose_latents[:, :, center_indices][:, :, :context_frames-1].to(device, dtype) + pose_limit = context_frames - 1 + partial_wananim_pose_latents = wananim_pose_latents[:, :, center_indices][:, :, :pose_limit].to(device, dtype) + + # ============ Prefix: prepend face/pose for indices 0-5 ============ + if has_prefix and c[0] != 0: + if partial_wananim_face_pixels is not None: + prefix_face = wananim_face_pixels[:, :, :24].to(device, dtype) + partial_wananim_face_pixels = torch.cat([prefix_face, partial_wananim_face_pixels], dim=2) + if partial_wananim_pose_latents is not None: + prefix_pose = wananim_pose_latents[:, :, :6].to(device, dtype) + partial_wananim_pose_latents = torch.cat([prefix_pose, partial_wananim_pose_latents], dim=2) + # ================================================================ partial_flashvsr_LQ_latent = None if LQ_images is not None: @@ -2024,17 +2103,38 @@ def predict_with_cfg(z, cfg_scale, positive_embeds, negative_embeds, timestep, i orig_model_input_frames = partial_latent_model_input.shape[1] + # ============ Replace msk placeholder channels ============ + # image_cond is optional in this path. Only build image_cond_in + # when partial_img_emb is available. + if partial_img_emb is not None: + # image_cond structure: [image_cond_mask (4) + image_embeds (32)] = 36 channels + # Replace the first 4 channels of image_cond_mask with window_msk + image_cond_in = partial_img_emb.clone() + image_cond_in[:4] = window_msk # Replace the first 4 channels + else: + image_cond_in = None + # ======================================== + original_seq_len = seq_len + if has_prefix and c[0] != 0: + seq_len = original_seq_len + 6 * base_patches_per_frame noise_pred_context, _, new_teacache = predict_with_cfg( partial_latent_model_input, cfg[idx], positive, text_embeds["negative_prompt_embeds"], - partial_timestep, idx, partial_img_emb, clip_fea, partial_control_latents, partial_vace_context, partial_unianim_data,partial_audio_proj, + partial_timestep, idx, image_cond_in, clip_fea, partial_control_latents, partial_vace_context, partial_unianim_data,partial_audio_proj, partial_control_camera_latents, partial_add_cond, current_teacache, context_window=c, fantasy_portrait_input=partial_fantasy_portrait_input, mtv_motion_tokens=partial_mtv_motion_tokens, s2v_audio_input=partial_s2v_audio_input, s2v_motion_frames=[1, 0], s2v_pose=partial_s2v_pose, humo_image_cond=humo_image_cond, humo_image_cond_neg=humo_image_cond_neg, humo_audio=humo_audio, humo_audio_neg=humo_audio_neg, wananim_face_pixels=partial_wananim_face_pixels, wananim_pose_latents=partial_wananim_pose_latents, multitalk_audio_embeds=multitalk_audio_embeds, uni3c_data=uni3c_data, flashvsr_LQ_latent=partial_flashvsr_LQ_latent) + seq_len = original_seq_len # restore seq_len + + # ============ Prefix: slice off prepended predictions ============ + if has_prefix and c[0] != 0: + noise_pred_context = noise_pred_context[:, 6:] + # ============================================================== + if cache_args is not None: self.window_tracker.cache_states[window_id] = new_teacache @@ -2211,6 +2311,8 @@ def predict_with_cfg(z, cfg_scale, positive_embeds, negative_embeds, timestep, i ref_images = image_embeds.get("ref_image", None) bg_images = image_embeds.get("bg_images", None) pose_images = image_embeds.get("pose_images", None) + prefix_ctx = image_embeds.get("prefix_ctx", None) + prefix_T = image_embeds.get("prefix_T", 0) current_ref_images = image_embeds.get("start_ref_image", None) if current_ref_images is not None: @@ -2262,10 +2364,10 @@ def predict_with_cfg(z, cfg_scale, positive_embeds, negative_embeds, timestep, i self.cache_state = [None, None] - noise = torch.randn(16, latent_window_size + 1, lat_h, lat_w, dtype=torch.float32, device=torch.device("cpu"), generator=seed_g).to(device) + noise = torch.randn(16, latent_window_size + 1 + prefix_T, lat_h, lat_w, dtype=torch.float32, device=torch.device("cpu"), generator=seed_g).to(device) seq_len = math.ceil((noise.shape[2] * noise.shape[3]) / 4 * noise.shape[1]) - if current_ref_images is not None or bg_images is not None or ref_latent is not None: + if current_ref_images is not None or bg_images is not None or ref_latent is not None or has_transition or prefix_ctx is not None: if offload: offload_transformer(transformer, remove_lora=False) offloaded = True @@ -2278,14 +2380,57 @@ def predict_with_cfg(z, cfg_scale, positive_embeds, negative_embeds, timestep, i bg_image_slice = bg_images_in[:, start:end].to(device) else: bg_image_slice = torch.zeros(3, frame_window_size-refert_num, lat_h * 8, lat_w * 8, device=device, dtype=vae.dtype) - if mask_reft_len == 0: + if mask_reft_len == 0 and not has_transition: temporal_ref_latents = vae.encode([bg_image_slice], device,tiled=tiled_vae)[0] else: - concatenated = torch.cat([current_ref_images.to(device, dtype=vae.dtype), bg_image_slice[:, mask_reft_len:]], dim=1) - temporal_ref_latents = vae.encode([concatenated.to(device, vae.dtype)], device,tiled=tiled_vae, pbar=False)[0] - msk[:, :mask_reft_len] = 1 - - if msk.shape[1] != temporal_ref_latents.shape[1]: + # Build concatenated image + concat_parts = [] + + # 1. Add start_ref_image (if any) + if current_ref_images is not None: + concat_parts.append(current_ref_images.to(device, dtype=vae.dtype)) + + # 2. Add the remaining part of bg_image_slice + bg_start_idx = mask_reft_len if current_ref_images is not None else 0 + if bg_image_slice.shape[1] > bg_start_idx: + concat_parts.append(bg_image_slice[:, bg_start_idx:]) + + # 3. VAE encode the image part + if len(concat_parts) > 0: + image_concat = torch.cat(concat_parts, dim=1) + temporal_ref_latents = vae.encode([image_concat.to(device, vae.dtype)], device, tiled=tiled_vae, pbar=False)[0] + else: + temporal_ref_latents = None + + # 4. Initialize mask + msk = torch.zeros(4, latent_window_size, lat_h, lat_w, device=device, dtype=dtype) + + # 4a. Hard mask for start_ref_image + if current_ref_images is not None: + msk[:, :mask_reft_len] = 1 + + # 4b. Hard mask for transition_latent (replaces first 8 target latent frames ON THE FIRST CHUNK ONLY) + if has_transition and start == 0: + if temporal_ref_latents is not None: + # ============ Mod: Replace first transition_len frames of target latent ============ + # Keep target latent after transition_len + remaining_target = temporal_ref_latents[:, transition_len:] + # transition_latent replaces the first transition_len target latent frames + temporal_ref_latents = torch.cat([transition_latent.to(device, dtype=transition_latent.dtype), remaining_target], dim=1) + log.info(f"Replaced first {transition_len} target latent frames with transition_video in Chunk 0") + # ======================================================================== + else: + # Only transition_latent, no image part + temporal_ref_latents = transition_latent.to(device, dtype=dtype) + + # Set hard mask. + transition_start = 0 + msk[:, transition_start:transition_start+transition_len] = transition_mask_values.to(device).view(1, -1, 1, 1) + + # [FIX] Unified shape alignment for temporal_ref_latents and msk + # Handles VAE time downsampling offset: (T-1)//4 + 1 vs latent_window_size + # Moved outside branches to ensure execution for both transition and non-transition modes + if temporal_ref_latents is not None and msk.shape[1] != temporal_ref_latents.shape[1]: if temporal_ref_latents.shape[1] < msk.shape[1]: pad_len = msk.shape[1] - temporal_ref_latents.shape[1] pad_tensor = temporal_ref_latents[:, -1:].repeat(1, pad_len, 1, 1) @@ -2293,9 +2438,10 @@ def predict_with_cfg(z, cfg_scale, positive_embeds, negative_embeds, timestep, i else: temporal_ref_latents = temporal_ref_latents[:, :msk.shape[1]] - if ref_latent is not None: + if ref_latent is not None or prefix_ctx is not None: + ref_part = prefix_ctx if prefix_ctx is not None else ref_latent temporal_ref_latents = torch.cat([msk, temporal_ref_latents], dim=0) # 4+C T H W - image_cond_in = torch.cat([ref_latent.to(device), temporal_ref_latents], dim=1) # 4+C T+trefs H W + image_cond_in = torch.cat([ref_part.to(device), temporal_ref_latents], dim=1) # 4+C T+trefs H W del temporal_ref_latents, msk, bg_image_slice else: image_cond_in = torch.cat([torch.tile(torch.zeros_like(noise[:1]), [4, 1, 1, 1]), torch.zeros_like(noise)], dim=0).to(device) @@ -2315,6 +2461,16 @@ def predict_with_cfg(z, cfg_scale, positive_embeds, negative_embeds, timestep, i elif wananim_face_pixels is not None: face_images_in = face_images[:, :, start:end].to(device, torch.float32) if face_images is not None else None + # ============ Prefix prepend for looping: copy first frame of current chunk ============ + if prefix_T > 0: + if pose_input_slice is not None: + first_pose = pose_input_slice[:, :, :1].repeat(1, 1, prefix_T, 1, 1) + pose_input_slice = torch.cat([first_pose, pose_input_slice], dim=2) + if face_images_in is not None: + first_face = face_images_in[:, :, :1].repeat(1, 1, prefix_T * 4, 1, 1) + face_images_in = torch.cat([first_face, face_images_in], dim=2) + # ======================================================================================== + if samples is not None: input_samples = samples["samples"] if input_samples is not None: @@ -2407,6 +2563,9 @@ def predict_with_cfg(z, cfg_scale, positive_embeds, negative_embeds, timestep, i timestep, i, cache_state=self.cache_state, image_cond=image_cond_in, clip_fea=clip_fea, wananim_face_pixels=face_images_in, wananim_pose_latents=pose_input_slice, uni3c_data=uni3c_data_input, ) + if prefix_T > 0: + noise_pred = noise_pred[:, prefix_T:] + latent = latent[:, prefix_T:] if callback is not None: callback_latent = (latent_model_input.to(device) - noise_pred.to(device) * t.to(device) / 1000).detach().permute(1,0,2,3) callback(step_iteration_count, callback_latent, None, estimated_iterations*(len(timesteps))) @@ -2636,6 +2795,8 @@ def predict_with_cfg(z, cfg_scale, positive_embeds, negative_embeds, timestep, i "looped": is_looped, "end_image": end_image if not fun_or_fl2v_model else None, "has_ref": has_ref, + "has_prefix": has_prefix, + "canvas_expansion_px": canvas_expansion_px, "drop_last": drop_last, "generator_state": seed_g.get_state(), "original_image": original_image.cpu() if original_image is not None else None,