nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Wan2.1 causal, few-step inference runnable

Browse Source
This commit is contained in:
JerryZhou54
2025-08-28 03:56:37 +00:00
parent 1575102965
commit 772ead0d34
5 changed files with 42 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/configs/pipelines/wan.py
+2
View File
@@ -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])
fastvideo/layers/rotary_embedding.py
-1
View File
@@ -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(
fastvideo/models/dits/causal_wanvideo.py
+7 -16
View File
@@ -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,
fastvideo/pipelines/pipeline_batch_info.py
+4
View File
@@ -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."""
fastvideo/pipelines/stages/causal_denoising.py
+29 -27
View File
@@ -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