diff --git a/README.md b/README.md index 88cda23..4b6fa83 100644 --- a/README.md +++ b/README.md @@ -72,12 +72,12 @@ Generates condition latents from source video, reference video, and/or reference Provides SCAIL-2 ref / pose / mask conditioning through `WanAnimatePlus SCAIL_2 Embeds`. -- Encodes `ref_image`, `pose_images`, `pose_image_mask`, `prefix_mask`, and `reference_image_mask` +- Encodes `ref_image`, `pose_images`, `pose_image_mask`, `prefix_frames`, `prefix_mask`, `bg_image`, and `reference_image_mask` - Aligns SCAIL-2 inputs to 32-pixel multiples before VAE encoding - Supports animation and replacement modes -- Supports optional `prefix_frames` and `transition_video` hard-freeze conditioning -- Uses a single SCAIL-2 prefix layout: 37 front pixel frames when `prefix_frames` is connected, with transition frames placed at frames 17-36 when `transition_video` is also connected; transition-only uses 21 front pixel frames -- Writes `prefix_mask` into the prefix pixel-mask frames before SCAIL-2 mask latent encoding +- Supports single-frame prefix reference encoding and optional `transition_video` hard-freeze conditioning +- By default, `prefix_frames` are encoded as full-resolution reference latents and do not expand the output canvas; disable `single_frame_prefix_encoding` to use the legacy 37 front pixel-frame prefix layout +- In single-frame prefix mode, `prefix_mask` follows the same reference-mask path as `reference_image_mask` - Supports context-window sampling; non-first windows can see prepended prefix/transition context without fusing those prepended predictions ## Installation @@ -167,19 +167,24 @@ Creates SCAIL-2 conditioning for WanAnimatePlus sampling. Use this node with SCA | `vae` | VAE model for encoding | | `width` / `height` / `num_frames` | Target dimensions; width and height are aligned to multiples of 32 | | `ref_image` | Reference image for SCAIL-2 conditioning | +| `bg_image` | Optional single background image for animation mode; prepended as the first `prefix_frames` item with an internal white mask and ignored in replacement mode | | `pose_images` | Driving pose video/images, encoded at half resolution | | `pose_image_mask` | Colored per-identity pose mask sequence | | `prefix_mask` | Optional colored mask images matching `prefix_frames`; expanded as `1+4+4...` and written into the prefix mask frames before mask latent encoding | | `reference_image_mask` | Colored reference mask image | | `replacement_mode` | Enables SCAIL-2 replacement-mode RoPE and reference-mask compositing | -| `prefix_frames` | Optional frames to hard-freeze at the front of the latent sequence | +| `preserve_main_ref_background` | Animation mode only; keeps the main reference image background when enabled, or uses `reference_image_mask` as a black-background alpha crop when disabled. Ignored in replacement mode | +| `single_frame_prefix_encoding` | Encodes `prefix_frames` as individual full-resolution reference latents instead of expanding the canvas; enabled by default | +| `prefix_frames` | Optional prefix images. In default single-frame mode these become reference-stream latents; with single-frame mode disabled they hard-freeze the front canvas | | `transition_video` | Optional transition frames to hard-freeze at the front of the latent sequence | | `clip_embeds` | Optional CLIP vision features from `WanAnimatePlus ClipVisionEncode` | | `force_offload` / `tiled_vae` | Memory controls for VAE encoding | -For short generations, context windows are optional. For long generations or low VRAM, context windows are recommended. In context-window mode, the sampler prepends protected prefix/transition latents for model context and removes those prepended predictions before overlap fusion. +For short generations, context windows are optional. For long generations or low VRAM, context windows are recommended. In context-window mode, single-frame prefix references remain visible through the SCAIL-2 reference stream. Legacy canvas prefixes and transition latents are prepended for model context and removed before overlap fusion. -For SCAIL-2, `prefix_frames` always expands the front canvas by 37 pixel frames. If only `transition_video` is connected, the front canvas expands by 21 pixel frames. The separate 45-frame outfit layout belongs to the regular AnimateEmbeds node and is not used by `WanAnimatePlus SCAIL_2 Embeds`. +For SCAIL-2, the default `single_frame_prefix_encoding` mode does not expand or trim the output for `prefix_frames`. If `transition_video` is connected, the front canvas expands by 21 pixel frames and those 21 frames are trimmed after decoding. With `single_frame_prefix_encoding` disabled, `prefix_frames` use the legacy 37 front pixel-frame canvas layout, with transition frames placed at frames 17-36 when `transition_video` is also connected. + +When `bg_image` is connected in animation mode, it consumes the first prefix slot and is handled as the first prefix image. The node internally adds a white mask for that background image only; user-provided `prefix_frames` and `prefix_mask` are then limited to four images each. In replacement mode, `bg_image` is ignored. ## Project Structure diff --git a/nodes.py b/nodes.py index b40cf41..72e8749 100644 --- a/nodes.py +++ b/nodes.py @@ -1653,10 +1653,11 @@ class WanAnimatePlusSCAIL2Embeds: "optional": { "clip_embeds": ("WANVIDIMAGE_CLIPEMBEDS", {"tooltip": "Clip vision encoded image"}), "ref_image": ("IMAGE", {"tooltip": "Reference image for SCAIL conditioning. If a sequence is connected, only the first frame is used."}), + "bg_image": ("IMAGE", {"tooltip": "Optional single background image for animation mode. Prepended before prefix_frames with an internal white mask and ignored in replacement mode."}), "pose_images": ("IMAGE", {"tooltip": "Driving pose video. Encoded at half resolution for SCAIL."}), - "prefix_frames": ("IMAGE", {"tooltip": "Optional frames to hard-freeze at the beginning of the SCAIL-2 latent sequence."}), - "prefix_mask": ("IMAGE", {"tooltip": "Optional colored mask images matching prefix_frames. Expanded as 1+4+4... and written into the prefix mask frames before mask latent encoding. In replacement mode it also alpha-crops matching prefix frames to black background."}), - "transition_video": ("IMAGE", {"tooltip": "Optional transition frames to hard-freeze after prefix_frames at the beginning of the SCAIL-2 latent sequence."}), + "prefix_frames": ("IMAGE", {"tooltip": "Optional prefix images. In single-frame prefix mode these are encoded as reference latents; in legacy mode they hard-freeze the beginning of the canvas."}), + "prefix_mask": ("IMAGE", {"tooltip": "Optional colored mask images matching prefix_frames. In single-frame prefix mode this follows the reference-mask path; in legacy canvas-prefix mode it is expanded as 1+4+4... and written into the prefix mask frames."}), + "transition_video": ("IMAGE", {"tooltip": "Optional transition frames to hard-freeze at the beginning of the canvas. In legacy canvas-prefix mode, transition frames are placed after the prefix frames."}), "pose_image_mask": ("IMAGE", {"tooltip": "SCAIL-2 colored per-identity driving pose mask. Background is normalized to black in animation mode and white in replacement mode."}), "reference_image_mask": ("IMAGE", {"tooltip": "SCAIL-2 colored per-identity reference mask image. Background is normalized to white in animation mode and black in replacement mode."}), "tiled_vae": ("BOOLEAN", {"default": False, "tooltip": "Use tiled VAE encoding for reduced memory use"}), @@ -1669,7 +1670,9 @@ class WanAnimatePlusSCAIL2Embeds: 'hm-mvgd-hm', 'hm-mkl-hm', ], {"default": 'disabled', "tooltip": "Color match transition_video to ref_image."}), - "prefix_alpha_crop": ("BOOLEAN", {"default": False, "tooltip": "Off writes prefix_mask with white background and keeps prefix frames unchanged. On writes prefix_mask with black background and alpha-crops prefix_frames."}), + "prefix_alpha_crop": ("BOOLEAN", {"default": False, "tooltip": "Off keeps prefix masks as white-background reference masks in animation mode. On uses black-background masks and alpha-crops prefix_frames. Replacement mode always uses black-background reference masks."}), + "preserve_main_ref_background": ("BOOLEAN", {"default": True, "tooltip": "Animation mode only. Keep the main reference image background. When off, reference_image_mask is normalized to black background and used to alpha-crop ref_image. Ignored in replacement mode."}), + "single_frame_prefix_encoding": ("BOOLEAN", {"default": True, "tooltip": "Encode prefix images as individual full-resolution reference latents instead of expanding the canvas."}), "by wuwukasi（bilibili）": ("BOOLEAN", {"default": True, "label_on": "ON", "label_off": "ON", "tooltip": "Follow wuwukasi on bilibili"}), } } @@ -1775,6 +1778,20 @@ class WanAnimatePlusSCAIL2Embeds: bg = mask.amin(dim=-1, keepdim=True) >= 0.95 return torch.where(bg, torch.zeros_like(mask), mask) + @staticmethod + def _alpha_crop_with_mask(images, masks): + if images is None or masks is None or images.shape[0] == 0 or masks.shape[0] == 0: + return images + count = min(images.shape[0], masks.shape[0]) + crop_mask = masks[:count, :, :, :3] + if crop_mask.shape[1] != images.shape[1] or crop_mask.shape[2] != images.shape[2]: + crop_mask = WanAnimatePlusSCAIL2Embeds._resize_bhwc(crop_mask, images.shape[2], images.shape[1], mode="nearest-exact") + crop_mask = WanAnimatePlusSCAIL2Embeds._normalize_mask_background(crop_mask, white_background=False) + is_char = (crop_mask[..., :3].max(dim=-1, keepdim=True).values > 0.1).to(device=images.device, dtype=images.dtype) + out = images.clone() + out[:count] = out[:count] * is_char + return out + @staticmethod def _color_match_frames(frames, ref_frame, method): if method == "disabled" or ref_frame is None or frames is None or frames.shape[0] == 0: @@ -1819,23 +1836,121 @@ class WanAnimatePlusSCAIL2Embeds: padded = padded[:t_lat * 4] return padded.view(t_lat, 28, h_lat, w_lat).movedim(0, 1).contiguous() + @staticmethod + def _empty_ref_mask_like(ref_latent): + return torch.zeros( + 28, ref_latent.shape[1], ref_latent.shape[-2], ref_latent.shape[-1], + device=ref_latent.device, dtype=ref_latent.dtype, + ) + def process(self, vae, width, height, num_frames, force_offload, pose_strength, ref_strength, pose_start_percent, pose_end_percent, ref_start_percent, ref_end_percent, replacement_mode, - clip_embeds=None, ref_image=None, pose_images=None, prefix_frames=None, prefix_mask=None, + clip_embeds=None, ref_image=None, bg_image=None, pose_images=None, prefix_frames=None, prefix_mask=None, transition_video=None, pose_image_mask=None, reference_image_mask=None, tiled_vae=False, - transition_colormatch='disabled', prefix_alpha_crop=False, **kwargs): + transition_colormatch='disabled', prefix_alpha_crop=False, preserve_main_ref_background=True, + single_frame_prefix_encoding=True, **kwargs): W = (width // 32) * 32 H = (height // 32) * 32 num_frames = ((num_frames - 1) // 4) * 4 + 1 + bg_prefix_mask_pixel_frames = 0 + bg_prefix_mask_index = None + crop_main_ref_background = (not replacement_mode) and (not preserve_main_ref_background) + + user_prefix_frames = prefix_frames + user_prefix_mask = prefix_mask if user_prefix_frames is not None else None + bg_prefix_active = bg_image is not None and not replacement_mode + if bg_image is not None and replacement_mode: + log.info("SCAIL-2 bg_image is ignored in replacement mode") + + bg_prefix_image = None + bg_mask = None + user_prefix_count = 0 + if bg_prefix_active: + if bg_image.shape[0] > 1: + log.warning("SCAIL-2 bg_image accepts one image; using the first frame") + bg_prefix_image = bg_image[:1, :, :, :3] + + if user_prefix_frames is not None: + if user_prefix_frames.shape[0] > 4: + log.warning("SCAIL-2 bg_image uses one prefix slot; truncating prefix_frames to 4 images") + user_prefix_frames = user_prefix_frames[:4, :, :, :3] + user_prefix_count = user_prefix_frames.shape[0] + if bg_prefix_image.shape[1] != user_prefix_frames.shape[1] or bg_prefix_image.shape[2] != user_prefix_frames.shape[2]: + bg_prefix_image = self._resize_bhwc(bg_prefix_image, user_prefix_frames.shape[2], user_prefix_frames.shape[1]) + bg_prefix_image = bg_prefix_image.to(device=user_prefix_frames.device, dtype=user_prefix_frames.dtype) + + if user_prefix_mask is not None: + if user_prefix_mask.shape[0] > 4: + log.warning("SCAIL-2 bg_image uses one prefix-mask slot; truncating prefix_mask to 4 images") + user_prefix_mask = user_prefix_mask[:4, :, :, :3] + bg_mask = torch.ones( + 1, user_prefix_mask.shape[1], user_prefix_mask.shape[2], 3, + device=user_prefix_mask.device, dtype=user_prefix_mask.dtype, + ) + else: + bg_mask = torch.ones_like(bg_prefix_image[:, :, :, :3]) + + if user_prefix_frames is not None and user_prefix_mask is not None and (replacement_mode or prefix_alpha_crop): + crop_count = min(user_prefix_frames.shape[0], user_prefix_mask.shape[0], 5) + if crop_count > 0: + cropped_prefix_frames = user_prefix_frames.clone() + cropped_prefix_frames[:crop_count] = self._alpha_crop_with_mask( + cropped_prefix_frames[:crop_count], + user_prefix_mask[:crop_count], + ) + user_prefix_frames = cropped_prefix_frames + + prefix_frames = user_prefix_frames + prefix_mask_for_prefix = user_prefix_mask + if bg_prefix_active: + if single_frame_prefix_encoding: + if user_prefix_frames is not None: + prefix_frames = torch.cat([user_prefix_frames, bg_prefix_image], dim=0) + bg_prefix_mask_index = prefix_frames.shape[0] - 1 + else: + prefix_frames = bg_prefix_image + bg_prefix_mask_index = 0 + if user_prefix_mask is not None: + if user_prefix_mask.shape[0] < user_prefix_count: + empty_user_masks = torch.zeros( + user_prefix_count - user_prefix_mask.shape[0], + user_prefix_mask.shape[1], user_prefix_mask.shape[2], 3, + device=user_prefix_mask.device, dtype=user_prefix_mask.dtype, + ) + user_prefix_mask = torch.cat([user_prefix_mask, empty_user_masks], dim=0) + elif user_prefix_mask.shape[0] > user_prefix_count: + user_prefix_mask = user_prefix_mask[:user_prefix_count] + prefix_mask_for_prefix = torch.cat([user_prefix_mask, bg_mask], dim=0) + elif user_prefix_count > 0: + empty_user_masks = torch.zeros( + user_prefix_count, bg_mask.shape[1], bg_mask.shape[2], 3, + device=bg_mask.device, dtype=bg_mask.dtype, + ) + prefix_mask_for_prefix = torch.cat([empty_user_masks, bg_mask], dim=0) + else: + prefix_mask_for_prefix = bg_mask + else: + if user_prefix_frames is not None: + prefix_frames = torch.cat([bg_prefix_image, user_prefix_frames], dim=0) + else: + prefix_frames = bg_prefix_image + if user_prefix_mask is not None: + prefix_mask_for_prefix = torch.cat([bg_mask, user_prefix_mask], dim=0) + else: + prefix_mask_for_prefix = bg_mask + bg_prefix_mask_pixel_frames = 1 + bg_prefix_mask_index = 0 + canvas_expansion_px = 0 freeze_canvas = None - if prefix_frames is not None: + canvas_prefix_frames = prefix_frames if not single_frame_prefix_encoding else None + if canvas_prefix_frames is not None: canvas_expansion_px = 37 elif transition_video is not None: canvas_expansion_px = 21 transition_px_range = None if transition_video is not None: - transition_px_range = (17, 37) if prefix_frames is not None else (0, 21) + transition_px_range = (17, 37) if canvas_prefix_frames is not None else (0, 21) if canvas_expansion_px: num_frames += canvas_expansion_px @@ -1847,16 +1962,16 @@ class WanAnimatePlusSCAIL2Embeds: pose_image_mask = torch.cat([_sample_reversed_prefix_frames(pose_image_mask, canvas_expansion_px), pose_image_mask], dim=0) prefix_mask_pixels = prefix_mask_frame_mask = scail_sam_keep_mask = scail_transition_keep_mask = None - if prefix_mask is not None: + if prefix_mask_for_prefix is not None: if prefix_frames is None: log.warning("SCAIL-2 prefix_mask was provided without prefix_frames. Ignoring prefix_mask.") else: prefix_count = min(prefix_frames.shape[0], 5) - if prefix_mask.shape[0] > prefix_count: - log.warning(f"SCAIL-2 prefix_mask has {prefix_mask.shape[0]} images, but prefix_frames has {prefix_count}. Truncating prefix_mask.") - mask_count = min(prefix_mask.shape[0], prefix_count) + if prefix_mask_for_prefix.shape[0] > prefix_count: + log.warning(f"SCAIL-2 prefix_mask has {prefix_mask_for_prefix.shape[0]} images, but prefix_frames has {prefix_count}. Truncating prefix_mask.") + mask_count = min(prefix_mask_for_prefix.shape[0], prefix_count) if mask_count > 0: - prefix_mask_pixels = self._build_prefix_pixels(prefix_mask[:mask_count]) + prefix_mask_pixels = self._build_prefix_pixels(prefix_mask_for_prefix[:mask_count]) lat_h = H // vae.upsampling_factor lat_w = W // vae.upsampling_factor @@ -1875,17 +1990,62 @@ class WanAnimatePlusSCAIL2Embeds: "ref_mask_flag": not replacement_mode, } + ref_latents = [] + ref_masks = [] + ref_mask_condition_used = False + prefix_ref_count = 0 + prefix_ref_black_background = replacement_mode or prefix_alpha_crop + if single_frame_prefix_encoding and prefix_frames is not None: + prefix_ref_images = prefix_frames[:5, :, :, :3] + prefix_ref_source_indices = list(range(prefix_ref_images.shape[0])) + prefix_ref_count = prefix_ref_images.shape[0] + prefix_ref_masks = [None] * prefix_ref_count + if prefix_mask_for_prefix is not None: + for out_idx, src_idx in enumerate(prefix_ref_source_indices): + if src_idx >= prefix_mask_for_prefix.shape[0]: + continue + mask = prefix_mask_for_prefix[src_idx:src_idx + 1, :, :, :3] + if src_idx != bg_prefix_mask_index: + mask = self._normalize_mask_background(mask, white_background=not prefix_ref_black_background) + prefix_ref_masks[out_idx] = mask + ref_mask_condition_used = any(mask is not None for mask in prefix_ref_masks) + for i in range(prefix_ref_images.shape[0]): + prefix_latent = self._encode_scail_20ch(vae, prefix_ref_images[i:i + 1], W, H, tiled_vae, scale=ref_strength).to(offload_device) + ref_latents.append(prefix_latent) + if i < len(prefix_ref_masks) and prefix_ref_masks[i] is not None: + ref_mask = self._resize_bhwc(prefix_ref_masks[i], W, H, mode="bicubic") + ref_masks.append(self._extract_mask_to_28ch(ref_mask).to(offload_device, vae.dtype)) + else: + ref_masks.append(self._empty_ref_mask_like(prefix_latent)) + log.info(f"SCAIL-2 prefix reference latents: {prefix_ref_count}") + if ref_image is not None: ref_image = ref_image[:1, :, :, :3] - if replacement_mode and reference_image_mask is not None: + if (replacement_mode or crop_main_ref_background) and reference_image_mask is not None: ref_mask = self._resize_bhwc(reference_image_mask[:1, :, :, :3], W, H, mode="nearest-exact") ref_mask = self._normalize_mask_background(ref_mask, white_background=False) is_char = (ref_mask[..., :3].max(dim=-1, keepdim=True).values > 0.1).to(ref_image.dtype) - ref_image = self._resize_bhwc(ref_image, W, H) * is_char + ref_image = self._resize_bhwc(ref_image, W, H) + is_char = is_char.to(device=ref_image.device, dtype=ref_image.dtype) + ref_image = ref_image * is_char ref_latent = self._encode_scail_20ch(vae, ref_image, W, H, tiled_vae, scale=ref_strength).to(offload_device) + ref_latents.append(ref_latent) + if reference_image_mask is not None: + ref_mask_condition_used = True + ref_mask = self._fit_sequence(reference_image_mask[:, :, :, :3], 1) + ref_mask = self._normalize_mask_background(ref_mask, white_background=not (replacement_mode or crop_main_ref_background)) + ref_mask = self._resize_bhwc(ref_mask, W, H, mode="bicubic") + ref_masks.append(self._extract_mask_to_28ch(ref_mask).to(offload_device, vae.dtype)) + else: + ref_masks.append(self._empty_ref_mask_like(ref_latent)) + log.info(f"SCAIL-2 reference latent shape: {ref_latent.shape}") + + if ref_latents: + ref_latent = torch.cat(ref_latents, dim=1) scail_embeds["ref_latent_pos"] = ref_latent scail_embeds["ref_latent_neg"] = ref_latent - log.info(f"SCAIL-2 reference latent shape: {ref_latent.shape}") + scail_embeds["ref_prefix_latents"] = prefix_ref_count + log.info(f"SCAIL-2 combined reference latent shape: {ref_latent.shape}") scail_freeze_latents = scail_freeze_mask = scail_condition_zero_mask = None actual_prefix_px = 0 @@ -1895,19 +2055,9 @@ class WanAnimatePlusSCAIL2Embeds: transition_match_ref = self._resize_bhwc(ref_image[:1, :, :, :3], W, H) if ref_image is not None else None if transition_colormatch != "disabled" and transition_match_ref is None: log.warning("SCAIL-2 transition_colormatch is enabled but ref_image is not connected. Skipping color match.") - if prefix_frames is not None: - prefix_pixels = self._build_prefix_pixels(prefix_frames) + if canvas_prefix_frames is not None: + prefix_pixels = self._build_prefix_pixels(canvas_prefix_frames) actual_prefix_px = min(prefix_pixels.shape[0], 17) - if prefix_alpha_crop and prefix_mask_pixels is not None: - crop_len = min(actual_prefix_px, prefix_mask_pixels.shape[0]) - if crop_len > 0: - prefix_alpha = prefix_mask_pixels[:crop_len, :, :, :3] - if prefix_alpha.shape[1] != prefix_pixels.shape[1] or prefix_alpha.shape[2] != prefix_pixels.shape[2]: - prefix_alpha = self._resize_bhwc(prefix_alpha, prefix_pixels.shape[2], prefix_pixels.shape[1], mode="nearest-exact") - prefix_alpha = self._normalize_mask_background(prefix_alpha, white_background=False) - is_char = (prefix_alpha[..., :3].max(dim=-1, keepdim=True).values > 0.1).to(prefix_pixels.dtype) - prefix_pixels = prefix_pixels.clone() - prefix_pixels[:crop_len] = prefix_pixels[:crop_len] * is_char prefix_pixels = self._resize_bhwc(prefix_pixels[:actual_prefix_px], W, H) freeze_canvas[:actual_prefix_px] = prefix_pixels.to(device, dtype=freeze_canvas.dtype) freeze_frame_mask[:actual_prefix_px] = 1.0 @@ -1930,7 +2080,7 @@ class WanAnimatePlusSCAIL2Embeds: scail_freeze_mask = scail_freeze_mask.to(offload_device) scail_condition_zero_mask = (scail_freeze_mask > 0).any(dim=(1, 2)) log.info(f"SCAIL-2 freeze latents shape: {scail_freeze_latents.shape}") - if prefix_frames is not None and actual_prefix_px > 0: + if canvas_prefix_frames is not None and actual_prefix_px > 0: prefix_frame_mask = torch.zeros(canvas_expansion_px, device=device, dtype=vae.dtype) prefix_frame_mask[:actual_prefix_px] = 1.0 scail_prefix_prepend_latents = int(self._frame_mask_to_latent_mask(prefix_frame_mask, target_latents).sum().item()) @@ -1967,7 +2117,7 @@ class WanAnimatePlusSCAIL2Embeds: if pose_image_mask is not None: pose_image_mask = self._fit_sequence(pose_image_mask[:, :, :, :3], num_frames) pose_image_mask = self._normalize_mask_background(pose_image_mask, white_background=replacement_mode) - elif prefix_mask_pixels is not None: + elif prefix_mask_pixels is not None and not single_frame_prefix_encoding: pose_bg = 1.0 if replacement_mode else 0.0 pose_image_mask = torch.full( ( @@ -1976,12 +2126,18 @@ class WanAnimatePlusSCAIL2Embeds: pose_bg, device=prefix_mask_pixels.device, dtype=prefix_mask_pixels.dtype, ) - if pose_image_mask is not None and prefix_mask_pixels is not None: + if pose_image_mask is not None and prefix_mask_pixels is not None and not single_frame_prefix_encoding: prefix_mask_len = min(prefix_mask_pixels.shape[0], pose_image_mask.shape[0]) prefix_mask_in = prefix_mask_pixels[:prefix_mask_len, :, :, :3] if prefix_mask_in.shape[1] != pose_image_mask.shape[1] or prefix_mask_in.shape[2] != pose_image_mask.shape[2]: prefix_mask_in = self._resize_bhwc(prefix_mask_in, pose_image_mask.shape[2], pose_image_mask.shape[1], mode="nearest-exact") - prefix_mask_in = self._normalize_mask_background(prefix_mask_in, white_background=not prefix_alpha_crop) + bg_keep = min(bg_prefix_mask_pixel_frames, prefix_mask_in.shape[0]) + if prefix_mask_in.shape[0] > bg_keep: + user_prefix_mask_in = self._normalize_mask_background(prefix_mask_in[bg_keep:], white_background=not prefix_alpha_crop) + if bg_keep > 0: + prefix_mask_in = torch.cat([prefix_mask_in[:bg_keep], user_prefix_mask_in], dim=0) + else: + prefix_mask_in = user_prefix_mask_in pose_image_mask = pose_image_mask.clone() pose_image_mask[:prefix_mask_len] = prefix_mask_in.to(device=pose_image_mask.device, dtype=pose_image_mask.dtype) prefix_mask_frame_mask = torch.zeros(num_frames, device=pose_image_mask.device, dtype=pose_image_mask.dtype) @@ -2010,16 +2166,13 @@ class WanAnimatePlusSCAIL2Embeds: scail_embeds["sam_latents"] = sam_latents log.info(f"SCAIL-2 driving mask latents shape: {sam_latents.shape}") - if reference_image_mask is not None: - ref_mask = self._fit_sequence(reference_image_mask[:, :, :, :3], 1) - ref_mask = self._normalize_mask_background(ref_mask, white_background=not replacement_mode) - ref_mask = self._resize_bhwc(ref_mask, W, H, mode="bicubic") - ref_mask_1f = self._extract_mask_to_28ch(ref_mask).to(offload_device, vae.dtype) + if ref_masks and ref_mask_condition_used: + ref_mask_prefix = torch.cat(ref_masks, dim=1) zeros = torch.zeros( - 28, target_latents, ref_mask_1f.shape[-2], ref_mask_1f.shape[-1], - device=offload_device, dtype=ref_mask_1f.dtype, + 28, target_latents, ref_mask_prefix.shape[-2], ref_mask_prefix.shape[-1], + device=offload_device, dtype=ref_mask_prefix.dtype, ) - scail_embeds["ref_mask_latents"] = torch.cat([ref_mask_1f, zeros], dim=1) + scail_embeds["ref_mask_latents"] = torch.cat([ref_mask_prefix, zeros], dim=1) log.info(f"SCAIL-2 reference mask latents shape: {scail_embeds['ref_mask_latents'].shape}") if force_offload: diff --git a/nodes_sampler.py b/nodes_sampler.py index 4d611ad..71cff37 100644 --- a/nodes_sampler.py +++ b/nodes_sampler.py @@ -1130,6 +1130,15 @@ class WanVideoSampler: tsr_k = experimental_args.get("tsr_k", 1.0) tsr_sigma = experimental_args.get("tsr_sigma", 1.0) + scail2_fast_path_blocked_by_experimental_sampling = ( + use_cfg_zero_star + or use_tangential + or raag_alpha > 0.0 + or use_fresca + or use_tsr + or bool(transformer.video_attention_split_steps) + ) + # Rotary positional embeddings (RoPE) # RoPE base freq scaling as used with CineScale @@ -1145,6 +1154,10 @@ class WanVideoSampler: if context_latents is not None and "comfy" not in rope_function: log.info("Bernini context requires comfy RoPE; overriding rope_function to comfy") rope_function = "comfy" + scail_embeds_for_rope = image_embeds.get("scail_embeds", None) + if scail_embeds_for_rope is not None and "comfy" not in rope_function: + log.info("SCAIL-2 requires comfy RoPE; overriding rope_function to comfy") + rope_function = "comfy" freqs = None @@ -1328,6 +1341,7 @@ class WanVideoSampler: log.info(f" {k}: {v.shape if isinstance(v, torch.Tensor) else v}") scail_data = scail_embeds.copy() scail_data = dict_to_device(scail_data, device, dtype) + scail2_fast_path_fallback_logged = False # WanMove @@ -1369,6 +1383,7 @@ class WanVideoSampler: context_latents=None, context_roles=None, context_window_start=0, scail_context_prepend_latents=0,): nonlocal transformer nonlocal audio_cfg_scale + nonlocal scail2_fast_path_fallback_logged autocast_enabled = ("fp8" in model["quantization"] and not transformer.patched_linear) with torch.autocast(device_type=mm.get_autocast_device(device), dtype=dtype) if autocast_enabled else nullcontext(): @@ -1813,7 +1828,187 @@ class WanVideoSampler: and base_params['num_memory_frames'] == 0 ) strict_plain_t2v_fast_path = context_latents is None and context_window is None and fast_path_no_extra_conditions - base_params["simple_t2v"] = context_latents is not None or strict_plain_t2v_fast_path + base_params["simple_t2v"] = (context_latents is not None and base_params['scail_input'] is None) or strict_plain_t2v_fast_path + scail2_fast_path = ( + base_params['scail_input'] is not None + and guidance_mode == "cfg" + and context_latents is None + and base_params['y'] is None + and image_cond_neg is None + and attn_cond is None + and attn_cond_neg is None + and not reverse_time + and not enhance_enabled + and cache_args is None + and freeinit_args is None + and not scail2_fast_path_blocked_by_experimental_sampling + and slg_args is None + and getattr(transformer, "slg_blocks", None) is None + and not getattr(transformer, "is_longcat", False) + and getattr(transformer, "audio_model", None) is None + and getattr(transformer, "vace_layers", None) is None + and not batched_cfg + and not bidirectional_sampling + and loop_args is None + and qwenvl_embeds_pos is None + and qwenvl_embeds_neg is None + and control_latents is None + and vace_data is None + and unianim_data is None + and audio_proj is None + and control_camera_latents is None + and add_cond is None + and multitalk_audio_embeds is None + and fantasy_portrait_input is None + and mtv_motion_tokens is None + and s2v_audio_input is None + and s2v_ref_motion is None + and s2v_pose is None + and humo_image_cond is None + and humo_audio is None + and wananim_pose_latents is None + and wananim_face_pixels is None + and uni3c_data is None + and latent_model_input_ovi is None + and flashvsr_LQ_latent is None + and base_params['fun_ref'] is None + and base_params['fun_camera'] is None + and base_params['audio_proj'] is None + and base_params['uni3c_data'] is None + and base_params['controlnet'] is None + and base_params['add_cond'] is None + and not base_params['nag_params'] + and base_params['nag_context'] is None + and base_params['multitalk_audio'] is None + and base_params['ref_target_masks'] is None + and base_params['inner_t'] is None + and base_params['standin_input'] is None + and base_params['fantasy_portrait_input'] is None + and base_params['phantom_ref'] is None + and base_params['mtv_motion_tokens'] is None + and base_params['mtv_motion_rotary_emb'] is None + and base_params['s2v_audio_input'] is None + and base_params['s2v_ref_latent'] is None + and base_params['s2v_ref_motion'] is None + and base_params['s2v_pose'] is None + and base_params['humo_audio'] is None + and base_params['wananim_pose_latents'] is None + and base_params['wananim_face_pixel_values'] is None + and base_params['lynx_embeds'] is None + and base_params['x_ovi'] is None + and base_params['flashvsr_LQ_latent'] is None + and base_params['longcat_num_cond_latents'] == 0 + and base_params['longcat_num_ref_latents'] == 0 + and base_params['longcat_avatar_options'] is None + and base_params['sdancer_input'] is None + and base_params['one_to_all_input'] is None + and base_params['dual_control_input'] is None + and base_params['rope_negative_offset'] == 0 + and base_params['num_memory_frames'] == 0 + and "comfy" in getattr(transformer, "rope_func", "") + ) + if base_params['scail_input'] is not None and not scail2_fast_path: + if not scail2_fast_path_fallback_logged: + scail2_fast_path_fallback_logged = True + fallback_reasons = [] + if guidance_mode != "cfg": + fallback_reasons.append(f"guidance_mode={guidance_mode!r}") + if context_latents is not None: + fallback_reasons.append("Bernini context_latents connected") + if base_params['y'] is not None: + fallback_reasons.append("image condition active") + if image_cond_neg is not None: + fallback_reasons.append("negative image condition active") + if attn_cond is not None or attn_cond_neg is not None: + fallback_reasons.append("attention condition active") + if reverse_time: + fallback_reasons.append("reverse_time enabled") + if enhance_enabled: + fallback_reasons.append("enhance enabled") + if cache_args is not None: + fallback_reasons.append("cache enabled") + if freeinit_args is not None: + fallback_reasons.append("FreeInit enabled") + if scail2_fast_path_blocked_by_experimental_sampling: + fallback_reasons.append("experimental sampling enabled") + if slg_args is not None or getattr(transformer, "slg_blocks", None) is not None: + fallback_reasons.append("SLG enabled") + if getattr(transformer, "is_longcat", False): + fallback_reasons.append("LongCat model active") + if getattr(transformer, "audio_model", None) is not None: + fallback_reasons.append("OVI audio model active") + if getattr(transformer, "vace_layers", None) is not None: + fallback_reasons.append("VACE layers active") + if batched_cfg: + fallback_reasons.append("batched CFG enabled") + if bidirectional_sampling: + fallback_reasons.append("bidirectional sampling enabled") + if loop_args is not None: + fallback_reasons.append("loop enabled") + if qwenvl_embeds_pos is not None or qwenvl_embeds_neg is not None: + fallback_reasons.append("QwenVL embeddings connected") + if control_latents is not None or control_camera_latents is not None: + fallback_reasons.append("control latents connected") + if vace_data is not None: + fallback_reasons.append("VACE input connected") + if unianim_data is not None: + fallback_reasons.append("UniAnimate input connected") + if audio_proj is not None: + fallback_reasons.append("audio projection connected") + if add_cond is not None: + fallback_reasons.append("additional condition connected") + if multitalk_audio_embeds is not None: + fallback_reasons.append("MultiTalk input connected") + if fantasy_portrait_input is not None: + fallback_reasons.append("FantasyPortrait input connected") + if mtv_motion_tokens is not None: + fallback_reasons.append("MTV motion tokens connected") + if s2v_audio_input is not None or s2v_ref_motion is not None or s2v_pose is not None: + fallback_reasons.append("S2V input connected") + if humo_image_cond is not None or humo_audio is not None: + fallback_reasons.append("HuMo input connected") + if wananim_pose_latents is not None or wananim_face_pixels is not None: + fallback_reasons.append("WanAnimate input connected") + if uni3c_data is not None: + fallback_reasons.append("Uni3C input connected") + if latent_model_input_ovi is not None: + fallback_reasons.append("OVI latent input connected") + if flashvsr_LQ_latent is not None: + fallback_reasons.append("FlashVSR input connected") + if base_params['fun_ref'] is not None or base_params['fun_camera'] is not None: + fallback_reasons.append("Fun reference/camera active") + if base_params['controlnet'] is not None: + fallback_reasons.append("ControlNet active") + if base_params['nag_params'] or base_params['nag_context'] is not None: + fallback_reasons.append("NAG active") + if base_params['ref_target_masks'] is not None: + fallback_reasons.append("reference target masks active") + if base_params['inner_t'] is not None: + fallback_reasons.append("EchoShot active") + if base_params['standin_input'] is not None: + fallback_reasons.append("Stand-In input connected") + if base_params['phantom_ref'] is not None: + fallback_reasons.append("Phantom reference connected") + if base_params['lynx_embeds'] is not None: + fallback_reasons.append("Lynx input connected") + if base_params['longcat_num_cond_latents'] != 0 or base_params['longcat_num_ref_latents'] != 0 or base_params['longcat_avatar_options'] is not None: + fallback_reasons.append("LongCat options active") + if base_params['sdancer_input'] is not None: + fallback_reasons.append("SteadyDancer input connected") + if base_params['one_to_all_input'] is not None: + fallback_reasons.append("One-to-All input connected") + if base_params['dual_control_input'] is not None: + fallback_reasons.append("Dual Control input connected") + if base_params['rope_negative_offset'] != 0 or base_params['num_memory_frames'] != 0: + fallback_reasons.append("StoryMem/memory RoPE active") + if "comfy" not in getattr(transformer, "rope_func", ""): + fallback_reasons.append(f"rope_func={getattr(transformer, 'rope_func', None)!r}") + if not fallback_reasons: + fallback_reasons.append("unsupported sampler/model option") + log.info("SCAIL-2 fast path disabled; using normal path. Reason(s): " + "; ".join(fallback_reasons)) + base_params["simple_scail2"] = scail2_fast_path + if scail2_fast_path: + base_params["simple_t2v"] = False batch_size = 1 diff --git a/wanvideo/modules/model.py b/wanvideo/modules/model.py index d73affa..561efd0 100644 --- a/wanvideo/modules/model.py +++ b/wanvideo/modules/model.py @@ -7,6 +7,7 @@ # - Added variable WanAnimate anchor count support for pose and face embedding offsets. # - Added a streamlined simple T2V/Bernini fast path with block-swap prefetch support. # - Added simple T2V text/time embedding caches, cache-key helpers, and guarded no-op device moves. +# - Added a streamlined SCAIL-2 CFG fast path for the local WanAnimatePlus sampler. # Licensed under the Apache License, Version 2.0 import math import torch @@ -2388,7 +2389,10 @@ class WanModel(torch.nn.Module): pose_img_ids = torch.zeros((pose_f_len_full, pose_h_len_full, pose_w_len_full, 3), device=device, dtype=dtype) global_h_offset, global_w_offset = 0, 120 # global spatial offset to separate pose from main frames spatially (SCAIL uses 120 as offset) - pose_t_start = 1 if scail_ref_frame_shape is not None else t_start + freq_offset + scail_ref_t_len = 0 + if scail_ref_frame_shape is not None: + scail_ref_t_len = ((scail_ref_frame_shape[-3] + (self.patch_size[0] // 2)) // self.patch_size[0]) + pose_t_start = scail_ref_t_len if scail_ref_frame_shape is not None else t_start + freq_offset pose_img_ids[:, :, :, 0] = pose_img_ids[:, :, :, 0] + torch.linspace(pose_t_start, pose_t_start + (pose_f_len_full - 1), steps=pose_f_len_full, device=device, dtype=dtype).reshape(-1, 1, 1) pose_img_ids[:, :, :, 1] = pose_img_ids[:, :, :, 1] + torch.linspace(global_h_offset + freq_offset, global_h_offset + pose_h_len_full - 1, steps=pose_h_len_full, device=device, dtype=dtype).reshape(1, -1, 1) pose_img_ids[:, :, :, 2] = pose_img_ids[:, :, :, 2] + torch.linspace(global_w_offset + freq_offset, global_w_offset + pose_w_len_full - 1, steps=pose_w_len_full, device=device, dtype=dtype).reshape(1, 1, -1) @@ -2515,6 +2519,21 @@ class WanModel(torch.nn.Module): List[Tensor]: List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8] """ + # Fast SCAIL-2 path: local SCAIL-2 CFG workflow without unrelated control branches. + if kwargs.get("simple_scail2", False): + return self._forward_scail2_fast( + x, t, context, seq_len, freqs, scail_input, + is_uncond=is_uncond, + current_step_percentage=current_step_percentage, + current_step=current_step, + last_step=last_step, + clip_fea=clip_fea, + control_lora_enabled=control_lora_enabled, + transformer_options=transformer_options, + ntk_alphas=ntk_alphas, + context_window_start=kwargs.get("context_window_start", 0), + ) + # Fast T2V path: bypass all extra checks for simple text-to-video inference if kwargs.get("simple_t2v", False): return self._forward_t2v(x, t, context, seq_len, freqs, @@ -2652,11 +2671,12 @@ class WanModel(torch.nn.Module): if ref_latent is not None and scail_input['ref_start_percent'] <= current_step_percentage <= scail_input['ref_end_percent']: scail_ref_mask_latents = scail_input.get("ref_mask_latents", None) scail_ref_frame_shape = ref_latent.shape + scail_ref_frames = ref_latent.shape[1] x = [torch.cat([v, u], dim=1) for v, u in zip([ref_latent], x)] seq_len += math.ceil((ref_latent.shape[-1] * ref_latent.shape[-2]) / 4 * ref_latent.shape[-3]) - F += 1 - prefix_frames = 1 - suffix_frames += 1 + F += scail_ref_frames + prefix_frames += scail_ref_frames + suffix_frames += scail_ref_frames #uni3c controlnet if uni3c_data is not None: @@ -3621,6 +3641,249 @@ class WanModel(torch.nn.Module): x = [u[:, prefix_frames:suffix_frames, ...].float() for u in x] return (x, x_ovi, pred_id) if pred_id is not None else (x, x_ovi, None) + def _forward_scail2_fast(self, x, t, context, seq_len, freqs, scail_input, + is_uncond=False, current_step_percentage=0.0, current_step=0, last_step=False, + clip_fea=None, control_lora_enabled=False, transformer_options={}, + ntk_alphas=None, context_window_start=0): + """Streamlined SCAIL-2 forward for the local CFG sampler path.""" + device = self.main_device + ntk = ntk_alphas if ntk_alphas is not None else [1.0, 1.0, 1.0] + + if self.lora_scheduling_enabled: + update_lora_step(self, current_step) + + if freqs is not None and freqs.device != device: + freqs = freqs.to(device) + + _, F, H, W = x[0].shape + suffix_frames = x[0].shape[1] + prefix_frames = 0 + pose_frame_shape = None + scail_ref_frame_shape = None + scail_ref_mask_latents = None + scail_ref_mask_flag = None + + if scail_input is not None: + scail_ref_mask_flag = scail_input.get("ref_mask_flag", None) + ref_latent = scail_input.get("ref_latent_pos", None) if not is_uncond else scail_input.get("ref_latent_neg", None) + if ref_latent is not None and scail_input["ref_start_percent"] <= current_step_percentage <= scail_input["ref_end_percent"]: + scail_ref_mask_latents = scail_input.get("ref_mask_latents", None) + scail_ref_frame_shape = ref_latent.shape + scail_ref_frames = ref_latent.shape[1] + x = [torch.cat([v, u], dim=1) for v, u in zip([ref_latent], x)] + seq_len += math.ceil((ref_latent.shape[-1] * ref_latent.shape[-2]) / 4 * ref_latent.shape[-3]) + F += scail_ref_frames + prefix_frames += scail_ref_frames + suffix_frames += scail_ref_frames + + patch_embedding = self.expanded_patch_embedding if control_lora_enabled else self.original_patch_embedding + _module_to_if_needed(patch_embedding, device) + x = [patch_embedding(u.unsqueeze(0).to(torch.float32)).to(x[0].dtype) for u in x] + + if scail_ref_mask_latents is not None: + if not hasattr(self, "patch_embedding_mask"): + raise ValueError("SCAIL-2 mask latents were provided, but the model has no patch_embedding_mask module") + _module_to_if_needed(self.patch_embedding_mask, device) + scail_mask_x = self.patch_embedding_mask(scail_ref_mask_latents.unsqueeze(0).to(torch.float32)).to(x[0].dtype) + if scail_mask_x.shape[2:] != x[0].shape[2:]: + raise ValueError(f"SCAIL-2 ref_mask_latents shape {scail_mask_x.shape[2:]} does not match latent token grid {x[0].shape[2:]}") + x = [u + scail_mask_x for u in x] + + grid_sizes = torch.stack([torch.tensor(u.shape[2:], device=device, dtype=torch.long) for u in x]) + original_grid_sizes = grid_sizes.clone() + x = [u.flatten(2).transpose(1, 2) for u in x] + self.original_seq_len = x[0].shape[1] + + if scail_input is not None: + scail_pose_latents = scail_input.get("pose_latent", None) + if scail_pose_latents is not None and scail_input["pose_start_percent"] <= current_step_percentage <= scail_input["pose_end_percent"]: + _module_to_if_needed(self.patch_embedding_pose, device) + scail_x = self.patch_embedding_pose(scail_pose_latents.unsqueeze(0).to(torch.float32)).to(x[0].dtype) + scail_sam_latents = scail_input.get("sam_latents", None) + if scail_sam_latents is not None: + if not hasattr(self, "patch_embedding_mask"): + raise ValueError("SCAIL-2 sam_latents were provided, but the model has no patch_embedding_mask module") + _module_to_if_needed(self.patch_embedding_mask, device) + scail_mask_x = self.patch_embedding_mask(scail_sam_latents.unsqueeze(0).to(torch.float32)).to(scail_x.dtype) + if scail_mask_x.shape[2:] != scail_x.shape[2:]: + raise ValueError(f"SCAIL-2 sam_latents shape {scail_mask_x.shape[2:]} does not match pose token grid {scail_x.shape[2:]}") + scail_x = scail_x + scail_mask_x + scail_x = scail_x.flatten(2).transpose(1, 2) + x = [torch.cat([u, scail_x], dim=1) for u in x] + seq_len += scail_x.shape[1] + pose_frame_shape = scail_pose_latents.shape + + seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.int32) + assert seq_lens.max() <= seq_len, f"max seq len {seq_lens.max()} exceeds provided seq_len {seq_len}" + x = torch.cat([torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1) for u in x]) + x = x.to(self.base_dtype) + + if freqs is None and "comfy" in self.rope_func: + cache_key = ( + "scail2_fast", + F, H, W, + tuple(pose_frame_shape) if pose_frame_shape is not None else None, + tuple(scail_ref_frame_shape) if scail_ref_frame_shape is not None else None, + None if scail_ref_mask_flag is None else bool(scail_ref_mask_flag), + self.rope_embedder.k, + tuple(ntk), + context_window_start, + ) + if self.cached_freqs is not None and hasattr(self, "cached_key") and self.cached_key == cache_key: + freqs = self.cached_freqs + else: + freqs = self.rope_encode_comfy( + F, H, W, + ntk_alphas=ntk, + pose_frame_shape=pose_frame_shape, + scail_ref_frame_shape=scail_ref_frame_shape, + ref_mask_flag=scail_ref_mask_flag, + context_window_start=context_window_start, + device=device, + dtype=x.dtype, + ) + self.cached_freqs = freqs + self.cached_key = cache_key + if freqs is None: + raise ValueError("SCAIL-2 fast path requires Comfy RoPE or precomputed freqs") + if freqs.device != device: + freqs = freqs.to(device) + + _module_to_if_needed(self.time_embedding, device) + time_embed_dtype = self.time_embedding[0].weight.dtype + if time_embed_dtype not in [torch.float16, torch.bfloat16, torch.float32]: + time_embed_dtype = self.base_dtype + time_projection_enabled = hasattr(self, "time_projection") + if time_projection_enabled: + _module_to_if_needed(self.time_projection, device) + cached_time = getattr(self, "_scail2_time_embed_cache", None) + if (cached_time is not None and cached_time[0] is t and cached_time[1] == device + and cached_time[2] == time_embed_dtype and cached_time[3] == time_projection_enabled): + e_raw, e = cached_time[4], cached_time[5] + else: + e_raw = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, t.flatten()).to(device=device, dtype=time_embed_dtype)) + if time_projection_enabled: + e = self.time_projection(e_raw).unflatten(1, (6, self.dim)) + else: + e = e_raw.reshape(1, 6, self.dim) + if e.dim() < 3 or e.shape[1] != 6: + e = e.reshape(1, 6, self.dim) + self._scail2_time_embed_cache = (t, device, time_embed_dtype, time_projection_enabled, e_raw, e) + + clip_embed = None + if clip_fea is not None and hasattr(self, "img_emb"): + _module_to_if_needed(self.img_emb, device) + clip_embed = self.img_emb(clip_fea.to(device)) + if self.offload_img_emb: + self.img_emb.to(self.offload_device, non_blocking=self.use_non_blocking) + + if hasattr(self, "text_embedding") and isinstance(context, list) and context != []: + text_embed_dtype = self.text_embedding[0].weight.dtype + if text_embed_dtype not in [torch.float16, torch.bfloat16, torch.float32]: + text_embed_dtype = self.base_dtype + text_cache_key = _t2v_context_cache_key(context, text_embed_dtype, device, self.text_len) + text_cache = getattr(self, "_scail2_text_embed_cache", None) + if text_cache is None: + text_cache = {} + self._scail2_text_embed_cache = text_cache + cached_context = text_cache.get(text_cache_key) + if cached_context is None: + _module_to_if_needed(self.text_embedding, device) + padded_context = torch.stack([torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) for u in context]).to(text_embed_dtype) + cached_context = self.text_embedding(padded_context.to(device)) + if len(text_cache) >= 4: + text_cache.clear() + text_cache[text_cache_key] = cached_context + if self.offload_txt_emb: + self.text_embedding.to(self.offload_device, non_blocking=self.use_non_blocking) + context = cached_context + elif isinstance(context, list): + context = None + + block_kwargs = dict( + e=e, + seq_lens=seq_lens, + grid_sizes=grid_sizes, + freqs=freqs, + context=context, + clip_embed=clip_embed, + current_step=torch.tensor(current_step), + last_step=torch.tensor(last_step, dtype=torch.bool), + chunked_self_attention=False, + seq_chunks=1, + num_latent_frames=F, + frame_tokens=x.shape[1] // F, + original_seq_len=self.original_seq_len, + transformer_options=transformer_options, + ) + + if torch.cuda.is_available(): + cuda_stream = None + events = [torch.cuda.Event() for _ in self.blocks] + swap_start_idx = len(self.blocks) - self.blocks_to_swap if self.blocks_to_swap > 0 else len(self.blocks) + else: + cuda_stream = None + events = None + swap_start_idx = len(self.blocks) + + attn_override_blocks = attention_mode = None + attention_mode_override_active = False + attention_mode_override = transformer_options.get("attention_mode_override", None) + if attention_mode_override is not None: + attn_override_blocks = attention_mode_override.get("blocks", range(len(self.blocks))) + if attention_mode_override["start_step"] <= current_step < attention_mode_override["end_step"]: + attention_mode_override_active = True + if attention_mode_override["verbose"]: + tqdm.write(f"Applying attention mode override: {attention_mode_override['mode']} at step {current_step} on blocks: {attn_override_blocks if attn_override_blocks is not None else 'all'}") + + for b, block in enumerate(self.blocks): + mm.throw_exception_if_processing_interrupted() + if attention_mode_override_active and b in attn_override_blocks: + attention_mode = attention_mode_override["mode"] + else: + attention_mode = None + + if self.prefetch_blocks > 0: + for prefetch_offset in range(1, self.prefetch_blocks + 1): + prefetch_idx = b + prefetch_offset + if prefetch_idx < len(self.blocks) and self.blocks_to_swap > 0 and prefetch_idx >= swap_start_idx: + context_mgr = torch.cuda.stream(cuda_stream) if torch.cuda.is_available() and cuda_stream is not None else nullcontext() + with context_mgr: + self.blocks[prefetch_idx].to(device, non_blocking=self.use_non_blocking) + if events is not None: + events[prefetch_idx].record(cuda_stream) + + if self.block_swap_debug: + transfer_start = time.perf_counter() + if b >= swap_start_idx and self.blocks_to_swap > 0: + if self.prefetch_blocks > 0 and events is not None: + if not events[b].query(): + events[b].synchronize() + block.to(device) + if self.block_swap_debug: + transfer_end = time.perf_counter() + transfer_time = transfer_end - transfer_start + compute_start = time.perf_counter() + + x, _, _, _ = block(x, attention_mode_override=attention_mode, **block_kwargs) + + if self.block_swap_debug: + compute_end = time.perf_counter() + compute_time = compute_end - compute_start + to_cpu_transfer_start = time.perf_counter() + if b >= swap_start_idx and self.blocks_to_swap > 0: + block.to(self.offload_device, non_blocking=self.use_non_blocking) + if self.block_swap_debug: + to_cpu_transfer_end = time.perf_counter() + to_cpu_transfer_time = to_cpu_transfer_end - to_cpu_transfer_start + log.info(f"Block {b}: transfer_time={transfer_time:.4f}s, compute_time={compute_time:.4f}s, to_cpu_transfer_time={to_cpu_transfer_time:.4f}s") + + x = x[:, :self.original_seq_len] + x = self.head(x, e_raw.to(device), temp_length=F) + x = self.unpatchify(x, original_grid_sizes) + x = [u[:, prefix_frames:suffix_frames, ...].float() for u in x] + return x, None, None + def _forward_t2v(self, x, t, context, seq_len, freqs, context_latents=None, context_frame_shapes=None, context_window_start=0, ntk_alphas=None): """Streamlined T2V forward — no block swap, no extra checks, ~15x faster.""" @@ -3675,6 +3938,8 @@ class WanModel(torch.nn.Module): # Time embedding (must produce [batch, 6, dim] for block.get_mod compatibility) _module_to_if_needed(self.time_embedding, device) time_embed_dtype = self.time_embedding[0].weight.dtype + if time_embed_dtype not in [torch.float16, torch.bfloat16, torch.float32]: + time_embed_dtype = self.base_dtype time_projection_enabled = hasattr(self, 'time_projection') if time_projection_enabled: _module_to_if_needed(self.time_projection, device) @@ -3683,7 +3948,7 @@ class WanModel(torch.nn.Module): and cached_time[2] == time_embed_dtype and cached_time[3] == time_projection_enabled): e_raw, e = cached_time[4], cached_time[5] else: - e_raw = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, t.flatten()).to(time_embed_dtype)) + e_raw = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, t.flatten()).to(device=device, dtype=time_embed_dtype)) if time_projection_enabled: e = self.time_projection(e_raw).unflatten(1, (6, self.dim)) # [batch, 6, dim] else: