nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge branch 'main' of github.com:hao-ai-lab/FastVideo

Browse Source
This commit is contained in:
“BrianChen1129”
2025-07-27 20:02:15 +00:00
parent 5c94cbad94 636524e87f
commit ebae1321cf
8 changed files with 50 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.pre-commit-config.yaml
+1 -1
View File
@@ -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:
fastvideo/configs/pipelines/registry.py
+1
View File
@@ -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
}
fastvideo/fastvideo_args.py
+11
View File
@@ -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",
fastvideo/training/distillation_pipeline.py
+9 -16
View File
@@ -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
fastvideo/training/training_pipeline.py
+8 -44
View File
@@ -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:
fastvideo/training/wan_i2v_distillation_pipeline.py
+3 -15
View File
@@ -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
fastvideo/training/wan_i2v_training_pipeline.py
+1 -15
View File
@@ -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,
fastvideo/worker/multiproc_executor.py
+16 -8
View File
@@ -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):