diff --git a/fastvideo/v1/dataset/parquet_dataset_map_style.py b/fastvideo/v1/dataset/parquet_dataset_map_style.py index 067f04340..eca514ef4 100644 --- a/fastvideo/v1/dataset/parquet_dataset_map_style.py +++ b/fastvideo/v1/dataset/parquet_dataset_map_style.py @@ -200,10 +200,6 @@ class LatentsParquetMapStyleDataset(Dataset): super().__init__() self.path = path self.cfg_rate = cfg_rate - if cfg_rate > 0.0: - raise ValueError( - "cfg_rate > 0.0 is not supported for now because it will trigger bug when num_data_workers > 0" - ) logger.info("Initializing LatentsParquetMapStyleDataset with path: %s", path) self.parquet_files, self.lengths = get_parquet_files_and_length(path) @@ -247,7 +243,7 @@ class LatentsParquetMapStyleDataset(Dataset): [self.lengths[0]]) all_latents_list, all_embs_list, all_masks_list, caption_text_list = collate_latents_embs_masks( - [row_dict], self.text_padding_length, self.keys) + [row_dict], self.text_padding_length, self.keys, cfg_rate=0.0) all_latents, all_embs, all_masks, caption_text = all_latents_list[ 0], all_embs_list[0], all_masks_list[0], caption_text_list[0] # add batch dimension @@ -268,7 +264,7 @@ class LatentsParquetMapStyleDataset(Dataset): ] all_latents, all_embs, all_masks, caption_text = collate_latents_embs_masks( - rows, self.text_padding_length, self.keys) + rows, self.text_padding_length, self.keys, self.cfg_rate) return all_latents, all_embs, all_masks, caption_text def __len__(self): diff --git a/fastvideo/v1/dataset/preprocessing_datasets.py b/fastvideo/v1/dataset/preprocessing_datasets.py index 612753d38..c2fad482c 100644 --- a/fastvideo/v1/dataset/preprocessing_datasets.py +++ b/fastvideo/v1/dataset/preprocessing_datasets.py @@ -461,7 +461,7 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, self.text_encoding_stage = TextEncodingStage( tokenizer=tokenizer, text_max_length=args.text_max_length, - cfg_rate=args.cfg) + cfg_rate=args.training_cfg_rate) def _load_raw_data(self) -> List[Dict]: """Load raw data from JSON files.""" diff --git a/fastvideo/v1/dataset/utils.py b/fastvideo/v1/dataset/utils.py index 9d36e8233..49231c774 100644 --- a/fastvideo/v1/dataset/utils.py +++ b/fastvideo/v1/dataset/utils.py @@ -1,3 +1,4 @@ +import random from typing import Any, Dict, List import numpy as np @@ -20,7 +21,7 @@ def pad(t: torch.Tensor, padding_length: int) -> torch.Tensor: return t[:padding_length], torch.ones(padding_length) -def get_torch_tensors_from_row_dict(row_dict, keys) -> Dict[str, Any]: +def get_torch_tensors_from_row_dict(row_dict, keys, cfg_rate) -> Dict[str, Any]: """ Get the latents and prompts from a row dictionary. """ @@ -42,7 +43,10 @@ def get_torch_tensors_from_row_dict(row_dict, keys) -> Dict[str, Any]: bytes = row_dict[f"{key}_bytes"] # TODO (peiyuan): read precision - data = np.frombuffer(bytes, dtype=np.float32).reshape(shape).copy() + if key == 'text_embedding' and random.random() < cfg_rate: + data = np.zeros((512, 4096), dtype=np.float32) + else: + data = np.frombuffer(bytes, dtype=np.float32).reshape(shape).copy() data = torch.from_numpy(data) if len(data.shape) == 3: B, L, D = data.shape @@ -53,8 +57,11 @@ def get_torch_tensors_from_row_dict(row_dict, keys) -> Dict[str, Any]: def collate_latents_embs_masks( - batch_to_process, text_padding_length, - keys) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, List[str]]: + batch_to_process, + text_padding_length, + keys, + cfg_rate=0.0 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, List[str]]: # Initialize tensors to hold padded embeddings and masks all_latents = [] all_embs = [] @@ -63,7 +70,7 @@ def collate_latents_embs_masks( # Process each row individually for i, row in enumerate(batch_to_process): # Get tensors from row - data = get_torch_tensors_from_row_dict(row, keys) + data = get_torch_tensors_from_row_dict(row, keys, cfg_rate) latents, emb = data["vae_latent"], data["text_embedding"] padded_emb, mask = pad(emb, text_padding_length) diff --git a/fastvideo/v1/fastvideo_args.py b/fastvideo/v1/fastvideo_args.py index 355e60c5b..3cdeecbe2 100644 --- a/fastvideo/v1/fastvideo_args.py +++ b/fastvideo/v1/fastvideo_args.py @@ -384,7 +384,7 @@ class TrainingArgs(FastVideoArgs): # diffusion setting ema_decay: float = 0.0 ema_start_step: int = 0 - cfg: float = 0.0 + training_cfg_rate: float = 0.0 precondition_outputs: bool = False # validation & logs @@ -528,7 +528,7 @@ class TrainingArgs(FastVideoArgs): type=int, default=0, help="Step to start EMA") - parser.add_argument("--cfg", + parser.add_argument("--training-cfg-rate", type=float, help="Classifier-free guidance scale") parser.add_argument( diff --git a/fastvideo/v1/pipelines/preprocess/v1_preprocess.py b/fastvideo/v1/pipelines/preprocess/v1_preprocess.py index fece08186..6c394b5bf 100644 --- a/fastvideo/v1/pipelines/preprocess/v1_preprocess.py +++ b/fastvideo/v1/pipelines/preprocess/v1_preprocess.py @@ -91,7 +91,7 @@ if __name__ == "__main__": type=str, default="google/t5-v1_1-xxl") parser.add_argument("--cache_dir", type=str, default="./cache_dir") - parser.add_argument("--cfg", type=float, default=0.0) + parser.add_argument("--training_cfg_rate", type=float, default=0.0) parser.add_argument( "--output_dir", type=str, diff --git a/fastvideo/v1/tests/nightly/test_e2e_overfit_single_sample.py b/fastvideo/v1/tests/nightly/test_e2e_overfit_single_sample.py index 4e86c9920..a8b6cce1a 100644 --- a/fastvideo/v1/tests/nightly/test_e2e_overfit_single_sample.py +++ b/fastvideo/v1/tests/nightly/test_e2e_overfit_single_sample.py @@ -122,7 +122,7 @@ def run_training(): "--checkpoints_total_limit", "3", "--allow_tf32", "--ema_start_step", "0", - "--cfg", "0.0", + "--training_cfg_rate", "0.0", "--output_dir", LOCAL_OUTPUT_DIR, "--tracker_project_name", "wan_finetune_overfit_ci", "--num_height", "480", diff --git a/fastvideo/v1/tests/training/VSA/test_training_loss_VSA.py b/fastvideo/v1/tests/training/VSA/test_training_loss_VSA.py index 36b082f61..3d5324c41 100644 --- a/fastvideo/v1/tests/training/VSA/test_training_loss_VSA.py +++ b/fastvideo/v1/tests/training/VSA/test_training_loss_VSA.py @@ -54,7 +54,7 @@ def run_worker(): "--checkpoints_total_limit", "3", "--allow_tf32", "--ema_start_step", "0", - "--cfg", "0.0", + "--training_cfg_rate", "0.0", "--output_dir", "data/wan_finetune_test_VSA", "--tracker_project_name", "wan_finetune_ci_VSA", "--wandb_run_name", wandb_name, diff --git a/fastvideo/v1/tests/training/Vanilla/test_training_loss.py b/fastvideo/v1/tests/training/Vanilla/test_training_loss.py index c8cd721cd..12a0f7762 100644 --- a/fastvideo/v1/tests/training/Vanilla/test_training_loss.py +++ b/fastvideo/v1/tests/training/Vanilla/test_training_loss.py @@ -59,7 +59,7 @@ def run_worker(): "--checkpoints_total_limit", "3", "--allow_tf32", "--ema_start_step", "0", - "--cfg", "0.0", + "--training_cfg_rate", "0.0", "--output_dir", "data/wan_finetune_test", "--tracker_project_name", "wan_finetune_ci", "--wandb_run_name", wandb_name, diff --git a/fastvideo/v1/training/training_pipeline.py b/fastvideo/v1/training/training_pipeline.py index f762d00a3..a52d6e504 100644 --- a/fastvideo/v1/training/training_pipeline.py +++ b/fastvideo/v1/training/training_pipeline.py @@ -105,6 +105,7 @@ class TrainingPipeline(ComposedPipelineBase, ABC): training_args.data_path, training_args.train_batch_size, num_data_workers=training_args.dataloader_num_workers, + cfg_rate=training_args.training_cfg_rate, drop_last=True, text_padding_length=training_args.pipeline_config. text_encoder_configs[0].arch_config. @@ -534,9 +535,9 @@ class TrainingPipeline(ComposedPipelineBase, ABC): training_args.validation_preprocessed_path, batch_size=1, num_data_workers=0, + cfg_rate=0.0, drop_last=False, - drop_first_row=sampling_param.negative_prompt is not None, - cfg_rate=training_args.cfg) + drop_first_row=sampling_param.negative_prompt is not None) if sampling_param.negative_prompt: _, negative_prompt_embeds, negative_prompt_attention_mask, _ = validation_dataset.get_validation_negative_prompt( ) diff --git a/scripts/finetune/finetune_v1.sh b/scripts/finetune/finetune_v1.sh index 06d5dea56..df9c55586 100644 --- a/scripts/finetune/finetune_v1.sh +++ b/scripts/finetune/finetune_v1.sh @@ -36,7 +36,7 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\ --checkpoints_total_limit 3\ --allow_tf32\ --ema_start_step 0\ - --cfg 0.0\ + --training_cfg_rate 0.0\ --output_dir="$DATA_DIR/outputs/wan_finetune"\ --tracker_project_name wan_finetune \ --num_height 480 \ diff --git a/scripts/finetune/finetune_v1_VSA.sh b/scripts/finetune/finetune_v1_VSA.sh index cc551d14a..e4ded3e05 100644 --- a/scripts/finetune/finetune_v1_VSA.sh +++ b/scripts/finetune/finetune_v1_VSA.sh @@ -42,7 +42,7 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS \ --checkpoints_total_limit 3 \ --allow_tf32 \ --ema_start_step 0 \ - --cfg 0.0 \ + --training_cfg_rate 0.0 \ --output_dir "$DATA_DIR/outputs/wan_finetune" \ --tracker_project_name VSA_finetune \ --num_height 448 \