diff --git a/examples/distill/SFWan2.1-T2V/distill_dmd_t2v_1.3B.sh b/examples/distill/SFWan2.1-T2V/distill_dmd_t2v_1.3B.sh index 61d8a01f4..ee6ffe88e 100644 --- a/examples/distill/SFWan2.1-T2V/distill_dmd_t2v_1.3B.sh +++ b/examples/distill/SFWan2.1-T2V/distill_dmd_t2v_1.3B.sh @@ -16,23 +16,23 @@ export NCCL_P2P_DISABLE=1 export TORCH_NCCL_ENABLE_MONITORING=0 # different cache dir for different processes export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID} -export MASTER_PORT=29503 +export MASTER_PORT=29501 export TOKENIZERS_PARALLELISM=false -export WANDB_API_KEY="2f25ad37933894dbf0966c838c0b8494987f9f2f" +export WANDB_API_KEY="50632ebd88ffd970521cec9ab4a1a2d7e85bfc45" export WANDB_BASE_URL="https://api.wandb.ai" -export WANDB_MODE=online +export WANDB_MODE=offline export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN # Configs -NUM_GPUS=1 +NUM_GPUS=4 # Model paths for Self-Forcing DMD distillation: GENERATOR_MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers" REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers" # Teacher model FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # Critic model -DATA_DIR="data/test-text-preprocessing/Node_0_GPU_1_File_1/combined_parquet_dataset/" -VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json" +DATA_DIR="data/mixkit-64_processed/Node_0_GPU_1_File_1/combined_parquet_dataset" +VALIDATION_DATASET_FILE="data/mixkit-64_processed/validation.json" # export CUDA_VISIBLE_DEVICES=4,5 # IP=[MASTER NODE IP] @@ -84,7 +84,7 @@ dataset_args=( validation_args=( --log_validation --validation_dataset_file "$VALIDATION_DATASET_FILE" - --validation_steps 50 + --validation_steps 100 --validation_sampling_steps "4" --validation_guidance_scale "6.0" # not used for dmd inference ) @@ -129,7 +129,7 @@ dmd_args=( # Self-forcing specific arguments self_forcing_args=( --independent_first_frame False # Whether to treat first frame independently - --same_step_across_blocks False # Whether to use same denoising step across all blocks + --same_step_across_blocks True # Whether to use same denoising step across all blocks --last_step_only False # Whether to only use the last denoising step --context_noise 0 # Amount of noise to add during context caching (0 = no noise) --validate_cache_structure False # Set to True for debugging KV cache issues diff --git a/examples/inference/basic/basic.py b/examples/inference/basic/basic.py index 316621eaf..67d94ca76 100644 --- a/examples/inference/basic/basic.py +++ b/examples/inference/basic/basic.py @@ -9,9 +9,9 @@ def main(): # If a local path is provided, FastVideo will make a best effort # attempt to identify the optimal arguments. generator = VideoGenerator.from_pretrained( - "Wan-AI/Wan2.1-T2V-14B-Diffusers", + "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", # FastVideo will automatically handle distributed setup - num_gpus=4, + num_gpus=1, use_fsdp_inference=True, dit_cpu_offload=False, vae_cpu_offload=False, @@ -25,7 +25,9 @@ def main(): # sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg" # Generate videos with the same simple API, regardless of GPU count prompt = ( - "A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open." + "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) # video = generator.generate_video(prompt, sampling_param=sampling_param, output_path="wan_t2v_videos/") @@ -33,9 +35,13 @@ def main(): # Generate another video with a different prompt, without reloading the # model! prompt2 = ( - "The video shows a green and orange object being flattened as if it were under a hydraulic press, with the press moving down and compressing the object") + "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) if __name__ == "__main__": - main() + main() \ No newline at end of file diff --git a/fastvideo/layers/layernorm.py b/fastvideo/layers/layernorm.py index d93e1d966..a409ddfac 100644 --- a/fastvideo/layers/layernorm.py +++ b/fastvideo/layers/layernorm.py @@ -212,9 +212,9 @@ class ScaleResidualLayerNormScaleShift(nn.Module): frame_seqlen = normalized.shape[1] // num_frames modulated = ( normalized.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * - (1.0 + scale) + shift).flatten(1, 2) + (1 + scale) + shift).flatten(1, 2) else: - modulated = normalized * (1.0 + scale) + shift + modulated = normalized * (1 + scale) + shift return modulated, residual_output @@ -267,11 +267,11 @@ class LayerNormScaleShift(nn.Module): frame_seqlen = normalized.shape[1] // num_frames output = ( normalized.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * - (1.0 + scale) + shift).flatten(1, 2) + (1 + scale) + shift).flatten(1, 2) else: # scale.shape: [batch_size, 1, inner_dim] # shift.shape: [batch_size, 1, inner_dim] - output = normalized * (1.0 + scale) + shift + output = normalized * (1 + scale) + shift if self.compute_dtype == torch.float32: output = output.to(x.dtype) diff --git a/fastvideo/models/dits/causal_wanvideo.py b/fastvideo/models/dits/causal_wanvideo.py index 0656b9740..ac78d3223 100644 --- a/fastvideo/models/dits/causal_wanvideo.py +++ b/fastvideo/models/dits/causal_wanvideo.py @@ -147,8 +147,8 @@ class CausalWanSelfAttention(nn.Module): # 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"] = kv_cache["k"].detach() - kv_cache["v"] = kv_cache["v"].detach() + # kv_cache["k"] = kv_cache["k"].detach() + # kv_cache["v"] = kv_cache["v"].detach() # logger.info("kv_cache['k'] is in comp graph: %s", kv_cache["k"].requires_grad or kv_cache["k"].grad_fn is not None) kv_cache["k"][:, local_start_index:local_end_index] = roped_key kv_cache["v"][:, local_start_index:local_end_index] = v @@ -179,7 +179,7 @@ class CausalWanTransformerBlock(nn.Module): super().__init__() # 1. Self-attention - self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.norm1 = nn.LayerNorm(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) @@ -212,8 +212,7 @@ class CausalWanTransformerBlock(nn.Module): norm_type="layer", eps=eps, elementwise_affine=True, - dtype=torch.float32, - compute_dtype=torch.float32) + dtype=torch.float32) # 2. Cross-attention # Only T2V for now @@ -226,8 +225,7 @@ class CausalWanTransformerBlock(nn.Module): norm_type="layer", eps=eps, elementwise_affine=False, - dtype=torch.float32, - compute_dtype=torch.float32) + dtype=torch.float32) # 3. Feed-forward self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh") @@ -252,29 +250,34 @@ class CausalWanTransformerBlock(nn.Module): if hidden_states.dim() == 4: hidden_states = hidden_states.squeeze(1) num_frames = temb.shape[1] - frame_seqlen = hidden_states.shape[1] // num_frames + frame_seqlen = hidden_states.shape[1] // num_frames bs, seq_length, _ = hidden_states.shape orig_dtype = hidden_states.dtype # assert orig_dtype != torch.float32 - e = self.scale_shift_table + temb.float() + e = self.scale_shift_table + temb # e.shape: [batch_size, num_frames, 6, inner_dim] assert e.shape == (bs, num_frames, 6, self.hidden_dim) shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk( 6, dim=2) # *_msa.shape: [batch_size, num_frames, 1, inner_dim] - assert shift_msa.dtype == torch.float32 + # assert shift_msa.dtype == torch.float32 + + # logger.info("temb sum: %s, dtype: %s", temb.float().sum().item(), temb.dtype) + # logger.info("scale_msa sum: %s, dtype: %s", scale_msa.float().sum().item(), scale_msa.dtype) + # logger.info("shift_msa sum: %s, dtype: %s", shift_msa.float().sum().item(), shift_msa.dtype) # 1. Self-attention - norm_hidden_states = (self.norm1(hidden_states.float()).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * - (1 + scale_msa) + shift_msa).flatten(1, 2).to(orig_dtype) + norm_hidden_states = (self.norm1(hidden_states).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * + (1 + scale_msa) + shift_msa).flatten(1, 2) + # logger.info("norm_hidden_states sum: %s, shape: %s", norm_hidden_states.float().sum().item(), norm_hidden_states.shape) 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) + query = self.norm_q.forward_native(query) if self.norm_k is not None: - key = self.norm_k(key) + key = self.norm_k.forward_native(key) query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) @@ -288,8 +291,6 @@ class CausalWanTransformerBlock(nn.Module): 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, @@ -298,13 +299,10 @@ class CausalWanTransformerBlock(nn.Module): 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 @@ -367,8 +365,7 @@ class CausalWanTransformer3DModel(BaseDiT): norm_type="layer", eps=config.eps, elementwise_affine=False, - dtype=torch.float32, - compute_dtype=torch.float32) + dtype=torch.float32) self.proj_out = nn.Linear( inner_dim, config.out_channels * math.prod(config.patch_size)) self.scale_shift_table = nn.Parameter( @@ -490,12 +487,16 @@ class CausalWanTransformer3DModel(BaseDiT): ) 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 + freqs_cis = (freqs_cos, + freqs_sin) if freqs_cos is not None else None hidden_states = self.patch_embedding(hidden_states) + grid_sizes = torch.stack( + [torch.tensor(hidden_states[0].shape[1:], dtype=torch.long)]) hidden_states = hidden_states.flatten(2).transpose(1, 2) + encoder_hidden_states = torch.cat([encoder_hidden_states, encoder_hidden_states.new_zeros(1, self.text_len - encoder_hidden_states.size(1), encoder_hidden_states.size(2))], dim=1) + temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image) timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0, sizes=timestep.shape) @@ -542,14 +543,9 @@ class CausalWanTransformer3DModel(BaseDiT): 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) + output = self.unpatchify(hidden_states, grid_sizes) - return output + return torch.stack(output) def _forward_train(self, hidden_states: torch.Tensor, @@ -658,3 +654,30 @@ class CausalWanTransformer3DModel(BaseDiT): return self._forward_inference(*args, **kwargs) else: return self._forward_train(*args, **kwargs) + + + def unpatchify(self, x, grid_sizes): + r""" + + + Args: + x (List[Tensor]): + List of patchified features, each with shape [L, C_out * prod(patch_size)] + grid_sizes (Tensor): + Original spatial-temporal grid dimensions before patching, + + + Returns: + Tensor: + Reconstructed video tensors with shape [B, C_out, F, H / 8, W / 8] + """ + + c = self.out_channels + out = [] + for u, v in zip(x, grid_sizes.tolist()): + u = u[:math.prod(v)].view(*v, *self.patch_size, c) + u = u.permute(6, 0, 3, 1, 4, 2, 5) + # u = torch.einsum('fhwpqrc->cfphqwr', u.contiguous()) + u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)]) + out.append(u) + return out diff --git a/fastvideo/models/dits/wanvideo.py b/fastvideo/models/dits/wanvideo.py index ac79e62e8..3262a9a1e 100644 --- a/fastvideo/models/dits/wanvideo.py +++ b/fastvideo/models/dits/wanvideo.py @@ -1,3 +1,5 @@ +import torch +import torch.nn as nn # SPDX-License-Identifier: Apache-2.0 import math @@ -37,16 +39,14 @@ class WanImageEmbedding(torch.nn.Module): def __init__(self, in_features: int, out_features: int): super().__init__() - self.norm1 = FP32LayerNorm(in_features) + self.norm1 = nn.LayerNorm(in_features) self.ff = MLP(in_features, in_features, out_features, act_type="gelu") - self.norm2 = FP32LayerNorm(out_features) + self.norm2 = nn.LayerNorm(out_features) - def forward(self, - encoder_hidden_states_image: torch.Tensor) -> torch.Tensor: - dtype = encoder_hidden_states_image.dtype + def forward(self, encoder_hidden_states_image: torch.Tensor) -> torch.Tensor: hidden_states = self.norm1(encoder_hidden_states_image) hidden_states = self.ff(hidden_states) - hidden_states = self.norm2(hidden_states).to(dtype) + hidden_states = self.norm2(hidden_states) return hidden_states @@ -62,7 +62,7 @@ class WanTimeTextImageEmbedding(nn.Module): super().__init__() self.time_embedder = TimestepEmbedder( - dim, frequency_embedding_size=time_freq_dim, act_layer="silu") + dim, frequency_embedding_size=time_freq_dim, act_layer="silu", freq_dtype=torch.float64) self.time_modulation = ModulateProjection(dim, factor=6, act_layer="silu") @@ -156,12 +156,12 @@ class WanT2VCrossAttention(WanSelfAttention): 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) + q = self.norm_q.forward_native(self.to_q(x)[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) + k = self.norm_k.forward_native(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 @@ -169,7 +169,7 @@ class WanT2VCrossAttention(WanSelfAttention): k = crossattn_cache["k"] v = crossattn_cache["v"] else: - k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d) + k = self.norm_k.forward_native(self.to_k(context)[0]).view(b, -1, n, d) v = self.to_v(context)[0].view(b, -1, n, d) # compute attention @@ -213,10 +213,10 @@ class WanI2VCrossAttention(WanSelfAttention): 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) + q = self.norm_q.forward_native(self.to_q(x)[0]).view(b, -1, n, d) + k = self.norm_k.forward_native(self.to_k(context)[0]).view(b, -1, n, d) v = self.to_v(context)[0].view(b, -1, n, d) - k_img = self.norm_added_k(self.add_k_proj(context_img)[0]).view( + k_img = self.norm_added_k.forward_native(self.add_k_proj(context_img)[0]).view( b, -1, n, d) v_img = self.add_v_proj(context_img)[0].view(b, -1, n, d) img_x = self.attn(q, k_img, v_img) @@ -247,7 +247,7 @@ class WanTransformerBlock(nn.Module): super().__init__() # 1. Self-attention - self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.norm1 = nn.LayerNorm(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) @@ -278,29 +278,29 @@ class WanTransformerBlock(nn.Module): norm_type="layer", eps=eps, elementwise_affine=True, - dtype=torch.float32, - compute_dtype=torch.float32) + 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, + self.attn2 = WanI2VCrossAttention(dim, + num_heads, + qk_norm=qk_norm, eps=eps) + else: # T2V - self.attn2 = WanT2VCrossAttention(dim, - num_heads, - qk_norm=qk_norm, + 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) + dim, + norm_type="layer", + eps=eps, + elementwise_affine=False, + dtype=torch.float32) # 3. Feed-forward self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh") @@ -319,12 +319,11 @@ class WanTransformerBlock(nn.Module): hidden_states = hidden_states.squeeze(1) bs, seq_length, _ = hidden_states.shape orig_dtype = hidden_states.dtype - # assert orig_dtype != torch.float32 if temb.dim() == 4: # temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v) shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( - self.scale_shift_table.unsqueeze(0) + temb.float() + self.scale_shift_table.unsqueeze(0) + temb ).chunk(6, dim=2) # batch_size, seq_len, 1, inner_dim shift_msa = shift_msa.squeeze(2) @@ -335,22 +334,20 @@ class WanTransformerBlock(nn.Module): c_gate_msa = c_gate_msa.squeeze(2) else: # temb: batch_size, 6, inner_dim (wan2.1/wan2.2 14B) - e = self.scale_shift_table + temb.float() + e = self.scale_shift_table + temb 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) + norm_hidden_states = self.norm1(hidden_states) * (1 + scale_msa) + shift_msa 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) + query = self.norm_q.forward_native(query) if self.norm_k is not None: - key = self.norm_k(key) + key = self.norm_k.forward_native(key) query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) @@ -370,26 +367,20 @@ class WanTransformerBlock(nn.Module): 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, + 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 WanTransformerBlock_VSA(nn.Module): def __init__(self, @@ -406,7 +397,7 @@ class WanTransformerBlock_VSA(nn.Module): super().__init__() # 1. Self-attention - self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.norm1 = nn.LayerNorm(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) @@ -438,8 +429,7 @@ class WanTransformerBlock_VSA(nn.Module): norm_type="layer", eps=eps, elementwise_affine=True, - dtype=torch.float32, - compute_dtype=torch.float32) + dtype=torch.float32) # 2. Cross-attention if added_kv_proj_dim is not None: @@ -459,8 +449,7 @@ class WanTransformerBlock_VSA(nn.Module): norm_type="layer", eps=eps, elementwise_affine=False, - dtype=torch.float32, - compute_dtype=torch.float32) + dtype=torch.float32) # 3. Feed-forward self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh") @@ -480,23 +469,22 @@ class WanTransformerBlock_VSA(nn.Module): bs, seq_length, _ = hidden_states.shape orig_dtype = hidden_states.dtype # assert orig_dtype != torch.float32 - e = self.scale_shift_table + temb.float() + e = self.scale_shift_table + temb 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) + norm_hidden_states = (self.norm1(hidden_states) * + (1 + scale_msa) + shift_msa) query, _ = self.to_q(norm_hidden_states) key, _ = self.to_k(norm_hidden_states) value, _ = self.to_v(norm_hidden_states) gate_compress, _ = self.to_gate_compress(norm_hidden_states) if self.norm_q is not None: - query = self.norm_q(query) + query = self.norm_q.forward_native(query) if self.norm_k is not None: - key = self.norm_k(key) + key = self.norm_k.forward_native(key) query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) @@ -521,8 +509,6 @@ class WanTransformerBlock_VSA(nn.Module): 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, @@ -530,17 +516,15 @@ class WanTransformerBlock_VSA(nn.Module): 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 WanTransformer3DModel(CachableDiT): _fsdp_shard_conditions = WanVideoConfig()._fsdp_shard_conditions _compile_conditions = WanVideoConfig()._compile_conditions @@ -598,8 +582,7 @@ class WanTransformer3DModel(CachableDiT): norm_type="layer", eps=config.eps, elementwise_affine=False, - dtype=torch.float32, - compute_dtype=torch.float32) + dtype=torch.float32) self.proj_out = nn.Linear( inner_dim, config.out_channels * math.prod(config.patch_size)) self.scale_shift_table = nn.Parameter( @@ -659,10 +642,12 @@ class WanTransformer3DModel(CachableDiT): 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 + freqs_cis = (freqs_cos, + freqs_sin) if freqs_cos is not None else None hidden_states = self.patch_embedding(hidden_states) + grid_sizes = torch.stack( + [torch.tensor(hidden_states[0].shape[1:], dtype=torch.long)]) hidden_states = hidden_states.flatten(2).transpose(1, 2) # timestep shape: batch_size, or batch_size, seq_len (wan 2.2 ti2v) @@ -672,6 +657,8 @@ class WanTransformer3DModel(CachableDiT): else: ts_seq_len = None + encoder_hidden_states = torch.cat([encoder_hidden_states, encoder_hidden_states.new_zeros(1, self.text_len - encoder_hidden_states.size(1), encoder_hidden_states.size(2))], dim=1) + temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( timestep, encoder_hidden_states, encoder_hidden_states_image, timestep_seq_len=ts_seq_len) if ts_seq_len is not None: @@ -728,14 +715,35 @@ class WanTransformer3DModel(CachableDiT): 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) + output = self.unpatchify(hidden_states, grid_sizes) - return output + return torch.stack(output) + + def unpatchify(self, x, grid_sizes): + r""" + + + Args: + x (List[Tensor]): + List of patchified features, each with shape [L, C_out * prod(patch_size)] + grid_sizes (Tensor): + Original spatial-temporal grid dimensions before patching, + + + Returns: + Tensor: + Reconstructed video tensors with shape [B, C_out, F, H / 8, W / 8] + """ + + c = self.out_channels + out = [] + for u, v in zip(x, grid_sizes.tolist()): + u = u[:math.prod(v)].view(*v, *self.patch_size, c) + u = u.permute(6, 0, 3, 1, 4, 2, 5) + # u = torch.einsum('fhwpqrc->cfphqwr', u.contiguous()) + u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)]) + out.append(u) + return out def maybe_cache_states(self, hidden_states: torch.Tensor, original_hidden_states: torch.Tensor) -> None: @@ -827,5 +835,4 @@ class WanTransformer3DModel(CachableDiT): if self.is_even: return hidden_states + self.previous_residual_even else: - return hidden_states + self.previous_residual_odd - \ No newline at end of file + return hidden_states + self.previous_residual_odd \ No newline at end of file diff --git a/fastvideo/tests/transformers/test_ori_causal_wanvideo.py b/fastvideo/tests/transformers/test_ori_causal_wanvideo.py new file mode 100644 index 000000000..831831d07 --- /dev/null +++ b/fastvideo/tests/transformers/test_ori_causal_wanvideo.py @@ -0,0 +1,292 @@ +# SPDX-License-Identifier: Apache-2.0 +import os +import math + +import numpy as np +import pytest +import torch +from fastvideo.wan.modules.causal_model import CausalWanModel + +from fastvideo.configs.pipelines import PipelineConfig +from fastvideo.forward_context import set_forward_context +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.models.loader.component_loader import TransformerLoader +from fastvideo.models.dits.causal_wanvideo import CausalWanTransformer3DModel +from fastvideo.utils import maybe_download_model +from fastvideo.configs.models.dits import WanVideoConfig +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + + +logger = init_logger(__name__) + +os.environ["MASTER_ADDR"] = "localhost" +os.environ["MASTER_PORT"] = "29503" + +BASE_MODEL_PATH = "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers" +MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, + local_dir=os.path.join( + 'data', BASE_MODEL_PATH)) +TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer") + + +@pytest.mark.usefixtures("distributed_setup") +def test_ori_causal_wan_transformer(): + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + precision = torch.float32 + precision_str = "fp32" + args = FastVideoArgs(model_path=TRANSFORMER_PATH, + dit_cpu_offload=True, + pipeline_config=PipelineConfig(dit_config=WanVideoConfig(), dit_precision=precision_str)) + args.device = device + + loader = TransformerLoader() + model2 = loader.load(TRANSFORMER_PATH, args).to(dtype=precision) + + model1 = CausalWanModel.from_pretrained( + "/mnt/weka/home/hao.zhang/wei/Self-Forcing/wan_models/Wan2.1-T2V-1.3B", device=device, + torch_dtype=precision).to(device, dtype=precision).requires_grad_(False) + causal_state_dict = torch.load("/mnt/weka/home/hao.zhang/wei/Self-Forcing/checkpoints/self_forcing_dmd.pt")["generator_ema"] + new_state_dict = {} + for k, v in causal_state_dict.items(): + if k.startswith("model."): + new_state_dict[k.replace("model.", "")] = v + causal_state_dict = new_state_dict + model1.load_state_dict(causal_state_dict) + + total_params = sum(p.numel() for p in model1.parameters()) + # Calculate weight sum for model1 (converting to float64 to avoid overflow) + weight_sum_model1 = sum( + p.to(torch.float64).sum().item() for p in model1.parameters()) + # Also calculate mean for more stable comparison + weight_mean_model1 = weight_sum_model1 / total_params + logger.info("Model 1 weight sum: %s", weight_sum_model1) + logger.info("Model 1 weight mean: %s", weight_mean_model1) + + # Calculate weight sum for model2 (converting to float64 to avoid overflow) + total_params_model2 = sum(p.numel() for p in model2.parameters()) + weight_sum_model2 = sum( + p.to(torch.float64).sum().item() for p in model2.parameters()) + # Also calculate mean for more stable comparison + weight_mean_model2 = weight_sum_model2 / total_params_model2 + logger.info("Model 2 weight sum: %s", weight_sum_model2) + logger.info("Model 2 weight mean: %s", weight_mean_model2) + + weight_sum_diff = abs(weight_sum_model1 - weight_sum_model2) + logger.info("Weight sum difference: %s", weight_sum_diff) + weight_mean_diff = abs(weight_mean_model1 - weight_mean_model2) + logger.info("Weight mean difference: %s", weight_mean_diff) + + # Set both models to eval mode + model1 = model1.eval() + model2 = model2.eval() + + # Create identical inputs for both models + batch_size = 1 + text_seq_len = 30 + + # Video latents [B, C, T, H, W] + hidden_states = torch.randn(batch_size, + 16, + 12, + 160, + 90, + device=device, + dtype=precision) + block_sizes = [3 for _ in range(4)] + timesteps = [1000, 750, 500, 250] + + # Text embeddings [B, L, D] (including global token) + encoder_hidden_states = torch.randn(batch_size, + text_seq_len + 1, + 4096, + device=device, + dtype=precision) + + output1 = _causal_inference(model1, hidden_states.clone(), encoder_hidden_states.clone(), block_sizes, timesteps, precision) + logger.info("Finish inference for model1") + output2 = _causal_inference(model2, hidden_states.clone(), encoder_hidden_states.clone(), block_sizes, timesteps, precision) + + # Check if outputs have the same shape + assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}" + assert output1.dtype == output2.dtype, f"Output dtype don't match: {output1.dtype} vs {output2.dtype}" + + logger.info("Output 1 Sum: %s", output1.float().sum().item()) + logger.info("Output 2 Sum: %s", output2.float().sum().item()) + + # Check if outputs are similar (allowing for small numerical differences) + max_diff = torch.max(torch.abs(output1 - output2)) + mean_diff = torch.mean(torch.abs(output1 - output2)) + logger.info("Max Diff: %s", max_diff.item()) + logger.info("Mean Diff: %s", mean_diff.item()) + assert max_diff < 1e-4, f"Maximum difference between outputs: {max_diff.item()}" + # mean diff + assert mean_diff < 1e-4, f"Mean difference between outputs: {mean_diff.item()}" + +def _causal_inference(transformer, latents, prompt_embeds, block_sizes, timesteps, target_dtype): + forward_batch = ForwardBatch( + data_type="dummy", + ) + start_index = 0 + pos_start_base = 0 + frame_seq_length = latents.shape[-1] * latents.shape[-2] // (WanVideoConfig().arch_config.patch_size[-1] * WanVideoConfig().arch_config.patch_size[-2]) + seq_len = frame_seq_length * latents.shape[2] + kv_cache1 = _initialize_kv_cache(transformer, batch_size=latents.shape[0], + kv_cache_size=frame_seq_length * latents.shape[2], + dtype=target_dtype, + device=latents.device) + crossattn_cache = _initialize_crossattn_cache( + transformer, + batch_size=latents.shape[0], + max_text_len=WanVideoConfig().arch_config.text_len, + dtype=target_dtype, + device=latents.device) + for current_num_frames, t_cur in zip(block_sizes, timesteps): + # logger.info(f"Current frame idx: {start_index}, Current timestep: {t_cur}") + # logger.info(f"k cache sum: {sum(kv_cache['k'].float().sum().item() for kv_cache in kv_cache1)}, v cache sum: {sum(kv_cache['v'].float().sum().item() for kv_cache in kv_cache1)}") + # logger.info(f"latents sum: {latents.float().sum().item()}, encoder_hidden_states sum: {prompt_embeds.float().sum().item()}") + current_latents = latents[:, :, start_index:start_index + + current_num_frames, :, :] + + attn_metadata = None + + with set_forward_context(current_timestep=0, + attn_metadata=attn_metadata, + forward_batch=forward_batch): + # Run transformer; follow DMD stage pattern + t_expanded_noise = t_cur * torch.ones( + (current_latents.shape[0], 1), + device=current_latents.device, + dtype=torch.long) + if isinstance(transformer, CausalWanModel): + pred_noise_btchw = transformer( + x=[current_latents[0]], + context=[prompt_embeds[0]], + t=t_expanded_noise, + seq_len=seq_len, + kv_cache=kv_cache1, + crossattn_cache=crossattn_cache, + current_start=(pos_start_base + start_index) * + frame_seq_length + ) + elif isinstance(transformer, CausalWanTransformer3DModel): + pred_noise_btchw = transformer( + current_latents, + prompt_embeds, + t_expanded_noise, + kv_cache=kv_cache1, + crossattn_cache=crossattn_cache, + current_start=(pos_start_base + start_index) * + frame_seq_length, + start_frame=start_index + ) + + # Write back and advance + latents[:, :, start_index:start_index + + current_num_frames, :, :] = pred_noise_btchw.clone() + + # Re-run with context timestep to update KV cache using clean context + context_noise = 0 + t_context = torch.ones([latents.shape[0]], + device=latents.device, + dtype=torch.long) * int(context_noise) + context_bcthw = pred_noise_btchw.to(target_dtype) + with set_forward_context(current_timestep=0, + attn_metadata=attn_metadata, + forward_batch=forward_batch): + t_expanded_context = t_context.unsqueeze(1) + if isinstance(transformer, CausalWanModel): + _ = transformer( + x=[context_bcthw[0]], + context=[prompt_embeds[0]], + t=t_expanded_context, + seq_len=seq_len, + kv_cache=kv_cache1, + crossattn_cache=crossattn_cache, + current_start=(pos_start_base + start_index) * + frame_seq_length + ) + elif isinstance(transformer, CausalWanTransformer3DModel): + _ = transformer( + context_bcthw, + prompt_embeds, + t_expanded_context, + kv_cache=kv_cache1, + crossattn_cache=crossattn_cache, + current_start=(pos_start_base + start_index) * + frame_seq_length, + start_frame=start_index + ) + start_index += current_num_frames + + return latents + +def _initialize_kv_cache(transformer, batch_size, kv_cache_size, dtype, device) -> None: + """ + Initialize a Per-GPU KV cache aligned with the Wan model assumptions. + """ + kv_cache1 = [] + if isinstance(transformer, CausalWanModel): + num_attention_heads = transformer.num_heads + attention_head_dim = transformer.dim // transformer.num_heads + elif isinstance(transformer, CausalWanTransformer3DModel): + num_attention_heads = transformer.num_attention_heads + attention_head_dim = transformer.attention_head_dim + + for _ in range(len(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), + }) + + return kv_cache1 + +def _initialize_crossattn_cache(transformer, batch_size, max_text_len, dtype, + device) -> None: + """ + Initialize a Per-GPU cross-attention cache aligned with the Wan model assumptions. + """ + crossattn_cache = [] + if isinstance(transformer, CausalWanModel): + num_attention_heads = transformer.num_heads + attention_head_dim = transformer.dim // transformer.num_heads + elif isinstance(transformer, CausalWanTransformer3DModel): + num_attention_heads = transformer.num_attention_heads + attention_head_dim = transformer.attention_head_dim + + for _ in range(len(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, + }) + return crossattn_cache \ No newline at end of file diff --git a/fastvideo/tests/transformers/test_ori_wanvideo.py b/fastvideo/tests/transformers/test_ori_wanvideo.py new file mode 100644 index 000000000..e3611bcf0 --- /dev/null +++ b/fastvideo/tests/transformers/test_ori_wanvideo.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +import os +import math + +import numpy as np +import pytest +import torch +from fastvideo.wan.modules.model import WanModel + +from fastvideo.configs.pipelines import PipelineConfig +from fastvideo.forward_context import set_forward_context +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.models.loader.component_loader import TransformerLoader +from fastvideo.utils import maybe_download_model +from fastvideo.configs.models.dits import WanVideoConfig +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + + +logger = init_logger(__name__) + +os.environ["MASTER_ADDR"] = "localhost" +os.environ["MASTER_PORT"] = "29503" + +BASE_MODEL_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, + local_dir=os.path.join( + 'data', BASE_MODEL_PATH)) +TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer") + + +@pytest.mark.usefixtures("distributed_setup") +def test_ori_wan_transformer(): + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + precision = torch.bfloat16 + precision_str = "bf16" + args = FastVideoArgs(model_path=TRANSFORMER_PATH, + dit_cpu_offload=True, + pipeline_config=PipelineConfig(dit_config=WanVideoConfig(), dit_precision=precision_str)) + args.device = device + + loader = TransformerLoader() + model2 = loader.load(TRANSFORMER_PATH, args).to(dtype=precision) + + model1 = WanModel.from_pretrained( + "/mnt/weka/home/hao.zhang/wei/Self-Forcing/wan_models/Wan2.1-T2V-1.3B", device=device, + torch_dtype=precision).to(device, dtype=precision).requires_grad_(False) + + total_params = sum(p.numel() for p in model1.parameters()) + # Calculate weight sum for model1 (converting to float64 to avoid overflow) + weight_sum_model1 = sum( + p.to(torch.float64).sum().item() for p in model1.parameters()) + # Also calculate mean for more stable comparison + weight_mean_model1 = weight_sum_model1 / total_params + logger.info("Model 1 weight sum: %s", weight_sum_model1) + logger.info("Model 1 weight mean: %s", weight_mean_model1) + + # Calculate weight sum for model2 (converting to float64 to avoid overflow) + total_params_model2 = sum(p.numel() for p in model2.parameters()) + weight_sum_model2 = sum( + p.to(torch.float64).sum().item() for p in model2.parameters()) + # Also calculate mean for more stable comparison + weight_mean_model2 = weight_sum_model2 / total_params_model2 + logger.info("Model 2 weight sum: %s", weight_sum_model2) + logger.info("Model 2 weight mean: %s", weight_mean_model2) + + weight_sum_diff = abs(weight_sum_model1 - weight_sum_model2) + logger.info("Weight sum difference: %s", weight_sum_diff) + weight_mean_diff = abs(weight_mean_model1 - weight_mean_model2) + logger.info("Weight mean difference: %s", weight_mean_diff) + + # Set both models to eval mode + model1 = model1.eval() + model2 = model2.eval() + + # Create identical inputs for both models + batch_size = 1 + text_seq_len = 30 + seq_len = math.ceil((160 * 90) / + (2 * 2) * + 21) + + # Video latents [B, C, T, H, W] + hidden_states = torch.randn(batch_size, + 16, + 21, + 160, + 90, + device=device, + dtype=precision) + + # Text embeddings [B, L, D] (including global token) + encoder_hidden_states = torch.randn(batch_size, + text_seq_len + 1, + 4096, + device=device, + dtype=precision) + + # Timestep + timestep = torch.tensor([500], device=device, dtype=precision) + + forward_batch = ForwardBatch( + data_type="dummy", + ) + + # with torch.amp.autocast('cuda', dtype=precision): + output1 = model1( + x=[hidden_states[0]], + context=[encoder_hidden_states[0]], + t=timestep, + seq_len=seq_len, + ) + with set_forward_context( + current_timestep=0, + attn_metadata=None, + forward_batch=forward_batch, + ): + output2 = model2(hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep) + + # Check if outputs have the same shape + assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}" + assert output1.dtype == output2.dtype, f"Output dtype don't match: {output1.dtype} vs {output2.dtype}" + + # Check if outputs are similar (allowing for small numerical differences) + max_diff = torch.max(torch.abs(output1 - output2)) + mean_diff = torch.mean(torch.abs(output1 - output2)) + logger.info("Max Diff: %s", max_diff.item()) + logger.info("Mean Diff: %s", mean_diff.item()) + assert max_diff < 1e-4, f"Maximum difference between outputs: {max_diff.item()}" + # mean diff + assert mean_diff < 1e-4, f"Mean difference between outputs: {mean_diff.item()}" \ No newline at end of file diff --git a/fastvideo/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py index 085eb480b..7ca40a3bb 100644 --- a/fastvideo/training/distillation_pipeline.py +++ b/fastvideo/training/distillation_pipeline.py @@ -730,12 +730,11 @@ class DistillationPipeline(TrainingPipeline): "encoder_hidden_states": training_batch.encoder_hidden_states, "encoder_attention_mask": training_batch.encoder_attention_mask, } - if getattr(self, "negative_prompt_embeds", None) is not None: - unconditional_dict = { - "encoder_hidden_states": self.negative_prompt_embeds, - "encoder_attention_mask": self.negative_prompt_attention_mask, - } - training_batch.unconditional_dict = unconditional_dict + unconditional_dict = { + "encoder_hidden_states": self.negative_prompt_embeds, + "encoder_attention_mask": self.negative_prompt_attention_mask, + } + training_batch.unconditional_dict = unconditional_dict training_batch.dmd_latent_vis_dict = {} training_batch.fake_score_latent_vis_dict = {} diff --git a/fastvideo/training/self_forcing_distillation_pipeline.py b/fastvideo/training/self_forcing_distillation_pipeline.py index e4d563d9b..06c882961 100644 --- a/fastvideo/training/self_forcing_distillation_pipeline.py +++ b/fastvideo/training/self_forcing_distillation_pipeline.py @@ -319,7 +319,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline): encoder_hidden_states_image=training_batch_temp.input_kwargs.get('encoder_hidden_states_image'), kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, - current_start=current_start_frame * self.frame_seq_length + current_start=current_start_frame * self.frame_seq_length, + start_frame=current_start_frame ).permute(0, 2, 1, 3, 4) denoised_pred = pred_noise_to_pred_video( @@ -348,7 +349,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline): encoder_hidden_states_image=training_batch_temp.input_kwargs.get('encoder_hidden_states_image'), kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, - current_start=current_start_frame * self.frame_seq_length + current_start=current_start_frame * self.frame_seq_length, + start_frame=current_start_frame ).permute(0, 2, 1, 3, 4) else: training_batch_temp = self._build_distill_input_kwargs( @@ -361,7 +363,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline): encoder_hidden_states_image=training_batch_temp.input_kwargs.get('encoder_hidden_states_image'), kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, - current_start=current_start_frame * self.frame_seq_length + current_start=current_start_frame * self.frame_seq_length, + start_frame=current_start_frame ).permute(0, 2, 1, 3, 4) denoised_pred = pred_noise_to_pred_video( @@ -393,7 +396,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline): encoder_hidden_states_image=training_batch_temp.input_kwargs.get('encoder_hidden_states_image'), kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, - current_start=current_start_frame * self.frame_seq_length + current_start=current_start_frame * self.frame_seq_length, + start_frame=current_start_frame ) # Step 3.4: update the start and end frame indices