[Feature] [Training] Add i2v training (#559)

Author: SolitaryThinker

examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.sh

@@ -69,7 +69,7 @@ miscellaneous_args=(
   --inference_mode False
   --allow_tf32
   --checkpoints_total_limit 3
-  --cfg 0.0
+  --training_cfg_rate 0.1
   --multi_phased_distill_schedule "4000-1"
   --not_apply_cfg_solver
   --dit_precision "fp32"

examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.slurm

@@ -106,7 +106,7 @@ miscellaneous_args=(
   --inference_mode False
   --allow_tf32
   --checkpoints_total_limit 3
-  --cfg 0.0
+  --training_cfg_rate 0.1
   --multi_phased_distill_schedule "4000-1"
   --not_apply_cfg_solver
   --dit_precision "fp32"

examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.sh

@@ -69,7 +69,7 @@ miscellaneous_args=(
   --inference_mode False
   --allow_tf32
   --checkpoints_total_limit 3
-  --cfg 0.0
+  --training_cfg_rate 0.1
   --multi_phased_distill_schedule "4000-1"
   --not_apply_cfg_solver
   --dit_precision "fp32"

examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.slurm

@@ -103,7 +103,7 @@ miscellaneous_args=(
   --inference_mode False
   --allow_tf32
   --checkpoints_total_limit 3
-  --cfg 0.0
+  --training_cfg_rate 0.1
   --multi_phased_distill_schedule "4000-1"
   --not_apply_cfg_solver
   --dit_precision "fp32"

fastvideo/v1/dataset/parquet_dataset_iterable_style.py

@@ -4,6 +4,7 @@
 from typing import Dict, List, Tuple
 
 import numpy as np
+import pyarrow as pa
 import pyarrow.parquet as pq
 import torch
 import tqdm
@@ -70,10 +71,12 @@ def __init__(self,
                  drop_last: bool = True,
                  text_padding_length: int = 512,
                  seed: int = 42,
-                 read_batch_size: int = 32):
+                 read_batch_size: int = 32,
+                 parquet_schema: pa.Schema = None):
         super().__init__()
         self.path = str(path)
         self.batch_size = batch_size
+        self.parquet_schema = parquet_schema
         self.cfg_rate = cfg_rate
         self.text_padding_length = text_padding_length
         self.seed = seed

fastvideo/v1/dataset/parquet_dataset_map_style.py

@@ -201,10 +201,6 @@ def __init__(
         self.path = path
         self.cfg_rate = cfg_rate
         self.parquet_schema = parquet_schema
-        if cfg_rate > 0.0:
-            raise ValueError(
-                "cfg_rate > 0.0 is not supported for now because it will trigger bug when num_data_workers > 0"
-            )
         logger.info("Initializing LatentsParquetMapStyleDataset with path: %s",
                     path)
         self.parquet_files, self.lengths = get_parquet_files_and_length(path)

fastvideo/v1/dataset/utils.py

@@ -1,5 +1,5 @@
 import random
-from typing import Any, Dict, List
+from typing import Any, Dict, List, cast
 
 import numpy as np
 import torch
@@ -108,7 +108,7 @@ def collate_rows_from_parquet_schema(rows,
         Dict containing batched tensors and metadata
     """
     if not rows:
-        return {}
+        return cast(Dict[str, Any], {})
 
     # Initialize containers for different data types
     batch_data: Dict[str, Any] = {}

fastvideo/v1/pipelines/pipeline_batch_info.py

@@ -39,6 +39,7 @@ class ForwardBatch:
     image_path: Optional[str] = None
     image_embeds: List[torch.Tensor] = field(default_factory=list)
     pil_image: Optional[PIL.Image.Image] = None
+    preprocessed_image: Optional[torch.Tensor] = None
 
     # Text inputs
     prompt: Optional[Union[str, List[str]]] = None
@@ -150,6 +151,10 @@ class TrainingBatch:
     latents: Optional[torch.Tensor] = None
     encoder_hidden_states: Optional[torch.Tensor] = None
     encoder_attention_mask: Optional[torch.Tensor] = None
+    # i2v
+    preprocessed_image: Optional[torch.Tensor] = None
+    image_embeds: Optional[torch.Tensor] = None
+    image_latents: Optional[torch.Tensor] = None
     infos: Optional[List[Dict[str, Any]]] = None
 
     # Transformer inputs