diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index a00649198..2a7327a78 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -60,7 +60,7 @@ repos: rev: v1.15.0 hooks: - id: mypy - args: [--python-version, '3.10', --follow-imports, "skip" ] + args: [--python-version, '3.10', --follow-imports, "skip", "--disable-error-code", "union-attr", "--disable-error-code", "override" ] additional_dependencies: [types-cachetools, types-setuptools, types-PyYAML, types-requests] - repo: local hooks: diff --git a/fastvideo/configs/pipelines/registry.py b/fastvideo/configs/pipelines/registry.py index e2d4bbab3..86620abac 100644 --- a/fastvideo/configs/pipelines/registry.py +++ b/fastvideo/configs/pipelines/registry.py @@ -27,6 +27,7 @@ PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = { "Wan-AI/Wan2.1-T2V-14B-Diffusers": WanT2V720PConfig, "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": FastWanT2V480PConfig, "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VConfig, + "FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers": WanT2V720PConfig, # Add other specific weight variants } diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index 48dd39774..fa66be947 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -139,6 +139,9 @@ class FastVideoArgs: # VSA parameters VSA_sparsity: float = 0.0 # inference/validation sparsity + # Master port for distributed training/inference + master_port: int | None = None + # Stage verification enable_stage_verification: bool = True @@ -360,6 +363,14 @@ class FastVideoArgs: help="Validation sparsity for VSA", ) + # Master port for distributed training/inference + parser.add_argument( + "--master-port", + type=int, + default=FastVideoArgs.master_port, + help="Master port for distributed training/inference", + ) + # Stage verification parser.add_argument( "--enable-stage-verification", diff --git a/fastvideo/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py index eb7551739..09e6f1045 100644 --- a/fastvideo/training/distillation_pipeline.py +++ b/fastvideo/training/distillation_pipeline.py @@ -77,7 +77,6 @@ class DistillationPipeline(TrainingPipeline): self.vae = self.get_module("vae") self.vae.requires_grad_(False) - assert self.training_args is not None self.timestep_shift = self.training_args.pipeline_config.flow_shift self.noise_scheduler = FlowMatchEulerDiscreteScheduler( shift=self.timestep_shift) @@ -182,7 +181,7 @@ class DistillationPipeline(TrainingPipeline): return training_batch def _generator_forward(self, training_batch: TrainingBatch) -> torch.Tensor: - assert training_batch.latents is not None + latents = training_batch.latents dtype = latents.dtype index = torch.randint(0, @@ -211,7 +210,7 @@ class DistillationPipeline(TrainingPipeline): pred_noise = self.transformer(**training_batch.input_kwargs).permute( 0, 2, 1, 3, 4) - assert training_batch.noise_latents is not None + pred_video = pred_noise_to_pred_video( pred_noise=pred_noise.flatten(0, 1), noise_input_latent=training_batch.noise_latents.flatten(0, 1), @@ -228,7 +227,7 @@ class DistillationPipeline(TrainingPipeline): self.num_train_timestep, [1], device=self.device, dtype=torch.long) - assert self.timestep_shift is not None + timestep = shift_timestep( timestep, self.timestep_shift, # type: ignore @@ -256,7 +255,6 @@ class DistillationPipeline(TrainingPipeline): fake_score_pred_noise = self.fake_score_transformer( **training_batch.input_kwargs).permute(0, 2, 1, 3, 4) - assert training_batch.noise_latents is not None pred_fake_video = pred_noise_to_pred_video( pred_noise=fake_score_pred_noise.flatten(0, 1), noise_input_latent=training_batch.noise_latents.flatten(0, 1), @@ -271,7 +269,6 @@ class DistillationPipeline(TrainingPipeline): real_score_pred_noise_cond = self.real_score_transformer( **training_batch.input_kwargs).permute(0, 2, 1, 3, 4) - assert training_batch.noise_latents is not None pred_real_video_cond = pred_noise_to_pred_video( pred_noise=real_score_pred_noise_cond.flatten(0, 1), noise_input_latent=training_batch.noise_latents.flatten(0, 1), @@ -286,7 +283,6 @@ class DistillationPipeline(TrainingPipeline): real_score_pred_noise_uncond = self.real_score_transformer( **training_batch.input_kwargs).permute(0, 2, 1, 3, 4) - assert training_batch.noise_latents is not None pred_real_video_uncond = pred_noise_to_pred_video( pred_noise=real_score_pred_noise_uncond.flatten(0, 1), noise_input_latent=training_batch.noise_latents.flatten(0, 1), @@ -320,7 +316,7 @@ class DistillationPipeline(TrainingPipeline): self.num_train_timestep, [1], device=self.device, dtype=torch.long) - assert self.timestep_shift is not None + fake_score_timestep = shift_timestep( fake_score_timestep, self.timestep_shift, # type: ignore @@ -360,7 +356,7 @@ class DistillationPipeline(TrainingPipeline): def _clip_model_grad_norm_(self, training_batch: TrainingBatch, transformer) -> TrainingBatch: - assert self.training_args is not None + max_grad_norm = self.training_args.max_grad_norm if max_grad_norm is not None: @@ -392,7 +388,7 @@ class DistillationPipeline(TrainingPipeline): training_batch.conditional_dict = conditional_dict training_batch.unconditional_dict = unconditional_dict - assert training_batch.latents is not None + training_batch.latents = training_batch.latents.permute(0, 2, 1, 3, 4) self.video_latent_shape = training_batch.latents.shape training_batch.raw_latent_shape = training_batch.latents.shape @@ -488,7 +484,7 @@ class DistillationPipeline(TrainingPipeline): def _resume_from_checkpoint(self) -> None: #TODO(yongqi) """Resume training from checkpoint with distillation models.""" - assert self.training_args is not None + logger.info("Loading distillation checkpoint from %s", self.training_args.resume_from_checkpoint) @@ -518,12 +514,12 @@ class DistillationPipeline(TrainingPipeline): self.generator_update_interval) assert isinstance(self.training_args, TrainingArgs) logger.info(" Max gradient norm: %s", self.training_args.max_grad_norm) - assert self.real_score_transformer is not None + logger.info( " Real score transformer parameters: %s B", sum(p.numel() for p in self.real_score_transformer.parameters()) / 1e9) - assert self.fake_score_transformer is not None + logger.info( " Fake score transformer parameters: %s B", sum(p.numel() @@ -531,7 +527,6 @@ class DistillationPipeline(TrainingPipeline): @torch.no_grad() def _log_validation(self, transformer, training_args, global_step) -> None: - assert training_args is not None training_args.inference_mode = True training_args.dit_cpu_offload = True if not training_args.log_validation: @@ -672,8 +667,6 @@ class DistillationPipeline(TrainingPipeline): def train(self) -> None: """Main training loop with distillation-specific logging.""" - assert self.training_args is not None - assert self.training_args.seed is not None, "seed must be set" seed = self.training_args.seed diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index be61ff314..2be1f43ee 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -97,12 +97,11 @@ class TrainingPipeline(LoRAPipeline, ABC): self.sp_world_size = self.sp_group.world_size self.local_rank = world_group.local_rank self.transformer = self.get_module("transformer") - assert training_args.seed is not None self.seed = training_args.seed - assert self.transformer is not None self.set_schemas() # Set random seeds for deterministic training + assert self.seed is not None, "seed must be set" set_random_seed(self.seed) self.transformer.train() if training_args.enable_gradient_checkpointing_type is not None: @@ -151,10 +150,6 @@ class TrainingPipeline(LoRAPipeline, ABC): self.noise_scheduler = noise_scheduler - assert training_args.gradient_accumulation_steps is not None - assert training_args.sp_size is not None - assert training_args.train_sp_batch_size is not None - assert training_args.max_train_steps is not None self.num_update_steps_per_epoch = math.ceil( len(self.train_dataloader) / training_args.gradient_accumulation_steps * training_args.sp_size / @@ -183,10 +178,6 @@ class TrainingPipeline(LoRAPipeline, ABC): return training_batch def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch: - assert self.training_args is not None - assert self.train_loader_iter is not None - assert self.train_dataloader is not None - batch = next(self.train_loader_iter, None) # type: ignore if batch is None: self.current_epoch += 1 @@ -222,11 +213,6 @@ class TrainingPipeline(LoRAPipeline, ABC): def _prepare_dit_inputs(self, training_batch: TrainingBatch) -> TrainingBatch: - assert self.training_args is not None - assert training_batch.latents is not None - assert training_batch.encoder_hidden_states is not None - assert training_batch.encoder_attention_mask is not None - assert self.noise_random_generator is not None latents = training_batch.latents batch_size = latents.shape[0] noise = torch.randn(latents.shape, @@ -268,13 +254,11 @@ class TrainingPipeline(LoRAPipeline, ABC): def _build_attention_metadata( self, training_batch: TrainingBatch) -> TrainingBatch: - assert self.training_args is not None - assert training_batch.timesteps is not None - assert training_batch.raw_latent_shape is not None latents_shape = training_batch.raw_latent_shape patch_size = self.training_args.pipeline_config.dit_config.patch_size current_vsa_sparsity = training_batch.current_vsa_sparsity - + assert latents_shape is not None + assert training_batch.timesteps is not None if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN": dit_seq_shape = [ latents_shape[2] // patch_size[0], @@ -292,12 +276,6 @@ class TrainingPipeline(LoRAPipeline, ABC): def _build_input_kwargs(self, training_batch: TrainingBatch) -> TrainingBatch: - assert self.training_args is not None - assert training_batch.noisy_model_input is not None - assert training_batch.encoder_hidden_states is not None - assert training_batch.encoder_attention_mask is not None - assert training_batch.timesteps is not None - training_batch.input_kwargs = { "hidden_states": training_batch.noisy_model_input, @@ -315,19 +293,10 @@ class TrainingPipeline(LoRAPipeline, ABC): def _transformer_forward_and_compute_loss( self, training_batch: TrainingBatch) -> TrainingBatch: - assert self.transformer is not None - assert self.training_args is not None - assert training_batch.noisy_model_input is not None - assert training_batch.latents is not None - assert training_batch.noise is not None - assert training_batch.sigmas is not None - if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN": assert training_batch.attn_metadata is not None else: assert training_batch.attn_metadata is None - - assert training_batch.input_kwargs is not None input_kwargs = training_batch.input_kwargs # if 'hunyuan' in self.training_args.model_type: @@ -341,7 +310,10 @@ class TrainingPipeline(LoRAPipeline, ABC): attn_metadata=training_batch.attn_metadata): model_pred = self.transformer(**input_kwargs) if self.training_args.precondition_outputs: + assert training_batch.sigmas is not None model_pred = training_batch.noisy_model_input - model_pred * training_batch.sigmas + assert training_batch.latents is not None + assert training_batch.noise is not None target = training_batch.latents if self.training_args.precondition_outputs else training_batch.noise - training_batch.latents # make sure no implicit broadcasting happens @@ -361,7 +333,6 @@ class TrainingPipeline(LoRAPipeline, ABC): return training_batch def _clip_grad_norm(self, training_batch: TrainingBatch) -> TrainingBatch: - assert self.training_args is not None max_grad_norm = self.training_args.max_grad_norm # TODO(will): perhaps move this into transformer api so that we can do @@ -383,8 +354,6 @@ class TrainingPipeline(LoRAPipeline, ABC): return training_batch def train_one_step(self, training_batch: TrainingBatch) -> TrainingBatch: - assert self.training_args is not None - training_batch = self._prepare_training(training_batch) for _ in range(self.training_args.gradient_accumulation_steps): @@ -423,7 +392,6 @@ class TrainingPipeline(LoRAPipeline, ABC): return training_batch def _resume_from_checkpoint(self) -> None: - assert self.training_args is not None logger.info("Loading checkpoint from %s", self.training_args.resume_from_checkpoint) resumed_step = load_checkpoint( @@ -439,12 +407,11 @@ class TrainingPipeline(LoRAPipeline, ABC): self.init_steps = 0 def train(self) -> None: - + assert self.seed is not None, "seed must be set" set_random_seed(self.seed) logger.info('rank: %s: start training', self.global_rank, local_main_process_only=False) - assert self.training_args is not None if not self.post_init_called: self.post_init() num_trainable_params = _get_trainable_params(self.transformer) @@ -550,9 +517,6 @@ class TrainingPipeline(LoRAPipeline, ABC): cleanup_dist_env_and_memory() def _log_training_info(self) -> None: - assert self.training_args is not None - 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 * @@ -593,6 +557,7 @@ class TrainingPipeline(LoRAPipeline, ABC): sampling_param.width = training_args.num_width sampling_param.num_inference_steps = num_inference_steps sampling_param.data_type = "video" + assert self.seed is not None sampling_param.seed = self.seed latents_size = [(sampling_param.num_frames - 1) // 4 + 1, @@ -618,7 +583,6 @@ class TrainingPipeline(LoRAPipeline, ABC): """ Generate a validation video and log it to wandb to check the quality during training. """ - assert training_args is not None training_args.inference_mode = True training_args.dit_cpu_offload = True if not training_args.log_validation: diff --git a/fastvideo/training/wan_i2v_distillation_pipeline.py b/fastvideo/training/wan_i2v_distillation_pipeline.py index 7dd8f57c9..6282a78bc 100644 --- a/fastvideo/training/wan_i2v_distillation_pipeline.py +++ b/fastvideo/training/wan_i2v_distillation_pipeline.py @@ -71,9 +71,6 @@ class WanI2VDistillationPipeline(DistillationPipeline): self.validation_pipeline = validation_pipeline def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch: - assert self.training_args is not None - assert self.train_dataloader is not None - batch = next(self.train_loader_iter, None) # type: ignore if batch is None: self.current_epoch += 1 @@ -118,6 +115,7 @@ class WanI2VDistillationPipeline(DistillationPipeline): sampling_param.image_path = validation_batch['video_path'] sampling_param.num_inference_steps = num_inference_steps sampling_param.data_type = "video" + assert self.seed is not None sampling_param.seed = self.seed latents_size = [(sampling_param.num_frames - 1) // 4 + 1, @@ -141,13 +139,6 @@ class WanI2VDistillationPipeline(DistillationPipeline): def _prepare_dit_inputs(self, training_batch: TrainingBatch) -> TrainingBatch: """Override to properly handle I2V concatenation - call parent first, then concatenate image conditioning.""" - assert self.training_args is not None - assert training_batch.latents is not None - assert training_batch.encoder_hidden_states is not None - assert training_batch.encoder_attention_mask is not None - assert self.noise_random_generator is not None - assert training_batch.image_latents is not None - # First, call parent method to prepare noise, timesteps, etc. for video latents training_batch = super()._prepare_dit_inputs(training_batch) @@ -189,11 +180,8 @@ class WanI2VDistillationPipeline(DistillationPipeline): def _build_distill_input_kwargs( self, noise_input: torch.Tensor, timestep: torch.Tensor, - text_dict: dict[str, torch.Tensor] | None, + text_dict: dict[str, torch.Tensor], training_batch: TrainingBatch) -> TrainingBatch: - assert training_batch.image_embeds is not None - assert training_batch.image_latents is not None - # Image Embeds for conditioning image_embeds = training_batch.image_embeds assert torch.isnan(image_embeds).sum() == 0 @@ -204,7 +192,7 @@ class WanI2VDistillationPipeline(DistillationPipeline): [noise_input, training_batch.image_latents.permute(0, 2, 1, 3, 4)], dim=2) - assert text_dict is not None + training_batch.input_kwargs = { "hidden_states": noisy_model_input.permute(0, 2, 1, 3, 4), # bs, c, t, h, w diff --git a/fastvideo/training/wan_i2v_training_pipeline.py b/fastvideo/training/wan_i2v_training_pipeline.py index 303ee20a4..fc5774c4c 100644 --- a/fastvideo/training/wan_i2v_training_pipeline.py +++ b/fastvideo/training/wan_i2v_training_pipeline.py @@ -62,9 +62,6 @@ class WanI2VTrainingPipeline(TrainingPipeline): dit_cpu_offload=True) def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch: - assert self.training_args is not None - assert self.train_dataloader is not None - batch = next(self.train_loader_iter, None) # type: ignore if batch is None: self.current_epoch += 1 @@ -102,12 +99,6 @@ class WanI2VTrainingPipeline(TrainingPipeline): def _prepare_dit_inputs(self, training_batch: TrainingBatch) -> TrainingBatch: """Override to properly handle I2V concatenation - call parent first, then concatenate image conditioning.""" - assert self.training_args is not None - assert training_batch.latents is not None - assert training_batch.encoder_hidden_states is not None - assert training_batch.encoder_attention_mask is not None - assert self.noise_random_generator is not None - assert training_batch.image_latents is not None # First, call parent method to prepare noise, timesteps, etc. for video latents training_batch = super()._prepare_dit_inputs(training_batch) @@ -144,12 +135,6 @@ class WanI2VTrainingPipeline(TrainingPipeline): def _build_input_kwargs(self, training_batch: TrainingBatch) -> TrainingBatch: - assert self.training_args is not None - assert training_batch.noisy_model_input is not None - assert training_batch.encoder_hidden_states is not None - assert training_batch.encoder_attention_mask is not None - assert training_batch.timesteps is not None - assert training_batch.image_embeds is not None # Image Embeds for conditioning image_embeds = training_batch.image_embeds @@ -186,6 +171,7 @@ class WanI2VTrainingPipeline(TrainingPipeline): sampling_param.image_path = validation_batch['video_path'] sampling_param.num_inference_steps = num_inference_steps sampling_param.data_type = "video" + assert self.seed is not None sampling_param.seed = self.seed latents_size = [(sampling_param.num_frames - 1) // 4 + 1, diff --git a/fastvideo/worker/multiproc_executor.py b/fastvideo/worker/multiproc_executor.py index 1aaca215b..1d770d322 100644 --- a/fastvideo/worker/multiproc_executor.py +++ b/fastvideo/worker/multiproc_executor.py @@ -31,15 +31,23 @@ class MultiprocExecutor(Executor): self.workers: list[BaseProcess] = [] self.worker_pipes = [] - self.master_port = None - for port in range(29503, 65535): - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - if s.connect_ex(('localhost', port)) != 0: - self.master_port = port - break - if self.master_port is None: - raise ValueError("No unused port found to use as master port") + # Check if master_port is provided in fastvideo_args + if hasattr( + self.fastvideo_args, + 'master_port') and self.fastvideo_args.master_port is not None: + self.master_port = self.fastvideo_args.master_port + logger.info("Using provided master port: %s", self.master_port) + else: + # Auto-find available port + for port in range(29503, 65535): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + if s.connect_ex(('localhost', port)) != 0: + self.master_port = port + break + else: + raise ValueError("No unused port found to use as master port") + logger.info("Auto-selected master port: %s", self.master_port) # Create pipes and start workers for rank in range(self.world_size):