diff --git a/examples/distill/Wan2.1-I2V/crush_smol/distill_i2v_14B.sh b/examples/distill/Wan2.1-I2V/crush_smol/distill_i2v_14B.sh index 6ba2449dd..6b2aaf880 100644 --- a/examples/distill/Wan2.1-I2V/crush_smol/distill_i2v_14B.sh +++ b/examples/distill/Wan2.1-I2V/crush_smol/distill_i2v_14B.sh @@ -6,7 +6,7 @@ export TOKENIZERS_PARALLELISM=false # export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers" -DATA_DIR="data/crush-smol_processed_wan21_i2v_14b/combined_parquet_dataset/" +DATA_DIR="data/crush-smol_processed_wan21_i2v_14b_single/combined_parquet_dataset/" VALIDATION_DATASET_FILE="examples/distill/Wan2.1-I2V/crush_smol/validation.json" NUM_GPUS=8 export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde' @@ -17,7 +17,7 @@ export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde' training_args=( --tracker_project_name "wan_i2v_distill" --output_dir "checkpoints/wan_i2v_distill" - --wandb_run_name "wan21_i2v_14b_ode_3steps" + --wandb_run_name "no_back_wan21_i2v_14b_ode_3steps_8gpu" --max_train_steps 2000 --train_batch_size 1 --train_sp_batch_size 1 @@ -56,7 +56,7 @@ validation_args=( # --log_visualization --log_visualization --validation_dataset_file "$VALIDATION_DATASET_FILE" - --validation_steps 50 + --validation_steps 1 --validation_sampling_steps "3" --validation_guidance_scale "1.0" # not used for dmd inference ) diff --git a/examples/distill/Wan2.1-I2V/crush_smol/preprocess_wan_data_i2v.sh b/examples/distill/Wan2.1-I2V/crush_smol/preprocess_wan_data_i2v.sh index 3a2c96c47..e72b0b60d 100644 --- a/examples/distill/Wan2.1-I2V/crush_smol/preprocess_wan_data_i2v.sh +++ b/examples/distill/Wan2.1-I2V/crush_smol/preprocess_wan_data_i2v.sh @@ -4,7 +4,7 @@ GPU_NUM=1 # 2,4,8 MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers" MODEL_TYPE="wan" DATA_MERGE_PATH="data/crush-smol/merge.txt" -OUTPUT_DIR="data/crush-smol_processed_wan21_i2v_14b/" +OUTPUT_DIR="data/crush-smol_processed_wan21_i2v_14b_single/" torchrun --nproc_per_node=$GPU_NUM \ fastvideo/pipelines/preprocess/v1_preprocess.py \ @@ -18,7 +18,7 @@ torchrun --nproc_per_node=$GPU_NUM \ --dataloader_num_workers 0 \ --output_dir=$OUTPUT_DIR \ --train_fps 16 \ - --samples_per_file 8 \ + --samples_per_file 8\ --flush_frequency 8 \ --video_length_tolerance_range 5 \ --preprocess_task "i2v" \ No newline at end of file diff --git a/examples/distill/Wan2.1-I2V/crush_smol/validation.json b/examples/distill/Wan2.1-I2V/crush_smol/validation.json index a7e327a04..30c3844a5 100644 --- a/examples/distill/Wan2.1-I2V/crush_smol/validation.json +++ b/examples/distill/Wan2.1-I2V/crush_smol/validation.json @@ -1,32 +1,5 @@ { "data": [ - { - "caption": "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.", - "image_path": null, - "video_path": "validation_dataset/yYcK4nANZz4-Scene-034.mp4", - "num_inference_steps": 40, - "height": 480, - "width": 832, - "num_frames": 77 - }, - { - "caption": "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.", - "image_path": null, - "video_path": "validation_dataset/yYcK4nANZz4-Scene-027.mp4", - "num_inference_steps": 40, - "height": 480, - "width": 832, - "num_frames": 77 - }, - { - "caption": "A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table.", - "image_path": null, - "video_path": "validation_dataset/yYcK4nANZz4-Scene-030.mp4", - "num_inference_steps": 40, - "height": 480, - "width": 832, - "num_frames": 77 - }, { "caption": "A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open.", "image_path": null, @@ -35,42 +8,6 @@ "height": 480, "width": 832, "num_frames": 77 - }, - { - "caption": "The video shows a green and orange object being flattened as if it were under a hydraulic press, with the press moving down and compressing the object.", - "image_path": null, - "video_path": "validation_dataset/1gGQy4nxyUo-Scene-056.mp4", - "num_inference_steps": 40, - "height": 480, - "width": 832, - "num_frames": 77 - }, - { - "caption": "The video shows a cylindrical object with a cityscape image being flattened as if it were under a hydraulic press. The object is placed on a metal platform, and a large, striped cylinder presses down on it, causing it to collapse and release a liquid inside. The background features a green wall with a yellow and red warning sign.", - "image_path": null, - "video_path": "validation_dataset/1gGQy4nxyUo-Scene-059.mp4", - "num_inference_steps": 40, - "height": 480, - "width": 832, - "num_frames": 77 - }, - { - "caption": "The video shows a close-up of an orange being flattened as if it were under a hydraulic press, with the press moving down and compressing the fruit until it is completely flattened.", - "image_path": null, - "video_path": "validation_dataset/EJqsC21GSBY-Scene-059.mp4", - "num_inference_steps": 40, - "height": 480, - "width": 832, - "num_frames": 77 - }, - { - "caption": "A colorful puzzle ball is being crushed by a large metal cylinder, which flattens the objects as if they were under a hydraulic press.", - "image_path": null, - "video_path": "validation_dataset/GBSfpTcKegk-Scene-003.mp4", - "num_inference_steps": 40, - "height": 480, - "width": 832, - "num_frames": 77 } ] } \ No newline at end of file diff --git a/fastvideo/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py index 8e83185c7..e05fe83cf 100644 --- a/fastvideo/training/distillation_pipeline.py +++ b/fastvideo/training/distillation_pipeline.py @@ -237,10 +237,11 @@ class DistillationPipeline(TrainingPipeline): dtype = latents.dtype # Step 1: Randomly sample a target timestep index from denoising_step_list - target_timestep_idx = torch.randint(0, - len(self.denoising_step_list), [1], - device=self.device, - dtype=torch.long) + # target_timestep_idx = torch.randint(0, + # len(self.denoising_step_list), [1], + # device=self.device, + # dtype=torch.long) + target_timestep_idx = torch.tensor([2], device=self.device, dtype=torch.long) target_timestep_idx_int = target_timestep_idx.item() target_timestep = self.denoising_step_list[target_timestep_idx] @@ -272,8 +273,19 @@ class DistillationPipeline(TrainingPipeline): training_batch_temp = self._build_distill_input_kwargs( current_noise_latents, current_timestep_tensor, training_batch.conditional_dict, training_batch) + # pred_flow = self.transformer( + # **training_batch_temp.input_kwargs).permute( + # 0, 2, 1, 3, 4) + # t_expand = current_timestep_tensor.repeat(current_noise_latents.shape[0]) pred_flow = self.transformer( - **training_batch_temp.input_kwargs).permute( + training_batch_temp.input_kwargs["hidden_states"], + training_batch_temp.input_kwargs["encoder_hidden_states"], + current_timestep_tensor, + encoder_hidden_states_image=training_batch_temp.input_kwargs["encoder_hidden_states_image"], + guidance=None, + # **training_batch.conditional_dict, + # **training_batch.unconditional_dict, + ).permute( 0, 2, 1, 3, 4) pred_clean = pred_noise_to_pred_video( pred_noise=pred_flow.flatten(0, 1), diff --git a/fastvideo/training/wan_i2v_distillation_pipeline.py b/fastvideo/training/wan_i2v_distillation_pipeline.py index 47b5d73de..122df390d 100644 --- a/fastvideo/training/wan_i2v_distillation_pipeline.py +++ b/fastvideo/training/wan_i2v_distillation_pipeline.py @@ -181,10 +181,24 @@ class WanI2VDistillationPipeline(DistillationPipeline): text_dict: dict[str, torch.Tensor], training_batch: TrainingBatch) -> TrainingBatch: # Image Embeds for conditioning - image_embeds = training_batch.image_embeds - assert torch.isnan(image_embeds).sum() == 0 + # image_embeds = training_batch.image_embeds + # assert torch.isnan(image_embeds).sum() == 0 + # image_embeds = image_embeds.to(get_local_torch_device(), + # dtype=torch.bfloat16) + from fastvideo.models.vision_utils import load_video + image = load_video("examples/distill/Wan2.1-I2V/crush_smol/validation_dataset/1gGQy4nxyUo-Scene-016.mp4")[0] + image_processor = self.validation_pipeline.get_module("image_processor") + image_encoder = self.validation_pipeline.get_module("image_encoder") + image_inputs = image_processor( + images=image, return_tensors="pt").to(get_local_torch_device()) + from fastvideo.forward_context import set_forward_context + with set_forward_context(current_timestep=0, attn_metadata=None): + outputs = image_encoder(**image_inputs) + image_embeds = outputs.last_hidden_state image_embeds = image_embeds.to(get_local_torch_device(), dtype=torch.bfloat16) + assert torch.isnan(image_embeds).sum() == 0 + # self.validation_pipeline. noisy_model_input = torch.cat( [noise_input, @@ -197,7 +211,7 @@ class WanI2VDistillationPipeline(DistillationPipeline): "encoder_hidden_states": text_dict["encoder_hidden_states"], "encoder_attention_mask": text_dict["encoder_attention_mask"], "timestep": timestep, - "encoder_hidden_states_image": image_embeds, + "encoder_hidden_states_image": image_embeds.permute(0, 2, 1, 3, 4), "return_dict": False, } training_batch.noise_latents = noise_input