diff --git a/fastvideo/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py
index 6963c67a9..962dc4643 100644
--- a/fastvideo/training/distillation_pipeline.py
+++ b/fastvideo/training/distillation_pipeline.py
@@ -1318,6 +1318,7 @@ class DistillationPipeline(TrainingPipeline):
                         x = torchvision.utils.make_grid(x, nrow=6)
                         x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
                         frames.append((x * 255).numpy().astype(np.uint8))
+                    frames = self._post_process_validation_frames(frames, batch)
                     videos.append(frames)
                     audios.append(output_batch.extra.get("audio"))
                     audio_sample_rates.append(
diff --git a/fastvideo/training/matrixgame_self_forcing_distillation_pipeline.py b/fastvideo/training/matrixgame_self_forcing_distillation_pipeline.py
index 33b154a89..d28ebdd08 100644
--- a/fastvideo/training/matrixgame_self_forcing_distillation_pipeline.py
+++ b/fastvideo/training/matrixgame_self_forcing_distillation_pipeline.py
@@ -142,18 +142,18 @@ class MatrixGameSelfForcingDistillationPipeline(SelfForcingDistillationPipeline
             getattr(training_args, "warp_denoising_step", False)
             or getattr(pipeline_config, "warp_denoising_step", False)
         )
-        training_args.warp_denoising_step = warp_denoising_step
-        pipeline_config.warp_denoising_step = warp_denoising_step
-
-        self.use_diagonal_denoising = bool(
-            getattr(pipeline_config, "use_diagonal_denoising", False)
-        )
-        pipeline_config.use_diagonal_denoising = self.use_diagonal_denoising
-
-        raw_warmup_mid_steps = getattr(
-            pipeline_config, "diagonal_warmup_mid_steps", None
-        )
-        pipeline_config.diagonal_warmup_mid_steps = raw_warmup_mid_steps
+        training_args.warp_denoising_step = warp_denoising_step
+        pipeline_config.warp_denoising_step = warp_denoising_step
+
+        self.use_diagonal_denoising = bool(
+            getattr(pipeline_config, "use_diagonal_denoising", False)
+        )
+        pipeline_config.use_diagonal_denoising = self.use_diagonal_denoising
+
+        raw_warmup_mid_steps = getattr(
+            pipeline_config, "diagonal_warmup_mid_steps", None
+        )
+        pipeline_config.diagonal_warmup_mid_steps = raw_warmup_mid_steps
 
         self.denoising_step_list = resolve_denoising_steps(
             pipeline_config.dmd_denoising_steps,
@@ -1356,13 +1356,13 @@ class MatrixGameSelfForcingDistillationPipeline(SelfForcingDistillationPipeline
                     key_names[i]: bool(keyboard_cond[frame_idx, i])
                     for i in range(min(len(key_names), keyboard_cond.shape[1]))
                 }
-                draw_keys_on_frame(frame, keys, mode='universal')
+                draw_keys_on_frame(frame, keys)
 
             # Draw mouse overlay
             if mouse_cond is not None and frame_idx < len(mouse_cond):
                 pitch = float(mouse_cond[frame_idx, 0])
                 yaw = float(mouse_cond[frame_idx, 1])
-                draw_mouse_on_frame(frame, pitch, yaw)
+                draw_mouse_on_frame(frame, yaw=yaw, pitch=pitch)
 
             processed_frames.append(frame)
 
@@ -1388,4 +1388,4 @@ if __name__ == "__main__":
     parser = TrainingArgs.add_cli_args(parser)
     parser = FastVideoArgs.add_cli_args(parser)
     args = parser.parse_args()
-    main(args)
+    main(args)
\ No newline at end of file