# Clone the project repository
git clone https://github.com/WeChatCV/Wan-Alpha.git
cd Wan-Alpha

# Create and activate Conda environment
conda create -n Wan-Alpha python=3.11 -y
conda activate Wan-Alpha

# Install dependencies
pip install -r requirements.txt

torchrun --nproc_per_node=8 --master_port=29501 generate_dora_lightx2v_mask.py --size 832*480\
         --ckpt_dir "path/to/your/Wan-2.1/Wan2.1-T2V-14B" \
         --dit_fsdp --t5_fsdp --ulysses_size 8 \
         --vae_lora_checkpoint "path/to/your/decoder.bin" \
         --lora_path "path/to/your/t2v.safetensors" \
         --lightx2v_path "path/to/your/lightx2v_T2V_14B_cfg_step_distill_v2_lora_rank64_bf16.safetensors" \
         --sample_guide_scale 1.0 \
         --frame_num 81 \
         --sample_steps 4 \
         --lora_ratio 1.0 \
         --lora_prefix "" \
         --alpha_shift_mean 0.05 \
         --cache_path_mask "path/to/your/gauss_mask" \
         --prompt_file ./data/prompt.txt \
         --output_dir ./output 

# An example of prompt.
This video has a transparent background. Close-up shot. A colorful parrot flying. Realistic style.

save_path="./checkpoints"
mkdir -p $save_path

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python3 train_vae.py \
  --train_dataset_path data.csv \
  --val_dataset_path val.csv \
  --output_path $save_path \
  --vae_path "Wan-2.1/Wan2.1-T2V-14B/Wan2.1_VAE.pth" \
  --num_frames 17 \
  --height 272 \
  --width 272 \
  --training_strategy deepspeed_stage_2 \
  --max_epochs 100 \
  --learning_rate 1e-4 \
  --lora_rank 128 \
  --batch_size 2 \
  --model_id_with_origin_paths "Wan2.1-T2V-1.3B:diffusion_pytorch_model*.safetensors" \
  --job_name "prompt" \
  --use_gradient_checkpointing_offload \
  --dataloader_num_workers 8 2>&1 | tee -a ${save_path}/train.log

job_name="12"
output_path="./checkpoints"
mkdir -p ${output_path}

# cache data
for id in {0..7}; do
  CUDA_VISIBLE_DEVICES=$id accelerate launch \
    --mixed_precision='bf16' \
    --num_processes=1 \
    --num_machines=$WORLD_SIZE \
    --machine_rank=$RANK \
    --main_process_ip=$MASTER_ADDR \
    --main_process_port=$MASTER_PORT \
    examples/wanvideo/model_training/prepare_trans_alpha_mask.py \
    --dataset_metadata_path data.csv \
    --height 640 \
    --width 624 \
    --data_file_keys video_fgr,video_pha \
    --dataset_repeat 1 \
    --model_id_with_origin_paths "Wan2.1-T2V-14B:models_t5_umt5-xxl-enc-bf16.pth" \
    --learning_rate 1e-4 \
    --num_epochs 1 \
    --remove_prefix_in_ckpt "pipe.t2v." \
    --output_path "" \
    --lora_base_model "dit" \
    --lora_target_modules "q,k,v,o,ffn.0,ffn.2" \
    --lora_rank 32 \
    --job_name $job_name \
    --job_id $id \
    --job_num 8 \
    --new_vae_path "VAE/pytorch_model.bin" \
    --use_gradient_checkpointing_offload  2>&1 | tee -a ${output_path}/train_prepare.log &
done
trap 'kill 0' SIGINT
wait


# train
accelerate launch \
  --mixed_precision='bf16' \
  --num_processes=$RANK_NUM \
  --num_machines=$WORLD_SIZE \
  --machine_rank=$RANK \
  --main_process_ip=$MASTER_ADDR \
  --main_process_port=$MASTER_PORT \
  --config_file $ACCELERATE_CONFIG_FILE \
  examples/wanvideo/model_training/train_gauss_ellipse.py \
  --dataset_metadata_path data.csv \
  --height 640 \
  --width 624 \
  --dataset_repeat 1 \
  --model_id_with_origin_paths "Wan2.1-T2V-14B:diffusion_pytorch_model*.safetensors" \
  --learning_rate 1e-4 \
  --num_epochs 20 \
  --remove_prefix_in_ckpt "pipe.dit." \
  --output_path $output_path \
  --lora_base_model "dit" \
  --lora_target_modules "q,k,v,o,ffn.0,ffn.2" \
  --lora_rank 32 \
  --initial_learnable_value 0.05 \
  --job_name $job_name \
  --use_gradient_checkpointing_offload 2>&1 | tee -a ${output_path}/train.log