From 663ea33ff113aaf353757d897beafa7d2259b4db Mon Sep 17 00:00:00 2001 From: William Lin Date: Tue, 26 Aug 2025 22:05:40 -0700 Subject: [PATCH 01/13] [bugfix] Fix wrong HF model string for FastWan2.2 5B (#763) --- fastvideo/configs/sample/registry.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/fastvideo/configs/sample/registry.py b/fastvideo/configs/sample/registry.py index 740cbbde9..cb646bff7 100644 --- a/fastvideo/configs/sample/registry.py +++ b/fastvideo/configs/sample/registry.py @@ -35,7 +35,8 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = { "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VSamplingParam, "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": FastWanT2V480PConfig, "Wan-AI/Wan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam, - "FastVideo/FastWan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam, + "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers": + Wan2_2_TI2V_5B_SamplingParam, "Wan-AI/Wan2.2-T2V-A14B-Diffusers": Wan2_2_T2V_A14B_SamplingParam, "Wan-AI/Wan2.2-I2V-A14B-Diffusers": Wan2_2_I2V_A14B_SamplingParam, # Add other specific weight variants From 7c554e5da82c3ba9356281d10318d9826eed1208 Mon Sep 17 00:00:00 2001 From: Zhang Peiyuan Date: Wed, 27 Aug 2025 16:12:47 -0700 Subject: [PATCH 02/13] Update Community Link (#765) --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index a61ab260c..43b72cb63 100644 --- a/README.md +++ b/README.md @@ -7,7 +7,7 @@ FastVideo features an end-to-end unified pipeline for accelerating diffusion models, starting from data preprocessing to model training, finetuning, distillation, and inference. FastVideo is designed to be modular and extensible, allowing users to easily add new optimizations and techniques. Whether it is training-free optimizations or post-training optimizations, FastVideo has you covered.

- | 🕹ī¸ Online Demo | Documentation | Quick Start | 🤗 FastWan | đŸŸŖđŸ’Ŧ Slack | đŸŸŖđŸ’Ŧ WeChat | + | 🕹ī¸ Online Demo | Documentation | Quick Start | 🤗 FastWan | đŸŸŖđŸ’Ŧ Slack | đŸŸŖđŸ’Ŧ WeChat |

From 36371c568968a4a67e8bd2d108c191078cfbe997 Mon Sep 17 00:00:00 2001 From: JerryZhou54 Date: Tue, 26 Aug 2025 19:49:57 +0000 Subject: [PATCH 03/13] checkpoint --- .../configs/models/dits/causal_wanvideo.py | 631 ++++++++++++++++++ fastvideo/configs/models/dits/wanvideo.py | 4 + 2 files changed, 635 insertions(+) create mode 100644 fastvideo/configs/models/dits/causal_wanvideo.py diff --git a/fastvideo/configs/models/dits/causal_wanvideo.py b/fastvideo/configs/models/dits/causal_wanvideo.py new file mode 100644 index 000000000..405962448 --- /dev/null +++ b/fastvideo/configs/models/dits/causal_wanvideo.py @@ -0,0 +1,631 @@ +# SPDX-License-Identifier: Apache-2.0 + +import math +from typing import Any + +import numpy as np +import torch +import torch.nn as nn + +from torch.nn.attention.flex_attention import create_block_mask, flex_attention +from torch.nn.attention.flex_attention import BlockMask +# wan 1.3B model has a weird channel / head configurations and require max-autotune to work with flexattention +# see https://github.com/pytorch/pytorch/issues/133254 +# change to default for other models +flex_attention = torch.compile( + flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs") +import torch.distributed as dist + +import fastvideo.envs as envs +from fastvideo.attention import (DistributedAttention, + LocalAttention) +from fastvideo.configs.models.dits import WanVideoConfig +from fastvideo.distributed.parallel_state import get_sp_world_size +from fastvideo.forward_context import get_forward_context +from fastvideo.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift, + RMSNorm, ScaleResidual, + ScaleResidualLayerNormScaleShift) +from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.layers.mlp import MLP +from fastvideo.layers.rotary_embedding import (_apply_rotary_emb, + get_rotary_pos_embed) +from fastvideo.layers.visual_embedding import (ModulateProjection, PatchEmbed, + TimestepEmbedder) +from fastvideo.logger import init_logger +from fastvideo.models.dits.base import BaseDiT +from fastvideo.platforms import AttentionBackendEnum, current_platform + +logger = init_logger(__name__) + + +class WanImageEmbedding(torch.nn.Module): + + def __init__(self, in_features: int, out_features: int): + super().__init__() + + self.norm1 = FP32LayerNorm(in_features) + self.ff = MLP(in_features, in_features, out_features, act_type="gelu") + self.norm2 = FP32LayerNorm(out_features) + + def forward(self, + encoder_hidden_states_image: torch.Tensor) -> torch.Tensor: + dtype = encoder_hidden_states_image.dtype + hidden_states = self.norm1(encoder_hidden_states_image) + hidden_states = self.ff(hidden_states) + hidden_states = self.norm2(hidden_states).to(dtype) + return hidden_states + + +class WanTimeTextImageEmbedding(nn.Module): + + def __init__( + self, + dim: int, + time_freq_dim: int, + text_embed_dim: int, + image_embed_dim: int | None = None, + ): + super().__init__() + + self.time_embedder = TimestepEmbedder( + dim, frequency_embedding_size=time_freq_dim, act_layer="silu") + self.time_modulation = ModulateProjection(dim, + factor=6, + act_layer="silu") + self.text_embedder = MLP(text_embed_dim, + dim, + dim, + bias=True, + act_type="gelu_pytorch_tanh") + + self.image_embedder = None + if image_embed_dim is not None: + self.image_embedder = WanImageEmbedding(image_embed_dim, dim) + + def forward( + self, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_hidden_states_image: torch.Tensor | None = None, + ): + temb = self.time_embedder(timestep) + timestep_proj = self.time_modulation(temb) + + encoder_hidden_states = self.text_embedder(encoder_hidden_states) + if encoder_hidden_states_image is not None: + assert self.image_embedder is not None + encoder_hidden_states_image = self.image_embedder( + encoder_hidden_states_image) + + return temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image + + +class WanSelfAttention(nn.Module): + + def __init__(self, + dim: int, + num_heads: int, + window_size=(-1, -1), + qk_norm=True, + eps=1e-6, + parallel_attention=False) -> None: + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.window_size = window_size + self.qk_norm = qk_norm + self.eps = eps + self.parallel_attention = parallel_attention + + # layers + self.to_q = ReplicatedLinear(dim, dim) + self.to_k = ReplicatedLinear(dim, dim) + self.to_v = ReplicatedLinear(dim, dim) + self.to_out = ReplicatedLinear(dim, dim) + self.norm_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + + # Scaled dot product attention + self.attn = LocalAttention( + num_heads=num_heads, + head_size=self.head_dim, + dropout_rate=0, + softmax_scale=None, + causal=False, + supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA)) + + def forward(self, x: torch.Tensor, context: torch.Tensor, + context_lens: int): + r""" + Args: + x(Tensor): Shape [B, L, num_heads, C / num_heads] + seq_lens(Tensor): Shape [B] + grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) + freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] + """ + pass + + +class WanT2VCrossAttention(WanSelfAttention): + + def forward(self, x, context, context_lens): + r""" + Args: + x(Tensor): Shape [B, L1, C] + context(Tensor): Shape [B, L2, C] + context_lens(Tensor): Shape [B] + """ + b, n, d = x.size(0), self.num_heads, self.head_dim + + # compute query, key, value + q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d) + k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d) + v = self.to_v(context)[0].view(b, -1, n, d) + + # compute attention + x = self.attn(q, k, v) + + # output + x = x.flatten(2) + x, _ = self.to_out(x) + return x + +class CausalWanTransformerBlock(nn.Module): + + def __init__(self, + dim: int, + ffn_dim: int, + num_heads: int, + qk_norm: str = "rms_norm_across_heads", + cross_attn_norm: bool = False, + eps: float = 1e-6, + added_kv_proj_dim: int | None = None, + supported_attention_backends: tuple[AttentionBackendEnum, ...] + | None = None, + prefix: str = ""): + super().__init__() + + # 1. Self-attention + self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.to_q = ReplicatedLinear(dim, dim, bias=True) + self.to_k = ReplicatedLinear(dim, dim, bias=True) + self.to_v = ReplicatedLinear(dim, dim, bias=True) + + self.to_out = ReplicatedLinear(dim, dim, bias=True) + self.attn1 = DistributedAttention( + num_heads=num_heads, + head_size=dim // num_heads, + causal=False, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.attn1") + self.hidden_dim = dim + self.num_attention_heads = num_heads + dim_head = dim // num_heads + if qk_norm == "rms_norm": + self.norm_q = RMSNorm(dim_head, eps=eps) + self.norm_k = RMSNorm(dim_head, eps=eps) + elif qk_norm == "rms_norm_across_heads": + # LTX applies qk norm across all heads + self.norm_q = RMSNorm(dim, eps=eps) + self.norm_k = RMSNorm(dim, eps=eps) + else: + print("QK Norm type not supported") + raise Exception + assert cross_attn_norm is True + self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift( + dim, + norm_type="layer", + eps=eps, + elementwise_affine=True, + dtype=torch.float32, + compute_dtype=torch.float32) + + # 2. Cross-attention + if added_kv_proj_dim is not None: + # I2V + self.attn2 = WanI2VCrossAttention(dim, + num_heads, + qk_norm=qk_norm, + eps=eps) + else: + # T2V + self.attn2 = WanT2VCrossAttention(dim, + num_heads, + qk_norm=qk_norm, + eps=eps) + self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift( + dim, + norm_type="layer", + eps=eps, + elementwise_affine=False, + dtype=torch.float32, + compute_dtype=torch.float32) + + # 3. Feed-forward + self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh") + self.mlp_residual = ScaleResidual() + + self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + freqs_cis: tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + if hidden_states.dim() == 4: + hidden_states = hidden_states.squeeze(1) + bs, seq_length, _ = hidden_states.shape + orig_dtype = hidden_states.dtype + # assert orig_dtype != torch.float32 + e = self.scale_shift_table + temb.float() + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk( + 6, dim=1) + assert shift_msa.dtype == torch.float32 + + # 1. Self-attention + norm_hidden_states = (self.norm1(hidden_states.float()) * + (1 + scale_msa) + shift_msa).to(orig_dtype) + query, _ = self.to_q(norm_hidden_states) + key, _ = self.to_k(norm_hidden_states) + value, _ = self.to_v(norm_hidden_states) + + if self.norm_q is not None: + query = self.norm_q(query) + if self.norm_k is not None: + key = self.norm_k(key) + + query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) + key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) + value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) + + # Apply rotary embeddings + cos, sin = freqs_cis + query, key = _apply_rotary_emb(query, cos, sin, + is_neox_style=False), _apply_rotary_emb( + key, cos, sin, is_neox_style=False) + + attn_output, _ = self.attn1(query, key, value) + attn_output = attn_output.flatten(2) + attn_output, _ = self.to_out(attn_output) + attn_output = attn_output.squeeze(1) + + null_shift = null_scale = torch.tensor([0], device=hidden_states.device) + norm_hidden_states, hidden_states = self.self_attn_residual_norm( + hidden_states, attn_output, gate_msa, null_shift, null_scale) + norm_hidden_states, hidden_states = norm_hidden_states.to( + orig_dtype), hidden_states.to(orig_dtype) + + # 2. Cross-attention + attn_output = self.attn2(norm_hidden_states, + context=encoder_hidden_states, + context_lens=None) + norm_hidden_states, hidden_states = self.cross_attn_residual_norm( + hidden_states, attn_output, 1, c_shift_msa, c_scale_msa) + norm_hidden_states, hidden_states = norm_hidden_states.to( + orig_dtype), hidden_states.to(orig_dtype) + + # 3. Feed-forward + ff_output = self.ffn(norm_hidden_states) + hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa) + hidden_states = hidden_states.to(orig_dtype) + + return hidden_states + +class CausalWanTransformer3DModel(BaseDiT): + _fsdp_shard_conditions = WanVideoConfig()._fsdp_shard_conditions + _compile_conditions = WanVideoConfig()._compile_conditions + _supported_attention_backends = WanVideoConfig( + )._supported_attention_backends + param_names_mapping = WanVideoConfig().param_names_mapping + reverse_param_names_mapping = WanVideoConfig().reverse_param_names_mapping + lora_param_names_mapping = WanVideoConfig().lora_param_names_mapping + + def __init__(self, config: WanVideoConfig, hf_config: dict[str, + Any]) -> None: + super().__init__(config=config, hf_config=hf_config) + + inner_dim = config.num_attention_heads * config.attention_head_dim + self.hidden_size = config.hidden_size + self.num_attention_heads = config.num_attention_heads + self.in_channels = config.in_channels + self.out_channels = config.out_channels + self.num_channels_latents = config.num_channels_latents + self.patch_size = config.patch_size + self.text_len = config.text_len + self.local_attn_size = config.local_attn_size + + # 1. Patch & position embedding + self.patch_embedding = PatchEmbed(in_chans=config.in_channels, + embed_dim=inner_dim, + patch_size=config.patch_size, + flatten=False) + + # 2. Condition embeddings + self.condition_embedder = WanTimeTextImageEmbedding( + dim=inner_dim, + time_freq_dim=config.freq_dim, + text_embed_dim=config.text_dim, + image_embed_dim=config.image_dim, + ) + + # 3. Transformer blocks + self.blocks = nn.ModuleList([ + CausalWanTransformerBlock(inner_dim, + config.ffn_dim, + config.num_attention_heads, + config.qk_norm, + config.cross_attn_norm, + config.eps, + config.added_kv_proj_dim, + self._supported_attention_backends, + prefix=f"{config.prefix}.blocks.{i}") + for i in range(config.num_layers) + ]) + + # 4. Output norm & projection + self.norm_out = LayerNormScaleShift(inner_dim, + norm_type="layer", + eps=config.eps, + elementwise_affine=False, + dtype=torch.float32, + compute_dtype=torch.float32) + self.proj_out = nn.Linear( + inner_dim, config.out_channels * math.prod(config.patch_size)) + self.scale_shift_table = nn.Parameter( + torch.randn(1, 2, inner_dim) / inner_dim**0.5) + + self.gradient_checkpointing = False + + # Causal-specific + self.block_mask = None + self.num_frame_per_block = 1 + self.independent_first_frame = False + + self.__post_init__() + + @staticmethod + def _prepare_blockwise_causal_attn_mask( + device: torch.device | str, num_frames: int = 21, + frame_seqlen: int = 1560, num_frame_per_block=1, local_attn_size=-1 + ) -> BlockMask: + """ + we will divide the token sequence into the following format + [1 latent frame] [1 latent frame] ... [1 latent frame] + We use flexattention to construct the attention mask + """ + total_length = num_frames * frame_seqlen + + # we do right padding to get to a multiple of 128 + padded_length = math.ceil(total_length / 128) * 128 - total_length + + ends = torch.zeros(total_length + padded_length, + device=device, dtype=torch.long) + + # Block-wise causal mask will attend to all elements that are before the end of the current chunk + frame_indices = torch.arange( + start=0, + end=total_length, + step=frame_seqlen * num_frame_per_block, + device=device + ) + + for tmp in frame_indices: + ends[tmp:tmp + frame_seqlen * num_frame_per_block] = tmp + \ + frame_seqlen * num_frame_per_block + + def attention_mask(b, h, q_idx, kv_idx): + if local_attn_size == -1: + return (kv_idx < ends[q_idx]) | (q_idx == kv_idx) + else: + return ((kv_idx < ends[q_idx]) & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))) | (q_idx == kv_idx) + # return ((kv_idx < total_length) & (q_idx < total_length)) | (q_idx == kv_idx) # bidirectional mask + + block_mask = create_block_mask(attention_mask, B=None, H=None, Q_LEN=total_length + padded_length, + KV_LEN=total_length + padded_length, _compile=False, device=device) + + if not dist.is_initialized() or dist.get_rank() == 0: + print( + f" cache a block wise causal mask with block size of {num_frame_per_block} frames") + print(block_mask) + + # import imageio + # import numpy as np + # from torch.nn.attention.flex_attention import create_mask + + # mask = create_mask(attention_mask, B=None, H=None, Q_LEN=total_length + + # padded_length, KV_LEN=total_length + padded_length, device=device) + # import cv2 + # mask = cv2.resize(mask[0, 0].cpu().float().numpy(), (1024, 1024)) + # imageio.imwrite("mask_%d.jpg" % (0), np.uint8(255. * mask)) + + return block_mask + + def _forward_inference( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + timestep: torch.LongTensor, + encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] + | None = None, + kv_cache: dict = None, + crossattn_cache: dict = None, + current_start: int = 0, + cache_start: int = 0, + **kwargs) -> torch.Tensor: + r""" + Run the diffusion model with kv caching. + See Algorithm 2 of CausVid paper https://arxiv.org/abs/2412.07772 for details. + This function will be run for num_frame times. + Process the latent frames one by one (1560 tokens each) + """ + + orig_dtype = hidden_states.dtype + if not isinstance(encoder_hidden_states, torch.Tensor): + encoder_hidden_states = encoder_hidden_states[0] + if isinstance(encoder_hidden_states_image, + list) and len(encoder_hidden_states_image) > 0: + encoder_hidden_states_image = encoder_hidden_states_image[0] + else: + encoder_hidden_states_image = None + + batch_size, num_channels, num_frames, height, width = hidden_states.shape + p_t, p_h, p_w = self.patch_size + post_patch_num_frames = num_frames // p_t + post_patch_height = height // p_h + post_patch_width = width // p_w + + # Get rotary embeddings + d = self.hidden_size // self.num_attention_heads + rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] + freqs_cos, freqs_sin = get_rotary_pos_embed( + (post_patch_num_frames * get_sp_world_size(), post_patch_height, + post_patch_width), + self.hidden_size, + self.num_attention_heads, + rope_dim_list, + dtype=torch.float32 if current_platform.is_mps() else torch.float64, + rope_theta=10000) + freqs_cos = freqs_cos.to(hidden_states.device) + freqs_sin = freqs_sin.to(hidden_states.device) + freqs_cis = (freqs_cos.float(), + freqs_sin.float()) if freqs_cos is not None else None + + hidden_states = self.patch_embedding(hidden_states) + hidden_states = hidden_states.flatten(2).transpose(1, 2) + + temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( + timestep, encoder_hidden_states, encoder_hidden_states_image) + timestep_proj = timestep_proj.unflatten(1, (6, -1)) + + if encoder_hidden_states_image is not None: + encoder_hidden_states = torch.concat( + [encoder_hidden_states_image, encoder_hidden_states], dim=1) + + encoder_hidden_states = encoder_hidden_states.to( + orig_dtype) if current_platform.is_mps( + ) else encoder_hidden_states # cast to orig_dtype for MPS + + assert encoder_hidden_states.dtype == orig_dtype + + # 4. Transformer blocks + if torch.is_grad_enabled() and self.gradient_checkpointing: + for block in self.blocks: + hidden_states = self._gradient_checkpointing_func( + block, hidden_states, encoder_hidden_states, + timestep_proj, freqs_cis) + else: + for block in self.blocks: + hidden_states = block(hidden_states, encoder_hidden_states, + timestep_proj, freqs_cis) + + # 5. Output norm, projection & unpatchify + shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, + dim=1) + hidden_states = self.norm_out(hidden_states, shift, scale) + hidden_states = self.proj_out(hidden_states) + + hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, + post_patch_height, + post_patch_width, p_t, p_h, p_w, + -1) + hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) + output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) + + return output + + def _forward_train(self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + timestep: torch.LongTensor, + encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] + | None = None, + **kwargs) -> torch.Tensor: + + orig_dtype = hidden_states.dtype + if not isinstance(encoder_hidden_states, torch.Tensor): + encoder_hidden_states = encoder_hidden_states[0] + if isinstance(encoder_hidden_states_image, + list) and len(encoder_hidden_states_image) > 0: + encoder_hidden_states_image = encoder_hidden_states_image[0] + else: + encoder_hidden_states_image = None + + batch_size, num_channels, num_frames, height, width = hidden_states.shape + p_t, p_h, p_w = self.patch_size + post_patch_num_frames = num_frames // p_t + post_patch_height = height // p_h + post_patch_width = width // p_w + + # Get rotary embeddings + d = self.hidden_size // self.num_attention_heads + rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] + freqs_cos, freqs_sin = get_rotary_pos_embed( + (post_patch_num_frames * get_sp_world_size(), post_patch_height, + post_patch_width), + self.hidden_size, + self.num_attention_heads, + rope_dim_list, + dtype=torch.float32 if current_platform.is_mps() else torch.float64, + rope_theta=10000) + freqs_cos = freqs_cos.to(hidden_states.device) + freqs_sin = freqs_sin.to(hidden_states.device) + freqs_cis = (freqs_cos.float(), + freqs_sin.float()) if freqs_cos is not None else None + + hidden_states = self.patch_embedding(hidden_states) + hidden_states = hidden_states.flatten(2).transpose(1, 2) + + temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( + timestep, encoder_hidden_states, encoder_hidden_states_image) + timestep_proj = timestep_proj.unflatten(1, (6, -1)) + + if encoder_hidden_states_image is not None: + encoder_hidden_states = torch.concat( + [encoder_hidden_states_image, encoder_hidden_states], dim=1) + + encoder_hidden_states = encoder_hidden_states.to( + orig_dtype) if current_platform.is_mps( + ) else encoder_hidden_states # cast to orig_dtype for MPS + + assert encoder_hidden_states.dtype == orig_dtype + + # 4. Transformer blocks + if torch.is_grad_enabled() and self.gradient_checkpointing: + for block in self.blocks: + hidden_states = self._gradient_checkpointing_func( + block, hidden_states, encoder_hidden_states, + timestep_proj, freqs_cis) + else: + for block in self.blocks: + hidden_states = block(hidden_states, encoder_hidden_states, + timestep_proj, freqs_cis) + + # 5. Output norm, projection & unpatchify + shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, + dim=1) + hidden_states = self.norm_out(hidden_states, shift, scale) + hidden_states = self.proj_out(hidden_states) + + hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, + post_patch_height, + post_patch_width, p_t, p_h, p_w, + -1) + hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) + output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) + + return output + + def forward( + self, + *args, + **kwargs + ): + if kwargs.get('kv_cache', None) is not None: + return self._forward_inference(*args, **kwargs) + else: + return self._forward_train(*args, **kwargs) diff --git a/fastvideo/configs/models/dits/wanvideo.py b/fastvideo/configs/models/dits/wanvideo.py index 8cd25cf6c..d31a3a607 100644 --- a/fastvideo/configs/models/dits/wanvideo.py +++ b/fastvideo/configs/models/dits/wanvideo.py @@ -92,6 +92,10 @@ class WanVideoArchConfig(DiTArchConfig): pos_embed_seq_len: int | None = None exclude_lora_layers: list[str] = field(default_factory=lambda: ["embedder"]) + # Causal Wan + local_attn_size: int = -1 # Window size for temporal local attention (-1 indicates global attention) + sink_size: int = 0 # Size of the attention sink, we keep the first `sink_size` frames unchanged when rolling the KV cache + def __post_init__(self): super().__post_init__() self.out_channels = self.out_channels or self.in_channels From 363cf0d173d24483ad946e8ec41b1c49a12f581e Mon Sep 17 00:00:00 2001 From: JerryZhou54 Date: Tue, 26 Aug 2025 23:31:46 +0000 Subject: [PATCH 04/13] First version of causalwan --- .../configs/models/dits/causal_wanvideo.py | 271 ++++++++++-------- fastvideo/layers/rotary_embedding.py | 9 + fastvideo/models/dits/wanvideo.py | 18 +- 3 files changed, 170 insertions(+), 128 deletions(-) diff --git a/fastvideo/configs/models/dits/causal_wanvideo.py b/fastvideo/configs/models/dits/causal_wanvideo.py index 405962448..244d3ad8c 100644 --- a/fastvideo/configs/models/dits/causal_wanvideo.py +++ b/fastvideo/configs/models/dits/causal_wanvideo.py @@ -29,83 +29,20 @@ from fastvideo.layers.linear import ReplicatedLinear from fastvideo.layers.mlp import MLP from fastvideo.layers.rotary_embedding import (_apply_rotary_emb, get_rotary_pos_embed) -from fastvideo.layers.visual_embedding import (ModulateProjection, PatchEmbed, - TimestepEmbedder) +from fastvideo.layers.visual_embedding import (PatchEmbed) from fastvideo.logger import init_logger from fastvideo.models.dits.base import BaseDiT +from fastvideo.models.dits.wanvideo import WanT2VCrossAttention, WanTimeTextImageEmbedding from fastvideo.platforms import AttentionBackendEnum, current_platform logger = init_logger(__name__) - - -class WanImageEmbedding(torch.nn.Module): - - def __init__(self, in_features: int, out_features: int): - super().__init__() - - self.norm1 = FP32LayerNorm(in_features) - self.ff = MLP(in_features, in_features, out_features, act_type="gelu") - self.norm2 = FP32LayerNorm(out_features) - - def forward(self, - encoder_hidden_states_image: torch.Tensor) -> torch.Tensor: - dtype = encoder_hidden_states_image.dtype - hidden_states = self.norm1(encoder_hidden_states_image) - hidden_states = self.ff(hidden_states) - hidden_states = self.norm2(hidden_states).to(dtype) - return hidden_states - - -class WanTimeTextImageEmbedding(nn.Module): - - def __init__( - self, - dim: int, - time_freq_dim: int, - text_embed_dim: int, - image_embed_dim: int | None = None, - ): - super().__init__() - - self.time_embedder = TimestepEmbedder( - dim, frequency_embedding_size=time_freq_dim, act_layer="silu") - self.time_modulation = ModulateProjection(dim, - factor=6, - act_layer="silu") - self.text_embedder = MLP(text_embed_dim, - dim, - dim, - bias=True, - act_type="gelu_pytorch_tanh") - - self.image_embedder = None - if image_embed_dim is not None: - self.image_embedder = WanImageEmbedding(image_embed_dim, dim) - - def forward( - self, - timestep: torch.Tensor, - encoder_hidden_states: torch.Tensor, - encoder_hidden_states_image: torch.Tensor | None = None, - ): - temb = self.time_embedder(timestep) - timestep_proj = self.time_modulation(temb) - - encoder_hidden_states = self.text_embedder(encoder_hidden_states) - if encoder_hidden_states_image is not None: - assert self.image_embedder is not None - encoder_hidden_states_image = self.image_embedder( - encoder_hidden_states_image) - - return temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image - - -class WanSelfAttention(nn.Module): +class CausalWanSelfAttention(nn.Module): def __init__(self, dim: int, num_heads: int, - window_size=(-1, -1), + local_attn_size: int = -1, + sink_size: int = 0, qk_norm=True, eps=1e-6, parallel_attention=False) -> None: @@ -114,10 +51,12 @@ class WanSelfAttention(nn.Module): self.dim = dim self.num_heads = num_heads self.head_dim = dim // num_heads - self.window_size = window_size + self.local_attn_size = local_attn_size + self.sink_size = sink_size self.qk_norm = qk_norm self.eps = eps self.parallel_attention = parallel_attention + self.max_attention_size = 32760 if local_attn_size == -1 else local_attn_size * 1560 # layers self.to_q = ReplicatedLinear(dim, dim) @@ -137,8 +76,15 @@ class WanSelfAttention(nn.Module): supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA)) - def forward(self, x: torch.Tensor, context: torch.Tensor, - context_lens: int): + def forward(self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + freqs_cis: tuple[torch.Tensor, torch.Tensor], + block_mask: BlockMask, + kv_cache: dict | None = None, + current_start: int = 0, + cache_start: int | None = None): r""" Args: x(Tensor): Shape [B, L, num_heads, C / num_heads] @@ -146,31 +92,82 @@ class WanSelfAttention(nn.Module): grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] """ - pass + if cache_start is None: + cache_start = current_start + cos, sin = freqs_cis + roped_query = _apply_rotary_emb(q, cos, sin, is_neox_style=False) + roped_key = _apply_rotary_emb(k, cos, sin, is_neox_style=False) -class WanT2VCrossAttention(WanSelfAttention): + if kv_cache is None: + # Padding for flex attention + padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1] + padded_roped_query = torch.cat( + [roped_query, + torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]], + device=q.device, dtype=v.dtype)], + dim=1 + ) - def forward(self, x, context, context_lens): - r""" - Args: - x(Tensor): Shape [B, L1, C] - context(Tensor): Shape [B, L2, C] - context_lens(Tensor): Shape [B] - """ - b, n, d = x.size(0), self.num_heads, self.head_dim + padded_roped_key = torch.cat( + [roped_key, torch.zeros([k.shape[0], padded_length, k.shape[2], k.shape[3]], + device=k.device, dtype=v.dtype)], + dim=1 + ) - # compute query, key, value - q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d) - k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d) - v = self.to_v(context)[0].view(b, -1, n, d) + padded_v = torch.cat( + [v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]], + device=v.device, dtype=v.dtype)], + dim=1 + ) - # compute attention - x = self.attn(q, k, v) + x = flex_attention( + query=padded_roped_query.transpose(2, 1), + key=padded_roped_key.transpose(2, 1), + value=padded_v.transpose(2, 1), + block_mask=block_mask + )[:, :, :-padded_length].transpose(2, 1) + else: + frame_seqlen = q.shape[1] + current_end = current_start + roped_query.shape[1] + sink_tokens = self.sink_size * frame_seqlen + # If we are using local attention and the current KV cache size is larger than the local attention size, we need to truncate the KV cache + kv_cache_size = kv_cache["k"].shape[1] + num_new_tokens = roped_query.shape[1] + if self.local_attn_size != -1 and (current_end > kv_cache["global_end_index"].item()) and ( + num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size): + # Calculate the number of new tokens added in this step + # Shift existing cache content left to discard oldest tokens + # Clone the source slice to avoid overlapping memory error + num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size + num_rolled_tokens = kv_cache["local_end_index"].item() - num_evicted_tokens - sink_tokens + kv_cache["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \ + kv_cache["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() + kv_cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \ + kv_cache["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() + # Insert the new keys/values at the end + local_end_index = kv_cache["local_end_index"].item() + current_end - \ + kv_cache["global_end_index"].item() - num_evicted_tokens + local_start_index = local_end_index - num_new_tokens + kv_cache["k"][:, local_start_index:local_end_index] = roped_key + kv_cache["v"][:, local_start_index:local_end_index] = v + else: + # Assign new keys/values directly up to current_end + local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item() + local_start_index = local_end_index - num_new_tokens + kv_cache["k"][:, local_start_index:local_end_index] = roped_key + kv_cache["v"][:, local_start_index:local_end_index] = v + x = self.attn( + roped_query, + kv_cache["k"][:, max(0, local_end_index - self.max_attention_size):local_end_index], + kv_cache["v"][:, max(0, local_end_index - self.max_attention_size):local_end_index] + ) + kv_cache["global_end_index"].fill_(current_end) + kv_cache["local_end_index"].fill_(local_end_index) # output x = x.flatten(2) - x, _ = self.to_out(x) + x = self.to_out(x) return x class CausalWanTransformerBlock(nn.Module): @@ -179,12 +176,13 @@ class CausalWanTransformerBlock(nn.Module): dim: int, ffn_dim: int, num_heads: int, + local_attn_size: int = -1, + sink_size: int = 0, qk_norm: str = "rms_norm_across_heads", cross_attn_norm: bool = False, eps: float = 1e-6, added_kv_proj_dim: int | None = None, - supported_attention_backends: tuple[AttentionBackendEnum, ...] - | None = None, + supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, prefix: str = ""): super().__init__() @@ -195,14 +193,16 @@ class CausalWanTransformerBlock(nn.Module): self.to_v = ReplicatedLinear(dim, dim, bias=True) self.to_out = ReplicatedLinear(dim, dim, bias=True) - self.attn1 = DistributedAttention( - num_heads=num_heads, - head_size=dim // num_heads, - causal=False, - supported_attention_backends=supported_attention_backends, - prefix=f"{prefix}.attn1") + self.attn1 = CausalWanSelfAttention( + dim, + num_heads, + local_attn_size=local_attn_size, + sink_size=sink_size, + qk_norm=qk_norm, + eps=eps) self.hidden_dim = dim self.num_attention_heads = num_heads + self.local_attn_size = local_attn_size dim_head = dim // num_heads if qk_norm == "rms_norm": self.norm_q = RMSNorm(dim_head, eps=eps) @@ -224,18 +224,11 @@ class CausalWanTransformerBlock(nn.Module): compute_dtype=torch.float32) # 2. Cross-attention - if added_kv_proj_dim is not None: - # I2V - self.attn2 = WanI2VCrossAttention(dim, - num_heads, - qk_norm=qk_norm, - eps=eps) - else: - # T2V - self.attn2 = WanT2VCrossAttention(dim, - num_heads, - qk_norm=qk_norm, - eps=eps) + # Only T2V for now + self.attn2 = WanT2VCrossAttention(dim, + num_heads, + qk_norm=qk_norm, + eps=eps) self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift( dim, norm_type="layer", @@ -256,6 +249,11 @@ class CausalWanTransformerBlock(nn.Module): encoder_hidden_states: torch.Tensor, temb: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor], + block_mask: BlockMask, + kv_cache: dict | None = None, + crossattn_cache: dict | None = None, + current_start: int = 0, + cache_start: int | None = None, ) -> torch.Tensor: if hidden_states.dim() == 4: hidden_states = hidden_states.squeeze(1) @@ -283,13 +281,7 @@ class CausalWanTransformerBlock(nn.Module): key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) - # Apply rotary embeddings - cos, sin = freqs_cis - query, key = _apply_rotary_emb(query, cos, sin, - is_neox_style=False), _apply_rotary_emb( - key, cos, sin, is_neox_style=False) - - attn_output, _ = self.attn1(query, key, value) + attn_output, _ = self.attn1(query, key, value, freqs_cis, block_mask, kv_cache, current_start, cache_start) attn_output = attn_output.flatten(2) attn_output, _ = self.to_out(attn_output) attn_output = attn_output.squeeze(1) @@ -303,7 +295,8 @@ class CausalWanTransformerBlock(nn.Module): # 2. Cross-attention attn_output = self.attn2(norm_hidden_states, context=encoder_hidden_states, - context_lens=None) + context_lens=None, + crossattn_cache=crossattn_cache) norm_hidden_states, hidden_states = self.cross_attn_residual_norm( hidden_states, attn_output, 1, c_shift_msa, c_scale_msa) norm_hidden_states, hidden_states = norm_hidden_states.to( @@ -358,6 +351,8 @@ class CausalWanTransformer3DModel(BaseDiT): CausalWanTransformerBlock(inner_dim, config.ffn_dim, config.num_attention_heads, + config.local_attn_size, + config.sink_size, config.qk_norm, config.cross_attn_norm, config.eps, @@ -489,7 +484,9 @@ class CausalWanTransformer3DModel(BaseDiT): self.num_attention_heads, rope_dim_list, dtype=torch.float32 if current_platform.is_mps() else torch.float64, - rope_theta=10000) + rope_theta=10000, + start_frame=current_start # Assume that current_start is 0 when kv_cache is None + ) freqs_cos = freqs_cos.to(hidden_states.device) freqs_sin = freqs_sin.to(hidden_states.device) freqs_cis = (freqs_cos.float(), @@ -513,15 +510,27 @@ class CausalWanTransformer3DModel(BaseDiT): assert encoder_hidden_states.dtype == orig_dtype # 4. Transformer blocks - if torch.is_grad_enabled() and self.gradient_checkpointing: - for block in self.blocks: + for block_index, block in enumerate(self.blocks): + if torch.is_grad_enabled() and self.gradient_checkpointing: + causal_kwargs = { + "kv_cache": kv_cache[block_index], + "current_start": current_start, + "cache_start": cache_start + } hidden_states = self._gradient_checkpointing_func( block, hidden_states, encoder_hidden_states, - timestep_proj, freqs_cis) - else: - for block in self.blocks: + timestep_proj, freqs_cis, + **causal_kwargs) + else: + causal_kwargs = { + "kv_cache": kv_cache[block_index], + "crossattn_cache": crossattn_cache[block_index], + "current_start": current_start, + "cache_start": cache_start + } hidden_states = block(hidden_states, encoder_hidden_states, - timestep_proj, freqs_cis) + timestep_proj, freqs_cis, + **causal_kwargs) # 5. Output norm, projection & unpatchify shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, @@ -577,6 +586,16 @@ class CausalWanTransformer3DModel(BaseDiT): freqs_cis = (freqs_cos.float(), freqs_sin.float()) if freqs_cos is not None else None + # Construct blockwise causal attn mask + if self.block_mask is None: + self.block_mask = self._prepare_blockwise_causal_attn_mask( + device=hidden_states.device, + num_frames=num_frames, + frame_seqlen=post_patch_height * post_patch_width, + num_frame_per_block=self.num_frame_per_block, + local_attn_size=self.local_attn_size + ) + hidden_states = self.patch_embedding(hidden_states) hidden_states = hidden_states.flatten(2).transpose(1, 2) @@ -599,11 +618,13 @@ class CausalWanTransformer3DModel(BaseDiT): for block in self.blocks: hidden_states = self._gradient_checkpointing_func( block, hidden_states, encoder_hidden_states, - timestep_proj, freqs_cis) + timestep_proj, freqs_cis, + block_mask=self.block_mask) else: for block in self.blocks: hidden_states = block(hidden_states, encoder_hidden_states, - timestep_proj, freqs_cis) + timestep_proj, freqs_cis, + block_mask=self.block_mask) # 5. Output norm, projection & unpatchify shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, diff --git a/fastvideo/layers/rotary_embedding.py b/fastvideo/layers/rotary_embedding.py index 6abe90609..4dd789db5 100644 --- a/fastvideo/layers/rotary_embedding.py +++ b/fastvideo/layers/rotary_embedding.py @@ -29,7 +29,9 @@ import torch from fastvideo.distributed.parallel_state import get_sp_group from fastvideo.layers.custom_op import CustomOp +from fastvideo.logger import init_logger +logger = init_logger(__name__) def _rotate_neox(x: torch.Tensor) -> torch.Tensor: x1 = x[..., :x.shape[-1] // 2] @@ -267,6 +269,7 @@ def get_nd_rotary_pos_embed( sp_rank: int = 0, sp_world_size: int = 1, dtype: torch.dtype = torch.float32, + start_frame: int = 0, ) -> tuple[torch.Tensor, torch.Tensor]: """ This is a n-d version of precompute_freqs_cis, which is a RoPE for tokens with n-d structure. @@ -292,6 +295,10 @@ def get_nd_rotary_pos_embed( full_grid = get_meshgrid_nd( start, *args, dim=len(rope_dim_list)) # [3, W, H, D] / [2, W, H] + if start_frame > 0: + full_grid[0] += start_frame + logger.info(full_grid[0]) + # Shard the grid if using sequence parallelism (sp_world_size > 1) assert shard_dim < len( rope_dim_list @@ -370,6 +377,7 @@ def get_rotary_pos_embed( interpolation_factor=1.0, shard_dim: int = 0, dtype: torch.dtype = torch.float32, + start_frame: int = 0, ) -> tuple[torch.Tensor, torch.Tensor]: """ Generate rotary positional embeddings for the given sizes. @@ -413,6 +421,7 @@ def get_rotary_pos_embed( sp_rank=sp_rank, sp_world_size=sp_world_size, dtype=dtype, + start_frame=start_frame, ) return freqs_cos, freqs_sin diff --git a/fastvideo/models/dits/wanvideo.py b/fastvideo/models/dits/wanvideo.py index 1cae2da01..977a181d5 100644 --- a/fastvideo/models/dits/wanvideo.py +++ b/fastvideo/models/dits/wanvideo.py @@ -145,7 +145,7 @@ class WanSelfAttention(nn.Module): class WanT2VCrossAttention(WanSelfAttention): - def forward(self, x, context, context_lens): + def forward(self, x, context, context_lens, crossattn_cache=None): r""" Args: x(Tensor): Shape [B, L1, C] @@ -156,8 +156,20 @@ class WanT2VCrossAttention(WanSelfAttention): # compute query, key, value q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d) - k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d) - v = self.to_v(context)[0].view(b, -1, n, d) + + if crossattn_cache is not None: + if not crossattn_cache["is_init"]: + crossattn_cache["is_init"] = True + k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d) + v = self.to_v(context)[0].view(b, -1, n, d) + crossattn_cache["k"] = k + crossattn_cache["v"] = v + else: + k = crossattn_cache["k"] + v = crossattn_cache["v"] + else: + k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d) + v = self.to_v(context)[0].view(b, -1, n, d) # compute attention x = self.attn(q, k, v) From a464f96b95de9459557e1d17803dea0574e126f1 Mon Sep 17 00:00:00 2001 From: JerryZhou54 Date: Tue, 26 Aug 2025 23:32:30 +0000 Subject: [PATCH 05/13] Fix file dir --- fastvideo/models/dits/causal_wanvideo.py | 652 +++++++++++++++++++++++ 1 file changed, 652 insertions(+) create mode 100644 fastvideo/models/dits/causal_wanvideo.py diff --git a/fastvideo/models/dits/causal_wanvideo.py b/fastvideo/models/dits/causal_wanvideo.py new file mode 100644 index 000000000..244d3ad8c --- /dev/null +++ b/fastvideo/models/dits/causal_wanvideo.py @@ -0,0 +1,652 @@ +# SPDX-License-Identifier: Apache-2.0 + +import math +from typing import Any + +import numpy as np +import torch +import torch.nn as nn + +from torch.nn.attention.flex_attention import create_block_mask, flex_attention +from torch.nn.attention.flex_attention import BlockMask +# wan 1.3B model has a weird channel / head configurations and require max-autotune to work with flexattention +# see https://github.com/pytorch/pytorch/issues/133254 +# change to default for other models +flex_attention = torch.compile( + flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs") +import torch.distributed as dist + +import fastvideo.envs as envs +from fastvideo.attention import (DistributedAttention, + LocalAttention) +from fastvideo.configs.models.dits import WanVideoConfig +from fastvideo.distributed.parallel_state import get_sp_world_size +from fastvideo.forward_context import get_forward_context +from fastvideo.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift, + RMSNorm, ScaleResidual, + ScaleResidualLayerNormScaleShift) +from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.layers.mlp import MLP +from fastvideo.layers.rotary_embedding import (_apply_rotary_emb, + get_rotary_pos_embed) +from fastvideo.layers.visual_embedding import (PatchEmbed) +from fastvideo.logger import init_logger +from fastvideo.models.dits.base import BaseDiT +from fastvideo.models.dits.wanvideo import WanT2VCrossAttention, WanTimeTextImageEmbedding +from fastvideo.platforms import AttentionBackendEnum, current_platform + +logger = init_logger(__name__) +class CausalWanSelfAttention(nn.Module): + + def __init__(self, + dim: int, + num_heads: int, + local_attn_size: int = -1, + sink_size: int = 0, + qk_norm=True, + eps=1e-6, + parallel_attention=False) -> None: + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.local_attn_size = local_attn_size + self.sink_size = sink_size + self.qk_norm = qk_norm + self.eps = eps + self.parallel_attention = parallel_attention + self.max_attention_size = 32760 if local_attn_size == -1 else local_attn_size * 1560 + + # layers + self.to_q = ReplicatedLinear(dim, dim) + self.to_k = ReplicatedLinear(dim, dim) + self.to_v = ReplicatedLinear(dim, dim) + self.to_out = ReplicatedLinear(dim, dim) + self.norm_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + + # Scaled dot product attention + self.attn = LocalAttention( + num_heads=num_heads, + head_size=self.head_dim, + dropout_rate=0, + softmax_scale=None, + causal=False, + supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA)) + + def forward(self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + freqs_cis: tuple[torch.Tensor, torch.Tensor], + block_mask: BlockMask, + kv_cache: dict | None = None, + current_start: int = 0, + cache_start: int | None = None): + r""" + Args: + x(Tensor): Shape [B, L, num_heads, C / num_heads] + seq_lens(Tensor): Shape [B] + grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) + freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] + """ + if cache_start is None: + cache_start = current_start + + cos, sin = freqs_cis + roped_query = _apply_rotary_emb(q, cos, sin, is_neox_style=False) + roped_key = _apply_rotary_emb(k, cos, sin, is_neox_style=False) + + if kv_cache is None: + # Padding for flex attention + padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1] + padded_roped_query = torch.cat( + [roped_query, + torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]], + device=q.device, dtype=v.dtype)], + dim=1 + ) + + padded_roped_key = torch.cat( + [roped_key, torch.zeros([k.shape[0], padded_length, k.shape[2], k.shape[3]], + device=k.device, dtype=v.dtype)], + dim=1 + ) + + padded_v = torch.cat( + [v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]], + device=v.device, dtype=v.dtype)], + dim=1 + ) + + x = flex_attention( + query=padded_roped_query.transpose(2, 1), + key=padded_roped_key.transpose(2, 1), + value=padded_v.transpose(2, 1), + block_mask=block_mask + )[:, :, :-padded_length].transpose(2, 1) + else: + frame_seqlen = q.shape[1] + current_end = current_start + roped_query.shape[1] + sink_tokens = self.sink_size * frame_seqlen + # If we are using local attention and the current KV cache size is larger than the local attention size, we need to truncate the KV cache + kv_cache_size = kv_cache["k"].shape[1] + num_new_tokens = roped_query.shape[1] + if self.local_attn_size != -1 and (current_end > kv_cache["global_end_index"].item()) and ( + num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size): + # Calculate the number of new tokens added in this step + # Shift existing cache content left to discard oldest tokens + # Clone the source slice to avoid overlapping memory error + num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size + num_rolled_tokens = kv_cache["local_end_index"].item() - num_evicted_tokens - sink_tokens + kv_cache["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \ + kv_cache["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() + kv_cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \ + kv_cache["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() + # Insert the new keys/values at the end + local_end_index = kv_cache["local_end_index"].item() + current_end - \ + kv_cache["global_end_index"].item() - num_evicted_tokens + local_start_index = local_end_index - num_new_tokens + kv_cache["k"][:, local_start_index:local_end_index] = roped_key + kv_cache["v"][:, local_start_index:local_end_index] = v + else: + # Assign new keys/values directly up to current_end + local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item() + local_start_index = local_end_index - num_new_tokens + kv_cache["k"][:, local_start_index:local_end_index] = roped_key + kv_cache["v"][:, local_start_index:local_end_index] = v + x = self.attn( + roped_query, + kv_cache["k"][:, max(0, local_end_index - self.max_attention_size):local_end_index], + kv_cache["v"][:, max(0, local_end_index - self.max_attention_size):local_end_index] + ) + kv_cache["global_end_index"].fill_(current_end) + kv_cache["local_end_index"].fill_(local_end_index) + + # output + x = x.flatten(2) + x = self.to_out(x) + return x + +class CausalWanTransformerBlock(nn.Module): + + def __init__(self, + dim: int, + ffn_dim: int, + num_heads: int, + local_attn_size: int = -1, + sink_size: int = 0, + qk_norm: str = "rms_norm_across_heads", + cross_attn_norm: bool = False, + eps: float = 1e-6, + added_kv_proj_dim: int | None = None, + supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, + prefix: str = ""): + super().__init__() + + # 1. Self-attention + self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.to_q = ReplicatedLinear(dim, dim, bias=True) + self.to_k = ReplicatedLinear(dim, dim, bias=True) + self.to_v = ReplicatedLinear(dim, dim, bias=True) + + self.to_out = ReplicatedLinear(dim, dim, bias=True) + self.attn1 = CausalWanSelfAttention( + dim, + num_heads, + local_attn_size=local_attn_size, + sink_size=sink_size, + qk_norm=qk_norm, + eps=eps) + self.hidden_dim = dim + self.num_attention_heads = num_heads + self.local_attn_size = local_attn_size + dim_head = dim // num_heads + if qk_norm == "rms_norm": + self.norm_q = RMSNorm(dim_head, eps=eps) + self.norm_k = RMSNorm(dim_head, eps=eps) + elif qk_norm == "rms_norm_across_heads": + # LTX applies qk norm across all heads + self.norm_q = RMSNorm(dim, eps=eps) + self.norm_k = RMSNorm(dim, eps=eps) + else: + print("QK Norm type not supported") + raise Exception + assert cross_attn_norm is True + self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift( + dim, + norm_type="layer", + eps=eps, + elementwise_affine=True, + dtype=torch.float32, + compute_dtype=torch.float32) + + # 2. Cross-attention + # Only T2V for now + self.attn2 = WanT2VCrossAttention(dim, + num_heads, + qk_norm=qk_norm, + eps=eps) + self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift( + dim, + norm_type="layer", + eps=eps, + elementwise_affine=False, + dtype=torch.float32, + compute_dtype=torch.float32) + + # 3. Feed-forward + self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh") + self.mlp_residual = ScaleResidual() + + self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + freqs_cis: tuple[torch.Tensor, torch.Tensor], + block_mask: BlockMask, + kv_cache: dict | None = None, + crossattn_cache: dict | None = None, + current_start: int = 0, + cache_start: int | None = None, + ) -> torch.Tensor: + if hidden_states.dim() == 4: + hidden_states = hidden_states.squeeze(1) + bs, seq_length, _ = hidden_states.shape + orig_dtype = hidden_states.dtype + # assert orig_dtype != torch.float32 + e = self.scale_shift_table + temb.float() + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk( + 6, dim=1) + assert shift_msa.dtype == torch.float32 + + # 1. Self-attention + norm_hidden_states = (self.norm1(hidden_states.float()) * + (1 + scale_msa) + shift_msa).to(orig_dtype) + query, _ = self.to_q(norm_hidden_states) + key, _ = self.to_k(norm_hidden_states) + value, _ = self.to_v(norm_hidden_states) + + if self.norm_q is not None: + query = self.norm_q(query) + if self.norm_k is not None: + key = self.norm_k(key) + + query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) + key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) + value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) + + attn_output, _ = self.attn1(query, key, value, freqs_cis, block_mask, kv_cache, current_start, cache_start) + attn_output = attn_output.flatten(2) + attn_output, _ = self.to_out(attn_output) + attn_output = attn_output.squeeze(1) + + null_shift = null_scale = torch.tensor([0], device=hidden_states.device) + norm_hidden_states, hidden_states = self.self_attn_residual_norm( + hidden_states, attn_output, gate_msa, null_shift, null_scale) + norm_hidden_states, hidden_states = norm_hidden_states.to( + orig_dtype), hidden_states.to(orig_dtype) + + # 2. Cross-attention + attn_output = self.attn2(norm_hidden_states, + context=encoder_hidden_states, + context_lens=None, + crossattn_cache=crossattn_cache) + norm_hidden_states, hidden_states = self.cross_attn_residual_norm( + hidden_states, attn_output, 1, c_shift_msa, c_scale_msa) + norm_hidden_states, hidden_states = norm_hidden_states.to( + orig_dtype), hidden_states.to(orig_dtype) + + # 3. Feed-forward + ff_output = self.ffn(norm_hidden_states) + hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa) + hidden_states = hidden_states.to(orig_dtype) + + return hidden_states + +class CausalWanTransformer3DModel(BaseDiT): + _fsdp_shard_conditions = WanVideoConfig()._fsdp_shard_conditions + _compile_conditions = WanVideoConfig()._compile_conditions + _supported_attention_backends = WanVideoConfig( + )._supported_attention_backends + param_names_mapping = WanVideoConfig().param_names_mapping + reverse_param_names_mapping = WanVideoConfig().reverse_param_names_mapping + lora_param_names_mapping = WanVideoConfig().lora_param_names_mapping + + def __init__(self, config: WanVideoConfig, hf_config: dict[str, + Any]) -> None: + super().__init__(config=config, hf_config=hf_config) + + inner_dim = config.num_attention_heads * config.attention_head_dim + self.hidden_size = config.hidden_size + self.num_attention_heads = config.num_attention_heads + self.in_channels = config.in_channels + self.out_channels = config.out_channels + self.num_channels_latents = config.num_channels_latents + self.patch_size = config.patch_size + self.text_len = config.text_len + self.local_attn_size = config.local_attn_size + + # 1. Patch & position embedding + self.patch_embedding = PatchEmbed(in_chans=config.in_channels, + embed_dim=inner_dim, + patch_size=config.patch_size, + flatten=False) + + # 2. Condition embeddings + self.condition_embedder = WanTimeTextImageEmbedding( + dim=inner_dim, + time_freq_dim=config.freq_dim, + text_embed_dim=config.text_dim, + image_embed_dim=config.image_dim, + ) + + # 3. Transformer blocks + self.blocks = nn.ModuleList([ + CausalWanTransformerBlock(inner_dim, + config.ffn_dim, + config.num_attention_heads, + config.local_attn_size, + config.sink_size, + config.qk_norm, + config.cross_attn_norm, + config.eps, + config.added_kv_proj_dim, + self._supported_attention_backends, + prefix=f"{config.prefix}.blocks.{i}") + for i in range(config.num_layers) + ]) + + # 4. Output norm & projection + self.norm_out = LayerNormScaleShift(inner_dim, + norm_type="layer", + eps=config.eps, + elementwise_affine=False, + dtype=torch.float32, + compute_dtype=torch.float32) + self.proj_out = nn.Linear( + inner_dim, config.out_channels * math.prod(config.patch_size)) + self.scale_shift_table = nn.Parameter( + torch.randn(1, 2, inner_dim) / inner_dim**0.5) + + self.gradient_checkpointing = False + + # Causal-specific + self.block_mask = None + self.num_frame_per_block = 1 + self.independent_first_frame = False + + self.__post_init__() + + @staticmethod + def _prepare_blockwise_causal_attn_mask( + device: torch.device | str, num_frames: int = 21, + frame_seqlen: int = 1560, num_frame_per_block=1, local_attn_size=-1 + ) -> BlockMask: + """ + we will divide the token sequence into the following format + [1 latent frame] [1 latent frame] ... [1 latent frame] + We use flexattention to construct the attention mask + """ + total_length = num_frames * frame_seqlen + + # we do right padding to get to a multiple of 128 + padded_length = math.ceil(total_length / 128) * 128 - total_length + + ends = torch.zeros(total_length + padded_length, + device=device, dtype=torch.long) + + # Block-wise causal mask will attend to all elements that are before the end of the current chunk + frame_indices = torch.arange( + start=0, + end=total_length, + step=frame_seqlen * num_frame_per_block, + device=device + ) + + for tmp in frame_indices: + ends[tmp:tmp + frame_seqlen * num_frame_per_block] = tmp + \ + frame_seqlen * num_frame_per_block + + def attention_mask(b, h, q_idx, kv_idx): + if local_attn_size == -1: + return (kv_idx < ends[q_idx]) | (q_idx == kv_idx) + else: + return ((kv_idx < ends[q_idx]) & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))) | (q_idx == kv_idx) + # return ((kv_idx < total_length) & (q_idx < total_length)) | (q_idx == kv_idx) # bidirectional mask + + block_mask = create_block_mask(attention_mask, B=None, H=None, Q_LEN=total_length + padded_length, + KV_LEN=total_length + padded_length, _compile=False, device=device) + + if not dist.is_initialized() or dist.get_rank() == 0: + print( + f" cache a block wise causal mask with block size of {num_frame_per_block} frames") + print(block_mask) + + # import imageio + # import numpy as np + # from torch.nn.attention.flex_attention import create_mask + + # mask = create_mask(attention_mask, B=None, H=None, Q_LEN=total_length + + # padded_length, KV_LEN=total_length + padded_length, device=device) + # import cv2 + # mask = cv2.resize(mask[0, 0].cpu().float().numpy(), (1024, 1024)) + # imageio.imwrite("mask_%d.jpg" % (0), np.uint8(255. * mask)) + + return block_mask + + def _forward_inference( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + timestep: torch.LongTensor, + encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] + | None = None, + kv_cache: dict = None, + crossattn_cache: dict = None, + current_start: int = 0, + cache_start: int = 0, + **kwargs) -> torch.Tensor: + r""" + Run the diffusion model with kv caching. + See Algorithm 2 of CausVid paper https://arxiv.org/abs/2412.07772 for details. + This function will be run for num_frame times. + Process the latent frames one by one (1560 tokens each) + """ + + orig_dtype = hidden_states.dtype + if not isinstance(encoder_hidden_states, torch.Tensor): + encoder_hidden_states = encoder_hidden_states[0] + if isinstance(encoder_hidden_states_image, + list) and len(encoder_hidden_states_image) > 0: + encoder_hidden_states_image = encoder_hidden_states_image[0] + else: + encoder_hidden_states_image = None + + batch_size, num_channels, num_frames, height, width = hidden_states.shape + p_t, p_h, p_w = self.patch_size + post_patch_num_frames = num_frames // p_t + post_patch_height = height // p_h + post_patch_width = width // p_w + + # Get rotary embeddings + d = self.hidden_size // self.num_attention_heads + rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] + freqs_cos, freqs_sin = get_rotary_pos_embed( + (post_patch_num_frames * get_sp_world_size(), post_patch_height, + post_patch_width), + self.hidden_size, + self.num_attention_heads, + rope_dim_list, + dtype=torch.float32 if current_platform.is_mps() else torch.float64, + rope_theta=10000, + start_frame=current_start # Assume that current_start is 0 when kv_cache is None + ) + freqs_cos = freqs_cos.to(hidden_states.device) + freqs_sin = freqs_sin.to(hidden_states.device) + freqs_cis = (freqs_cos.float(), + freqs_sin.float()) if freqs_cos is not None else None + + hidden_states = self.patch_embedding(hidden_states) + hidden_states = hidden_states.flatten(2).transpose(1, 2) + + temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( + timestep, encoder_hidden_states, encoder_hidden_states_image) + timestep_proj = timestep_proj.unflatten(1, (6, -1)) + + if encoder_hidden_states_image is not None: + encoder_hidden_states = torch.concat( + [encoder_hidden_states_image, encoder_hidden_states], dim=1) + + encoder_hidden_states = encoder_hidden_states.to( + orig_dtype) if current_platform.is_mps( + ) else encoder_hidden_states # cast to orig_dtype for MPS + + assert encoder_hidden_states.dtype == orig_dtype + + # 4. Transformer blocks + for block_index, block in enumerate(self.blocks): + if torch.is_grad_enabled() and self.gradient_checkpointing: + causal_kwargs = { + "kv_cache": kv_cache[block_index], + "current_start": current_start, + "cache_start": cache_start + } + hidden_states = self._gradient_checkpointing_func( + block, hidden_states, encoder_hidden_states, + timestep_proj, freqs_cis, + **causal_kwargs) + else: + causal_kwargs = { + "kv_cache": kv_cache[block_index], + "crossattn_cache": crossattn_cache[block_index], + "current_start": current_start, + "cache_start": cache_start + } + hidden_states = block(hidden_states, encoder_hidden_states, + timestep_proj, freqs_cis, + **causal_kwargs) + + # 5. Output norm, projection & unpatchify + shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, + dim=1) + hidden_states = self.norm_out(hidden_states, shift, scale) + hidden_states = self.proj_out(hidden_states) + + hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, + post_patch_height, + post_patch_width, p_t, p_h, p_w, + -1) + hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) + output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) + + return output + + def _forward_train(self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + timestep: torch.LongTensor, + encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] + | None = None, + **kwargs) -> torch.Tensor: + + orig_dtype = hidden_states.dtype + if not isinstance(encoder_hidden_states, torch.Tensor): + encoder_hidden_states = encoder_hidden_states[0] + if isinstance(encoder_hidden_states_image, + list) and len(encoder_hidden_states_image) > 0: + encoder_hidden_states_image = encoder_hidden_states_image[0] + else: + encoder_hidden_states_image = None + + batch_size, num_channels, num_frames, height, width = hidden_states.shape + p_t, p_h, p_w = self.patch_size + post_patch_num_frames = num_frames // p_t + post_patch_height = height // p_h + post_patch_width = width // p_w + + # Get rotary embeddings + d = self.hidden_size // self.num_attention_heads + rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] + freqs_cos, freqs_sin = get_rotary_pos_embed( + (post_patch_num_frames * get_sp_world_size(), post_patch_height, + post_patch_width), + self.hidden_size, + self.num_attention_heads, + rope_dim_list, + dtype=torch.float32 if current_platform.is_mps() else torch.float64, + rope_theta=10000) + freqs_cos = freqs_cos.to(hidden_states.device) + freqs_sin = freqs_sin.to(hidden_states.device) + freqs_cis = (freqs_cos.float(), + freqs_sin.float()) if freqs_cos is not None else None + + # Construct blockwise causal attn mask + if self.block_mask is None: + self.block_mask = self._prepare_blockwise_causal_attn_mask( + device=hidden_states.device, + num_frames=num_frames, + frame_seqlen=post_patch_height * post_patch_width, + num_frame_per_block=self.num_frame_per_block, + local_attn_size=self.local_attn_size + ) + + hidden_states = self.patch_embedding(hidden_states) + hidden_states = hidden_states.flatten(2).transpose(1, 2) + + temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( + timestep, encoder_hidden_states, encoder_hidden_states_image) + timestep_proj = timestep_proj.unflatten(1, (6, -1)) + + if encoder_hidden_states_image is not None: + encoder_hidden_states = torch.concat( + [encoder_hidden_states_image, encoder_hidden_states], dim=1) + + encoder_hidden_states = encoder_hidden_states.to( + orig_dtype) if current_platform.is_mps( + ) else encoder_hidden_states # cast to orig_dtype for MPS + + assert encoder_hidden_states.dtype == orig_dtype + + # 4. Transformer blocks + if torch.is_grad_enabled() and self.gradient_checkpointing: + for block in self.blocks: + hidden_states = self._gradient_checkpointing_func( + block, hidden_states, encoder_hidden_states, + timestep_proj, freqs_cis, + block_mask=self.block_mask) + else: + for block in self.blocks: + hidden_states = block(hidden_states, encoder_hidden_states, + timestep_proj, freqs_cis, + block_mask=self.block_mask) + + # 5. Output norm, projection & unpatchify + shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, + dim=1) + hidden_states = self.norm_out(hidden_states, shift, scale) + hidden_states = self.proj_out(hidden_states) + + hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, + post_patch_height, + post_patch_width, p_t, p_h, p_w, + -1) + hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) + output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) + + return output + + def forward( + self, + *args, + **kwargs + ): + if kwargs.get('kv_cache', None) is not None: + return self._forward_inference(*args, **kwargs) + else: + return self._forward_train(*args, **kwargs) From b5519f1f91167f87c93a8b7ea4e57916f35946cb Mon Sep 17 00:00:00 2001 From: Wei Zhou <69577934+JerryZhou54@users.noreply.github.com> Date: Tue, 26 Aug 2025 16:33:25 -0700 Subject: [PATCH 06/13] Delete fastvideo/configs/models/dits/causal_wanvideo.py --- .../configs/models/dits/causal_wanvideo.py | 652 ------------------ 1 file changed, 652 deletions(-) delete mode 100644 fastvideo/configs/models/dits/causal_wanvideo.py diff --git a/fastvideo/configs/models/dits/causal_wanvideo.py b/fastvideo/configs/models/dits/causal_wanvideo.py deleted file mode 100644 index 244d3ad8c..000000000 --- a/fastvideo/configs/models/dits/causal_wanvideo.py +++ /dev/null @@ -1,652 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -import math -from typing import Any - -import numpy as np -import torch -import torch.nn as nn - -from torch.nn.attention.flex_attention import create_block_mask, flex_attention -from torch.nn.attention.flex_attention import BlockMask -# wan 1.3B model has a weird channel / head configurations and require max-autotune to work with flexattention -# see https://github.com/pytorch/pytorch/issues/133254 -# change to default for other models -flex_attention = torch.compile( - flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs") -import torch.distributed as dist - -import fastvideo.envs as envs -from fastvideo.attention import (DistributedAttention, - LocalAttention) -from fastvideo.configs.models.dits import WanVideoConfig -from fastvideo.distributed.parallel_state import get_sp_world_size -from fastvideo.forward_context import get_forward_context -from fastvideo.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift, - RMSNorm, ScaleResidual, - ScaleResidualLayerNormScaleShift) -from fastvideo.layers.linear import ReplicatedLinear -from fastvideo.layers.mlp import MLP -from fastvideo.layers.rotary_embedding import (_apply_rotary_emb, - get_rotary_pos_embed) -from fastvideo.layers.visual_embedding import (PatchEmbed) -from fastvideo.logger import init_logger -from fastvideo.models.dits.base import BaseDiT -from fastvideo.models.dits.wanvideo import WanT2VCrossAttention, WanTimeTextImageEmbedding -from fastvideo.platforms import AttentionBackendEnum, current_platform - -logger = init_logger(__name__) -class CausalWanSelfAttention(nn.Module): - - def __init__(self, - dim: int, - num_heads: int, - local_attn_size: int = -1, - sink_size: int = 0, - qk_norm=True, - eps=1e-6, - parallel_attention=False) -> None: - assert dim % num_heads == 0 - super().__init__() - self.dim = dim - self.num_heads = num_heads - self.head_dim = dim // num_heads - self.local_attn_size = local_attn_size - self.sink_size = sink_size - self.qk_norm = qk_norm - self.eps = eps - self.parallel_attention = parallel_attention - self.max_attention_size = 32760 if local_attn_size == -1 else local_attn_size * 1560 - - # layers - self.to_q = ReplicatedLinear(dim, dim) - self.to_k = ReplicatedLinear(dim, dim) - self.to_v = ReplicatedLinear(dim, dim) - self.to_out = ReplicatedLinear(dim, dim) - self.norm_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() - self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() - - # Scaled dot product attention - self.attn = LocalAttention( - num_heads=num_heads, - head_size=self.head_dim, - dropout_rate=0, - softmax_scale=None, - causal=False, - supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, - AttentionBackendEnum.TORCH_SDPA)) - - def forward(self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - freqs_cis: tuple[torch.Tensor, torch.Tensor], - block_mask: BlockMask, - kv_cache: dict | None = None, - current_start: int = 0, - cache_start: int | None = None): - r""" - Args: - x(Tensor): Shape [B, L, num_heads, C / num_heads] - seq_lens(Tensor): Shape [B] - grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) - freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] - """ - if cache_start is None: - cache_start = current_start - - cos, sin = freqs_cis - roped_query = _apply_rotary_emb(q, cos, sin, is_neox_style=False) - roped_key = _apply_rotary_emb(k, cos, sin, is_neox_style=False) - - if kv_cache is None: - # Padding for flex attention - padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1] - padded_roped_query = torch.cat( - [roped_query, - torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]], - device=q.device, dtype=v.dtype)], - dim=1 - ) - - padded_roped_key = torch.cat( - [roped_key, torch.zeros([k.shape[0], padded_length, k.shape[2], k.shape[3]], - device=k.device, dtype=v.dtype)], - dim=1 - ) - - padded_v = torch.cat( - [v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]], - device=v.device, dtype=v.dtype)], - dim=1 - ) - - x = flex_attention( - query=padded_roped_query.transpose(2, 1), - key=padded_roped_key.transpose(2, 1), - value=padded_v.transpose(2, 1), - block_mask=block_mask - )[:, :, :-padded_length].transpose(2, 1) - else: - frame_seqlen = q.shape[1] - current_end = current_start + roped_query.shape[1] - sink_tokens = self.sink_size * frame_seqlen - # If we are using local attention and the current KV cache size is larger than the local attention size, we need to truncate the KV cache - kv_cache_size = kv_cache["k"].shape[1] - num_new_tokens = roped_query.shape[1] - if self.local_attn_size != -1 and (current_end > kv_cache["global_end_index"].item()) and ( - num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size): - # Calculate the number of new tokens added in this step - # Shift existing cache content left to discard oldest tokens - # Clone the source slice to avoid overlapping memory error - num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size - num_rolled_tokens = kv_cache["local_end_index"].item() - num_evicted_tokens - sink_tokens - kv_cache["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \ - kv_cache["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() - kv_cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \ - kv_cache["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() - # Insert the new keys/values at the end - local_end_index = kv_cache["local_end_index"].item() + current_end - \ - kv_cache["global_end_index"].item() - num_evicted_tokens - local_start_index = local_end_index - num_new_tokens - kv_cache["k"][:, local_start_index:local_end_index] = roped_key - kv_cache["v"][:, local_start_index:local_end_index] = v - else: - # Assign new keys/values directly up to current_end - local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item() - local_start_index = local_end_index - num_new_tokens - kv_cache["k"][:, local_start_index:local_end_index] = roped_key - kv_cache["v"][:, local_start_index:local_end_index] = v - x = self.attn( - roped_query, - kv_cache["k"][:, max(0, local_end_index - self.max_attention_size):local_end_index], - kv_cache["v"][:, max(0, local_end_index - self.max_attention_size):local_end_index] - ) - kv_cache["global_end_index"].fill_(current_end) - kv_cache["local_end_index"].fill_(local_end_index) - - # output - x = x.flatten(2) - x = self.to_out(x) - return x - -class CausalWanTransformerBlock(nn.Module): - - def __init__(self, - dim: int, - ffn_dim: int, - num_heads: int, - local_attn_size: int = -1, - sink_size: int = 0, - qk_norm: str = "rms_norm_across_heads", - cross_attn_norm: bool = False, - eps: float = 1e-6, - added_kv_proj_dim: int | None = None, - supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, - prefix: str = ""): - super().__init__() - - # 1. Self-attention - self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) - self.to_q = ReplicatedLinear(dim, dim, bias=True) - self.to_k = ReplicatedLinear(dim, dim, bias=True) - self.to_v = ReplicatedLinear(dim, dim, bias=True) - - self.to_out = ReplicatedLinear(dim, dim, bias=True) - self.attn1 = CausalWanSelfAttention( - dim, - num_heads, - local_attn_size=local_attn_size, - sink_size=sink_size, - qk_norm=qk_norm, - eps=eps) - self.hidden_dim = dim - self.num_attention_heads = num_heads - self.local_attn_size = local_attn_size - dim_head = dim // num_heads - if qk_norm == "rms_norm": - self.norm_q = RMSNorm(dim_head, eps=eps) - self.norm_k = RMSNorm(dim_head, eps=eps) - elif qk_norm == "rms_norm_across_heads": - # LTX applies qk norm across all heads - self.norm_q = RMSNorm(dim, eps=eps) - self.norm_k = RMSNorm(dim, eps=eps) - else: - print("QK Norm type not supported") - raise Exception - assert cross_attn_norm is True - self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift( - dim, - norm_type="layer", - eps=eps, - elementwise_affine=True, - dtype=torch.float32, - compute_dtype=torch.float32) - - # 2. Cross-attention - # Only T2V for now - self.attn2 = WanT2VCrossAttention(dim, - num_heads, - qk_norm=qk_norm, - eps=eps) - self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift( - dim, - norm_type="layer", - eps=eps, - elementwise_affine=False, - dtype=torch.float32, - compute_dtype=torch.float32) - - # 3. Feed-forward - self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh") - self.mlp_residual = ScaleResidual() - - self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) - - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor, - temb: torch.Tensor, - freqs_cis: tuple[torch.Tensor, torch.Tensor], - block_mask: BlockMask, - kv_cache: dict | None = None, - crossattn_cache: dict | None = None, - current_start: int = 0, - cache_start: int | None = None, - ) -> torch.Tensor: - if hidden_states.dim() == 4: - hidden_states = hidden_states.squeeze(1) - bs, seq_length, _ = hidden_states.shape - orig_dtype = hidden_states.dtype - # assert orig_dtype != torch.float32 - e = self.scale_shift_table + temb.float() - shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk( - 6, dim=1) - assert shift_msa.dtype == torch.float32 - - # 1. Self-attention - norm_hidden_states = (self.norm1(hidden_states.float()) * - (1 + scale_msa) + shift_msa).to(orig_dtype) - query, _ = self.to_q(norm_hidden_states) - key, _ = self.to_k(norm_hidden_states) - value, _ = self.to_v(norm_hidden_states) - - if self.norm_q is not None: - query = self.norm_q(query) - if self.norm_k is not None: - key = self.norm_k(key) - - query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) - key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) - value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) - - attn_output, _ = self.attn1(query, key, value, freqs_cis, block_mask, kv_cache, current_start, cache_start) - attn_output = attn_output.flatten(2) - attn_output, _ = self.to_out(attn_output) - attn_output = attn_output.squeeze(1) - - null_shift = null_scale = torch.tensor([0], device=hidden_states.device) - norm_hidden_states, hidden_states = self.self_attn_residual_norm( - hidden_states, attn_output, gate_msa, null_shift, null_scale) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype), hidden_states.to(orig_dtype) - - # 2. Cross-attention - attn_output = self.attn2(norm_hidden_states, - context=encoder_hidden_states, - context_lens=None, - crossattn_cache=crossattn_cache) - norm_hidden_states, hidden_states = self.cross_attn_residual_norm( - hidden_states, attn_output, 1, c_shift_msa, c_scale_msa) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype), hidden_states.to(orig_dtype) - - # 3. Feed-forward - ff_output = self.ffn(norm_hidden_states) - hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa) - hidden_states = hidden_states.to(orig_dtype) - - return hidden_states - -class CausalWanTransformer3DModel(BaseDiT): - _fsdp_shard_conditions = WanVideoConfig()._fsdp_shard_conditions - _compile_conditions = WanVideoConfig()._compile_conditions - _supported_attention_backends = WanVideoConfig( - )._supported_attention_backends - param_names_mapping = WanVideoConfig().param_names_mapping - reverse_param_names_mapping = WanVideoConfig().reverse_param_names_mapping - lora_param_names_mapping = WanVideoConfig().lora_param_names_mapping - - def __init__(self, config: WanVideoConfig, hf_config: dict[str, - Any]) -> None: - super().__init__(config=config, hf_config=hf_config) - - inner_dim = config.num_attention_heads * config.attention_head_dim - self.hidden_size = config.hidden_size - self.num_attention_heads = config.num_attention_heads - self.in_channels = config.in_channels - self.out_channels = config.out_channels - self.num_channels_latents = config.num_channels_latents - self.patch_size = config.patch_size - self.text_len = config.text_len - self.local_attn_size = config.local_attn_size - - # 1. Patch & position embedding - self.patch_embedding = PatchEmbed(in_chans=config.in_channels, - embed_dim=inner_dim, - patch_size=config.patch_size, - flatten=False) - - # 2. Condition embeddings - self.condition_embedder = WanTimeTextImageEmbedding( - dim=inner_dim, - time_freq_dim=config.freq_dim, - text_embed_dim=config.text_dim, - image_embed_dim=config.image_dim, - ) - - # 3. Transformer blocks - self.blocks = nn.ModuleList([ - CausalWanTransformerBlock(inner_dim, - config.ffn_dim, - config.num_attention_heads, - config.local_attn_size, - config.sink_size, - config.qk_norm, - config.cross_attn_norm, - config.eps, - config.added_kv_proj_dim, - self._supported_attention_backends, - prefix=f"{config.prefix}.blocks.{i}") - for i in range(config.num_layers) - ]) - - # 4. Output norm & projection - self.norm_out = LayerNormScaleShift(inner_dim, - norm_type="layer", - eps=config.eps, - elementwise_affine=False, - dtype=torch.float32, - compute_dtype=torch.float32) - self.proj_out = nn.Linear( - inner_dim, config.out_channels * math.prod(config.patch_size)) - self.scale_shift_table = nn.Parameter( - torch.randn(1, 2, inner_dim) / inner_dim**0.5) - - self.gradient_checkpointing = False - - # Causal-specific - self.block_mask = None - self.num_frame_per_block = 1 - self.independent_first_frame = False - - self.__post_init__() - - @staticmethod - def _prepare_blockwise_causal_attn_mask( - device: torch.device | str, num_frames: int = 21, - frame_seqlen: int = 1560, num_frame_per_block=1, local_attn_size=-1 - ) -> BlockMask: - """ - we will divide the token sequence into the following format - [1 latent frame] [1 latent frame] ... [1 latent frame] - We use flexattention to construct the attention mask - """ - total_length = num_frames * frame_seqlen - - # we do right padding to get to a multiple of 128 - padded_length = math.ceil(total_length / 128) * 128 - total_length - - ends = torch.zeros(total_length + padded_length, - device=device, dtype=torch.long) - - # Block-wise causal mask will attend to all elements that are before the end of the current chunk - frame_indices = torch.arange( - start=0, - end=total_length, - step=frame_seqlen * num_frame_per_block, - device=device - ) - - for tmp in frame_indices: - ends[tmp:tmp + frame_seqlen * num_frame_per_block] = tmp + \ - frame_seqlen * num_frame_per_block - - def attention_mask(b, h, q_idx, kv_idx): - if local_attn_size == -1: - return (kv_idx < ends[q_idx]) | (q_idx == kv_idx) - else: - return ((kv_idx < ends[q_idx]) & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))) | (q_idx == kv_idx) - # return ((kv_idx < total_length) & (q_idx < total_length)) | (q_idx == kv_idx) # bidirectional mask - - block_mask = create_block_mask(attention_mask, B=None, H=None, Q_LEN=total_length + padded_length, - KV_LEN=total_length + padded_length, _compile=False, device=device) - - if not dist.is_initialized() or dist.get_rank() == 0: - print( - f" cache a block wise causal mask with block size of {num_frame_per_block} frames") - print(block_mask) - - # import imageio - # import numpy as np - # from torch.nn.attention.flex_attention import create_mask - - # mask = create_mask(attention_mask, B=None, H=None, Q_LEN=total_length + - # padded_length, KV_LEN=total_length + padded_length, device=device) - # import cv2 - # mask = cv2.resize(mask[0, 0].cpu().float().numpy(), (1024, 1024)) - # imageio.imwrite("mask_%d.jpg" % (0), np.uint8(255. * mask)) - - return block_mask - - def _forward_inference( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor | list[torch.Tensor], - timestep: torch.LongTensor, - encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] - | None = None, - kv_cache: dict = None, - crossattn_cache: dict = None, - current_start: int = 0, - cache_start: int = 0, - **kwargs) -> torch.Tensor: - r""" - Run the diffusion model with kv caching. - See Algorithm 2 of CausVid paper https://arxiv.org/abs/2412.07772 for details. - This function will be run for num_frame times. - Process the latent frames one by one (1560 tokens each) - """ - - orig_dtype = hidden_states.dtype - if not isinstance(encoder_hidden_states, torch.Tensor): - encoder_hidden_states = encoder_hidden_states[0] - if isinstance(encoder_hidden_states_image, - list) and len(encoder_hidden_states_image) > 0: - encoder_hidden_states_image = encoder_hidden_states_image[0] - else: - encoder_hidden_states_image = None - - batch_size, num_channels, num_frames, height, width = hidden_states.shape - p_t, p_h, p_w = self.patch_size - post_patch_num_frames = num_frames // p_t - post_patch_height = height // p_h - post_patch_width = width // p_w - - # Get rotary embeddings - d = self.hidden_size // self.num_attention_heads - rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] - freqs_cos, freqs_sin = get_rotary_pos_embed( - (post_patch_num_frames * get_sp_world_size(), post_patch_height, - post_patch_width), - self.hidden_size, - self.num_attention_heads, - rope_dim_list, - dtype=torch.float32 if current_platform.is_mps() else torch.float64, - rope_theta=10000, - start_frame=current_start # Assume that current_start is 0 when kv_cache is None - ) - freqs_cos = freqs_cos.to(hidden_states.device) - freqs_sin = freqs_sin.to(hidden_states.device) - freqs_cis = (freqs_cos.float(), - freqs_sin.float()) if freqs_cos is not None else None - - hidden_states = self.patch_embedding(hidden_states) - hidden_states = hidden_states.flatten(2).transpose(1, 2) - - temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( - timestep, encoder_hidden_states, encoder_hidden_states_image) - timestep_proj = timestep_proj.unflatten(1, (6, -1)) - - if encoder_hidden_states_image is not None: - encoder_hidden_states = torch.concat( - [encoder_hidden_states_image, encoder_hidden_states], dim=1) - - encoder_hidden_states = encoder_hidden_states.to( - orig_dtype) if current_platform.is_mps( - ) else encoder_hidden_states # cast to orig_dtype for MPS - - assert encoder_hidden_states.dtype == orig_dtype - - # 4. Transformer blocks - for block_index, block in enumerate(self.blocks): - if torch.is_grad_enabled() and self.gradient_checkpointing: - causal_kwargs = { - "kv_cache": kv_cache[block_index], - "current_start": current_start, - "cache_start": cache_start - } - hidden_states = self._gradient_checkpointing_func( - block, hidden_states, encoder_hidden_states, - timestep_proj, freqs_cis, - **causal_kwargs) - else: - causal_kwargs = { - "kv_cache": kv_cache[block_index], - "crossattn_cache": crossattn_cache[block_index], - "current_start": current_start, - "cache_start": cache_start - } - hidden_states = block(hidden_states, encoder_hidden_states, - timestep_proj, freqs_cis, - **causal_kwargs) - - # 5. Output norm, projection & unpatchify - shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, - dim=1) - hidden_states = self.norm_out(hidden_states, shift, scale) - hidden_states = self.proj_out(hidden_states) - - hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, - post_patch_height, - post_patch_width, p_t, p_h, p_w, - -1) - hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) - output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) - - return output - - def _forward_train(self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor | list[torch.Tensor], - timestep: torch.LongTensor, - encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] - | None = None, - **kwargs) -> torch.Tensor: - - orig_dtype = hidden_states.dtype - if not isinstance(encoder_hidden_states, torch.Tensor): - encoder_hidden_states = encoder_hidden_states[0] - if isinstance(encoder_hidden_states_image, - list) and len(encoder_hidden_states_image) > 0: - encoder_hidden_states_image = encoder_hidden_states_image[0] - else: - encoder_hidden_states_image = None - - batch_size, num_channels, num_frames, height, width = hidden_states.shape - p_t, p_h, p_w = self.patch_size - post_patch_num_frames = num_frames // p_t - post_patch_height = height // p_h - post_patch_width = width // p_w - - # Get rotary embeddings - d = self.hidden_size // self.num_attention_heads - rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] - freqs_cos, freqs_sin = get_rotary_pos_embed( - (post_patch_num_frames * get_sp_world_size(), post_patch_height, - post_patch_width), - self.hidden_size, - self.num_attention_heads, - rope_dim_list, - dtype=torch.float32 if current_platform.is_mps() else torch.float64, - rope_theta=10000) - freqs_cos = freqs_cos.to(hidden_states.device) - freqs_sin = freqs_sin.to(hidden_states.device) - freqs_cis = (freqs_cos.float(), - freqs_sin.float()) if freqs_cos is not None else None - - # Construct blockwise causal attn mask - if self.block_mask is None: - self.block_mask = self._prepare_blockwise_causal_attn_mask( - device=hidden_states.device, - num_frames=num_frames, - frame_seqlen=post_patch_height * post_patch_width, - num_frame_per_block=self.num_frame_per_block, - local_attn_size=self.local_attn_size - ) - - hidden_states = self.patch_embedding(hidden_states) - hidden_states = hidden_states.flatten(2).transpose(1, 2) - - temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( - timestep, encoder_hidden_states, encoder_hidden_states_image) - timestep_proj = timestep_proj.unflatten(1, (6, -1)) - - if encoder_hidden_states_image is not None: - encoder_hidden_states = torch.concat( - [encoder_hidden_states_image, encoder_hidden_states], dim=1) - - encoder_hidden_states = encoder_hidden_states.to( - orig_dtype) if current_platform.is_mps( - ) else encoder_hidden_states # cast to orig_dtype for MPS - - assert encoder_hidden_states.dtype == orig_dtype - - # 4. Transformer blocks - if torch.is_grad_enabled() and self.gradient_checkpointing: - for block in self.blocks: - hidden_states = self._gradient_checkpointing_func( - block, hidden_states, encoder_hidden_states, - timestep_proj, freqs_cis, - block_mask=self.block_mask) - else: - for block in self.blocks: - hidden_states = block(hidden_states, encoder_hidden_states, - timestep_proj, freqs_cis, - block_mask=self.block_mask) - - # 5. Output norm, projection & unpatchify - shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, - dim=1) - hidden_states = self.norm_out(hidden_states, shift, scale) - hidden_states = self.proj_out(hidden_states) - - hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, - post_patch_height, - post_patch_width, p_t, p_h, p_w, - -1) - hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) - output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) - - return output - - def forward( - self, - *args, - **kwargs - ): - if kwargs.get('kv_cache', None) is not None: - return self._forward_inference(*args, **kwargs) - else: - return self._forward_train(*args, **kwargs) From cf1942fd47809a9601a3b58f46b8ea37f0e79081 Mon Sep 17 00:00:00 2001 From: William Lin Date: Tue, 26 Aug 2025 23:13:20 -0700 Subject: [PATCH 07/13] [dev] Will/causal infer (#764) --- examples/inference/basic/basic_causal.py | 66 ++ .../basic/wan/wan_causal_dmd_pipeline.py | 69 ++ .../basic/wan/wan_causal_pipeline.py | 71 +++ fastvideo/pipelines/stages/__init__.py | 4 + .../pipelines/stages/causal_denoising.py | 601 ++++++++++++++++++ .../checkpoint_conversion/wan_to_diffusers.py | 88 +++ 6 files changed, 899 insertions(+) create mode 100644 examples/inference/basic/basic_causal.py create mode 100644 fastvideo/pipelines/basic/wan/wan_causal_dmd_pipeline.py create mode 100644 fastvideo/pipelines/basic/wan/wan_causal_pipeline.py create mode 100644 fastvideo/pipelines/stages/causal_denoising.py create mode 100644 scripts/checkpoint_conversion/wan_to_diffusers.py diff --git a/examples/inference/basic/basic_causal.py b/examples/inference/basic/basic_causal.py new file mode 100644 index 000000000..956fc2b2e --- /dev/null +++ b/examples/inference/basic/basic_causal.py @@ -0,0 +1,66 @@ +import os +import time +from fastvideo import VideoGenerator, SamplingParam + +OUTPUT_PATH = "video_samples_causal" +def main(): + # FastVideo will automatically use the optimal default arguments for the + # model. + # If a local path is provided, FastVideo will make a best effort + # attempt to identify the optimal arguments. + load_start_time = time.perf_counter() + model_name = "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers" + generator = VideoGenerator.from_pretrained( + model_name, + # FastVideo will automatically handle distributed setup + num_gpus=1, + use_fsdp_inference=True, + text_encoder_cpu_offload=False, + dit_cpu_offload=False, + ) + load_end_time = time.perf_counter() + load_time = load_end_time - load_start_time + + + sampling_param = SamplingParam.from_pretrained(model_name) + + start_time = time.perf_counter() + prompt = ( + "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " + "wide with interest. The playful yet serene atmosphere is complemented by soft " + "natural light filtering through the petals. Mid-shot, warm and cheerful tones." + ) + video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param) + end_time = time.perf_counter() + gen_time1 = end_time - start_time + e2e_gen_time = end_time - load_start_time + + # Generate another video with a different prompt, without reloading the + # model! + gen_time_list = [] + for i in range(10): + start_time = time.perf_counter() + prompt2 = ( + "A majestic lion strides across the golden savanna, its powerful frame " + "glistening under the warm afternoon sun. The tall grass ripples gently in " + "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " + "embodying the raw energy of the wild. Low angle, steady tracking shot, " + "cinematic.") + video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param) + end_time = time.perf_counter() + gen_time2 = end_time - start_time + gen_time_list.append(gen_time2) + # print(f"Time taken for video {i}: {gen_time2} seconds") + + print(f"Time taken for first video: {gen_time1} seconds") + print(f"Average time taken for 10 videos: {sum(gen_time_list) / len(gen_time_list)} seconds") + for i, gen_time in enumerate(gen_time_list): + print(f"Time taken for video {i}: {gen_time} seconds") + + # print(f"Time taken for second video: {gen_time2} seconds") + print(f"Time taken to load model: {load_time} seconds") + print(f"Time taken for first e2e generation: {e2e_gen_time} seconds") + + +if __name__ == "__main__": + main() diff --git a/fastvideo/pipelines/basic/wan/wan_causal_dmd_pipeline.py b/fastvideo/pipelines/basic/wan/wan_causal_dmd_pipeline.py new file mode 100644 index 000000000..a46898530 --- /dev/null +++ b/fastvideo/pipelines/basic/wan/wan_causal_dmd_pipeline.py @@ -0,0 +1,69 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Wan causal DMD pipeline implementation. + +This module wires the causal DMD denoising stage into the modular pipeline. +""" + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler) +from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline + +# isort: off +from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, + CausalDMDDenosingStage, InputValidationStage, + LatentPreparationStage, TextEncodingStage, + TimestepPreparationStage) +# isort: on + +logger = init_logger(__name__) + + +class WanCausalDMDPipeline(LoRAPipeline, ComposedPipelineBase): + + _required_config_modules = [ + "text_encoder", "tokenizer", "vae", "transformer", "scheduler" + ] + + def initialize_pipeline(self, fastvideo_args: FastVideoArgs): + self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler( + shift=fastvideo_args.pipeline_config.flow_shift) + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: + """Set up pipeline stages with proper dependency injection.""" + + self.add_stage(stage_name="input_validation_stage", + stage=InputValidationStage()) + + self.add_stage(stage_name="prompt_encoding_stage", + stage=TextEncodingStage( + text_encoders=[self.get_module("text_encoder")], + tokenizers=[self.get_module("tokenizer")], + )) + + self.add_stage(stage_name="conditioning_stage", + stage=ConditioningStage()) + + self.add_stage(stage_name="timestep_preparation_stage", + stage=TimestepPreparationStage( + scheduler=self.get_module("scheduler"))) + + self.add_stage(stage_name="latent_preparation_stage", + stage=LatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer", None))) + + self.add_stage(stage_name="denoising_stage", + stage=CausalDMDDenosingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"))) + + self.add_stage(stage_name="decoding_stage", + stage=DecodingStage(vae=self.get_module("vae"))) + + +EntryClass = WanCausalDMDPipeline + + diff --git a/fastvideo/pipelines/basic/wan/wan_causal_pipeline.py b/fastvideo/pipelines/basic/wan/wan_causal_pipeline.py new file mode 100644 index 000000000..35983417c --- /dev/null +++ b/fastvideo/pipelines/basic/wan/wan_causal_pipeline.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Wan causal pipeline implementation. + +This module contains an implementation of the Wan causal pipeline +using the modular pipeline architecture. +""" + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler) +from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline + +# isort: off +from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, + CausalDenoisingStage, + InputValidationStage, + LatentPreparationStage, + TextEncodingStage, + TimestepPreparationStage) +# isort: on + +logger = init_logger(__name__) + + +class WanCausalPipeline(LoRAPipeline, ComposedPipelineBase): + + _required_config_modules = [ + "text_encoder", "tokenizer", "vae", "transformer", "scheduler" + ] + + def initialize_pipeline(self, fastvideo_args: FastVideoArgs): + + self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler( + shift=fastvideo_args.pipeline_config.flow_shift) + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: + """Set up pipeline stages with proper dependency injection.""" + + self.add_stage(stage_name="input_validation_stage", + stage=InputValidationStage()) + + self.add_stage(stage_name="prompt_encoding_stage", + stage=TextEncodingStage( + text_encoders=[self.get_module("text_encoder")], + tokenizers=[self.get_module("tokenizer")], + )) + + self.add_stage(stage_name="conditioning_stage", + stage=ConditioningStage()) + + self.add_stage(stage_name="timestep_preparation_stage", + stage=TimestepPreparationStage( + scheduler=self.get_module("scheduler"))) + + self.add_stage(stage_name="latent_preparation_stage", + stage=LatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer", None))) + + self.add_stage(stage_name="denoising_stage", + stage=CausalDenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"))) + + self.add_stage(stage_name="decoding_stage", + stage=DecodingStage(vae=self.get_module("vae"))) + + +EntryClass = WanCausalPipeline diff --git a/fastvideo/pipelines/stages/__init__.py b/fastvideo/pipelines/stages/__init__.py index 1852edc49..29571910b 100644 --- a/fastvideo/pipelines/stages/__init__.py +++ b/fastvideo/pipelines/stages/__init__.py @@ -11,6 +11,8 @@ from fastvideo.pipelines.stages.conditioning import ConditioningStage from fastvideo.pipelines.stages.decoding import DecodingStage from fastvideo.pipelines.stages.denoising import (DenoisingStage, DmdDenoisingStage) +from fastvideo.pipelines.stages.causal_denoising import ( + CausalDMDDenosingStage, CausalDenoisingStage) from fastvideo.pipelines.stages.encoding import EncodingStage from fastvideo.pipelines.stages.image_encoding import (ImageEncodingStage, ImageVAEEncodingStage) @@ -30,6 +32,8 @@ __all__ = [ "ConditioningStage", "DenoisingStage", "DmdDenoisingStage", + "CausalDenoisingStage", + "CausalDMDDenosingStage", "EncodingStage", "DecodingStage", "ImageEncodingStage", diff --git a/fastvideo/pipelines/stages/causal_denoising.py b/fastvideo/pipelines/stages/causal_denoising.py new file mode 100644 index 000000000..07289ab75 --- /dev/null +++ b/fastvideo/pipelines/stages/causal_denoising.py @@ -0,0 +1,601 @@ +import torch # type: ignore + +from fastvideo.distributed import get_local_torch_device +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.forward_context import set_forward_context +from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler) +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.denoising import DenoisingStage + +try: + from fastvideo.attention.backends.sliding_tile_attn import ( + SlidingTileAttentionBackend) + st_attn_available = True +except ImportError: + st_attn_available = False + SlidingTileAttentionBackend = None # type: ignore + +try: + from fastvideo.attention.backends.video_sparse_attn import ( + VideoSparseAttentionBackend) + vsa_available = True +except ImportError: + vsa_available = False + VideoSparseAttentionBackend = None # type: ignore + +# class CausalDenoisingStage(LoRAPipeline, ComposedPipelineBase): +# """ +# Flow Causal denoising stage. +# """ + +# def forward( +# self, +# batch: ForwardBatch, +# fastvideo_args: FastVideoArgs, +# ) -> ForwardBatch: +# # Setup precision and autocast settings +# target_dtype = torch.bfloat16 +# autocast_enabled = (target_dtype != torch.float32 +# ) and not fastvideo_args.disable_autocast + +# # Gather args for causal block-wise processing +# try: +# num_frame_per_block = getattr(fastvideo_args.pipeline_config, +# "num_frame_per_block", 1) +# except Exception: +# num_frame_per_block = 1 +# try: +# independent_first_frame = getattr(fastvideo_args.pipeline_config, +# "independent_first_frame", False) +# except Exception: +# independent_first_frame = False + +# # Required fields +# timesteps = batch.timesteps +# if timesteps is None: +# raise ValueError("Timesteps must be provided") + +# # Prepare image latents and embeddings for I2V generation +# image_embeds = batch.image_embeds +# if len(image_embeds) > 0: +# assert torch.isnan(image_embeds[0]).sum() == 0 +# image_embeds = [ +# image_embed.to(target_dtype) for image_embed in image_embeds +# ] + +# image_kwargs = self.prepare_extra_func_kwargs( +# self.transformer.forward, +# { +# "encoder_hidden_states_image": image_embeds, +# "mask_strategy": dict_to_3d_list( +# None, t_max=50, l_max=60, h_max=24) +# }, +# ) + +# pos_cond_kwargs = self.prepare_extra_func_kwargs( +# self.transformer.forward, +# { +# "encoder_hidden_states_2": batch.clip_embedding_pos, +# "encoder_attention_mask": batch.prompt_attention_mask, +# }, +# ) + +# neg_cond_kwargs = self.prepare_extra_func_kwargs( +# self.transformer.forward, +# { +# "encoder_hidden_states_2": batch.clip_embedding_neg, +# "encoder_attention_mask": batch.negative_attention_mask, +# }, +# ) + +# # Prepare STA parameters +# if st_attn_available and self.attn_backend == SlidingTileAttentionBackend: +# self.prepare_sta_param(batch, fastvideo_args) + +# # Get latents and embeddings +# latents = batch.latents +# assert latents is not None, "latents must be provided" +# prompt_embeds = batch.prompt_embeds +# assert torch.isnan(prompt_embeds[0]).sum() == 0 + +# # Determine blocks +# b, c, t, h, w = latents.shape +# if not independent_first_frame or (independent_first_frame +# and batch.image_latent is not None): +# if t % num_frame_per_block != 0: +# raise ValueError( +# "num_frames must be divisible by num_frame_per_block for causal denoising" +# ) +# num_blocks = t // num_frame_per_block +# block_sizes = [num_frame_per_block] * num_blocks +# start_index = 0 +# else: +# if (t - 1) % num_frame_per_block != 0: +# raise ValueError( +# "(num_frames - 1) must be divisible by num_frame_per_block when independent_first_frame=True" +# ) +# num_blocks = (t - 1) // num_frame_per_block +# block_sizes = [1] + [num_frame_per_block] * num_blocks +# start_index = 0 + +# # Run causal block-wise denoising loop +# with self.progress_bar(total=len(block_sizes) * +# len(timesteps)) as progress_bar: +# for current_num_frames in block_sizes: +# # Build the latent input for this block +# current_latents = latents[:, :, start_index:start_index + +# current_num_frames, :, :] + +# # Run spatial denoising loop over timesteps +# for i, t_cur in enumerate(timesteps): +# # Expand latents for I2V +# latent_model_input = current_latents.to(target_dtype) +# if batch.image_latent is not None and independent_first_frame and start_index == 0: +# # Prepend the image latent only for the first block when required +# latent_model_input = torch.cat([ +# batch.image_latent.to(target_dtype), +# latent_model_input +# ], +# dim=2) + +# latent_model_input = self.scheduler.scale_model_input( +# latent_model_input, t_cur) + +# # Prepare inputs for transformer +# t_expand = t_cur.repeat(latent_model_input.shape[0]) +# guidance_expand = (torch.tensor( +# [fastvideo_args.pipeline_config.embedded_cfg_scale] * +# latent_model_input.shape[0], +# dtype=torch.float32, +# device=get_local_torch_device(), +# ).to(target_dtype) * +# 1000.0 if fastvideo_args.pipeline_config. +# embedded_cfg_scale is not None else None) + +# # Predict noise residual +# with torch.autocast(device_type="cuda", +# dtype=target_dtype, +# enabled=autocast_enabled): +# if (st_attn_available and self.attn_backend == SlidingTileAttentionBackend) or \ +# (vsa_available and self.attn_backend == VideoSparseAttentionBackend): +# self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls( +# ) +# if self.attn_metadata_builder_cls is not None: +# self.attn_metadata_builder = self.attn_metadata_builder_cls( +# ) +# attn_metadata = self.attn_metadata_builder.build( # type: ignore +# current_timestep=i, # type: ignore +# raw_latent_shape=(current_num_frames, h, +# w), # type: ignore +# patch_size=fastvideo_args.pipeline_config. +# dit_config.patch_size, # type: ignore +# STA_param=batch.STA_param, # type: ignore +# VSA_sparsity=fastvideo_args. +# VSA_sparsity, # type: ignore +# device=get_local_torch_device(), +# ) +# assert attn_metadata is not None, "attn_metadata cannot be None" +# else: +# attn_metadata = None +# else: +# attn_metadata = None + +# batch.is_cfg_negative = False +# with set_forward_context(current_timestep=i, +# attn_metadata=attn_metadata, +# forward_batch=batch): +# noise_pred = self.transformer( +# latent_model_input, +# prompt_embeds, +# t_expand, +# guidance=guidance_expand, +# **image_kwargs, +# **pos_cond_kwargs, +# ) + +# if batch.do_classifier_free_guidance: +# batch.is_cfg_negative = True +# with set_forward_context( +# current_timestep=i, +# attn_metadata=attn_metadata, +# forward_batch=batch): +# noise_pred_uncond = self.transformer( +# latent_model_input, +# batch.negative_prompt_embeds, +# t_expand, +# guidance=guidance_expand, +# **image_kwargs, +# **neg_cond_kwargs, +# ) +# noise_pred_text = noise_pred +# noise_pred = noise_pred_uncond + batch.guidance_scale * ( +# noise_pred_text - noise_pred_uncond) + +# if batch.guidance_rescale > 0.0: +# noise_pred = self.rescale_noise_cfg( +# noise_pred, +# noise_pred_text, +# guidance_rescale=batch.guidance_rescale, +# ) + +# # Compute the previous noisy sample for this block only +# current_latents = self.scheduler.step( +# noise_pred, +# t_cur, +# current_latents, +# return_dict=False)[0] + +# if progress_bar is not None: +# progress_bar.update() + +# # After finishing timesteps, write denoised block back +# latents[:, :, start_index:start_index + +# current_num_frames, :, :] = current_latents + +# # Advance to next block +# start_index += current_num_frames + +# # Update batch with final latents +# batch.latents = latents +# return batch + + +class CausalDMDDenosingStage(DenoisingStage): + """ + Denoising stage for causal diffusion. + """ + + def __init__(self, transformer, scheduler) -> None: + super().__init__(transformer, scheduler) + self.scheduler = FlowMatchEulerDiscreteScheduler(shift=8.0) + # KV and cross-attention cache state (initialized on first forward) + self.kv_cache1 = None + self.crossattn_cache = None + # Model-dependent constants (aligned with causal_inference.py assumptions) + self.num_transformer_blocks = 30 + self.frame_seq_length = 1560 + try: + self.local_attn_size = getattr(self.transformer.model, + "local_attn_size", + -1) # type: ignore + except Exception: + self.local_attn_size = -1 + + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + target_dtype = torch.bfloat16 + autocast_enabled = (target_dtype != torch.float32 + ) and not fastvideo_args.disable_autocast + + # Args + try: + num_frame_per_block = getattr(fastvideo_args.pipeline_config, + "num_frame_per_block", 1) + except Exception: + num_frame_per_block = 1 + try: + independent_first_frame = getattr(fastvideo_args.pipeline_config, + "independent_first_frame", False) + except Exception: + independent_first_frame = False + + # Timesteps for DMD + timesteps = torch.tensor( + fastvideo_args.pipeline_config.dmd_denoising_steps, + dtype=torch.long, + device=get_local_torch_device()) + + # Image kwargs (kept empty unless caller provides compatible args) + image_kwargs = {} + + pos_cond_kwargs = self.prepare_extra_func_kwargs( + self.transformer.forward, + { + # "encoder_hidden_states_2": batch.clip_embedding_pos, + "encoder_attention_mask": batch.prompt_attention_mask, + }, + ) + + # STA + if st_attn_available and self.attn_backend == SlidingTileAttentionBackend: + self.prepare_sta_param(batch, fastvideo_args) + + # Latents and prompts + assert batch.latents is not None, "latents must be provided" + latents = batch.latents # [B, C, T, H, W] + b, c, t, h, w = latents.shape + prompt_embeds = batch.prompt_embeds + assert torch.isnan(prompt_embeds[0]).sum() == 0 + + # Initialize or reset caches + if self.kv_cache1 is None: + self._initialize_kv_cache(batch_size=latents.shape[0], + dtype=latents.dtype, + device=latents.device) + self._initialize_crossattn_cache(batch_size=latents.shape[0], + dtype=latents.dtype, + device=latents.device) + else: + # reset cross-attention cache + for block_index in range(self.num_transformer_blocks): + self.crossattn_cache[block_index][ + "is_init"] = False # type: ignore + # reset kv cache pointers + for block_index in range(len(self.kv_cache1)): + self.kv_cache1[block_index][ + "global_end_index"] = torch.tensor( # type: ignore + [0], + dtype=torch.long, + device=latents.device) + self.kv_cache1[block_index][ + "local_end_index"] = torch.tensor( # type: ignore + [0], + dtype=torch.long, + device=latents.device) + + # Optional: cache context features from provided image latents prior to generation + current_start_frame = 0 + if getattr(batch, "image_latent", None) is not None: + image_latent = batch.image_latent + assert image_latent is not None + input_frames = image_latent.shape[2] + # timestep zero (or configured context noise) for cache warm-up + t_zero = torch.zeros([latents.shape[0]], + device=latents.device, + dtype=torch.long) + if independent_first_frame and input_frames >= 1: + # warm-up with the very first frame independently + image_first_btchw = image_latent[:, :, :1, :, :].to( + target_dtype).permute(0, 2, 1, 3, 4) + with torch.autocast(device_type="cuda", + dtype=target_dtype, + enabled=autocast_enabled): + _ = self.transformer( + image_first_btchw, + prompt_embeds, + t_zero, + kv_cache=self.kv_cache1, + crossattn_cache=self.crossattn_cache, + current_start=current_start_frame * + self.frame_seq_length, + **image_kwargs, + **pos_cond_kwargs, + ) + current_start_frame += 1 + remaining_frames = input_frames - 1 + else: + remaining_frames = input_frames + + # process remaining input frames in blocks of num_frame_per_block + while remaining_frames > 0: + block = min(num_frame_per_block, remaining_frames) + ref_btchw = image_latent[:, :, current_start_frame: + current_start_frame + + block, :, :].to(target_dtype).permute( + 0, 2, 1, 3, 4) + with torch.autocast(device_type="cuda", + dtype=target_dtype, + enabled=autocast_enabled): + _ = self.transformer( + ref_btchw, + prompt_embeds, + t_zero, + kv_cache=self.kv_cache1, + crossattn_cache=self.crossattn_cache, + current_start=current_start_frame * + self.frame_seq_length, + **image_kwargs, + **pos_cond_kwargs, + ) + current_start_frame += block + remaining_frames -= block + + # Base position offset from any cache warm-up + pos_start_base = current_start_frame + + # Determine block sizes + if not independent_first_frame or (independent_first_frame + and batch.image_latent is not None): + if t % num_frame_per_block != 0: + raise ValueError( + "num_frames must be divisible by num_frame_per_block for causal DMD denoising" + ) + num_blocks = t // num_frame_per_block + block_sizes = [num_frame_per_block] * num_blocks + start_index = 0 + else: + if (t - 1) % num_frame_per_block != 0: + raise ValueError( + "(num_frames - 1) must be divisible by num_frame_per_block when independent_first_frame=True" + ) + num_blocks = (t - 1) // num_frame_per_block + block_sizes = [1] + [num_frame_per_block] * num_blocks + start_index = 0 + + # DMD loop in causal blocks + with self.progress_bar(total=len(block_sizes) * + len(timesteps)) as progress_bar: + for current_num_frames in block_sizes: + current_latents = latents[:, :, start_index:start_index + + current_num_frames, :, :] + # use BTCHW for DMD conversion routines + noise_latents_btchw = current_latents.permute(0, 2, 1, 3, 4) + video_raw_latent_shape = noise_latents_btchw.shape + + for i, t_cur in enumerate(timesteps): + # Copy for pred conversion + noise_latents = noise_latents_btchw.clone() + latent_model_input = current_latents.to(target_dtype) + + if batch.image_latent is not None and independent_first_frame and start_index == 0: + latent_model_input = torch.cat([ + latent_model_input, + batch.image_latent.to(target_dtype) + ], + dim=2) + + # Prepare inputs + t_expand = t_cur.repeat(latent_model_input.shape[0]) + + # Attention metadata if needed + if (vsa_available and self.attn_backend + == VideoSparseAttentionBackend): + self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls( + ) + if self.attn_metadata_builder_cls is not None: + self.attn_metadata_builder = self.attn_metadata_builder_cls( + ) + attn_metadata = self.attn_metadata_builder.build( # type: ignore + current_timestep=i, # type: ignore + raw_latent_shape=(current_num_frames, h, + w), # type: ignore + patch_size=fastvideo_args.pipeline_config. + dit_config.patch_size, # type: ignore + STA_param=batch.STA_param, # type: ignore + VSA_sparsity=fastvideo_args. + VSA_sparsity, # type: ignore + device=get_local_torch_device(), # type: ignore + ) # type: ignore + assert attn_metadata is not None, "attn_metadata cannot be None" + else: + attn_metadata = None + else: + attn_metadata = None + + with torch.autocast(device_type="cuda", + dtype=target_dtype, + enabled=autocast_enabled): + with set_forward_context(current_timestep=i, + attn_metadata=attn_metadata, + forward_batch=batch): + # Run transformer; follow DMD stage pattern + pred_noise_btchw = self.transformer( + latent_model_input.permute(0, 2, 1, 3, 4), + prompt_embeds, + t_expand, + kv_cache=self.kv_cache1, + crossattn_cache=self.crossattn_cache, + current_start=(pos_start_base + start_index) * + self.frame_seq_length, + **image_kwargs, + **pos_cond_kwargs, + ).permute(0, 2, 1, 3, 4) + + # Convert pred noise to pred video with FM Euler scheduler utilities + from fastvideo.training.training_utils import ( + pred_noise_to_pred_video) + pred_video_btchw = pred_noise_to_pred_video( + pred_noise=pred_noise_btchw.flatten(0, 1), + noise_input_latent=noise_latents.flatten(0, 1), + timestep=t_expand, + scheduler=self.scheduler).unflatten( + 0, pred_noise_btchw.shape[:2]) + + if i < len(timesteps) - 1: + next_timestep = timesteps[i + 1] * torch.ones( + [1], + dtype=torch.long, + device=pred_video_btchw.device) + noise = torch.randn( + video_raw_latent_shape, + dtype=pred_video_btchw.dtype, + generator=(batch.generator[0] if isinstance( + batch.generator, list) else + batch.generator)).to(self.device) + noise_btchw = noise + current_latents_btchw = self.scheduler.add_noise( + pred_video_btchw.flatten(0, 1), + noise_btchw.flatten(0, 1), + next_timestep).unflatten(0, + pred_video_btchw.shape[:2]) + current_latents = current_latents_btchw.permute( + 0, 2, 1, 3, 4) + else: + current_latents = pred_video_btchw.permute( + 0, 2, 1, 3, 4) + + if progress_bar is not None: + progress_bar.update() + + # Write back and advance + latents[:, :, start_index:start_index + + current_num_frames, :, :] = current_latents + + # Re-run with context timestep to update KV cache using clean context + context_noise = getattr(fastvideo_args.pipeline_config, + "context_noise", 0) + t_context = torch.ones([latents.shape[0]], + device=latents.device, + dtype=torch.long) * int(context_noise) + context_btchw = current_latents.to(target_dtype).permute( + 0, 2, 1, 3, 4) + with torch.autocast(device_type="cuda", + dtype=target_dtype, + enabled=autocast_enabled): + _ = self.transformer( + context_btchw, + prompt_embeds, + t_context, + kv_cache=self.kv_cache1, + crossattn_cache=self.crossattn_cache, + current_start=(pos_start_base + start_index) * + self.frame_seq_length, + **image_kwargs, + **pos_cond_kwargs, + ) + start_index += current_num_frames + + batch.latents = latents + return batch + + def _initialize_kv_cache(self, batch_size, dtype, device): + """ + Initialize a Per-GPU KV cache aligned with the Wan model assumptions. + """ + kv_cache1 = [] + if self.local_attn_size != -1: + kv_cache_size = self.local_attn_size * self.frame_seq_length + else: + kv_cache_size = 32760 + + for _ in range(self.num_transformer_blocks): + kv_cache1.append({ + "k": + torch.zeros([batch_size, kv_cache_size, 12, 128], + dtype=dtype, + device=device), + "v": + torch.zeros([batch_size, kv_cache_size, 12, 128], + dtype=dtype, + device=device), + "global_end_index": + torch.tensor([0], dtype=torch.long, device=device), + "local_end_index": + torch.tensor([0], dtype=torch.long, device=device), + }) + + self.kv_cache1 = kv_cache1 + + def _initialize_crossattn_cache(self, batch_size, dtype, device): + """ + Initialize a Per-GPU cross-attention cache aligned with the Wan model assumptions. + """ + crossattn_cache = [] + for _ in range(self.num_transformer_blocks): + crossattn_cache.append({ + "k": + torch.zeros([batch_size, 512, 12, 128], + dtype=dtype, + device=device), + "v": + torch.zeros([batch_size, 512, 12, 128], + dtype=dtype, + device=device), + "is_init": + False, + }) + self.crossattn_cache = crossattn_cache diff --git a/scripts/checkpoint_conversion/wan_to_diffusers.py b/scripts/checkpoint_conversion/wan_to_diffusers.py new file mode 100644 index 000000000..11407039e --- /dev/null +++ b/scripts/checkpoint_conversion/wan_to_diffusers.py @@ -0,0 +1,88 @@ +from huggingface_hub import save_torch_state_dict, load_state_dict_from_file +from safetensors.torch import save_file +import torch +import re +from collections import OrderedDict + +_param_names_mapping: dict = { + r"^text_embedding\.0\.(.*)$": + r"condition_embedder.text_embedder.linear_1.\1", + r"^text_embedding\.2\.(.*)$": + r"condition_embedder.text_embedder.linear_2.\1", + r"^time_embedding\.0\.(.*)$": + r"condition_embedder.time_embedder.linear_1.\1", + r"^time_embedding\.2\.(.*)$": + r"condition_embedder.time_embedder.linear_2.\1", + r"^time_projection\.1\.(.*)$": + r"condition_embedder.time_proj.\1", + r"^img_emb\.proj\.0\.(.*)$": + r"condition_embedder.image_embedder.norm1.\1", + r"^img_emb\.proj\.1\.(.*)$": + r"condition_embedder.image_embedder.ff.net.0.proj.\1", + r"^img_emb\.proj\.3\.(.*)$": + r"condition_embedder.image_embedder.ff.net.2.\1", + r"^img_emb\.proj\.4\.(.*)$": + r"condition_embedder.image_embedder.norm2.\1", + r"^head\.modulation": + r"scale_shift_table", + r"^head\.head\.(.*)$": + r"proj_out.\1", + r"^blocks\.(\d+)\.self_attn\.q\.(.*)$": + r"blocks.\1.attn1.to_q.\2", + r"^blocks\.(\d+)\.self_attn\.k\.(.*)$": + r"blocks.\1.attn1.to_k.\2", + r"^blocks\.(\d+)\.self_attn\.v\.(.*)$": + r"blocks.\1.attn1.to_v.\2", + r"^blocks\.(\d+)\.self_attn\.o\.(.*)$": + r"blocks.\1.attn1.to_out.0.\2", + r"^blocks\.(\d+)\.self_attn\.norm_q\.(.*)$": + r"blocks.\1.attn1.norm_q.\2", + r"^blocks\.(\d+)\.self_attn\.norm_k\.(.*)$": + r"blocks.\1.attn1.norm_k.\2", + r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$": + r"blocks.\1.attn2.to_q.\2", + r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$": + r"blocks.\1.attn2.to_k.\2", + r"^blocks\.(\d+)\.cross_attn\.k_img\.(.*)$": + r"blocks.\1.attn2.add_k_proj.\2", + r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$": + r"blocks.\1.attn2.to_v.\2", + r"^blocks\.(\d+)\.cross_attn\.v_img\.(.*)$": + r"blocks.\1.attn2.add_v_proj.\2", + r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$": + r"blocks.\1.attn2.to_out.0.\2", + r"^blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$": + r"blocks.\1.attn2.norm_q.\2", + r"^blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$": + r"blocks.\1.attn2.norm_k.\2", + r"^blocks\.(\d+)\.cross_attn\.norm_k_img\.(.*)$": + r"blocks.\1.attn2.norm_added_k.\2", + r"^blocks\.(\d+)\.ffn\.0\.(.*)$": + r"blocks.\1.ffn.net.0.proj.\2", + r"^blocks\.(\d+)\.ffn\.2\.(.*)$": + r"blocks.\1.ffn.net.2.\2", + r"^blocks\.(\d+)\.modulation": + r"blocks.\1.scale_shift_table", + r"^blocks\.(\d+)\.norm3\.(.*)$": + r"blocks.\1.norm2.\2", +} + +state_dict = load_state_dict_from_file("FastVideo/diffusion_pytorch_model.safetensors") +new_state_dict = OrderedDict() +for k, v in state_dict.items(): + new_key = k + for pattern, replacement in _param_names_mapping.items(): + if re.match(pattern, k): + new_key = re.sub(pattern, replacement, k) + break # Stop at the first match + new_state_dict[new_key] = v + if "norm_added_k" in new_key: + dummy_key = new_key.replace("norm_added_k", "norm_added_q") + dummy_value = torch.zeros_like(v) + new_state_dict[dummy_key] = dummy_value +del state_dict + +save_torch_state_dict( + new_state_dict, + "Wan2.1-Fun-1.3B-InP-Diffusers/transformer" +) \ No newline at end of file From b198ba56076a5ce7df61bc2fb9de65e8fd9c077f Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Wed, 27 Aug 2025 02:12:43 -0700 Subject: [PATCH 08/13] checkpoint --- ...causal.py => basic_self_forcing_causal.py} | 0 fastvideo/configs/pipelines/registry.py | 2 + fastvideo/configs/pipelines/wan.py | 8 +++ fastvideo/configs/sample/registry.py | 17 ++++- fastvideo/configs/sample/wan.py | 9 +++ fastvideo/models/loader/fsdp_load.py | 1 + fastvideo/models/registry.py | 2 + .../basic/wan/wan_causal_pipeline.py | 71 ------------------- fastvideo/pipelines/pipeline_registry.py | 1 + fastvideo/pipelines/stages/__init__.py | 4 +- .../checkpoint_conversion/wan_to_diffusers.py | 7 +- scripts/huggingface/upload_hf.py | 4 +- 12 files changed, 46 insertions(+), 80 deletions(-) rename examples/inference/basic/{basic_causal.py => basic_self_forcing_causal.py} (100%) delete mode 100644 fastvideo/pipelines/basic/wan/wan_causal_pipeline.py diff --git a/examples/inference/basic/basic_causal.py b/examples/inference/basic/basic_self_forcing_causal.py similarity index 100% rename from examples/inference/basic/basic_causal.py rename to examples/inference/basic/basic_self_forcing_causal.py diff --git a/fastvideo/configs/pipelines/registry.py b/fastvideo/configs/pipelines/registry.py index cd8fb9938..93628a75d 100644 --- a/fastvideo/configs/pipelines/registry.py +++ b/fastvideo/configs/pipelines/registry.py @@ -9,6 +9,7 @@ from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig from fastvideo.configs.pipelines.wan import (FastWan2_1_T2V_480P_Config, FastWan2_2_TI2V_5B_Config, + SelfForcingWanT2V480PConfig, WanI2V480PConfig, WanI2V720PConfig, WanT2V480PConfig, WanT2V720PConfig) from fastvideo.logger import init_logger @@ -34,6 +35,7 @@ PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = { "Wan-AI/Wan2.2-TI2V-5B-Diffusers": WanT2V720PConfig, "Wan-AI/Wan2.2-T2V-A14B-Diffusers": WanT2V480PConfig, "Wan-AI/Wan2.2-I2V-A14B-Diffusers": WanI2V480PConfig, + "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers": SelfForcingWanT2V480PConfig, # Add other specific weight variants } diff --git a/fastvideo/configs/pipelines/wan.py b/fastvideo/configs/pipelines/wan.py index 317127f66..579da6dfd 100644 --- a/fastvideo/configs/pipelines/wan.py +++ b/fastvideo/configs/pipelines/wan.py @@ -138,3 +138,11 @@ class Wan2_2_T2V_A14B_Config(WanT2V480PConfig): @dataclass class Wan2_2_I2V_A14B_Config(WanT2V480PConfig): pass + + +# ============================================= +# ============= Causal Self-Forcing ============= +# ============================================= +@dataclass +class SelfForcingWanT2V480PConfig(WanT2V480PConfig): + is_causal: bool = True diff --git a/fastvideo/configs/sample/registry.py b/fastvideo/configs/sample/registry.py index cb646bff7..32f1820fb 100644 --- a/fastvideo/configs/sample/registry.py +++ b/fastvideo/configs/sample/registry.py @@ -17,6 +17,7 @@ from fastvideo.configs.sample.wan import ( WanI2V_14B_720P_SamplingParam, WanT2V_1_3B_SamplingParam, WanT2V_14B_SamplingParam, + SelfForcingWanT2V480PConfig, ) # isort: on from fastvideo.logger import init_logger @@ -28,17 +29,29 @@ logger = init_logger(__name__) SAMPLING_PARAM_REGISTRY: dict[str, Any] = { "FastVideo/FastHunyuan-diffusers": FastHunyuanSamplingParam, "hunyuanvideo-community/HunyuanVideo": HunyuanSamplingParam, + "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VSamplingParam, + + # Wan2.1 "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V_1_3B_SamplingParam, "Wan-AI/Wan2.1-T2V-14B-Diffusers": WanT2V_14B_SamplingParam, "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V_14B_480P_SamplingParam, "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V_14B_720P_SamplingParam, - "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VSamplingParam, - "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": FastWanT2V480PConfig, + + # Wan2.2 "Wan-AI/Wan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam, "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers": Wan2_2_TI2V_5B_SamplingParam, "Wan-AI/Wan2.2-T2V-A14B-Diffusers": Wan2_2_T2V_A14B_SamplingParam, "Wan-AI/Wan2.2-I2V-A14B-Diffusers": Wan2_2_I2V_A14B_SamplingParam, + + # FastWan2.1 + "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": FastWanT2V480PConfig, + + # FastWan2.2 + "FastVideo/FastWan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam, + + # Causal Self-Forcing Wan2.1 + "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers": SelfForcingWanT2V480PConfig, # Add other specific weight variants } diff --git a/fastvideo/configs/sample/wan.py b/fastvideo/configs/sample/wan.py index 23cb0cdbf..4c435fdc8 100644 --- a/fastvideo/configs/sample/wan.py +++ b/fastvideo/configs/sample/wan.py @@ -141,3 +141,12 @@ class Wan2_2_I2V_A14B_SamplingParam(Wan2_2_Base_SamplingParam): guidance_scale_2: float = 3.5 num_inference_steps: int = 40 fps: int = 16 + + +# ============================================= +# ============= Causal Self-Forcing ============= +# ============================================= +@dataclass +class SelfForcingWanT2V480PConfig(WanT2V_1_3B_SamplingParam): + num_frames_per_block: int = 3 + sliding_window_num_frames: int = 21 diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index 117b1a728..72763c18b 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -248,6 +248,7 @@ def load_model_from_full_model_state_dict( sharded_sd = {} custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict( full_sd_iterator, param_names_mapping) # type: ignore + print(custom_param_sd.keys()) for target_param_name, full_tensor in custom_param_sd.items(): meta_sharded_param = meta_sd.get(target_param_name) if meta_sharded_param is None: diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index 5c585fa16..68be9bc13 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -25,12 +25,14 @@ _TEXT_TO_VIDEO_DIT_MODELS = { "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoTransformer3DModel"), "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"), + "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"), "StepVideoModel": ("dits", "stepvideo", "StepVideoModel") } _IMAGE_TO_VIDEO_DIT_MODELS = { # "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoDiT"), "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"), + "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"), } _TEXT_ENCODER_MODELS = { diff --git a/fastvideo/pipelines/basic/wan/wan_causal_pipeline.py b/fastvideo/pipelines/basic/wan/wan_causal_pipeline.py deleted file mode 100644 index 35983417c..000000000 --- a/fastvideo/pipelines/basic/wan/wan_causal_pipeline.py +++ /dev/null @@ -1,71 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -Wan causal pipeline implementation. - -This module contains an implementation of the Wan causal pipeline -using the modular pipeline architecture. -""" - -from fastvideo.fastvideo_args import FastVideoArgs -from fastvideo.logger import init_logger -from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( - FlowMatchEulerDiscreteScheduler) -from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline - -# isort: off -from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, - CausalDenoisingStage, - InputValidationStage, - LatentPreparationStage, - TextEncodingStage, - TimestepPreparationStage) -# isort: on - -logger = init_logger(__name__) - - -class WanCausalPipeline(LoRAPipeline, ComposedPipelineBase): - - _required_config_modules = [ - "text_encoder", "tokenizer", "vae", "transformer", "scheduler" - ] - - def initialize_pipeline(self, fastvideo_args: FastVideoArgs): - - self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler( - shift=fastvideo_args.pipeline_config.flow_shift) - - def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: - """Set up pipeline stages with proper dependency injection.""" - - self.add_stage(stage_name="input_validation_stage", - stage=InputValidationStage()) - - self.add_stage(stage_name="prompt_encoding_stage", - stage=TextEncodingStage( - text_encoders=[self.get_module("text_encoder")], - tokenizers=[self.get_module("tokenizer")], - )) - - self.add_stage(stage_name="conditioning_stage", - stage=ConditioningStage()) - - self.add_stage(stage_name="timestep_preparation_stage", - stage=TimestepPreparationStage( - scheduler=self.get_module("scheduler"))) - - self.add_stage(stage_name="latent_preparation_stage", - stage=LatentPreparationStage( - scheduler=self.get_module("scheduler"), - transformer=self.get_module("transformer", None))) - - self.add_stage(stage_name="denoising_stage", - stage=CausalDenoisingStage( - transformer=self.get_module("transformer"), - scheduler=self.get_module("scheduler"))) - - self.add_stage(stage_name="decoding_stage", - stage=DecodingStage(vae=self.get_module("vae"))) - - -EntryClass = WanCausalPipeline diff --git a/fastvideo/pipelines/pipeline_registry.py b/fastvideo/pipelines/pipeline_registry.py index d80fd40b5..9955c3729 100644 --- a/fastvideo/pipelines/pipeline_registry.py +++ b/fastvideo/pipelines/pipeline_registry.py @@ -21,6 +21,7 @@ _PIPELINE_NAME_TO_ARCHITECTURE_NAME: dict[str, str] = { "WanPipeline": "wan", "WanDMDPipeline": "wan", "WanImageToVideoPipeline": "wan", + "WanCausalDMDPipeline": "wan", "StepVideoPipeline": "stepvideo", "HunyuanVideoPipeline": "hunyuan", } diff --git a/fastvideo/pipelines/stages/__init__.py b/fastvideo/pipelines/stages/__init__.py index 29571910b..9566b426f 100644 --- a/fastvideo/pipelines/stages/__init__.py +++ b/fastvideo/pipelines/stages/__init__.py @@ -7,12 +7,11 @@ complete diffusion pipelines. """ from fastvideo.pipelines.stages.base import PipelineStage +from fastvideo.pipelines.stages.causal_denoising import CausalDMDDenosingStage from fastvideo.pipelines.stages.conditioning import ConditioningStage from fastvideo.pipelines.stages.decoding import DecodingStage from fastvideo.pipelines.stages.denoising import (DenoisingStage, DmdDenoisingStage) -from fastvideo.pipelines.stages.causal_denoising import ( - CausalDMDDenosingStage, CausalDenoisingStage) from fastvideo.pipelines.stages.encoding import EncodingStage from fastvideo.pipelines.stages.image_encoding import (ImageEncodingStage, ImageVAEEncodingStage) @@ -32,7 +31,6 @@ __all__ = [ "ConditioningStage", "DenoisingStage", "DmdDenoisingStage", - "CausalDenoisingStage", "CausalDMDDenosingStage", "EncodingStage", "DecodingStage", diff --git a/scripts/checkpoint_conversion/wan_to_diffusers.py b/scripts/checkpoint_conversion/wan_to_diffusers.py index 11407039e..1fec923db 100644 --- a/scripts/checkpoint_conversion/wan_to_diffusers.py +++ b/scripts/checkpoint_conversion/wan_to_diffusers.py @@ -1,4 +1,5 @@ from huggingface_hub import save_torch_state_dict, load_state_dict_from_file +# from safetensors import safetensors from safetensors.torch import save_file import torch import re @@ -67,7 +68,8 @@ _param_names_mapping: dict = { r"blocks.\1.norm2.\2", } -state_dict = load_state_dict_from_file("FastVideo/diffusion_pytorch_model.safetensors") +state_dict = load_state_dict_from_file("checkpoints/self_forcing_dmd.pt") +# state_dict = state_dict["generator_ema"] new_state_dict = OrderedDict() for k, v in state_dict.items(): new_key = k @@ -84,5 +86,6 @@ del state_dict save_torch_state_dict( new_state_dict, - "Wan2.1-Fun-1.3B-InP-Diffusers/transformer" + "new_diffusers/", + max_shard_size="10GB" ) \ No newline at end of file diff --git a/scripts/huggingface/upload_hf.py b/scripts/huggingface/upload_hf.py index 1e4ce88b7..9bbad77f1 100644 --- a/scripts/huggingface/upload_hf.py +++ b/scripts/huggingface/upload_hf.py @@ -3,7 +3,7 @@ from huggingface_hub import HfApi api = HfApi() api.upload_folder( - folder_path="Wan2.2-TI2V-5B-Diffusers", - repo_id="FastVideo/FastWan2.2-TI2V-5B-Diffusers", + folder_path="wow", + repo_id="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers", repo_type="model", ) From 157510296559ca3ba597c48fdb8bee6999e6e6b2 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Wed, 27 Aug 2025 02:56:31 -0700 Subject: [PATCH 09/13] debugging model --- fastvideo/models/loader/fsdp_load.py | 6 ++ .../checkpoint_conversion/wan_to_diffusers.py | 78 ++++++++++++++++++- 2 files changed, 80 insertions(+), 4 deletions(-) diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index 72763c18b..12e5786c2 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -249,6 +249,12 @@ def load_model_from_full_model_state_dict( custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict( full_sd_iterator, param_names_mapping) # type: ignore print(custom_param_sd.keys()) + print("--------------------------------") + print("--------------------------------") + print("--------------------------------") + print("--------------------------------") + print("--------------------------------") + print(meta_sd.keys()) for target_param_name, full_tensor in custom_param_sd.items(): meta_sharded_param = meta_sd.get(target_param_name) if meta_sharded_param is None: diff --git a/scripts/checkpoint_conversion/wan_to_diffusers.py b/scripts/checkpoint_conversion/wan_to_diffusers.py index 1fec923db..ea52ed687 100644 --- a/scripts/checkpoint_conversion/wan_to_diffusers.py +++ b/scripts/checkpoint_conversion/wan_to_diffusers.py @@ -68,15 +68,85 @@ _param_names_mapping: dict = { r"blocks.\1.norm2.\2", } +# The following mapping has an extra 'patch_embedding' field and also contains +# the 'model' prefixes +_self_forcing_to_diffusers_param_names_mapping: dict = { + r"^model.patch_embedding\.(.*)$": + r"patch_embedding.\1", + r"^model.text_embedding\.0\.(.*)$": + r"condition_embedder.text_embedder.linear_1.\1", + r"^model.text_embedding\.2\.(.*)$": + r"condition_embedder.text_embedder.linear_2.\1", + r"^model.time_embedding\.0\.(.*)$": + r"condition_embedder.time_embedder.linear_1.\1", + r"^model.time_embedding\.2\.(.*)$": + r"condition_embedder.time_embedder.linear_2.\1", + r"^model.time_projection\.1\.(.*)$": + r"condition_embedder.time_proj.\1", + r"^model.img_emb\.proj\.0\.(.*)$": + r"condition_embedder.image_embedder.norm1.\1", + r"^model.img_emb\.proj\.1\.(.*)$": + r"condition_embedder.image_embedder.ff.net.0.proj.\1", + r"^model.img_emb\.proj\.3\.(.*)$": + r"condition_embedder.image_embedder.ff.net.2.\1", + r"^model.img_emb\.proj\.4\.(.*)$": + r"condition_embedder.image_embedder.norm2.\1", + r"^model.head\.modulation": + r"scale_shift_table", + r"^model.head\.head\.(.*)$": + r"proj_out.\1", + r"^model.blocks\.(\d+)\.self_attn\.q\.(.*)$": + r"blocks.\1.attn1.to_q.\2", + r"^model.blocks\.(\d+)\.self_attn\.k\.(.*)$": + r"blocks.\1.attn1.to_k.\2", + r"^model.blocks\.(\d+)\.self_attn\.v\.(.*)$": + r"blocks.\1.attn1.to_v.\2", + r"^model.blocks\.(\d+)\.self_attn\.o\.(.*)$": + r"blocks.\1.attn1.to_out.0.\2", + r"^model.blocks\.(\d+)\.self_attn\.norm_q\.(.*)$": + r"blocks.\1.attn1.norm_q.\2", + r"^model.blocks\.(\d+)\.self_attn\.norm_k\.(.*)$": + r"blocks.\1.attn1.norm_k.\2", + r"^model.blocks\.(\d+)\.cross_attn\.q\.(.*)$": + r"blocks.\1.attn2.to_q.\2", + r"^model.blocks\.(\d+)\.cross_attn\.k\.(.*)$": + r"blocks.\1.attn2.to_k.\2", + r"^model.blocks\.(\d+)\.cross_attn\.k_img\.(.*)$": + r"blocks.\1.attn2.add_k_proj.\2", + r"^model.blocks\.(\d+)\.cross_attn\.v\.(.*)$": + r"blocks.\1.attn2.to_v.\2", + r"^model.blocks\.(\d+)\.cross_attn\.v_img\.(.*)$": + r"blocks.\1.attn2.add_v_proj.\2", + r"^model.blocks\.(\d+)\.cross_attn\.o\.(.*)$": + r"blocks.\1.attn2.to_out.0.\2", + r"^model.blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$": + r"blocks.\1.attn2.norm_q.\2", + r"^model.blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$": + r"blocks.\1.attn2.norm_k.\2", + r"^model.blocks\.(\d+)\.cross_attn\.norm_k_img\.(.*)$": + r"blocks.\1.attn2.norm_added_k.\2", + r"^model.blocks\.(\d+)\.ffn\.0\.(.*)$": + r"blocks.\1.ffn.net.0.proj.\2", + r"^model.blocks\.(\d+)\.ffn\.2\.(.*)$": + r"blocks.\1.ffn.net.2.\2", + r"^model.blocks\.(\d+)\.modulation": + r"blocks.\1.scale_shift_table", + r"^model.blocks\.(\d+)\.norm3\.(.*)$": + r"blocks.\1.norm2.\2", +} + state_dict = load_state_dict_from_file("checkpoints/self_forcing_dmd.pt") -# state_dict = state_dict["generator_ema"] +state_dict = state_dict["generator_ema"] new_state_dict = OrderedDict() for k, v in state_dict.items(): new_key = k - for pattern, replacement in _param_names_mapping.items(): + for pattern, replacement in _self_forcing_to_diffusers_param_names_mapping.items(): if re.match(pattern, k): new_key = re.sub(pattern, replacement, k) break # Stop at the first match + else: + # print(f"No match found for {k}") + raise ValueError(f"No match found for {k}") new_state_dict[new_key] = v if "norm_added_k" in new_key: dummy_key = new_key.replace("norm_added_k", "norm_added_q") @@ -86,6 +156,6 @@ del state_dict save_torch_state_dict( new_state_dict, - "new_diffusers/", + "new2/", max_shard_size="10GB" -) \ No newline at end of file +) From 772ead0d34d48ed0da15505c8420dbf650976fc0 Mon Sep 17 00:00:00 2001 From: JerryZhou54 Date: Thu, 28 Aug 2025 01:20:57 +0000 Subject: [PATCH 10/13] Wan2.1 causal, few-step inference runnable --- fastvideo/configs/pipelines/wan.py | 2 + fastvideo/layers/rotary_embedding.py | 1 - fastvideo/models/dits/causal_wanvideo.py | 23 +++----- fastvideo/pipelines/pipeline_batch_info.py | 4 ++ .../pipelines/stages/causal_denoising.py | 56 ++++++++++--------- 5 files changed, 42 insertions(+), 44 deletions(-) diff --git a/fastvideo/configs/pipelines/wan.py b/fastvideo/configs/pipelines/wan.py index 579da6dfd..2d90a98f3 100644 --- a/fastvideo/configs/pipelines/wan.py +++ b/fastvideo/configs/pipelines/wan.py @@ -146,3 +146,5 @@ class Wan2_2_I2V_A14B_Config(WanT2V480PConfig): @dataclass class SelfForcingWanT2V480PConfig(WanT2V480PConfig): is_causal: bool = True + dmd_denoising_steps: list[int] | None = field( + default_factory=lambda: [1000, 750, 500, 250]) diff --git a/fastvideo/layers/rotary_embedding.py b/fastvideo/layers/rotary_embedding.py index 4dd789db5..e3de7018f 100644 --- a/fastvideo/layers/rotary_embedding.py +++ b/fastvideo/layers/rotary_embedding.py @@ -297,7 +297,6 @@ def get_nd_rotary_pos_embed( if start_frame > 0: full_grid[0] += start_frame - logger.info(full_grid[0]) # Shard the grid if using sequence parallelism (sp_world_size > 1) assert shard_dim < len( diff --git a/fastvideo/models/dits/causal_wanvideo.py b/fastvideo/models/dits/causal_wanvideo.py index 244d3ad8c..d2955a03c 100644 --- a/fastvideo/models/dits/causal_wanvideo.py +++ b/fastvideo/models/dits/causal_wanvideo.py @@ -58,14 +58,6 @@ class CausalWanSelfAttention(nn.Module): self.parallel_attention = parallel_attention self.max_attention_size = 32760 if local_attn_size == -1 else local_attn_size * 1560 - # layers - self.to_q = ReplicatedLinear(dim, dim) - self.to_k = ReplicatedLinear(dim, dim) - self.to_v = ReplicatedLinear(dim, dim) - self.to_out = ReplicatedLinear(dim, dim) - self.norm_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() - self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() - # Scaled dot product attention self.attn = LocalAttention( num_heads=num_heads, @@ -96,8 +88,8 @@ class CausalWanSelfAttention(nn.Module): cache_start = current_start cos, sin = freqs_cis - roped_query = _apply_rotary_emb(q, cos, sin, is_neox_style=False) - roped_key = _apply_rotary_emb(k, cos, sin, is_neox_style=False) + roped_query = _apply_rotary_emb(q, cos, sin, is_neox_style=False).type_as(v) + roped_key = _apply_rotary_emb(k, cos, sin, is_neox_style=False).type_as(v) if kv_cache is None: # Padding for flex attention @@ -165,9 +157,6 @@ class CausalWanSelfAttention(nn.Module): kv_cache["global_end_index"].fill_(current_end) kv_cache["local_end_index"].fill_(local_end_index) - # output - x = x.flatten(2) - x = self.to_out(x) return x class CausalWanTransformerBlock(nn.Module): @@ -281,7 +270,7 @@ class CausalWanTransformerBlock(nn.Module): key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) - attn_output, _ = self.attn1(query, key, value, freqs_cis, block_mask, kv_cache, current_start, cache_start) + attn_output = self.attn1(query, key, value, freqs_cis, block_mask, kv_cache, current_start, cache_start) attn_output = attn_output.flatten(2) attn_output, _ = self.to_out(attn_output) attn_output = attn_output.squeeze(1) @@ -515,7 +504,8 @@ class CausalWanTransformer3DModel(BaseDiT): causal_kwargs = { "kv_cache": kv_cache[block_index], "current_start": current_start, - "cache_start": cache_start + "cache_start": cache_start, + "block_mask": self.block_mask } hidden_states = self._gradient_checkpointing_func( block, hidden_states, encoder_hidden_states, @@ -526,7 +516,8 @@ class CausalWanTransformer3DModel(BaseDiT): "kv_cache": kv_cache[block_index], "crossattn_cache": crossattn_cache[block_index], "current_start": current_start, - "cache_start": cache_start + "cache_start": cache_start, + "block_mask": self.block_mask } hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index 7fc9fa56c..f9850f3f0 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -173,6 +173,10 @@ class ForwardBatch: logging_info: PipelineLoggingInfo = field( default_factory=PipelineLoggingInfo) + # Causal Wan parameters + num_frames_per_block: int = 3 + sliding_window_num_frames: int = 21 + def __post_init__(self): """Initialize dependent fields after dataclass initialization.""" diff --git a/fastvideo/pipelines/stages/causal_denoising.py b/fastvideo/pipelines/stages/causal_denoising.py index 07289ab75..462920d1a 100644 --- a/fastvideo/pipelines/stages/causal_denoising.py +++ b/fastvideo/pipelines/stages/causal_denoising.py @@ -273,10 +273,10 @@ class CausalDMDDenosingStage(DenoisingStage): # Args try: - num_frame_per_block = getattr(fastvideo_args.pipeline_config, - "num_frame_per_block", 1) + num_frames_per_block = getattr(batch, + "num_frames_per_block", 1) except Exception: - num_frame_per_block = 1 + num_frames_per_block = 1 try: independent_first_frame = getattr(fastvideo_args.pipeline_config, "independent_first_frame", False) @@ -314,10 +314,10 @@ class CausalDMDDenosingStage(DenoisingStage): # Initialize or reset caches if self.kv_cache1 is None: self._initialize_kv_cache(batch_size=latents.shape[0], - dtype=latents.dtype, + dtype=target_dtype, device=latents.device) self._initialize_crossattn_cache(batch_size=latents.shape[0], - dtype=latents.dtype, + dtype=target_dtype, device=latents.device) else: # reset cross-attention cache @@ -372,7 +372,7 @@ class CausalDMDDenosingStage(DenoisingStage): # process remaining input frames in blocks of num_frame_per_block while remaining_frames > 0: - block = min(num_frame_per_block, remaining_frames) + block = min(num_frames_per_block, remaining_frames) ref_btchw = image_latent[:, :, current_start_frame: current_start_frame + block, :, :].to(target_dtype).permute( @@ -400,20 +400,20 @@ class CausalDMDDenosingStage(DenoisingStage): # Determine block sizes if not independent_first_frame or (independent_first_frame and batch.image_latent is not None): - if t % num_frame_per_block != 0: + if t % num_frames_per_block != 0: raise ValueError( - "num_frames must be divisible by num_frame_per_block for causal DMD denoising" + "num_frames must be divisible by num_frames_per_block for causal DMD denoising" ) - num_blocks = t // num_frame_per_block - block_sizes = [num_frame_per_block] * num_blocks + num_blocks = t // num_frames_per_block + block_sizes = [num_frames_per_block] * num_blocks start_index = 0 else: - if (t - 1) % num_frame_per_block != 0: + if (t - 1) % num_frames_per_block != 0: raise ValueError( "(num_frames - 1) must be divisible by num_frame_per_block when independent_first_frame=True" ) - num_blocks = (t - 1) // num_frame_per_block - block_sizes = [1] + [num_frame_per_block] * num_blocks + num_blocks = (t - 1) // num_frames_per_block + block_sizes = [1] + [num_frames_per_block] * num_blocks start_index = 0 # DMD loop in causal blocks @@ -474,7 +474,7 @@ class CausalDMDDenosingStage(DenoisingStage): forward_batch=batch): # Run transformer; follow DMD stage pattern pred_noise_btchw = self.transformer( - latent_model_input.permute(0, 2, 1, 3, 4), + latent_model_input, prompt_embeds, t_expand, kv_cache=self.kv_cache1, @@ -531,22 +531,24 @@ class CausalDMDDenosingStage(DenoisingStage): t_context = torch.ones([latents.shape[0]], device=latents.device, dtype=torch.long) * int(context_noise) - context_btchw = current_latents.to(target_dtype).permute( - 0, 2, 1, 3, 4) + context_bcthw = current_latents.to(target_dtype) with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled): - _ = self.transformer( - context_btchw, - prompt_embeds, - t_context, - kv_cache=self.kv_cache1, - crossattn_cache=self.crossattn_cache, - current_start=(pos_start_base + start_index) * - self.frame_seq_length, - **image_kwargs, - **pos_cond_kwargs, - ) + with set_forward_context(current_timestep=0, + attn_metadata=attn_metadata, + forward_batch=batch): + _ = self.transformer( + context_bcthw, + prompt_embeds, + t_context, + kv_cache=self.kv_cache1, + crossattn_cache=self.crossattn_cache, + current_start=(pos_start_base + start_index) * + self.frame_seq_length, + **image_kwargs, + **pos_cond_kwargs, + ) start_index += current_num_frames batch.latents = latents From 1a79b30ea4989bd20d78d85d141b509f1ff6db9c Mon Sep 17 00:00:00 2001 From: JerryZhou54 Date: Thu, 28 Aug 2025 03:18:58 +0000 Subject: [PATCH 11/13] Fix inter-block issues --- fastvideo/models/dits/causal_wanvideo.py | 3 ++- fastvideo/pipelines/stages/causal_denoising.py | 6 ++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/fastvideo/models/dits/causal_wanvideo.py b/fastvideo/models/dits/causal_wanvideo.py index d2955a03c..97ac3a5da 100644 --- a/fastvideo/models/dits/causal_wanvideo.py +++ b/fastvideo/models/dits/causal_wanvideo.py @@ -440,6 +440,7 @@ class CausalWanTransformer3DModel(BaseDiT): crossattn_cache: dict = None, current_start: int = 0, cache_start: int = 0, + start_frame: int = 0, **kwargs) -> torch.Tensor: r""" Run the diffusion model with kv caching. @@ -474,7 +475,7 @@ class CausalWanTransformer3DModel(BaseDiT): rope_dim_list, dtype=torch.float32 if current_platform.is_mps() else torch.float64, rope_theta=10000, - start_frame=current_start # Assume that current_start is 0 when kv_cache is None + start_frame=start_frame # Assume that start_frame is 0 when kv_cache is None ) freqs_cos = freqs_cos.to(hidden_states.device) freqs_sin = freqs_sin.to(hidden_states.device) diff --git a/fastvideo/pipelines/stages/causal_denoising.py b/fastvideo/pipelines/stages/causal_denoising.py index 462920d1a..2dda08546 100644 --- a/fastvideo/pipelines/stages/causal_denoising.py +++ b/fastvideo/pipelines/stages/causal_denoising.py @@ -481,6 +481,7 @@ class CausalDMDDenosingStage(DenoisingStage): crossattn_cache=self.crossattn_cache, current_start=(pos_start_base + start_index) * self.frame_seq_length, + start_frame=start_index, **image_kwargs, **pos_cond_kwargs, ).permute(0, 2, 1, 3, 4) @@ -507,12 +508,12 @@ class CausalDMDDenosingStage(DenoisingStage): batch.generator, list) else batch.generator)).to(self.device) noise_btchw = noise - current_latents_btchw = self.scheduler.add_noise( + noise_latents_btchw = self.scheduler.add_noise( pred_video_btchw.flatten(0, 1), noise_btchw.flatten(0, 1), next_timestep).unflatten(0, pred_video_btchw.shape[:2]) - current_latents = current_latents_btchw.permute( + current_latents = noise_latents_btchw.permute( 0, 2, 1, 3, 4) else: current_latents = pred_video_btchw.permute( @@ -546,6 +547,7 @@ class CausalDMDDenosingStage(DenoisingStage): crossattn_cache=self.crossattn_cache, current_start=(pos_start_base + start_index) * self.frame_seq_length, + start_frame=start_index, **image_kwargs, **pos_cond_kwargs, ) From 8dd0b6536d77314aa419d164f4cbfdda8b432ca1 Mon Sep 17 00:00:00 2001 From: JerryZhou54 Date: Thu, 28 Aug 2025 04:07:11 +0000 Subject: [PATCH 12/13] Fix pre-commit tests --- fastvideo/configs/models/dits/wanvideo.py | 4 +- fastvideo/layers/rotary_embedding.py | 1 + .../basic/wan/wan_causal_dmd_pipeline.py | 8 +- .../pipelines/stages/causal_denoising.py | 80 +++++++++---------- 4 files changed, 47 insertions(+), 46 deletions(-) diff --git a/fastvideo/configs/models/dits/wanvideo.py b/fastvideo/configs/models/dits/wanvideo.py index d31a3a607..a08f81942 100644 --- a/fastvideo/configs/models/dits/wanvideo.py +++ b/fastvideo/configs/models/dits/wanvideo.py @@ -93,8 +93,8 @@ class WanVideoArchConfig(DiTArchConfig): exclude_lora_layers: list[str] = field(default_factory=lambda: ["embedder"]) # Causal Wan - local_attn_size: int = -1 # Window size for temporal local attention (-1 indicates global attention) - sink_size: int = 0 # Size of the attention sink, we keep the first `sink_size` frames unchanged when rolling the KV cache + local_attn_size: int = -1 # Window size for temporal local attention (-1 indicates global attention) + sink_size: int = 0 # Size of the attention sink, we keep the first `sink_size` frames unchanged when rolling the KV cache def __post_init__(self): super().__post_init__() diff --git a/fastvideo/layers/rotary_embedding.py b/fastvideo/layers/rotary_embedding.py index e3de7018f..1270f3151 100644 --- a/fastvideo/layers/rotary_embedding.py +++ b/fastvideo/layers/rotary_embedding.py @@ -33,6 +33,7 @@ from fastvideo.logger import init_logger logger = init_logger(__name__) + def _rotate_neox(x: torch.Tensor) -> torch.Tensor: x1 = x[..., :x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2:] diff --git a/fastvideo/pipelines/basic/wan/wan_causal_dmd_pipeline.py b/fastvideo/pipelines/basic/wan/wan_causal_dmd_pipeline.py index a46898530..97426420d 100644 --- a/fastvideo/pipelines/basic/wan/wan_causal_dmd_pipeline.py +++ b/fastvideo/pipelines/basic/wan/wan_causal_dmd_pipeline.py @@ -13,8 +13,10 @@ from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline # isort: off from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, - CausalDMDDenosingStage, InputValidationStage, - LatentPreparationStage, TextEncodingStage, + CausalDMDDenosingStage, + InputValidationStage, + LatentPreparationStage, + TextEncodingStage, TimestepPreparationStage) # isort: on @@ -65,5 +67,3 @@ class WanCausalDMDPipeline(LoRAPipeline, ComposedPipelineBase): EntryClass = WanCausalDMDPipeline - - diff --git a/fastvideo/pipelines/stages/causal_denoising.py b/fastvideo/pipelines/stages/causal_denoising.py index 2dda08546..0876cbda3 100644 --- a/fastvideo/pipelines/stages/causal_denoising.py +++ b/fastvideo/pipelines/stages/causal_denoising.py @@ -250,8 +250,8 @@ class CausalDMDDenosingStage(DenoisingStage): super().__init__(transformer, scheduler) self.scheduler = FlowMatchEulerDiscreteScheduler(shift=8.0) # KV and cross-attention cache state (initialized on first forward) - self.kv_cache1 = None - self.crossattn_cache = None + self.kv_cache1: list | None = None + self.crossattn_cache: list | None = None # Model-dependent constants (aligned with causal_inference.py assumptions) self.num_transformer_blocks = 30 self.frame_seq_length = 1560 @@ -273,8 +273,7 @@ class CausalDMDDenosingStage(DenoisingStage): # Args try: - num_frames_per_block = getattr(batch, - "num_frames_per_block", 1) + num_frames_per_block = getattr(batch, "num_frames_per_block", 1) except Exception: num_frames_per_block = 1 try: @@ -290,7 +289,7 @@ class CausalDMDDenosingStage(DenoisingStage): device=get_local_torch_device()) # Image kwargs (kept empty unless caller provides compatible args) - image_kwargs = {} + image_kwargs: dict = {} pos_cond_kwargs = self.prepare_extra_func_kwargs( self.transformer.forward, @@ -320,6 +319,7 @@ class CausalDMDDenosingStage(DenoisingStage): dtype=target_dtype, device=latents.device) else: + assert self.crossattn_cache is not None # reset cross-attention cache for block_index in range(self.num_transformer_blocks): self.crossattn_cache[block_index][ @@ -468,23 +468,23 @@ class CausalDMDDenosingStage(DenoisingStage): with torch.autocast(device_type="cuda", dtype=target_dtype, - enabled=autocast_enabled): - with set_forward_context(current_timestep=i, - attn_metadata=attn_metadata, - forward_batch=batch): - # Run transformer; follow DMD stage pattern - pred_noise_btchw = self.transformer( - latent_model_input, - prompt_embeds, - t_expand, - kv_cache=self.kv_cache1, - crossattn_cache=self.crossattn_cache, - current_start=(pos_start_base + start_index) * - self.frame_seq_length, - start_frame=start_index, - **image_kwargs, - **pos_cond_kwargs, - ).permute(0, 2, 1, 3, 4) + enabled=autocast_enabled), \ + set_forward_context(current_timestep=i, + attn_metadata=attn_metadata, + forward_batch=batch): + # Run transformer; follow DMD stage pattern + pred_noise_btchw = self.transformer( + latent_model_input, + prompt_embeds, + t_expand, + kv_cache=self.kv_cache1, + crossattn_cache=self.crossattn_cache, + current_start=(pos_start_base + start_index) * + self.frame_seq_length, + start_frame=start_index, + **image_kwargs, + **pos_cond_kwargs, + ).permute(0, 2, 1, 3, 4) # Convert pred noise to pred video with FM Euler scheduler utilities from fastvideo.training.training_utils import ( @@ -535,28 +535,28 @@ class CausalDMDDenosingStage(DenoisingStage): context_bcthw = current_latents.to(target_dtype) with torch.autocast(device_type="cuda", dtype=target_dtype, - enabled=autocast_enabled): - with set_forward_context(current_timestep=0, - attn_metadata=attn_metadata, - forward_batch=batch): - _ = self.transformer( - context_bcthw, - prompt_embeds, - t_context, - kv_cache=self.kv_cache1, - crossattn_cache=self.crossattn_cache, - current_start=(pos_start_base + start_index) * - self.frame_seq_length, - start_frame=start_index, - **image_kwargs, - **pos_cond_kwargs, - ) + enabled=autocast_enabled), \ + set_forward_context(current_timestep=0, + attn_metadata=attn_metadata, + forward_batch=batch): + _ = self.transformer( + context_bcthw, + prompt_embeds, + t_context, + kv_cache=self.kv_cache1, + crossattn_cache=self.crossattn_cache, + current_start=(pos_start_base + start_index) * + self.frame_seq_length, + start_frame=start_index, + **image_kwargs, + **pos_cond_kwargs, + ) start_index += current_num_frames batch.latents = latents return batch - def _initialize_kv_cache(self, batch_size, dtype, device): + def _initialize_kv_cache(self, batch_size, dtype, device) -> None: """ Initialize a Per-GPU KV cache aligned with the Wan model assumptions. """ @@ -584,7 +584,7 @@ class CausalDMDDenosingStage(DenoisingStage): self.kv_cache1 = kv_cache1 - def _initialize_crossattn_cache(self, batch_size, dtype, device): + def _initialize_crossattn_cache(self, batch_size, dtype, device) -> None: """ Initialize a Per-GPU cross-attention cache aligned with the Wan model assumptions. """ From 8acc7c3655133950179ba685fad0ac287821d01b Mon Sep 17 00:00:00 2001 From: JerryZhou54 Date: Thu, 28 Aug 2025 04:10:51 +0000 Subject: [PATCH 13/13] small fix --- fastvideo/models/loader/fsdp_load.py | 7 - .../pipelines/stages/causal_denoising.py | 219 +----------------- 2 files changed, 1 insertion(+), 225 deletions(-) diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index 12e5786c2..117b1a728 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -248,13 +248,6 @@ def load_model_from_full_model_state_dict( sharded_sd = {} custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict( full_sd_iterator, param_names_mapping) # type: ignore - print(custom_param_sd.keys()) - print("--------------------------------") - print("--------------------------------") - print("--------------------------------") - print("--------------------------------") - print("--------------------------------") - print(meta_sd.keys()) for target_param_name, full_tensor in custom_param_sd.items(): meta_sharded_param = meta_sd.get(target_param_name) if meta_sharded_param is None: diff --git a/fastvideo/pipelines/stages/causal_denoising.py b/fastvideo/pipelines/stages/causal_denoising.py index 0876cbda3..be00b1f2e 100644 --- a/fastvideo/pipelines/stages/causal_denoising.py +++ b/fastvideo/pipelines/stages/causal_denoising.py @@ -7,6 +7,7 @@ from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( FlowMatchEulerDiscreteScheduler) from fastvideo.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.pipelines.stages.denoising import DenoisingStage +from fastvideo.training.training_utils import pred_noise_to_pred_video try: from fastvideo.attention.backends.sliding_tile_attn import ( @@ -24,222 +25,6 @@ except ImportError: vsa_available = False VideoSparseAttentionBackend = None # type: ignore -# class CausalDenoisingStage(LoRAPipeline, ComposedPipelineBase): -# """ -# Flow Causal denoising stage. -# """ - -# def forward( -# self, -# batch: ForwardBatch, -# fastvideo_args: FastVideoArgs, -# ) -> ForwardBatch: -# # Setup precision and autocast settings -# target_dtype = torch.bfloat16 -# autocast_enabled = (target_dtype != torch.float32 -# ) and not fastvideo_args.disable_autocast - -# # Gather args for causal block-wise processing -# try: -# num_frame_per_block = getattr(fastvideo_args.pipeline_config, -# "num_frame_per_block", 1) -# except Exception: -# num_frame_per_block = 1 -# try: -# independent_first_frame = getattr(fastvideo_args.pipeline_config, -# "independent_first_frame", False) -# except Exception: -# independent_first_frame = False - -# # Required fields -# timesteps = batch.timesteps -# if timesteps is None: -# raise ValueError("Timesteps must be provided") - -# # Prepare image latents and embeddings for I2V generation -# image_embeds = batch.image_embeds -# if len(image_embeds) > 0: -# assert torch.isnan(image_embeds[0]).sum() == 0 -# image_embeds = [ -# image_embed.to(target_dtype) for image_embed in image_embeds -# ] - -# image_kwargs = self.prepare_extra_func_kwargs( -# self.transformer.forward, -# { -# "encoder_hidden_states_image": image_embeds, -# "mask_strategy": dict_to_3d_list( -# None, t_max=50, l_max=60, h_max=24) -# }, -# ) - -# pos_cond_kwargs = self.prepare_extra_func_kwargs( -# self.transformer.forward, -# { -# "encoder_hidden_states_2": batch.clip_embedding_pos, -# "encoder_attention_mask": batch.prompt_attention_mask, -# }, -# ) - -# neg_cond_kwargs = self.prepare_extra_func_kwargs( -# self.transformer.forward, -# { -# "encoder_hidden_states_2": batch.clip_embedding_neg, -# "encoder_attention_mask": batch.negative_attention_mask, -# }, -# ) - -# # Prepare STA parameters -# if st_attn_available and self.attn_backend == SlidingTileAttentionBackend: -# self.prepare_sta_param(batch, fastvideo_args) - -# # Get latents and embeddings -# latents = batch.latents -# assert latents is not None, "latents must be provided" -# prompt_embeds = batch.prompt_embeds -# assert torch.isnan(prompt_embeds[0]).sum() == 0 - -# # Determine blocks -# b, c, t, h, w = latents.shape -# if not independent_first_frame or (independent_first_frame -# and batch.image_latent is not None): -# if t % num_frame_per_block != 0: -# raise ValueError( -# "num_frames must be divisible by num_frame_per_block for causal denoising" -# ) -# num_blocks = t // num_frame_per_block -# block_sizes = [num_frame_per_block] * num_blocks -# start_index = 0 -# else: -# if (t - 1) % num_frame_per_block != 0: -# raise ValueError( -# "(num_frames - 1) must be divisible by num_frame_per_block when independent_first_frame=True" -# ) -# num_blocks = (t - 1) // num_frame_per_block -# block_sizes = [1] + [num_frame_per_block] * num_blocks -# start_index = 0 - -# # Run causal block-wise denoising loop -# with self.progress_bar(total=len(block_sizes) * -# len(timesteps)) as progress_bar: -# for current_num_frames in block_sizes: -# # Build the latent input for this block -# current_latents = latents[:, :, start_index:start_index + -# current_num_frames, :, :] - -# # Run spatial denoising loop over timesteps -# for i, t_cur in enumerate(timesteps): -# # Expand latents for I2V -# latent_model_input = current_latents.to(target_dtype) -# if batch.image_latent is not None and independent_first_frame and start_index == 0: -# # Prepend the image latent only for the first block when required -# latent_model_input = torch.cat([ -# batch.image_latent.to(target_dtype), -# latent_model_input -# ], -# dim=2) - -# latent_model_input = self.scheduler.scale_model_input( -# latent_model_input, t_cur) - -# # Prepare inputs for transformer -# t_expand = t_cur.repeat(latent_model_input.shape[0]) -# guidance_expand = (torch.tensor( -# [fastvideo_args.pipeline_config.embedded_cfg_scale] * -# latent_model_input.shape[0], -# dtype=torch.float32, -# device=get_local_torch_device(), -# ).to(target_dtype) * -# 1000.0 if fastvideo_args.pipeline_config. -# embedded_cfg_scale is not None else None) - -# # Predict noise residual -# with torch.autocast(device_type="cuda", -# dtype=target_dtype, -# enabled=autocast_enabled): -# if (st_attn_available and self.attn_backend == SlidingTileAttentionBackend) or \ -# (vsa_available and self.attn_backend == VideoSparseAttentionBackend): -# self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls( -# ) -# if self.attn_metadata_builder_cls is not None: -# self.attn_metadata_builder = self.attn_metadata_builder_cls( -# ) -# attn_metadata = self.attn_metadata_builder.build( # type: ignore -# current_timestep=i, # type: ignore -# raw_latent_shape=(current_num_frames, h, -# w), # type: ignore -# patch_size=fastvideo_args.pipeline_config. -# dit_config.patch_size, # type: ignore -# STA_param=batch.STA_param, # type: ignore -# VSA_sparsity=fastvideo_args. -# VSA_sparsity, # type: ignore -# device=get_local_torch_device(), -# ) -# assert attn_metadata is not None, "attn_metadata cannot be None" -# else: -# attn_metadata = None -# else: -# attn_metadata = None - -# batch.is_cfg_negative = False -# with set_forward_context(current_timestep=i, -# attn_metadata=attn_metadata, -# forward_batch=batch): -# noise_pred = self.transformer( -# latent_model_input, -# prompt_embeds, -# t_expand, -# guidance=guidance_expand, -# **image_kwargs, -# **pos_cond_kwargs, -# ) - -# if batch.do_classifier_free_guidance: -# batch.is_cfg_negative = True -# with set_forward_context( -# current_timestep=i, -# attn_metadata=attn_metadata, -# forward_batch=batch): -# noise_pred_uncond = self.transformer( -# latent_model_input, -# batch.negative_prompt_embeds, -# t_expand, -# guidance=guidance_expand, -# **image_kwargs, -# **neg_cond_kwargs, -# ) -# noise_pred_text = noise_pred -# noise_pred = noise_pred_uncond + batch.guidance_scale * ( -# noise_pred_text - noise_pred_uncond) - -# if batch.guidance_rescale > 0.0: -# noise_pred = self.rescale_noise_cfg( -# noise_pred, -# noise_pred_text, -# guidance_rescale=batch.guidance_rescale, -# ) - -# # Compute the previous noisy sample for this block only -# current_latents = self.scheduler.step( -# noise_pred, -# t_cur, -# current_latents, -# return_dict=False)[0] - -# if progress_bar is not None: -# progress_bar.update() - -# # After finishing timesteps, write denoised block back -# latents[:, :, start_index:start_index + -# current_num_frames, :, :] = current_latents - -# # Advance to next block -# start_index += current_num_frames - -# # Update batch with final latents -# batch.latents = latents -# return batch - class CausalDMDDenosingStage(DenoisingStage): """ @@ -487,8 +272,6 @@ class CausalDMDDenosingStage(DenoisingStage): ).permute(0, 2, 1, 3, 4) # Convert pred noise to pred video with FM Euler scheduler utilities - from fastvideo.training.training_utils import ( - pred_noise_to_pred_video) pred_video_btchw = pred_noise_to_pred_video( pred_noise=pred_noise_btchw.flatten(0, 1), noise_input_latent=noise_latents.flatten(0, 1),