fix scheduler

fastvideo/v1/fastvideo_args.py

@@ -385,6 +385,7 @@ class TrainingArgs(FastVideoArgs):
     num_height: int = 0
     num_width: int = 0
     num_frames: int = 0
+    flow_shift: int = 0
 
     train_batch_size: int = 0
     num_latent_t: int = 0
@@ -503,6 +504,10 @@ class TrainingArgs(FastVideoArgs):
                             type=int,
                             required=True,
                             help="Number of frames")
+        parser.add_argument("--flow-shift",
+                            type=int,
+                            required=True,
+                            help="Flow shift")
 
         # Training batch and model configuration
         parser.add_argument("--train-batch-size",

fastvideo/v1/training/training_pipeline.py

@@ -429,7 +429,10 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
             self.seed)
         logger.info("Initialized random seeds with seed: %s", self.seed)
 
-        self.noise_scheduler = FlowMatchEulerDiscreteScheduler()
+        self.noise_scheduler = FlowMatchEulerDiscreteScheduler(
+            num_train_timesteps=self.noise_scheduler.num_train_timesteps,
+            shift=self.training_args.flow_shift,
+        )
 
         if self.training_args.resume_from_checkpoint:
             self._resume_from_checkpoint()