Merge branch 'main' of github.com:hao-ai-lab/FastVideo

fastvideo/pipelines/basic/wan/wan_i2v_dmd_pipeline.py

@@ -12,12 +12,10 @@ from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
 from fastvideo.pipelines.lora_pipeline import LoRAPipeline
 
 # isort: off
-from fastvideo.pipelines.stages import (ImageEncodingStage, ConditioningStage,
-                                        DecodingStage, DmdDenoisingStage,
-                                        EncodingStage, InputValidationStage,
-                                        LatentPreparationStage,
-                                        TextEncodingStage,
-                                        TimestepPreparationStage)
+from fastvideo.pipelines.stages import (
+    ImageEncodingStage, ConditioningStage, DecodingStage, DmdDenoisingStage,
+    ImageVAEEncodingStage, InputValidationStage, LatentPreparationStage,
+    TextEncodingStage, TimestepPreparationStage)
 # isort: on
 from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
     FlowMatchEulerDiscreteScheduler)
@@ -67,7 +65,7 @@ class WanImageToVideoDmdPipeline(LoRAPipeline, ComposedPipelineBase):
                            transformer=self.get_module("transformer")))
 
         self.add_stage(stage_name="image_latent_preparation_stage",
-                       stage=EncodingStage(vae=self.get_module("vae")))
+                       stage=ImageVAEEncodingStage(vae=self.get_module("vae")))
 
         self.add_stage(stage_name="denoising_stage",
                        stage=DmdDenoisingStage(

fastvideo/training/distillation_pipeline.py

@@ -241,6 +241,7 @@ class DistillationPipeline(TrainingPipeline):
                                             len(self.denoising_step_list), [1],
                                             device=self.device,
                                             dtype=torch.long)
+        target_timestep_idx_int = target_timestep_idx.item()
         target_timestep = self.denoising_step_list[target_timestep_idx]
 
         # Step 2: Simulate the multi-step inference process up to the target timestep
@@ -258,6 +259,8 @@ class DistillationPipeline(TrainingPipeline):
 
         # Only run intermediate steps if target_timestep_idx > 0
         max_target_idx = len(self.denoising_step_list) - 1
+        noise_latents = []
+        noise_latent_index = target_timestep_idx_int - 1
         if max_target_idx > 0:
             # Run student model for all steps before the target timestep
             with torch.no_grad():
@@ -294,11 +297,19 @@ class DistillationPipeline(TrainingPipeline):
                     current_noise_latents = self.noise_scheduler.add_noise(
                         pred_clean.flatten(0, 1), noise.flatten(0, 1),
                         next_timestep_tensor).unflatten(0, pred_clean.shape[:2])
+                    latent_copy = current_noise_latents.clone()
+                    noise_latents.append(latent_copy)
 
         # Step 3: Use the simulated noisy input for the final training step
         # For timestep index 0, this is pure noise
         # For timestep index k > 0, this is the result after k denoising steps + noise at target level
-        noisy_input = current_noise_latents
+        if noise_latent_index >= 0:
+            assert noise_latent_index < len(
+                self.denoising_step_list
+            ) - 1, "noise_latent_index is out of bounds"
+            noisy_input = noise_latents[noise_latent_index]
+        else:
+            noisy_input = current_noise_latents
 
         # Step 4: Final student prediction (this is what we train on)
         training_batch = self._build_distill_input_kwargs(
@@ -802,7 +813,7 @@ class DistillationPipeline(TrainingPipeline):
         dmd_latents_vis_dict = training_batch.dmd_latent_vis_dict
         fake_score_latents_vis_dict = training_batch.fake_score_latent_vis_dict
         fake_score_log_keys = ['generator_pred_video']
-        dmd_log_keys = ['faker_score_pred_video']
+        dmd_log_keys = ['faker_score_pred_video', 'real_score_pred_video']
 
         for latent_key in fake_score_log_keys:
             latents = fake_score_latents_vis_dict[latent_key]