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

[Training] Use inference pipeline for training validation (#585)

Browse Source
This commit is contained in:
William Lin
2025-07-04 17:23:55 -05:00
committed by GitHub
parent ed1e8d6bad
commit 66fdcc8e76
34 changed files with 540 additions and 242 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/README.md
+10
View File
@@ -0,0 +1,10 @@
This directory contain e2e examples scripts for finetuning Wan2.1 I2V.
Execute the following commands from `FastVideo/` to run training:
- Download crush-smol dataset:
`bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/download_dataset.sh`
- Preprocess the videos and captions into latents:
`bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/preprocess_wan_data_i2v.sh`
- Edit the following file and run finetuning:
`bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.sh`
examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/download_dataset.sh
+3
View File
@@ -0,0 +1,3 @@
#!/bin/bash
python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"
examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.sh
+93
View File
@@ -0,0 +1,93 @@
#!/bin/bash
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
export TOKENIZERS_PARALLELISM=false
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
MODEL_PATH="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers"
DATA_DIR="data/crush-smol_processed_i2v_1_3b_inp/combined_parquet_dataset/"
VALIDATION_DATASET_FILE="examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json"
NUM_GPUS=8
# export CUDA_VISIBLE_DEVICES=4,5
# IP=[MASTER NODE IP]
# Training arguments
training_args=(
--tracker_project_name "wan_i2v_finetune"
--output_dir "$DATA_DIR/outputs/wan_i2v_finetune"
--max_train_steps 2000
--train_batch_size 4
--train_sp_batch_size 1
--gradient_accumulation_steps 1
--num_latent_t 8
--num_height 480
--num_width 832
--num_frames 77
)
# Parallel arguments
parallel_args=(
--num_gpus $NUM_GPUS
--sp_size 4
--tp_size 4
--hsdp_replicate_dim 2
--hsdp_shard_dim 4
)
# Model arguments
model_args=(
--model_path $MODEL_PATH
--pretrained_model_name_or_path $MODEL_PATH
)
# Dataset arguments
dataset_args=(
--data_path "$DATA_DIR"
--dataloader_num_workers 1
)
# Validation arguments
validation_args=(
--log_validation
--validation_dataset_file "$VALIDATION_DATASET_FILE"
--validation_steps 100
--validation_sampling_steps "40"
--validation_guidance_scale "1.0"
)
# Optimizer arguments
optimizer_args=(
--learning_rate 2e-5
--mixed_precision "bf16"
--checkpointing_steps 2000
--weight_decay 1e-4
--max_grad_norm 1.0
)
# Miscellaneous arguments
miscellaneous_args=(
--inference_mode False
--allow_tf32
--checkpoints_total_limit 3
--training_cfg_rate 0.1
--multi_phased_distill_schedule "4000-1"
--not_apply_cfg_solver
--dit_precision "fp32"
--num_euler_timesteps 50
--ema_start_step 0
--enable_gradient_checkpointing_type "full"
)
# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
torchrun \
--nnodes 1 \
--nproc_per_node $NUM_GPUS \
fastvideo/v1/training/wan_i2v_training_pipeline.py \
"${parallel_args[@]}" \
"${model_args[@]}" \
"${dataset_args[@]}" \
"${training_args[@]}" \
"${optimizer_args[@]}" \
"${validation_args[@]}" \
"${miscellaneous_args[@]}"
examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.slurm
+131
View File
@@ -0,0 +1,131 @@
#!/bin/bash
#SBATCH --job-name=i2v
#SBATCH --partition=main
#SBATCH --qos=hao
#SBATCH --nodes=4
#SBATCH --ntasks=4
#SBATCH --ntasks-per-node=1
#SBATCH --gres=gpu:8
#SBATCH --cpus-per-task=128
#SBATCH --nodelist=fs-mbz-gpu-[100-850]
#SBATCH --mem=1440G
#SBATCH --output=i2v_output/i2v_%j.out
#SBATCH --error=i2v_output/i2v_%j.err
#SBATCH --exclusive
set -e -x
# Environment Setup
source ~/conda/miniconda/bin/activate
conda activate will-fv
# Basic Info
export WANDB_MODE="online"
export NCCL_P2P_DISABLE=1
export TORCH_NCCL_ENABLE_MONITORING=0
# different cache dir for different processes
export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
export MASTER_PORT=29500
export NODE_RANK=$SLURM_PROCID
nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
export MASTER_ADDR=${nodes[0]}
export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
export TOKENIZERS_PARALLELISM=false
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
echo "MASTER_ADDR: $MASTER_ADDR"
echo "NODE_RANK: $NODE_RANK"
# Configs
NUM_GPUS=8
MODEL_PATH="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers"
DATA_DIR="data/crush-smol_processed_i2v_1_3b_inp/combined_parquet_dataset/"
VALIDATION_DATASET_FILE="examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json"
# export CUDA_VISIBLE_DEVICES=4,5
# IP=[MASTER NODE IP]
# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
# Training arguments
training_args=(
--tracker_project_name wan_i2v_finetune
--output_dir="$DATA_DIR/outputs/wan_i2v_finetune_2n"
--max_train_steps=2000
--train_batch_size=2
--train_sp_batch_size 1
--gradient_accumulation_steps=1
--num_latent_t 8
--num_height 480
--num_width 832
--num_frames 77
)
# Parallel arguments
parallel_args=(
--num_gpus $NUM_GPUS
--sp_size $NUM_GPUS
--tp_size $NUM_GPUS
--hsdp_replicate_dim $SLURM_JOB_NUM_NODES
--hsdp_shard_dim $NUM_GPUS
)
# Model arguments
model_args=(
--model_path $MODEL_PATH
--pretrained_model_name_or_path $MODEL_PATH
)
# Dataset arguments
dataset_args=(
--data_path "$DATA_DIR"
--dataloader_num_workers 10
)
# Validation arguments
validation_args=(
--log_validation
--validation_dataset_file "$VALIDATION_DATASET_FILE"
--validation_steps 100
--validation_sampling_steps "40"
--validation_guidance_scale "1.0"
)
# Optimizer arguments
optimizer_args=(
--learning_rate=1e-5
--mixed_precision="bf16"
--checkpointing_steps=1000
--weight_decay 1e-4
--max_grad_norm 1.0
)
# Miscellaneous arguments
miscellaneous_args=(
--inference_mode False
--allow_tf32
--checkpoints_total_limit 3
--training_cfg_rate 0.1
--multi_phased_distill_schedule "4000-1"
--not_apply_cfg_solver
--dit_precision "fp32"
--num_euler_timesteps 50
--ema_start_step 0
--enable_gradient_checkpointing_type "full"
)
srun torchrun \
--nnodes $SLURM_JOB_NUM_NODES \
--nproc_per_node $NUM_GPUS \
--node_rank $SLURM_PROCID \
--rdzv_backend=c10d \
--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
fastvideo/v1/training/wan_i2v_training_pipeline.py \
"${parallel_args[@]}" \
"${model_args[@]}" \
"${dataset_args[@]}" \
"${training_args[@]}" \
"${optimizer_args[@]}" \
"${validation_args[@]}" \
"${miscellaneous_args[@]}"
examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/preprocess_wan_data_i2v.sh
+26
View File
@@ -0,0 +1,26 @@
#!/bin/bash
GPU_NUM=1 # 2,4,8
MODEL_PATH="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers"
MODEL_TYPE="wan"
DATA_MERGE_PATH="data/crush-smol/merge.txt"
OUTPUT_DIR="data/crush-smol_processed_i2v_1_3b_inp/"
VALIDATION_PATH="examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json"
torchrun --nproc_per_node=$GPU_NUM \
fastvideo/v1/pipelines/preprocess/v1_preprocess.py \
--model_path $MODEL_PATH \
--data_merge_path $DATA_MERGE_PATH \
--preprocess_video_batch_size 8 \
--max_height 480 \
--max_width 832 \
--num_frames 77 \
--dataloader_num_workers 0 \
--output_dir=$OUTPUT_DIR \
--model_type $MODEL_TYPE \
--train_fps 16 \
--validation_dataset_file $VALIDATION_PATH \
--samples_per_file 8 \
--flush_frequency 8 \
--video_length_tolerance_range 5 \
--preprocess_task "i2v"
examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json
+31
View File
@@ -0,0 +1,31 @@
{
"data": [
{
"caption": "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.",
"image_path": null,
"video_path": "validation_dataset/yYcK4nANZz4-Scene-034.mp4",
"num_inference_steps": 40,
"height": 480,
"width": 832,
"num_frames": 77
},
{
"caption": "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.",
"image_path": null,
"video_path": "validation_dataset/yYcK4nANZz4-Scene-027.mp4",
"num_inference_steps": 40,
"height": 480,
"width": 832,
"num_frames": 77
},
{
"caption": "A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table.",
"image_path": null,
"video_path": "validation_dataset/yYcK4nANZz4-Scene-030.mp4",
"num_inference_steps": 40,
"height": 480,
"width": 832,
"num_frames": 77
}
]
}
examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation_dataset/yYcK4nANZz4-Scene-027.mp4
BIN
View File
Binary file not shown.
examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation_dataset/yYcK4nANZz4-Scene-030.mp4
BIN
View File
Binary file not shown.
examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation_dataset/yYcK4nANZz4-Scene-034.mp4
BIN
View File
Binary file not shown.
examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.sh
+4 -2
View File
@@ -2,11 +2,12 @@
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
export TOKENIZERS_PARALLELISM=false
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
DATA_DIR="data/crush-smol_processed_i2v/combined_parquet_dataset/"
VALIDATION_DIR="data/crush-smol_processed_i2v/validation_parquet_dataset/"
VALIDATION_DATASET_FILE="examples/training/finetune/wan_i2v_14b_480p/crush_smol/validation.json"
NUM_GPUS=8
# export CUDA_VISIBLE_DEVICES=4,5
# IP=[MASTER NODE IP]
@@ -49,7 +50,7 @@ dataset_args=(
# Validation arguments
validation_args=(
--log_validation
--validation_preprocessed_path "$VALIDATION_DIR"
--validation_dataset_file "$VALIDATION_DATASET_FILE"
--validation_steps 100
--validation_sampling_steps "40"
--validation_guidance_scale "1.0"
@@ -75,6 +76,7 @@ miscellaneous_args=(
--dit_precision "fp32"
--num_euler_timesteps 50
--ema_start_step 0
--enable_gradient_checkpointing_type "full"
)
# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.slurm
+3 -2
View File
@@ -42,7 +42,7 @@ NUM_GPUS=8
MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
DATA_DIR="data/crush-smol_processed_i2v/combined_parquet_dataset/"
VALIDATION_DIR="data/crush-smol_processed_i2v/validation_parquet_dataset/"
VALIDATION_DATASET_FILE="examples/training/finetune/wan_i2v_14b_480p/crush_smol/validation.json"
# export CUDA_VISIBLE_DEVICES=4,5
# IP=[MASTER NODE IP]
@@ -86,7 +86,7 @@ dataset_args=(
# Validation arguments
validation_args=(
--log_validation
--validation_preprocessed_path "$VALIDATION_DIR"
--validation_dataset_file "$VALIDATION_DATASET_FILE"
--validation_steps 100
--validation_sampling_steps "40"
--validation_guidance_scale "1.0"
@@ -112,6 +112,7 @@ miscellaneous_args=(
--dit_precision "fp32"
--num_euler_timesteps 50
--ema_start_step 0
--enable_gradient_checkpointing_type "full"
)
srun torchrun \
examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.sh
+4 -2
View File
@@ -2,11 +2,12 @@
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
export TOKENIZERS_PARALLELISM=false
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/"
VALIDATION_DIR="data/crush-smol_processed_t2v/validation_parquet_dataset/"
VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1_3b/crush_smol/validation.json"
NUM_GPUS=4
# export CUDA_VISIBLE_DEVICES=4,5
@@ -49,7 +50,7 @@ dataset_args=(
# Validation arguments
validation_args=(
--log_validation
--validation_preprocessed_path $VALIDATION_DIR
--validation_dataset_file $VALIDATION_DATASET_FILE
--validation_steps 50
--validation_sampling_steps "50"
--validation_guidance_scale "1.0"
@@ -75,6 +76,7 @@ miscellaneous_args=(
--dit_precision "fp32"
--num_euler_timesteps 50
--ema_start_step 0
--enable_gradient_checkpointing_type "full"
)
torchrun \
examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.slurm
+3 -2
View File
@@ -41,7 +41,7 @@ echo "NODE_RANK: $NODE_RANK"
NUM_GPUS=8
MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/"
VALIDATION_DIR="data/crush-smol_processed_t2v/validation_parquet_dataset/"
VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1_3b/crush_smol/validation.json"
# export CUDA_VISIBLE_DEVICES=4,5
# IP=[MASTER NODE IP]
@@ -83,7 +83,7 @@ dataset_args=(
# Validation arguments
validation_args=(
--log_validation
--validation_preprocessed_path "$VALIDATION_DIR"
--validation_dataset_file "$VALIDATION_DATASET_FILE"
--validation_steps 100
--validation_sampling_steps "50"
--validation_guidance_scale "1.0"
@@ -109,6 +109,7 @@ miscellaneous_args=(
--dit_precision "fp32"
--num_euler_timesteps 50
--ema_start_step 0
--enable_gradient_checkpointing_type "full"
)
srun torchrun \
fastvideo/v1/configs/pipelines/registry.py
+1
View File
@@ -23,6 +23,7 @@ PIPE_NAME_TO_CONFIG: Dict[str, Type[PipelineConfig]] = {
"FastVideo/FastHunyuan-diffusers": FastHunyuanConfig,
"hunyuanvideo-community/HunyuanVideo": HunyuanConfig,
"Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V480PConfig,
"weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers": WanI2V480PConfig,
"Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V480PConfig,
"Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V720PConfig,
"Wan-AI/Wan2.1-T2V-14B-Diffusers": WanT2V720PConfig,
fastvideo/v1/dataset/__init__.py
+4 -1
View File
@@ -12,7 +12,10 @@ from fastvideo.v1.dataset.validation_dataset import ValidationDataset
def getdataset(args) -> VideoCaptionMergedDataset:
temporal_sample = TemporalRandomCrop(args.num_frames) # 16 x
if args.do_temporal_sample:
temporal_sample = TemporalRandomCrop(args.num_frames) # 16 x
else:
temporal_sample = None
norm_fun = Lambda(lambda x: 2.0 * x - 1.0)
resize_topcrop = [
CenterCropResizeVideo((args.max_height, args.max_width), top_crop=True),
fastvideo/v1/dataset/preprocessing_datasets.py
+6 -4
View File
@@ -281,10 +281,12 @@ class FrameSamplingStage(DatasetFilterStage):
frame_interval).astype(int)
# Temporal crop if too long
if len(frame_indices
) > self.num_frames and temporal_sample_fn is not None:
begin_index, end_index = temporal_sample_fn(len(frame_indices))
frame_indices = frame_indices[begin_index:end_index]
if len(frame_indices) > self.num_frames:
if temporal_sample_fn is not None:
begin_index, end_index = temporal_sample_fn(len(frame_indices))
frame_indices = frame_indices[begin_index:end_index]
else:
frame_indices = frame_indices[:self.num_frames]
batch.sample_frame_index = frame_indices.tolist()
batch.sample_num_frames = len(frame_indices)
fastvideo/v1/dataset/validation_dataset.py
+65 -4
View File
@@ -4,15 +4,17 @@ import os
import pathlib
import datasets
import torch
from torch.utils.data import IterableDataset
from fastvideo.v1.distributed import (get_sp_world_size, get_world_rank,
get_world_size)
from fastvideo.v1.logger import init_logger
from fastvideo.v1.models.vision_utils import load_image, load_video
logger = init_logger(__name__)
class ValidationDataset(torch.utils.data.IterableDataset):
class ValidationDataset(IterableDataset):
def __init__(self, filename: str):
super().__init__()
@@ -48,10 +50,65 @@ class ValidationDataset(torch.utils.data.IterableDataset):
f"Unsupported file format {self.filename.suffix} for validation dataset. Supported formats are: {_SUPPORTED_FILE_FORMATS}"
)
self._data = data.to_iterable_dataset()
# Get distributed training info
self.global_rank = get_world_rank()
self.world_size = get_world_size()
self.sp_world_size = get_sp_world_size()
self.num_sp_groups = self.world_size // self.sp_world_size
# Convert to list to get total samples
self.all_samples = list(data)
self.original_total_samples = len(self.all_samples)
# Extend samples to be a multiple of DP degree (num_sp_groups)
remainder = self.original_total_samples % self.num_sp_groups
if remainder != 0:
samples_to_add = self.num_sp_groups - remainder
# Duplicate samples cyclically to reach the target
additional_samples = []
for i in range(samples_to_add):
additional_samples.append(
self.all_samples[i % self.original_total_samples])
self.all_samples.extend(additional_samples)
self.total_samples = len(self.all_samples)
# Calculate which SP group this rank belongs to
self.sp_group_id = self.global_rank // self.sp_world_size
# Now all SP groups will have equal number of samples
self.samples_per_sp_group = self.total_samples // self.num_sp_groups
# Calculate start and end indices for this SP group
self.start_idx = self.sp_group_id * self.samples_per_sp_group
self.end_idx = self.start_idx + self.samples_per_sp_group
# Get samples for this SP group
self.sp_group_samples = self.all_samples[self.start_idx:self.end_idx]
logger.info(
"Rank %s (SP group %s): "
"Original samples: %s, "
"Extended samples: %s, "
"SP group samples: %s, "
"Range: [%s:%s]",
self.global_rank,
self.sp_group_id,
self.original_total_samples,
self.total_samples,
len(self.sp_group_samples),
self.start_idx,
self.end_idx,
local_main_process_only=False)
def __len__(self):
"""Return the number of samples for this SP group."""
return len(self.sp_group_samples)
def __iter__(self):
for sample in self._data:
for sample in self.sp_group_samples:
# For consistency reasons, we mandate that "caption" is always present in the validation dataset.
# However, since the model specifications use "prompt", we create an alias here.
sample["prompt"] = sample["caption"]
@@ -64,6 +121,7 @@ class ValidationDataset(torch.utils.data.IterableDataset):
if sample.get("image_path", None) is not None:
image_path = sample["image_path"]
image_path = os.path.join(self.dir, image_path)
sample["image_path"] = image_path
if not pathlib.Path(image_path).is_file(
) and not image_path.startswith("http"):
logger.warning("Image file %s does not exist.", image_path)
@@ -73,6 +131,7 @@ class ValidationDataset(torch.utils.data.IterableDataset):
if sample.get("video_path", None) is not None:
video_path = sample["video_path"]
video_path = os.path.join(self.dir, video_path)
sample["video_path"] = video_path
if not pathlib.Path(video_path).is_file(
) and not video_path.startswith("http"):
logger.warning("Video file %s does not exist.", video_path)
@@ -82,6 +141,7 @@ class ValidationDataset(torch.utils.data.IterableDataset):
if sample.get("control_image_path", None) is not None:
control_image_path = sample["control_image_path"]
control_image_path = os.path.join(self.dir, control_image_path)
sample["control_image_path"] = control_image_path
if not pathlib.Path(control_image_path).is_file(
) and not control_image_path.startswith("http"):
logger.warning("Control Image file %s does not exist.",
@@ -92,6 +152,7 @@ class ValidationDataset(torch.utils.data.IterableDataset):
if sample.get("control_video_path", None) is not None:
control_video_path = sample["control_video_path"]
control_video_path = os.path.join(self.dir, control_video_path)
sample["control_video_path"] = control_video_path
if not pathlib.Path(control_video_path).is_file(
) and not control_video_path.startswith("http"):
logger.warning("Control Video file %s does not exist.",
fastvideo/v1/distributed/parallel_state.py
+7 -3
View File
@@ -750,6 +750,7 @@ def init_distributed_environment(
distributed_init_method: str = "env://",
local_rank: int = 0,
backend: str = "nccl",
device_id: Optional[torch.device] = None,
):
logger.debug(
"world_size=%d rank=%d local_rank=%d "
@@ -764,7 +765,8 @@ def init_distributed_environment(
backend=backend,
init_method=distributed_init_method,
world_size=world_size,
rank=rank)
rank=rank,
device_id=device_id)
# set the local rank
# local_rank is not available in torch ProcessGroup,
# see https://github.com/pytorch/pytorch/issues/122816
@@ -921,13 +923,15 @@ def maybe_init_distributed_environment_and_model_parallel(
local_rank = int(os.environ.get("LOCAL_RANK", 0))
world_size = int(os.environ.get("WORLD_SIZE", 1))
rank = int(os.environ.get("RANK", 0))
device = torch.device(f"cuda:{local_rank}")
torch.cuda.set_device(device)
torch.cuda.set_device(local_rank)
init_distributed_environment(
world_size=world_size,
rank=rank,
local_rank=local_rank,
distributed_init_method=distributed_init_method)
distributed_init_method=distributed_init_method,
device_id=device)
initialize_model_parallel(tensor_model_parallel_size=tp_size,
sequence_model_parallel_size=sp_size)
fastvideo/v1/pipelines/preprocess/preprocess_pipeline_i2v.py
+22 -23
View File
@@ -89,11 +89,10 @@ class PreprocessPipeline_I2V(BasePreprocessPipeline):
"""Get CLIP features from the first frame of each video."""
first_frame = valid_data["pixel_values"][:, :, 0, :, :].permute(
0, 2, 3, 1) # (B, C, T, H, W) -> (B, H, W, C)
batch_size, _, num_frames, height, width = valid_data[
"pixel_values"].shape
latent_height = height // self.get_module(
"vae").spatial_compression_ratio
latent_width = width // self.get_module("vae").spatial_compression_ratio
_, _, num_frames, height, width = valid_data["pixel_values"].shape
# latent_height = height // self.get_module(
# "vae").spatial_compression_ratio
# latent_width = width // self.get_module("vae").spatial_compression_ratio
processed_images = []
# Frame has values between -1 and 1
@@ -157,26 +156,26 @@ class PreprocessPipeline_I2V(BasePreprocessPipeline):
latent_condition = latent_condition * self.get_module(
"vae").scaling_factor
mask_lat_size = torch.ones(batch_size, 1, num_frames, latent_height,
latent_width)
mask_lat_size[:, :, list(range(1, num_frames))] = 0
first_frame_mask = mask_lat_size[:, :, 0:1]
first_frame_mask = torch.repeat_interleave(
first_frame_mask,
dim=2,
repeats=self.get_module("vae").temporal_compression_ratio)
mask_lat_size = torch.concat(
[first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2)
mask_lat_size = mask_lat_size.view(
batch_size, -1,
self.get_module("vae").temporal_compression_ratio, latent_height,
latent_width)
mask_lat_size = mask_lat_size.transpose(1, 2)
mask_lat_size = mask_lat_size.to(latent_condition.device)
# mask_lat_size = torch.ones(batch_size, 1, num_frames, latent_height,
# latent_width)
# mask_lat_size[:, :, list(range(1, num_frames))] = 0
# first_frame_mask = mask_lat_size[:, :, 0:1]
# first_frame_mask = torch.repeat_interleave(
# first_frame_mask,
# dim=2,
# repeats=self.get_module("vae").temporal_compression_ratio)
# mask_lat_size = torch.concat(
# [first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2)
# mask_lat_size = mask_lat_size.view(
# batch_size, -1,
# self.get_module("vae").temporal_compression_ratio, latent_height,
# latent_width)
# mask_lat_size = mask_lat_size.transpose(1, 2)
# mask_lat_size = mask_lat_size.to(latent_condition.device)
image_latent = torch.concat([mask_lat_size, latent_condition], dim=1)
# image_latent = torch.concat([mask_lat_size, latent_condition], dim=1)
features["first_frame_latent"] = image_latent
features["first_frame_latent"] = latent_condition
return features
fastvideo/v1/pipelines/preprocess/v1_preprocess.py
+3
View File
@@ -79,6 +79,9 @@ if __name__ == "__main__":
parser.add_argument("--text_max_length", type=int, default=256)
parser.add_argument("--speed_factor", type=float, default=1.0)
parser.add_argument("--drop_short_ratio", type=float, default=1.0)
parser.add_argument("--do_temporal_sample",
default=False,
action="store_true")
# text encoder & vae & diffusion model
parser.add_argument("--text_encoder_name",
type=str,
fastvideo/v1/pipelines/stages/denoising.py
-4
View File
@@ -300,10 +300,6 @@ class DenoisingStage(PipelineStage):
if st_attn_available and self.attn_backend == SlidingTileAttentionBackend and fastvideo_args.STA_mode == STA_Mode.STA_SEARCHING:
self.save_sta_search_results(batch)
if fastvideo_args.use_cpu_offload:
self.transformer.to('cpu')
torch.cuda.empty_cache()
return batch
def prepare_extra_func_kwargs(self, func, kwargs) -> Dict[str, Any]:
fastvideo/v1/pipelines/stages/encoding.py
+12 -16
View File
@@ -56,23 +56,16 @@ class EncodingStage(PipelineStage):
latent_height = batch.height // self.vae.spatial_compression_ratio
latent_width = batch.width // self.vae.spatial_compression_ratio
image = batch.preprocessed_image
# TODO(will)
if image is None:
assert batch.pil_image is not None
image = batch.pil_image
image = self.preprocess(
image,
vae_scale_factor=self.vae.spatial_compression_ratio,
height=batch.height,
width=batch.width).to(get_local_torch_device(),
dtype=torch.float32)
assert batch.pil_image is not None
image = batch.pil_image
image = self.preprocess(
image,
vae_scale_factor=self.vae.spatial_compression_ratio,
height=batch.height,
width=batch.width).to(get_local_torch_device(), dtype=torch.float32)
image = image.unsqueeze(2)
image = image.unsqueeze(2)
else:
# assumes image is loaded from parquet file and used for validation
image = image.transpose(1, 2)
logger.info("image: %s", image.shape)
video_condition = torch.cat([
image,
image.new_zeros(image.shape[0], image.shape[1],
@@ -143,6 +136,9 @@ class EncodingStage(PipelineStage):
if hasattr(self, 'maybe_free_model_hooks'):
self.maybe_free_model_hooks()
self.vae.to("cpu")
torch.cuda.empty_cache()
return batch
def retrieve_latents(self,
fastvideo/v1/pipelines/stages/image_encoding.py
+1 -2
View File
@@ -54,8 +54,7 @@ class ImageEncodingStage(PipelineStage):
Returns:
The batch with encoded prompt embeddings.
"""
if fastvideo_args.use_cpu_offload:
self.image_encoder = self.image_encoder.to(get_local_torch_device())
self.image_encoder = self.image_encoder.to(get_local_torch_device())
image = batch.pil_image
fastvideo/v1/pipelines/stages/input_validation.py
+5 -2
View File
@@ -7,7 +7,7 @@ import torch
from fastvideo.v1.fastvideo_args import FastVideoArgs
from fastvideo.v1.logger import init_logger
from fastvideo.v1.models.vision_utils import load_image
from fastvideo.v1.models.vision_utils import load_image, load_video
from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
from fastvideo.v1.pipelines.stages.base import PipelineStage
from fastvideo.v1.pipelines.stages.validators import (StageValidators,
@@ -94,7 +94,10 @@ class InputValidationStage(PipelineStage):
# for i2v, get image from image_path
if batch.image_path is not None:
image = load_image(batch.image_path)
if batch.image_path.endswith(".mp4"):
image = load_video(batch.image_path)[0]
else:
image = load_image(batch.image_path)
batch.pil_image = image
return batch
fastvideo/v1/pipelines/wan/wan_i2v_pipeline.py
-32
View File
@@ -76,36 +76,4 @@ class WanImageToVideoPipeline(LoRAPipeline, ComposedPipelineBase):
stage=DecodingStage(vae=self.get_module("vae")))
class WanImageToVideoValidationPipeline(ComposedPipelineBase):
"""
I2V Validation pipeline for Wan2.1, assumes that the input are preprocess latents.
"""
_required_config_modules = ["vae", "scheduler", "transformer"]
def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
self.modules["scheduler"] = FlowUniPCMultistepScheduler(
shift=fastvideo_args.pipeline_config.flow_shift)
def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
self.add_stage(stage_name="timestep_preparation_stage",
stage=TimestepPreparationStage(
scheduler=self.get_module("scheduler")))
self.add_stage(stage_name="latent_preparation_stage",
stage=LatentPreparationStage(
scheduler=self.get_module("scheduler"),
transformer=self.get_module("transformer")))
self.add_stage(stage_name="image_latent_preparation_stage",
stage=EncodingStage(vae=self.get_module("vae")))
self.add_stage(stage_name="denoising_stage",
stage=DenoisingStage(
transformer=self.get_module("transformer"),
scheduler=self.get_module("scheduler")))
self.add_stage(stage_name="decoding_stage",
stage=DecodingStage(vae=self.get_module("vae")))
EntryClass = WanImageToVideoPipeline
fastvideo/v1/pipelines/wan/wan_pipeline.py
-30
View File
@@ -67,34 +67,4 @@ class WanPipeline(LoRAPipeline, ComposedPipelineBase):
stage=DecodingStage(vae=self.get_module("vae")))
class WanValidationPipeline(ComposedPipelineBase):
"""
Validation pipeline for Wan2.1, assumes that the input are preprocess latents.
"""
_required_config_modules = ["vae", "scheduler", "transformer"]
def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
self.modules["scheduler"] = FlowUniPCMultistepScheduler(
shift=fastvideo_args.pipeline_config.flow_shift)
def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
"""Set up pipeline stages with proper dependency injection."""
self.add_stage(stage_name="timestep_preparation_stage",
stage=TimestepPreparationStage(
scheduler=self.get_module("scheduler")))
self.add_stage(stage_name="latent_preparation_stage",
stage=LatentPreparationStage(
scheduler=self.get_module("scheduler"),
transformer=self.get_module("transformer", None)))
self.add_stage(stage_name="denoising_stage",
stage=DenoisingStage(
transformer=self.get_module("transformer"),
scheduler=self.get_module("scheduler")))
self.add_stage(stage_name="decoding_stage",
stage=DecodingStage(vae=self.get_module("vae")))
EntryClass = WanPipeline
fastvideo/v1/tests/training/VSA/test_training_loss_VSA.py
+1 -1
View File
@@ -32,7 +32,7 @@ def run_worker():
"--inference_mode", "False",
"--pretrained_model_name_or_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
"--data_path", "data/mini_dataset_i2v_VSA/combined_parquet_dataset",
"--validation_preprocessed_path", "data/mini_dataset_i2v_VSA/validation_parquet_dataset",
"--validation_dataset_file", "examples/training/finetune/wan_t2v_1_3b/crush_smol/validation.json",
"--train_batch_size", "1",
"--num_latent_t", "4",
"--num_gpus", "2",
fastvideo/v1/tests/training/Vanilla/test_training_loss.py
+1 -1
View File
@@ -38,7 +38,7 @@ def run_worker():
"--inference_mode", "False",
"--pretrained_model_name_or_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
"--data_path", "data/crush-smol_processed_t2v/combined_parquet_dataset",
"--validation_preprocessed_path", "data/crush-smol_processed_t2v/validation_parquet_dataset",
"--validation_dataset_file", "examples/training/finetune/wan_t2v_1_3b/crush_smol/validation.json",
"--train_batch_size", "2",
"--num_latent_t", "4",
"--num_gpus", "4",
fastvideo/v1/training/training_pipeline.py
+51 -62
View File
@@ -14,6 +14,7 @@ import torchvision
from diffusers import FlowMatchEulerDiscreteScheduler
from diffusers.optimization import get_scheduler
from einops import rearrange
from torch.utils.data import DataLoader
from torchdata.stateful_dataloader import StatefulDataLoader
from tqdm.auto import tqdm
@@ -24,6 +25,7 @@ from fastvideo.v1.configs.sample import SamplingParam
from fastvideo.v1.dataset import build_parquet_map_style_dataloader
from fastvideo.v1.dataset.dataloader.schema import (
pyarrow_schema_t2v, pyarrow_schema_t2v_validation)
from fastvideo.v1.dataset.validation_dataset import ValidationDataset
from fastvideo.v1.distributed import (cleanup_dist_env_and_memory,
get_local_torch_device, get_sp_group,
get_world_group)
@@ -38,7 +40,7 @@ from fastvideo.v1.training.training_utils import (
clip_grad_norm_while_handling_failing_dtensor_cases,
compute_density_for_timestep_sampling, get_sigmas, load_checkpoint,
normalize_dit_input, save_checkpoint, shard_latents_across_sp)
from fastvideo.v1.utils import is_vsa_available, set_random_seed
from fastvideo.v1.utils import is_vsa_available, set_random_seed, shallow_asdict
import wandb # isort: skip
@@ -416,6 +418,9 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
self.init_steps = 0
def train(self) -> None:
logger.info('rank: %s: start training',
self.global_rank,
local_main_process_only=False)
assert self.training_args is not None
# Set random seeds for deterministic training
@@ -548,60 +553,40 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
logger.info("VSA validation sparsity: %s",
self.training_args.VSA_sparsity)
def _prepare_validation_inputs(
self, sampling_param: SamplingParam, training_args: TrainingArgs,
validation_batch: Dict[str, Any], num_inference_steps: int,
negative_prompt_embeds: torch.Tensor | None,
negative_prompt_attention_mask: torch.Tensor | None
) -> ForwardBatch:
def _prepare_validation_batch(self, sampling_param: SamplingParam,
training_args: TrainingArgs,
validation_batch: Dict[str, Any],
num_inference_steps: int) -> ForwardBatch:
sampling_param.prompt = validation_batch['prompt']
sampling_param.height = training_args.num_height
sampling_param.width = training_args.num_width
sampling_param.num_inference_steps = num_inference_steps
sampling_param.data_type = "video"
sampling_param.seed = self.seed
assert len(validation_batch['info_list']
) == 1, "Only batch size 1 is supported for validation"
prompt = validation_batch['info_list'][0]['prompt']
prompt_embeds = validation_batch['text_embedding']
prompt_attention_mask = validation_batch['text_attention_mask']
prompt_embeds = prompt_embeds.to(get_local_torch_device())
prompt_attention_mask = prompt_attention_mask.to(
get_local_torch_device())
# Calculate sizes
latents_size = [(sampling_param.num_frames - 1) // 4 + 1,
sampling_param.height // 8, sampling_param.width // 8]
n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
temporal_compression_factor = training_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio
num_frames = (training_args.num_latent_t -
1) * temporal_compression_factor + 1
# Prepare batch for validation
sampling_param.num_frames = num_frames
batch = ForwardBatch(
prompt=prompt,
data_type="video",
**shallow_asdict(sampling_param),
latents=None,
seed=self.seed, # Use deterministic seed
generator=torch.Generator(device="cpu").manual_seed(self.seed),
prompt_embeds=[prompt_embeds],
prompt_attention_mask=[prompt_attention_mask],
negative_prompt_embeds=[negative_prompt_embeds],
negative_attention_mask=[negative_prompt_attention_mask],
height=training_args.num_height,
width=training_args.num_width,
num_frames=num_frames,
num_inference_steps=
num_inference_steps, # Use the current validation step
guidance_scale=sampling_param.guidance_scale,
n_tokens=n_tokens,
eta=0.0,
VSA_sparsity=training_args.VSA_sparsity,
)
return batch
@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.use_cpu_offload = False
training_args.use_cpu_offload = True
if not training_args.log_validation:
return
if self.validation_pipeline is None:
@@ -617,48 +602,53 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
logger.info("Using validation seed: %s", self.seed)
# Prepare validation prompts
logger.info('fastvideo_args.validation_preprocessed_path: %s',
training_args.validation_preprocessed_path)
validation_dataset, validation_dataloader = build_parquet_map_style_dataloader(
training_args.validation_preprocessed_path,
batch_size=1,
parquet_schema=self.validation_dataset_schema,
num_data_workers=0,
cfg_rate=0.0,
drop_last=False,
drop_first_row=sampling_param.negative_prompt is not None)
if sampling_param.negative_prompt:
negative_prompt_embeds, negative_prompt_attention_mask, negative_prompt = validation_dataset.get_validation_negative_prompt(
)
logger.info("Using negative_prompt: %s", negative_prompt)
logger.info('rank: %s: fastvideo_args.validation_dataset_file: %s',
self.global_rank,
training_args.validation_dataset_file,
local_main_process_only=False)
validation_dataset = ValidationDataset(
training_args.validation_dataset_file)
validation_dataloader = DataLoader(validation_dataset,
batch_size=None,
num_workers=0)
transformer.eval()
validation_steps = training_args.validation_sampling_steps.split(",")
validation_steps = [int(step) for step in validation_steps]
validation_steps = [step for step in validation_steps if step > 0]
# Log validation results for this step
world_group = get_world_group()
num_sp_groups = world_group.world_size // self.sp_group.world_size
# Process each validation prompt for each validation step
for num_inference_steps in validation_steps:
logger.info("rank: %s: num_inference_steps: %s",
self.global_rank,
num_inference_steps,
local_main_process_only=False)
step_videos: List[np.ndarray] = []
step_captions: List[str] = []
for validation_batch in validation_dataloader:
batch = self._prepare_validation_inputs(
sampling_param, training_args, validation_batch,
num_inference_steps, negative_prompt_embeds,
negative_prompt_attention_mask)
batch = self._prepare_validation_batch(sampling_param,
training_args,
validation_batch,
num_inference_steps)
logger.info("rank: %s: rank_in_sp_group: %s, batch.prompt: %s",
self.global_rank,
self.rank_in_sp_group,
batch.prompt,
local_main_process_only=False)
assert batch.prompt is not None and isinstance(
batch.prompt, str)
step_captions.append(batch.prompt)
# Run validation inference
with torch.no_grad(), torch.autocast("cuda",
dtype=torch.bfloat16):
output_batch = self.validation_pipeline.forward(
batch, training_args)
samples = output_batch.output
output_batch = self.validation_pipeline.forward(
batch, training_args)
samples = output_batch.output
if self.rank_in_sp_group != 0:
continue
@@ -672,10 +662,7 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
frames.append((x * 255).numpy().astype(np.uint8))
step_videos.append(frames)
# Log validation results for this step
world_group = get_world_group()
num_sp_groups = world_group.world_size // self.sp_group.world_size
torch.distributed.barrier()
# Only sp_group leaders (rank_in_sp_group == 0) need to send their
# results to global rank 0
if self.rank_in_sp_group == 0:
@@ -716,8 +703,10 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
# Other sp_group leaders send their results to global rank 0
world_group.send_object(step_videos, dst=0)
world_group.send_object(step_captions, dst=0)
torch.distributed.barrier()
# Re-enable gradients for training
training_args.inference_mode = False
transformer.train()
gc.collect()
torch.cuda.empty_cache()
fastvideo/v1/training/wan_i2v_training_pipeline.py
+44 -42
View File
@@ -15,10 +15,9 @@ from fastvideo.v1.models.schedulers.scheduling_flow_unipc_multistep import (
FlowUniPCMultistepScheduler)
from fastvideo.v1.pipelines.pipeline_batch_info import (ForwardBatch,
TrainingBatch)
from fastvideo.v1.pipelines.wan.wan_i2v_pipeline import (
WanImageToVideoValidationPipeline)
from fastvideo.v1.pipelines.wan.wan_i2v_pipeline import WanImageToVideoPipeline
from fastvideo.v1.training.training_pipeline import TrainingPipeline
from fastvideo.v1.utils import is_vsa_available
from fastvideo.v1.utils import is_vsa_available, shallow_asdict
vsa_available = is_vsa_available()
@@ -50,15 +49,18 @@ class WanI2VTrainingPipeline(TrainingPipeline):
args_copy = deepcopy(training_args)
args_copy.inference_mode = True
args_copy.pipeline_config.vae_config.load_encoder = False
validation_pipeline = WanImageToVideoValidationPipeline.from_pretrained(
args_copy.use_cpu_offload = True
# args_copy.pipeline_config.vae_config.load_encoder = False
# validation_pipeline = WanImageToVideoValidationPipeline.from_pretrained(
validation_pipeline = WanImageToVideoPipeline.from_pretrained(
training_args.model_path,
args=None,
inference_mode=True,
loaded_modules={"transformer": self.get_module("transformer")},
tp_size=training_args.tp_size,
sp_size=training_args.sp_size,
num_gpus=training_args.num_gpus)
num_gpus=training_args.num_gpus,
use_cpu_offload=True)
self.validation_pipeline = validation_pipeline
@@ -117,8 +119,29 @@ class WanI2VTrainingPipeline(TrainingPipeline):
image_latents = training_batch.image_latents.to(
get_local_torch_device(), dtype=torch.bfloat16)
temporal_compression_ratio = 4
num_frames = (self.training_args.num_latent_t -
1) * temporal_compression_ratio + 1
batch_size, num_channels, _, latent_height, latent_width = image_latents.shape
mask_lat_size = torch.ones(batch_size, 1, num_frames, latent_height,
latent_width)
mask_lat_size[:, :, 1:] = 0
first_frame_mask = mask_lat_size[:, :, :1]
first_frame_mask = torch.repeat_interleave(
first_frame_mask, dim=2, repeats=temporal_compression_ratio)
mask_lat_size = torch.cat([first_frame_mask, mask_lat_size[:, :, 1:]],
dim=2)
mask_lat_size = mask_lat_size.view(batch_size, -1,
temporal_compression_ratio,
latent_height, latent_width)
mask_lat_size = mask_lat_size.transpose(1, 2)
mask_lat_size = mask_lat_size.to(
image_latents.device).to(dtype=torch.bfloat16)
training_batch.noisy_model_input = torch.cat(
[training_batch.noisy_model_input, image_latents], dim=1)
[training_batch.noisy_model_input, mask_lat_size, image_latents],
dim=1)
return training_batch
@@ -156,55 +179,34 @@ class WanI2VTrainingPipeline(TrainingPipeline):
}
return training_batch
def _prepare_validation_inputs(
self, sampling_param: SamplingParam, training_args: TrainingArgs,
validation_batch: Dict[str, Any], num_inference_steps: int,
negative_prompt_embeds: torch.Tensor | None,
negative_prompt_attention_mask: torch.Tensor | None
) -> ForwardBatch:
embeddings = validation_batch['text_embedding']
masks = validation_batch['text_attention_mask']
clip_features = validation_batch['clip_feature']
pil_image = validation_batch['pil_image']
infos = validation_batch['info_list']
prompt = infos[0]['prompt']
def _prepare_validation_batch(self, sampling_param: SamplingParam,
training_args: TrainingArgs,
validation_batch: Dict[str, Any],
num_inference_steps: int) -> ForwardBatch:
sampling_param.prompt = validation_batch['prompt']
sampling_param.height = training_args.num_height
sampling_param.width = training_args.num_width
sampling_param.image_path = validation_batch['video_path']
sampling_param.num_inference_steps = num_inference_steps
sampling_param.data_type = "video"
sampling_param.seed = self.seed
prompt_embeds = embeddings.to(get_local_torch_device())
prompt_attention_mask = masks.to(get_local_torch_device())
clip_features = clip_features.to(get_local_torch_device())
# Calculate sizes
latents_size = [(sampling_param.num_frames - 1) // 4 + 1,
sampling_param.height // 8, sampling_param.width // 8]
n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
temporal_compression_factor = training_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio
num_frames = (training_args.num_latent_t -
1) * temporal_compression_factor + 1
# Prepare batch for validation
sampling_param.num_frames = num_frames
batch = ForwardBatch(
prompt=prompt,
data_type="video",
**shallow_asdict(sampling_param),
latents=None,
seed=self.seed, # Use deterministic seed
generator=torch.Generator(device="cpu").manual_seed(self.seed),
prompt_embeds=[prompt_embeds],
prompt_attention_mask=[prompt_attention_mask],
negative_prompt_embeds=[negative_prompt_embeds],
negative_attention_mask=[negative_prompt_attention_mask],
image_embeds=[clip_features],
preprocessed_image=pil_image,
height=training_args.num_height,
width=training_args.num_width,
num_frames=num_frames,
num_inference_steps=
num_inference_steps, # Use the current validation step
guidance_scale=sampling_param.guidance_scale,
n_tokens=n_tokens,
eta=0.0,
VSA_sparsity=training_args.VSA_sparsity,
)
return batch
fastvideo/v1/training/wan_training_pipeline.py
+3 -2
View File
@@ -6,7 +6,7 @@ from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs
from fastvideo.v1.logger import init_logger
from fastvideo.v1.models.schedulers.scheduling_flow_unipc_multistep import (
FlowUniPCMultistepScheduler)
from fastvideo.v1.pipelines.wan.wan_pipeline import WanValidationPipeline
from fastvideo.v1.pipelines.wan.wan_pipeline import WanPipeline
from fastvideo.v1.training.training_pipeline import TrainingPipeline
from fastvideo.v1.utils import is_vsa_available
@@ -36,8 +36,9 @@ class WanTrainingPipeline(TrainingPipeline):
args_copy = deepcopy(training_args)
args_copy.inference_mode = True
args_copy.use_cpu_offload = True
args_copy.pipeline_config.vae_config.load_encoder = False
validation_pipeline = WanValidationPipeline.from_pretrained(
validation_pipeline = WanPipeline.from_pretrained(
training_args.model_path,
args=None,
inference_mode=True,
pyproject.toml
-1
View File
@@ -79,7 +79,6 @@ exclude = ["assets*", "docker*", "docs", "scripts*"]
exclude = ["assets*", "docker*", "docs", "scripts*"]
[tool.mypy]
warn_return_any = true
warn_unused_configs = true
ignore_missing_imports = true
disallow_untyped_calls = true
scripts/finetune/finetune_v1.sh
+3 -2
View File
@@ -1,9 +1,10 @@
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
export TOKENIZERS_PARALLELISM=false
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
DATA_DIR=[your data dir]
VALIDATION_DIR=[your validation dir]
VALIDATION_DATASET_FILE=[your validation dataset file]
NUM_GPUS=4
# export CUDA_VISIBLE_DEVICES=4,5
# IP=[MASTER NODE IP]
@@ -15,7 +16,7 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\
--inference_mode False\
--pretrained_model_name_or_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
--data_path "$DATA_DIR"\
--validation_preprocessed_path "$VALIDATION_DIR"\
--validation_dataset_file "$VALIDATION_DATASET_FILE"\
--train_batch_size=4 \
--num_latent_t 20 \
--sp_size 4 \
scripts/finetune/finetune_v1_VSA.sh
+3 -2
View File
@@ -1,10 +1,11 @@
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=offline
export WANDB_API_KEY='your_wandb_api_key'
export TOKENIZERS_PARALLELISM=false
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
export TRITON_CACHE_DIR=/tmp/triton_cache
DATA_DIR=~/train/
VALIDATION_DIR=~latents/test/
VALIDATION_DATASET_FILE=[your validation dataset file]
NUM_GPUS=8
export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
# export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
@@ -21,7 +22,7 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS \
--pretrained_model_name_or_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
--cache_dir "/home/ray/.cache" \
--data_path "$DATA_DIR" \
--validation_preprocessed_path "$VALIDATION_DIR" \
--validation_dataset_file "$VALIDATION_DATASET_FILE" \
--train_batch_size 1 \
--num_latent_t 16 \
--sp_size 1 \