[Fmt] fmt code

Author: Eigensystem

fastvideo/v1/dataset/parquet_datasets.py

@@ -192,7 +192,6 @@ def get_validation_negative_prompt(
                 lat = lat[:, self.rank_in_sp_group, :, :, :]
             return lat, emb, mask, info
 
-
     def __len__(self):
         if self.local_indices is None:
             try:

fastvideo/v1/fastvideo_args.py

@@ -386,7 +386,8 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
 
     @classmethod
     def from_cli_args(cls, args: argparse.Namespace) -> "FastVideoArgs":
-        assert getattr(args, 'model_path', None) is not None, "model_path must be set in args"
+        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
@@ -425,7 +426,9 @@ def from_cli_args(cls, args: argparse.Namespace) -> "FastVideoArgs":
                 else:
                     kwargs[attr] = Mode.INFERENCE
             elif attr == 'device_str':
-                kwargs[attr] = getattr(args, 'device', None) or "cuda" if torch.cuda.is_available() else "cpu"
+                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)
@@ -652,9 +655,12 @@ def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs":
                     else:
                         kwargs[attr] = Mode(mode_value)
                 else:
-                    kwargs[attr] = Mode.TRAINING  # Default to training for TrainingArgs
+                    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"
+                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)

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, Mode
+from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs
 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
@@ -161,7 +161,7 @@ def from_pretrained(cls,
         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():
@@ -185,7 +185,6 @@ def from_pretrained(cls,
         else:
             raise ValueError(f"Invalid mode: {mode_str}")
 
-
         fastvideo_args.check_fastvideo_args()
 
         logger.info("fastvideo_args in from_pretrained: %s", fastvideo_args)

fastvideo/v1/training/distillation_pipeline.py

@@ -6,18 +6,18 @@
 import numpy as np
 import torch
 import torchvision
-import wandb
 from diffusers.optimization import get_scheduler
 from einops import rearrange
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from torch.distributed.fsdp import ShardingStrategy
 from torchdata.stateful_dataloader import StatefulDataLoader
 
+import wandb
 from fastvideo.distill.solver import EulerSolver
 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, Mode
+from fastvideo.v1.fastvideo_args import FastVideoArgs, Mode, TrainingArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.pipelines import ComposedPipelineBase
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
@@ -54,6 +54,7 @@ def reshard_fsdp(model):
         if m._has_params and m.sharding_strategy is not ShardingStrategy.NO_SHARD:
             torch.distributed.fsdp._runtime_utils._reshard(m, m._handle, True)
 
+
 class DistillationPipeline(ComposedPipelineBase, ABC):
     """
     A pipeline for distillation training. All distillation pipelines should inherit from this class.
@@ -77,10 +78,12 @@ def initialize_distillation_pipeline(self, fastvideo_args: TrainingArgs):
 
         # Initialize teacher model without deepcopy to avoid FSDP issues
         logger.info("Creating teacher model...")
-        from fastvideo.v1.models.loader.component_loader import TransformerLoader
+        from fastvideo.v1.models.loader.component_loader import (
+            TransformerLoader)
         teacher_loader = TransformerLoader()
         transformer_path = os.path.join(self.model_path, "transformer")
-        self.teacher_transformer = teacher_loader.load(transformer_path, "", fastvideo_args)
+        self.teacher_transformer = teacher_loader.load(transformer_path, "",
+                                                       fastvideo_args)
         self.teacher_transformer.requires_grad_(False)
         self.teacher_transformer.eval()
         logger.info("Teacher model initialized")
@@ -89,7 +92,8 @@ def initialize_distillation_pipeline(self, fastvideo_args: TrainingArgs):
         if fastvideo_args.use_ema:
             logger.info("Creating EMA model...")
             ema_loader = TransformerLoader()
-            self.ema_transformer = ema_loader.load(transformer_path, "", fastvideo_args)
+            self.ema_transformer = ema_loader.load(transformer_path, "",
+                                                   fastvideo_args)
             self.ema_transformer.requires_grad_(False)
             self.ema_transformer.eval()
             logger.info("EMA model initialized")
@@ -326,5 +330,3 @@ def log_validation(self, transformer, fastvideo_args, global_step):
 
         gc.collect()
         torch.cuda.empty_cache()
-
-

fastvideo/v1/training/wan_distillation_pipeline.py

@@ -4,23 +4,25 @@
 from copy import deepcopy
 
 import torch
-import wandb
 from tqdm.auto import tqdm
 
+import wandb
 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, 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
-from fastvideo.v1.training.training_utils import (
-    clip_grad_norm_while_handling_failing_dtensor_cases,
-    save_checkpoint, normalize_dit_input)
 from fastvideo.v1.pipelines.wan.wan_pipeline import WanValidationPipeline
-from fastvideo.v1.training.distillation_pipeline import DistillationPipeline, reshard_fsdp
+from fastvideo.v1.training.distillation_pipeline import (DistillationPipeline,
+                                                         reshard_fsdp)
+from fastvideo.v1.training.training_utils import (
+    clip_grad_norm_while_handling_failing_dtensor_cases, normalize_dit_input,
+    save_checkpoint)
 
 logger = init_logger(__name__)
 
+
 def get_norm(model_pred, norms, gradient_accumulation_steps):
     """Calculate and aggregate model prediction norms."""
     fro_norm = (
@@ -44,6 +46,7 @@ def get_norm(model_pred, norms, gradient_accumulation_steps):
     norms["absolute mean"] += absolute_mean.item()
     norms["absolute max"] += absolute_max.item()
 
+
 class WanDistillationPipeline(DistillationPipeline):
     """
     A distillation pipeline for Wan.
@@ -124,15 +127,14 @@ def distill_one_step(
             noise = torch.randn_like(latents)
 
             indices = torch.randint(0,
-                                  num_euler_timesteps, (batch_size, ),
-                                  device=latents.device).long()
+                                    num_euler_timesteps, (batch_size, ),
+                                    device=latents.device).long()
 
             if sp_size > 1:
                 self.sp_group.broadcast(indices, src=0)
 
             # Add noise according to flow matching
-            sigmas = extract_into_tensor(solver.sigmas, indices,
-                                         latents.shape)
+            sigmas = extract_into_tensor(solver.sigmas, indices, latents.shape)
             sigmas_prev = extract_into_tensor(solver.sigmas_prev, indices,
                                               latents.shape)
 
@@ -186,16 +188,23 @@ def distill_one_step(
                     # Get teacher model prediction on unconditional embedding
                     with torch.autocast("cuda", dtype=torch.bfloat16):
                         input_kwargs = {
-                            "hidden_states": noisy_model_input,
-                            "encoder_hidden_states": uncond_prompt_embed.unsqueeze(0).expand(
-                                    batch_size, -1, -1),
-                            "timestep": timesteps,
-                            "encoder_attention_mask": uncond_prompt_mask.unsqueeze(0).expand(batch_size, -1),
-                            "return_dict": False,
+                            "hidden_states":
+                            noisy_model_input,
+                            "encoder_hidden_states":
+                            uncond_prompt_embed.unsqueeze(0).expand(
+                                batch_size, -1, -1),
+                            "timestep":
+                            timesteps,
+                            "encoder_attention_mask":
+                            uncond_prompt_mask.unsqueeze(0).expand(
+                                batch_size, -1),
+                            "return_dict":
+                            False,
                         }
                         with set_forward_context(current_timestep=timesteps,
                                                  attn_metadata=None):
-                            uncond_teacher_output = teacher_transformer(**input_kwargs)[0]
+                            uncond_teacher_output = teacher_transformer(
+                                **input_kwargs)[0]
                 teacher_output = uncond_teacher_output + distill_cfg * (
                     cond_teacher_output - uncond_teacher_output)
                 x_prev = solver.euler_step(noisy_model_input, teacher_output,
@@ -305,19 +314,24 @@ def forward(
         uncond_prompt_mask = self.uncond_prompt_mask
 
         # Train!
-        total_batch_size = (self.world_size * self.training_args.gradient_accumulation_steps /
-                            self.training_args.sp_size * self.training_args.train_sp_batch_size)
+        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 distillation training *****")
         logger.info(f"  Resume training from step {init_steps}")
         logger.info(
-            f"  Instantaneous batch size per device = {self.training_args.train_batch_size}")
+            f"  Instantaneous batch size per device = {self.training_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 = {self.training_args.gradient_accumulation_steps}"
         )
-        logger.info(f"  Total optimization steps = {self.training_args.max_train_steps}")
+        logger.info(
+            f"  Total optimization steps = {self.training_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"
         )
@@ -354,12 +368,13 @@ def get_num_phases(multi_phased_distill_schedule, step):
                     return int(phase)
             return int(phase)
 
-        for step in range(init_steps + 1, self.training_args.max_train_steps + 1):
+        for step in range(init_steps + 1,
+                          self.training_args.max_train_steps + 1):
             start_time = time.perf_counter()
 
             assert self.training_args.multi_phased_distill_schedule is not None
-            num_phases = get_num_phases(self.training_args.multi_phased_distill_schedule,
-                                        step)
+            num_phases = get_num_phases(
+                self.training_args.multi_phased_distill_schedule, step)
             try:
                 loss, grad_norm, pred_norm = self.distill_one_step(
                     self.transformer,
@@ -407,7 +422,6 @@ def get_num_phases(multi_phased_distill_schedule, step):
                 step -= 1
                 continue
 
-
             if self.rank <= 0:
                 wandb.log(
                     {
@@ -441,10 +455,10 @@ def get_num_phases(multi_phased_distill_schedule, step):
                 else:
                     if self.training_args.use_ema:
                         save_checkpoint(self.ema_transformer, self.rank,
-                                           self.training_args.output_dir, step)
+                                        self.training_args.output_dir, step)
                     else:
                         save_checkpoint(self.transformer, self.rank,
-                                           self.training_args.output_dir, step)
+                                        self.training_args.output_dir, step)
                 self.sp_group.barrier()
 
             if self.training_args.log_validation and step % self.training_args.validation_steps == 0:
@@ -454,8 +468,9 @@ def get_num_phases(multi_phased_distill_schedule, step):
         if self.training_args.use_lora:
             raise NotImplementedError("LoRA is not supported now")
         else:
-            save_checkpoint(self.transformer, self.rank, self.training_args.output_dir,
-                               self.training_args.max_train_steps)
+            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()