From ff4e183adac5cc6dd8cfd639745e393d39e1123c Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Sat, 5 Jul 2025 13:32:01 -0700 Subject: [PATCH] run lint 3.10 --- .../v1/training/wan_i2v_training_pipeline.py | 20 +++++++++---------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/fastvideo/v1/training/wan_i2v_training_pipeline.py b/fastvideo/v1/training/wan_i2v_training_pipeline.py index 69799f6d6..f43c1ce04 100644 --- a/fastvideo/v1/training/wan_i2v_training_pipeline.py +++ b/fastvideo/v1/training/wan_i2v_training_pipeline.py @@ -148,21 +148,19 @@ class WanI2VTrainingPipeline(TrainingPipeline): first_frame_mask = mask_lat_size[:, :, 0:1] first_frame_mask = torch.repeat_interleave( - first_frame_mask, - dim=2, - repeats=temporal_compression_ratio) + first_frame_mask, dim=2, repeats=temporal_compression_ratio) mask_lat_size = torch.concat( [first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2) - mask_lat_size = mask_lat_size.view( - batch_size, -1, - temporal_compression_ratio, latent_height, - latent_width) + mask_lat_size = mask_lat_size.view(batch_size, -1, + temporal_compression_ratio, + latent_height, latent_width) mask_lat_size = mask_lat_size.transpose(1, 2) - mask_lat_size = mask_lat_size.to(image_latents.device).to( - dtype=torch.bfloat16) + mask_lat_size = mask_lat_size.to( + image_latents.device).to(dtype=torch.bfloat16) + + mask_lat_size = shard_latents_across_sp(mask_lat_size, + self.training_args.num_latent_t) - mask_lat_size = shard_latents_across_sp(mask_lat_size, self.training_args.num_latent_t) - training_batch.noisy_model_input = torch.cat( [training_batch.noisy_model_input, mask_lat_size, image_latents], dim=1)