From 772ead0d34d48ed0da15505c8420dbf650976fc0 Mon Sep 17 00:00:00 2001 From: JerryZhou54 Date: Thu, 28 Aug 2025 01:20:57 +0000 Subject: [PATCH] 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