diff --git a/training/args.py b/training/args.py index dee56f3..9e38d31 100644 --- a/training/args.py +++ b/training/args.py @@ -78,6 +78,7 @@ def _get_dataset_args(parser: argparse.ArgumentParser) -> None: nargs="+", type=int, default=[49], + help="CogVideoX1.5 need to guarantee that ((num_frames - 1) // self.vae_scale_factor_temporal + 1) % patch_size_t != 0, such as 53" ) parser.add_argument( "--load_tensors", diff --git a/training/cogvideox_image_to_video_lora.py b/training/cogvideox_image_to_video_lora.py index 4523f73..eea8ffc 100644 --- a/training/cogvideox_image_to_video_lora.py +++ b/training/cogvideox_image_to_video_lora.py @@ -787,6 +787,7 @@ def main(args): num_frames=num_frames, vae_scale_factor_spatial=VAE_SCALE_FACTOR_SPATIAL, patch_size=model_config.patch_size, + patch_size_t=model_config.patch_size_t if hasattr(model_config, "patch_size_t") else None, attention_head_dim=model_config.attention_head_dim, device=accelerator.device, ) diff --git a/training/cogvideox_text_to_video_lora.py b/training/cogvideox_text_to_video_lora.py index cf678a5..a9a2448 100644 --- a/training/cogvideox_text_to_video_lora.py +++ b/training/cogvideox_text_to_video_lora.py @@ -696,6 +696,7 @@ def main(args): num_frames=num_frames, vae_scale_factor_spatial=VAE_SCALE_FACTOR_SPATIAL, patch_size=model_config.patch_size, + patch_size_t=model_config.patch_size_t if hasattr(model_config, "patch_size_t") else None, attention_head_dim=model_config.attention_head_dim, device=accelerator.device, ) diff --git a/training/cogvideox_text_to_video_sft.py b/training/cogvideox_text_to_video_sft.py index 9b44217..bf9acad 100644 --- a/training/cogvideox_text_to_video_sft.py +++ b/training/cogvideox_text_to_video_sft.py @@ -662,6 +662,7 @@ def main(args): num_frames=num_frames, vae_scale_factor_spatial=VAE_SCALE_FACTOR_SPATIAL, patch_size=model_config.patch_size, + patch_size_t=model_config.patch_size_t if hasattr(model_config, "patch_size_t") else None, attention_head_dim=model_config.attention_head_dim, device=accelerator.device, ) diff --git a/training/utils.py b/training/utils.py index b2922c8..9513805 100644 --- a/training/utils.py +++ b/training/utils.py @@ -198,6 +198,7 @@ def prepare_rotary_positional_embeddings( num_frames: int, vae_scale_factor_spatial: int = 8, patch_size: int = 2, + patch_size_t: int = None, attention_head_dim: int = 64, device: Optional[torch.device] = None, base_height: int = 480, @@ -207,14 +208,28 @@ def prepare_rotary_positional_embeddings( grid_width = width // (vae_scale_factor_spatial * patch_size) base_size_width = base_width // (vae_scale_factor_spatial * patch_size) base_size_height = base_height // (vae_scale_factor_spatial * patch_size) + if patch_size_t is None: + # CogVideoX 1.0 + grid_crops_coords = get_resize_crop_region_for_grid((grid_height, grid_width), base_size_width, base_size_height) + freqs_cos, freqs_sin = get_3d_rotary_pos_embed( + embed_dim=attention_head_dim, + crops_coords=grid_crops_coords, + grid_size=(grid_height, grid_width), + temporal_size=num_frames, + ) + else: + # CogVideoX 1.5 + base_num_frames = (num_frames + patch_size_t - 1) // patch_size_t + + freqs_cos, freqs_sin = get_3d_rotary_pos_embed( + embed_dim=attention_head_dim, + crops_coords=None, + grid_size=(grid_height, grid_width), + temporal_size=base_num_frames, + grid_type="slice", + max_size=(base_size_height, base_size_width), + ) - grid_crops_coords = get_resize_crop_region_for_grid((grid_height, grid_width), base_size_width, base_size_height) - freqs_cos, freqs_sin = get_3d_rotary_pos_embed( - embed_dim=attention_head_dim, - crops_coords=grid_crops_coords, - grid_size=(grid_height, grid_width), - temporal_size=num_frames, - ) freqs_cos = freqs_cos.to(device=device) freqs_sin = freqs_sin.to(device=device)