diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index c263353f9..19fae38d8 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -906,6 +906,7 @@ class TrainingArgs(FastVideoArgs): # Self-forcing specific arguments num_frame_per_block: int = 3 + chunkwise_timestep: bool = True independent_first_frame: bool = False enable_gradient_masking: bool = True gradient_mask_last_n_frames: int = 21 @@ -1318,6 +1319,10 @@ class TrainingArgs(FastVideoArgs): type=int, default=TrainingArgs.num_frame_per_block, help="Number of frames per block for causal generation") + parser.add_argument( + "--chunkwise-timestep", + action=StoreBoolean, + help="Whether to share one timestep within block") parser.add_argument( "--independent-first-frame", action=StoreBoolean, diff --git a/fastvideo/training/matrixgame_ar_diffusion_pipeline.py b/fastvideo/training/matrixgame_ar_diffusion_pipeline.py index dc0953dda..3bc321b78 100644 --- a/fastvideo/training/matrixgame_ar_diffusion_pipeline.py +++ b/fastvideo/training/matrixgame_ar_diffusion_pipeline.py @@ -92,6 +92,9 @@ class MatrixGameARDiffusionPipeline(TrainingPipeline): self.num_frame_per_block = self._resolve_num_frame_per_block( training_args ) + self.chunkwise_timestep = bool( + getattr(training_args, "chunkwise_timestep", False) + ) self.timestep_shift = training_args.pipeline_config.flow_shift self.ar_noise_scheduler = DiffusionForcingScheduler( shift=self.timestep_shift, sigma_min=0.0, extra_one_step=True @@ -102,9 +105,11 @@ class MatrixGameARDiffusionPipeline(TrainingPipeline): logger.info( "MatrixGame AR diffusion pipeline initialized with " - "num_frame_per_block=%d, diffusion_forcing_shift=%.1f", + "num_frame_per_block=%d, diffusion_forcing_shift=%.1f, " + "chunkwise_timestep=%s", self.num_frame_per_block, self.timestep_shift, + self.chunkwise_timestep, ) def initialize_validation_pipeline(self, training_args: TrainingArgs): @@ -195,6 +200,28 @@ class MatrixGameARDiffusionPipeline(TrainingPipeline): device=device, dtype=torch.long, ) + if not getattr(self, "chunkwise_timestep", False): + return timestep + + chunk_size = int(num_frame_per_block) + if chunk_size <= 0: + raise ValueError( + f"num_frame_per_block must be > 0, got {chunk_size}" + ) + + num_chunks = (num_frame + chunk_size - 1) // chunk_size + padded_num_frames = num_chunks * chunk_size + if padded_num_frames != num_frame: + pad = timestep[:, -1:].expand( + batch_size, + padded_num_frames - num_frame, + ) + timestep = torch.cat([timestep, pad], dim=1) + timestep = timestep.reshape(batch_size, num_chunks, chunk_size) + timestep[:, :, 1:] = timestep[:, :, :1] + timestep = timestep.reshape(batch_size, padded_num_frames)[ + :, :num_frame + ] return timestep def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch: @@ -292,6 +319,13 @@ class MatrixGameARDiffusionPipeline(TrainingPipeline): device=latents_btchw.device, dtype=latents_btchw.dtype, ) + if ( + self.training_args.sp_size > 1 + and self.sp_group is not None + and hasattr(self.sp_group, "broadcast") + ): + self.sp_group.broadcast(timesteps, src=0) + self.sp_group.broadcast(noise, src=0) noisy_latents = self.ar_noise_scheduler.add_noise( latents_btchw.flatten(0, 1),