From dfe0b288e182c4babb86bcb92d7fa66eb51f76d6 Mon Sep 17 00:00:00 2001
From: Yongqi Chen <144848849+BrianChen1129@users.noreply.github.com>
Date: Fri, 1 Aug 2025 23:15:37 -0400
Subject: [PATCH] [Bugfix] Add i2v vae loading (#686)

---
 fastvideo/training/wan_i2v_training_pipeline.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/fastvideo/training/wan_i2v_training_pipeline.py b/fastvideo/training/wan_i2v_training_pipeline.py
index 18810b0c1..9f07191ba 100644
--- a/fastvideo/training/wan_i2v_training_pipeline.py
+++ b/fastvideo/training/wan_i2v_training_pipeline.py
@@ -27,7 +27,7 @@ class WanI2VTrainingPipeline(TrainingPipeline):
     """
     A training pipeline for Wan.
     """
-    _required_config_modules = ["scheduler", "transformer"]
+    _required_config_modules = ["scheduler", "transformer", "vae"]
 
     def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
         self.modules["scheduler"] = FlowUniPCMultistepScheduler(