[Fix] adapt to main

Eigensystem · Eigensystem · commit 555b7a0b88b9 · 2025-06-07T17:56:28.000+08:00
diff --git a/fastvideo/v1/fastvideo_args.py b/fastvideo/v1/fastvideo_args.py
@@ -6,15 +6,25 @@
 import dataclasses
 from contextlib import contextmanager
 from dataclasses import field
+from enum import Enum
 from typing import Any, Callable, List, Optional, Tuple
 
+import torch
+
 from fastvideo.v1.configs.models import DiTConfig, EncoderConfig, VAEConfig
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.utils import FlexibleArgumentParser, StoreBoolean
 
 logger = init_logger(__name__)
 
 
+class Mode(Enum):
+    """Enumeration for FastVideo execution modes."""
+    INFERENCE = "inference"
+    TRAINING = "training"
+    DISTILL = "distill"
+
+
 def preprocess_text(prompt: str) -> str:
     return prompt
 
@@ -34,7 +44,7 @@ class FastVideoArgs:
     # Distributed executor backend
     distributed_executor_backend: str = "mp"
 
-    mode: str = "inference"  # Options: "inference", "training", "distill"
+    mode: Mode = Mode.INFERENCE
 
     # HuggingFace specific parameters
     trust_remote_code: bool = False
@@ -115,15 +125,15 @@ class FastVideoArgs:
 
     @property
     def training_mode(self) -> bool:
-        return self.mode == "training"
+        return self.mode == Mode.TRAINING
 
     @property
     def distill_mode(self) -> bool:
-        return self.mode == "distill"
+        return self.mode == Mode.DISTILL
 
     @property
     def inference_mode(self) -> bool:
-        return self.mode == "inference"
+        return self.mode == Mode.INFERENCE
 
     def __post_init__(self):
         pass
@@ -160,8 +170,8 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
         parser.add_argument(
             "--mode",
             type=str,
-            default=FastVideoArgs.mode,
-            choices=["inference", "training", "distill"],
+            default=FastVideoArgs.mode.value,
+            choices=[mode.value for mode in Mode],
             help="The mode to use",
         )
 
@@ -376,9 +386,16 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
 
     @classmethod
     def from_cli_args(cls, args: argparse.Namespace) -> "FastVideoArgs":
-        args.tp_size = args.tensor_parallel_size
-        args.sp_size = args.sequence_parallel_size
-        args.flow_shift = getattr(args, "shift", args.flow_shift)
+        assert getattr(args, 'model_path', None) is not None, "model_path must be set in args"
+        # Handle attribute mapping with safe getattr
+        if hasattr(args, 'tensor_parallel_size'):
+            args.tp_size = args.tensor_parallel_size
+        if hasattr(args, 'sequence_parallel_size'):
+            args.sp_size = args.sequence_parallel_size
+        if hasattr(args, 'shift'):
+            args.flow_shift = args.shift
+        elif hasattr(args, 'flow_shift'):
+            args.flow_shift = args.flow_shift
 
         # Get all fields from the dataclass
         attrs = [attr.name for attr in dataclasses.fields(cls)]
@@ -397,6 +414,18 @@ def from_cli_args(cls, args: argparse.Namespace) -> "FastVideoArgs":
                 kwargs[attr] = args.data_parallel_shards
             elif attr == 'flow_shift' and hasattr(args, 'shift'):
                 kwargs[attr] = args.shift
+            elif attr == 'mode':
+                # Convert string mode to Mode enum
+                mode_value = getattr(args, attr, None)
+                if mode_value:
+                    if isinstance(mode_value, Mode):
+                        kwargs[attr] = mode_value
+                    else:
+                        kwargs[attr] = Mode(mode_value)
+                else:
+                    kwargs[attr] = Mode.INFERENCE
+            elif attr == 'device_str':
+                kwargs[attr] = getattr(args, 'device', None) or "cuda" if torch.cuda.is_available() else "cpu"
             # Use getattr with default value from the dataclass for potentially missing attributes
             else:
                 default_value = getattr(cls, attr, None)
@@ -595,9 +624,6 @@ class TrainingArgs(FastVideoArgs):
     # master_weight_type
     master_weight_type: str = ""
 
-    # For fast checking in LoRA pipeline
-    training_mode: bool = True
-
     @classmethod
     def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs":
         # Get all fields from the dataclass
@@ -617,6 +643,18 @@ def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs":
                 kwargs[attr] = args.data_parallel_size
             elif attr == 'dp_shards' and hasattr(args, 'data_parallel_shards'):
                 kwargs[attr] = args.data_parallel_shards
+            elif attr == 'mode':
+                # Convert string mode to Mode enum
+                mode_value = getattr(args, attr, None)
+                if mode_value:
+                    if isinstance(mode_value, Mode):
+                        kwargs[attr] = mode_value
+                    else:
+                        kwargs[attr] = Mode(mode_value)
+                else:
+                    kwargs[attr] = Mode.TRAINING  # Default to training for TrainingArgs
+            elif attr == 'device_str':
+                kwargs[attr] = getattr(args, 'device', None) or "cuda" if torch.cuda.is_available() else "cpu"
             # Use getattr with default value from the dataclass for potentially missing attributes
             else:
                 default_value = getattr(cls, attr, None)
diff --git a/fastvideo/v1/pipelines/composed_pipeline_base.py b/fastvideo/v1/pipelines/composed_pipeline_base.py
@@ -18,7 +18,7 @@
 from fastvideo.v1.distributed import (init_distributed_environment,
                                       initialize_model_parallel,
                                       model_parallel_is_initialized)
-from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs
+from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs, Mode
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.models.loader.component_loader import PipelineComponentLoader
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
@@ -94,7 +94,6 @@ def __init__(self,
                 self.initialize_validation_pipeline(self.training_args)
             self.initialize_training_pipeline(self.training_args)
 
-        # TODO(jinzhe): discuss this
         if fastvideo_args.distill_mode:
             assert self.training_args is not None
             if self.training_args.log_validation:
@@ -159,38 +158,33 @@ def from_pretrained(cls,
             config_args = shallow_asdict(config)
             config_args.update(kwargs)
 
-        if args.mode == "inference":
-            fastvideo_args = FastVideoArgs(model_path=model_path,
-                                           device_str=device or "cuda" if
-                                           torch.cuda.is_available() else "cpu",
-                                           **config_args)
-
-            fastvideo_args.model_path = model_path
-            fastvideo_args.device_str = device or "cuda" if torch.cuda.is_available(
-            ) else "cpu"
+        args.model_path = model_path
+        # Handle both string mode and Mode enum values
+        mode_str = args.mode if isinstance(args.mode, str) else args.mode.value
+        
+        if mode_str == "inference":
+            fastvideo_args = FastVideoArgs.from_cli_args(args)
             for key, value in config_args.items():
                 setattr(fastvideo_args, key, value)
-        else:
+
+        elif mode_str == "training" or mode_str == "distill":
             assert args is not None, "args must be provided for training mode"
             fastvideo_args = TrainingArgs.from_cli_args(args)
-            # TODO(will): fix this so that its not so ugly
-            fastvideo_args.model_path = model_path
-            fastvideo_args.device_str = device or "cuda" if torch.cuda.is_available(
-            ) else "cpu"
             for key, value in config_args.items():
                 setattr(fastvideo_args, key, value)
 
             fastvideo_args.num_gpus = int(os.environ.get("WORLD_SIZE", 1))
             fastvideo_args.use_cpu_offload = False
             # make sure we are in training mode
-            fastvideo_args.mode = args.mode
             # we hijack the precision to be the master weight type so that the
             # model is loaded with the correct precision. Subsequently we will
             # use FSDP2's MixedPrecisionPolicy to set the precision for the
             # fwd, bwd, and other operations' precision.
             # fastvideo_args.precision = fastvideo_args.master_weight_type
             assert fastvideo_args.master_weight_type == 'fp32', 'only fp32 is supported for training'
-            # assert fastvideo_args.precision == 'fp32', 'only fp32 is supported for training'
+        else:
+            raise ValueError(f"Invalid mode: {mode_str}")
+
 
         fastvideo_args.check_fastvideo_args()
 
diff --git a/fastvideo/v1/training/distillation_pipeline.py b/fastvideo/v1/training/distillation_pipeline.py
@@ -17,7 +17,7 @@
 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.fastvideo_args import FastVideoArgs, TrainingArgs, Mode
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.pipelines import ComposedPipelineBase
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
@@ -150,8 +150,6 @@ def initialize_distillation_pipeline(self, fastvideo_args: TrainingArgs):
         train_dataset = ParquetVideoTextDataset(
             fastvideo_args.data_path,
             batch_size=fastvideo_args.train_batch_size,
-            rank=self.rank,
-            world_size=self.world_size,
             cfg_rate=fastvideo_args.cfg,
             num_latent_t=fastvideo_args.num_latent_t)
 
@@ -203,7 +201,7 @@ def distill_one_step(self, transformer, model_type, teacher_transformer,
 
     def log_validation(self, transformer, fastvideo_args, global_step):
         """Log validation results during training."""
-        fastvideo_args.mode = "inference"
+        fastvideo_args.mode = Mode.INFERENCE
         fastvideo_args.use_cpu_offload = False
         if not fastvideo_args.log_validation:
             return
@@ -218,8 +216,6 @@ def log_validation(self, transformer, fastvideo_args, global_step):
         validation_dataset = ParquetVideoTextDataset(
             fastvideo_args.validation_prompt_dir,
             batch_size=1,
-            rank=0,
-            world_size=1,
             cfg_rate=0,
             num_latent_t=fastvideo_args.num_latent_t)
 
@@ -324,7 +320,7 @@ def log_validation(self, transformer, fastvideo_args, global_step):
             wandb.log(logs, step=global_step)
 
         # Re-enable gradients for training
-        fastvideo_args.mode = "distill"
+        fastvideo_args.mode = Mode.DISTILL
         transformer.requires_grad_(True)
         transformer.train()
 
diff --git a/fastvideo/v1/training/wan_distillation_pipeline.py b/fastvideo/v1/training/wan_distillation_pipeline.py
@@ -9,7 +9,7 @@
 
 from fastvideo.distill.solver import extract_into_tensor
 from fastvideo.v1.distributed import cleanup_dist_env_and_memory, get_sp_group
-from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs
+from fastvideo.v1.fastvideo_args import FastVideoArgs, Mode, 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
@@ -60,7 +60,7 @@ def initialize_validation_pipeline(self, fastvideo_args: FastVideoArgs):
         logger.info("Initializing validation pipeline...")
         args_copy = deepcopy(fastvideo_args)
 
-        args_copy.mode = "inference"
+        args_copy.mode = Mode.INFERENCE
         args_copy.vae_config.load_encoder = False
         validation_pipeline = WanValidationPipeline.from_pretrained(
             fastvideo_args.model_path, args=args_copy)
diff --git a/scripts/distill/distill_v1.sh b/scripts/distill/distill_v1.sh
@@ -21,6 +21,8 @@ torchrun --nnodes 1 --nproc_per_node $num_gpus\
     --train_batch_size=1 \
     --num_latent_t 4 \
     --sp_size $num_gpus \
+    --dp_size $num_gpus \
+    --dp_shards $num_gpus \
     --train_sp_batch_size 1 \
     --dataloader_num_workers $num_gpus \
     --gradient_accumulation_steps=1 \
