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 6b2aaf880..a65ed8623 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_single/combined_parquet_dataset/" +DATA_DIR="data/crush-smol_processed_wan21_i2v_14b/combined_parquet_dataset/" VALIDATION_DATASET_FILE="examples/distill/Wan2.1-I2V/crush_smol/validation.json" NUM_GPUS=8 export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde' @@ -56,7 +56,7 @@ validation_args=( # --log_visualization --log_visualization --validation_dataset_file "$VALIDATION_DATASET_FILE" - --validation_steps 1 + --validation_steps 50 --validation_sampling_steps "3" --validation_guidance_scale "1.0" # not used for dmd inference ) diff --git a/examples/distill/Wan2.1-I2V/crush_smol/validation.json b/examples/distill/Wan2.1-I2V/crush_smol/validation.json index 30c3844a5..a7e327a04 100644 --- a/examples/distill/Wan2.1-I2V/crush_smol/validation.json +++ b/examples/distill/Wan2.1-I2V/crush_smol/validation.json @@ -1,5 +1,32 @@ { "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, @@ -8,6 +35,42 @@ "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/examples/distill/Wan2.1-I2V/crush_smol/validation_orig.json b/examples/distill/Wan2.1-I2V/crush_smol/validation_orig.json deleted file mode 100644 index a7e327a04..000000000 --- a/examples/distill/Wan2.1-I2V/crush_smol/validation_orig.json +++ /dev/null @@ -1,76 +0,0 @@ -{ - "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, - "video_path": "validation_dataset/1gGQy4nxyUo-Scene-016.mp4", - "num_inference_steps": 40, - "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/examples/distill/Wan2.1-I2V/crush_smol/validation_repeat.json b/examples/distill/Wan2.1-I2V/crush_smol/validation_repeat.json new file mode 100644 index 000000000..30c3844a5 --- /dev/null +++ b/examples/distill/Wan2.1-I2V/crush_smol/validation_repeat.json @@ -0,0 +1,13 @@ +{ + "data": [ + { + "caption": "A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open.", + "image_path": null, + "video_path": "validation_dataset/1gGQy4nxyUo-Scene-016.mp4", + "num_inference_steps": 40, + "height": 480, + "width": 832, + "num_frames": 77 + } + ] +} \ No newline at end of file diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index ba79b95d0..d94fe38ad 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -118,6 +118,8 @@ class ForwardBatch: # Misc save_video: bool = True return_frames: bool = False + return_trajectory_latents: bool = False + trajectory_latents: list[torch.Tensor] = field(default_factory=list) # TeaCache parameters enable_teacache: bool = False diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index 7e75033b9..f1c08acc4 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -677,6 +677,7 @@ class DmdDenoisingStage(DenoisingStage): # TODO(yongqi) hard code prepare latents latents = torch.randn( latents.permute(0, 2, 1, 3, 4).shape, + # latents.shape, dtype=torch.bfloat16, device="cuda", generator=torch.Generator(device="cuda").manual_seed(42)) @@ -801,6 +802,9 @@ class DmdDenoisingStage(DenoisingStage): scheduler=self.scheduler).unflatten( 0, pred_noise.shape[:2]) + if batch.return_trajectory_latents: + batch.trajectory_latents.append(pred_video) + if i < len(timesteps) - 1: next_timestep = timesteps[i + 1] * torch.ones( [1], dtype=torch.long, device=pred_video.device) diff --git a/fastvideo/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py index e05fe83cf..17d50953d 100644 --- a/fastvideo/training/distillation_pipeline.py +++ b/fastvideo/training/distillation_pipeline.py @@ -241,7 +241,9 @@ class DistillationPipeline(TrainingPipeline): # 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 = 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] @@ -279,14 +281,15 @@ class DistillationPipeline(TrainingPipeline): # t_expand = current_timestep_tensor.repeat(current_noise_latents.shape[0]) pred_flow = self.transformer( training_batch_temp.input_kwargs["hidden_states"], - training_batch_temp.input_kwargs["encoder_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"], + 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) + ).permute(0, 2, 1, 3, 4) pred_clean = pred_noise_to_pred_video( pred_noise=pred_flow.flatten(0, 1), noise_input_latent=current_noise_latents.flatten(0, 1), @@ -338,6 +341,63 @@ class DistillationPipeline(TrainingPipeline): "generator_timestep"] = target_timestep.float().detach() return pred_video + def _generator_forward_with_validation_pipeline( + self, training_batch: TrainingBatch) -> torch.Tensor: + dtype = training_batch.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.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] + # self.training_args.pipeline_config.dmd_denoising_steps, + # from copy import deepcopy + # args = deepcopy(self.training_args) + # args.pipeline_config.dmd_denoising_steps = self.training_args.pipeline_config.dmd_denoising_steps[:target_timestep_idx_int] + + # Step 2: Simulate the multi-step inference process up to the target timestep + # Start from pure noise like in inference + current_noise_latents = torch.randn(self.video_latent_shape, + device=self.device, + dtype=dtype) + current_timestep = self.denoising_step_list[2] + current_timestep_tensor = current_timestep * torch.ones( + 1, device=self.device, dtype=torch.long) + training_batch_temp = self._build_distill_input_kwargs( + current_noise_latents, current_timestep_tensor, + training_batch.conditional_dict, training_batch) + manual_batch = ForwardBatch( + data_type="video", + # **shallow_asdict(self.training_args), + latents=current_noise_latents.permute(0, 2, 1, 3, 4), + generator=torch.Generator(device="cpu").manual_seed(self.seed), + image_embeds=[ + training_batch_temp.input_kwargs["encoder_hidden_states_image"] + ], + prompt_embeds=[ + training_batch_temp.input_kwargs["encoder_hidden_states"] + ], + image_latent=training_batch_temp.input_kwargs["image_latent"], + timesteps=current_timestep_tensor, + return_trajectory_latents=True, + # n_tokens=n_tokens, + # eta=0.0, + # VSA_sparsity=training_args.VSA_sparsity, + ) + batch = self.validation_pipeline.denoising_stage.forward( + batch=manual_batch, fastvideo_args=self.training_args) + + latents = batch.trajectory_latents[target_timestep_idx_int] + latents = latents + + training_batch.dmd_latent_vis_dict[ + "generator_timestep"] = target_timestep.float().detach() + + return latents + def _dmd_forward(self, generator_pred_video: torch.Tensor, training_batch: TrainingBatch) -> torch.Tensor: """Compute DMD (Diffusion Model Distillation) loss.""" @@ -443,8 +503,12 @@ class DistillationPipeline(TrainingPipeline): current_timestep=training_batch.timesteps, attn_metadata=training_batch.attn_metadata_vsa): if self.training_args.simulate_generator_forward: - generator_pred_video = self._generator_multi_step_simulation_forward( - training_batch) + if True: + generator_pred_video = self._generator_forward_with_validation_pipeline( + training_batch) + else: + generator_pred_video = self._generator_multi_step_simulation_forward( + training_batch) else: generator_pred_video = self._generator_forward(training_batch) @@ -578,8 +642,12 @@ class DistillationPipeline(TrainingPipeline): current_timestep=batch_gen.timesteps, attn_metadata=batch_gen.attn_metadata_vsa): if self.training_args.simulate_generator_forward: - generator_pred_video = self._generator_multi_step_simulation_forward( - batch_gen) + if True: + generator_pred_video = self._generator_forward_with_validation_pipeline( + batch_gen) + else: + generator_pred_video = self._generator_multi_step_simulation_forward( + batch_gen) else: generator_pred_video = self._generator_forward( batch_gen) diff --git a/fastvideo/training/wan_i2v_distillation_pipeline.py b/fastvideo/training/wan_i2v_distillation_pipeline.py index 122df390d..7c9246423 100644 --- a/fastvideo/training/wan_i2v_distillation_pipeline.py +++ b/fastvideo/training/wan_i2v_distillation_pipeline.py @@ -181,23 +181,23 @@ 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 = 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 = training_batch.image_embeds + assert torch.isnan(image_embeds).sum() == 0 image_embeds = image_embeds.to(get_local_torch_device(), dtype=torch.bfloat16) - assert torch.isnan(image_embeds).sum() == 0 + # 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( @@ -210,8 +210,9 @@ class WanI2VDistillationPipeline(DistillationPipeline): 4), # bs, c, t, h, w "encoder_hidden_states": text_dict["encoder_hidden_states"], "encoder_attention_mask": text_dict["encoder_attention_mask"], + "image_latent": training_batch.image_latents, "timestep": timestep, - "encoder_hidden_states_image": image_embeds.permute(0, 2, 1, 3, 4), + "encoder_hidden_states_image": image_embeds, "return_dict": False, } training_batch.noise_latents = noise_input