| 12 | |
| 13 | |
| 14 | class CausalDiffusion(BaseModel): |
| 15 | def __init__(self, args, device): |
| 16 | """ |
| 17 | Initialize the Diffusion loss module. |
| 18 | """ |
| 19 | super().__init__(args, device) |
| 20 | self.num_frame_per_block = getattr(args, "num_frame_per_block", 1) |
| 21 | if self.num_frame_per_block > 1: |
| 22 | self.generator.model.num_frame_per_block = self.num_frame_per_block |
| 23 | self.independent_first_frame = getattr(args, "independent_first_frame", False) |
| 24 | if self.independent_first_frame and not getattr(args, "i2v", False): |
| 25 | self.generator.model.independent_first_frame = True |
| 26 | |
| 27 | if args.gradient_checkpointing: |
| 28 | self.generator.enable_gradient_checkpointing() |
| 29 | |
| 30 | # Step 2: Initialize all hyperparameters |
| 31 | self.num_train_timestep = args.num_train_timestep |
| 32 | self.min_step = int(0.02 * self.num_train_timestep) |
| 33 | self.max_step = int(0.98 * self.num_train_timestep) |
| 34 | self.guidance_scale = args.guidance_scale |
| 35 | self.timestep_shift = getattr(args, "timestep_shift", 1.0) |
| 36 | self.teacher_forcing = getattr(args, "teacher_forcing", False) |
| 37 | # Noise augmentation in teacher forcing, we add small noise to clean context latents |
| 38 | self.noise_augmentation_max_timestep = getattr(args, "noise_augmentation_max_timestep", 0) |
| 39 | |
| 40 | self.args = args |
| 41 | self.device = device |
| 42 | self.inference_pipeline = None |
| 43 | |
| 44 | # Error recycling (SVI-style error buffer) |
| 45 | # When ``enable_position_bucketing`` is true, each rank holds a 2D |
| 46 | # buffer ``(local_block_position × timestep)``. The pos dimension only |
| 47 | # covers the LOCAL slice of the sequence this rank is responsible for |
| 48 | # (no cross-SP-rank pos sharing — those positions are simply not |
| 49 | # reachable by this rank during forward), so memory cost scales as |
| 50 | # ``num_blocks_global / sp_size`` instead of ``num_blocks_global``. |
| 51 | # ``global_block_offset`` is recorded for logging only. |
| 52 | # During the first ``buffer_warmup_iter`` global steps, errors are |
| 53 | # all-gathered across the DP group (ranks with the same SP rank but |
| 54 | # different DP replicas), so each rank's local pos buckets fill up |
| 55 | # ``dp_size`` × faster without any wasted bandwidth. |
| 56 | self.error_buffer = None |
| 57 | self.noise_error_buffer = None |
| 58 | self.er_num_blocks = 0 # local; >0 means 2D position-bucketed |
| 59 | self.er_block_offset = 0 # global block offset for THIS rank |
| 60 | er_cfg = getattr(args, "error_recycling", None) |
| 61 | if er_cfg is not None and getattr(er_cfg, "enabled", False): |
| 62 | from utils.error_buffer import build_error_buffer |
| 63 | cfg_dict = er_cfg if isinstance(er_cfg, dict) else dict(er_cfg) |
| 64 | cfg_dict.setdefault("num_train_timesteps", self.num_train_timestep) |
| 65 | sp_size = int(getattr(args, "sequence_parallel_size", 1) or 1) |
| 66 | if cfg_dict.get("enable_position_bucketing", False): |
| 67 | shape = list(getattr(args, "image_or_video_shape", [1, 0])) |
| 68 | total_frames = int(shape[1]) if len(shape) > 1 else 0 |
| 69 | assert total_frames > 0 and self.num_frame_per_block > 0, ( |
| 70 | "enable_position_bucketing=true requires " |
| 71 | "image_or_video_shape[1] and num_frame_per_block to be set." |