Pass pre-commit tests

Co-authored-by: Zihang-He <z6he@ucsd.edu>

.github/workflows/pre-commit.yml

@@ -10,7 +10,7 @@ jobs:
     - uses: actions/checkout@v4
     - uses: actions/setup-python@v5
       with:
-        python-version: "3.10"
+        python-version: "3.12"
     - run: echo "::add-matcher::.github/workflows/matchers/actionlint.json"
     - run: echo "::add-matcher::.github/workflows/matchers/mypy.json"
     - uses: pre-commit/action@v3.0.1

.pre-commit-config.yaml

@@ -33,7 +33,7 @@ repos:
     args: [--in-place, --verbose]
     additional_dependencies: [toml] # TODO: Remove when yapf is upgraded
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.11.4
+  rev: v0.11.12
   hooks:
   - id: ruff
     args: [--output-format, github, --fix]
@@ -48,7 +48,7 @@ repos:
   hooks:
   - id: isort
 - repo: https://github.com/jackdewinter/pymarkdown
-  rev: v0.9.29
+  rev: v0.9.30
   hooks:
   - id: pymarkdown
     args: [fix]

docs/source/inference/optimizations.md

@@ -72,6 +72,12 @@ FastVideo will automatically detect and use `FA3` if it is installed when using
 pip install st_attn==0.0.4
 ```
 
+Then download STA mask strategy from Hugging Face
+
+```bash
+python scripts/huggingface/download_hf.py --repo_id=FastVideo/STA_Mask_Strategy --local_dir=assets/ --repo_type=dataset
+```
+
 Please see [this page](#sta-installation) for more installation instructions.
 
 (optimizations-sage)=

fastvideo/data_preprocess/preprocess.py

@@ -15,12 +15,8 @@ from fastvideo.v1.pipelines.preprocess_pipeline import PreprocessPipeline
 
 logger = init_logger(__name__)
 
-BASE_MODEL_PATH = "/workspace/data/Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
-                                  local_dir=os.path.join(
-                                      'data', BASE_MODEL_PATH))
-
 def main(args):
+    args.model_path = maybe_download_model(args.model_path)
     # Assume using torchrun
     local_rank = int(os.getenv("RANK", 0))
     rank = int(os.environ.get("RANK", 0))
@@ -31,7 +27,7 @@ def main(args):
     if not dist.is_initialized():
         dist.init_process_group(backend="nccl", init_method="env://", world_size=world_size, rank=local_rank)
 
-    pipeline_config = PipelineConfig.from_pretrained(MODEL_PATH)
+    pipeline_config = PipelineConfig.from_pretrained(args.model_path)
     kwargs = {
         "use_cpu_offload": False,
         "vae_precision": "fp32",
@@ -39,7 +35,7 @@ def main(args):
     }
     pipeline_config_args = shallow_asdict(pipeline_config)
     pipeline_config_args.update(kwargs)
-    fastvideo_args = FastVideoArgs(model_path=MODEL_PATH,
+    fastvideo_args = FastVideoArgs(model_path=args.model_path,
                                    num_gpus=world_size,
                                    device_str="cuda",
                                    **pipeline_config_args,
@@ -47,7 +43,7 @@ def main(args):
     fastvideo_args.check_fastvideo_args()
     fastvideo_args.device = torch.device(f"cuda:{local_rank}")
 
-    pipeline = PreprocessPipeline(MODEL_PATH, fastvideo_args)
+    pipeline = PreprocessPipeline(args.model_path, fastvideo_args)
     pipeline.forward(batch=None, fastvideo_args=fastvideo_args, args=args)
 
 
@@ -77,6 +73,17 @@ if __name__ == "__main__":
         default=8,
         help="Batch size (per device) for the training dataloader.",
     )
+    parser.add_argument(
+        "--samples_per_file",
+        type=int,
+        default=64
+    )
+    parser.add_argument(
+        "--flush_frequency",
+        type=int,
+        default=256,
+        help="how often to save to parquet files"
+    )
     parser.add_argument("--num_latent_t", type=int, default=28, help="Number of latent timesteps.")
     parser.add_argument("--max_height", type=int, default=480)
     parser.add_argument("--max_width", type=int, default=848)
@@ -108,4 +115,4 @@ if __name__ == "__main__":
     )
 
     args = parser.parse_args()
-    main(args)
+    main(args)

fastvideo/distill/solver.py

@@ -7,7 +7,7 @@ from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.schedulers.scheduling_utils import SchedulerMixin
 from diffusers.utils import BaseOutput, logging
 
-# from fastvideo.models.mochi_hf.pipeline_mochi import linear_quadratic_schedule
+from fastvideo.models.mochi_hf.pipeline_mochi import linear_quadratic_schedule
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
@@ -38,7 +38,6 @@ class PCMFMScheduler(SchedulerMixin, ConfigMixin):
         linear_range=0.5,
     ):
         if linear_quadratic:
-            raise NotImplementedError("Linear quadratic schedule is not implemented")
             linear_steps = int(num_train_timesteps * linear_range)
             sigmas = linear_quadratic_schedule(num_train_timesteps, linear_quadratic_threshold, linear_steps)
             sigmas = torch.tensor(sigmas).to(dtype=torch.float32)

fastvideo/models/mochi_hf/mochi_latents_utils.py

@@ -31,7 +31,7 @@ mochi_latents_std = torch.tensor([
 mochi_scaling_factor = 1.0
 
 
-def normalize_dit_input(model_type, latents, args=None):
+def normalize_dit_input(model_type, latents):
     if model_type == "mochi":
         latents_mean = mochi_latents_mean.to(latents.device, latents.dtype)
         latents_std = mochi_latents_std.to(latents.device, latents.dtype)
@@ -41,16 +41,5 @@ def normalize_dit_input(model_type, latents, args=None):
         return latents * 0.476986
     elif model_type == "hunyuan":
         return latents * 0.476986
-    elif model_type == "wan":
-        from fastvideo.v1.configs.models.vaes.wanvae import WanVAEConfig
-        vae_config = WanVAEConfig()
-        latents_mean = torch.tensor(vae_config.arch_config.latents_mean)
-        latents_std = 1.0 / torch.tensor(vae_config.arch_config.latents_std)
-
-
-        latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(device=latents.device)
-        latents_std = latents_std.view(1, -1, 1, 1, 1).to(device=latents.device)
-        latents = ((latents.float() - latents_mean) * latents_std).to(latents)
-        return latents
     else:
         raise NotImplementedError(f"model_type {model_type} not supported")

fastvideo/v1/configs/models/vaes/wanvae.py

@@ -63,7 +63,7 @@ class WanVAEArchConfig(VAEArchConfig):
 
 @dataclass
 class WanVAEConfig(VAEConfig):
-    arch_config: VAEArchConfig = field(default_factory=WanVAEArchConfig)
+    arch_config: WanVAEArchConfig = field(default_factory=WanVAEArchConfig)
     use_feature_cache: bool = True
 
     use_tiling: bool = False

fastvideo/v1/dataset/__init__.py

@@ -1,3 +1,5 @@
+import os
+
 from torchvision import transforms
 from torchvision.transforms import Lambda
 from transformers import AutoTokenizer
@@ -25,8 +27,8 @@ def getdataset(args, start_idx=0) -> T2V_dataset:
         *resize_topcrop,
         norm_fun,
     ])
-    # tokenizer = AutoTokenizer.from_pretrained("/storage/ongoing/new/Open-Sora-Plan/cache_dir/mt5-xxl", cache_dir=args.cache_dir)
-    tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_name,
+    tokenizer_path = os.path.join(args.model_path, "tokenizer")
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path,
                                               cache_dir=args.cache_dir)
     if args.dataset == "t2v":
         return T2V_dataset(args,

fastvideo/v1/dataset/parquet_datasets.py

@@ -43,7 +43,8 @@ class ParquetVideoTextDataset(Dataset):
         self.cfg_rate = cfg_rate
         self.num_latent_t = num_latent_t
         self.local_indices = None
-        self.plan_output_dir = os.path.join(self.path, "data_plan.json")
+        self.plan_output_dir = os.path.join(
+            self.path, f"data_plan_{self.world_size}_{self.sp_world_size}.json")
 
         ranks = get_sp_group().ranks
         group_ranks: List[List] = [[] for _ in range(self.world_size)]
@@ -54,6 +55,7 @@ class ParquetVideoTextDataset(Dataset):
             # This will be useful when resume training
             if os.path.exists(self.plan_output_dir):
                 print(f"Using existing plan from {self.plan_output_dir}")
+                dist.barrier()
                 return
 
             # Find all parquet files recursively, and record num_rows for each file
@@ -87,6 +89,7 @@ class ParquetVideoTextDataset(Dataset):
 
             with open(self.plan_output_dir, "w") as f:
                 json.dump(plan, f)
+        dist.barrier()
 
     def __len__(self):
         if self.local_indices is None:

fastvideo/v1/distributed/parallel_state.py

@@ -655,7 +655,7 @@ class GroupCoordinator:
                 tensor_dict[key] = value
         return tensor_dict
 
-    def barrier(self):
+    def barrier(self) -> None:
         """Barrier synchronization among the group.
         NOTE: don't use `device_group` here! `barrier` in NCCL is
         terrible because it is internally a broadcast operation with

fastvideo/v1/entrypoints/data_preprocessor.py

@@ -1,29 +0,0 @@
-from torch.utils.data import DataLoader
-from torch.utils.data.distributed import DistributedSampler
-
-from fastvideo.v1.pipelines.wan.wan_latent_pipeline import WanLatentPipeline
-
-
-def main():
-    print("Starting data preprocessor")
-    pipeline = WanLatentPipeline.from_pretrained(
-        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
-
-    train_dataset = getdataset(args)
-    sampler = DistributedSampler(train_dataset,
-                                 rank=local_rank,
-                                 num_replicas=world_size,
-                                 shuffle=True)
-    train_dataloader = DataLoader(
-        train_dataset,
-        sampler=sampler,
-        batch_size=args.train_batch_size,
-        num_workers=args.dataloader_num_workers,
-    )
-
-    for batch in train_dataloader:
-        pipeline(batch)
-
-
-if __name__ == "__main__":
-    main()

fastvideo/v1/fastvideo_args.py

@@ -70,7 +70,7 @@ class FastVideoArgs:
     # Text encoder configuration
     DEFAULT_TEXT_ENCODER_PRECISIONS = (
         "fp16",
-        # "fp16",
+        "fp16",
     )
     text_encoder_precisions: Tuple[str, ...] = field(
         default_factory=lambda: FastVideoArgs.DEFAULT_TEXT_ENCODER_PRECISIONS)
@@ -478,7 +478,7 @@ class TrainingArgs(FastVideoArgs):
     output_dir: str = ""
     checkpoints_total_limit: int = 0
     checkpointing_steps: int = 0
-    resume_from_checkpoint: str = ""
+    resume_from_checkpoint: bool = False
     logging_dir: str = ""
 
     # optimizer & scheduler

fastvideo/v1/models/loader/component_loader.py

@@ -394,30 +394,23 @@ class TransformerLoader(ComponentLoader):
         default_dtype = PRECISION_TO_TYPE[fastvideo_args.precision]
 
         # Load the model using FSDP loader
-        logger.info("Loading model from %s, default_dtype: %s", cls_name, default_dtype)
-        # model = load_fsdp_model(model_cls=model_cls,
-        #                         init_params={
-        #                             "config": dit_config,
-        #                             "hf_config": hf_config
-        #                         },
-        #                         weight_dir_list=safetensors_list,
-        #                         device=fastvideo_args.device,
-        #                         cpu_offload=fastvideo_args.use_cpu_offload,
-        #                         default_dtype=default_dtype)
-        model = load_fsdp_model(model_cls=model_cls,
-                                init_params={
-                                    "config": dit_config,
-                                    "hf_config": hf_config
-                                },
-                                weight_dir_list=safetensors_list,
-                                device=fastvideo_args.device,
-                                cpu_offload=fastvideo_args.use_cpu_offload,
-                                default_dtype=default_dtype,
-                                # TODO(will): make these configurable
-                                param_dtype=torch.bfloat16,
-                                reduce_dtype=torch.float32,
-                                output_dtype=None,
-                                )
+        logger.info("Loading model from %s, default_dtype: %s", cls_name,
+                    default_dtype)
+        model = load_fsdp_model(
+            model_cls=model_cls,
+            init_params={
+                "config": dit_config,
+                "hf_config": hf_config
+            },
+            weight_dir_list=safetensors_list,
+            device=fastvideo_args.device,
+            cpu_offload=fastvideo_args.use_cpu_offload,
+            default_dtype=default_dtype,
+            # TODO(will): make these configurable
+            param_dtype=torch.bfloat16,
+            reduce_dtype=torch.float32,
+            output_dtype=None,
+        )
         if fastvideo_args.enable_torch_compile:
             logger.info("Torch Compile enabled for DiT")
             for n, m in reversed(list(model.named_modules())):

fastvideo/v1/models/loader/fsdp_load.py

@@ -14,14 +14,15 @@ from typing import (Any, Callable, DefaultDict, Dict, Generator, Hashable, List,
 import torch
 from torch import nn
 from torch.distributed import DeviceMesh, init_device_mesh
-from torch.distributed.fsdp import CPUOffloadPolicy, fully_shard, MixedPrecisionPolicy
 from torch.distributed._tensor import distribute_tensor
+from torch.distributed.fsdp import (CPUOffloadPolicy, MixedPrecisionPolicy,
+                                    fully_shard)
 from torch.nn.modules.module import _IncompatibleKeys
 
 from fastvideo.v1.distributed.parallel_state import (
     get_sequence_model_parallel_world_size)
-from fastvideo.v1.models.loader.weight_utils import safetensors_weights_iterator
 from fastvideo.v1.logger import init_logger
+from fastvideo.v1.models.loader.weight_utils import safetensors_weights_iterator
 
 logger = init_logger(__name__)
 
@@ -89,11 +90,6 @@ def get_param_names_mapping(
 
 
 # TODO(PY): add compile option
-    # param_dtype: torch.dtype,
-    # reduce_dtype: torch.dtype,
-    # output_dtype: torch.dtype,
-    # pp_enabled: bool = False,
-    # cpu_offload: bool = False,
 def load_fsdp_model(
     model_cls: Type[nn.Module],
     init_params: Dict[str, Any],
@@ -106,9 +102,11 @@ def load_fsdp_model(
     output_dtype: Optional[torch.dtype] = None,
 ) -> torch.nn.Module:
 
-    mp_policy = MixedPrecisionPolicy(param_dtype, reduce_dtype, output_dtype, cast_forward_inputs=True)
+    mp_policy = MixedPrecisionPolicy(param_dtype,
+                                     reduce_dtype,
+                                     output_dtype,
+                                     cast_forward_inputs=True)
 
-    # with set_default_dtype(default_dtype), torch.device("meta"):
     with set_default_dtype(default_dtype), torch.device("meta"):
         model = model_cls(**init_params)
 
@@ -231,10 +229,6 @@ def load_fsdp_model_from_full_model_state_dict(
         NotImplementedError: If got FSDP with more than 1D.
     """
     meta_sharded_sd = model.state_dict()
-    # s = fully_shard.state(model)
-    # logger.info(f"type(s): {type(s)}")
-    # logger.info(f"s: {s}")
-    # import pdb; pdb.set_trace()
 
     sharded_sd = {}
     to_merge_params: DefaultDict[Hashable, Dict[Any, Any]] = defaultdict(dict)

fastvideo/v1/pipelines/composed_pipeline_base.py

@@ -40,6 +40,8 @@ class ComposedPipelineBase(ABC):
 
     is_video_pipeline: bool = False  # To be overridden by video pipelines
     _required_config_modules: List[str] = []
+    training_args: Optional[TrainingArgs] = None
+    fastvideo_args: Optional[FastVideoArgs] = None
 
     # TODO(will): args should support both inference args and training args
     def __init__(self,
@@ -51,7 +53,15 @@ class ComposedPipelineBase(ABC):
         Initialize the pipeline. After __init__, the pipeline should be ready to
         use. The pipeline should be stateless and not hold any batch state.
         """
-        self.fastvideo_args = fastvideo_args
+
+        if fastvideo_args.training_mode:
+            assert isinstance(fastvideo_args, TrainingArgs)
+            self.training_args = fastvideo_args
+            assert self.training_args is not None
+        else:
+            self.fastvideo_args = fastvideo_args
+            assert self.fastvideo_args is not None
+
         self.model_path = model_path
         self._stages: List[PipelineStage] = []
         self._stage_name_mapping: Dict[str, PipelineStage] = {}
@@ -77,27 +87,22 @@ class ComposedPipelineBase(ABC):
         self.modules = self.load_modules(fastvideo_args)
 
         if fastvideo_args.training_mode:
-            if fastvideo_args.log_validation:
-                self.initialize_validation_pipeline(fastvideo_args)
-            self.initialize_training_pipeline(fastvideo_args)
+            assert self.training_args is not None
+            if self.training_args.log_validation:
+                self.initialize_validation_pipeline(self.training_args)
+            self.initialize_training_pipeline(self.training_args)
 
         self.initialize_pipeline(fastvideo_args)
 
-        # logger.info("Creating pipeline stages...")
-        # self.create_pipeline_stages(fastvideo_args)
-
-        if fastvideo_args.training_mode:
-            logger.info("Creating training pipeline stages...")
-            self.create_training_stages(fastvideo_args)
-        else:
+        if not fastvideo_args.training_mode:
             logger.info("Creating pipeline stages...")
             self.create_pipeline_stages(fastvideo_args)
 
-    def initialize_training_pipeline(self, fastvideo_args: FastVideoArgs):
+    def initialize_training_pipeline(self, training_args: TrainingArgs):
         raise NotImplementedError(
             "if training_mode is True, the pipeline must implement this method")
 
-    def initialize_validation_pipeline(self, fastvideo_args: FastVideoArgs):
+    def initialize_validation_pipeline(self, training_args: TrainingArgs):
         raise NotImplementedError(
             "if log_validation is True, the pipeline must implement this method"
         )
@@ -134,7 +139,7 @@ class ComposedPipelineBase(ABC):
             config_args = shallow_asdict(config)
             config_args.update(kwargs)
 
-        if args.inference_mode:
+        if args is None or args.inference_mode:
             fastvideo_args = FastVideoArgs(model_path=model_path,
                                            device_str=device or "cuda" if
                                            torch.cuda.is_available() else "cpu",
@@ -155,7 +160,6 @@ class ComposedPipelineBase(ABC):
             for key, value in config_args.items():
                 setattr(fastvideo_args, key, value)
 
-            # we use cpu offload for training 
             fastvideo_args.use_cpu_offload = False
             # make sure we are in training mode
             fastvideo_args.inference_mode = False
@@ -168,7 +172,7 @@ class ComposedPipelineBase(ABC):
 
         fastvideo_args.check_fastvideo_args()
 
-        logger.info(f"fastvideo_args in from_pretrained: {fastvideo_args}")
+        logger.info("fastvideo_args in from_pretrained: %s", fastvideo_args)
 
         return cls(model_path,
                    fastvideo_args,
@@ -190,6 +194,8 @@ class ComposedPipelineBase(ABC):
         init_distributed_environment(world_size=world_size,
                                      rank=rank,
                                      local_rank=local_rank)
+        assert fastvideo_args.tp_size is not None, "tp_size must be set"
+        assert fastvideo_args.sp_size is not None, "sp_size must be set"
         initialize_model_parallel(
             tensor_model_parallel_size=fastvideo_args.tp_size,
             sequence_model_parallel_size=fastvideo_args.sp_size)
@@ -244,14 +250,7 @@ class ComposedPipelineBase(ABC):
         """
         raise NotImplementedError
 
-    # @abstractmethod
-    # def create_validation_stages(self, fastvideo_args: FastVideoArgs):
-    #     """
-    #     Create the validation pipeline stages.
-    #     """
-    #     raise NotImplementedError
-
-    def create_training_stages(self, fastvideo_args: FastVideoArgs):
+    def create_training_stages(self, training_args: TrainingArgs):
         """
         Create the training pipeline stages.
         """

fastvideo/v1/pipelines/preprocess_pipeline.py

@@ -9,6 +9,7 @@ import gc
 import multiprocessing
 import os
 from concurrent.futures import ProcessPoolExecutor
+from typing import Any, Dict
 
 import numpy as np
 import pyarrow as pa
@@ -52,8 +53,8 @@ class PreprocessPipeline(ComposedPipelineBase):
         args,
     ):
         # Initialize class variables for data sharing
-        self.video_data = {}  # Store video metadata and paths
-        self.latent_data = {}  # Store latent tensors
+        self.video_data: Dict[str, Any] = {}  # Store video metadata and paths
+        self.latent_data: Dict[str, Any] = {}  # Store latent tensors
         self.preprocess_validation_text(fastvideo_args, args)
         self.preprocess_video_and_text(fastvideo_args, args)
 
@@ -126,7 +127,6 @@ class PreprocessPipeline(ComposedPipelineBase):
                         valid_data["pixel_values"].to(
                             fastvideo_args.device)).mean
 
-                # Get corresponding captions for this batch
                 batch_captions = valid_data["text"]
 
                 batch = ForwardBatch(
@@ -135,6 +135,7 @@ class PreprocessPipeline(ComposedPipelineBase):
                     prompt_embeds=[],
                     prompt_attention_mask=[],
                 )
+                assert hasattr(self, "prompt_encoding_stage")
                 result_batch = self.prompt_encoding_stage(batch, fastvideo_args)
                 prompt_embeds, prompt_attention_mask = result_batch.prompt_embeds[
                     0], result_batch.prompt_attention_mask[0]
@@ -142,7 +143,7 @@ class PreprocessPipeline(ComposedPipelineBase):
 
                 # Get sequence lengths from attention masks (number of 1s)
                 seq_lens = prompt_attention_mask.sum(dim=1)
-                # Create a list to store non-padded embeddings and masks
+
                 non_padded_embeds = []
                 non_padded_masks = []
 
@@ -266,7 +267,7 @@ class PreprocessPipeline(ComposedPipelineBase):
                     self.all_tables = []
                 self.all_tables.append(table)
 
-                logger.info(f"Collected batch with {len(table)} samples")
+                logger.info("Collected batch with %s samples", len(table))
 
             if num_processed_samples >= args.flush_frequency:
                 assert hasattr(self, 'all_tables') and self.all_tables
@@ -295,7 +296,7 @@ class PreprocessPipeline(ComposedPipelineBase):
                 print(
                     f"Using {num_workers} workers to process {total_chunks} chunks"
                 )
-                logger.info(f"Chunks per worker: {chunks_per_worker}")
+                logger.info("Chunks per worker: %s", chunks_per_worker)
 
                 # Prepare work ranges
                 work_ranges = []
@@ -319,30 +320,28 @@ class PreprocessPipeline(ComposedPipelineBase):
                         try:
                             written = future.result()
                             total_written += written
-                            logger.info(
-                                f"Processed chunk with {written} samples")
+                            logger.info("Processed chunk with %s samples",
+                                        written)
                         except Exception as e:
                             work_range = futures[future]
                             failed_ranges.append(work_range)
-                            logger.error(
-                                f"Failed to process range {work_range[0]}-{work_range[1]}: {str(e)}"
-                            )
+                            logger.error("Failed to process range %s-%s: %s",
+                                         work_range[0], work_range[1], str(e))
 
                 # Retry failed ranges sequentially
                 if failed_ranges:
-                    logger.warning(
-                        f"Retrying {len(failed_ranges)} failed ranges sequentially"
-                    )
+                    logger.warning("Retrying %s failed ranges sequentially",
+                                   len(failed_ranges))
                     for work_range in failed_ranges:
                         try:
                             total_written += self.process_chunk_range(
                                 work_range)
                         except Exception as e:
                             logger.error(
-                                f"Failed to process range {work_range[0]}-{work_range[1]} after retry: {str(e)}"
-                            )
+                                "Failed to process range %s-%s after retry: %s",
+                                work_range[0], work_range[1], str(e))
 
-                logger.info(f"Total samples written: {total_written}")
+                logger.info("Total samples written: %s", total_written)
 
                 num_processed_samples = 0
                 self.all_tables = []
@@ -353,7 +352,6 @@ class PreprocessPipeline(ComposedPipelineBase):
                                               "validation_parquet_dataset")
         os.makedirs(validation_parquet_dir, exist_ok=True)
 
-
         with open(args.validation_prompt_txt, encoding="utf-8") as file:
             lines = file.readlines()
         prompts = [line.strip() for line in lines]
@@ -374,6 +372,7 @@ class PreprocessPipeline(ComposedPipelineBase):
                     prompt_embeds=[],
                     prompt_attention_mask=[],
                 )
+                assert hasattr(self, "prompt_encoding_stage")
                 result_batch = self.prompt_encoding_stage(batch, fastvideo_args)
             prompt_embeds = result_batch.prompt_embeds[0]
             prompt_attention_mask = result_batch.prompt_attention_mask[0]
@@ -382,15 +381,15 @@ class PreprocessPipeline(ComposedPipelineBase):
 
             # Get the sequence length from attention mask (number of 1s)
             seq_len = prompt_attention_mask.sum().item()
-            # Slice the embeddings to keep only the non-padding parts
+
             text_embedding = prompt_embeds[0, :seq_len].cpu().numpy()
             text_attention_mask = prompt_attention_mask[
                 0, :seq_len].cpu().numpy().astype(np.uint8)
 
             # Log the shapes after removing padding
             logger.info(
-                f"Shape after removing padding - Embeddings: {text_embedding.shape}, Mask: {text_attention_mask.shape}"
-            )
+                "Shape after removing padding - Embeddings: %s, Mask: %s",
+                text_embedding.shape, text_attention_mask.shape)
 
             # Create record for Parquet dataset
             record = {
@@ -415,7 +414,7 @@ class PreprocessPipeline(ComposedPipelineBase):
             }
             batch_data.append(record)
 
-            logger.info(f"Saved validation sample: {file_name}")
+            logger.info("Saved validation sample: %s", file_name)
 
         if batch_data:
             # Add progress bar for writing to Parquet dataset
@@ -468,7 +467,7 @@ class PreprocessPipeline(ComposedPipelineBase):
             write_pbar.update(1)
             write_pbar.close()
 
-            logger.info(f"Total validation samples: {len(table)}")
+            logger.info("Total validation samples: %s", len(table))
 
             work_range = (0, 1, table, 0, validation_parquet_dir, len(table))
 
@@ -485,30 +484,28 @@ class PreprocessPipeline(ComposedPipelineBase):
                     except Exception as e:
                         work_range = futures[future]
                         failed_ranges.append(work_range)
-                        logger.error(
-                            f"Failed to process range {work_range[0]}-{work_range[1]}: {str(e)}"
-                        )
+                        logger.error("Failed to process range %s-%s: %s",
+                                     work_range[0], work_range[1], str(e))
 
-            # Retry failed ranges sequentially
             if failed_ranges:
-                logger.warning(
-                    f"Retrying {len(failed_ranges)} failed ranges sequentially")
+                logger.warning("Retrying %s failed ranges sequentially",
+                               len(failed_ranges))
                 for work_range in failed_ranges:
                     try:
                         total_written += self.process_chunk_range(work_range)
                     except Exception as e:
                         logger.error(
-                            f"Failed to process range {work_range[0]}-{work_range[1]} after retry: {str(e)}"
-                        )
+                            "Failed to process range %s-%s after retry: %s",
+                            work_range[0], work_range[1], str(e))
 
-            logger.info(f"Total validation samples written: {total_written}")
+            logger.info("Total validation samples written: %s", total_written)
 
             # Clear memory
             del table
             gc.collect()  # Force garbage collection
 
     @staticmethod
-    def process_chunk_range(args):
+    def process_chunk_range(args: Any) -> int:
         start_idx, end_idx, table, worker_id, output_dir, samples_per_file = args
         try:
             total_written = 0
@@ -554,10 +551,9 @@ class PreprocessPipeline(ComposedPipelineBase):
 
             return total_written
         except Exception as e:
-            logger.error(
-                f"Error processing chunks {start_idx}-{end_idx} for worker {worker_id}: {str(e)}"
-            )
+            logger.error("Error processing chunks %s-%s for worker %s: %s",
+                         start_idx, end_idx, worker_id, str(e))
             raise
 
 
-EntryClass = PreprocessPipeline
+EntryClass = PreprocessPipeline

fastvideo/v1/pipelines/stages/denoising.py

@@ -23,7 +23,6 @@ from fastvideo.v1.logger import init_logger
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.v1.pipelines.stages.base import PipelineStage
 from fastvideo.v1.platforms import _Backend
-from fastvideo.v1.utils import PRECISION_TO_TYPE
 
 st_attn_available = False
 spec = importlib.util.find_spec("st_attn")
@@ -74,6 +73,7 @@ class DenoisingStage(PipelineStage):
         )
 
         # Setup precision and autocast settings
+        # TODO(will): make the precision configurable for inference
         # target_dtype = PRECISION_TO_TYPE[fastvideo_args.precision]
         target_dtype = torch.bfloat16
         autocast_enabled = (target_dtype != torch.float32
@@ -84,7 +84,6 @@ class DenoisingStage(PipelineStage):
         ), get_sequence_model_parallel_rank()
         sp_group = world_size > 1
         if sp_group:
-            # b c t h w -> b t n s h w
             latents = rearrange(batch.latents,
                                 "b t (n s) h w -> b t n s h w",
                                 n=world_size).contiguous()
@@ -190,7 +189,7 @@ class DenoisingStage(PipelineStage):
 
                 # Predict noise residual
                 with torch.autocast(device_type="cuda",
-                                    dtype=torch.bfloat16,
+                                    dtype=target_dtype,
                                     enabled=autocast_enabled):
 
                     # TODO(will-refactor): all of this should be in the stage's init

fastvideo/v1/pipelines/training_pipeline.py

@@ -1,841 +0,0 @@
-import gc
-import os
-import sys
-import time
-import traceback
-from abc import ABC, abstractmethod
-from collections import deque
-from copy import deepcopy
-
-import imageio
-import numpy as np
-import torch
-import torchvision
-from diffusers.optimization import get_scheduler
-from einops import rearrange
-from torchdata.stateful_dataloader import StatefulDataLoader
-from tqdm.auto import tqdm
-
-# import torch.distributed as dist
-import wandb
-from fastvideo.models.mochi_hf.mochi_latents_utils import normalize_dit_input
-from fastvideo.v1.configs.sample import SamplingParam
-from fastvideo.v1.dataset.parquet_datasets import ParquetVideoTextDataset
-from fastvideo.v1.distributed import cleanup_dist_env_and_memory, get_sp_group
-from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs
-from fastvideo.v1.forward_context import set_forward_context
-from fastvideo.v1.logger import init_logger
-from fastvideo.v1.pipelines import ComposedPipelineBase
-from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.v1.pipelines.training_utils import (
-    _clip_grad_norm_while_handling_failing_dtensor_cases,
-    compute_density_for_timestep_sampling, get_sigmas, save_checkpoint)
-from fastvideo.v1.pipelines.wan.wan_pipeline import WanValidationPipeline
-
-logger = init_logger(__name__)
-
-# Manual gradient checking flag - set to True to enable gradient verification
-ENABLE_GRADIENT_CHECK = False
-GRADIENT_CHECK_DTYPE = torch.bfloat16
-
-
-class TrainingPipeline(ComposedPipelineBase, ABC):
-    """
-    A pipeline for training a model. All training pipelines should inherit from this class.
-    All reusable components and code should be implemented in this class.
-    """
-    _required_config_modules = ["scheduler", "transformer"]
-
-    def initialize_training_pipeline(self, fastvideo_args: TrainingArgs):
-        logger.info("Initializing training pipeline...")
-        self.device = fastvideo_args.device
-        self.sp_group = get_sp_group()
-        self.world_size = self.sp_group.world_size
-        self.rank = self.sp_group.rank
-        self.local_rank = self.sp_group.local_rank
-        self.transformer = self.get_module("transformer")
-        assert self.transformer is not None
-
-        self.transformer.requires_grad_(True)
-        self.transformer.train()
-
-        args = fastvideo_args
-
-        noise_scheduler = self.modules["scheduler"]
-        params_to_optimize = self.transformer.parameters()
-        params_to_optimize = list(
-            filter(lambda p: p.requires_grad, params_to_optimize))
-
-        optimizer = torch.optim.AdamW(
-            params_to_optimize,
-            lr=args.learning_rate,
-            betas=(0.9, 0.999),
-            weight_decay=args.weight_decay,
-            eps=1e-8,
-        )
-
-        init_steps = 0
-        logger.info("optimizer: %s", optimizer)
-
-        # todo add lr scheduler
-        lr_scheduler = get_scheduler(
-            args.lr_scheduler,
-            optimizer=optimizer,
-            num_warmup_steps=args.lr_warmup_steps * self.world_size,
-            num_training_steps=args.max_train_steps * self.world_size,
-            num_cycles=args.lr_num_cycles,
-            power=args.lr_power,
-            last_epoch=init_steps - 1,
-        )
-
-        train_dataset = ParquetVideoTextDataset(
-            args.data_path,
-            batch_size=args.train_batch_size,
-            rank=self.rank,
-            world_size=self.world_size,
-            cfg_rate=args.cfg,
-            num_latent_t=args.num_latent_t)
-
-        train_dataloader = StatefulDataLoader(
-            train_dataset,
-            batch_size=args.train_batch_size,
-            num_workers=args.
-            dataloader_num_workers,  # Reduce number of workers to avoid memory issues
-            prefetch_factor=2,
-            shuffle=False,
-            pin_memory=True,
-            drop_last=True)
-
-        self.lr_scheduler = lr_scheduler
-        self.train_dataset = train_dataset
-        self.train_dataloader = train_dataloader
-        self.init_steps = init_steps
-        self.optimizer = optimizer
-        self.noise_scheduler = noise_scheduler
-        # self.noise_random_generator = noise_random_generator
-
-        # num_update_steps_per_epoch = math.ceil(
-        #     len(train_dataloader) / args.gradient_accumulation_steps *
-        #     args.sp_size / args.train_sp_batch_size)
-        # args.num_train_epochs = math.ceil(args.max_train_steps /
-        #                                   num_update_steps_per_epoch)
-
-        if self.rank <= 0:
-            project = args.tracker_project_name or "fastvideo"
-            wandb.init(project=project, config=args)
-
-    @abstractmethod
-    def initialize_validation_pipeline(self, fastvideo_args: FastVideoArgs):
-        raise NotImplementedError(
-            "Training pipelines must implement this method")
-
-    @abstractmethod
-    def train_one_step(self, transformer, model_type, optimizer, lr_scheduler,
-                       loader, noise_scheduler, noise_random_generator,
-                       gradient_accumulation_steps, sp_size,
-                       precondition_outputs, max_grad_norm, weighting_scheme,
-                       logit_mean, logit_std, mode_scale):
-        """
-        Train one step of the model.
-        """
-        raise NotImplementedError(
-            "Training pipeline must implement this method")
-
-    def log_validation(self, transformer, fastvideo_args, global_step):
-        fastvideo_args.inference_mode = True
-        fastvideo_args.use_cpu_offload = False
-        if not fastvideo_args.log_validation:
-            return
-        if self.validation_pipeline is None:
-            raise ValueError("Validation pipeline is not set")
-
-        # Create sampling parameters if not provided
-        sampling_param = SamplingParam.from_pretrained(
-            fastvideo_args.model_path)
-
-        # Prepare validation prompts
-        print('fastvideo_args.validation_prompt_dir',
-              fastvideo_args.validation_prompt_dir)
-        validation_dataset = ParquetVideoTextDataset(
-            fastvideo_args.validation_prompt_dir,
-            batch_size=1,
-            rank=0,
-            world_size=1,
-            cfg_rate=0,
-            num_latent_t=args.num_latent_t)
-
-        validation_dataloader = StatefulDataLoader(
-            validation_dataset,
-            batch_size=1,
-            num_workers=1,  # Reduce number of workers to avoid memory issues
-            prefetch_factor=2,
-            shuffle=False,
-            pin_memory=True,
-            drop_last=False)
-
-        transformer.requires_grad_(False)
-        for p in transformer.parameters():
-            p.requires_grad = False
-        transformer.eval()
-
-        # Add the transformer to the validation pipeline
-        self.validation_pipeline.add_module("transformer", transformer)
-        self.validation_pipeline.latent_preparation_stage.transformer = transformer
-        self.validation_pipeline.denoising_stage.transformer = transformer
-
-        # Process each validation prompt
-        videos = []
-        captions = []
-        for _, embeddings, masks, infos in validation_dataloader:
-            logger.info(f"infos: {infos}")
-            caption = infos['caption']
-            captions.append(caption)
-            prompt_embeds = embeddings.to(fastvideo_args.device).to(torch.bfloat16)
-            prompt_attention_mask = masks.to(fastvideo_args.device).to(torch.bfloat16)
-
-            # Calculate sizes
-            latents_size = [(sampling_param.num_frames - 1) // 4 + 1,
-                            sampling_param.height // 8,
-                            sampling_param.width // 8]
-            n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
-            logger.info('embed dtype', prompt_embeds.dtype)
-
-            num_frames = (fastvideo_args.num_latent_t - 1) * 4 + 1
-            logger.info(f"validation num_frames: {num_frames}")
-            # Prepare batch for validation
-            # print('shape of embeddings', prompt_embeds.shape)
-            batch = ForwardBatch(
-                # **shallow_asdict(sampling_param),
-                data_type="video",
-                latents=None,
-                # seed=sampling_param.seed,
-                # data_type="video",
-                prompt_embeds=[prompt_embeds],
-                prompt_attention_mask=[prompt_attention_mask],
-                # make sure we use the same height, width, and num_frames as the training pipeline
-                height=args.num_height,
-                width=args.num_width,
-                num_frames=num_frames,
-                # num_inference_steps=fastvideo_args.validation_sampling_steps,
-                num_inference_steps=50,
-                # guidance_scale=fastvideo_args.validation_guidance_scale,
-                guidance_scale=1,
-                n_tokens=n_tokens,
-                do_classifier_free_guidance=False,
-                eta=0.0,
-                extra={},
-            )
-
-            # Run validation inference
-            with torch.autocast("cuda", dtype=torch.bfloat16):
-                with torch.inference_mode():
-                    output_batch = self.validation_pipeline.forward(
-                        batch, fastvideo_args)
-                    samples = output_batch.output
-
-            # Process outputs
-            video = rearrange(samples, "b c t h w -> t b c h w")
-            frames = []
-            for x in video:
-                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))
-            videos.append(frames)
-
-        # Log validation results
-        rank = int(os.environ.get("RANK", 0))
-
-        if rank == 0:
-            video_filenames = []
-            video_captions = []
-            for i, video in enumerate(videos):
-                caption = captions[i]
-                os.makedirs(fastvideo_args.output_dir, exist_ok=True)
-                filename = os.path.join(
-                    fastvideo_args.output_dir,
-                    f"validation_step_{global_step}_video_{i}.mp4")
-                imageio.mimsave(filename, video, fps=sampling_param.fps)
-                video_filenames.append(filename)
-                video_captions.append(
-                    caption)  # Store the caption for each video
-
-            logs = {
-                "validation_videos": [
-                    wandb.Video(filename,
-                                caption=caption) for filename, caption in zip(
-                                    video_filenames, video_captions)
-                ]
-            }
-            wandb.log(logs, step=global_step)
-
-        # Re-enable gradients for training
-        transformer.requires_grad_(True)
-        transformer.train()
-
-        gc.collect()
-        torch.cuda.empty_cache()
-
-    def gradient_check_parameters(self,
-                                  transformer,
-                                  latents,
-                                  encoder_hidden_states,
-                                  encoder_attention_mask,
-                                  timesteps,
-                                  target,
-                                  eps=5e-2,
-                                  max_params_to_check=2000):
-        """
-        Verify gradients using finite differences for FSDP models with GRADIENT_CHECK_DTYPE.
-        Uses standard tolerances for GRADIENT_CHECK_DTYPE precision.
-        """
-        # Move all inputs to CPU and clear GPU memory
-        inputs_cpu = {
-            'latents': latents.cpu(),
-            'encoder_hidden_states': encoder_hidden_states.cpu(),
-            'encoder_attention_mask': encoder_attention_mask.cpu(),
-            'timesteps': timesteps.cpu(),
-            'target': target.cpu()
-        }
-        del latents, encoder_hidden_states, encoder_attention_mask, timesteps, target
-        torch.cuda.empty_cache()
-
-        def compute_loss():
-            # Move inputs to GPU, compute loss, cleanup
-            inputs_gpu = {
-                k:
-                v.to(self.fastvideo_args.device,
-                     dtype=GRADIENT_CHECK_DTYPE
-                     if k != 'encoder_attention_mask' else None)
-                for k, v in inputs_cpu.items()
-            }
-
-            # Use GRADIENT_CHECK_DTYPE for more accurate gradient checking
-            # with torch.autocast(enabled=False, device_type="cuda"):
-            with torch.autocast("cuda", dtype=GRADIENT_CHECK_DTYPE):
-                with set_forward_context(
-                        current_timestep=inputs_gpu['timesteps'],
-                        attn_metadata=None):
-                    model_pred = transformer(
-                        hidden_states=inputs_gpu['latents'],
-                        encoder_hidden_states=inputs_gpu[
-                            'encoder_hidden_states'],
-                        timestep=inputs_gpu['timesteps'],
-                        encoder_attention_mask=inputs_gpu[
-                            'encoder_attention_mask'],
-                        return_dict=False)[0]
-
-                if self.fastvideo_args.precondition_outputs:
-                    sigmas = get_sigmas(self.noise_scheduler,
-                                        inputs_gpu['latents'].device,
-                                        inputs_gpu['timesteps'],
-                                        n_dim=inputs_gpu['latents'].ndim,
-                                        dtype=inputs_gpu['latents'].dtype)
-                    model_pred = inputs_gpu['latents'] - model_pred * sigmas
-                    target_adjusted = inputs_gpu['target']
-                else:
-                    target_adjusted = inputs_gpu['target']
-
-                loss = torch.mean((model_pred - target_adjusted)**2)
-
-            # Cleanup and return
-            loss_cpu = loss.cpu()
-            del inputs_gpu, model_pred, target_adjusted
-            if 'sigmas' in locals(): del sigmas
-            torch.cuda.empty_cache()
-            return loss_cpu.to(self.fastvideo_args.device)
-
-        try:
-            # Get analytical gradients
-            transformer.zero_grad()
-            analytical_loss = compute_loss()
-            analytical_loss.backward()
-
-            # Check gradients for selected parameters
-            absolute_errors = []
-            param_count = 0
-            rank = int(os.environ.get("RANK", 0))
-
-            sp_group = get_sp_group()
-            for name, param in transformer.named_parameters():
-                sp_group.barrier()
-                # skip scale_shift_table because it is not sharded
-                if 'scale_shift_table' in name:
-                    continue
-                if isinstance(param.grad, torch.distributed.tensor.DTensor):
-                    l = param.grad.full_tensor()
-                    distributed = True
-                else:
-                    l = param.grad
-                    distributed = False
-                    continue
-                # logger.info(f"rank: {rank}, name: {name}, param: {param.shape}, grad: {param.grad.shape}, distributed: {distributed}", local_main_process_only=False)
-                # logger.info(f"rank: {rank}, name: {name}, type of param: {type(param)}, type of grad: {type(param.grad)}", local_main_process_only=False)
-                # logger.info(f"rank: {rank}, name: {name}, param: {param}, grad: {param.grad}", local_main_process_only=False)
-                if not (param.requires_grad and param.grad is not None
-                        and param_count < max_params_to_check
-                        and l.abs().max() > 5e-4):
-                    continue
-                if not distributed:
-                    if rank != 0:
-                        continue
-
-                # Get local parameter and gradient tensors
-                local_param = param._local_tensor if hasattr(
-                    param, '_local_tensor') else param
-                local_grad = param.grad._local_tensor if hasattr(
-                    param.grad, '_local_tensor') else param.grad
-                # logger.info(f"rank: {rank}, local_param: {local_param.shape}, local_grad: {local_grad.shape}", local_main_process_only=False)
-
-                # Find first significant gradient element
-                flat_param = local_param.data.view(-1)
-                flat_grad = local_grad.view(-1)
-                # logger.info(f"rank: {rank}, flat_param: {flat_param.shape}, flat_grad: {flat_grad.shape}", local_main_process_only=False)
-                check_idx = next((i for i in range(min(10, flat_param.numel()))
-                                  if abs(flat_grad[i]) > 1e-4), 0)
-                # logger.info(f"rank: {rank}, check_idx: {check_idx}", local_main_process_only=False)
-
-                # Store original values
-                orig_value = flat_param[check_idx].item()
-                analytical_grad = flat_grad[check_idx].item()
-
-                # Compute numerical gradient
-                for delta in [eps, -eps]:
-                    with torch.no_grad():
-                        # only have a single rank modify the parameter
-                        if rank == 0:
-                            flat_param[check_idx] = orig_value + delta
-                        loss = compute_loss()
-                        if delta > 0: loss_plus = loss.item()
-                        else: loss_minus = loss.item()
-
-                # Restore parameter and compute error
-                with torch.no_grad():
-                    flat_param[check_idx] = orig_value
-
-                numerical_grad = (loss_plus - loss_minus) / (2 * eps)
-                abs_error = abs(analytical_grad - numerical_grad)
-                rel_error = abs_error / max(abs(analytical_grad),
-                                            abs(numerical_grad), 1e-3)
-                absolute_errors.append(abs_error)
-
-                if self.rank <= 0:
-                    logger.info(
-                        f"{name}[{check_idx}]: analytical={analytical_grad:.6f}, "
-                        f"numerical={numerical_grad:.6f}, abs_error={abs_error:.2e}, rel_error={rel_error:.2%}"
-                    )
-
-                # param_count += 1
-
-            # Compute and log statistics
-            if self.rank <= 0:
-                if absolute_errors:
-                    min_err, max_err, mean_err = min(absolute_errors), max(
-                        absolute_errors
-                    ), sum(absolute_errors) / len(absolute_errors)
-                    logger.info(
-                        f"Gradient check stats: min={min_err:.2e}, max={max_err:.2e}, mean={mean_err:.2e}"
-                    )
-
-                    wandb.log({
-                        "grad_check/min_abs_error":
-                        min_err,
-                        "grad_check/max_abs_error":
-                        max_err,
-                        "grad_check/mean_abs_error":
-                        mean_err,
-                        "grad_check/analytical_loss":
-                        analytical_loss.item(),
-                    })
-                return max_err
-
-            return float('inf')
-
-        except Exception as e:
-            logger.error(f"Gradient check failed: {e}")
-            traceback.print_exc()
-            return float('inf')
-
-    def setup_gradient_check(self, args, loader_iter, noise_scheduler,
-                             noise_random_generator):
-        """
-        Setup and perform gradient check on a fresh batch.
-        Args:
-            args: Training arguments
-            loader_iter: Data loader iterator
-            noise_scheduler: Noise scheduler for diffusion
-            noise_random_generator: Random number generator for noise
-        Returns:
-            float or None: Maximum gradient error or None if check is disabled/fails
-        """
-        if not ENABLE_GRADIENT_CHECK:
-            return None
-
-        try:
-            # Get a fresh batch and process it exactly like train_one_step
-            check_latents, check_encoder_hidden_states, check_encoder_attention_mask, check_infos = next(
-                loader_iter)
-
-            # Process exactly like in train_one_step but use GRADIENT_CHECK_DTYPE
-            check_latents = check_latents.to(self.fastvideo_args.device,
-                                             dtype=GRADIENT_CHECK_DTYPE)
-            check_encoder_hidden_states = check_encoder_hidden_states.to(
-                self.fastvideo_args.device, dtype=GRADIENT_CHECK_DTYPE)
-            check_latents = normalize_dit_input("wan", check_latents)
-            batch_size = check_latents.shape[0]
-            check_noise = torch.randn_like(check_latents)
-
-            check_u = compute_density_for_timestep_sampling(
-                weighting_scheme=args.weighting_scheme,
-                batch_size=batch_size,
-                generator=noise_random_generator,
-                logit_mean=args.logit_mean,
-                logit_std=args.logit_std,
-                mode_scale=args.mode_scale,
-            )
-            check_indices = (check_u *
-                             noise_scheduler.config.num_train_timesteps).long()
-            check_timesteps = noise_scheduler.timesteps[check_indices].to(
-                device=check_latents.device)
-
-            check_sigmas = get_sigmas(
-                noise_scheduler,
-                check_latents.device,
-                check_timesteps,
-                n_dim=check_latents.ndim,
-                dtype=check_latents.dtype,
-            )
-            check_noisy_model_input = (
-                1.0 - check_sigmas) * check_latents + check_sigmas * check_noise
-
-            # Compute target exactly like train_one_step
-            if args.precondition_outputs:
-                check_target = check_latents
-            else:
-                check_target = check_noise - check_latents
-
-            # Perform gradient check with the exact same inputs as training
-            max_grad_error = self.gradient_check_parameters(
-                transformer=self.transformer,
-                latents=
-                check_noisy_model_input,  # Use noisy input like in training
-                encoder_hidden_states=check_encoder_hidden_states,
-                encoder_attention_mask=check_encoder_attention_mask,
-                timesteps=check_timesteps,
-                target=check_target,
-                max_params_to_check=100  # Check more parameters
-            )
-
-            if max_grad_error > 5e-2:
-                logger.error(
-                    f"❌ Large gradient error detected: {max_grad_error:.2e}")
-            else:
-                logger.info(
-                    f"✅ Gradient check passed: max error {max_grad_error:.2e}")
-
-            return max_grad_error
-
-        except Exception as e:
-            logger.error(f"Gradient check setup failed: {e}")
-            traceback.print_exc()
-            return None
-
-
-class WanTrainingPipeline(TrainingPipeline):
-    """
-    A training pipeline for Wan.
-    """
-    _required_config_modules = ["scheduler", "transformer"]
-
-    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
-        pass
-
-    def create_training_stages(self, fastvideo_args: FastVideoArgs):
-        pass
-
-    def initialize_validation_pipeline(self, fastvideo_args: FastVideoArgs):
-        self.latents = None
-        logger.info("Initializing validation pipeline...")
-        args_copy = deepcopy(fastvideo_args)
-
-        args_copy.inference_mode = True
-        args_copy.vae_config.load_encoder = False
-        # TODO(will): clean this up
-        args_copy.precision = "bf16"
-        validation_pipeline = WanValidationPipeline.from_pretrained(
-            args.model_path, args=args_copy)
-
-        self.validation_pipeline = validation_pipeline
-
-    def train_one_step(
-        self,
-        transformer,
-        model_type,
-        optimizer,
-        lr_scheduler,
-        loader_iter,
-        noise_scheduler,
-        noise_random_generator,
-        gradient_accumulation_steps,
-        sp_size,
-        precondition_outputs,
-        max_grad_norm,
-        weighting_scheme,
-        logit_mean,
-        logit_std,
-        mode_scale,
-    ):
-        self.modules["transformer"].requires_grad_(True)
-        self.modules["transformer"].train()
-
-        total_loss = 0.0
-        optimizer.zero_grad()
-        for _ in range(gradient_accumulation_steps):
-            logger.info(f"Rank {self.rank}: Training step {_}", local_main_process_only=False)
-            if self.latents is None:
-                (
-                    self.latents,
-                    self.encoder_hidden_states,
-                    self.encoder_attention_mask,
-                    self.infos,
-                ) = next(loader_iter)
-            latents = self.latents
-            encoder_hidden_states = self.encoder_hidden_states
-            encoder_attention_mask = self.encoder_attention_mask
-            infos = self.infos
-            logger.info(f"Rank {self.rank}: Training step {_} loaded data", local_main_process_only=False)
-            latents = latents.to(self.fastvideo_args.device,
-                                 dtype=torch.bfloat16)
-            encoder_hidden_states = encoder_hidden_states.to(
-                self.fastvideo_args.device, dtype=torch.bfloat16)
-            latents = normalize_dit_input(model_type, latents)
-            batch_size = latents.shape[0]
-            noise = torch.randn_like(latents)
-            u = compute_density_for_timestep_sampling(
-                weighting_scheme=weighting_scheme,
-                batch_size=batch_size,
-                generator=noise_random_generator,
-                logit_mean=logit_mean,
-                logit_std=logit_std,
-                mode_scale=mode_scale,
-            )
-            indices = (u * noise_scheduler.config.num_train_timesteps).long()
-            timesteps = noise_scheduler.timesteps[indices].to(
-                device=latents.device)
-            if sp_size > 1:
-                # Make sure that the timesteps are the same across all sp processes.
-                sp_group = get_sp_group()
-                sp_group.broadcast(timesteps, src=0)
-            sigmas = get_sigmas(
-                noise_scheduler,
-                latents.device,
-                timesteps,
-                n_dim=latents.ndim,
-                dtype=latents.dtype,
-            )
-            noisy_model_input = (1.0 - sigmas) * latents + sigmas * noise
-            with torch.autocast("cuda", dtype=torch.bfloat16):
-                input_kwargs = {
-                    "hidden_states": noisy_model_input,
-                    "encoder_hidden_states": encoder_hidden_states,
-                    "timestep": timesteps,
-                    "encoder_attention_mask": encoder_attention_mask,  # B, L
-                    "return_dict": False,
-                }
-                if 'hunyuan' in model_type:
-                    input_kwargs["guidance"] = torch.tensor(
-                        [1000.0],
-                        device=noisy_model_input.device,
-                        dtype=torch.bfloat16)
-                with set_forward_context(current_timestep=timesteps,
-                                         attn_metadata=None):
-                    model_pred = transformer(**input_kwargs)[0]
-
-                if precondition_outputs:
-                    model_pred = noisy_model_input - model_pred * sigmas
-                if precondition_outputs:
-                    target = latents
-                else:
-                    target = noise - latents
-
-                loss = (torch.mean((model_pred.float() - target.float())**2) /
-                        gradient_accumulation_steps)
-
-            loss.backward()
-
-            avg_loss = loss.detach().clone()
-            sp_group = get_sp_group()
-            sp_group.all_reduce(avg_loss, op=torch.distributed.ReduceOp.AVG)
-            total_loss += avg_loss.item()
-
-        model_parts = [self.transformer]
-        grad_norm = _clip_grad_norm_while_handling_failing_dtensor_cases(
-            [p for m in model_parts for p in m.parameters()],
-            max_grad_norm,
-            foreach=None,
-        )
-
-        optimizer.step()
-        print('device after optimizer step',
-              next(transformer.named_parameters())[1].device)
-        lr_scheduler.step()
-        print('device after scheduler step',
-              next(transformer.named_parameters())[1].device)
-        return total_loss, grad_norm.item()
-
-    def forward(
-        self,
-        batch: ForwardBatch,
-        fastvideo_args: FastVideoArgs,
-    ):
-        args = fastvideo_args
-        self.fastvideo_args = args
-        train_dataloader = self.train_dataloader
-        init_steps = self.init_steps
-        lr_scheduler = self.lr_scheduler
-        optimizer = self.optimizer
-        noise_scheduler = self.noise_scheduler
-        noise_random_generator = None
-
-        from diffusers import FlowMatchEulerDiscreteScheduler
-        noise_scheduler = FlowMatchEulerDiscreteScheduler()
-
-        # Train!
-        total_batch_size = (self.world_size * args.gradient_accumulation_steps /
-                            args.sp_size * args.train_sp_batch_size)
-        logger.info("***** Running training *****")
-        # logger.info(f"  Num examples = {len(train_dataset)}")
-        # logger.info(f"  Dataloader size = {len(train_dataloader)}")
-        # logger.info(f"  Num Epochs = {args.num_train_epochs}")
-        logger.info(f"  Resume training from step {init_steps}")
-        logger.info(
-            f"  Instantaneous batch size per device = {args.train_batch_size}")
-        logger.info(
-            f"  Total train batch size (w. data & sequence parallel, accumulation) = {total_batch_size}"
-        )
-        logger.info(
-            f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}"
-        )
-        logger.info(f"  Total optimization steps = {args.max_train_steps}")
-        logger.info(
-            f"  Total training parameters per FSDP shard = {sum(p.numel() for p in self.transformer.parameters() if p.requires_grad) / 1e9} B"
-        )
-        # print dtype
-        logger.info(
-            f"  Master weight dtype: {self.transformer.parameters().__next__().dtype}"
-        )
-
-        # Potentially load in the weights and states from a previous save
-        if args.resume_from_checkpoint:
-            assert NotImplementedError(
-                "resume_from_checkpoint is not supported now.")
-            # TODO
-
-        progress_bar = tqdm(
-            range(0, args.max_train_steps),
-            initial=init_steps,
-            desc="Steps",
-            # Only show the progress bar once on each machine.
-            disable=self.local_rank > 0,
-        )
-
-        loader_iter = iter(train_dataloader)
-
-        step_times = deque(maxlen=100)
-
-        # todo future
-        for i in range(init_steps):
-            next(loader_iter)
-        # get gpu memory usage
-        gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2
-        logger.info(
-            f"GPU memory usage before train_one_step: {gpu_memory_usage} MB")
-
-        for step in range(init_steps + 1, args.max_train_steps + 1):
-            start_time = time.perf_counter()
-
-            loss, grad_norm = self.train_one_step(
-                self.transformer,
-                # args.model_type,
-                "wan",
-                optimizer,
-                lr_scheduler,
-                loader_iter,
-                noise_scheduler,
-                noise_random_generator,
-                args.gradient_accumulation_steps,
-                args.sp_size,
-                args.precondition_outputs,
-                args.max_grad_norm,
-                args.weighting_scheme,
-                args.logit_mean,
-                args.logit_std,
-                args.mode_scale,
-            )
-            gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2
-            logger.info(
-                f"GPU memory usage after train_one_step: {gpu_memory_usage} MB")
-
-            step_time = time.perf_counter() - start_time
-            step_times.append(step_time)
-            avg_step_time = sum(step_times) / len(step_times)
-
-            # Manual gradient checking - only at first step
-            if step == 1 and ENABLE_GRADIENT_CHECK:
-                logger.info(f"Performing gradient check at step {step}")
-                self.setup_gradient_check(args, loader_iter, noise_scheduler,
-                                          noise_random_generator)
-
-            progress_bar.set_postfix({
-                "loss": f"{loss:.4f}",
-                "step_time": f"{step_time:.2f}s",
-                "grad_norm": grad_norm,
-            })
-            progress_bar.update(1)
-            if self.rank <= 0:
-                wandb.log(
-                    {
-                        "train_loss": loss,
-                        "learning_rate": lr_scheduler.get_last_lr()[0],
-                        "step_time": step_time,
-                        "avg_step_time": avg_step_time,
-                        "grad_norm": grad_norm,
-                    },
-                    step=step,
-                )
-            if step % args.checkpointing_steps == 0:
-                # Your existing checkpoint saving code
-                save_checkpoint(self.transformer, self.rank, args.output_dir,
-                                step)
-                self.transformer.train()
-                self.sp_group.barrier()
-            if args.log_validation and step % args.validation_steps == 0:
-                self.log_validation(self.transformer, args, step)
-
-        save_checkpoint(self.transformer, self.rank, args.output_dir,
-                        args.max_train_steps)
-
-        if get_sp_group():
-            cleanup_dist_env_and_memory()
-
-
-def main(args):
-    logger.info("Starting training pipeline...")
-
-    pipeline = WanTrainingPipeline.from_pretrained(
-        args.pretrained_model_name_or_path, args=args)
-    args = pipeline.fastvideo_args
-    pipeline.forward(None, args)
-    logger.info("Training pipeline done")
-
-
-if __name__ == "__main__":
-    argv = sys.argv
-    from fastvideo.v1.fastvideo_args import TrainingArgs
-    from fastvideo.v1.utils import FlexibleArgumentParser
-    parser = FlexibleArgumentParser()
-    parser = TrainingArgs.add_cli_args(parser)
-    parser = FastVideoArgs.add_cli_args(parser)
-    args = parser.parse_args()
-    args.use_cpu_offload = False
-    print(args)
-    main(args)

fastvideo/v1/pipelines/wan/wan_latent_pipeline.py

@@ -1,19 +0,0 @@
-from fastvideo.v1.fastvideo_args import FastVideoArgs
-from fastvideo.v1.logger import init_logger
-from fastvideo.v1.pipelines.composed_pipeline_base import ComposedPipelineBase
-from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
-
-logger = init_logger(__name__)
-
-
-class WanLatentPipeline(ComposedPipelineBase):
-    _required_config_modules = ["text_encoder", "tokenizer", "vae"]
-
-    # def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
-
-    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
-        pass
-
-    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs):
-        logger.info("WAN Latent Pipeline forward")
-        pass

fastvideo/v1/pipelines/wan/wan_pipeline.py

@@ -15,7 +15,6 @@ from fastvideo.v1.pipelines.stages import (ConditioningStage, DecodingStage,
                                            TextEncodingStage,
                                            TimestepPreparationStage)
 
-
 logger = init_logger(__name__)
 
 

fastvideo/v1/training/training_pipeline.py

@@ -0,0 +1,515 @@
+import gc
+import os
+import traceback
+from abc import ABC, abstractmethod
+
+import imageio
+import numpy as np
+import torch
+import torchvision
+from diffusers.optimization import get_scheduler
+from einops import rearrange
+from torchdata.stateful_dataloader import StatefulDataLoader
+
+from fastvideo.v1.configs.sample import SamplingParam
+from fastvideo.v1.dataset.parquet_datasets import ParquetVideoTextDataset
+from fastvideo.v1.distributed import get_sp_group
+from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs
+from fastvideo.v1.forward_context import set_forward_context
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.pipelines import ComposedPipelineBase
+from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.v1.training.training_utils import (
+    compute_density_for_timestep_sampling, get_sigmas, normalize_dit_input)
+
+import wandb  # isort: skip
+
+logger = init_logger(__name__)
+
+# Note: if checking with float32, cannot use flash-attn.
+GRADIENT_CHECK_DTYPE = torch.bfloat16
+
+
+class TrainingPipeline(ComposedPipelineBase, ABC):
+    """
+    A pipeline for training a model. All training pipelines should inherit from this class.
+    All reusable components and code should be implemented in this class.
+    """
+    _required_config_modules = ["scheduler", "transformer"]
+    validation_pipeline: ComposedPipelineBase
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        raise RuntimeError(
+            "create_pipeline_stages should not be called for training pipeline")
+
+    def initialize_training_pipeline(self, training_args: TrainingArgs):
+        logger.info("Initializing training pipeline...")
+        self.device = training_args.device
+        self.sp_group = get_sp_group()
+        self.world_size = self.sp_group.world_size
+        self.rank = self.sp_group.rank
+        self.local_rank = self.sp_group.local_rank
+        self.transformer = self.get_module("transformer")
+        assert self.transformer is not None
+
+        self.transformer.requires_grad_(True)
+        self.transformer.train()
+
+        noise_scheduler = self.modules["scheduler"]
+        params_to_optimize = self.transformer.parameters()
+        params_to_optimize = list(
+            filter(lambda p: p.requires_grad, params_to_optimize))
+
+        self.optimizer = torch.optim.AdamW(
+            params_to_optimize,
+            lr=training_args.learning_rate,
+            betas=(0.9, 0.999),
+            weight_decay=training_args.weight_decay,
+            eps=1e-8,
+        )
+
+        self.init_steps = 0
+        logger.info("optimizer: %s", self.optimizer)
+
+        self.lr_scheduler = get_scheduler(
+            training_args.lr_scheduler,
+            optimizer=self.optimizer,
+            num_warmup_steps=training_args.lr_warmup_steps * self.world_size,
+            num_training_steps=training_args.max_train_steps * self.world_size,
+            num_cycles=training_args.lr_num_cycles,
+            power=training_args.lr_power,
+            last_epoch=self.init_steps - 1,
+        )
+
+        self.train_dataset = ParquetVideoTextDataset(
+            training_args.data_path,
+            batch_size=training_args.train_batch_size,
+            rank=self.rank,
+            world_size=self.world_size,
+            cfg_rate=training_args.cfg,
+            num_latent_t=training_args.num_latent_t)
+
+        self.train_dataloader = StatefulDataLoader(
+            self.train_dataset,
+            batch_size=training_args.train_batch_size,
+            num_workers=training_args.
+            dataloader_num_workers,  # Reduce number of workers to avoid memory issues
+            prefetch_factor=2,
+            shuffle=False,
+            pin_memory=True,
+            pin_memory_device=f"cuda:{torch.cuda.current_device()}",
+            drop_last=True)
+
+        self.noise_scheduler = noise_scheduler
+
+        if self.rank <= 0:
+            project = training_args.tracker_project_name or "fastvideo"
+            wandb.init(project=project, config=training_args)
+
+    @abstractmethod
+    def initialize_validation_pipeline(self, training_args: TrainingArgs):
+        raise NotImplementedError(
+            "Training pipelines must implement this method")
+
+    @abstractmethod
+    def train_one_step(self, transformer, model_type, optimizer, lr_scheduler,
+                       loader, noise_scheduler, noise_random_generator,
+                       gradient_accumulation_steps, sp_size,
+                       precondition_outputs, max_grad_norm, weighting_scheme,
+                       logit_mean, logit_std, mode_scale):
+        """
+        Train one step of the model.
+        """
+        raise NotImplementedError(
+            "Training pipeline must implement this method")
+
+    def log_validation(self, transformer, training_args, global_step) -> None:
+        assert training_args is not None
+        training_args.inference_mode = True
+        training_args.use_cpu_offload = False
+        if not training_args.log_validation:
+            return
+        if self.validation_pipeline is None:
+            raise ValueError("Validation pipeline is not set")
+
+        # Create sampling parameters if not provided
+        sampling_param = SamplingParam.from_pretrained(training_args.model_path)
+
+        # Prepare validation prompts
+        logger.info('fastvideo_args.validation_prompt_dir: %s',
+                    training_args.validation_prompt_dir)
+        validation_dataset = ParquetVideoTextDataset(
+            training_args.validation_prompt_dir,
+            batch_size=1,
+            rank=0,
+            world_size=1,
+            cfg_rate=0,
+            num_latent_t=training_args.num_latent_t)
+
+        validation_dataloader = StatefulDataLoader(
+            validation_dataset,
+            batch_size=1,
+            num_workers=1,  # Reduce number of workers to avoid memory issues
+            prefetch_factor=2,
+            shuffle=False,
+            pin_memory=True,
+            drop_last=False)
+
+        transformer.requires_grad_(False)
+        for p in transformer.parameters():
+            p.requires_grad = False
+        transformer.eval()
+
+        # Add the transformer to the validation pipeline
+        self.validation_pipeline.add_module("transformer", transformer)
+        self.validation_pipeline.latent_preparation_stage.transformer = transformer  # type: ignore[attr-defined]
+        self.validation_pipeline.denoising_stage.transformer = transformer  # type: ignore[attr-defined]
+
+        # Process each validation prompt
+        videos = []
+        captions = []
+        for _, embeddings, masks, infos in validation_dataloader:
+            logger.info("infos: %s", infos)
+            caption = infos['caption']
+            captions.append(caption)
+            prompt_embeds = embeddings.to(training_args.device)
+            prompt_attention_mask = masks.to(training_args.device)
+
+            # Calculate sizes
+            latents_size = [(sampling_param.num_frames - 1) // 4 + 1,
+                            sampling_param.height // 8,
+                            sampling_param.width // 8]
+            n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
+
+            temporal_compression_factor = training_args.vae_config.arch_config.temporal_compression_ratio
+            num_frames = (training_args.num_latent_t -
+                          1) * temporal_compression_factor + 1
+            logger.info(
+                "validation num_frames: %s, temporal_compression_factor: %s, num_latent_t: %s",
+                num_frames, temporal_compression_factor,
+                training_args.num_latent_t)
+            # Prepare batch for validation
+            batch = ForwardBatch(
+                data_type="video",
+                latents=None,
+                # seed=sampling_param.seed,
+                prompt_embeds=[prompt_embeds],
+                prompt_attention_mask=[prompt_attention_mask],
+                # make sure we use the same height, width, and num_frames as the training pipeline
+                height=training_args.num_height,
+                width=training_args.num_width,
+                num_frames=num_frames,
+                # num_inference_steps=fastvideo_args.validation_sampling_steps,
+                num_inference_steps=sampling_param.num_inference_steps,
+                # guidance_scale=fastvideo_args.validation_guidance_scale,
+                guidance_scale=1,
+                n_tokens=n_tokens,
+                do_classifier_free_guidance=False,
+                eta=0.0,
+                extra={},
+            )
+
+            # Run validation inference
+            with torch.inference_mode(), torch.autocast("cuda",
+                                                        dtype=torch.bfloat16):
+                output_batch = self.validation_pipeline.forward(
+                    batch, training_args)
+                samples = output_batch.output
+
+            # Process outputs
+            video = rearrange(samples, "b c t h w -> t b c h w")
+            frames = []
+            for x in video:
+                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))
+            videos.append(frames)
+
+        # Log validation results
+        rank = int(os.environ.get("RANK", 0))
+
+        if rank == 0:
+            video_filenames = []
+            video_captions = []
+            for i, video in enumerate(videos):
+                caption = captions[i]
+                os.makedirs(training_args.output_dir, exist_ok=True)
+                filename = os.path.join(
+                    training_args.output_dir,
+                    f"validation_step_{global_step}_video_{i}.mp4")
+                imageio.mimsave(filename, video, fps=sampling_param.fps)
+                video_filenames.append(filename)
+                video_captions.append(
+                    caption)  # Store the caption for each video
+
+            logs = {
+                "validation_videos": [
+                    wandb.Video(filename,
+                                caption=caption) for filename, caption in zip(
+                                    video_filenames, video_captions)
+                ]
+            }
+            wandb.log(logs, step=global_step)
+
+        # Re-enable gradients for training
+        transformer.requires_grad_(True)
+        transformer.train()
+
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    def gradient_check_parameters(self,
+                                  transformer,
+                                  latents,
+                                  encoder_hidden_states,
+                                  encoder_attention_mask,
+                                  timesteps,
+                                  target,
+                                  eps=5e-2,
+                                  max_params_to_check=2000) -> float:
+        """
+        Verify gradients using finite differences for FSDP models with GRADIENT_CHECK_DTYPE.
+        Uses standard tolerances for GRADIENT_CHECK_DTYPE precision.
+        """
+        assert self.training_args is not None
+        # Move all inputs to CPU and clear GPU memory
+        inputs_cpu = {
+            'latents': latents.cpu(),
+            'encoder_hidden_states': encoder_hidden_states.cpu(),
+            'encoder_attention_mask': encoder_attention_mask.cpu(),
+            'timesteps': timesteps.cpu(),
+            'target': target.cpu()
+        }
+        del latents, encoder_hidden_states, encoder_attention_mask, timesteps, target
+        torch.cuda.empty_cache()
+
+        def compute_loss() -> torch.Tensor:
+            assert self.training_args is not None
+            # Move inputs to GPU, compute loss, cleanup
+            inputs_gpu = {
+                k:
+                v.to(self.training_args.device,
+                     dtype=GRADIENT_CHECK_DTYPE
+                     if k != 'encoder_attention_mask' else None)
+                for k, v in inputs_cpu.items()
+            }
+
+            # Use GRADIENT_CHECK_DTYPE for more accurate gradient checking
+            # with torch.autocast(enabled=False, device_type="cuda"):
+            with torch.autocast("cuda", dtype=GRADIENT_CHECK_DTYPE):
+                with set_forward_context(
+                        current_timestep=inputs_gpu['timesteps'],
+                        attn_metadata=None):
+                    model_pred = transformer(
+                        hidden_states=inputs_gpu['latents'],
+                        encoder_hidden_states=inputs_gpu[
+                            'encoder_hidden_states'],
+                        timestep=inputs_gpu['timesteps'],
+                        encoder_attention_mask=inputs_gpu[
+                            'encoder_attention_mask'],
+                        return_dict=False)[0]
+
+                if self.training_args.precondition_outputs:
+                    sigmas = get_sigmas(self.noise_scheduler,
+                                        inputs_gpu['latents'].device,
+                                        inputs_gpu['timesteps'],
+                                        n_dim=inputs_gpu['latents'].ndim,
+                                        dtype=inputs_gpu['latents'].dtype)
+                    model_pred = inputs_gpu['latents'] - model_pred * sigmas
+                    target_adjusted = inputs_gpu['target']
+                else:
+                    target_adjusted = inputs_gpu['target']
+
+                loss = torch.mean((model_pred - target_adjusted)**2)
+
+            # Cleanup and return
+            loss_cpu = loss.cpu()
+            del inputs_gpu, model_pred, target_adjusted
+            if 'sigmas' in locals():
+                del sigmas
+            torch.cuda.empty_cache()
+            return loss_cpu.to(self.training_args.device)
+
+        try:
+            # Get analytical gradients
+            transformer.zero_grad()
+            analytical_loss = compute_loss()
+            analytical_loss.backward()
+
+            # Check gradients for selected parameters
+            absolute_errors: list[float] = []
+            param_count = 0
+
+            rank = int(os.environ.get("RANK", 0))
+            sp_group = get_sp_group()
+            for name, param in transformer.named_parameters():
+                sp_group.barrier()
+                # skip scale_shift_table because it is not sharded
+                if 'scale_shift_table' in name:
+                    continue
+                if isinstance(param.grad, torch.distributed.tensor.DTensor):
+                    full_grad = param.grad.full_tensor()
+                    distributed = True
+                else:
+                    full_grad = param.grad
+                    distributed = False
+                    continue
+                if not (param.requires_grad and param.grad is not None
+                        and param_count < max_params_to_check
+                        and full_grad.abs().max() > 5e-4):
+                    continue
+                if not distributed and rank != 0:
+                    continue
+
+                # Get local parameter and gradient tensors
+                local_param = param._local_tensor if hasattr(
+                    param, '_local_tensor') else param
+                local_grad = param.grad._local_tensor if hasattr(
+                    param.grad, '_local_tensor') else param.grad
+
+                # Find first significant gradient element
+                flat_param = local_param.data.view(-1)
+                flat_grad = local_grad.view(-1)
+                check_idx = next((i for i in range(min(10, flat_param.numel()))
+                                  if abs(flat_grad[i]) > 1e-4), 0)
+
+                # Store original values
+                orig_value = flat_param[check_idx].item()
+                analytical_grad = flat_grad[check_idx].item()
+
+                # Compute numerical gradient
+                for delta in [eps, -eps]:
+                    with torch.no_grad():
+                        # only have a single rank modify the parameter
+                        # because we are using FSDP
+                        if rank <= 0:
+                            flat_param[check_idx] = orig_value + delta
+                        loss = compute_loss()
+                        if delta > 0:
+                            loss_plus = loss.item()
+                        else:
+                            loss_minus = loss.item()
+
+                # Restore parameter and compute error
+                with torch.no_grad():
+                    flat_param[check_idx] = orig_value
+
+                numerical_grad = (loss_plus - loss_minus) / (2 * eps)
+                abs_error = abs(analytical_grad - numerical_grad)
+                rel_error = abs_error / max(abs(analytical_grad),
+                                            abs(numerical_grad), 1e-3)
+                absolute_errors.append(abs_error)
+
+                if self.rank <= 0:
+                    logger.info(
+                        "%s[%s]: analytical=%.5f, numerical=%.5f, abs_error=%.2e, rel_error=%.2f%%",
+                        name, check_idx, analytical_grad, numerical_grad,
+                        abs_error, rel_error * 100)
+
+                # param_count += 1
+
+            # Compute and log statistics
+            if rank <= 0 and absolute_errors:
+                min_err, max_err, mean_err = min(absolute_errors), max(
+                    absolute_errors
+                ), sum(absolute_errors) / len(absolute_errors)
+                logger.info("Gradient check stats: min=%s, max=%s, mean=%s",
+                            min_err, max_err, mean_err)
+
+                wandb.log({
+                    "grad_check/min_abs_error": min_err,
+                    "grad_check/max_abs_error": max_err,
+                    "grad_check/mean_abs_error": mean_err,
+                    "grad_check/analytical_loss": analytical_loss.item(),
+                })
+                return max_err
+
+            return float('inf')
+
+        except Exception as e:
+            logger.error("Gradient check failed: %s", e)
+            traceback.print_exc()
+            return float('inf')
+
+    def setup_gradient_check(self, args, loader_iter, noise_scheduler,
+                             noise_random_generator) -> float | None:
+        """
+        Setup and perform gradient check on a fresh batch.
+        Args:
+            args: Training arguments
+            loader_iter: Data loader iterator
+            noise_scheduler: Noise scheduler for diffusion
+            noise_random_generator: Random number generator for noise
+        Returns:
+            float or None: Maximum gradient error or None if check is disabled/fails
+        """
+        assert self.training_args is not None
+
+        try:
+            # Get a fresh batch and process it exactly like train_one_step
+            check_latents, check_encoder_hidden_states, check_encoder_attention_mask, check_infos = next(
+                loader_iter)
+
+            # Process exactly like in train_one_step but use GRADIENT_CHECK_DTYPE
+            check_latents = check_latents.to(self.training_args.device,
+                                             dtype=GRADIENT_CHECK_DTYPE)
+            check_encoder_hidden_states = check_encoder_hidden_states.to(
+                self.training_args.device, dtype=GRADIENT_CHECK_DTYPE)
+            check_latents = normalize_dit_input("wan", check_latents)
+            batch_size = check_latents.shape[0]
+            check_noise = torch.randn_like(check_latents)
+
+            check_u = compute_density_for_timestep_sampling(
+                weighting_scheme=args.weighting_scheme,
+                batch_size=batch_size,
+                generator=noise_random_generator,
+                logit_mean=args.logit_mean,
+                logit_std=args.logit_std,
+                mode_scale=args.mode_scale,
+            )
+            check_indices = (check_u *
+                             noise_scheduler.config.num_train_timesteps).long()
+            check_timesteps = noise_scheduler.timesteps[check_indices].to(
+                device=check_latents.device)
+
+            check_sigmas = get_sigmas(
+                noise_scheduler,
+                check_latents.device,
+                check_timesteps,
+                n_dim=check_latents.ndim,
+                dtype=check_latents.dtype,
+            )
+            check_noisy_model_input = (
+                1.0 - check_sigmas) * check_latents + check_sigmas * check_noise
+
+            # Compute target exactly like train_one_step
+            if args.precondition_outputs:
+                check_target = check_latents
+            else:
+                check_target = check_noise - check_latents
+
+            # Perform gradient check with the exact same inputs as training
+            max_grad_error = self.gradient_check_parameters(
+                transformer=self.transformer,
+                latents=
+                check_noisy_model_input,  # Use noisy input like in training
+                encoder_hidden_states=check_encoder_hidden_states,
+                encoder_attention_mask=check_encoder_attention_mask,
+                timesteps=check_timesteps,
+                target=check_target,
+                max_params_to_check=100  # Check more parameters
+            )
+
+            if max_grad_error > 5e-2:
+                logger.error("❌ Large gradient error detected: %s",
+                             max_grad_error)
+            else:
+                logger.info("✅ Gradient check passed: max error %s",
+                            max_grad_error)
+
+            return max_grad_error
+
+        except Exception as e:
+            logger.error("Gradient check setup failed: %s", e)
+            traceback.print_exc()
+            return None

fastvideo/v1/training/training_utils.py

@@ -5,25 +5,24 @@ from typing import List, Optional, Tuple, Union
 
 import torch
 import torch.distributed as dist
-import torch.distributed.tensor
 from torch.distributed.fsdp import FullStateDictConfig
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from torch.distributed.fsdp import StateDictType
 
 from fastvideo.v1.logger import init_logger
 
-_HAS_ERRORED_CLIP_GRAD_NORM_WHILE_HANDLING_FAILING_DTENSOR_CASES = False
-
 logger = init_logger(__name__)
 
+_HAS_ERRORED_CLIP_GRAD_NORM_WHILE_HANDLING_FAILING_DTENSOR_CASES = False
+
 
 def compute_density_for_timestep_sampling(
     weighting_scheme: str,
     batch_size: int,
     generator,
-    logit_mean: float = None,
-    logit_std: float = None,
-    mode_scale: float = None,
+    logit_mean: Optional[float] = None,
+    logit_std: Optional[float] = None,
+    mode_scale: Optional[float] = None,
 ):
     """
     Compute the density for sampling the timesteps when doing SD3 training.
@@ -54,7 +53,7 @@ def get_sigmas(noise_scheduler,
                device,
                timesteps,
                n_dim=4,
-               dtype=torch.float32):
+               dtype=torch.float32) -> torch.Tensor:
     sigmas = noise_scheduler.sigmas.to(device=device, dtype=dtype)
     schedule_timesteps = noise_scheduler.timesteps.to(device)
     timesteps = timesteps.to(device)
@@ -67,7 +66,7 @@ def get_sigmas(noise_scheduler,
     return sigma
 
 
-def save_checkpoint(transformer, rank, output_dir, step):
+def save_checkpoint(transformer, rank, output_dir, step) -> None:
     # Configure FSDP to save full state dict
     FSDP.set_state_dict_type(
         transformer,
@@ -95,7 +94,25 @@ def save_checkpoint(transformer, rank, output_dir, step):
         logger.info("--> checkpoint saved at step %s to %s", step, weight_path)
 
 
-def _clip_grad_norm_while_handling_failing_dtensor_cases(
+def normalize_dit_input(model_type, latents, args=None) -> torch.Tensor:
+    if model_type == "hunyuan_hf" or model_type == "hunyuan":
+        return latents * 0.476986
+    elif model_type == "wan":
+        from fastvideo.v1.configs.models.vaes.wanvae import WanVAEConfig
+        vae_config = WanVAEConfig()
+        latents_mean = torch.tensor(vae_config.arch_config.latents_mean)
+        latents_std = 1.0 / torch.tensor(vae_config.arch_config.latents_std)
+
+        latents_mean = latents_mean.view(1, -1, 1, 1,
+                                         1).to(device=latents.device)
+        latents_std = latents_std.view(1, -1, 1, 1, 1).to(device=latents.device)
+        latents = ((latents.float() - latents_mean) * latents_std).to(latents)
+        return latents
+    else:
+        raise NotImplementedError(f"model_type {model_type} not supported")
+
+
+def clip_grad_norm_while_handling_failing_dtensor_cases(
     parameters: Union[torch.Tensor, List[torch.Tensor]],
     max_norm: float,
     norm_type: float = 2.0,
@@ -119,8 +136,8 @@ def _clip_grad_norm_while_handling_failing_dtensor_cases(
                 )
         except Exception as e:
             logger.warning(
-                f"An error occurred while clipping gradients: {e}. Gradient clipping will be skipped and gradient "
-                f"norm will not be logged.")
+                "An error occurred while clipping gradients: %s. Gradient clipping will be skipped and gradient "
+                "norm will not be logged.", e)
             _HAS_ERRORED_CLIP_GRAD_NORM_WHILE_HANDLING_FAILING_DTENSOR_CASES = True
     return None
 
@@ -180,6 +197,7 @@ def clip_grad_norm_(
         total_norm = total_norm.full_tensor()
 
     if pp_mesh is not None:
+        raise NotImplementedError("Pipeline parallel is not supported")
         if math.isinf(norm_type):
             dist.all_reduce(total_norm,
                             op=dist.ReduceOp.MAX,
@@ -239,10 +257,7 @@ def _get_total_norm(
     error_if_nonfinite: bool = False,
     foreach: Optional[bool] = None,
 ) -> torch.Tensor:
-    if isinstance(tensors, torch.Tensor):
-        tensors = [tensors]
-    else:
-        tensors = list(tensors)
+    tensors = [tensors] if isinstance(tensors, torch.Tensor) else list(tensors)
     norm_type = float(norm_type)
     if len(tensors) == 0:
         return torch.tensor(0.0)
@@ -295,8 +310,8 @@ def _group_tensors_by_device_and_dtype(
     with_indices: bool = False,
 ) -> dict[tuple[torch.device, torch.dtype], tuple[
         List[List[Optional[torch.Tensor]]], List[int]]]:
-    return torch._C._group_tensors_by_device_and_dtype(tensorlistlist,
-                                                       with_indices)
+    return torch._C._group_tensors_by_device_and_dtype(  # type: ignore[no-any-return]
+        tensorlistlist, with_indices)
 
 
 def _device_has_foreach_support(device: torch.device) -> bool:

fastvideo/v1/training/wan_training_pipeline.py

@@ -0,0 +1,317 @@
+import sys
+import time
+from collections import deque
+from copy import deepcopy
+
+import torch
+from diffusers import FlowMatchEulerDiscreteScheduler
+from tqdm.auto import tqdm
+
+from fastvideo.v1.distributed import cleanup_dist_env_and_memory, get_sp_group
+from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs
+from fastvideo.v1.forward_context import set_forward_context
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.v1.pipelines.wan.wan_pipeline import WanValidationPipeline
+from fastvideo.v1.training.training_pipeline import TrainingPipeline
+from fastvideo.v1.training.training_utils import (
+    clip_grad_norm_while_handling_failing_dtensor_cases,
+    compute_density_for_timestep_sampling, get_sigmas, normalize_dit_input,
+    save_checkpoint)
+
+import wandb  # isort: skip
+
+logger = init_logger(__name__)
+
+# Manual gradient checking flag - set to True to enable gradient verification
+ENABLE_GRADIENT_CHECK = False
+
+
+class WanTrainingPipeline(TrainingPipeline):
+    """
+    A training pipeline for Wan.
+    """
+    _required_config_modules = ["scheduler", "transformer"]
+
+    def create_training_stages(self, training_args: TrainingArgs):
+        """
+        May be used in future refactors.
+        """
+        pass
+
+    def initialize_validation_pipeline(self, training_args: TrainingArgs):
+        logger.info("Initializing validation pipeline...")
+        args_copy = deepcopy(training_args)
+
+        args_copy.inference_mode = True
+        args_copy.vae_config.load_encoder = False
+        validation_pipeline = WanValidationPipeline.from_pretrained(
+            args.model_path, args=None, inference_mode=True)
+
+        self.validation_pipeline = validation_pipeline
+
+    def train_one_step(
+        self,
+        transformer,
+        model_type,
+        optimizer,
+        lr_scheduler,
+        loader_iter,
+        noise_scheduler,
+        noise_random_generator,
+        gradient_accumulation_steps,
+        sp_size,
+        precondition_outputs,
+        max_grad_norm,
+        weighting_scheme,
+        logit_mean,
+        logit_std,
+        mode_scale,
+    ) -> tuple[float, float]:
+        assert self.training_args is not None
+        self.modules["transformer"].requires_grad_(True)
+        self.modules["transformer"].train()
+
+        total_loss = 0.0
+        optimizer.zero_grad()
+        for _ in range(gradient_accumulation_steps):
+            (
+                latents,
+                encoder_hidden_states,
+                encoder_attention_mask,
+                infos,
+            ) = next(loader_iter)
+            latents = latents.to(self.training_args.device,
+                                 dtype=torch.bfloat16)
+            encoder_hidden_states = encoder_hidden_states.to(
+                self.training_args.device, dtype=torch.bfloat16)
+            latents = normalize_dit_input(model_type, latents)
+            batch_size = latents.shape[0]
+            noise = torch.randn_like(latents)
+            u = compute_density_for_timestep_sampling(
+                weighting_scheme=weighting_scheme,
+                batch_size=batch_size,
+                generator=noise_random_generator,
+                logit_mean=logit_mean,
+                logit_std=logit_std,
+                mode_scale=mode_scale,
+            )
+            indices = (u * noise_scheduler.config.num_train_timesteps).long()
+            timesteps = noise_scheduler.timesteps[indices].to(
+                device=latents.device)
+            if sp_size > 1:
+                # Make sure that the timesteps are the same across all sp processes.
+                sp_group = get_sp_group()
+                sp_group.broadcast(timesteps, src=0)
+            sigmas = get_sigmas(
+                noise_scheduler,
+                latents.device,
+                timesteps,
+                n_dim=latents.ndim,
+                dtype=latents.dtype,
+            )
+            noisy_model_input = (1.0 - sigmas) * latents + sigmas * noise
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                input_kwargs = {
+                    "hidden_states": noisy_model_input,
+                    "encoder_hidden_states": encoder_hidden_states,
+                    "timestep": timesteps,
+                    "encoder_attention_mask": encoder_attention_mask,  # B, L
+                    "return_dict": False,
+                }
+                if 'hunyuan' in model_type:
+                    input_kwargs["guidance"] = torch.tensor(
+                        [1000.0],
+                        device=noisy_model_input.device,
+                        dtype=torch.bfloat16)
+                with set_forward_context(current_timestep=timesteps,
+                                         attn_metadata=None):
+                    model_pred = transformer(**input_kwargs)[0]
+
+                if precondition_outputs:
+                    model_pred = noisy_model_input - model_pred * sigmas
+                target = latents if precondition_outputs else noise - latents
+
+                loss = (torch.mean((model_pred.float() - target.float())**2) /
+                        gradient_accumulation_steps)
+
+            loss.backward()
+
+            avg_loss = loss.detach().clone()
+            sp_group = get_sp_group()
+            sp_group.all_reduce(avg_loss, op=torch.distributed.ReduceOp.AVG)
+            total_loss += avg_loss.item()
+
+        # TODO(will): perhaps move this into transformer api so that we can do
+        # the following:
+        # grad_norm = transformer.clip_grad_norm_(max_grad_norm)
+        if max_grad_norm is not None:
+            model_parts = [self.transformer]
+            grad_norm = clip_grad_norm_while_handling_failing_dtensor_cases(
+                [p for m in model_parts for p in m.parameters()],
+                max_grad_norm,
+                foreach=None,
+            )
+            grad_norm = grad_norm.item() if grad_norm is not None else 0.0
+        else:
+            grad_norm = 0.0
+
+        optimizer.step()
+        lr_scheduler.step()
+        return total_loss, grad_norm
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ):
+        assert self.training_args is not None
+        noise_random_generator = None
+
+        noise_scheduler = FlowMatchEulerDiscreteScheduler()
+
+        # Train!
+        assert self.training_args.sp_size is not None
+        assert self.training_args.gradient_accumulation_steps is not None
+        total_batch_size = (self.world_size *
+                            self.training_args.gradient_accumulation_steps /
+                            self.training_args.sp_size *
+                            self.training_args.train_sp_batch_size)
+        logger.info("***** Running training *****")
+        # logger.info(f"  Num examples = {len(train_dataset)}")
+        # logger.info(f"  Dataloader size = {len(train_dataloader)}")
+        # logger.info(f"  Num Epochs = {args.num_train_epochs}")
+        logger.info("  Resume training from step %s", self.init_steps)
+        logger.info("  Instantaneous batch size per device = %s",
+                    self.training_args.train_batch_size)
+        logger.info(
+            "  Total train batch size (w. data & sequence parallel, accumulation) = %s",
+            total_batch_size)
+        logger.info("  Gradient Accumulation steps = %s",
+                    self.training_args.gradient_accumulation_steps)
+        logger.info("  Total optimization steps = %s",
+                    self.training_args.max_train_steps)
+        logger.info(
+            "  Total training parameters per FSDP shard = %s B",
+            sum(p.numel()
+                for p in self.transformer.parameters() if p.requires_grad) /
+            1e9)
+        # print dtype
+        logger.info("  Master weight dtype: %s",
+                    self.transformer.parameters().__next__().dtype)
+
+        # Potentially load in the weights and states from a previous save
+        if self.training_args.resume_from_checkpoint:
+            assert NotImplementedError(
+                "resume_from_checkpoint is not supported now.")
+            # TODO
+
+        progress_bar = tqdm(
+            range(0, self.training_args.max_train_steps),
+            initial=self.init_steps,
+            desc="Steps",
+            # Only show the progress bar once on each machine.
+            disable=self.local_rank > 0,
+        )
+
+        loader_iter = iter(self.train_dataloader)
+
+        step_times: deque[float] = deque(maxlen=100)
+
+        # TODO(will): fix this
+        # for i in range(self.init_steps):
+        #     next(loader_iter)
+        # get gpu memory usage
+        gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2
+        logger.info("GPU memory usage before train_one_step: %s MB",
+                    gpu_memory_usage)
+
+        for step in range(self.init_steps + 1, args.max_train_steps + 1):
+            start_time = time.perf_counter()
+
+            loss, grad_norm = self.train_one_step(
+                self.transformer,
+                # args.model_type,
+                "wan",
+                self.optimizer,
+                self.lr_scheduler,
+                loader_iter,
+                noise_scheduler,
+                noise_random_generator,
+                self.training_args.gradient_accumulation_steps,
+                self.training_args.sp_size,
+                self.training_args.precondition_outputs,
+                self.training_args.max_grad_norm,
+                self.training_args.weighting_scheme,
+                self.training_args.logit_mean,
+                self.training_args.logit_std,
+                self.training_args.mode_scale,
+            )
+            gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2
+            logger.info("GPU memory usage after train_one_step: %s MB",
+                        gpu_memory_usage)
+
+            step_time = time.perf_counter() - start_time
+            step_times.append(step_time)
+            avg_step_time = sum(step_times) / len(step_times)
+
+            # Manual gradient checking - only at first step
+            if step == 1 and ENABLE_GRADIENT_CHECK:
+                logger.info("Performing gradient check at step %s", step)
+                self.setup_gradient_check(args, loader_iter, noise_scheduler,
+                                          noise_random_generator)
+
+            progress_bar.set_postfix({
+                "loss": f"{loss:.4f}",
+                "step_time": f"{step_time:.2f}s",
+                "grad_norm": grad_norm,
+            })
+            progress_bar.update(1)
+            if self.rank <= 0:
+                wandb.log(
+                    {
+                        "train_loss": loss,
+                        "learning_rate": self.lr_scheduler.get_last_lr()[0],
+                        "step_time": step_time,
+                        "avg_step_time": avg_step_time,
+                        "grad_norm": grad_norm,
+                    },
+                    step=step,
+                )
+            if step % self.training_args.checkpointing_steps == 0:
+                # Your existing checkpoint saving code
+                save_checkpoint(self.transformer, self.rank,
+                                self.training_args.output_dir, step)
+                self.transformer.train()
+                self.sp_group.barrier()
+            if self.training_args.log_validation and step % self.training_args.validation_steps == 0:
+                self.log_validation(self.transformer, self.training_args, step)
+
+        save_checkpoint(self.transformer, self.rank,
+                        self.training_args.output_dir,
+                        self.training_args.max_train_steps)
+
+        if get_sp_group():
+            cleanup_dist_env_and_memory()
+
+
+def main(args) -> None:
+    logger.info("Starting training pipeline...")
+
+    pipeline = WanTrainingPipeline.from_pretrained(
+        args.pretrained_model_name_or_path, args=args)
+    args = pipeline.training_args
+    pipeline.forward(None, args)
+    logger.info("Training pipeline done")
+
+
+if __name__ == "__main__":
+    argv = sys.argv
+    from fastvideo.v1.fastvideo_args import TrainingArgs
+    from fastvideo.v1.utils import FlexibleArgumentParser
+    parser = FlexibleArgumentParser()
+    parser = TrainingArgs.add_cli_args(parser)
+    parser = FastVideoArgs.add_cli_args(parser)
+    args = parser.parse_args()
+    args.use_cpu_offload = False
+    main(args)

scripts/dataset_preparation/prepare_json_file.py

@@ -1,11 +1,11 @@
 import json
 from pathlib import Path
-import csv
+
 import cv2
 
 
-def get_video_info(video_path, metadata):
-    """Extract video information using OpenCV and corresponding metadata"""
+def get_video_info(video_path, prompt_text):
+    """Extract video information using OpenCV and corresponding prompt text"""
     cap = cv2.VideoCapture(str(video_path))
 
     if not cap.isOpened():
@@ -23,66 +23,60 @@ def get_video_info(video_path, metadata):
 
     return {
         "path": video_path.name,
-        "title": metadata.get("Video Title", ""),
-        "description": metadata.get("Video Description", ""),
-        "video_url": metadata.get("Video URL", ""),
-        "download_url": metadata.get("Download URL", ""),
         "resolution": {
             "width": width,
             "height": height
         },
         "fps": fps,
         "duration": duration,
-        "cap": [metadata.get("Video Description", "")]
+        "cap": [prompt_text]
     }
 
 
-def read_csv_file(csv_path):
-    """Read and return the content of a CSV file"""
+def read_prompt_file(prompt_path):
+    """Read and return the content of a prompt file"""
     try:
-        with open(csv_path, 'r', encoding='utf-8') as f:
-            reader = csv.DictReader(f)
-            return list(reader)
+        with open(prompt_path, 'r', encoding='utf-8') as f:
+            return f.read().strip()
     except Exception as e:
-        print(f"Error reading CSV file {csv_path}: {e}")
+        print(f"Error reading prompt file {prompt_path}: {e}")
         return None
 
 
-def process_videos_from_csv(video_dir_path, csv_path, verbose=False):
-    """Process videos using metadata from CSV file
+def process_videos_and_prompts(video_dir_path, prompt_dir_path, verbose=False):
+    """Process videos and their corresponding prompt files
     
     Args:
         video_dir_path (str): Path to directory containing video files
-        csv_path (str): Path to CSV file containing video metadata
+        prompt_dir_path (str): Path to directory containing prompt files
         verbose (bool): Whether to print verbose processing information
     """
     video_dir = Path(video_dir_path)
-    csv_data = read_csv_file(csv_path)
+    prompt_dir = Path(prompt_dir_path)
     processed_data = []
 
     # Ensure directories exist
-    if not video_dir.exists():
-        print(f"Error: Video directory does not exist: {video_dir}")
-        return []
-
-    if csv_data is None:
+    if not video_dir.exists() or not prompt_dir.exists():
+        print(f"Error: One or both directories do not exist:\nVideos: {video_dir}\nPrompts: {prompt_dir}")
         return []
 
     # Process each video file
-    for row in csv_data:
-        video_filename = row.get("Filename")
-        if not video_filename:
+    for video_file in video_dir.glob('*.mp4'):
+        video_name = video_file.stem
+        prompt_file = prompt_dir / f"{video_name}.txt"
+
+        # Check if corresponding prompt file exists
+        if not prompt_file.exists():
+            print(f"Warning: No prompt file found for video {video_name}")
             continue
 
-        video_file = video_dir / video_filename
-
-        # Check if video file exists
-        if not video_file.exists():
-            print(f"Warning: Video file not found: {video_filename}")
+        # Read prompt content
+        prompt_text = read_prompt_file(prompt_file)
+        if prompt_text is None:
             continue
 
         # Process video and add to results
-        video_info = get_video_info(video_file, row)
+        video_info = get_video_info(video_file, prompt_text)
         if video_info:
             processed_data.append(video_info)
 
@@ -111,9 +105,9 @@ def parse_args():
     """Parse command line arguments"""
     import argparse
 
-    parser = argparse.ArgumentParser(description='Process videos using metadata from CSV file')
+    parser = argparse.ArgumentParser(description='Process videos and their corresponding prompt files')
     parser.add_argument('--video_dir', '-v', required=True, help='Directory containing video files')
-    parser.add_argument('--csv_path', '-c', required=True, help='Path to CSV file containing video metadata')
+    parser.add_argument('--prompt_dir', '-p', required=True, help='Directory containing prompt text files')
     parser.add_argument('--output_path',
                         '-o',
                         required=True,
@@ -127,8 +121,8 @@ if __name__ == "__main__":
     # Parse command line arguments
     args = parse_args()
 
-    # Process videos from CSV
-    processed_videos = process_videos_from_csv(args.video_dir, args.csv_path, args.verbose)
+    # Process videos and prompts
+    processed_videos = process_videos_and_prompts(args.video_dir, args.prompt_dir, args.verbose)
 
     if processed_videos:
         # Save results

scripts/dataset_preparation/resize_videos.py

@@ -24,9 +24,9 @@ def is_16_9_ratio(width: int, height: int, tolerance: float = 0.1) -> bool:
 def resize_video(args_tuple):
     """
     Resize a single video file.
-    args_tuple: (input_file, output_dir, width, height, fps, num_frames)
+    args_tuple: (input_file, output_dir, width, height, fps)
     """
-    input_file, output_dir, width, height, fps, num_frames = args_tuple
+    input_file, output_dir, width, height, fps = args_tuple
     video = None
     resized = None
     output_file = output_dir / f"{input_file.name}"
@@ -39,13 +39,6 @@ def resize_video(args_tuple):
     if not is_16_9_ratio(video.w, video.h):
         return (input_file.name, "skipped", "Not 16:9")
 
-    # Calculate target duration based on num_frames and fps
-    target_duration = num_frames / fps
-    
-    # Trim video if it's longer than target duration
-    if video.duration > target_duration:
-        video = video.subclip(0, target_duration)
-
     def process_frame(frame):
         frame_float = frame.astype(float) / 255.0
         resized = resize(frame_float, (height, width, 3), mode='reflect', anti_aliasing=True, preserve_range=True)
@@ -82,7 +75,7 @@ def process_folder(args):
     print(f"Target: {args.width}x{args.height} at {args.fps}fps")
 
     # Prepare arguments for parallel processing
-    process_args = [(video_file, output_path, args.width, args.height, args.fps, args.num_frames) for video_file in video_files]
+    process_args = [(video_file, output_path, args.width, args.height, args.fps) for video_file in video_files]
 
     successful = 0
     skipped = 0
@@ -122,7 +115,6 @@ def parse_args():
     parser.add_argument('--width', type=int, default=1280, help='Target width in pixels (default: 848)')
     parser.add_argument('--height', type=int, default=720, help='Target height in pixels (default: 480)')
     parser.add_argument('--fps', type=int, default=30, help='Target frames per second (default: 30)')
-    parser.add_argument('--num_frames', type=int, default=163, help='Target number of frames (default: 163)')
     parser.add_argument('--max_workers',
                         type=int,
                         default=4,

scripts/distill/distill_wan.sh

@@ -1,45 +0,0 @@
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-
-DATA_DIR=/workspace/data
-num_gpus=1
-IP=127.0.0.1
-
-torchrun --nnodes 1 --nproc_per_node $num_gpus \
-    --node_rank=0 \
-    --rdzv_id=456 \
-    --rdzv_backend=c10d \
-    --rdzv_endpoint=$IP:29500 \
-    fastvideo/distill_wan.py\
-    --seed 42\
-    --cache_dir "$DATA_DIR/.cache"\
-    --data_json_path "$DATA_DIR/HD-Mixkit-Finetune-Hunyuan/videos2caption.json"\
-    --validation_prompt_dir "$DATA_DIR/HD-Mixkit-Finetune-Hunyuan/validation"\
-    --train_batch_size=1 \
-    --num_latent_t 1 \
-    --sp_size $num_gpus \
-    --train_sp_batch_size 1\
-    --dataloader_num_workers 4\
-    --gradient_accumulation_steps=1\
-    --max_train_steps=320\
-    --learning_rate=1e-6\
-    --mixed_precision="bf16"\
-    --master_weight_type="bf16"\
-    --checkpointing_steps=64\
-    --validation_steps 64\
-    --validation_sampling_steps "2,4,8" \
-    --checkpoints_total_limit 3\
-    --allow_tf32\
-    --ema_start_step 0\
-    --cfg 0.0\
-    --log_validation\
-    --output_dir="$DATA_DIR/outputs/hy_phase1_shift17_bs_16_HD"\
-    --tracker_project_name Hunyuan_Distill \
-    --num_height 720 \
-    --num_width 1280 \
-    --num_frames  125 \
-    --shift 17 \
-    --validation_guidance_scale "1.0" \
-    --num_euler_timesteps 50 \
-    --multi_phased_distill_schedule "4000-1" \
-    --not_apply_cfg_solver 

scripts/finetune/finetune_v1.sh

@@ -2,18 +2,15 @@ export WANDB_BASE_URL="https://api.wandb.ai"
 export WANDB_MODE=online
 # export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
 
-DATA_DIR=data/cats_480_single_latents_parq/combined_parquet_dataset
-VALIDATION_DIR=data/cats_480_single_latents_parq/validation_parquet_dataset
-NUM_GPUS=4
+DATA_DIR=data/HD-Mixkit-Finetune-Wan/combined_parquet_dataset
+VALIDATION_DIR=data/HD-Mixkit-Finetune-Wan/validation_parquet_dataset
+NUM_GPUS=1
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]
 
 # If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
-    # --gradient_checkpointing\
-    # --pretrained_model_name_or_path hunyuanvideo-community/HunyuanVideo \
-    # --pretrained_model_name_or_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
 torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\
-    fastvideo/v1/pipelines/training_pipeline.py\
+    fastvideo/v1/training/wan_training_pipeline.py\
     --model_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
     --inference_mode False\
     --pretrained_model_name_or_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
@@ -21,17 +18,17 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\
     --data_path "$DATA_DIR"\
     --validation_prompt_dir "$VALIDATION_DIR"\
     --train_batch_size=1\
-    --num_latent_t 20 \
+    --num_latent_t 4 \
     --sp_size $NUM_GPUS \
     --tp_size $NUM_GPUS \
     --train_sp_batch_size 1\
-    --dataloader_num_workers 6\
+    --dataloader_num_workers 5\
     --gradient_accumulation_steps=1\
-    --max_train_steps=3000 \
+    --max_train_steps=120 \
     --learning_rate=1e-6\
     --mixed_precision="bf16"\
-    --checkpointing_steps=5000 \
-    --validation_steps 200\
+    --checkpointing_steps=50 \
+    --validation_steps 20\
     --validation_sampling_steps "2,4,8" \
     --log_validation \
     --checkpoints_total_limit 3\
@@ -50,4 +47,4 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\
     --weight_decay 0.01 \
     --not_apply_cfg_solver \
     --master_weight_type "fp32" \
-    --max_grad_norm 1.0 
+    --max_grad_norm 1.0

scripts/inference/v1_preprocess_wan.sh

@@ -1,27 +0,0 @@
-#!/bin/bash
-
-num_gpus=2
-export FASTVIDEO_ATTENTION_BACKEND=
-export MODEL_BASE=Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-# export MODEL_BASE=hunyuanvideo-community/HunyuanVideo
-# Note that the tp_size and sp_size should be the same and equal to the number
-# of GPUs. They are used for different parallel groups. sp_size is used for
-# dit model and tp_size is used for encoder models.
-torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
-    fastvideo/v1/entrypoints/data_preprocessor.py \
-    --sp_size $num_gpus \
-    --tp_size $num_gpus \
-    --height 480 \
-    --width 832 \
-    --num_frames 77 \
-    --num_inference_steps 50 \
-    --fps 16 \
-    --guidance_scale 3.0 \
-    --prompt_path ./assets/prompt.txt \
-    --neg_prompt "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" \
-    --seed 1024 \
-    --output_path outputs_video/ \
-    --model_path $MODEL_BASE \
-    --vae-sp \
-    --text-encoder-precision "fp32" \
-    --use-cpu-offload

scripts/preprocess/preprocess_wan_data.sh

@@ -1,10 +1,9 @@
 # export WANDB_MODE="offline"
 GPU_NUM=1 # 2,4,8
-MODEL_PATH="/workspace/data/Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-TEXT_ENCODER_PATH="/workspace/data/Wan-AI/Wan2.1-T2V-1.3B-Diffusers/tokenizer"
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 MODEL_TYPE="wan"
-DATA_MERGE_PATH="/workspace/data/Mixkit-Src/merge.txt"
-OUTPUT_DIR="/workspace/data/HD-Mixkit-Finetune-Wan"
+DATA_MERGE_PATH="your/path/to/Mixkit-Src/merge.txt"
+OUTPUT_DIR="your/path"
 VALIDATION_PATH="assets/prompt.txt"
 
 torchrun --nproc_per_node=$GPU_NUM \
@@ -12,13 +11,13 @@ torchrun --nproc_per_node=$GPU_NUM \
     --model_path $MODEL_PATH \
     --data_merge_path $DATA_MERGE_PATH \
     --preprocess_video_batch_size=4 \
-    --preprocess_text_batch_size=4 \
     --max_height=480 \
     --max_width=832 \
     --num_frames=81 \
     --dataloader_num_workers 1 \
     --output_dir=$OUTPUT_DIR \
     --model_type $MODEL_TYPE \
-    --text_encoder_name $TEXT_ENCODER_PATH \
     --train_fps 16 \
-    --validation_prompt_txt $VALIDATION_PATH
+    --validation_prompt_txt $VALIDATION_PATH \
+    --samples_per_file 108 \
+    --flush_frequency 108