From ce67cd3729e5f5aa744920d61102c0eb05f275df Mon Sep 17 00:00:00 2001 From: Wei Zhou <69577934+JerryZhou54@users.noreply.github.com> Date: Thu, 28 Aug 2025 16:47:49 -0700 Subject: [PATCH] [Feat] Support Self-Forcing's Causal Inference for Wan2.1 T2V 1.3B (#766) Co-authored-by: SolitaryThinker --- .../basic/basic_self_forcing_causal.py | 31 + fastvideo/configs/models/dits/wanvideo.py | 6 + fastvideo/configs/pipelines/registry.py | 2 + fastvideo/configs/pipelines/wan.py | 10 + fastvideo/configs/sample/registry.py | 17 +- fastvideo/configs/sample/wan.py | 8 + fastvideo/layers/rotary_embedding.py | 9 + fastvideo/models/dits/causal_wanvideo.py | 648 ++++++++++++++++++ fastvideo/models/dits/wanvideo.py | 18 +- fastvideo/models/registry.py | 2 + fastvideo/models/utils.py | 21 + .../basic/wan/wan_causal_dmd_pipeline.py | 69 ++ fastvideo/pipelines/pipeline_registry.py | 1 + fastvideo/pipelines/stages/__init__.py | 2 + .../pipelines/stages/causal_denoising.py | 409 +++++++++++ fastvideo/pipelines/stages/denoising.py | 3 +- fastvideo/training/distillation_pipeline.py | 4 +- fastvideo/training/training_utils.py | 21 - .../checkpoint_conversion/wan_to_diffusers.py | 161 +++++ 19 files changed, 1412 insertions(+), 30 deletions(-) create mode 100644 examples/inference/basic/basic_self_forcing_causal.py create mode 100644 fastvideo/models/dits/causal_wanvideo.py create mode 100644 fastvideo/pipelines/basic/wan/wan_causal_dmd_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_self_forcing_causal.py b/examples/inference/basic/basic_self_forcing_causal.py new file mode 100644 index 000000000..e2241a94e --- /dev/null +++ b/examples/inference/basic/basic_self_forcing_causal.py @@ -0,0 +1,31 @@ +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. + 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, + ) + + sampling_param = SamplingParam.from_pretrained(model_name) + + 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) + +if __name__ == "__main__": + main() diff --git a/fastvideo/configs/models/dits/wanvideo.py b/fastvideo/configs/models/dits/wanvideo.py index 8cd25cf6c..178e938c4 100644 --- a/fastvideo/configs/models/dits/wanvideo.py +++ b/fastvideo/configs/models/dits/wanvideo.py @@ -92,6 +92,12 @@ 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 + num_frames_per_block: int = 3 + sliding_window_num_frames: int = 21 + def __post_init__(self): super().__post_init__() self.out_channels = self.out_channels or self.in_channels 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..2d90a98f3 100644 --- a/fastvideo/configs/pipelines/wan.py +++ b/fastvideo/configs/pipelines/wan.py @@ -138,3 +138,13 @@ 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 + dmd_denoising_steps: list[int] | None = field( + default_factory=lambda: [1000, 750, 500, 250]) 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..d94c5f714 100644 --- a/fastvideo/configs/sample/wan.py +++ b/fastvideo/configs/sample/wan.py @@ -141,3 +141,11 @@ 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): + pass diff --git a/fastvideo/layers/rotary_embedding.py b/fastvideo/layers/rotary_embedding.py index 6abe90609..1270f3151 100644 --- a/fastvideo/layers/rotary_embedding.py +++ b/fastvideo/layers/rotary_embedding.py @@ -29,6 +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: @@ -267,6 +270,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 +296,9 @@ 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 + # 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/causal_wanvideo.py b/fastvideo/models/dits/causal_wanvideo.py new file mode 100644 index 000000000..993a8044c --- /dev/null +++ b/fastvideo/models/dits/causal_wanvideo.py @@ -0,0 +1,648 @@ +# 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 + + # 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).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 + 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) + + 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.attention_head_dim = config.attention_head_dim + 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, + start_frame: 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=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) + 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, + "block_mask": self.block_mask + } + 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, + "block_mask": self.block_mask + } + 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, + start_frame: int = 0, + **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, + start_frame=start_frame + ) + 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) 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) 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/models/utils.py b/fastvideo/models/utils.py index 36233720b..57cb7a74f 100644 --- a/fastvideo/models/utils.py +++ b/fastvideo/models/utils.py @@ -137,3 +137,24 @@ def modulate(x: torch.Tensor, else: return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze( 1) # type: ignore[union-attr] + + +def pred_noise_to_pred_video(pred_noise: torch.Tensor, + noise_input_latent: torch.Tensor, + timestep: torch.Tensor, + scheduler: Any) -> torch.Tensor: + """ + Convert predicted noise to clean latent. + """ + timestep = timestep.expand(noise_input_latent.shape[0]) + dtype = pred_noise.dtype + device = pred_noise.device + pred_noise = pred_noise.float().to(device) + noise_input_latent = noise_input_latent.float().to(device) + sigmas = scheduler.sigmas.float().to(device) + timesteps = scheduler.timesteps.float().to(device) + timestep_id = torch.argmin( + (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) + sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1) + pred_video = noise_input_latent - sigma_t * pred_noise + return pred_video.to(dtype) 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..97426420d --- /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/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 1852edc49..9566b426f 100644 --- a/fastvideo/pipelines/stages/__init__.py +++ b/fastvideo/pipelines/stages/__init__.py @@ -7,6 +7,7 @@ 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, @@ -30,6 +31,7 @@ __all__ = [ "ConditioningStage", "DenoisingStage", "DmdDenoisingStage", + "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..e25b62cf0 --- /dev/null +++ b/fastvideo/pipelines/stages/causal_denoising.py @@ -0,0 +1,409 @@ +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.logger import init_logger +from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler) +from fastvideo.models.utils import pred_noise_to_pred_video +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 + +logger = init_logger(__name__) + + +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: list | None = None + self.crossattn_cache: list | None = None + # Model-dependent constants (aligned with causal_inference.py assumptions) + self.num_transformer_blocks = self.transformer.config.arch_config.num_layers + self.num_frames_per_block = self.transformer.config.arch_config.num_frames_per_block + self.sliding_window_num_frames = self.transformer.config.arch_config.sliding_window_num_frames + + 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 + + latent_seq_length = batch.latents.shape[-1] * batch.latents.shape[-2] + patch_ratio = self.transformer.config.arch_config.patch_size[ + -1] * self.transformer.config.arch_config.patch_size[-2] + self.frame_seq_length = latent_seq_length // patch_ratio + # TODO(will): make this a parameter once we add i2v support + independent_first_frame = self.transformer.independent_first_frame + + # 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: dict = {} + + 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=target_dtype, + device=latents.device) + self._initialize_crossattn_cache( + batch_size=latents.shape[0], + max_text_len=fastvideo_args.pipeline_config. + text_encoder_configs[0].arch_config.text_len, + 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][ + "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(self.num_frames_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 % self.num_frames_per_block != 0: + raise ValueError( + "num_frames must be divisible by num_frames_per_block for causal DMD denoising" + ) + num_blocks = t // self.num_frames_per_block + block_sizes = [self.num_frames_per_block] * num_blocks + start_index = 0 + else: + if (t - 1) % self.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) // self.num_frames_per_block + block_sizes = [1] + [self.num_frames_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), \ + 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 + 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 + 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 = noise_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_bcthw = current_latents.to(target_dtype) + with torch.autocast(device_type="cuda", + dtype=target_dtype, + 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) -> None: + """ + Initialize a Per-GPU KV cache aligned with the Wan model assumptions. + """ + kv_cache1 = [] + num_attention_heads = self.transformer.num_attention_heads + attention_head_dim = self.transformer.attention_head_dim + if self.local_attn_size != -1: + kv_cache_size = self.local_attn_size * self.frame_seq_length + else: + kv_cache_size = self.frame_seq_length * self.sliding_window_num_frames + + for _ in range(self.num_transformer_blocks): + kv_cache1.append({ + "k": + torch.zeros([ + batch_size, kv_cache_size, num_attention_heads, + attention_head_dim + ], + dtype=dtype, + device=device), + "v": + torch.zeros([ + batch_size, kv_cache_size, num_attention_heads, + attention_head_dim + ], + 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, max_text_len, dtype, + device) -> None: + """ + Initialize a Per-GPU cross-attention cache aligned with the Wan model assumptions. + """ + crossattn_cache = [] + num_attention_heads = self.transformer.num_attention_heads + attention_head_dim = self.transformer.attention_head_dim + for _ in range(self.num_transformer_blocks): + crossattn_cache.append({ + "k": + torch.zeros([ + batch_size, max_text_len, num_attention_heads, + attention_head_dim + ], + dtype=dtype, + device=device), + "v": + torch.zeros([ + batch_size, max_text_len, num_attention_heads, + attention_head_dim + ], + dtype=dtype, + device=device), + "is_init": + False, + }) + self.crossattn_cache = crossattn_cache diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index 8f08b22a6..6cffa10bb 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -24,6 +24,7 @@ from fastvideo.logger import init_logger from fastvideo.models.loader.component_loader import TransformerLoader from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( FlowMatchEulerDiscreteScheduler) +from fastvideo.models.utils import pred_noise_to_pred_video from fastvideo.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.pipelines.stages.base import PipelineStage from fastvideo.pipelines.stages.validators import StageValidators as V @@ -778,8 +779,6 @@ class DmdDenoisingStage(DenoisingStage): **pos_cond_kwargs, ).permute(0, 2, 1, 3, 4) - from fastvideo.training.training_utils import ( - pred_noise_to_pred_video) pred_video = pred_noise_to_pred_video( pred_noise=pred_noise.flatten(0, 1), noise_input_latent=noise_latents.flatten(0, 1), diff --git a/fastvideo/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py index 53676deb0..47381a91e 100644 --- a/fastvideo/training/distillation_pipeline.py +++ b/fastvideo/training/distillation_pipeline.py @@ -29,6 +29,7 @@ from fastvideo.forward_context import set_forward_context from fastvideo.logger import init_logger from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( FlowMatchEulerDiscreteScheduler) +from fastvideo.models.utils import pred_noise_to_pred_video from fastvideo.pipelines import (ComposedPipelineBase, ForwardBatch, TrainingBatch) from fastvideo.training.activation_checkpoint import ( @@ -36,8 +37,7 @@ from fastvideo.training.activation_checkpoint import ( from fastvideo.training.training_pipeline import TrainingPipeline from fastvideo.training.training_utils import ( clip_grad_norm_while_handling_failing_dtensor_cases, get_scheduler, - load_distillation_checkpoint, pred_noise_to_pred_video, - save_distillation_checkpoint, shift_timestep) + load_distillation_checkpoint, save_distillation_checkpoint, shift_timestep) from fastvideo.utils import is_vsa_available, set_random_seed import wandb # isort: skip diff --git a/fastvideo/training/training_utils.py b/fastvideo/training/training_utils.py index dd794a197..bdeb81b4f 100644 --- a/fastvideo/training/training_utils.py +++ b/fastvideo/training/training_utils.py @@ -826,27 +826,6 @@ def custom_to_hf_state_dict( return new_state_dict -def pred_noise_to_pred_video(pred_noise: torch.Tensor, - noise_input_latent: torch.Tensor, - timestep: torch.Tensor, - scheduler: Any) -> torch.Tensor: - """ - Convert predicted noise to clean latent. - """ - timestep = timestep.expand(noise_input_latent.shape[0]) - dtype = pred_noise.dtype - device = pred_noise.device - pred_noise = pred_noise.float().to(device) - noise_input_latent = noise_input_latent.float().to(device) - sigmas = scheduler.sigmas.float().to(device) - timesteps = scheduler.timesteps.float().to(device) - timestep_id = torch.argmin( - (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) - sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1) - pred_video = noise_input_latent - sigma_t * pred_noise - return pred_video.to(dtype) - - def shift_timestep(timestep: torch.Tensor, shift: float, num_train_timestep: float) -> torch.Tensor: if shift == 1: diff --git a/scripts/checkpoint_conversion/wan_to_diffusers.py b/scripts/checkpoint_conversion/wan_to_diffusers.py new file mode 100644 index 000000000..ea52ed687 --- /dev/null +++ b/scripts/checkpoint_conversion/wan_to_diffusers.py @@ -0,0 +1,161 @@ +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 +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", +} + +# 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"] +new_state_dict = OrderedDict() +for k, v in state_dict.items(): + new_key = k + 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") + dummy_value = torch.zeros_like(v) + new_state_dict[dummy_key] = dummy_value +del state_dict + +save_torch_state_dict( + new_state_dict, + "new2/", + max_shard_size="10GB" +)