diff --git a/examples/inference/basic/basic_self_forcing_causal.py b/examples/inference/basic/basic_self_forcing_causal.py index e2241a94e..9d716e9bc 100644 --- a/examples/inference/basic/basic_self_forcing_causal.py +++ b/examples/inference/basic/basic_self_forcing_causal.py @@ -12,20 +12,33 @@ def main(): generator = VideoGenerator.from_pretrained( model_name, # FastVideo will automatically handle distributed setup - num_gpus=1, + num_gpus=2, use_fsdp_inference=True, text_encoder_cpu_offload=False, dit_cpu_offload=False, + num_frame_per_block=4, ) sampling_param = SamplingParam.from_pretrained(model_name) - + sampling_param.num_frames = 77 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) + return + import time + start_time = time.perf_counter() + for i in range(10): + prompt = ( + "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " + "wide with interest. The playful yet serene atmosphere is complemented by soft " + "natural light filtering through the petals. Mid-shot, warm and cheerful tones." + ) + video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param) + end_time = time.perf_counter() + print(f"Time taken to generate 10 videos: {end_time - start_time} seconds, average time per video: {(end_time - start_time) / 10} seconds") if __name__ == "__main__": main() diff --git a/examples/inference/basic/basic_self_forcing_causal_wan2_2_t2v.py b/examples/inference/basic/basic_self_forcing_causal_wan2_2_t2v.py index 6c11689ec..bfab29fee 100644 --- a/examples/inference/basic/basic_self_forcing_causal_wan2_2_t2v.py +++ b/examples/inference/basic/basic_self_forcing_causal_wan2_2_t2v.py @@ -1,6 +1,7 @@ # NOTE: This is still a work in progress, and the checkpoints are not released yet. from fastvideo import VideoGenerator +import time # from fastvideo.configs.sample import SamplingParam @@ -10,6 +11,7 @@ def main(): # model. # If a local path is provided, FastVideo will make a best effort # attempt to identify the optimal arguments. + load_start_time = time.perf_counter() generator = VideoGenerator.from_pretrained( "rand0nmr/SFWan2.2-T2V-A14B-Diffusers", # FastVideo will automatically handle distributed setup @@ -23,9 +25,10 @@ def main(): pin_cpu_memory=True, init_weights_from_safetensors="/mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_cfg2/checkpoint-246_weight_only/generator_inference_transformer/", init_weights_from_safetensors_2="/mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_cfg2/checkpoint-246_weight_only/generator_2_inference_transformer/", - num_frame_per_block=7, - # image_encoder_cpu_offload=False, + num_frame_per_block=4, ) + load_end_time = time.perf_counter() + load_time = load_end_time - load_start_time # sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers") # sampling_param.num_frames = 45 @@ -36,8 +39,16 @@ def main(): "wide with interest. The playful yet serene atmosphere is complemented by soft " "natural light filtering through the petals. Mid-shot, warm and cheerful tones." ) - _ = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, num_frames=81) + _ = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=False, num_frames=77) + # return + start_time = time.perf_counter() + for _ in range(10): + video2 = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=False, num_frames=77) + end_time = time.perf_counter() + gen_time = (end_time - start_time) / 10 + print(f"Time taken to load model: {load_time} seconds") + print(f"Time taken to generate 10 videos: {gen_time} seconds") if __name__ == "__main__": main() \ No newline at end of file diff --git a/fastvideo/attention/layer.py b/fastvideo/attention/layer.py index 7b086a3c6..f8687b890 100644 --- a/fastvideo/attention/layer.py +++ b/fastvideo/attention/layer.py @@ -11,6 +11,9 @@ from fastvideo.distributed.parallel_state import (get_sp_parallel_rank, from fastvideo.forward_context import ForwardContext, get_forward_context from fastvideo.platforms import AttentionBackendEnum from fastvideo.utils import get_compute_dtype +from fastvideo.logger import init_logger + +logger = init_logger(__name__) class DistributedAttention(nn.Module): @@ -91,6 +94,7 @@ class DistributedAttention(nn.Module): ctx_attn_metadata = forward_context.attn_metadata # Stack QKV + logger.info(f"rank: {local_rank}, q: {q.shape}, k: {k.shape}, v: {v.shape}", local_main_process_only=False) qkv = torch.cat([q, k, v], dim=0) # [3, seq_len, num_heads, head_dim] # Redistribute heads across sequence dimension diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index 759d5f30b..1a25bd771 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -174,6 +174,9 @@ class FastVideoArgs: init_weights_from_safetensors: str = "" # path to safetensors file for initial weight loading init_weights_from_safetensors_2: str = "" # path to safetensors file for initial weight loading for transformer_2 + # Self-forcing specific arguments + num_frame_per_block: int = 3 + # # DMD parameters # dmd_denoising_steps: List[int] | None = field(default=None) @@ -433,6 +436,13 @@ class FastVideoArgs: type=str, help="Path to safetensors file for initial weight loading") + # Self-forcing specific arguments + parser.add_argument( + "--num-frame-per-block", + type=int, + default=FastVideoArgs.num_frame_per_block, + help="Number of frames per block for causal generation") + # Add pipeline configuration arguments PipelineConfig.add_cli_args(parser) @@ -751,7 +761,6 @@ class TrainingArgs(FastVideoArgs): warp_denoising_step: bool = False # Self-forcing specific arguments - num_frame_per_block: int = 3 independent_first_frame: bool = False enable_gradient_masking: bool = True gradient_mask_last_n_frames: int = 21 @@ -1146,11 +1155,6 @@ class TrainingArgs(FastVideoArgs): ) # Self-forcing specific arguments - parser.add_argument( - "--num-frame-per-block", - type=int, - default=TrainingArgs.num_frame_per_block, - help="Number of frames per block for causal generation") parser.add_argument( "--independent-first-frame", action=StoreBoolean, diff --git a/fastvideo/pipelines/stages/causal_denoising.py b/fastvideo/pipelines/stages/causal_denoising.py index 777314547..81b7f6ee8 100644 --- a/fastvideo/pipelines/stages/causal_denoising.py +++ b/fastvideo/pipelines/stages/causal_denoising.py @@ -1,6 +1,8 @@ import torch # type: ignore -from fastvideo.distributed import get_local_torch_device +from fastvideo.distributed import (get_local_torch_device, sequence_model_parallel_all_gather, + get_sp_parallel_rank, get_sp_world_size) +from einops import rearrange from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.forward_context import set_forward_context from fastvideo.logger import init_logger @@ -47,6 +49,7 @@ class CausalDMDDenosingStage(DenoisingStage): # Model-dependent constants (aligned with causal_inference.py assumptions) self.num_transformer_blocks = len(self.transformer.blocks) self.num_frames_per_block = self.transformer.config.arch_config.num_frames_per_block + self.num_frames_per_block = 4 self.sliding_window_num_frames = self.transformer.config.arch_config.sliding_window_num_frames try: @@ -221,13 +224,34 @@ class CausalDMDDenosingStage(DenoisingStage): # 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: + for idx, current_num_frames in enumerate(block_sizes): + logger.info(f"================================================") + logger.info(f"Block {idx}") + logger.info(f"================================================") + logger.info(f"Current num frames: {current_num_frames}") + logger.info(f"Processing block {start_index} to {start_index + current_num_frames}") + logger.info(f"Latents shape: {latents.shape}") current_latents = latents[:, :, start_index:start_index + current_num_frames, :, :] + + sp_world_size, rank_in_sp_group = get_sp_world_size( + ), get_sp_parallel_rank() + sp_group = sp_world_size > 1 + # logger.info(f"SP group: {sp_group}") + logger.info(f"before SP: Current latents shape: {current_latents.shape}") + if sp_group: + current_latents = rearrange(current_latents, + "b c (n t) h w -> b c n t h w", + n=sp_world_size).contiguous() + current_latents = current_latents[:, :, rank_in_sp_group, :, :, :] + logger.info(f"after SP: Current latents shape: {current_latents.shape}") + # Per-rank temporal offset (in frames) within this block when using SP + frames_per_rank = (current_num_frames // sp_world_size) if sp_group else current_num_frames + rank_offset_frames = (rank_in_sp_group * frames_per_rank) if sp_group else 0 # 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 - + logger.info(f"before DMD: Noise latents shape: {noise_latents_btchw.shape}") for i, t_cur in enumerate(timesteps): if boundary_timestep is not None and t_cur < boundary_timestep: current_model = self.transformer_2 @@ -289,9 +313,9 @@ class CausalDMDDenosingStage(DenoisingStage): t_expanded_noise, kv_cache=_get_kv_cache(t_cur), crossattn_cache=crossattn_cache, - current_start=(pos_start_base + start_index) * + current_start=(pos_start_base + start_index + rank_offset_frames) * self.frame_seq_length, - start_frame=start_index, + start_frame=pos_start_base + start_index + rank_offset_frames, **image_kwargs, **pos_cond_kwargs, ).permute(0, 2, 1, 3, 4) @@ -353,16 +377,23 @@ class CausalDMDDenosingStage(DenoisingStage): progress_bar.update() # Write back and advance + logger.info(f"before all gather: Current latents shape: {current_latents.shape}") + context_bcthw = current_latents.clone().to(target_dtype) + if sp_group: + current_latents = sequence_model_parallel_all_gather(current_latents, dim=2) + logger.info(f"after all gather: Current latents shape: {current_latents.shape}") + logger.info(f"start_index: {start_index}, current_num_frames: {current_num_frames}") + logger.info(f"before write back: Latents shape: {latents.shape}") latents[:, :, start_index:start_index + current_num_frames, :, :] = current_latents - + logger.info(f"after write back: Latents shape: {latents.shape}") # 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) + # context_bcthw = current_latents.to(target_dtype) with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled), \ @@ -391,9 +422,9 @@ class CausalDMDDenosingStage(DenoisingStage): t_expanded_context, kv_cache=kv_cache1, crossattn_cache=crossattn_cache, - current_start=(pos_start_base + start_index) * + current_start=(pos_start_base + start_index + rank_offset_frames) * self.frame_seq_length, - start_frame=start_index, + start_frame=pos_start_base + start_index + rank_offset_frames, **image_kwargs, **pos_cond_kwargs, )