diff --git a/examples/wan2.1_causal_forcing/predict_t2v.py b/examples/wan2.1_causal_forcing/predict_t2v.py new file mode 100644 index 00000000..eed2ca86 --- /dev/null +++ b/examples/wan2.1_causal_forcing/predict_t2v.py @@ -0,0 +1,329 @@ +import os +import sys + +import numpy as np +import torch +from diffusers import FlowMatchEulerDiscreteScheduler +from omegaconf import OmegaConf +from PIL import Image + +current_file_path = os.path.abspath(__file__) +project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dirname(current_file_path)), os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))] +for project_root in project_roots: + sys.path.insert(0, project_root) if project_root not in sys.path else None + +from videox_fun.dist import set_multi_gpus_devices, shard_model +from videox_fun.models import (AutoencoderKLWan, AutoTokenizer, + WanT5EncoderModel, + WanTransformer3DModel_SelfForcing) +from videox_fun.pipeline import WanSelfForcingPipeline +from videox_fun.utils import (register_auto_device_hook, + safe_enable_group_offload) +from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler +from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler +from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8, + convert_weight_dtype_wrapper, + replace_parameters_by_name) +from videox_fun.utils.lora_utils import merge_lora, unmerge_lora +from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent, + save_videos_grid) + +# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload]. +# model_full_load means that the entire model will be moved to the GPU. +# +# model_full_load_and_qfloat8 means that the entire model will be moved to the GPU, +# and the transformer model has been quantized to float8, which can save more GPU memory. +# +# model_cpu_offload means that the entire model will be moved to the CPU after use, which can save some GPU memory. +# +# model_cpu_offload_and_qfloat8 indicates that the entire model will be moved to the CPU after use, +# and the transformer model has been quantized to float8, which can save more GPU memory. +# +# model_group_offload transfers internal layer groups between CPU/CUDA, +# balancing memory efficiency and speed between full-module and leaf-level offloading methods. +# +# sequential_cpu_offload means that each layer of the model will be moved to the CPU after use, +# resulting in slower speeds but saving a large amount of GPU memory. +GPU_memory_mode = "model_cpu_offload" +# Multi GPUs config +# Please ensure that the product of ulysses_degree and ring_degree equals the number of GPUs used. +# For example, if you are using 8 GPUs, you can set ulysses_degree = 2 and ring_degree = 4. +# If you are using 1 GPU, you can set ulysses_degree = 1 and ring_degree = 1. +ulysses_degree = 1 +ring_degree = 1 +# Use FSDP to save more GPU memory in multi gpus. +fsdp_dit = False +fsdp_text_encoder = True +# Compile will give a speedup in fixed resolution and need a little GPU memory. +# The compile_dit is not compatible with the fsdp_dit and sequential_cpu_offload. +compile_dit = False + +# Config and model path +config_path = "config/wan2.1/wan_civitai.yaml" +# model path +model_name = "models/Diffusion_Transformer/Wan2.1-T2V-1.3B" + +# Choose the sampler in "Flow", "Flow_Unipc", "Flow_DPM++". +sampler_name = "Flow" +# [NOTE]: Noise schedule shift parameter. Affects temporal dynamics. +# Used when the sampler is in "Flow_Unipc", "Flow_DPM++". +shift = 5.0 + +# Causal-Forcing checkpoint to overlay on top of the Wan2.1 base model. +transformer_path = "output_dir_wan2.1_causal_forcing_dmd/checkpoint-4000/diffusion_pytorch_model.safetensors" +use_ema = False +vae_path = None +lora_path = None + +# Other params +sample_size = [480, 832] +video_length = 81 +fps = 16 + +# Causal-Forcing causal inference config +# `num_frame_per_block`: 1 = frame-wise (matches the .../checkpoints/framewise ckpts); +# 3 = chunk-wise (matches the .../checkpoints/chunkwise ckpts). +num_frame_per_block = 1 +# Local attention window size (-1 for global attention) +local_attn_size = -1 +# Others +independent_first_frame = False +context_noise = 0.0 + +# Use torch.float16 if GPU does not support torch.bfloat16 +# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16 +weight_dtype = torch.bfloat16 +prompt = "A stylish woman walks down a Tokyo street filled with warm glowing neon and animated city signage. She wears a black leather jacket, a long red dress, and black boots, and carries a black purse. She wears sunglasses and red lipstick. She walks confidently and casually. The street is damp and reflective, creating a mirror effect of the colorful lights. Many pedestrians walk about." +negative_prompt = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" +# -------- Stage selector (uncomment ONE block) -------- +# All few-step distilled stages (2, 3) bake CFG into the student weights, so +# inference MUST use guidance_scale=1.0 — CF's CausalInferencePipeline does +# zero CFG (grep "unconditional/cfg/guidance" in pipeline/causal_inference.py +# returns 0). Using gs>1 stacks CFG on top of a CFG-baked model and produces +# over-saturated, AR-unstable outputs. +# +# Stage 1 — AR diffusion (`ar_diffusion.pt`): 50-step UniPC + CFG. +# guidance_scale = 3.0 +# num_inference_steps = 50 +# stochastic_sampling = False +# +# Stage 2 — CCD (`causal_cd.pt`): 4-step consistency-distilled. +# guidance_scale = 1.0 +# num_inference_steps = 4 +# stochastic_sampling = True +# +# Stage 3 — DMD (`causal_forcing.pt`): 2-step distribution-matching distilled. +# guidance_scale = 1.0 +# num_inference_steps = 2 +# stochastic_sampling = True +guidance_scale = 1.0 +num_inference_steps = 2 +stochastic_sampling = True +seed = 43 +lora_weight = 0.55 +save_path = "samples/wan-videos-causal-forcing" + +device = set_multi_gpus_devices(ulysses_degree, ring_degree) +config = OmegaConf.load(config_path) + +# Load transformer with causal inference support if enabled +transformer_additional_kwargs = OmegaConf.to_container(config['transformer_additional_kwargs']) +transformer_additional_kwargs['local_attn_size'] = local_attn_size + +transformer = WanTransformer3DModel_SelfForcing.from_pretrained( + os.path.join(model_name, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')), + transformer_additional_kwargs=transformer_additional_kwargs, + low_cpu_mem_usage=True, + torch_dtype=weight_dtype, +) + +if transformer_path is not None: + def _resolve_transformer_path(raw_path: str, prefer_ema: bool) -> str: + """Resolve a checkpoint path to the actual weights file to load. + + File paths are returned unchanged so external `.pt` ckpts (CF official) keep + working. For trainer-output dirs, prefer EMA when asked and available, + otherwise fall back to live `transformer/` weights. + """ + if os.path.isfile(raw_path): + return raw_path + if os.path.isdir(raw_path): + candidates = [] + if prefer_ema: + candidates.append(os.path.join(raw_path, "ema_transformer", "diffusion_pytorch_model.safetensors")) + candidates.append(os.path.join(raw_path, "transformer", "diffusion_pytorch_model.safetensors")) + candidates.append(os.path.join(raw_path, "diffusion_pytorch_model.safetensors")) + for c in candidates: + if os.path.isfile(c): + return c + raise FileNotFoundError( + f"transformer_path={raw_path!r} is neither a file nor a checkpoint dir " + f"with a known safetensors layout (transformer/ or ema_transformer/)." + ) + + _raw_transformer_path = transformer_path + transformer_path = _resolve_transformer_path(transformer_path, prefer_ema=use_ema) + if transformer_path != _raw_transformer_path: + print(f"use_ema={use_ema}: resolved {_raw_transformer_path} -> {transformer_path}") + print(f"From checkpoint: {transformer_path}") + if transformer_path.endswith("safetensors"): + from safetensors.torch import load_file, safe_open + state_dict = load_file(transformer_path) + else: + state_dict = torch.load(transformer_path, map_location="cpu") + + state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict + state_dict = state_dict["generator_ema"] if "generator_ema" in state_dict else state_dict + state_dict = state_dict["generator"] if "generator" in state_dict else state_dict + # Causal-Forcing's FSDP-saved ckpts (causal_cd.pt / causal_forcing.pt) keep the + # `model._fsdp_wrapped_module.` prefix; strip it before the generic `model.` strip + # so both kinds of ckpt land at bare parameter names. + if any("._fsdp_wrapped_module." in k for k in state_dict.keys()): + state_dict = {k.replace("model._fsdp_wrapped_module.", "model.", 1) if k.startswith("model._fsdp_wrapped_module.") else k: v for k, v in state_dict.items()} + if any(k.startswith("model.") for k in state_dict.keys()): + state_dict = {k.replace("model.", "", 1) if k.startswith("model.") else k: v for k, v in state_dict.items()} + + m, u = transformer.load_state_dict(state_dict, strict=False) + print(f"missing keys: {len(m)}, unexpected keys: {len(u)}") + +# Get Vae +vae = AutoencoderKLWan.from_pretrained( + os.path.join(model_name, config['vae_kwargs'].get('vae_subpath', 'vae')), + additional_kwargs=OmegaConf.to_container(config['vae_kwargs']), +).to(weight_dtype) + +if vae_path is not None: + print(f"From checkpoint: {vae_path}") + if vae_path.endswith("safetensors"): + from safetensors.torch import load_file, safe_open + state_dict = load_file(vae_path) + else: + state_dict = torch.load(vae_path, map_location="cpu") + state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict + + m, u = vae.load_state_dict(state_dict, strict=False) + print(f"missing keys: {len(m)}, unexpected keys: {len(u)}") + +# Get Tokenizer +tokenizer = AutoTokenizer.from_pretrained( + os.path.join(model_name, config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer')), +) + +# Get Text encoder +text_encoder = WanT5EncoderModel.from_pretrained( + os.path.join(model_name, config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')), + additional_kwargs=OmegaConf.to_container(config['text_encoder_kwargs']), + low_cpu_mem_usage=True, + torch_dtype=weight_dtype, +) + +# Get Scheduler +Chosen_Scheduler = scheduler_dict = { + "Flow": FlowMatchEulerDiscreteScheduler, + "Flow_Unipc": FlowUniPCMultistepScheduler, + "Flow_DPM++": FlowDPMSolverMultistepScheduler, +}[sampler_name] +if sampler_name == "Flow_Unipc" or sampler_name == "Flow_DPM++": + config['scheduler_kwargs']['shift'] = 1 +scheduler = Chosen_Scheduler( + **filter_kwargs(Chosen_Scheduler, OmegaConf.to_container(config['scheduler_kwargs'])) +) + +# Get Pipeline +pipeline = WanSelfForcingPipeline( + transformer=transformer, + vae=vae, + tokenizer=tokenizer, + text_encoder=text_encoder, + scheduler=scheduler, +) + +if ulysses_degree > 1 or ring_degree > 1: + from functools import partial + transformer.enable_multi_gpus_inference() + if fsdp_dit: + shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype) + pipeline.transformer = shard_fn(pipeline.transformer) + print("Add FSDP DIT") + if fsdp_text_encoder: + shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype) + pipeline.text_encoder = shard_fn(pipeline.text_encoder) + print("Add FSDP TEXT ENCODER") + +if compile_dit: + for i in range(len(pipeline.transformer.blocks)): + pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i]) + print("Add Compile") + +if GPU_memory_mode == "sequential_cpu_offload": + replace_parameters_by_name(transformer, ["modulation",], device=device) + transformer.freqs = transformer.freqs.to(device=device) + pipeline.enable_sequential_cpu_offload(device=device) +elif GPU_memory_mode == "model_group_offload": + register_auto_device_hook(pipeline.transformer) + safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True) +elif GPU_memory_mode == "model_cpu_offload_and_qfloat8": + convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device) + convert_weight_dtype_wrapper(transformer, weight_dtype) + pipeline.enable_model_cpu_offload(device=device) +elif GPU_memory_mode == "model_cpu_offload": + pipeline.enable_model_cpu_offload(device=device) +elif GPU_memory_mode == "model_full_load_and_qfloat8": + convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device) + convert_weight_dtype_wrapper(transformer, weight_dtype) + pipeline.to(device=device) +else: + pipeline.to(device=device) + +generator = torch.Generator(device=device).manual_seed(seed) + +if lora_path is not None: + pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) + +with torch.no_grad(): + video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1 + latent_frames = (video_length - 1) // vae.config.temporal_compression_ratio + 1 + + sample = pipeline( + prompt, + num_frames = video_length, + negative_prompt = negative_prompt, + height = sample_size[0], + width = sample_size[1], + generator = generator, + guidance_scale = guidance_scale, + num_inference_steps = num_inference_steps, + shift = shift, + num_frame_per_block = num_frame_per_block, + independent_first_frame = independent_first_frame, + context_noise = context_noise, + stochastic_sampling = stochastic_sampling, + ).videos + +if lora_path is not None: + pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) + +def save_results(): + if not os.path.exists(save_path): + os.makedirs(save_path, exist_ok=True) + + index = len([path for path in os.listdir(save_path)]) + 1 + prefix = str(index).zfill(8) + if video_length == 1: + video_path = os.path.join(save_path, prefix + ".png") + + image = sample[0, :, 0] + image = image.transpose(0, 1).transpose(1, 2) + image = (image * 255).numpy().astype(np.uint8) + image = Image.fromarray(image) + image.save(video_path) + else: + video_path = os.path.join(save_path, prefix + ".mp4") + save_videos_grid(sample, video_path, fps=fps) + +if ulysses_degree * ring_degree > 1: + import torch.distributed as dist + if dist.get_rank() == 0: + save_results() +else: + save_results() \ No newline at end of file diff --git a/scripts/wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION.md b/scripts/wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION.md new file mode 100644 index 00000000..92fea918 --- /dev/null +++ b/scripts/wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION.md @@ -0,0 +1,276 @@ +# Wan2.1 Causal-Forcing Stage 1: Autoregressive Diffusion Training Guide + +This document provides the complete workflow for **Causal-Forcing Stage 1 — Autoregressive Diffusion (AR Diffusion) training** on Wan2.1-T2V-1.3B. + +> **What is Causal-Forcing?** +> +> Causal-Forcing is a three-stage pipeline that compresses a video diffusion model into a **few-step causal autoregressive generator**: +> +> 1. **Stage 1 — AR Diffusion** (`train_ar_diffusion.py`): Train the model with **teacher forcing** on clean video latents, learning to denoise chunk-by-chunk (or frame-by-frame) in a causal autoregressive manner. This produces a strong AR backbone. +> 2. **Stage 2 — Causal Consistency Distillation** (`train_causal_consistency_distill.py`): Distill the multi-step AR model into a one-step-per-block consistency model. +> 3. **Stage 3 — Causal DMD** (`train_causal_dmd.py`): Further distill to a **2-step** generator using distribution matching with a 14B teacher. +> +> This README covers **Stage 1 only**. + +--- + +## Table of Contents +- [1. Environment Setup](#1-environment-setup) +- [2. Download Pretrained Models](#2-download-pretrained-models) +- [3. Prepare Training Data](#3-prepare-training-data) + - [3.1 Quick Demo Dataset](#31-quick-demo-dataset) + - [3.2 Dataset Structure](#32-dataset-structure) + - [3.3 metadata.json Format](#33-metadatajson-format) +- [4. Training](#4-training) + - [4.1 Quick Start](#41-quick-start) + - [4.2 Key Parameters](#42-key-parameters) + - [4.3 Causal-Forcing-Specific Parameters](#43-causal-forcing-specific-parameters) + - [4.4 Training with FSDP](#44-training-with-fsdp) +- [5. Use the Trained Checkpoint](#5-use-the-trained-checkpoint) +- [6. Additional Resources](#6-additional-resources) + +--- + +## 1. Environment Setup + +**Method 1: Using requirements.txt** + +```bash +pip install -r requirements.txt +``` + +**Method 2: Manual Installation** + +```bash +pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image +pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime +pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" +pip install yunchang xfuser modelscope openpyxl +pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y +pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 +``` + +**Method 3: Using Docker** + +```bash +# pull image +docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun + +# enter image +docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun +``` + +--- + +## 2. Download Pretrained Models + +Stage 1 initializes the Wan2.1-T2V-1.3B model and trains it with causal teacher forcing. + +```bash +# Create model directory +mkdir -p models/Diffusion_Transformer + +# Download Wan2.1 T2V base model +modelscope download --model Wan-AI/Wan2.1-T2V-1.3B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-1.3B +``` + +--- + +## 3. Prepare Training Data + +Stage 1 trains directly on **raw videos** with online VAE encoding — no ODE trajectory pairs are needed. + +### 3.1 Quick Demo Dataset + +We provide a small demo dataset for quick testing: + +```bash +# Download official demo dataset +modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo +``` + +### 3.2 Dataset Structure + +``` +📦 datasets/ +├── 📂 X-Fun-Videos-Demo/ +│ ├── 📂 train/ +│ │ ├── 📄 video001.mp4 +│ │ ├── 📄 video002.mp4 +│ │ └── 📄 ... +│ └── 📄 metadata.json +``` + +### 3.3 metadata.json Format + +**Relative paths** (recommended): +```json +[ + { + "file_path": "train/video001.mp4", + "text": "A beautiful sunset over the ocean, golden hour lighting", + "type": "video" + }, + { + "file_path": "train/video002.mp4", + "text": "A person walking through a forest, cinematic view", + "type": "video" + } +] +``` + +**Absolute paths**: +```json +[ + { + "file_path": "/mnt/data/videos/sunset.mp4", + "text": "A beautiful sunset over the ocean", + "type": "video" + } +] +``` + +--- + +## 4. Training + +### 4.1 Quick Start + +The ready-to-use launcher is [train_ar_diffusion.sh](./train_ar_diffusion.sh): + +```bash +export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" + +accelerate launch --mixed_precision="bf16" --use_fsdp \ + --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \ + --fsdp_transformer_layer_cls_to_wrap=CasualWanAttentionBlock \ + --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT \ + --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False \ + scripts/wan2.1_causal_forcing/train_ar_diffusion.py \ + --config_path="config/wan2.1/wan_civitai.yaml" \ + --pretrained_model_name_or_path=$MODEL_NAME \ + --train_data_dir=$DATASET_NAME \ + --train_data_meta=$DATASET_META_NAME \ + --image_sample_size=640 \ + --video_sample_size=640 \ + --token_sample_size=640 \ + --fix_sample_size 480 832 \ + --video_sample_stride=2 \ + --video_sample_n_frames=81 \ + --train_batch_size=1 \ + --gradient_accumulation_steps=1 \ + --dataloader_num_workers=8 \ + --num_train_epochs=100 \ + --checkpointing_steps=200 \ + --learning_rate=2e-06 \ + --lr_scheduler="constant_with_warmup" \ + --lr_warmup_steps=100 \ + --seed=42 \ + --output_dir="output_dir_wan2.1_causal_forcing_ar_diffusion" \ + --gradient_checkpointing \ + --mixed_precision="bf16" \ + --adam_weight_decay=3e-2 \ + --adam_epsilon=1e-10 \ + --vae_mini_batch=1 \ + --max_grad_norm=0.05 \ + --random_hw_adapt \ + --training_with_video_token_length \ + --enable_bucket \ + --num_frame_per_block=3 \ + --train_sampling_steps=1000 \ + --shift=5.0 \ + --use_timestep_weight \ + --trainable_modules "." +``` + +Or simply: + +```bash +bash scripts/wan2.1_causal_forcing/train_ar_diffusion.sh +``` + +### 4.2 Key Parameters + +| Parameter | Description | Example Value | +|-----------|-------------|---------------| +| `--pretrained_model_name_or_path` | Base model to initialize | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B/` | +| `--config_path` | Model config YAML | `config/wan2.1/wan_civitai.yaml` | +| `--train_data_dir` | Video root directory (prepended to relative `file_path`) | `""` | +| `--train_data_meta` | Path to `metadata.json` | `datasets/X-Fun-Videos-Demo/metadata.json` | +| `--train_batch_size` | Per-GPU batch size | 1 | +| `--gradient_accumulation_steps` | Gradient accumulation steps | 1 | +| `--dataloader_num_workers` | DataLoader workers | 8 | +| `--num_train_epochs` | Number of training epochs | 100 | +| `--checkpointing_steps` | Save checkpoint every N steps | 200 | +| `--learning_rate` | Initial learning rate | 2e-06 | +| `--lr_scheduler` | LR scheduler type | `constant_with_warmup` | +| `--lr_warmup_steps` | LR warmup steps | 100 | +| `--output_dir` | Output directory | `output_dir_wan2.1_causal_forcing_ar_diffusion` | +| `--gradient_checkpointing` | Enable activation checkpointing | - | +| `--mixed_precision` | `fp16` / `bf16` | `bf16` | +| `--adam_weight_decay` | AdamW weight decay | 3e-2 | +| `--adam_epsilon` | AdamW epsilon | 1e-10 | +| `--max_grad_norm` | Gradient clipping threshold | 0.05 | +| `--trainable_modules` | Trainable modules (`"."` = all) | `"."` | + +**Video Sampling Parameters**: + +| Parameter | Description | Example Value | +|-----------|-------------|---------------| +| `--image_sample_size` | Image sampling size | 640 | +| `--video_sample_size` | Video sampling size | 640 | +| `--token_sample_size` | Token sampling size | 640 | +| `--fix_sample_size` | Fixed `[height, width]` for output | `480 832` | +| `--video_sample_stride` | Frame sampling stride | 2 | +| `--video_sample_n_frames` | Number of video frames | 81 | +| `--random_hw_adapt` | Enable random resolution adaptation | - | +| `--training_with_video_token_length` | Enable token-length-based training | - | +| `--enable_bucket` | Enable aspect-ratio bucket sampling | - | +| `--vae_mini_batch` | VAE encoding mini-batch size (1 to avoid OOM) | 1 | + +### 4.3 Causal-Forcing-Specific Parameters + +| Parameter | Description | Example Value | +|-----------|-------------|---------------| +| `--num_frame_per_block` | Frames per causal block. `3` = chunkwise, `1` = framewise | 3 | +| `--independent_first_frame` | First frame is independent (`[1, N, N, ...]` block pattern, useful for I2V) | - | +| `--shift` | `FlowMatchEulerDiscreteScheduler` shift (Causal-Forcing default: 5.0) | 5.0 | +| `--train_sampling_steps` | Total scheduler timesteps for flow matching | 1000 | +| `--use_timestep_weight` | Apply per-timestep Gaussian loss weight (centered at T/2) | - | +| `--no_teacher_forcing` | Disable teacher forcing (use diffusion forcing instead) | - | +| `--noise_augmentation_max_timestep` | Add light noise to clean context tokens during teacher forcing (0 = off) | 0 | + +> **Teacher Forcing**: By default, Stage 1 uses teacher forcing — the model receives the **clean GT latent** as context (`clean_x`) when predicting the current block. This stabilizes early training. Disable with `--no_teacher_forcing` to train under diffusion forcing instead. + +### 4.4 Training with FSDP + +The script above already uses FSDP with `CasualWanAttentionBlock` auto-wrapping. For single-GPU training without FSDP: + +```bash +accelerate launch --mixed_precision="bf16" \ + scripts/wan2.1_causal_forcing/train_ar_diffusion.py \ + ... +``` + +--- + +## 5. Use the Trained Checkpoint + +The Stage 1 checkpoint is used to initialize **Stage 2 (Causal Consistency Distillation)**. Pass its path to `train_causal_consistency_distill.py` via `--transformer_path` and `--teacher_transformer_path`: + +```bash +--transformer_path="output_dir_wan2.1_causal_forcing_ar_diffusion/checkpoint-{N}/diffusion_pytorch_model.safetensors" \ +--teacher_transformer_path="output_dir_wan2.1_causal_forcing_ar_diffusion/checkpoint-{N}/diffusion_pytorch_model.safetensors" +``` + +See [train_causal_consistency_distill.sh](./train_causal_consistency_distill.sh) for the full Stage 2 workflow. + +--- + +## 6. Additional Resources + +- **Causal-Forcing Paper**: https://github.com/thu-ml/Causal-Forcing +- **Official GitHub**: https://github.com/aigc-apps/VideoX-Fun diff --git a/scripts/wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION_zh-CN.md b/scripts/wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION_zh-CN.md new file mode 100644 index 00000000..08aefe20 --- /dev/null +++ b/scripts/wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION_zh-CN.md @@ -0,0 +1,276 @@ +# Wan2.1 Causal-Forcing 第一阶段：自回归扩散（AR Diffusion）训练指南 + +本文档介绍 **Causal-Forcing 第一阶段 — 自回归扩散（AR Diffusion）训练** 的完整流程。 + +> **什么是 Causal-Forcing？** +> +> Causal-Forcing 是一个三阶段流水线，将视频扩散模型压缩为 **少步因果自回归生成器**： +> +> 1. **第一阶段 — AR Diffusion**（`train_ar_diffusion.py`）：在干净视频 latent 上使用 **teacher forcing** 训练，让模型学会按块（或按帧）因果自回归地去噪，得到一个强大的 AR 骨架。 +> 2. **第二阶段 — 因果一致性蒸馏**（`train_causal_consistency_distill.py`）：将多步 AR 模型蒸馏为每块一步的一致性模型。 +> 3. **第三阶段 — 因果 DMD**（`train_causal_dmd.py`）：利用 14B 教师模型做分布匹配，进一步蒸馏为 **2 步** 生成器。 +> +> 本文档仅覆盖 **第一阶段**。 + +--- + +## 目录 +- [一、环境配置](#一环境配置) +- [二、下载预训练模型](#二下载预训练模型) +- [三、准备训练数据](#三准备训练数据) + - [3.1 快速测试数据集](#31-快速测试数据集) + - [3.2 数据集结构](#32-数据集结构) + - [3.3 metadata.json 格式](#33-metadatajson-格式) +- [四、训练](#四训练) + - [4.1 快速开始](#41-快速开始) + - [4.2 主要参数说明](#42-主要参数说明) + - [4.3 Causal-Forcing 特有参数](#43-causal-forcing-特有参数) + - [4.4 使用 FSDP 训练](#44-使用-fsdp-训练) +- [五、使用训练好的 Checkpoint](#五使用训练好的-checkpoint) +- [六、更多资源](#六更多资源) + +--- + +## 一、环境配置 + +**方式 1：使用 requirements.txt** + +```bash +pip install -r requirements.txt +``` + +**方式 2：手动安装依赖** + +```bash +pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image +pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime +pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" +pip install yunchang xfuser modelscope openpyxl +pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y +pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 +``` + +**方式 3：使用 docker** + +```bash +# 拉取镜像 +docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun + +# 进入容器 +docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun +``` + +--- + +## 二、下载预训练模型 + +第一阶段以 Wan2.1-T2V-1.3B 为基础模型，在其上进行因果 teacher forcing 训练。 + +```bash +# 创建模型目录 +mkdir -p models/Diffusion_Transformer + +# 下载 Wan2.1 T2V 基础模型 +modelscope download --model Wan-AI/Wan2.1-T2V-1.3B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-1.3B +``` + +--- + +## 三、准备训练数据 + +第一阶段直接在 **原始视频** 上训练，VAE 在线编码 — **不需要** ODE 轨迹对或预计算 latent。 + +### 3.1 快速测试数据集 + +我们提供了一个小型 demo 数据集用于快速测试： + +```bash +# 下载官方示例数据集 +modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo +``` + +### 3.2 数据集结构 + +``` +📦 datasets/ +├── 📂 X-Fun-Videos-Demo/ +│ ├── 📂 train/ +│ │ ├── 📄 video001.mp4 +│ │ ├── 📄 video002.mp4 +│ │ └── 📄 ... +│ └── 📄 metadata.json +``` + +### 3.3 metadata.json 格式 + +**相对路径**（推荐）： +```json +[ + { + "file_path": "train/video001.mp4", + "text": "A beautiful sunset over the ocean, golden hour lighting", + "type": "video" + }, + { + "file_path": "train/video002.mp4", + "text": "A person walking through a forest, cinematic view", + "type": "video" + } +] +``` + +**绝对路径**： +```json +[ + { + "file_path": "/mnt/data/videos/sunset.mp4", + "text": "A beautiful sunset over the ocean", + "type": "video" + } +] +``` + +--- + +## 四、训练 + +### 4.1 快速开始 + +直接执行启动脚本 [train_ar_diffusion.sh](./train_ar_diffusion.sh)： + +```bash +export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" + +accelerate launch --mixed_precision="bf16" --use_fsdp \ + --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \ + --fsdp_transformer_layer_cls_to_wrap=CasualWanAttentionBlock \ + --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT \ + --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False \ + scripts/wan2.1_causal_forcing/train_ar_diffusion.py \ + --config_path="config/wan2.1/wan_civitai.yaml" \ + --pretrained_model_name_or_path=$MODEL_NAME \ + --train_data_dir=$DATASET_NAME \ + --train_data_meta=$DATASET_META_NAME \ + --image_sample_size=640 \ + --video_sample_size=640 \ + --token_sample_size=640 \ + --fix_sample_size 480 832 \ + --video_sample_stride=2 \ + --video_sample_n_frames=81 \ + --train_batch_size=1 \ + --gradient_accumulation_steps=1 \ + --dataloader_num_workers=8 \ + --num_train_epochs=100 \ + --checkpointing_steps=200 \ + --learning_rate=2e-06 \ + --lr_scheduler="constant_with_warmup" \ + --lr_warmup_steps=100 \ + --seed=42 \ + --output_dir="output_dir_wan2.1_causal_forcing_ar_diffusion" \ + --gradient_checkpointing \ + --mixed_precision="bf16" \ + --adam_weight_decay=3e-2 \ + --adam_epsilon=1e-10 \ + --vae_mini_batch=1 \ + --max_grad_norm=0.05 \ + --random_hw_adapt \ + --training_with_video_token_length \ + --enable_bucket \ + --num_frame_per_block=3 \ + --train_sampling_steps=1000 \ + --shift=5.0 \ + --use_timestep_weight \ + --trainable_modules "." +``` + +或者直接执行： + +```bash +bash scripts/wan2.1_causal_forcing/train_ar_diffusion.sh +``` + +### 4.2 主要参数说明 + +| 参数 | 说明 | 示例值 | +|------|------|-------| +| `--pretrained_model_name_or_path` | 初始化基础模型 | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B/` | +| `--config_path` | 模型配置 YAML | `config/wan2.1/wan_civitai.yaml` | +| `--train_data_dir` | 视频根目录（拼接到 `file_path` 之前） | `""` | +| `--train_data_meta` | `metadata.json` 路径 | `datasets/X-Fun-Videos-Demo/metadata.json` | +| `--train_batch_size` | 每卡 batch size | 1 | +| `--gradient_accumulation_steps` | 梯度累积步数 | 1 | +| `--dataloader_num_workers` | DataLoader 子进程数 | 8 | +| `--num_train_epochs` | 训练 epoch 数 | 100 | +| `--checkpointing_steps` | 每 N 步保存 checkpoint | 200 | +| `--learning_rate` | 初始学习率 | 2e-06 | +| `--lr_scheduler` | 学习率调度器 | `constant_with_warmup` | +| `--lr_warmup_steps` | 学习率预热步数 | 100 | +| `--output_dir` | 输出目录 | `output_dir_wan2.1_causal_forcing_ar_diffusion` | +| `--gradient_checkpointing` | 启用激活重计算 | - | +| `--mixed_precision` | `fp16` / `bf16` | `bf16` | +| `--adam_weight_decay` | AdamW 权重衰减 | 3e-2 | +| `--adam_epsilon` | AdamW epsilon | 1e-10 | +| `--max_grad_norm` | 梯度裁剪阈值 | 0.05 | +| `--trainable_modules` | 可训练模块（`"."` 表示全量） | `"."` | + +**视频采样参数**： + +| 参数 | 说明 | 示例值 | +|------|------|-------| +| `--image_sample_size` | 图像采样尺寸 | 640 | +| `--video_sample_size` | 视频采样尺寸 | 640 | +| `--token_sample_size` | Token 采样尺寸 | 640 | +| `--fix_sample_size` | 固定输出 `[高度, 宽度]` | `480 832` | +| `--video_sample_stride` | 帧采样步长 | 2 | +| `--video_sample_n_frames` | 视频帧数 | 81 | +| `--random_hw_adapt` | 启用随机分辨率自适应 | - | +| `--training_with_video_token_length` | 启用基于 token 长度的训练 | - | +| `--enable_bucket` | 启用宽高比 bucket 采样 | - | +| `--vae_mini_batch` | VAE 编码 mini-batch 大小（设为 1 可避免显存溢出） | 1 | + +### 4.3 Causal-Forcing 特有参数 + +| 参数 | 说明 | 示例值 | +|------|------|-------| +| `--num_frame_per_block` | 每个因果块的帧数。`3` = chunkwise，`1` = framewise | 3 | +| `--independent_first_frame` | 第一帧独立（`[1, N, N, ...]` 块模式，I2V 场景适用） | - | +| `--shift` | `FlowMatchEulerDiscreteScheduler` 的 shift 值（Causal-Forcing 默认 5.0） | 5.0 | +| `--train_sampling_steps` | flow matching 调度器总时间步数 | 1000 | +| `--use_timestep_weight` | 启用逐时间步高斯损失权重（以 T/2 为中心） | - | +| `--no_teacher_forcing` | 关闭 teacher forcing（改为 diffusion forcing） | - | +| `--noise_augmentation_max_timestep` | teacher forcing 时向干净上下文 token 添加轻微噪声（0 = 关闭） | 0 | + +> **Teacher Forcing 说明**：默认情况下，第一阶段使用 teacher forcing — 模型在预测当前块时接收 **干净的 GT latent** 作为上下文（`clean_x`）。这可以稳定早期训练。如需使用 diffusion forcing，可通过 `--no_teacher_forcing` 关闭。 + +### 4.4 使用 FSDP 训练 + +上述脚本已配置 FSDP，使用 `CasualWanAttentionBlock` 自动包装。单卡训练无需 FSDP： + +```bash +accelerate launch --mixed_precision="bf16" \ + scripts/wan2.1_causal_forcing/train_ar_diffusion.py \ + ... +``` + +--- + +## 五、使用训练好的 Checkpoint + +第一阶段输出的 checkpoint 用于初始化 **第二阶段（因果一致性蒸馏）**。在 `train_causal_consistency_distill.py` 中通过 `--transformer_path` 和 `--teacher_transformer_path` 指定： + +```bash +--transformer_path="output_dir_wan2.1_causal_forcing_ar_diffusion/checkpoint-{N}/diffusion_pytorch_model.safetensors" \ +--teacher_transformer_path="output_dir_wan2.1_causal_forcing_ar_diffusion/checkpoint-{N}/diffusion_pytorch_model.safetensors" +``` + +完整的第二阶段流程参见 [train_causal_consistency_distill.sh](./train_causal_consistency_distill.sh)。 + +--- + +## 六、更多资源 + +- **Causal-Forcing 论文**：https://github.com/thu-ml/Causal-Forcing +- **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun diff --git a/scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md b/scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md new file mode 100644 index 00000000..83659ab3 --- /dev/null +++ b/scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md @@ -0,0 +1,288 @@ +# Wan2.1 Causal-Forcing Stage 2: Causal Consistency Distillation (CCD) Training Guide + +This document provides the complete workflow for **Causal-Forcing Stage 2 — Causal Consistency Distillation (CCD)** on Wan2.1-T2V-1.3B. + +> **What is Causal Consistency Distillation?** +> +> CCD is the **second stage** of the Causal-Forcing pipeline, which compresses a video diffusion model into a **few-step causal autoregressive generator**: +> +> 1. **Stage 1 — AR Diffusion** (`train_ar_diffusion.py`): Train the model with teacher forcing on clean video latents, producing a strong AR backbone. +> 2. **Stage 2 — Causal Consistency Distillation** (`train_causal_consistency_distill.py`): Distill the multi-step AR model into a **one-step-per-block** consistency model using an EMA teacher with CFG guidance. +> 3. **Stage 3 — Causal DMD** (`train_causal_dmd.py`): Further distill to a **2-step** generator using distribution matching with a 14B teacher. +> +> This README covers **Stage 2 only**. See [README_TRAIN_AR_DIFFUSION.md](./README_TRAIN_AR_DIFFUSION.md) for Stage 1. + +--- + +## Table of Contents +- [1. Prerequisites](#1-prerequisites) +- [2. Environment Setup](#2-environment-setup) +- [3. Download Pretrained Models](#3-download-pretrained-models) +- [4. Prepare Training Data](#4-prepare-training-data) + - [4.1 Quick Demo Dataset](#41-quick-demo-dataset) + - [4.2 Dataset Structure](#42-dataset-structure) + - [4.3 metadata.json Format](#43-metadatajson-format) +- [5. Training](#5-training) + - [5.1 Quick Start](#51-quick-start) + - [5.2 Key Parameters](#52-key-parameters) + - [5.3 CCD-Specific Parameters](#53-ccd-specific-parameters) +- [6. Use the Trained Checkpoint](#6-use-the-trained-checkpoint) +- [7. Additional Resources](#7-additional-resources) + +--- + +## 1. Prerequisites + +Stage 2 requires a **Stage 1 AR Diffusion checkpoint** to initialize both the generator and the EMA teacher. + +```bash +# Example: Stage 1 checkpoint from AR Diffusion training +export STAGE1_CKPT="output_dir_wan2.1_causal_forcing_ar_diffusion/checkpoint-2000/diffusion_pytorch_model.safetensors" +``` + +See [README_TRAIN_AR_DIFFUSION.md](./README_TRAIN_AR_DIFFUSION.md) for how to produce this checkpoint. + +--- + +## 2. Environment Setup + +**Method 1: Using requirements.txt** + +```bash +pip install -r requirements.txt +``` + +**Method 2: Manual Installation** + +```bash +pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image +pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime +pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" +pip install yunchang xfuser modelscope openpyxl +pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y +pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 +``` + +**Method 3: Using Docker** + +```bash +# pull image +docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun + +# enter image +docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun +``` + +--- + +## 3. Download Pretrained Models + +CCD initializes the generator and teacher from Wan2.1-T2V-1.3B, then loads the Stage 1 checkpoint on top. + +```bash +# Create model directory +mkdir -p models/Diffusion_Transformer + +# Download Wan2.1 T2V base model +modelscope download --model Wan-AI/Wan2.1-T2V-1.3B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-1.3B +``` + +--- + +## 4. Prepare Training Data + +CCD trains directly on **raw videos** with online VAE encoding — same data format as Stage 1. + +### 4.1 Quick Demo Dataset + +We provide a small demo dataset for quick testing: + +```bash +# Download official demo dataset +modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo +``` + +### 4.2 Dataset Structure + +``` +📦 datasets/ +├── 📂 X-Fun-Videos-Demo/ +│ ├── 📂 train/ +│ │ ├── 📄 video001.mp4 +│ │ ├── 📄 video002.mp4 +│ │ └── 📄 ... +│ └── 📄 metadata.json +``` + +### 4.3 metadata.json Format + +**Relative paths** (recommended): +```json +[ + { + "file_path": "train/video001.mp4", + "text": "A beautiful sunset over the ocean, golden hour lighting", + "type": "video" + } +] +``` + +**Absolute paths**: +```json +[ + { + "file_path": "/mnt/data/videos/sunset.mp4", + "text": "A beautiful sunset over the ocean", + "type": "video" + } +] +``` + +--- + +## 5. Training + +### 5.1 Quick Start + +The ready-to-use launcher is [train_causal_consistency_distill.sh](./train_causal_consistency_distill.sh): + +```bash +export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" +export STAGE1_CKPT="output_dir_wan2.1_causal_forcing_ar_diffusion/checkpoint-2000/diffusion_pytorch_model.safetensors" + +accelerate launch --mixed_precision="bf16" scripts/wan2.1_causal_forcing/train_causal_consistency_distill.py \ + --config_path="config/wan2.1/wan_civitai.yaml" \ + --pretrained_model_name_or_path=$MODEL_NAME \ + --train_data_dir=$DATASET_NAME \ + --train_data_meta=$DATASET_META_NAME \ + --transformer_path=$STAGE1_CKPT \ + --teacher_transformer_path=$STAGE1_CKPT \ + --image_sample_size=640 \ + --video_sample_size=640 \ + --token_sample_size=640 \ + --fix_sample_size 480 832 \ + --video_sample_stride=2 \ + --video_sample_n_frames=81 \ + --train_batch_size=1 \ + --gradient_accumulation_steps=1 \ + --dataloader_num_workers=8 \ + --num_train_epochs=100 \ + --checkpointing_steps=500 \ + --learning_rate=2.0e-06 \ + --lr_scheduler="constant_with_warmup" \ + --lr_warmup_steps=100 \ + --seed=42 \ + --output_dir="output_dir_wan2.1_causal_forcing_ccd" \ + --gradient_checkpointing \ + --mixed_precision="bf16" \ + --adam_weight_decay=0.0 \ + --adam_beta1=0.0 \ + --adam_beta2=0.999 \ + --adam_epsilon=1e-10 \ + --vae_mini_batch=1 \ + --max_grad_norm=10.0 \ + --random_hw_adapt \ + --training_with_video_token_length \ + --enable_bucket \ + --num_frame_per_block=3 \ + --shift=5.0 \ + --discrete_cd_N=48 \ + --guidance_scale=3.0 \ + --ema_weight=0.99 \ + --ema_start_step=200 \ + --trainable_modules "." +``` + +Or simply: + +```bash +bash scripts/wan2.1_causal_forcing/train_causal_consistency_distill.sh +``` + +### 5.2 Key Parameters + +| Parameter | Description | Example Value | +|-----------|-------------|---------------| +| `--pretrained_model_name_or_path` | Base model to initialize | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B/` | +| `--config_path` | Model config YAML | `config/wan2.1/wan_civitai.yaml` | +| `--train_data_dir` | Video root directory (prepended to relative `file_path`) | `""` | +| `--train_data_meta` | Path to `metadata.json` | `datasets/X-Fun-Videos-Demo/metadata.json` | +| `--transformer_path` | Stage 1 checkpoint for generator init | `$STAGE1_CKPT` | +| `--teacher_transformer_path` | Stage 1 checkpoint for teacher init (defaults to generator) | `$STAGE1_CKPT` | +| `--train_batch_size` | Per-GPU batch size | 1 | +| `--gradient_accumulation_steps` | Gradient accumulation steps | 1 | +| `--dataloader_num_workers` | DataLoader workers | 8 | +| `--num_train_epochs` | Number of training epochs | 100 | +| `--checkpointing_steps` | Save checkpoint every N steps | 500 | +| `--learning_rate` | Initial learning rate | 2e-06 | +| `--lr_scheduler` | LR scheduler type | `constant_with_warmup` | +| `--lr_warmup_steps` | LR warmup steps | 100 | +| `--output_dir` | Output directory | `output_dir_wan2.1_causal_forcing_ccd` | +| `--gradient_checkpointing` | Enable activation checkpointing | - | +| `--mixed_precision` | `fp16` / `bf16` | `bf16` | +| `--adam_weight_decay` | AdamW weight decay | 0.0 | +| `--adam_beta1` | AdamW beta1 (CCD uses 0.0) | 0.0 | +| `--adam_beta2` | AdamW beta2 | 0.999 | +| `--adam_epsilon` | AdamW epsilon | 1e-10 | +| `--max_grad_norm` | Gradient clipping threshold | 10.0 | +| `--trainable_modules` | Trainable modules (`"."` = all) | `"."` | +| `--low_vram` | Enable low VRAM mode (offload VAE/text encoder) | - | + +**Video Sampling Parameters**: + +| Parameter | Description | Example Value | +|-----------|-------------|---------------| +| `--image_sample_size` | Image sampling size | 640 | +| `--video_sample_size` | Video sampling size | 640 | +| `--token_sample_size` | Token sampling size | 640 | +| `--fix_sample_size` | Fixed `[height, width]` for output | `480 832` | +| `--video_sample_stride` | Frame sampling stride | 2 | +| `--video_sample_n_frames` | Number of video frames | 81 | +| `--random_hw_adapt` | Enable random resolution adaptation | - | +| `--training_with_video_token_length` | Enable token-length-based training | - | +| `--enable_bucket` | Enable aspect-ratio bucket sampling | - | +| `--vae_mini_batch` | VAE encoding mini-batch size (1 to avoid OOM) | 1 | + +### 5.3 CCD-Specific Parameters + +| Parameter | Description | Example Value | +|-----------|-------------|---------------| +| `--num_frame_per_block` | Frames per causal block. `3` = chunkwise, `1` = framewise | 3 | +| `--independent_first_frame` | First frame is independent (`[1, N, N, ...]` block pattern, useful for I2V) | - | +| `--shift` | `FlowMatchEulerDiscreteScheduler` shift (Causal-Forcing default: 5.0) | 5.0 | +| `--discrete_cd_N` | Number of discrete timesteps for the consistency schedule | 48 | +| `--guidance_scale` | CFG guidance scale used by the EMA teacher | 3.0 | +| `--ema_weight` | EMA decay for the consistency-target generator copy (set <=0 to disable) | 0.99 | +| `--ema_start_step` | Steps to wait before EMA tracking starts | 200 | + +> **How CCD works**: For each training sample, CCD: +> 1. Loads clean video latents (via online VAE encoding). +> 2. Samples a timestep index from the discrete consistency schedule `[0, N-2]`. +> 3. Adds noise to the clean latents at timestep `t` and `t_next`. +> 4. Runs the **generator** on `x_t` to predict `x0`. +> 5. Runs the **EMA teacher** (with CFG guidance) on `x_{t_next}` to produce the consistency target. +> 6. Minimizes the L2 loss between the generator prediction and the teacher target. + +> **EMA Teacher**: The EMA copy tracks the generator with polyak updates (`ema = decay*ema + (1-decay)*gen`). Before `--ema_start_step`, the EMA mirrors the live generator. The teacher uses CFG with `--guidance_scale=3.0` to produce higher-quality targets. + +--- + +## 6. Use the Trained Checkpoint + +The Stage 2 CCD checkpoint is used to initialize **Stage 3 (Causal DMD)**. Pass its path to `train_causal_dmd.py`: + +```bash +--generator_path="output_dir_wan2.1_causal_forcing_ccd/checkpoint-{N}/diffusion_pytorch_model.safetensors" +``` + +See `train_causal_dmd.sh` for the full Stage 3 workflow. + +--- + +## 7. Additional Resources + +- **Causal-Forcing Paper**: https://github.com/thu-ml/Causal-Forcing +- **Official GitHub**: https://github.com/aigc-apps/VideoX-Fun diff --git a/scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md b/scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md new file mode 100644 index 00000000..1a7e70c2 --- /dev/null +++ b/scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md @@ -0,0 +1,288 @@ +# Wan2.1 Causal-Forcing 第二阶段：因果一致性蒸馏（CCD）训练指南 + +本文档介绍 **Causal-Forcing 第二阶段 — 因果一致性蒸馏（CCD）** 的完整流程。 + +> **什么是因果一致性蒸馏？** +> +> CCD 是 Causal-Forcing 流水线的 **第二阶段**，将视频扩散模型压缩为 **少步因果自回归生成器**： +> +> 1. **第一阶段 — AR Diffusion**（`train_ar_diffusion.py`）：在干净视频 latent 上使用 teacher forcing 训练，得到强大的 AR 骨架。 +> 2. **第二阶段 — 因果一致性蒸馏**（`train_causal_consistency_distill.py`）：利用 EMA teacher 和 CFG 引导，将多步 AR 模型蒸馏为 **每块一步** 的一致性模型。 +> 3. **第三阶段 — 因果 DMD**（`train_causal_dmd.py`）：利用 14B 教师模型做分布匹配，进一步蒸馏为 **2 步** 生成器。 +> +> 本文档仅覆盖 **第二阶段**。第一阶段参见 [README_TRAIN_AR_DIFFUSION_zh-CN.md](./README_TRAIN_AR_DIFFUSION_zh-CN.md)。 + +--- + +## 目录 +- [一、前置条件](#一前置条件) +- [二、环境配置](#二环境配置) +- [三、下载预训练模型](#三下载预训练模型) +- [四、准备训练数据](#四准备训练数据) + - [4.1 快速测试数据集](#41-快速测试数据集) + - [4.2 数据集结构](#42-数据集结构) + - [4.3 metadata.json 格式](#43-metadatajson-格式) +- [五、训练](#五训练) + - [5.1 快速开始](#51-快速开始) + - [5.2 主要参数说明](#52-主要参数说明) + - [5.3 CCD 特有参数](#53-ccd-特有参数) +- [六、使用训练好的 Checkpoint](#六使用训练好的-checkpoint) +- [七、更多资源](#七更多资源) + +--- + +## 一、前置条件 + +第二阶段需要 **第一阶段 AR Diffusion 的 checkpoint** 来初始化 generator 和 EMA teacher。 + +```bash +# 示例：第一阶段 AR Diffusion 训练输出的 checkpoint +export STAGE1_CKPT="output_dir_wan2.1_causal_forcing_ar_diffusion/checkpoint-2000/diffusion_pytorch_model.safetensors" +``` + +第一阶段的训练方法参见 [README_TRAIN_AR_DIFFUSION_zh-CN.md](./README_TRAIN_AR_DIFFUSION_zh-CN.md)。 + +--- + +## 二、环境配置 + +**方式 1：使用 requirements.txt** + +```bash +pip install -r requirements.txt +``` + +**方式 2：手动安装依赖** + +```bash +pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image +pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime +pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" +pip install yunchang xfuser modelscope openpyxl +pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y +pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 +``` + +**方式 3：使用 docker** + +```bash +# 拉取镜像 +docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun + +# 进入容器 +docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun +``` + +--- + +## 三、下载预训练模型 + +CCD 以 Wan2.1-T2V-1.3B 初始化 generator 和 teacher，然后加载第一阶段 checkpoint。 + +```bash +# 创建模型目录 +mkdir -p models/Diffusion_Transformer + +# 下载 Wan2.1 T2V 基础模型 +modelscope download --model Wan-AI/Wan2.1-T2V-1.3B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-1.3B +``` + +--- + +## 四、准备训练数据 + +CCD 直接在 **原始视频** 上训练，VAE 在线编码 — 数据格式与第一阶段相同。 + +### 4.1 快速测试数据集 + +我们提供了一个小型 demo 数据集用于快速测试： + +```bash +# 下载官方示例数据集 +modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo +``` + +### 4.2 数据集结构 + +``` +📦 datasets/ +├── 📂 X-Fun-Videos-Demo/ +│ ├── 📂 train/ +│ │ ├── 📄 video001.mp4 +│ │ ├── 📄 video002.mp4 +│ │ └── 📄 ... +│ └── 📄 metadata.json +``` + +### 4.3 metadata.json 格式 + +**相对路径**（推荐）： +```json +[ + { + "file_path": "train/video001.mp4", + "text": "A beautiful sunset over the ocean, golden hour lighting", + "type": "video" + } +] +``` + +**绝对路径**： +```json +[ + { + "file_path": "/mnt/data/videos/sunset.mp4", + "text": "A beautiful sunset over the ocean", + "type": "video" + } +] +``` + +--- + +## 五、训练 + +### 5.1 快速开始 + +直接执行启动脚本 [train_causal_consistency_distill.sh](./train_causal_consistency_distill.sh)： + +```bash +export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" +export STAGE1_CKPT="output_dir_wan2.1_causal_forcing_ar_diffusion/checkpoint-2000/diffusion_pytorch_model.safetensors" + +accelerate launch --mixed_precision="bf16" scripts/wan2.1_causal_forcing/train_causal_consistency_distill.py \ + --config_path="config/wan2.1/wan_civitai.yaml" \ + --pretrained_model_name_or_path=$MODEL_NAME \ + --train_data_dir=$DATASET_NAME \ + --train_data_meta=$DATASET_META_NAME \ + --transformer_path=$STAGE1_CKPT \ + --teacher_transformer_path=$STAGE1_CKPT \ + --image_sample_size=640 \ + --video_sample_size=640 \ + --token_sample_size=640 \ + --fix_sample_size 480 832 \ + --video_sample_stride=2 \ + --video_sample_n_frames=81 \ + --train_batch_size=1 \ + --gradient_accumulation_steps=1 \ + --dataloader_num_workers=8 \ + --num_train_epochs=100 \ + --checkpointing_steps=500 \ + --learning_rate=2.0e-06 \ + --lr_scheduler="constant_with_warmup" \ + --lr_warmup_steps=100 \ + --seed=42 \ + --output_dir="output_dir_wan2.1_causal_forcing_ccd" \ + --gradient_checkpointing \ + --mixed_precision="bf16" \ + --adam_weight_decay=0.0 \ + --adam_beta1=0.0 \ + --adam_beta2=0.999 \ + --adam_epsilon=1e-10 \ + --vae_mini_batch=1 \ + --max_grad_norm=10.0 \ + --random_hw_adapt \ + --training_with_video_token_length \ + --enable_bucket \ + --num_frame_per_block=3 \ + --shift=5.0 \ + --discrete_cd_N=48 \ + --guidance_scale=3.0 \ + --ema_weight=0.99 \ + --ema_start_step=200 \ + --trainable_modules "." +``` + +或者直接执行： + +```bash +bash scripts/wan2.1_causal_forcing/train_causal_consistency_distill.sh +``` + +### 5.2 主要参数说明 + +| 参数 | 说明 | 示例值 | +|------|------|-------| +| `--pretrained_model_name_or_path` | 初始化基础模型 | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B/` | +| `--config_path` | 模型配置 YAML | `config/wan2.1/wan_civitai.yaml` | +| `--train_data_dir` | 视频根目录（拼接到 `file_path` 之前） | `""` | +| `--train_data_meta` | `metadata.json` 路径 | `datasets/X-Fun-Videos-Demo/metadata.json` | +| `--transformer_path` | 第一阶段 checkpoint（generator 初始化） | `$STAGE1_CKPT` | +| `--teacher_transformer_path` | 第一阶段 checkpoint（teacher 初始化，默认同 generator） | `$STAGE1_CKPT` | +| `--train_batch_size` | 每卡 batch size | 1 | +| `--gradient_accumulation_steps` | 梯度累积步数 | 1 | +| `--dataloader_num_workers` | DataLoader 子进程数 | 8 | +| `--num_train_epochs` | 训练 epoch 数 | 100 | +| `--checkpointing_steps` | 每 N 步保存 checkpoint | 500 | +| `--learning_rate` | 初始学习率 | 2e-06 | +| `--lr_scheduler` | 学习率调度器 | `constant_with_warmup` | +| `--lr_warmup_steps` | 学习率预热步数 | 100 | +| `--output_dir` | 输出目录 | `output_dir_wan2.1_causal_forcing_ccd` | +| `--gradient_checkpointing` | 启用激活重计算 | - | +| `--mixed_precision` | `fp16` / `bf16` | `bf16` | +| `--adam_weight_decay` | AdamW 权重衰减 | 0.0 | +| `--adam_beta1` | AdamW beta1（CCD 使用 0.0） | 0.0 | +| `--adam_beta2` | AdamW beta2 | 0.999 | +| `--adam_epsilon` | AdamW epsilon | 1e-10 | +| `--max_grad_norm` | 梯度裁剪阈值 | 10.0 | +| `--trainable_modules` | 可训练模块（`"."` 表示全量） | `"."` | +| `--low_vram` | 启用低显存模式（VAE/文本编码器分时加载） | - | + +**视频采样参数**： + +| 参数 | 说明 | 示例值 | +|------|------|-------| +| `--image_sample_size` | 图像采样尺寸 | 640 | +| `--video_sample_size` | 视频采样尺寸 | 640 | +| `--token_sample_size` | Token 采样尺寸 | 640 | +| `--fix_sample_size` | 固定输出 `[高度, 宽度]` | `480 832` | +| `--video_sample_stride` | 帧采样步长 | 2 | +| `--video_sample_n_frames` | 视频帧数 | 81 | +| `--random_hw_adapt` | 启用随机分辨率自适应 | - | +| `--training_with_video_token_length` | 启用基于 token 长度的训练 | - | +| `--enable_bucket` | 启用宽高比 bucket 采样 | - | +| `--vae_mini_batch` | VAE 编码 mini-batch 大小（设为 1 可避免显存溢出） | 1 | + +### 5.3 CCD 特有参数 + +| 参数 | 说明 | 示例值 | +|------|------|-------| +| `--num_frame_per_block` | 每个因果块的帧数。`3` = chunkwise，`1` = framewise | 3 | +| `--independent_first_frame` | 第一帧独立（`[1, N, N, ...]` 块模式，I2V 场景适用） | - | +| `--shift` | `FlowMatchEulerDiscreteScheduler` 的 shift 值（Causal-Forcing 默认 5.0） | 5.0 | +| `--discrete_cd_N` | 一致性调度器的离散时间步数量 | 48 | +| `--guidance_scale` | EMA teacher 使用的 CFG 引导强度 | 3.0 | +| `--ema_weight` | EMA 衰减率（一致性目标 generator 副本，设为 <=0 禁用） | 0.99 | +| `--ema_start_step` | EMA 跟踪启动前的等待步数 | 200 | + +> **CCD 工作原理**：对每个训练样本，CCD 执行以下步骤： +> 1. 加载干净视频 latent（通过 VAE 在线编码）。 +> 2. 从离散一致性调度 `[0, N-2]` 中采样一个时间步索引。 +> 3. 在时间步 `t` 和 `t_next` 分别给干净 latent 加噪。 +> 4. **Generator** 在 `x_t` 上预测 `x0`。 +> 5. **EMA teacher**（带 CFG 引导）在 `x_{t_next}` 上生成一致性目标。 +> 6. 最小化 generator 预测与 teacher 目标之间的 L2 损失。 + +> **EMA Teacher**：EMA 副本通过 polyak 更新跟踪 generator（`ema = decay*ema + (1-decay)*gen`）。在 `--ema_start_step` 之前，EMA 直接镜像 generator。Teacher 使用 `--guidance_scale=3.0` 的 CFG 来生成更高质量的目标。 + +--- + +## 六、使用训练好的 Checkpoint + +第二阶段 CCD 输出的 checkpoint 用于初始化 **第三阶段（因果 DMD）**。在 `train_causal_dmd.py` 中指定路径： + +```bash +--generator_path="output_dir_wan2.1_causal_forcing_ccd/checkpoint-{N}/diffusion_pytorch_model.safetensors" +``` + +完整的第三阶段流程参见 `train_causal_dmd.sh`。 + +--- + +## 七、更多资源 + +- **Causal-Forcing 论文**：https://github.com/thu-ml/Causal-Forcing +- **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun diff --git a/scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL.md b/scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL.md new file mode 100644 index 00000000..d4e8d9f0 --- /dev/null +++ b/scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL.md @@ -0,0 +1,302 @@ +# Wan2.1 Causal-Forcing Stage 3: Distribution Matching Distillation (DMD) Training Guide + +This document provides the complete workflow for **Causal-Forcing Stage 3 — Distribution Matching Distillation (DMD)** on Wan2.1-T2V-1.3B. + +> **What is Distribution Matching Distillation?** +> +> DMD is the **third and final stage** of the Causal-Forcing pipeline, which further compresses a CCD model into a **2-step causal autoregressive generator** using a large (14B) teacher: +> +> 1. **Stage 1 — AR Diffusion** (`train_ar_diffusion.py`): Train the model with teacher forcing on clean video latents, producing a strong AR backbone. +> 2. **Stage 2 — Causal Consistency Distillation** (`train_causal_consistency_distill.py`): Distill the multi-step AR model into a **one-step-per-block** consistency model using an EMA teacher with CFG guidance. +> 3. **Stage 3 — Distribution Matching Distillation** (`train_causal_dmd.py`): Further distill to a **2-step frame-wise** generator using distribution matching with a **14B real-score teacher**. +> +> This README covers **Stage 3 only**. See [README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md](./README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md) for Stage 2 and [README_TRAIN_AR_DIFFUSION.md](./README_TRAIN_AR_DIFFUSION.md) for Stage 1. + +--- + +## Table of Contents +- [1. Prerequisites](#1-prerequisites) +- [2. Environment Setup](#2-environment-setup) +- [3. Download Pretrained Models](#3-download-pretrained-models) +- [4. Prepare Training Data](#4-prepare-training-data) + - [4.1 Quick Demo Dataset](#41-quick-demo-dataset) + - [4.2 Dataset Structure](#42-dataset-structure) + - [4.3 metadata.json Format](#43-metadatajson-format) +- [5. Training](#5-training) + - [5.1 Quick Start](#51-quick-start) + - [5.2 Key Parameters](#52-key-parameters) + - [5.3 DMD-Specific Parameters](#53-dmd-specific-parameters) +- [6. Use the Trained Checkpoint](#6-use-the-trained-checkpoint) +- [7. Additional Resources](#7-additional-resources) + +--- + +## 1. Prerequisites + +Stage 3 requires: + +1. A **Stage 2 CCD checkpoint** to initialize the generator (and critic). +2. A **Wan2.1-T2V-14B** model as the DMD real-score teacher. + +```bash +# Example: Stage 2 checkpoint from CCD training +export STAGE2_CKPT="output_dir_wan2.1_causal_forcing_ccd/checkpoint-2000/diffusion_pytorch_model.safetensors" +``` + +See [README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md](./README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md) for how to produce this checkpoint. + +--- + +## 2. Environment Setup + +**Method 1: Using requirements.txt** + +```bash +pip install -r requirements.txt +``` + +**Method 2: Manual Installation** + +```bash +pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image +pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime +pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" +pip install yunchang xfuser modelscope openpyxl +pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y +pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 +``` + +**Method 3: Using Docker** + +```bash +# pull image +docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun + +# enter image +docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun +``` + +--- + +## 3. Download Pretrained Models + +DMD requires **two** pretrained models: + +- **Wan2.1-T2V-1.3B**: base model for the generator/critic. +- **Wan2.1-T2V-14B**: the non-causal real-score teacher used by DMD to compute the real distribution score. + +```bash +# Create model directory +mkdir -p models/Diffusion_Transformer + +# Download Wan2.1 T2V 1.3B (student base model) +modelscope download --model Wan-AI/Wan2.1-T2V-1.3B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-1.3B + +# Download Wan2.1 T2V 14B (DMD real-score teacher) +modelscope download --model Wan-AI/Wan2.1-T2V-14B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-14B +``` + +The Stage 2 CCD checkpoint (generator/critic init) is loaded via `--transformer_path`. + +--- + +## 4. Prepare Training Data + +DMD uses a **TextDataset** (`train_mode="normal"`) — it only needs prompts, not video data, because the generator creates its own training samples via autoregressive rollout. However, a `metadata.json` with prompts is still required. + +### 4.1 Quick Demo Dataset + +```bash +# Download official demo dataset +modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo +``` + +### 4.2 Dataset Structure + +Same as Stage 1/2. See [README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md](./README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md#42-dataset-structure) for details. + +### 4.3 metadata.json Format + +DMD only uses the `"text"` field from each entry. The `file_path` and other fields are ignored when `--train_mode="normal"` (TextDataset mode). + +```json +[ + { + "text": "A beautiful sunset over the ocean, golden hour lighting" + }, + { + "text": "A person walking through a forest, cinematic view" + } +] +``` + +> **Note**: You can reuse any video metadata.json — only the `text` field matters for DMD. + +--- + +## 5. Training + +### 5.1 Quick Start + +The ready-to-use launcher is [train_causal_dmd.sh](./train_causal_dmd.sh): + +```bash +export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/" +export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" +export STAGE2_CKPT="output_dir_wan2.1_causal_forcing_ccd/checkpoint-2000/diffusion_pytorch_model.safetensors" + +accelerate launch --mixed_precision="bf16" --use_fsdp \ + --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \ + --fsdp_transformer_layer_cls_to_wrap=CasualWanAttentionBlock \ + --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT \ + --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False \ + scripts/wan2.1_causal_forcing/train_causal_dmd.py \ + --config_path="config/wan2.1/wan_civitai.yaml" \ + --pretrained_model_name_or_path=$MODEL_NAME \ + --real_score_pretrained_model_name_or_path=$REAL_SCORE_MODEL_NAME \ + --train_data_dir=$DATASET_NAME \ + --train_data_meta=$DATASET_META_NAME \ + --transformer_path=$STAGE2_CKPT \ + --image_sample_size=640 \ + --video_sample_size=640 \ + --token_sample_size=640 \ + --fix_sample_size 480 832 \ + --video_sample_stride=2 \ + --video_sample_n_frames=81 \ + --train_batch_size=1 \ + --gradient_accumulation_steps=1 \ + --dataloader_num_workers=8 \ + --num_train_epochs=100 \ + --checkpointing_steps=200 \ + --learning_rate=2.0e-06 \ + --learning_rate_critic=2.0e-06 \ + --lr_scheduler="constant_with_warmup" \ + --lr_warmup_steps=100 \ + --seed=42 \ + --output_dir="output_dir_wan2.1_causal_forcing_dmd" \ + --gradient_checkpointing \ + --mixed_precision="bf16" \ + --adam_weight_decay=0.0 \ + --adam_beta1=0.0 \ + --adam_beta2=0.999 \ + --adam_epsilon=1e-10 \ + --vae_mini_batch=1 \ + --max_grad_norm=10.0 \ + --random_hw_adapt \ + --training_with_video_token_length \ + --enable_bucket \ + --num_frame_per_block=1 \ + --use_kv_cache_training \ + --denoising_step_indices_list 1000 500 \ + --real_guidance_scale=3.0 \ + --fake_guidance_scale=0.0 \ + --gen_update_interval=5 \ + --resume_from_checkpoint="latest" \ + --trainable_modules "." +``` + +Or simply: + +```bash +bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh +``` + +### 5.2 Key Parameters + +| Parameter | Description | Example Value | +|-----------|-------------|---------------| +| `--pretrained_model_name_or_path` | Base model (1.3B) for generator/critic init | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B/` | +| `--real_score_pretrained_model_name_or_path` | 14B real-score teacher for DMD | `models/Diffusion_Transformer/Wan2.1-T2V-14B` | +| `--config_path` | Model config YAML | `config/wan2.1/wan_civitai.yaml` | +| `--train_data_dir` | Data root directory | `""` | +| `--train_data_meta` | Path to `metadata.json` (prompts only) | `datasets/X-Fun-Videos-Demo/metadata.json` | +| `--transformer_path` | Stage 2 CCD checkpoint for generator/critic init | `$STAGE2_CKPT` | +| `--train_batch_size` | Per-GPU batch size | 1 | +| `--gradient_accumulation_steps` | Gradient accumulation steps | 1 | +| `--dataloader_num_workers` | DataLoader workers | 8 | +| `--num_train_epochs` | Number of training epochs | 100 | +| `--checkpointing_steps` | Save checkpoint every N steps | 200 | +| `--learning_rate` | Generator learning rate | 2e-06 | +| `--learning_rate_critic` | Critic learning rate | 2e-06 | +| `--lr_scheduler` | LR scheduler type | `constant_with_warmup` | +| `--lr_warmup_steps` | LR warmup steps | 100 | +| `--output_dir` | Output directory | `output_dir_wan2.1_causal_forcing_dmd` | +| `--gradient_checkpointing` | Enable activation checkpointing | - | +| `--mixed_precision` | `fp16` / `bf16` | `bf16` | +| `--adam_weight_decay` | AdamW weight decay | 0.0 | +| `--adam_beta1` | AdamW beta1 (DMD uses 0.0) | 0.0 | +| `--adam_beta2` | AdamW beta2 | 0.999 | +| `--adam_epsilon` | AdamW epsilon | 1e-10 | +| `--max_grad_norm` | Gradient clipping threshold | 10.0 | +| `--trainable_modules` | Trainable modules (`"."` = all) | `"."` | +| `--low_vram` | Enable low VRAM mode (offload VAE/text encoder) | - | + +**Video Sampling Parameters**: + +| Parameter | Description | Example Value | +|-----------|-------------|---------------| +| `--image_sample_size` | Image sampling size | 640 | +| `--video_sample_size` | Video sampling size | 640 | +| `--token_sample_size` | Token sampling size | 640 | +| `--fix_sample_size` | Fixed `[height, width]` for output | `480 832` | +| `--video_sample_stride` | Frame sampling stride | 2 | +| `--video_sample_n_frames` | Number of video frames | 81 | +| `--random_hw_adapt` | Enable random resolution adaptation | - | +| `--training_with_video_token_length` | Enable token-length-based training | - | +| `--enable_bucket` | Enable aspect-ratio bucket sampling | - | +| `--vae_mini_batch` | VAE encoding mini-batch size (1 to avoid OOM) | 1 | + +### 5.3 DMD-Specific Parameters + +| Parameter | Description | Example Value | +|-----------|-------------|---------------| +| `--denoising_step_indices_list` | Denoising step indices (DMD core param). `[1000, 500]` = 2-step DMD | `1000 500` | +| `--real_guidance_scale` | CFG scale for the real-score (14B teacher) | 3.0 | +| `--fake_guidance_scale` | CFG scale for the fake-score (generator). 0.0 = no CFG | 0.0 | +| `--gen_update_interval` | Generator update interval (generator updates every N critic steps) | 5 | +| `--num_frame_per_block` | Frames per causal block. `1` = frame-wise (DMD default) | 1 | +| `--use_kv_cache_training` | Use KV cache block-by-block training (matches original Self-Forcing) | - | +| `--independent_first_frame` | First frame is independent (`[1, N, N, ...]` block pattern, useful for I2V) | - | +| `--context_noise` | Context noise level for KV cache update | 0 | +| `--use_teacher_forcing` | Enable teacher forcing (pass clean_x to transformer) | - | +| `--teacher_forcing_prob` | Probability of applying teacher forcing per step (1.0 = always) | 1.0 | +| `--train_mode` | Training mode: `normal` (TextDataset, prompt-only) or `i2v` | `normal` | +| `--resume_from_checkpoint` | Resume from checkpoint. Use `"latest"` to auto-select | `"latest"` | + +--- + +## 6. Use the Trained Checkpoint + +The Stage 3 DMD checkpoint is the **final model** in the Causal-Forcing pipeline. Use it for inference: + +```python +# In examples/wan2.1_causal_forcing/predict_t2v.py +transformer_path = "output_dir_wan2.1_causal_forcing_dmd/checkpoint-{N}/diffusion_pytorch_model.safetensors" + +# DMD Stage 3 inference config +guidance_scale = 1.0 # CFG is baked into distilled weights +num_inference_steps = 2 # 2-step DMD +stochastic_sampling = True +num_frame_per_block = 1 # Frame-wise generation +``` + +Or run: + +```bash +python examples/wan2.1_causal_forcing/predict_t2v.py +``` + +> **Stage Selector** in `predict_t2v.py` provides preset configs for all stages: +> - **Stage 1 (AR Diffusion)**: `guidance_scale=3.0`, `num_inference_steps=50`, `stochastic_sampling=False` +> - **Stage 2 (CCD)**: `guidance_scale=1.0`, `num_inference_steps=4`, `stochastic_sampling=True` +> - **Stage 3 (DMD)**: `guidance_scale=1.0`, `num_inference_steps=2`, `stochastic_sampling=True` + +--- + +## 7. Additional Resources + +- **Causal-Forcing Paper**: https://github.com/thu-ml/Causal-Forcing +- **Official GitHub**: https://github.com/aigc-apps/VideoX-Fun diff --git a/scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md b/scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md new file mode 100644 index 00000000..a1e72e8e --- /dev/null +++ b/scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md @@ -0,0 +1,301 @@ +# Wan2.1 Causal-Forcing Stage 3: 分布匹配蒸馏 (DMD) 训练指南 + +本文档提供了将 Wan2.1 进行 **Causal-Forcing Stage 3 — 分布匹配蒸馏 (DMD)** 的完整工作流。 + +> **什么是分布匹配蒸馏？** +> +> DMD 是 Causal-Forcing 流水线的**第三阶段（最终阶段）**，使用大规模（14B）teacher 将 CCD 模型进一步蒸馏为 **2 步因果自回归生成器**： +> +> 1. **Stage 1 — AR Diffusion** (`train_ar_diffusion.py`)：在干净视频 latent 上以 teacher forcing 训练模型，产出强 AR 基座。 +> 2. **Stage 2 — 因果一致性蒸馏 (CCD)** (`train_causal_consistency_distill.py`)：使用 EMA teacher + CFG 将多步 AR 模型蒸馏为**每块一步**的一致性模型。 +> 3. **Stage 3 — 分布匹配蒸馏 (DMD)** (`train_causal_dmd.py`)：使用 **14B real-score teacher** 的分布匹配进一步蒸馏为 **2 步逐帧**生成器。 +> +> 本文档仅覆盖 **Stage 3**。Stage 2 请参见 [README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md](./README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md)，Stage 1 请参见 [README_TRAIN_AR_DIFFUSION_zh-CN.md](./README_TRAIN_AR_DIFFUSION_zh-CN.md)。 + +--- + +## 目录 +- [一、前置条件](#一前置条件) +- [二、环境配置](#二环境配置) +- [三、下载预训练模型](#三下载预训练模型) +- [四、数据准备](#四数据准备) + - [4.1 快速测试数据集](#41-快速测试数据集) + - [4.2 数据集结构](#42-数据集结构) + - [4.3 metadata.json 格式](#43-metadatajson-格式) +- [五、训练](#五训练) + - [5.1 快速开始](#51-快速开始) + - [5.2 关键参数](#52-关键参数) + - [5.3 DMD 特有参数](#53-dmd-特有参数) +- [六、使用训练好的 Checkpoint](#六使用训练好的-checkpoint) +- [七、更多资源](#七更多资源) + +--- + +## 一、前置条件 + +Stage 3 需要： + +1. **Stage 2 CCD checkpoint**：用于初始化生成器（和判别器）。 +2. **Wan2.1-T2V-14B** 模型：作为 DMD real-score teacher。 + +```bash +# 示例：Stage 2 CCD 训练的 checkpoint +export STAGE2_CKPT="output_dir_wan2.1_causal_forcing_ccd/checkpoint-2000/diffusion_pytorch_model.safetensors" +``` + +如何产出该 checkpoint 请参见 [README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md](./README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md)。 + +--- + +## 二、环境配置 + +**方式 1：使用 requirements.txt** + +```bash +pip install -r requirements.txt +``` + +**方式 2：手动安装依赖** + +```bash +pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image +pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime +pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2" +pip install yunchang xfuser modelscope openpyxl +pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y +pip install opencv-python-headless +pip install deepspeed==0.17.0 numpy==1.26.4 +``` + +**方式 3：使用 Docker** + +```bash +# 拉取镜像 +docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun + +# 进入容器 +docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun +``` + +--- + +## 三、下载预训练模型 + +DMD 需要**两个**预训练模型： + +- **Wan2.1-T2V-1.3B**：生成器/判别器的基础模型。 +- **Wan2.1-T2V-14B**：非因果的 real-score teacher，用于计算真实分布得分。 + +```bash +# 创建模型目录 +mkdir -p models/Diffusion_Transformer + +# 下载 Wan2.1 T2V 1.3B（student 基础模型） +modelscope download --model Wan-AI/Wan2.1-T2V-1.3B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-1.3B + +# 下载 Wan2.1 T2V 14B（DMD real-score teacher） +modelscope download --model Wan-AI/Wan2.1-T2V-14B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-14B +``` + +Stage 2 CCD checkpoint（生成器/判别器初始化）通过 `--transformer_path` 加载。 + +--- + +## 四、数据准备 + +DMD 使用 **TextDataset**（`train_mode="normal"`）— 只需要提示词，不需要视频数据，因为生成器通过自回归 rollout 自行创建训练样本。但仍需提供包含提示词的 `metadata.json`。 + +### 4.1 快速测试数据集 + +```bash +# 下载官方示例数据集 +modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo +``` + +### 4.2 数据集结构 + +与 Stage 1/2 相同。详见 [README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md](./README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md#42-数据集结构)。 + +### 4.3 metadata.json 格式 + +DMD 仅使用每条记录的 `"text"` 字段。当 `--train_mode="normal"`（TextDataset 模式）时，`file_path` 和其他字段会被忽略。 + +```json +[ + { + "text": "A beautiful sunset over the ocean, golden hour lighting" + }, + { + "text": "A person walking through a forest, cinematic view" + } +] +``` + +> **说明**：你可以复用任何视频 metadata.json — DMD 只关心 `text` 字段。 + +--- + +## 五、训练 + +### 5.1 快速开始 + +可直接使用的启动脚本为 [train_causal_dmd.sh](./train_causal_dmd.sh)： + +```bash +export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/" +export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" +export STAGE2_CKPT="output_dir_wan2.1_causal_forcing_ccd/checkpoint-2000/diffusion_pytorch_model.safetensors" + +accelerate launch --mixed_precision="bf16" --use_fsdp \ + --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \ + --fsdp_transformer_layer_cls_to_wrap=CasualWanAttentionBlock \ + --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT \ + --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False \ + scripts/wan2.1_causal_forcing/train_causal_dmd.py \ + --config_path="config/wan2.1/wan_civitai.yaml" \ + --pretrained_model_name_or_path=$MODEL_NAME \ + --real_score_pretrained_model_name_or_path=$REAL_SCORE_MODEL_NAME \ + --train_data_dir=$DATASET_NAME \ + --train_data_meta=$DATASET_META_NAME \ + --transformer_path=$STAGE2_CKPT \ + --image_sample_size=640 \ + --video_sample_size=640 \ + --token_sample_size=640 \ + --fix_sample_size 480 832 \ + --video_sample_stride=2 \ + --video_sample_n_frames=81 \ + --train_batch_size=1 \ + --gradient_accumulation_steps=1 \ + --dataloader_num_workers=8 \ + --num_train_epochs=100 \ + --checkpointing_steps=200 \ + --learning_rate=2.0e-06 \ + --learning_rate_critic=2.0e-06 \ + --lr_scheduler="constant_with_warmup" \ + --lr_warmup_steps=100 \ + --seed=42 \ + --output_dir="output_dir_wan2.1_causal_forcing_dmd" \ + --gradient_checkpointing \ + --mixed_precision="bf16" \ + --adam_weight_decay=0.0 \ + --adam_beta1=0.0 \ + --adam_beta2=0.999 \ + --adam_epsilon=1e-10 \ + --vae_mini_batch=1 \ + --max_grad_norm=10.0 \ + --random_hw_adapt \ + --training_with_video_token_length \ + --enable_bucket \ + --num_frame_per_block=1 \ + --use_kv_cache_training \ + --denoising_step_indices_list 1000 500 \ + --real_guidance_scale=3.0 \ + --fake_guidance_scale=0.0 \ + --gen_update_interval=5 \ + --resume_from_checkpoint="latest" \ + --trainable_modules "." +``` + +或直接运行： + +```bash +bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh +``` + +### 5.2 关键参数 + +| 参数 | 说明 | 示例值 | +|------|------|--------| +| `--pretrained_model_name_or_path` | 基础模型（1.3B），用于生成器/判别器初始化 | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B/` | +| `--real_score_pretrained_model_name_or_path` | 14B real-score teacher | `models/Diffusion_Transformer/Wan2.1-T2V-14B` | +| `--config_path` | 模型配置 YAML | `config/wan2.1/wan_civitai.yaml` | +| `--train_data_dir` | 数据根目录 | `""` | +| `--train_data_meta` | `metadata.json` 路径（仅含提示词） | `datasets/X-Fun-Videos-Demo/metadata.json` | +| `--transformer_path` | Stage 2 CCD checkpoint（生成器/判别器初始化） | `$STAGE2_CKPT` | +| `--train_batch_size` | 每 GPU batch 大小 | 1 | +| `--gradient_accumulation_steps` | 梯度累积步数 | 1 | +| `--dataloader_num_workers` | DataLoader 子进程数 | 8 | +| `--num_train_epochs` | 训练 epoch 数 | 100 | +| `--checkpointing_steps` | 每 N 步保存 checkpoint | 200 | +| `--learning_rate` | 生成器学习率 | 2e-06 | +| `--learning_rate_critic` | 判别器学习率 | 2e-06 | +| `--lr_scheduler` | 学习率调度器 | `constant_with_warmup` | +| `--lr_warmup_steps` | 学习率预热步数 | 100 | +| `--output_dir` | 输出目录 | `output_dir_wan2.1_causal_forcing_dmd` | +| `--gradient_checkpointing` | 激活重计算 | - | +| `--mixed_precision` | 混合精度：`fp16/bf16` | `bf16` | +| `--adam_weight_decay` | AdamW 权重衰减 | 0.0 | +| `--adam_beta1` | AdamW beta1（DMD 使用 0.0） | 0.0 | +| `--adam_beta2` | AdamW beta2 | 0.999 | +| `--adam_epsilon` | AdamW epsilon | 1e-10 | +| `--max_grad_norm` | 梯度裁剪阈值 | 10.0 | +| `--trainable_modules` | 可训练模块（`"."` = 全部） | `"."` | +| `--low_vram` | 低显存模式（卸载 VAE/文本编码器） | - | + +**视频采样参数**： + +| 参数 | 说明 | 示例值 | +|------|------|--------| +| `--image_sample_size` | 图像采样尺寸 | 640 | +| `--video_sample_size` | 视频采样尺寸 | 640 | +| `--token_sample_size` | Token 采样尺寸 | 640 | +| `--fix_sample_size` | 固定输出 `[高度, 宽度]` | `480 832` | +| `--video_sample_stride` | 帧采样步幅 | 2 | +| `--video_sample_n_frames` | 视频帧数 | 81 | +| `--random_hw_adapt` | 启用随机分辨率适配 | - | +| `--training_with_video_token_length` | 启用基于 token 长度的训练 | - | +| `--enable_bucket` | 启用宽高比分桶采样 | - | +| `--vae_mini_batch` | VAE 编码迷你批次大小（设为 1 避免 OOM） | 1 | + +### 5.3 DMD 特有参数 + +| 参数 | 说明 | 示例值 | +|------|------|--------| +| `--denoising_step_indices_list` | 去噪步骤索引（DMD 核心参数）。`[1000, 500]` = 2 步 DMD | `1000 500` | +| `--real_guidance_scale` | real-score（14B teacher）的 CFG scale | 3.0 | +| `--fake_guidance_scale` | fake-score（生成器）的 CFG scale。0.0 = 无 CFG | 0.0 | +| `--gen_update_interval` | 生成器更新间隔（每 N 步判别器更新后更新 1 次生成器） | 5 | +| `--num_frame_per_block` | 每个因果块的帧数。`1` = 逐帧（DMD 默认值） | 1 | +| `--use_kv_cache_training` | 使用 KV 缓存逐块训练（匹配原始 Self-Forcing） | - | +| `--independent_first_frame` | 第一帧是否独立（`[1, N, N, ...]` 模式，适用于 I2V） | - | +| `--context_noise` | KV 缓存更新的上下文噪声级别 | 0 | +| `--use_teacher_forcing` | 启用 teacher forcing（将 clean_x 传给 transformer） | - | +| `--teacher_forcing_prob` | 每步应用 teacher forcing 的概率（1.0 = 始终） | 1.0 | +| `--train_mode` | 训练模式：`normal`（TextDataset，仅提示词）或 `i2v` | `normal` | +| `--resume_from_checkpoint` | 恢复训练。使用 `"latest"` 自动选择最新 checkpoint | `"latest"` | +--- + +## 六、使用训练好的 Checkpoint + +Stage 3 DMD checkpoint 是 Causal-Forcing 流水线的**最终模型**，直接用于推理： + +```python +# 在 examples/wan2.1_causal_forcing/predict_t2v.py 中 +transformer_path = "output_dir_wan2.1_causal_forcing_dmd/checkpoint-{N}/diffusion_pytorch_model.safetensors" + +# DMD Stage 3 推理配置 +guidance_scale = 1.0 # CFG 已烘焦到蒸馏权重中 +num_inference_steps = 2 # 2 步 DMD +stochastic_sampling = True +num_frame_per_block = 1 # 逐帧生成 +``` + +或运行： + +```bash +python examples/wan2.1_causal_forcing/predict_t2v.py +``` + +> **阶段选择器**：`predict_t2v.py` 中包含阶段选择器部分，提供不同阶段的预设配置： +> - **Stage 1 (AR Diffusion)**：`guidance_scale=3.0`、`num_inference_steps=50`、`stochastic_sampling=False` +> - **Stage 2 (CCD)**：`guidance_scale=1.0`、`num_inference_steps=4`、`stochastic_sampling=True` +> - **Stage 3 (DMD)**：`guidance_scale=1.0`、`num_inference_steps=2`、`stochastic_sampling=True` + +--- + +## 七、更多资源 + +- **Causal-Forcing 论文**：https://github.com/thu-ml/Causal-Forcing +- **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun diff --git a/scripts/wan2.1_causal_forcing/train_ar_diffusion.py b/scripts/wan2.1_causal_forcing/train_ar_diffusion.py new file mode 100644 index 00000000..04bd4e75 --- /dev/null +++ b/scripts/wan2.1_causal_forcing/train_ar_diffusion.py @@ -0,0 +1,1509 @@ +"""Causal-Forcing Stage 1 (Autoregressive Diffusion) training, modified from +scripts/wan2.1_self_forcing/train_distill.py and Causal-Forcing +(https://github.com/thu-ml/Causal-Forcing). +""" +#!/usr/bin/env python +# coding=utf-8 +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and + +import argparse +import gc +import logging +import math +import os +import pickle +import shutil +import sys + +import accelerate +import diffusers +import numpy as np +import torch +import torch.nn.functional as F +import torch.utils.checkpoint +import transformers +from accelerate import Accelerator +from accelerate.logging import get_logger +from accelerate.state import AcceleratorState +from accelerate.utils import ProjectConfiguration, set_seed +from diffusers import FlowMatchEulerDiscreteScheduler +from diffusers.optimization import get_scheduler +from diffusers.utils import check_min_version, is_wandb_available +from diffusers.utils.torch_utils import is_compiled_module +from einops import rearrange +from omegaconf import OmegaConf +from packaging import version +from torch.utils.tensorboard import SummaryWriter +from torchvision import transforms +from tqdm.auto import tqdm +from transformers import AutoTokenizer +from transformers.utils import ContextManagers + +import datasets + +current_file_path = os.path.abspath(__file__) +project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dirname(current_file_path)), os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))] +for project_root in project_roots: + sys.path.insert(0, project_root) if project_root not in sys.path else None + +from videox_fun.data import (ASPECT_RATIO_512, ASPECT_RATIO_RANDOM_CROP_512, + ASPECT_RATIO_RANDOM_CROP_PROB, + AspectRatioBatchImageVideoSampler, + ImageVideoDataset, RandomSampler, + get_closest_ratio) +from videox_fun.models import (AutoencoderKLWan, WanT5EncoderModel, + WanTransformer3DModel_SelfForcing) +from videox_fun.pipeline import WanSelfForcingPipeline +from videox_fun.utils.utils import save_videos_grid + +if is_wandb_available(): + import wandb + + +def filter_kwargs(cls, kwargs): + import inspect + sig = inspect.signature(cls.__init__) + valid_params = set(sig.parameters.keys()) - {'self', 'cls'} + filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_params} + return filtered_kwargs + + +def get_random_downsample_ratio(sample_size, image_ratio=[], + all_choices=False, rng=None): + def _create_special_list(length): + if length == 1: + return [1.0] + if length >= 2: + first_element = 0.75 + remaining_sum = 1.0 - first_element + other_elements_value = remaining_sum / (length - 1) + special_list = [first_element] + [other_elements_value] * (length - 1) + return special_list + + if sample_size >= 1536: + number_list = [1, 1.25, 1.5, 2, 2.5, 3] + image_ratio + elif sample_size >= 1024: + number_list = [1, 1.25, 1.5, 2] + image_ratio + elif sample_size >= 768: + number_list = [1, 1.25, 1.5] + image_ratio + elif sample_size >= 512: + number_list = [1] + image_ratio + else: + number_list = [1] + + if all_choices: + return number_list + + number_list_prob = np.array(_create_special_list(len(number_list))) + if rng is None: + return np.random.choice(number_list, p = number_list_prob) + else: + return rng.choice(number_list, p = number_list_prob) + +# Will error if the minimal version of diffusers is not installed. Remove at your own risks. +check_min_version("0.18.0.dev0") + +logger = get_logger(__name__, log_level="INFO") + +def log_validation(vae, text_encoder, tokenizer, transformer3d, args, config, accelerator, weight_dtype, global_step): + try: + is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine' + if is_deepspeed: + origin_config = transformer3d.config + transformer3d.config = accelerator.unwrap_model(transformer3d).config + with torch.no_grad(), torch.amp.autocast('cuda', dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + logger.info("Running validation... ") + scheduler_kwargs = OmegaConf.to_container(config['scheduler_kwargs']) + scheduler_kwargs['shift'] = args.shift + scheduler = FlowMatchEulerDiscreteScheduler( + **filter_kwargs(FlowMatchEulerDiscreteScheduler, scheduler_kwargs) + ) + pipeline = WanSelfForcingPipeline( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d, + scheduler=scheduler, + ) + pipeline = pipeline.to(accelerator.device) + + if args.seed is None: + generator = None + else: + rank_seed = args.seed + accelerator.process_index + generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed) + logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}") + + for i in range(len(args.validation_prompts)): + if args.fix_sample_size is not None: + height, width = args.fix_sample_size + else: + height, width = args.video_sample_size, args.video_sample_size + sample = pipeline( + args.validation_prompts[i], + num_frames = args.video_sample_n_frames, + negative_prompt = args.negative_prompt, + height = height, + width = width, + generator = generator, + guidance_scale = args.validation_guidance_scale, + num_inference_steps = args.validation_num_inference_steps, + shift = args.shift, + num_frame_per_block = args.num_frame_per_block, + independent_first_frame = args.independent_first_frame, + context_noise = 0, + stochastic_sampling = False, + ).videos + os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True) + save_videos_grid( + sample, + os.path.join( + args.output_dir, + f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.mp4" + ) + ) + + del pipeline + gc.collect() + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype) + if not args.enable_text_encoder_in_dataloader: + text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype) + if is_deepspeed: + transformer3d.config = origin_config + except Exception as e: + gc.collect() + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + print(f"Eval error on rank {accelerator.process_index} with info {e}") + vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype) + if not args.enable_text_encoder_in_dataloader: + text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype) + +def parse_args(): + parser = argparse.ArgumentParser(description="Causal-Forcing Stage 1: Autoregressive Diffusion training for Wan2.1.") + parser.add_argument( + "--pretrained_model_name_or_path", + type=str, + default=None, + required=True, + help="Path to pretrained model or model identifier from huggingface.co/models.", + ) + parser.add_argument( + "--train_data_dir", + type=str, + default=None, + help=( + "A folder containing the training data. " + ), + ) + parser.add_argument( + "--train_data_meta", + type=str, + default=None, + help=( + "A csv containing the training data. " + ), + ) + parser.add_argument( + "--validation_prompts", + type=str, + default=None, + nargs="+", + help=("A set of prompts evaluated every `--validation_epochs` and logged to `--report_to`."), + ) + parser.add_argument( + "--negative_prompt", + type=str, + default="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走", + help=("The negative prompt used for validation generation."), + ) + parser.add_argument( + "--output_dir", + type=str, + default="sd-model-finetuned", + help="The output directory where the model predictions and checkpoints will be written.", + ) + parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.") + parser.add_argument( + "--use_came", + action="store_true", + help="whether to use came", + ) + parser.add_argument( + "--train_batch_size", type=int, default=16, help="Batch size (per device) for the training dataloader." + ) + parser.add_argument( + "--vae_mini_batch", type=int, default=32, help="mini batch size for vae." + ) + parser.add_argument("--num_train_epochs", type=int, default=100) + parser.add_argument( + "--max_train_steps", + type=int, + default=None, + help="Total number of training steps to perform. If provided, overrides num_train_epochs.", + ) + parser.add_argument( + "--gradient_accumulation_steps", + type=int, + default=1, + help="Number of updates steps to accumulate before performing a backward/update pass.", + ) + parser.add_argument( + "--gradient_checkpointing", + action="store_true", + help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.", + ) + parser.add_argument( + "--learning_rate", + type=float, + default=1e-4, + help="Initial learning rate (after the potential warmup period) to use.", + ) + parser.add_argument( + "--scale_lr", + action="store_true", + default=False, + help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.", + ) + parser.add_argument( + "--lr_scheduler", + type=str, + default="constant", + help=( + 'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",' + ' "constant", "constant_with_warmup"]' + ), + ) + parser.add_argument( + "--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler." + ) + parser.add_argument( + "--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes." + ) + parser.add_argument( + "--allow_tf32", + action="store_true", + help=( + "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see" + " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices" + ), + ) + parser.add_argument( + "--dataloader_num_workers", + type=int, + default=0, + help=( + "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process." + ), + ) + parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.") + parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.") + parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.") + parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer") + parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.") + parser.add_argument( + "--logging_dir", + type=str, + default="logs", + help=( + "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to" + " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***." + ), + ) + parser.add_argument( + "--mixed_precision", + type=str, + default=None, + choices=["no", "fp16", "bf16"], + help=( + "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >=" + " 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the" + " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config." + ), + ) + parser.add_argument( + "--report_to", + type=str, + default="tensorboard", + help=( + 'The integration to report the results and logs to. Supported platforms are `"tensorboard"`' + ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.' + ), + ) + parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank") + parser.add_argument( + "--checkpointing_steps", + type=int, + default=500, + help=( + "Save a checkpoint of the training state every X updates. These checkpoints are only suitable for resuming" + " training using `--resume_from_checkpoint`." + ), + ) + parser.add_argument( + "--checkpoints_total_limit", + type=int, + default=None, + help=("Max number of checkpoints to store."), + ) + parser.add_argument( + "--resume_from_checkpoint", + type=str, + default=None, + help=( + "Whether training should be resumed from a previous checkpoint. Use a path saved by" + ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.' + ), + ) + parser.add_argument( + "--validation_epochs", + type=int, + default=5, + help="Run validation every X epochs.", + ) + parser.add_argument( + "--validation_steps", + type=int, + default=2000, + help="Run validation every X steps.", + ) + parser.add_argument( + "--validation_guidance_scale", + type=float, + default=3.0, + help="CFG scale used when sampling validation videos.", + ) + parser.add_argument( + "--validation_num_inference_steps", + type=int, + default=50, + help=( + "Number of denoising steps used for validation. AR diffusion validation runs a" + " full multi-step rollout, so a larger value (e.g. 50) is recommended." + ), + ) + parser.add_argument( + "--tracker_project_name", + type=str, + default="text2image-fine-tune", + help=( + "The `project_name` argument passed to Accelerator.init_trackers for" + " more information see https://huggingface.co/docs/accelerate/v0.17.0/en/package_reference/accelerator#accelerate.Accelerator" + ), + ) + + parser.add_argument( + "--enable_text_encoder_in_dataloader", action="store_true", help="Whether or not to use text encoder in dataloader." + ) + parser.add_argument( + "--enable_bucket", action="store_true", help="Whether enable bucket sample in datasets." + ) + parser.add_argument( + "--random_ratio_crop", action="store_true", help="Whether enable random ratio crop sample in datasets." + ) + parser.add_argument( + "--random_hw_adapt", action="store_true", help="Whether enable random adapt height and width in datasets." + ) + parser.add_argument( + "--training_with_video_token_length", action="store_true", help="The training stage of the model in training.", + ) + parser.add_argument( + "--train_sampling_steps", + type=int, + default=1000, + help="Total number of scheduler timesteps for sampling.", + ) + parser.add_argument( + "--token_sample_size", + type=int, + default=512, + help="Sample size of the token.", + ) + parser.add_argument( + "--video_sample_size", + type=int, + default=512, + help="Sample size of the video.", + ) + parser.add_argument( + "--image_sample_size", + type=int, + default=512, + help="Sample size of the image.", + ) + parser.add_argument( + "--fix_sample_size", + nargs=2, type=int, default=None, + help="Fix Sample size [height, width] when using bucket and collate_fn." + ) + parser.add_argument( + "--video_sample_stride", + type=int, + default=4, + help="Sample stride of the video.", + ) + parser.add_argument( + "--video_sample_n_frames", + type=int, + default=17, + help="Num frame of video.", + ) + parser.add_argument( + "--video_repeat", + type=int, + default=0, + help="Num of repeat video.", + ) + parser.add_argument( + "--config_path", + type=str, + default=None, + help=( + "The config of the model in training." + ), + ) + parser.add_argument( + "--transformer_path", + type=str, + default=None, + help=("If you want to load the weight from other transformers, input its path."), + ) + parser.add_argument( + "--vae_path", + type=str, + default=None, + help=("If you want to load the weight from other vaes, input its path."), + ) + + parser.add_argument( + '--trainable_modules', + nargs='+', + help='Enter a list of trainable modules' + ) + parser.add_argument( + '--trainable_modules_low_learning_rate', + nargs='+', + default=[], + help='Enter a list of trainable modules with lower learning rate' + ) + parser.add_argument( + '--tokenizer_max_length', + type=int, + default=512, + help='Max length of tokenizer' + ) + parser.add_argument( + "--use_deepspeed", action="store_true", help="Whether or not to use deepspeed." + ) + parser.add_argument( + "--use_fsdp", action="store_true", help="Whether or not to use fsdp." + ) + parser.add_argument( + "--low_vram", action="store_true", help="Whether enable low_vram mode." + ) + parser.add_argument( + "--num_frame_per_block", + type=int, + default=3, + help="Number of latent frames per causal block. 3 = chunk-wise, 1 = frame-wise." + ) + parser.add_argument( + "--independent_first_frame", + action="store_true", + help="Whether first frame is independent ([1, N, N, ...] pattern, useful for I2V)." + ) + parser.add_argument( + "--shift", + type=float, + default=5.0, + help="Shift value for FlowMatchEulerDiscreteScheduler. Causal-Forcing uses 5.0 by default." + ) + parser.add_argument( + "--no_teacher_forcing", + action="store_true", + help="Disable teacher forcing (train under diffusion forcing instead). Stage 1 defaults to TF." + ) + parser.add_argument( + "--noise_augmentation_max_timestep", + type=int, + default=0, + help=( + "If > 0, add light flow-matching noise (sampled in [0, value)) " + "to the clean context tokens during teacher forcing." + ), + ) + parser.add_argument( + "--use_timestep_weight", + action="store_true", + help="Apply the Causal-Forcing per-timestep loss weight (Gaussian centered at T/2)." + ) + args = parser.parse_args() + env_local_rank = int(os.environ.get("LOCAL_RANK", -1)) + if env_local_rank != -1 and env_local_rank != args.local_rank: + args.local_rank = env_local_rank + + return args + + +def main(): + args = parse_args() + + logging_dir = os.path.join(args.output_dir, args.logging_dir) + + config = OmegaConf.load(args.config_path) + accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir) + + accelerator = Accelerator( + gradient_accumulation_steps=args.gradient_accumulation_steps, + mixed_precision=args.mixed_precision, + log_with=args.report_to, + project_config=accelerator_project_config, + ) + + deepspeed_plugin = accelerator.state.deepspeed_plugin if hasattr(accelerator.state, "deepspeed_plugin") else None + fsdp_plugin = accelerator.state.fsdp_plugin if hasattr(accelerator.state, "fsdp_plugin") else None + if deepspeed_plugin is not None: + zero_stage = int(deepspeed_plugin.zero_stage) + fsdp_stage = 0 + print(f"Using DeepSpeed Zero stage: {zero_stage}") + + args.use_deepspeed = True + if zero_stage == 3: + print(f"Auto set save_state to True because zero_stage == 3") + args.save_state = True + elif fsdp_plugin is not None: + from torch.distributed.fsdp import ShardingStrategy + zero_stage = 0 + if fsdp_plugin.sharding_strategy is ShardingStrategy.FULL_SHARD: + fsdp_stage = 3 + elif fsdp_plugin.sharding_strategy is None: # The fsdp_plugin.sharding_strategy is None in FSDP 2. + fsdp_stage = 3 + elif fsdp_plugin.sharding_strategy is ShardingStrategy.SHARD_GRAD_OP: + fsdp_stage = 2 + else: + fsdp_stage = 0 + print(f"Using FSDP stage: {fsdp_stage}") + + args.use_fsdp = True + if fsdp_stage == 3: + print(f"Auto set save_state to True because fsdp_stage == 3") + args.save_state = True + else: + zero_stage = 0 + fsdp_stage = 0 + print("DeepSpeed is not enabled.") + + if accelerator.is_main_process: + writer = SummaryWriter(log_dir=logging_dir) + + # Make one log on every process with the configuration for debugging. + logging.basicConfig( + format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", + datefmt="%m/%d/%Y %H:%M:%S", + level=logging.INFO, + ) + logger.info(accelerator.state, main_process_only=False) + if accelerator.is_local_main_process: + datasets.utils.logging.set_verbosity_warning() + transformers.utils.logging.set_verbosity_warning() + diffusers.utils.logging.set_verbosity_info() + else: + datasets.utils.logging.set_verbosity_error() + transformers.utils.logging.set_verbosity_error() + diffusers.utils.logging.set_verbosity_error() + + # If passed along, set the training seed now. + if args.seed is not None: + set_seed(args.seed) + rng = np.random.default_rng(np.random.PCG64(args.seed + accelerator.process_index)) + torch_rng = torch.Generator(accelerator.device).manual_seed(args.seed + accelerator.process_index) + print(f"Init rng with seed {args.seed + accelerator.process_index}. Process_index is {accelerator.process_index}") + else: + rng = None + torch_rng = None + print(f"No seed provided; using global default RNG. Process_index is {accelerator.process_index}") + + # Handle the repository creation + if accelerator.is_main_process: + if args.output_dir is not None: + os.makedirs(args.output_dir, exist_ok=True) + + # For mixed precision training we cast all non-trainable weigths (vae, non-lora text_encoder and non-lora transformer3d) to half-precision + # as these weights are only used for inference, keeping weights in full precision is not required. + weight_dtype = torch.float32 + if accelerator.mixed_precision == "fp16": + weight_dtype = torch.float16 + args.mixed_precision = accelerator.mixed_precision + elif accelerator.mixed_precision == "bf16": + weight_dtype = torch.bfloat16 + args.mixed_precision = accelerator.mixed_precision + + # Load scheduler, tokenizer and models. + scheduler_kwargs = OmegaConf.to_container(config['scheduler_kwargs']) + scheduler_kwargs['shift'] = args.shift + noise_scheduler = FlowMatchEulerDiscreteScheduler( + **filter_kwargs(FlowMatchEulerDiscreteScheduler, scheduler_kwargs) + ) + + # Get Tokenizer + tokenizer = AutoTokenizer.from_pretrained( + os.path.join(args.pretrained_model_name_or_path, config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer')), + ) + + def deepspeed_zero_init_disabled_context_manager(): + """ + returns either a context list that includes one that will disable zero.Init or an empty context list + """ + deepspeed_plugin = AcceleratorState().deepspeed_plugin if accelerate.state.is_initialized() else None + if deepspeed_plugin is None: + return [] + + return [deepspeed_plugin.zero3_init_context_manager(enable=False)] + + # Currently Accelerate doesn't know how to handle multiple models under Deepspeed ZeRO stage 3. + # For this to work properly all models must be run through `accelerate.prepare`. But accelerate + # will try to assign the same optimizer with the same weights to all models during + # `deepspeed.initialize`, which of course doesn't work. + # + # For now the following workaround will partially support Deepspeed ZeRO-3, by excluding the 2 + # frozen models from being partitioned during `zero.Init` which gets called during + # `from_pretrained` So CLIPTextModel and AutoencoderKL will not enjoy the parameter sharding + # across multiple gpus and only UNet2DConditionModel will get ZeRO sharded. + with ContextManagers(deepspeed_zero_init_disabled_context_manager()): + # Get Text encoder + text_encoder = WanT5EncoderModel.from_pretrained( + os.path.join(args.pretrained_model_name_or_path, config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')), + additional_kwargs=OmegaConf.to_container(config['text_encoder_kwargs']), + low_cpu_mem_usage=True, + torch_dtype=weight_dtype, + ) + text_encoder = text_encoder.eval() + # Get Vae + vae = AutoencoderKLWan.from_pretrained( + os.path.join(args.pretrained_model_name_or_path, config['vae_kwargs'].get('vae_subpath', 'vae')), + additional_kwargs=OmegaConf.to_container(config['vae_kwargs']), + ) + vae.eval() + + # Get Transformer (causal generator). + # IMPORTANT: keep the trainable transformer in fp32. accelerate's + # mixed_precision="bf16" will autocast the forward to bf16 while keeping + # the master weights and Adam moments in fp32. If params live in bf16, + # every update (LR*grad ~ 1e-5 for LR=2e-6) falls below bf16 mantissa + # precision (~1e-3 relative) and is rounded to zero, so the model barely + # moves over thousands of steps and never learns the causal/KV-cache + # structure. CF stage 1 also keeps params in fp32 (its 10k ckpt on disk + # is fp32; with the bug above ours used to be bf16). + # + # NOTE: WanTransformer3DModel.from_pretrained defaults to torch_dtype=bf16 + # and ends with `model = model.to(torch_dtype)`, so an explicit + # torch_dtype=torch.float32 is required to keep the loaded weights fp32. + transformer3d = WanTransformer3DModel_SelfForcing.from_pretrained( + os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')), + transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']), + low_cpu_mem_usage=True, + torch_dtype=torch.float32, + ) + # Stage 1 only trains the causal generator; mark the block layout the + # downstream Self-Forcing pipeline expects at sampling time. + transformer3d.num_frame_per_block = args.num_frame_per_block + transformer3d.independent_first_frame = args.independent_first_frame + + # Freeze vae and text_encoder; transformer3d is toggled per-module below. + vae.requires_grad_(False) + text_encoder.requires_grad_(False) + transformer3d.requires_grad_(False) + + if args.transformer_path is not None: + print(f"From checkpoint: {args.transformer_path}") + if args.transformer_path.endswith("safetensors"): + from safetensors.torch import load_file, safe_open + state_dict = load_file(args.transformer_path) + else: + state_dict = torch.load(args.transformer_path, map_location="cpu") + state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict + state_dict = state_dict["generator_ema"] if "generator_ema" in state_dict else state_dict + state_dict = state_dict["generator"] if "generator" in state_dict else state_dict + if any(k.startswith("model.") for k in state_dict.keys()): + state_dict = {k.replace("model.", "", 1) if k.startswith("model.") else k: v for k, v in state_dict.items()} + + m, u = transformer3d.load_state_dict(state_dict, strict=False) + print(f"missing keys: {len(m)}, unexpected keys: {len(u)}") + assert len(u) == 0 + + if args.vae_path is not None: + print(f"From checkpoint: {args.vae_path}") + if args.vae_path.endswith("safetensors"): + from safetensors.torch import load_file, safe_open + state_dict = load_file(args.vae_path) + else: + state_dict = torch.load(args.vae_path, map_location="cpu") + state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict + + m, u = vae.load_state_dict(state_dict, strict=False) + print(f"missing keys: {len(m)}, unexpected keys: {len(u)}") + assert len(u) == 0 + + # A good trainable modules is showed below now. + # For 3D Patch: trainable_modules = ['ff.net', 'pos_embed', 'attn2', 'proj_out', 'timepositionalencoding', 'h_position', 'w_position'] + # For 2D Patch: trainable_modules = ['ff.net', 'attn2', 'timepositionalencoding', 'h_position', 'w_position'] + transformer3d.train() + if accelerator.is_main_process: + accelerator.print( + f"Trainable modules '{args.trainable_modules}'." + ) + for name, param in transformer3d.named_parameters(): + for trainable_module_name in args.trainable_modules + args.trainable_modules_low_learning_rate: + if trainable_module_name in name: + param.requires_grad = True + break + + # `accelerate` 0.16.0 will have better support for customized saving + if version.parse(accelerate.__version__) >= version.parse("0.16.0"): + # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format + if fsdp_stage != 0 or zero_stage == 3: + def save_model_hook(models, weights, output_dir): + accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True) + if accelerator.is_main_process: + from safetensors.torch import save_file + + safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors") + accelerate_state_dict = {k: v.to(dtype=weight_dtype) for k, v in accelerate_state_dict.items()} + save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"}) + + with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file: + pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file) + + def load_model_hook(models, input_dir): + pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl") + if os.path.exists(pkl_path): + with open(pkl_path, 'rb') as file: + loaded_number, _ = pickle.load(file) + batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0) + print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.") + else: + # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format + def save_model_hook(models, weights, output_dir): + if accelerator.is_main_process: + models[0].save_pretrained(os.path.join(output_dir, "transformer")) + if not args.use_deepspeed: + weights.pop() + + with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file: + pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file) + + def load_model_hook(models, input_dir): + for i in range(len(models)): + # pop models so that they are not loaded again + model = models.pop() + + # load diffusers style into model + load_model = WanTransformer3DModel.from_pretrained( + input_dir, subfolder="transformer" + ) + model.register_to_config(**load_model.config) + + model.load_state_dict(load_model.state_dict()) + del load_model + + pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl") + if os.path.exists(pkl_path): + with open(pkl_path, 'rb') as file: + loaded_number, _ = pickle.load(file) + batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0) + print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.") + + accelerator.register_save_state_pre_hook(save_model_hook) + accelerator.register_load_state_pre_hook(load_model_hook) + + if args.gradient_checkpointing: + transformer3d.enable_gradient_checkpointing() + + # Enable TF32 for faster training on Ampere GPUs, + # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices + if args.allow_tf32: + torch.backends.cuda.matmul.allow_tf32 = True + + if args.scale_lr: + args.learning_rate = ( + args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes + ) + + # Initialize the optimizer + if args.use_8bit_adam: + try: + import bitsandbytes as bnb + except ImportError: + raise ImportError( + "Please install bitsandbytes to use 8-bit Adam. You can do so by running `pip install bitsandbytes`" + ) + + optimizer_cls = bnb.optim.AdamW8bit + elif args.use_came: + try: + from came_pytorch import CAME + except Exception: + raise ImportError( + "Please install came_pytorch to use CAME. You can do so by running `pip install came_pytorch`" + ) + + optimizer_cls = CAME + else: + optimizer_cls = torch.optim.AdamW + + trainable_params = list(filter(lambda p: p.requires_grad, transformer3d.parameters())) + trainable_params_optim = [ + {'params': [], 'lr': args.learning_rate}, + {'params': [], 'lr': args.learning_rate / 2}, + ] + in_already = [] + for name, param in transformer3d.named_parameters(): + high_lr_flag = False + if name in in_already: + continue + for trainable_module_name in args.trainable_modules: + if trainable_module_name in name: + in_already.append(name) + high_lr_flag = True + trainable_params_optim[0]['params'].append(param) + if accelerator.is_main_process: + print(f"Set {name} to lr : {args.learning_rate}") + break + if high_lr_flag: + continue + for trainable_module_name in args.trainable_modules_low_learning_rate: + if trainable_module_name in name: + in_already.append(name) + trainable_params_optim[1]['params'].append(param) + if accelerator.is_main_process: + print(f"Set {name} to lr : {args.learning_rate / 2}") + break + + if args.use_came: + optimizer = optimizer_cls( + trainable_params_optim, + lr=args.learning_rate, + # weight_decay=args.adam_weight_decay, + betas=(0.9, 0.999, 0.9999), + eps=(1e-30, 1e-16) + ) + else: + optimizer = optimizer_cls( + trainable_params_optim, + lr=args.learning_rate, + betas=(args.adam_beta1, args.adam_beta2), + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + ) + + # Get the training dataset + sample_n_frames_bucket_interval = vae.config.temporal_compression_ratio + + if args.fix_sample_size is not None and args.enable_bucket: + args.video_sample_size = max(max(args.fix_sample_size), args.video_sample_size) + args.image_sample_size = max(max(args.fix_sample_size), args.image_sample_size) + args.training_with_video_token_length = False + args.random_hw_adapt = False + + # Get the dataset (Stage 1 always trains on raw videos with teacher forcing). + train_dataset = ImageVideoDataset( + args.train_data_meta, args.train_data_dir, + video_sample_size=args.video_sample_size, video_sample_stride=args.video_sample_stride, video_sample_n_frames=args.video_sample_n_frames, + video_repeat=args.video_repeat, + image_sample_size=args.image_sample_size, + enable_bucket=args.enable_bucket, enable_inpaint=False, + ) + + # Causal-Forcing needs at least one full block of latent frames per clip. + # Enforce the matching pixel-frame minimum at video-read time so too-short + # clips raise and get resampled instead of collapsing the block reshape. + if args.independent_first_frame: + # latent_frames - 1 must be >= num_frame_per_block + train_dataset.min_video_sample_n_frames = args.num_frame_per_block * sample_n_frames_bucket_interval + 1 + else: + # latent_frames must be >= num_frame_per_block + train_dataset.min_video_sample_n_frames = (args.num_frame_per_block - 1) * sample_n_frames_bucket_interval + 1 + + def get_length_to_frame_num(token_length): + if args.image_sample_size > args.video_sample_size: + sample_sizes = list(range(args.video_sample_size, args.image_sample_size + 1, 128)) + + if sample_sizes[-1] != args.image_sample_size: + sample_sizes.append(args.image_sample_size) + else: + sample_sizes = [args.image_sample_size] + + length_to_frame_num = { + sample_size: min(token_length / sample_size / sample_size, args.video_sample_n_frames) // sample_n_frames_bucket_interval * sample_n_frames_bucket_interval + 1 for sample_size in sample_sizes + } + + return length_to_frame_num + + aspect_ratio_sample_size = {key : [x / 512 * args.video_sample_size for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()} + batch_sampler_generator = torch.Generator().manual_seed(args.seed) + batch_sampler = AspectRatioBatchImageVideoSampler( + sampler=RandomSampler(train_dataset, generator=batch_sampler_generator), dataset=train_dataset.dataset, + batch_size=args.train_batch_size, train_folder = args.train_data_dir, drop_last=True, + aspect_ratios=aspect_ratio_sample_size, + ) + + def collate_fn(examples): + # Get token length + target_token_length = args.video_sample_n_frames * args.token_sample_size * args.token_sample_size + length_to_frame_num = get_length_to_frame_num(target_token_length) + + # Create new output + new_examples = {} + new_examples["target_token_length"] = target_token_length + new_examples["pixel_values"] = [] + new_examples["text"] = [] + + # Get downsample ratio in image and videos + pixel_value = examples[0]["pixel_values"] + data_type = examples[0]["data_type"] + f, h, w, c = np.shape(pixel_value) + if data_type == 'image': + random_downsample_ratio = 1 if not args.random_hw_adapt else get_random_downsample_ratio(args.image_sample_size, image_ratio=[args.image_sample_size / args.video_sample_size], rng=rng) + + aspect_ratio_sample_size = {key : [x / 512 * args.image_sample_size / random_downsample_ratio for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()} + aspect_ratio_random_crop_sample_size = {key : [x / 512 * args.image_sample_size / random_downsample_ratio for x in ASPECT_RATIO_RANDOM_CROP_512[key]] for key in ASPECT_RATIO_RANDOM_CROP_512.keys()} + + batch_video_length = args.video_sample_n_frames + sample_n_frames_bucket_interval + else: + if args.random_hw_adapt: + if args.training_with_video_token_length: + local_min_size = np.min(np.array([np.mean(np.array([np.shape(example["pixel_values"])[1], np.shape(example["pixel_values"])[2]])) for example in examples])) + # The video will be resized to a lower resolution than its own. + choice_list = [length for length in list(length_to_frame_num.keys()) if length < local_min_size * 1.25] + if len(choice_list) == 0: + choice_list = list(length_to_frame_num.keys()) + if rng is None: + local_video_sample_size = np.random.choice(choice_list) + else: + local_video_sample_size = rng.choice(choice_list) + batch_video_length = length_to_frame_num[local_video_sample_size] + random_downsample_ratio = args.video_sample_size / local_video_sample_size + else: + random_downsample_ratio = get_random_downsample_ratio( + args.video_sample_size, rng=rng) + batch_video_length = args.video_sample_n_frames + sample_n_frames_bucket_interval + else: + random_downsample_ratio = 1 + batch_video_length = args.video_sample_n_frames + sample_n_frames_bucket_interval + + aspect_ratio_sample_size = {key : [x / 512 * args.video_sample_size / random_downsample_ratio for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()} + aspect_ratio_random_crop_sample_size = {key : [x / 512 * args.video_sample_size / random_downsample_ratio for x in ASPECT_RATIO_RANDOM_CROP_512[key]] for key in ASPECT_RATIO_RANDOM_CROP_512.keys()} + + if args.fix_sample_size is not None: + fix_sample_size = [int(x / 16) * 16 for x in args.fix_sample_size] + elif args.random_ratio_crop: + if rng is None: + random_sample_size = aspect_ratio_random_crop_sample_size[ + np.random.choice(list(aspect_ratio_random_crop_sample_size.keys()), p = ASPECT_RATIO_RANDOM_CROP_PROB) + ] + else: + random_sample_size = aspect_ratio_random_crop_sample_size[ + rng.choice(list(aspect_ratio_random_crop_sample_size.keys()), p = ASPECT_RATIO_RANDOM_CROP_PROB) + ] + random_sample_size = [int(x / 16) * 16 for x in random_sample_size] + else: + closest_size, closest_ratio = get_closest_ratio(h, w, ratios=aspect_ratio_sample_size) + closest_size = [int(x / 16) * 16 for x in closest_size] + + min_example_length = min( + [example["pixel_values"].shape[0] for example in examples] + ) + batch_video_length = int(min(batch_video_length, min_example_length)) + + # Magvae needs the number of frames to be 4n + 1. + batch_video_length = (batch_video_length - 1) // sample_n_frames_bucket_interval * sample_n_frames_bucket_interval + 1 + + # Causal-Forcing needs the latent frame count to align with num_frame_per_block. + # Always keep at least one full causal block so the per-block timestep + # reshape (num_frames -> [-1, num_frame_per_block]) stays valid. + k = (batch_video_length - 1) // sample_n_frames_bucket_interval + if args.independent_first_frame: + # latent_frames - 1 = k must be a positive multiple of num_frame_per_block + k = max(k // args.num_frame_per_block, 1) * args.num_frame_per_block + else: + # latent_frames = k + 1 must be a positive multiple of num_frame_per_block + k = max((k + 1) // args.num_frame_per_block, 1) * args.num_frame_per_block - 1 + batch_video_length = k * sample_n_frames_bucket_interval + 1 + + for example in examples: + if args.fix_sample_size is not None: + # To 0~1 + pixel_values = torch.from_numpy(example["pixel_values"]).permute(0, 3, 1, 2).contiguous() + pixel_values = pixel_values / 255. + + # Get adapt hw for resize + fix_sample_size = list(map(lambda x: int(x), fix_sample_size)) + transform = transforms.Compose([ + transforms.Resize(fix_sample_size, interpolation=transforms.InterpolationMode.BILINEAR), # Image.BICUBIC + transforms.CenterCrop(fix_sample_size), + transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True), + ]) + elif args.random_ratio_crop: + # To 0~1 + pixel_values = torch.from_numpy(example["pixel_values"]).permute(0, 3, 1, 2).contiguous() + pixel_values = pixel_values / 255. + + # Get adapt hw for resize + b, c, h, w = pixel_values.size() + th, tw = random_sample_size + if th / tw > h / w: + nh = int(th) + nw = int(w / h * nh) + else: + nw = int(tw) + nh = int(h / w * nw) + + transform = transforms.Compose([ + transforms.Resize([nh, nw]), + transforms.CenterCrop([int(x) for x in random_sample_size]), + transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True), + ]) + else: + # To 0~1 + pixel_values = torch.from_numpy(example["pixel_values"]).permute(0, 3, 1, 2).contiguous() + pixel_values = pixel_values / 255. + + # Get adapt hw for resize + closest_size = list(map(lambda x: int(x), closest_size)) + if closest_size[0] / h > closest_size[1] / w: + resize_size = closest_size[0], int(w * closest_size[0] / h) + else: + resize_size = int(h * closest_size[1] / w), closest_size[1] + + transform = transforms.Compose([ + transforms.Resize(resize_size, interpolation=transforms.InterpolationMode.BILINEAR), # Image.BICUBIC + transforms.CenterCrop(closest_size), + transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True), + ]) + + new_examples["pixel_values"].append(transform(pixel_values)[:batch_video_length]) + new_examples["text"].append(example["text"]) + + # Limit the number of frames to the same + new_examples["pixel_values"] = torch.stack([example for example in new_examples["pixel_values"]]) + + # Encode prompts when enable_text_encoder_in_dataloader=True + if args.enable_text_encoder_in_dataloader: + prompt_ids = tokenizer( + new_examples['text'], + max_length=args.tokenizer_max_length, + padding="max_length", + add_special_tokens=True, + truncation=True, + return_tensors="pt" + ) + text_input_ids = prompt_ids.input_ids + prompt_attention_mask = prompt_ids.attention_mask + + seq_lens = prompt_attention_mask.gt(0).sum(dim=1).long() + prompt_embeds = text_encoder(text_input_ids.to("cpu"), attention_mask=prompt_attention_mask.to("cpu"))[0] + prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] + + new_examples['encoder_attention_mask'] = prompt_ids.attention_mask + new_examples['encoder_hidden_states'] = prompt_embeds + + return new_examples + + # DataLoaders creation: + train_dataloader = torch.utils.data.DataLoader( + train_dataset, + batch_sampler=batch_sampler, + collate_fn=collate_fn, + persistent_workers=True if args.dataloader_num_workers != 0 else False, + num_workers=args.dataloader_num_workers, + ) + + # Scheduler and math around the number of training steps. + overrode_max_train_steps = False + num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) + if args.max_train_steps is None: + args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch + overrode_max_train_steps = True + + lr_scheduler = get_scheduler( + args.lr_scheduler, + optimizer=optimizer, + num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes, + num_training_steps=args.max_train_steps * accelerator.num_processes, + ) + + # Prepare everything with our `accelerator`. + transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( + transformer3d, optimizer, train_dataloader, lr_scheduler + ) + if fsdp_stage != 0 or zero_stage != 0: + from functools import partial + + from videox_fun.dist import set_multi_gpus_devices, shard_model + shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype) + text_encoder = shard_fn(text_encoder) + + # Move text_encode and vae to gpu and cast to weight_dtype + vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype) + if not args.enable_text_encoder_in_dataloader: + text_encoder.to(accelerator.device if not args.low_vram else "cpu") + + # We need to recalculate our total training steps as the size of the training dataloader may have changed. + num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) + if overrode_max_train_steps: + args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch + # Afterwards we recalculate our number of training epochs + args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch) + + # We need to initialize the trackers we use, and also store our configuration. + # The trackers initializes automatically on the main process. + if accelerator.is_main_process: + tracker_config = dict(vars(args)) + keys_to_pop = [k for k, v in tracker_config.items() if isinstance(v, list)] + for k in keys_to_pop: + tracker_config.pop(k) + print(f"Removed tracker_config['{k}']") + accelerator.init_trackers(args.tracker_project_name, tracker_config) + + # Function for unwrapping if model was compiled with `torch.compile`. + def unwrap_model(model): + model = accelerator.unwrap_model(model) + model = model._orig_mod if is_compiled_module(model) else model + return model + + # Train! + total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps + + logger.info("***** Running training *****") + logger.info(f" Num examples = {len(train_dataset)}") + logger.info(f" Num Epochs = {args.num_train_epochs}") + logger.info(f" Instantaneous batch size per device = {args.train_batch_size}") + logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}") + logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}") + logger.info(f" Total optimization steps = {args.max_train_steps}") + global_step = 0 + first_epoch = 0 + + # Potentially load in the weights and states from a previous save + if args.resume_from_checkpoint: + if args.resume_from_checkpoint != "latest": + path = os.path.basename(args.resume_from_checkpoint) + else: + # Get the most recent checkpoint + dirs = os.listdir(args.output_dir) + dirs = [d for d in dirs if d.startswith("checkpoint")] + dirs = sorted(dirs, key=lambda x: int(x.split("-")[1])) + path = dirs[-1] if len(dirs) > 0 else None + + if path is None: + accelerator.print( + f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run." + ) + args.resume_from_checkpoint = None + initial_global_step = 0 + else: + global_step = int(path.split("-")[1]) + + initial_global_step = global_step + + pkl_path = os.path.join(os.path.join(args.output_dir, path), "sampler_pos_start.pkl") + if os.path.exists(pkl_path): + with open(pkl_path, 'rb') as file: + _, first_epoch = pickle.load(file) + else: + first_epoch = global_step // num_update_steps_per_epoch + print(f"Load pkl from {pkl_path}. Get first_epoch = {first_epoch}.") + + accelerator.print(f"Resuming from checkpoint {path}") + accelerator.load_state(os.path.join(args.output_dir, path)) + else: + initial_global_step = 0 + + progress_bar = tqdm( + range(0, args.max_train_steps), + initial=initial_global_step, + desc="Steps", + # Only show the progress bar once on each machine. + disable=not accelerator.is_local_main_process, + ) + + # Materialise the full per-step schedule once so that `add_noise` and the + # training-weight table can be looked up by timestep. + noise_scheduler.set_timesteps(args.train_sampling_steps, device=accelerator.device) + schedule_timesteps = noise_scheduler.timesteps.clone() + # Causal-Forcing per-timestep loss weight (Gaussian centered at T/2). Mirrors + # `FlowMatchScheduler.set_timesteps(training=True)` in Causal-Forcing. + _x = schedule_timesteps.float() + _y = torch.exp(-2 * ((_x - args.train_sampling_steps / 2) / args.train_sampling_steps) ** 2) + _y_shifted = _y - _y.min() + training_weights_table = _y_shifted * (args.train_sampling_steps / _y_shifted.sum().clamp_min(1e-12)) + training_weights_table = training_weights_table.to(accelerator.device) + + for epoch in range(first_epoch, args.num_train_epochs): + train_loss = 0.0 + batch_sampler.sampler.generator = torch.Generator().manual_seed(args.seed + epoch) + for step, batch in enumerate(train_dataloader): + # Data batch sanity check (only available when training on raw videos) + if epoch == first_epoch and step == 0: + pixel_values, texts = batch['pixel_values'].cpu(), batch['text'] + pixel_values = rearrange(pixel_values, "b f c h w -> b c f h w") + os.makedirs(os.path.join(args.output_dir, "sanity_check"), exist_ok=True) + for idx, (pixel_value, text) in enumerate(zip(pixel_values, texts)): + pixel_value = pixel_value[None, ...] + gif_name = '-'.join(text.replace('/', '').split()[:10]) if not text == '' else f'{global_step}-{idx}' + save_videos_grid(pixel_value, f"{args.output_dir}/sanity_check/{gif_name[:10]}.mp4", rescale=True) + + with torch.amp.autocast('cuda', dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + # Convert prompts to text embeddings. + if args.enable_text_encoder_in_dataloader: + prompt_embeds = batch['encoder_hidden_states'].to(device=accelerator.device) + else: + with torch.no_grad(): + prompt_ids = tokenizer( + batch['text'], + padding="max_length", + max_length=args.tokenizer_max_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt" + ) + text_input_ids = prompt_ids.input_ids + prompt_attention_mask = prompt_ids.attention_mask + + seq_lens = prompt_attention_mask.gt(0).sum(dim=1).long() + # Under --low_vram the encoder is parked on CPU; move it to + # the device just for this forward pass and park it back. + if args.low_vram: + text_encoder.to(accelerator.device) + prompt_embeds = text_encoder(text_input_ids.to(accelerator.device), attention_mask=prompt_attention_mask.to(accelerator.device))[0] + prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] + if args.low_vram: + text_encoder.to('cpu') + torch.cuda.empty_cache() + + # Convert pixel videos to clean latents via the VAE. + pixel_values = batch["pixel_values"].to(weight_dtype) + if args.low_vram: + torch.cuda.empty_cache() + vae.to(accelerator.device) + + with torch.no_grad(): + def _batch_encode_vae(pixel_values): + pixel_values = rearrange(pixel_values, "b f c h w -> b c f h w") + bs = args.vae_mini_batch + new_pixel_values = [] + for i in range(0, pixel_values.shape[0], bs): + pixel_values_bs = pixel_values[i : i + bs] + pixel_values_bs = vae.encode(pixel_values_bs)[0] + pixel_values_bs = pixel_values_bs.sample() + new_pixel_values.append(pixel_values_bs) + return torch.cat(new_pixel_values, dim=0) + clean_latents = _batch_encode_vae(pixel_values) + + if args.low_vram: + vae.to('cpu') + torch.cuda.empty_cache() + + with accelerator.accumulate(transformer3d): + def get_sigmas(timesteps, n_dim=4, dtype=torch.float32): + sigmas = noise_scheduler.sigmas.to(device=accelerator.device, dtype=dtype) + schedule_timesteps = noise_scheduler.timesteps.to(accelerator.device) + timesteps = timesteps.to(accelerator.device) + step_indices = torch.argmin((schedule_timesteps.unsqueeze(0) - timesteps.reshape(-1).unsqueeze(1)).abs(), dim=1) + sigma = sigmas[step_indices] + while len(sigma.shape) < n_dim: + sigma = sigma.unsqueeze(-1) + return sigma + + def add_noise(latents, noise, timesteps): + """Per-frame flow-matching add_noise; supports timesteps of shape [B, F].""" + sigmas = get_sigmas(timesteps, n_dim=2, dtype=latents.dtype) + sigmas = sigmas.reshape(latents.shape[0], latents.shape[2], 1, 1, 1).permute(0, 2, 1, 3, 4) + return (1.0 - sigmas) * latents + sigmas * noise + + def get_per_block_timestep(min_step, max_step, batch_size, num_frames): + """One timestep per causal block; frames inside the same block share it.""" + timestep = torch.randint( + min_step, max_step, [batch_size, num_frames], + device=accelerator.device, dtype=torch.long, generator=torch_rng, + ) + if args.independent_first_frame: + ts_rest = timestep[:, 1:] + ts_rest = ts_rest.reshape(ts_rest.shape[0], -1, args.num_frame_per_block) + ts_rest[:, :, 1:] = ts_rest[:, :, 0:1] + ts_rest = ts_rest.reshape(ts_rest.shape[0], -1) + timestep = torch.cat([timestep[:, 0:1], ts_rest], dim=1) + else: + timestep = timestep.reshape(timestep.shape[0], -1, args.num_frame_per_block) + timestep[:, :, 1:] = timestep[:, :, 0:1] + timestep = timestep.reshape(timestep.shape[0], -1) + return timestep + + bsz, channel, num_frames, height, width = clean_latents.shape + + # Sample per-block timesteps and build the noisy input. + # Mirrors `BaseModel._get_timestep` in Causal-Forcing. + step_index = get_per_block_timestep( + int(0.02 * args.train_sampling_steps), + int(0.98 * args.train_sampling_steps), + bsz, num_frames, + ) + timestep = schedule_timesteps.to(accelerator.device)[step_index].to(dtype=torch.float32) + + noise = torch.randn(clean_latents.shape, dtype=weight_dtype, device=accelerator.device, generator=torch_rng) + noisy_latents = add_noise(clean_latents, noise, timestep) + training_target = noise - clean_latents # flow-matching target + + # Optional noise augmentation on the clean context tokens. + if (not args.no_teacher_forcing) and args.noise_augmentation_max_timestep > 0: + aug_index = get_per_block_timestep( + 0, args.noise_augmentation_max_timestep, bsz, num_frames, + ) + aug_t = schedule_timesteps.to(accelerator.device)[aug_index].to(dtype=torch.float32) + clean_latents_aug = add_noise(clean_latents, noise, aug_t) + else: + clean_latents_aug = clean_latents + aug_t = None + + with torch.amp.autocast('cuda', dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + # The transformer expects List[Tensor] of shape [C, F, H, W] per item. + noisy_input_list = [noisy_latents[i] for i in range(bsz)] + clean_x_list = [clean_latents_aug[i] for i in range(bsz)] if not args.no_teacher_forcing else None + patch_h, patch_w = accelerator.unwrap_model(transformer3d).config.patch_size[1:] + full_seq_len = num_frames * height * width // (patch_h * patch_w) + + flow_pred = transformer3d( + x=noisy_input_list, + context=prompt_embeds, + t=timestep.to(torch.int64), + seq_len=full_seq_len, + clean_x=clean_x_list, + aug_t=(aug_t.to(torch.int64) if aug_t is not None else None), + ) + if isinstance(flow_pred, list): + flow_pred = torch.stack(flow_pred, dim=0) + + # Per-frame flow-matching MSE; optionally weighted by the Causal-Forcing + # Gaussian timestep weight. + loss_per_frame = F.mse_loss( + flow_pred.float(), training_target.float(), reduction='none' + ).mean(dim=(1, 3, 4)) # [B, F] + if args.use_timestep_weight: + flat_t = timestep.reshape(-1).to(training_weights_table.dtype) + schedule_for_lookup = schedule_timesteps.to(accelerator.device, dtype=training_weights_table.dtype) + weight_idx = torch.argmin((schedule_for_lookup.unsqueeze(0) - flat_t.unsqueeze(1)).abs(), dim=1) + weights = training_weights_table[weight_idx].reshape(timestep.shape) + loss = (loss_per_frame * weights.to(loss_per_frame.dtype)).mean() + else: + loss = loss_per_frame.mean() + + avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean() + train_loss += avg_loss.item() / args.gradient_accumulation_steps + + accelerator.backward(loss) + if accelerator.sync_gradients: + accelerator.clip_grad_norm_(trainable_params, args.max_grad_norm) + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + + # Checks if the accelerator has performed an optimization step behind the scenes + if accelerator.sync_gradients: + + progress_bar.update(1) + global_step += 1 + accelerator.log({"train_loss": train_loss}, step=global_step) + train_loss = 0.0 + + if global_step % args.checkpointing_steps == 0: + if args.use_deepspeed or args.use_fsdp or accelerator.is_main_process: + # _before_ saving state, check if this save would set us over the `checkpoints_total_limit` + if args.checkpoints_total_limit is not None: + checkpoints = os.listdir(args.output_dir) + checkpoints = [d for d in checkpoints if d.startswith("checkpoint")] + checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1])) + + # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints + if len(checkpoints) >= args.checkpoints_total_limit: + num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1 + removing_checkpoints = checkpoints[0:num_to_remove] + + logger.info( + f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints" + ) + logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}") + + for removing_checkpoint in removing_checkpoints: + removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint) + shutil.rmtree(removing_checkpoint) + + gc.collect() + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") + accelerator.save_state(save_path) + logger.info(f"Saved state to {save_path}") + + if args.validation_prompts is not None and global_step % args.validation_steps == 0: + log_validation( + vae, + text_encoder, + tokenizer, + transformer3d, + args, + config, + accelerator, + weight_dtype, + global_step, + ) + + logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]} + progress_bar.set_postfix(**logs) + + if global_step >= args.max_train_steps: + break + + if args.validation_prompts is not None and epoch % args.validation_epochs == 0: + log_validation( + vae, + text_encoder, + tokenizer, + transformer3d, + args, + config, + accelerator, + weight_dtype, + global_step, + ) + + # Create the pipeline using the trained modules and save it. + accelerator.wait_for_everyone() + if accelerator.is_main_process: + transformer3d = unwrap_model(transformer3d) + + if args.use_deepspeed or args.use_fsdp or accelerator.is_main_process: + gc.collect() + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") + accelerator.save_state(save_path) + logger.info(f"Saved state to {save_path}") + + accelerator.end_training() + + +if __name__ == "__main__": + main() diff --git a/scripts/wan2.1_causal_forcing/train_ar_diffusion.sh b/scripts/wan2.1_causal_forcing/train_ar_diffusion.sh new file mode 100644 index 00000000..5914fbe4 --- /dev/null +++ b/scripts/wan2.1_causal_forcing/train_ar_diffusion.sh @@ -0,0 +1,51 @@ +export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/" +export DATASET_NAME="datasets/internal_datasets/" +export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. +# export NCCL_IB_DISABLE=1 +# export NCCL_P2P_DISABLE=1 +NCCL_DEBUG=INFO + +# Causal-Forcing Stage 1: Autoregressive Diffusion Training (teacher forcing). +# - num_frame_per_block=3 -> chunkwise; set to 1 for the framewise variant. +# - shift=5.0 mirrors the Causal-Forcing default scheduler shift. +accelerate launch --mixed_precision="bf16" --use_fsdp \ + --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \ + --fsdp_transformer_layer_cls_to_wrap=CasualWanAttentionBlock \ + --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT \ + --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False \ + scripts/wan2.1_causal_forcing/train_ar_diffusion.py \ + --config_path="config/wan2.1/wan_civitai.yaml" \ + --pretrained_model_name_or_path=$MODEL_NAME \ + --train_data_dir=$DATASET_NAME \ + --train_data_meta=$DATASET_META_NAME \ + --image_sample_size=640 \ + --video_sample_size=640 \ + --token_sample_size=640 \ + --fix_sample_size 480 832 \ + --video_sample_stride=2 \ + --video_sample_n_frames=81 \ + --train_batch_size=1 \ + --gradient_accumulation_steps=1 \ + --dataloader_num_workers=8 \ + --num_train_epochs=100 \ + --checkpointing_steps=200 \ + --learning_rate=2e-06 \ + --lr_scheduler="constant_with_warmup" \ + --lr_warmup_steps=100 \ + --seed=42 \ + --output_dir="output_dir_wan2.1_causal_forcing_ar_diffusion" \ + --gradient_checkpointing \ + --mixed_precision="bf16" \ + --adam_weight_decay=3e-2 \ + --adam_epsilon=1e-10 \ + --vae_mini_batch=1 \ + --max_grad_norm=0.05 \ + --random_hw_adapt \ + --training_with_video_token_length \ + --enable_bucket \ + --num_frame_per_block=3 \ + --train_sampling_steps=1000 \ + --shift=5.0 \ + --use_timestep_weight \ + --trainable_modules "." diff --git a/scripts/wan2.1_causal_forcing/train_causal_consistency_distill.py b/scripts/wan2.1_causal_forcing/train_causal_consistency_distill.py new file mode 100644 index 00000000..d00f8073 --- /dev/null +++ b/scripts/wan2.1_causal_forcing/train_causal_consistency_distill.py @@ -0,0 +1,1651 @@ +"""Causal-Forcing Stage 2 (Causal Consistency Distillation) training, modified from +scripts/wan2.1_self_forcing/train_distill.py and Causal-Forcing +(https://github.com/thu-ml/Causal-Forcing). +""" +#!/usr/bin/env python +# coding=utf-8 +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and + +import argparse +import gc +import logging +import math +import os +import pickle +import shutil +import sys + +import accelerate +import diffusers +import numpy as np +import torch +import torch.nn.functional as F +import torch.utils.checkpoint +import transformers +from accelerate import Accelerator +from accelerate.logging import get_logger +from accelerate.state import AcceleratorState +from accelerate.utils import ProjectConfiguration, set_seed +from diffusers import FlowMatchEulerDiscreteScheduler +from diffusers.optimization import get_scheduler +from diffusers.utils import check_min_version, is_wandb_available +from diffusers.utils.torch_utils import is_compiled_module +from einops import rearrange +from omegaconf import OmegaConf +from packaging import version +from torch.utils.tensorboard import SummaryWriter +from torchvision import transforms +from tqdm.auto import tqdm +from transformers import AutoTokenizer +from transformers.utils import ContextManagers + +import datasets + +current_file_path = os.path.abspath(__file__) +project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dirname(current_file_path)), os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))] +for project_root in project_roots: + sys.path.insert(0, project_root) if project_root not in sys.path else None + +from videox_fun.data import (ASPECT_RATIO_512, ASPECT_RATIO_RANDOM_CROP_512, + ASPECT_RATIO_RANDOM_CROP_PROB, + AspectRatioBatchImageVideoSampler, + ImageVideoDataset, RandomSampler, + get_closest_ratio) +from videox_fun.models import (AutoencoderKLWan, WanT5EncoderModel, + WanTransformer3DModel_SelfForcing) +from videox_fun.pipeline import WanSelfForcingPipeline +from videox_fun.utils.utils import save_videos_grid + +if is_wandb_available(): + import wandb + + +def filter_kwargs(cls, kwargs): + import inspect + sig = inspect.signature(cls.__init__) + valid_params = set(sig.parameters.keys()) - {'self', 'cls'} + filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_params} + return filtered_kwargs + + +def get_random_downsample_ratio(sample_size, image_ratio=[], + all_choices=False, rng=None): + def _create_special_list(length): + if length == 1: + return [1.0] + if length >= 2: + first_element = 0.75 + remaining_sum = 1.0 - first_element + other_elements_value = remaining_sum / (length - 1) + special_list = [first_element] + [other_elements_value] * (length - 1) + return special_list + + if sample_size >= 1536: + number_list = [1, 1.25, 1.5, 2, 2.5, 3] + image_ratio + elif sample_size >= 1024: + number_list = [1, 1.25, 1.5, 2] + image_ratio + elif sample_size >= 768: + number_list = [1, 1.25, 1.5] + image_ratio + elif sample_size >= 512: + number_list = [1] + image_ratio + else: + number_list = [1] + + if all_choices: + return number_list + + number_list_prob = np.array(_create_special_list(len(number_list))) + if rng is None: + return np.random.choice(number_list, p = number_list_prob) + else: + return rng.choice(number_list, p = number_list_prob) + +# Will error if the minimal version of diffusers is not installed. Remove at your own risks. +check_min_version("0.18.0.dev0") + +logger = get_logger(__name__, log_level="INFO") + +def log_validation(vae, text_encoder, tokenizer, transformer3d, args, config, accelerator, weight_dtype, global_step): + try: + is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine' + if is_deepspeed: + origin_config = transformer3d.config + transformer3d.config = accelerator.unwrap_model(transformer3d).config + with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + logger.info("Running validation... ") + scheduler_kwargs = OmegaConf.to_container(config['scheduler_kwargs']) + scheduler_kwargs['shift'] = args.shift + scheduler = FlowMatchEulerDiscreteScheduler( + **filter_kwargs(FlowMatchEulerDiscreteScheduler, scheduler_kwargs) + ) + pipeline = WanSelfForcingPipeline( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d, + scheduler=scheduler, + ) + pipeline = pipeline.to(accelerator.device) + + if args.seed is None: + generator = None + else: + rank_seed = args.seed + accelerator.process_index + generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed) + logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}") + + for i in range(len(args.validation_prompts)): + if args.fix_sample_size is not None: + height, width = args.fix_sample_size + else: + height, width = args.video_sample_size, args.video_sample_size + sample = pipeline( + args.validation_prompts[i], + num_frames = args.video_sample_n_frames, + negative_prompt = args.negative_prompt, + height = height, + width = width, + generator = generator, + guidance_scale = args.validation_guidance_scale, + num_inference_steps = args.validation_num_inference_steps, + shift = args.shift, + num_frame_per_block = args.num_frame_per_block, + independent_first_frame = args.independent_first_frame, + context_noise = 0, + stochastic_sampling = False, + ).videos + os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True) + save_videos_grid( + sample, + os.path.join( + args.output_dir, + f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.mp4" + ) + ) + + del pipeline + gc.collect() + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype) + if not args.enable_text_encoder_in_dataloader: + text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype) + if is_deepspeed: + transformer3d.config = origin_config + except Exception as e: + gc.collect() + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + print(f"Eval error on rank {accelerator.process_index} with info {e}") + vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype) + if not args.enable_text_encoder_in_dataloader: + text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype) + +def parse_args(): + parser = argparse.ArgumentParser(description="Causal-Forcing Stage 2 (Option B): Causal Consistency Distillation Initialization for Wan2.1.") + parser.add_argument( + "--pretrained_model_name_or_path", + type=str, + default=None, + required=True, + help="Path to pretrained model or model identifier from huggingface.co/models.", + ) + parser.add_argument( + "--train_data_dir", + type=str, + default=None, + help=( + "A folder containing the training data. " + ), + ) + parser.add_argument( + "--train_data_meta", + type=str, + default=None, + help=( + "A csv containing the training data. " + ), + ) + parser.add_argument( + "--validation_prompts", + type=str, + default=None, + nargs="+", + help=("A set of prompts evaluated every `--validation_epochs` and logged to `--report_to`."), + ) + parser.add_argument( + "--negative_prompt", + type=str, + default="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走", + help=("The negative prompt used for validation generation."), + ) + parser.add_argument( + "--output_dir", + type=str, + default="sd-model-finetuned", + help="The output directory where the model predictions and checkpoints will be written.", + ) + parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.") + parser.add_argument( + "--use_came", + action="store_true", + help="whether to use came", + ) + parser.add_argument( + "--train_batch_size", type=int, default=16, help="Batch size (per device) for the training dataloader." + ) + parser.add_argument( + "--vae_mini_batch", type=int, default=32, help="mini batch size for vae." + ) + parser.add_argument("--num_train_epochs", type=int, default=100) + parser.add_argument( + "--max_train_steps", + type=int, + default=None, + help="Total number of training steps to perform. If provided, overrides num_train_epochs.", + ) + parser.add_argument( + "--gradient_accumulation_steps", + type=int, + default=1, + help="Number of updates steps to accumulate before performing a backward/update pass.", + ) + parser.add_argument( + "--gradient_checkpointing", + action="store_true", + help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.", + ) + parser.add_argument( + "--learning_rate", + type=float, + default=1e-4, + help="Initial learning rate (after the potential warmup period) to use.", + ) + parser.add_argument( + "--scale_lr", + action="store_true", + default=False, + help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.", + ) + parser.add_argument( + "--lr_scheduler", + type=str, + default="constant", + help=( + 'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",' + ' "constant", "constant_with_warmup"]' + ), + ) + parser.add_argument( + "--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler." + ) + parser.add_argument( + "--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes." + ) + parser.add_argument( + "--allow_tf32", + action="store_true", + help=( + "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see" + " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices" + ), + ) + parser.add_argument( + "--dataloader_num_workers", + type=int, + default=0, + help=( + "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process." + ), + ) + parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.") + parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.") + parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.") + parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer") + parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.") + parser.add_argument( + "--logging_dir", + type=str, + default="logs", + help=( + "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to" + " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***." + ), + ) + parser.add_argument( + "--mixed_precision", + type=str, + default=None, + choices=["no", "fp16", "bf16"], + help=( + "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >=" + " 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the" + " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config." + ), + ) + parser.add_argument( + "--report_to", + type=str, + default="tensorboard", + help=( + 'The integration to report the results and logs to. Supported platforms are `"tensorboard"`' + ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.' + ), + ) + parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank") + parser.add_argument( + "--checkpointing_steps", + type=int, + default=500, + help=( + "Save a checkpoint of the training state every X updates. These checkpoints are only suitable for resuming" + " training using `--resume_from_checkpoint`." + ), + ) + parser.add_argument( + "--checkpoints_total_limit", + type=int, + default=None, + help=("Max number of checkpoints to store."), + ) + parser.add_argument( + "--resume_from_checkpoint", + type=str, + default=None, + help=( + "Whether training should be resumed from a previous checkpoint. Use a path saved by" + ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.' + ), + ) + parser.add_argument( + "--validation_epochs", + type=int, + default=5, + help="Run validation every X epochs.", + ) + parser.add_argument( + "--validation_steps", + type=int, + default=2000, + help="Run validation every X steps.", + ) + parser.add_argument( + "--validation_guidance_scale", + type=float, + default=3.0, + help="CFG scale used when sampling validation videos.", + ) + parser.add_argument( + "--validation_num_inference_steps", + type=int, + default=50, + help=( + "Number of denoising steps used for validation. AR diffusion validation runs a" + " full multi-step rollout, so a larger value (e.g. 50) is recommended." + ), + ) + parser.add_argument( + "--tracker_project_name", + type=str, + default="text2image-fine-tune", + help=( + "The `project_name` argument passed to Accelerator.init_trackers for" + " more information see https://huggingface.co/docs/accelerate/v0.17.0/en/package_reference/accelerator#accelerate.Accelerator" + ), + ) + + parser.add_argument( + "--enable_text_encoder_in_dataloader", action="store_true", help="Whether or not to use text encoder in dataloader." + ) + parser.add_argument( + "--enable_bucket", action="store_true", help="Whether enable bucket sample in datasets." + ) + parser.add_argument( + "--random_ratio_crop", action="store_true", help="Whether enable random ratio crop sample in datasets." + ) + parser.add_argument( + "--random_hw_adapt", action="store_true", help="Whether enable random adapt height and width in datasets." + ) + parser.add_argument( + "--training_with_video_token_length", action="store_true", help="The training stage of the model in training.", + ) + parser.add_argument( + "--train_sampling_steps", + type=int, + default=1000, + help="Total number of scheduler timesteps for sampling.", + ) + parser.add_argument( + "--token_sample_size", + type=int, + default=512, + help="Sample size of the token.", + ) + parser.add_argument( + "--video_sample_size", + type=int, + default=512, + help="Sample size of the video.", + ) + parser.add_argument( + "--image_sample_size", + type=int, + default=512, + help="Sample size of the image.", + ) + parser.add_argument( + "--fix_sample_size", + nargs=2, type=int, default=None, + help="Fix Sample size [height, width] when using bucket and collate_fn." + ) + parser.add_argument( + "--video_sample_stride", + type=int, + default=4, + help="Sample stride of the video.", + ) + parser.add_argument( + "--video_sample_n_frames", + type=int, + default=17, + help="Num frame of video.", + ) + parser.add_argument( + "--video_repeat", + type=int, + default=0, + help="Num of repeat video.", + ) + parser.add_argument( + "--config_path", + type=str, + default=None, + help=( + "The config of the model in training." + ), + ) + parser.add_argument( + "--transformer_path", + type=str, + default=None, + help=("If you want to load the weight from other transformers, input its path."), + ) + parser.add_argument( + "--vae_path", + type=str, + default=None, + help=("If you want to load the weight from other vaes, input its path."), + ) + + parser.add_argument( + '--trainable_modules', + nargs='+', + help='Enter a list of trainable modules' + ) + parser.add_argument( + '--trainable_modules_low_learning_rate', + nargs='+', + default=[], + help='Enter a list of trainable modules with lower learning rate' + ) + parser.add_argument( + '--tokenizer_max_length', + type=int, + default=512, + help='Max length of tokenizer' + ) + parser.add_argument( + "--use_deepspeed", action="store_true", help="Whether or not to use deepspeed." + ) + parser.add_argument( + "--use_fsdp", action="store_true", help="Whether or not to use fsdp." + ) + parser.add_argument( + "--low_vram", action="store_true", help="Whether enable low_vram mode." + ) + parser.add_argument( + "--num_frame_per_block", + type=int, + default=3, + help="Number of latent frames per causal block. 3 = chunk-wise, 1 = frame-wise." + ) + parser.add_argument( + "--independent_first_frame", + action="store_true", + help="Whether first frame is independent ([1, N, N, ...] pattern, useful for I2V)." + ) + parser.add_argument( + "--shift", + type=float, + default=5.0, + help="Shift value for FlowMatchEulerDiscreteScheduler. Causal-Forcing uses 5.0 by default." + ) + parser.add_argument( + "--discrete_cd_N", + type=int, + default=48, + help="Number of discrete timesteps used for the consistency schedule (`discrete_cd_N` in Causal-Forcing). Default: 48." + ) + parser.add_argument( + "--guidance_scale", + type=float, + default=3.0, + help="Classifier-free guidance scale applied to the frozen teacher when generating the one-step ODE target. Default: 3.0." + ) + parser.add_argument( + "--ema_weight", + type=float, + default=0.99, + help="EMA decay for the consistency-target generator copy. Set <=0 to disable EMA and use the live generator as the target." + ) + parser.add_argument( + "--ema_start_step", + type=int, + default=200, + help="Number of optimizer steps to wait before EMA tracking starts. Before this point the EMA copy mirrors the live generator." + ) + parser.add_argument( + "--teacher_transformer_path", + type=str, + default=None, + help="Optional path to a separate teacher (Stage 1 AR-diffusion) safetensors checkpoint. If unset, the teacher is initialised from the same weights as the generator." + ) + args = parser.parse_args() + env_local_rank = int(os.environ.get("LOCAL_RANK", -1)) + if env_local_rank != -1 and env_local_rank != args.local_rank: + args.local_rank = env_local_rank + + return args + + +def main(): + args = parse_args() + + logging_dir = os.path.join(args.output_dir, args.logging_dir) + + config = OmegaConf.load(args.config_path) + accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir) + + accelerator = Accelerator( + gradient_accumulation_steps=args.gradient_accumulation_steps, + mixed_precision=args.mixed_precision, + log_with=args.report_to, + project_config=accelerator_project_config, + ) + + deepspeed_plugin = accelerator.state.deepspeed_plugin if hasattr(accelerator.state, "deepspeed_plugin") else None + fsdp_plugin = accelerator.state.fsdp_plugin if hasattr(accelerator.state, "fsdp_plugin") else None + if deepspeed_plugin is not None: + zero_stage = int(deepspeed_plugin.zero_stage) + fsdp_stage = 0 + print(f"Using DeepSpeed Zero stage: {zero_stage}") + + args.use_deepspeed = True + if zero_stage == 3: + print(f"Auto set save_state to True because zero_stage == 3") + args.save_state = True + elif fsdp_plugin is not None: + from torch.distributed.fsdp import ShardingStrategy + zero_stage = 0 + if fsdp_plugin.sharding_strategy is ShardingStrategy.FULL_SHARD: + fsdp_stage = 3 + elif fsdp_plugin.sharding_strategy is None: # The fsdp_plugin.sharding_strategy is None in FSDP 2. + fsdp_stage = 3 + elif fsdp_plugin.sharding_strategy is ShardingStrategy.SHARD_GRAD_OP: + fsdp_stage = 2 + else: + fsdp_stage = 0 + print(f"Using FSDP stage: {fsdp_stage}") + + args.use_fsdp = True + if fsdp_stage == 3: + print(f"Auto set save_state to True because fsdp_stage == 3") + args.save_state = True + else: + zero_stage = 0 + fsdp_stage = 0 + print("DeepSpeed is not enabled.") + + if accelerator.is_main_process: + writer = SummaryWriter(log_dir=logging_dir) + + # Make one log on every process with the configuration for debugging. + logging.basicConfig( + format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", + datefmt="%m/%d/%Y %H:%M:%S", + level=logging.INFO, + ) + logger.info(accelerator.state, main_process_only=False) + if accelerator.is_local_main_process: + datasets.utils.logging.set_verbosity_warning() + transformers.utils.logging.set_verbosity_warning() + diffusers.utils.logging.set_verbosity_info() + else: + datasets.utils.logging.set_verbosity_error() + transformers.utils.logging.set_verbosity_error() + diffusers.utils.logging.set_verbosity_error() + + # If passed along, set the training seed now. + if args.seed is not None: + set_seed(args.seed) + rng = np.random.default_rng(np.random.PCG64(args.seed + accelerator.process_index)) + torch_rng = torch.Generator(accelerator.device).manual_seed(args.seed + accelerator.process_index) + print(f"Init rng with seed {args.seed + accelerator.process_index}. Process_index is {accelerator.process_index}") + else: + rng = None + torch_rng = None + print(f"No seed provided; using global default RNG. Process_index is {accelerator.process_index}") + + # Handle the repository creation + if accelerator.is_main_process: + if args.output_dir is not None: + os.makedirs(args.output_dir, exist_ok=True) + + # For mixed precision training we cast all non-trainable weigths (vae, non-lora text_encoder and non-lora transformer3d) to half-precision + # as these weights are only used for inference, keeping weights in full precision is not required. + weight_dtype = torch.float32 + if accelerator.mixed_precision == "fp16": + weight_dtype = torch.float16 + args.mixed_precision = accelerator.mixed_precision + elif accelerator.mixed_precision == "bf16": + weight_dtype = torch.bfloat16 + args.mixed_precision = accelerator.mixed_precision + + # Load scheduler, tokenizer and models. + scheduler_kwargs = OmegaConf.to_container(config['scheduler_kwargs']) + scheduler_kwargs['shift'] = args.shift + noise_scheduler = FlowMatchEulerDiscreteScheduler( + **filter_kwargs(FlowMatchEulerDiscreteScheduler, scheduler_kwargs) + ) + + # Get Tokenizer + tokenizer = AutoTokenizer.from_pretrained( + os.path.join(args.pretrained_model_name_or_path, config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer')), + ) + + def deepspeed_zero_init_disabled_context_manager(): + """ + returns either a context list that includes one that will disable zero.Init or an empty context list + """ + deepspeed_plugin = AcceleratorState().deepspeed_plugin if accelerate.state.is_initialized() else None + if deepspeed_plugin is None: + return [] + + return [deepspeed_plugin.zero3_init_context_manager(enable=False)] + + # Currently Accelerate doesn't know how to handle multiple models under Deepspeed ZeRO stage 3. + # For this to work properly all models must be run through `accelerate.prepare`. But accelerate + # will try to assign the same optimizer with the same weights to all models during + # `deepspeed.initialize`, which of course doesn't work. + # + # For now the following workaround will partially support Deepspeed ZeRO-3, by excluding the 2 + # frozen models from being partitioned during `zero.Init` which gets called during + # `from_pretrained` So CLIPTextModel and AutoencoderKL will not enjoy the parameter sharding + # across multiple gpus and only UNet2DConditionModel will get ZeRO sharded. + with ContextManagers(deepspeed_zero_init_disabled_context_manager()): + # Get Text encoder + text_encoder = WanT5EncoderModel.from_pretrained( + os.path.join(args.pretrained_model_name_or_path, config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')), + additional_kwargs=OmegaConf.to_container(config['text_encoder_kwargs']), + low_cpu_mem_usage=True, + torch_dtype=weight_dtype, + ) + text_encoder = text_encoder.eval() + # Get Vae + vae = AutoencoderKLWan.from_pretrained( + os.path.join(args.pretrained_model_name_or_path, config['vae_kwargs'].get('vae_subpath', 'vae')), + additional_kwargs=OmegaConf.to_container(config['vae_kwargs']), + ) + vae.eval() + + # Get Transformer (causal generator) + # IMPORTANT: keep the trainable transformer in fp32. accelerate's + # mixed_precision="bf16" will autocast the forward to bf16 while keeping + # the master weights and Adam moments in fp32. If params live in bf16, + # every update (LR*grad ~ 1e-5 for LR=2e-6) falls below bf16 mantissa + # precision (~1e-3 relative) and is rounded to zero — Stage 1 hit exactly + # this and CCD has the same LR + same Adam betas so it would hit it again. + # CF official keeps fp32 master weights via FSDP's default behavior + # (MixedPrecision(param_dtype=bf16) with no compute_dtype set). + transformer3d = WanTransformer3DModel_SelfForcing.from_pretrained( + os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')), + transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']), + low_cpu_mem_usage=True, + torch_dtype=torch.float32, + ) + # Stage 2 CCD trains the causal generator with the same block layout the + # downstream Self-Forcing pipeline expects at sampling time. + transformer3d.num_frame_per_block = args.num_frame_per_block + transformer3d.independent_first_frame = args.independent_first_frame + + # Freeze vae and text_encoder; transformer3d is toggled per-module below. + vae.requires_grad_(False) + text_encoder.requires_grad_(False) + transformer3d.requires_grad_(False) + + if args.transformer_path is not None: + print(f"From checkpoint: {args.transformer_path}") + if args.transformer_path.endswith("safetensors"): + from safetensors.torch import load_file, safe_open + state_dict = load_file(args.transformer_path) + else: + state_dict = torch.load(args.transformer_path, map_location="cpu") + state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict + state_dict = state_dict["generator_ema"] if "generator_ema" in state_dict else state_dict + state_dict = state_dict["generator"] if "generator" in state_dict else state_dict + if any(k.startswith("model.") for k in state_dict.keys()): + state_dict = {k.replace("model.", "", 1) if k.startswith("model.") else k: v for k, v in state_dict.items()} + + m, u = transformer3d.load_state_dict(state_dict, strict=False) + print(f"missing keys: {len(m)}, unexpected keys: {len(u)}") + assert len(u) == 0 + + # EMA target: a non-trainable copy of the generator used as the + # right-hand-side of the consistency objective. We use a real module + # (instead of diffusers' EMAModel) so the CCD forward can call it directly, + # which keeps the implementation parallel to Causal-Forcing's + # `model/naive_consistency.py::NaiveConsistency`. + use_ema = args.ema_weight is not None and args.ema_weight > 0 + if use_ema: + # EMA must shadow the fp32 generator: the polyak update + # `ema = decay*ema + (1-decay)*gen` mixes two values one mantissa apart, + # which underflows in bf16 for decay=0.99 (the smaller-magnitude term + # gets rounded away). + ema_transformer3d = WanTransformer3DModel_SelfForcing.from_pretrained( + os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')), + transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']), + low_cpu_mem_usage=True, + torch_dtype=torch.float32, + ) + ema_transformer3d.num_frame_per_block = args.num_frame_per_block + ema_transformer3d.independent_first_frame = args.independent_first_frame + ema_transformer3d.requires_grad_(False) + ema_transformer3d.eval() + ema_transformer3d.load_state_dict(transformer3d.state_dict(), strict=True) + else: + ema_transformer3d = None + + # Teacher: frozen Stage 1 AR-diffusion model used to produce the one-step + # ODE target. If `--teacher_transformer_path` is unset, it shares weights + # with the (still-frozen) generator that was just loaded. + # Teacher is frozen but kept fp32 for numerical parity with the generator + # at init (we mirror the generator's state_dict below when no separate + # teacher path is given). Cast to bf16 happens via autocast at call time. + teacher_transformer3d = WanTransformer3DModel_SelfForcing.from_pretrained( + os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')), + transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']), + low_cpu_mem_usage=True, + torch_dtype=torch.float32, + ) + teacher_transformer3d.num_frame_per_block = args.num_frame_per_block + teacher_transformer3d.independent_first_frame = args.independent_first_frame + teacher_transformer3d.requires_grad_(False) + teacher_transformer3d.eval() + if args.teacher_transformer_path is not None: + print(f"From checkpoint: {args.teacher_transformer_path}") + if args.teacher_transformer_path.endswith("safetensors"): + from safetensors.torch import load_file, safe_open + state_dict = load_file(args.teacher_transformer_path) + else: + state_dict = torch.load(args.teacher_transformer_path, map_location="cpu") + state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict + state_dict = state_dict["generator_ema"] if "generator_ema" in state_dict else state_dict + state_dict = state_dict["generator"] if "generator" in state_dict else state_dict + if any(k.startswith("model.") for k in state_dict.keys()): + state_dict = {k.replace("model.", "", 1) if k.startswith("model.") else k: v for k, v in state_dict.items()} + + m, u = teacher_transformer3d.load_state_dict(state_dict, strict=False) + print(f"missing keys: {len(m)}, unexpected keys: {len(u)}") + assert len(u) == 0 + elif args.transformer_path is not None: + # Mirror the generator weights into the teacher. + teacher_transformer3d.load_state_dict(transformer3d.state_dict(), strict=True) + + if args.vae_path is not None: + print(f"From checkpoint: {args.vae_path}") + if args.vae_path.endswith("safetensors"): + from safetensors.torch import load_file, safe_open + state_dict = load_file(args.vae_path) + else: + state_dict = torch.load(args.vae_path, map_location="cpu") + state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict + + m, u = vae.load_state_dict(state_dict, strict=False) + print(f"missing keys: {len(m)}, unexpected keys: {len(u)}") + assert len(u) == 0 + + # A good trainable modules is showed below now. + # For 3D Patch: trainable_modules = ['ff.net', 'pos_embed', 'attn2', 'proj_out', 'timepositionalencoding', 'h_position', 'w_position'] + # For 2D Patch: trainable_modules = ['ff.net', 'attn2', 'timepositionalencoding', 'h_position', 'w_position'] + transformer3d.train() + if accelerator.is_main_process: + accelerator.print( + f"Trainable modules '{args.trainable_modules}'." + ) + for name, param in transformer3d.named_parameters(): + for trainable_module_name in args.trainable_modules + args.trainable_modules_low_learning_rate: + if trainable_module_name in name: + param.requires_grad = True + break + + # `accelerate` 0.16.0 will have better support for customized saving + if version.parse(accelerate.__version__) >= version.parse("0.16.0"): + # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format + def _save_ema_pretrained(output_dir): + if ema_transformer3d is not None and accelerator.is_main_process: + ema_transformer3d.save_pretrained(os.path.join(output_dir, "ema_transformer")) + + def _load_ema_pretrained(input_dir): + if ema_transformer3d is None: + return + ema_path = os.path.join(input_dir, "ema_transformer") + if not os.path.exists(ema_path): + return + ema_loaded = WanTransformer3DModel_SelfForcing.from_pretrained( + ema_path, + transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']), + low_cpu_mem_usage=True, + ) + ema_transformer3d.load_state_dict(ema_loaded.state_dict(), strict=True) + del ema_loaded + print(f"Loaded EMA generator from {ema_path}.") + + if fsdp_stage != 0 or zero_stage == 3: + def save_model_hook(models, weights, output_dir): + accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True) + if accelerator.is_main_process: + from safetensors.torch import save_file + + safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors") + accelerate_state_dict = {k: v.to(dtype=weight_dtype) for k, v in accelerate_state_dict.items()} + save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"}) + + with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file: + pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file) + _save_ema_pretrained(output_dir) + + def load_model_hook(models, input_dir): + pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl") + if os.path.exists(pkl_path): + with open(pkl_path, 'rb') as file: + loaded_number, _ = pickle.load(file) + batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0) + print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.") + _load_ema_pretrained(input_dir) + else: + # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format + def save_model_hook(models, weights, output_dir): + if accelerator.is_main_process: + models[0].save_pretrained(os.path.join(output_dir, "transformer")) + if not args.use_deepspeed: + weights.pop() + + with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file: + pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file) + _save_ema_pretrained(output_dir) + + def load_model_hook(models, input_dir): + for i in range(len(models)): + # pop models so that they are not loaded again + model = models.pop() + + # load diffusers style into model + load_model = WanTransformer3DModel.from_pretrained( + input_dir, subfolder="transformer" + ) + model.register_to_config(**load_model.config) + + model.load_state_dict(load_model.state_dict()) + del load_model + + pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl") + if os.path.exists(pkl_path): + with open(pkl_path, 'rb') as file: + loaded_number, _ = pickle.load(file) + batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0) + print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.") + _load_ema_pretrained(input_dir) + + accelerator.register_save_state_pre_hook(save_model_hook) + accelerator.register_load_state_pre_hook(load_model_hook) + + if args.gradient_checkpointing: + transformer3d.enable_gradient_checkpointing() + + # Enable TF32 for faster training on Ampere GPUs, + # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices + if args.allow_tf32: + torch.backends.cuda.matmul.allow_tf32 = True + + if args.scale_lr: + args.learning_rate = ( + args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes + ) + + # Initialize the optimizer + if args.use_8bit_adam: + try: + import bitsandbytes as bnb + except ImportError: + raise ImportError( + "Please install bitsandbytes to use 8-bit Adam. You can do so by running `pip install bitsandbytes`" + ) + + optimizer_cls = bnb.optim.AdamW8bit + elif args.use_came: + try: + from came_pytorch import CAME + except Exception: + raise ImportError( + "Please install came_pytorch to use CAME. You can do so by running `pip install came_pytorch`" + ) + + optimizer_cls = CAME + else: + optimizer_cls = torch.optim.AdamW + + trainable_params = list(filter(lambda p: p.requires_grad, transformer3d.parameters())) + trainable_params_optim = [ + {'params': [], 'lr': args.learning_rate}, + {'params': [], 'lr': args.learning_rate / 2}, + ] + in_already = [] + for name, param in transformer3d.named_parameters(): + high_lr_flag = False + if name in in_already: + continue + for trainable_module_name in args.trainable_modules: + if trainable_module_name in name: + in_already.append(name) + high_lr_flag = True + trainable_params_optim[0]['params'].append(param) + if accelerator.is_main_process: + print(f"Set {name} to lr : {args.learning_rate}") + break + if high_lr_flag: + continue + for trainable_module_name in args.trainable_modules_low_learning_rate: + if trainable_module_name in name: + in_already.append(name) + trainable_params_optim[1]['params'].append(param) + if accelerator.is_main_process: + print(f"Set {name} to lr : {args.learning_rate / 2}") + break + + if args.use_came: + optimizer = optimizer_cls( + trainable_params_optim, + lr=args.learning_rate, + # weight_decay=args.adam_weight_decay, + betas=(0.9, 0.999, 0.9999), + eps=(1e-30, 1e-16) + ) + else: + optimizer = optimizer_cls( + trainable_params_optim, + lr=args.learning_rate, + betas=(args.adam_beta1, args.adam_beta2), + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + ) + + # Get the training dataset + sample_n_frames_bucket_interval = vae.config.temporal_compression_ratio + + if args.fix_sample_size is not None and args.enable_bucket: + args.video_sample_size = max(max(args.fix_sample_size), args.video_sample_size) + args.image_sample_size = max(max(args.fix_sample_size), args.image_sample_size) + args.training_with_video_token_length = False + args.random_hw_adapt = False + + # Get the dataset (Stage 1 always trains on raw videos with teacher forcing). + train_dataset = ImageVideoDataset( + args.train_data_meta, args.train_data_dir, + video_sample_size=args.video_sample_size, video_sample_stride=args.video_sample_stride, video_sample_n_frames=args.video_sample_n_frames, + video_repeat=args.video_repeat, + image_sample_size=args.image_sample_size, + enable_bucket=args.enable_bucket, enable_inpaint=False, + ) + + def get_length_to_frame_num(token_length): + if args.image_sample_size > args.video_sample_size: + sample_sizes = list(range(args.video_sample_size, args.image_sample_size + 1, 128)) + + if sample_sizes[-1] != args.image_sample_size: + sample_sizes.append(args.image_sample_size) + else: + sample_sizes = [args.image_sample_size] + + length_to_frame_num = { + sample_size: min(token_length / sample_size / sample_size, args.video_sample_n_frames) // sample_n_frames_bucket_interval * sample_n_frames_bucket_interval + 1 for sample_size in sample_sizes + } + + return length_to_frame_num + + aspect_ratio_sample_size = {key : [x / 512 * args.video_sample_size for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()} + batch_sampler_generator = torch.Generator().manual_seed(args.seed) + batch_sampler = AspectRatioBatchImageVideoSampler( + sampler=RandomSampler(train_dataset, generator=batch_sampler_generator), dataset=train_dataset.dataset, + batch_size=args.train_batch_size, train_folder = args.train_data_dir, drop_last=True, + aspect_ratios=aspect_ratio_sample_size, + ) + + def collate_fn(examples): + # Get token length + target_token_length = args.video_sample_n_frames * args.token_sample_size * args.token_sample_size + length_to_frame_num = get_length_to_frame_num(target_token_length) + + # Create new output + new_examples = {} + new_examples["target_token_length"] = target_token_length + new_examples["pixel_values"] = [] + new_examples["text"] = [] + + # Get downsample ratio in image and videos + pixel_value = examples[0]["pixel_values"] + data_type = examples[0]["data_type"] + f, h, w, c = np.shape(pixel_value) + if data_type == 'image': + random_downsample_ratio = 1 if not args.random_hw_adapt else get_random_downsample_ratio(args.image_sample_size, image_ratio=[args.image_sample_size / args.video_sample_size], rng=rng) + + aspect_ratio_sample_size = {key : [x / 512 * args.image_sample_size / random_downsample_ratio for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()} + aspect_ratio_random_crop_sample_size = {key : [x / 512 * args.image_sample_size / random_downsample_ratio for x in ASPECT_RATIO_RANDOM_CROP_512[key]] for key in ASPECT_RATIO_RANDOM_CROP_512.keys()} + + batch_video_length = args.video_sample_n_frames + sample_n_frames_bucket_interval + else: + if args.random_hw_adapt: + if args.training_with_video_token_length: + local_min_size = np.min(np.array([np.mean(np.array([np.shape(example["pixel_values"])[1], np.shape(example["pixel_values"])[2]])) for example in examples])) + # The video will be resized to a lower resolution than its own. + choice_list = [length for length in list(length_to_frame_num.keys()) if length < local_min_size * 1.25] + if len(choice_list) == 0: + choice_list = list(length_to_frame_num.keys()) + if rng is None: + local_video_sample_size = np.random.choice(choice_list) + else: + local_video_sample_size = rng.choice(choice_list) + batch_video_length = length_to_frame_num[local_video_sample_size] + random_downsample_ratio = args.video_sample_size / local_video_sample_size + else: + random_downsample_ratio = get_random_downsample_ratio( + args.video_sample_size, rng=rng) + batch_video_length = args.video_sample_n_frames + sample_n_frames_bucket_interval + else: + random_downsample_ratio = 1 + batch_video_length = args.video_sample_n_frames + sample_n_frames_bucket_interval + + aspect_ratio_sample_size = {key : [x / 512 * args.video_sample_size / random_downsample_ratio for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()} + aspect_ratio_random_crop_sample_size = {key : [x / 512 * args.video_sample_size / random_downsample_ratio for x in ASPECT_RATIO_RANDOM_CROP_512[key]] for key in ASPECT_RATIO_RANDOM_CROP_512.keys()} + + if args.fix_sample_size is not None: + fix_sample_size = [int(x / 16) * 16 for x in args.fix_sample_size] + elif args.random_ratio_crop: + if rng is None: + random_sample_size = aspect_ratio_random_crop_sample_size[ + np.random.choice(list(aspect_ratio_random_crop_sample_size.keys()), p = ASPECT_RATIO_RANDOM_CROP_PROB) + ] + else: + random_sample_size = aspect_ratio_random_crop_sample_size[ + rng.choice(list(aspect_ratio_random_crop_sample_size.keys()), p = ASPECT_RATIO_RANDOM_CROP_PROB) + ] + random_sample_size = [int(x / 16) * 16 for x in random_sample_size] + else: + closest_size, closest_ratio = get_closest_ratio(h, w, ratios=aspect_ratio_sample_size) + closest_size = [int(x / 16) * 16 for x in closest_size] + + min_example_length = min( + [example["pixel_values"].shape[0] for example in examples] + ) + batch_video_length = int(min(batch_video_length, min_example_length)) + + # Magvae needs the number of frames to be 4n + 1. + batch_video_length = (batch_video_length - 1) // sample_n_frames_bucket_interval * sample_n_frames_bucket_interval + 1 + + # Causal-Forcing needs the latent frame count to align with num_frame_per_block. + k = (batch_video_length - 1) // sample_n_frames_bucket_interval + if args.independent_first_frame: + # latent_frames - 1 = k must be divisible by num_frame_per_block + k = (k // args.num_frame_per_block) * args.num_frame_per_block + else: + # latent_frames = k + 1 must be divisible by num_frame_per_block + k = ((k + 1) // args.num_frame_per_block) * args.num_frame_per_block - 1 + batch_video_length = k * sample_n_frames_bucket_interval + 1 + + if batch_video_length <= 0: + batch_video_length = 1 + + for example in examples: + if args.fix_sample_size is not None: + # To 0~1 + pixel_values = torch.from_numpy(example["pixel_values"]).permute(0, 3, 1, 2).contiguous() + pixel_values = pixel_values / 255. + + # Get adapt hw for resize + fix_sample_size = list(map(lambda x: int(x), fix_sample_size)) + transform = transforms.Compose([ + transforms.Resize(fix_sample_size, interpolation=transforms.InterpolationMode.BILINEAR), # Image.BICUBIC + transforms.CenterCrop(fix_sample_size), + transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True), + ]) + elif args.random_ratio_crop: + # To 0~1 + pixel_values = torch.from_numpy(example["pixel_values"]).permute(0, 3, 1, 2).contiguous() + pixel_values = pixel_values / 255. + + # Get adapt hw for resize + b, c, h, w = pixel_values.size() + th, tw = random_sample_size + if th / tw > h / w: + nh = int(th) + nw = int(w / h * nh) + else: + nw = int(tw) + nh = int(h / w * nw) + + transform = transforms.Compose([ + transforms.Resize([nh, nw]), + transforms.CenterCrop([int(x) for x in random_sample_size]), + transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True), + ]) + else: + # To 0~1 + pixel_values = torch.from_numpy(example["pixel_values"]).permute(0, 3, 1, 2).contiguous() + pixel_values = pixel_values / 255. + + # Get adapt hw for resize + closest_size = list(map(lambda x: int(x), closest_size)) + if closest_size[0] / h > closest_size[1] / w: + resize_size = closest_size[0], int(w * closest_size[0] / h) + else: + resize_size = int(h * closest_size[1] / w), closest_size[1] + + transform = transforms.Compose([ + transforms.Resize(resize_size, interpolation=transforms.InterpolationMode.BILINEAR), # Image.BICUBIC + transforms.CenterCrop(closest_size), + transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True), + ]) + + new_examples["pixel_values"].append(transform(pixel_values)[:batch_video_length]) + new_examples["text"].append(example["text"]) + + # Limit the number of frames to the same + new_examples["pixel_values"] = torch.stack([example for example in new_examples["pixel_values"]]) + + # Encode prompts when enable_text_encoder_in_dataloader=True + if args.enable_text_encoder_in_dataloader: + prompt_ids = tokenizer( + new_examples['text'], + max_length=args.tokenizer_max_length, + padding="max_length", + add_special_tokens=True, + truncation=True, + return_tensors="pt" + ) + text_input_ids = prompt_ids.input_ids + prompt_attention_mask = prompt_ids.attention_mask + + seq_lens = prompt_attention_mask.gt(0).sum(dim=1).long() + prompt_embeds = text_encoder(text_input_ids.to("cpu"), attention_mask=prompt_attention_mask.to("cpu"))[0] + prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] + + new_examples['encoder_attention_mask'] = prompt_ids.attention_mask + new_examples['encoder_hidden_states'] = prompt_embeds + + return new_examples + + # DataLoaders creation: + train_dataloader = torch.utils.data.DataLoader( + train_dataset, + batch_sampler=batch_sampler, + collate_fn=collate_fn, + persistent_workers=True if args.dataloader_num_workers != 0 else False, + num_workers=args.dataloader_num_workers, + ) + + # Scheduler and math around the number of training steps. + overrode_max_train_steps = False + num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) + if args.max_train_steps is None: + args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch + overrode_max_train_steps = True + + lr_scheduler = get_scheduler( + args.lr_scheduler, + optimizer=optimizer, + num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes, + num_training_steps=args.max_train_steps * accelerator.num_processes, + ) + + # Prepare everything with our `accelerator`. + transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( + transformer3d, optimizer, train_dataloader, lr_scheduler + ) + if fsdp_stage != 0 or zero_stage != 0: + from functools import partial + + from videox_fun.dist import set_multi_gpus_devices, shard_model + shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype) + text_encoder = shard_fn(text_encoder) + + # Move text_encode and vae to gpu and cast to weight_dtype + vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype) + if not args.enable_text_encoder_in_dataloader: + text_encoder.to(accelerator.device if not args.low_vram else "cpu") + + # Move the frozen teacher / EMA copies to the same device as the generator. + # They are never wrapped by accelerator.prepare(), so we keep them in + # inference mode on the local device only. + # DO NOT pass dtype=weight_dtype here — we deliberately keep both modules + # in fp32 to match the generator's master-weight precision: + # * Teacher's forward is autocast to bf16 anyway, so fp32 storage is + # just a parity/precision-safety choice (extra ~5GB on 1.3B is fine + # on a GB200). + # * EMA polyak update below (`v.mul_(decay).add_(..., alpha=1-decay)`) + # runs in v.dtype. With decay=0.99 the (1-decay)*delta term has + # magnitude ~LR*grad*0.01 ~ 1e-8, which underflows in bf16. EMA MUST + # be fp32 or the consistency target stops tracking the live generator. + teacher_transformer3d.to(accelerator.device) + if ema_transformer3d is not None: + ema_transformer3d.to(accelerator.device) + + # We need to recalculate our total training steps as the size of the training dataloader may have changed. + num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) + if overrode_max_train_steps: + args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch + # Afterwards we recalculate our number of training epochs + args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch) + + # We need to initialize the trackers we use, and also store our configuration. + # The trackers initializes automatically on the main process. + if accelerator.is_main_process: + tracker_config = dict(vars(args)) + keys_to_pop = [k for k, v in tracker_config.items() if isinstance(v, list)] + for k in keys_to_pop: + tracker_config.pop(k) + print(f"Removed tracker_config['{k}']") + accelerator.init_trackers(args.tracker_project_name, tracker_config) + + # Function for unwrapping if model was compiled with `torch.compile`. + def unwrap_model(model): + model = accelerator.unwrap_model(model) + model = model._orig_mod if is_compiled_module(model) else model + return model + + # Train! + total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps + + logger.info("***** Running training *****") + logger.info(f" Num examples = {len(train_dataset)}") + logger.info(f" Num Epochs = {args.num_train_epochs}") + logger.info(f" Instantaneous batch size per device = {args.train_batch_size}") + logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}") + logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}") + logger.info(f" Total optimization steps = {args.max_train_steps}") + global_step = 0 + first_epoch = 0 + + # Potentially load in the weights and states from a previous save + if args.resume_from_checkpoint: + if args.resume_from_checkpoint != "latest": + path = os.path.basename(args.resume_from_checkpoint) + else: + # Get the most recent checkpoint + dirs = os.listdir(args.output_dir) + dirs = [d for d in dirs if d.startswith("checkpoint")] + dirs = sorted(dirs, key=lambda x: int(x.split("-")[1])) + path = dirs[-1] if len(dirs) > 0 else None + + if path is None: + accelerator.print( + f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run." + ) + args.resume_from_checkpoint = None + initial_global_step = 0 + else: + global_step = int(path.split("-")[1]) + + initial_global_step = global_step + + pkl_path = os.path.join(os.path.join(args.output_dir, path), "sampler_pos_start.pkl") + if os.path.exists(pkl_path): + with open(pkl_path, 'rb') as file: + _, first_epoch = pickle.load(file) + else: + first_epoch = global_step // num_update_steps_per_epoch + print(f"Load pkl from {pkl_path}. Get first_epoch = {first_epoch}.") + + accelerator.print(f"Resuming from checkpoint {path}") + accelerator.load_state(os.path.join(args.output_dir, path)) + else: + initial_global_step = 0 + + progress_bar = tqdm( + range(0, args.max_train_steps), + initial=initial_global_step, + desc="Steps", + # Only show the progress bar once on each machine. + disable=not accelerator.is_local_main_process, + ) + + # Materialise the discrete consistency schedule with `args.discrete_cd_N` + # timesteps. The CCD objective walks adjacent (t, t_next) pairs from this + # grid — same convention as `FlowMatchScheduler.set_timesteps(num_inference_steps=discrete_cd_N)` + # in Causal-Forcing's `model/naive_consistency.py`. + noise_scheduler.set_timesteps(args.discrete_cd_N, device=accelerator.device) + cd_timesteps = noise_scheduler.timesteps.clone() # length N, sorted high -> low + cd_sigmas = noise_scheduler.sigmas.clone() # length N+1 (final 0 entry) + # Lookup table used by `add_noise` / x0-conversion below. + schedule_timesteps = cd_timesteps.to(accelerator.device) + + # Pre-encode the unconditional (negative) prompt once; reused for every + # teacher CFG step. Under `--low_vram` the text encoder is parked on CPU, + # so move it to the device for this one-off encode then put it back. + with torch.no_grad(): + if args.low_vram: + text_encoder.to(accelerator.device) + neg_inputs = tokenizer( + [args.negative_prompt], padding="max_length", + max_length=args.tokenizer_max_length, truncation=True, + add_special_tokens=True, return_tensors="pt", + ) + neg_seq_len = neg_inputs.attention_mask.gt(0).sum(dim=1).long() + neg_embed = text_encoder(neg_inputs.input_ids.to(accelerator.device), attention_mask=neg_inputs.attention_mask.to(accelerator.device))[0] + negative_prompt_embed = neg_embed[0, :neg_seq_len[0]].detach() + if args.low_vram: + text_encoder.to("cpu") + torch.cuda.empty_cache() + + for epoch in range(first_epoch, args.num_train_epochs): + train_loss = 0.0 + batch_sampler.sampler.generator = torch.Generator().manual_seed(args.seed + epoch) + for step, batch in enumerate(train_dataloader): + # Data batch sanity check + if epoch == first_epoch and step == 0: + pixel_values, texts = batch['pixel_values'].cpu(), batch['text'] + pixel_values = rearrange(pixel_values, "b f c h w -> b c f h w") + os.makedirs(os.path.join(args.output_dir, "sanity_check"), exist_ok=True) + for idx, (pixel_value, text) in enumerate(zip(pixel_values, texts)): + pixel_value = pixel_value[None, ...] + gif_name = '-'.join(text.replace('/', '').split()[:10]) if not text == '' else f'{global_step}-{idx}' + save_videos_grid(pixel_value, f"{args.output_dir}/sanity_check/{gif_name[:10]}.mp4", rescale=True) + + with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + # Convert prompts to text embeddings. + if args.enable_text_encoder_in_dataloader: + prompt_embeds = batch['encoder_hidden_states'].to(device=accelerator.device) + else: + with torch.no_grad(): + prompt_ids = tokenizer( + batch['text'], + padding="max_length", + max_length=args.tokenizer_max_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt" + ) + text_input_ids = prompt_ids.input_ids + prompt_attention_mask = prompt_ids.attention_mask + + seq_lens = prompt_attention_mask.gt(0).sum(dim=1).long() + # Under --low_vram the encoder is parked on CPU; move it to + # the device just for this forward pass and park it back. + if args.low_vram: + text_encoder.to(accelerator.device) + prompt_embeds = text_encoder(text_input_ids.to(accelerator.device), attention_mask=prompt_attention_mask.to(accelerator.device))[0] + prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] + if args.low_vram: + text_encoder.to('cpu') + torch.cuda.empty_cache() + + # Convert pixel videos to clean latents via the VAE. + pixel_values = batch["pixel_values"].to(weight_dtype) + if args.low_vram: + torch.cuda.empty_cache() + vae.to(accelerator.device) + + with torch.no_grad(): + def _batch_encode_vae(pixel_values): + pixel_values = rearrange(pixel_values, "b f c h w -> b c f h w") + bs = args.vae_mini_batch + new_pixel_values = [] + for i in range(0, pixel_values.shape[0], bs): + pixel_values_bs = pixel_values[i : i + bs] + pixel_values_bs = vae.encode(pixel_values_bs)[0] + pixel_values_bs = pixel_values_bs.sample() + new_pixel_values.append(pixel_values_bs) + return torch.cat(new_pixel_values, dim=0) + clean_latents = _batch_encode_vae(pixel_values) + + if args.low_vram: + vae.to('cpu') + torch.cuda.empty_cache() + + with accelerator.accumulate(transformer3d): + def get_sigmas(timesteps, n_dim=4, dtype=torch.float32): + sigmas = noise_scheduler.sigmas.to(device=accelerator.device, dtype=dtype) + sched_t = schedule_timesteps.to(accelerator.device) + timesteps = timesteps.to(accelerator.device) + step_indices = torch.argmin((sched_t.unsqueeze(0) - timesteps.reshape(-1).unsqueeze(1)).abs(), dim=1) + sigma = sigmas[step_indices] + while len(sigma.shape) < n_dim: + sigma = sigma.unsqueeze(-1) + return sigma + + def add_noise(latents, noise, timesteps): + """Per-frame flow-matching add_noise; supports timesteps of shape [B, F].""" + sigmas = get_sigmas(timesteps, n_dim=2, dtype=latents.dtype) + sigmas = sigmas.reshape(latents.shape[0], latents.shape[2], 1, 1, 1).permute(0, 2, 1, 3, 4) + return (1.0 - sigmas) * latents + sigmas * noise + + def _flow_to_x0(flow_pred, xt, timesteps): + """Recover x0 from flow-matching velocity: x0 = xt - sigma_t * flow_pred.""" + sigmas = get_sigmas(timesteps, n_dim=2, dtype=flow_pred.dtype) + sigmas = sigmas.reshape(flow_pred.shape[0], flow_pred.shape[2], 1, 1, 1).permute(0, 2, 1, 3, 4) + return xt - sigmas * flow_pred + + bsz, channel, num_frames, height, width = clean_latents.shape + + # Sample one CCD timestep index per batch item from the + # discrete schedule [0, N - 2] and share it across all causal + # frames — mirrors `NaiveConsistency.generator_loss`. + timestep_idx = torch.randint( + 0, args.discrete_cd_N - 1, [bsz], + device=accelerator.device, generator=torch_rng, + ) + t_scalar = cd_timesteps.to(accelerator.device)[timestep_idx].float() # [B] + t_next_scalar = cd_timesteps.to(accelerator.device)[timestep_idx + 1].float() # [B] + timestep = t_scalar.unsqueeze(1).expand(bsz, num_frames).contiguous() + timestep_next = t_next_scalar.unsqueeze(1).expand(bsz, num_frames).contiguous() + + noise = torch.randn(clean_latents.shape, dtype=weight_dtype, device=accelerator.device, generator=torch_rng) + noisy_latents = add_noise(clean_latents, noise, timestep) + + # Build the per-item lists once: every forward in this step is + # teacher-forced on the clean latent, matching Stage 2 of + # Causal-Forcing++. + noisy_input_list = [noisy_latents[i] for i in range(bsz)] + clean_x_list = [clean_latents[i] for i in range(bsz)] + patch_h, patch_w = accelerator.unwrap_model(transformer3d).config.patch_size[1:] + full_seq_len = num_frames * height * width // (patch_h * patch_w) + + # --- Teacher CFG forward + one ODE step to latent_t_next --- + with torch.no_grad(): + uncond_prompt_embeds = [negative_prompt_embed for _ in range(bsz)] + with torch.cuda.amp.autocast(dtype=weight_dtype): + flow_cond = teacher_transformer3d( + x=noisy_input_list, context=prompt_embeds, + t=timestep.to(torch.int64), seq_len=full_seq_len, + clean_x=clean_x_list, + ) + flow_uncond = teacher_transformer3d( + x=noisy_input_list, context=uncond_prompt_embeds, + t=timestep.to(torch.int64), seq_len=full_seq_len, + clean_x=clean_x_list, + ) + if isinstance(flow_cond, list): + flow_cond = torch.stack(flow_cond, dim=0) + if isinstance(flow_uncond, list): + flow_uncond = torch.stack(flow_uncond, dim=0) + flow_teacher = flow_uncond + args.guidance_scale * (flow_cond - flow_uncond) + + # Per-frame Euler step on the flow-matching ODE. + # `dt = (t - t_next) / num_train_timestep`; divisor matches the + # `/1000` in `NaiveConsistency.generator_loss` (the CF + # FlowMatchScheduler uses `num_train_timesteps=1000`). + dt = ((timestep - timestep_next) / 1000.0).reshape(bsz, num_frames, 1, 1, 1).permute(0, 2, 1, 3, 4) + latent_t_next = noisy_latents - dt.to(noisy_latents.dtype) * flow_teacher.to(noisy_latents.dtype) + + # --- Generator at (latent_t, t) -> x0_pred_t (with grad) --- + with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + flow_g = transformer3d( + x=noisy_input_list, context=prompt_embeds, + t=timestep.to(torch.int64), seq_len=full_seq_len, + clean_x=clean_x_list, + ) + if isinstance(flow_g, list): + flow_g = torch.stack(flow_g, dim=0) + x0_pred_t = _flow_to_x0(flow_g, noisy_latents, timestep) + + # --- EMA-target at (latent_t_next, t_next) -> x0_pred_t_next (no grad) --- + with torch.no_grad(): + ema_model = ema_transformer3d if ema_transformer3d is not None else accelerator.unwrap_model(transformer3d) + noisy_next_list = [latent_t_next[i] for i in range(bsz)] + with torch.cuda.amp.autocast(dtype=weight_dtype): + flow_ema = ema_model( + x=noisy_next_list, context=prompt_embeds, + t=timestep_next.to(torch.int64), seq_len=full_seq_len, + clean_x=clean_x_list, + ) + if isinstance(flow_ema, list): + flow_ema = torch.stack(flow_ema, dim=0) + x0_pred_t_next = _flow_to_x0(flow_ema, latent_t_next, timestep_next).detach() + + # CCD objective: MSE between the two x0 predictions. + loss = F.mse_loss(x0_pred_t.float(), x0_pred_t_next.float()) + + avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean() + train_loss += avg_loss.item() / args.gradient_accumulation_steps + + accelerator.backward(loss) + if accelerator.sync_gradients: + accelerator.clip_grad_norm_(trainable_params, args.max_grad_norm) + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + + # EMA update for the consistency target. Before `--ema_start_step` + # we keep the EMA copy in lock-step with the live generator, so + # that early-training instability does not bias the right-hand + # side of the CCD objective. + if ema_transformer3d is not None and accelerator.sync_gradients: + live_state = accelerator.unwrap_model(transformer3d).state_dict() + if global_step < args.ema_start_step: + ema_transformer3d.load_state_dict(live_state, strict=True) + else: + decay = args.ema_weight + ema_state = ema_transformer3d.state_dict() + with torch.no_grad(): + for k, v in ema_state.items(): + if v.dtype.is_floating_point: + v.mul_(decay).add_(live_state[k].to(v.dtype, copy=False), alpha=1.0 - decay) + else: + v.copy_(live_state[k]) + + # Checks if the accelerator has performed an optimization step behind the scenes + if accelerator.sync_gradients: + + progress_bar.update(1) + global_step += 1 + accelerator.log({"train_loss": train_loss}, step=global_step) + train_loss = 0.0 + + if global_step % args.checkpointing_steps == 0: + if args.use_deepspeed or args.use_fsdp or accelerator.is_main_process: + # _before_ saving state, check if this save would set us over the `checkpoints_total_limit` + if args.checkpoints_total_limit is not None: + checkpoints = os.listdir(args.output_dir) + checkpoints = [d for d in checkpoints if d.startswith("checkpoint")] + checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1])) + + # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints + if len(checkpoints) >= args.checkpoints_total_limit: + num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1 + removing_checkpoints = checkpoints[0:num_to_remove] + + logger.info( + f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints" + ) + logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}") + + for removing_checkpoint in removing_checkpoints: + removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint) + shutil.rmtree(removing_checkpoint) + + gc.collect() + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") + accelerator.save_state(save_path) + logger.info(f"Saved state to {save_path}") + + if args.validation_prompts is not None and global_step % args.validation_steps == 0: + log_validation( + vae, + text_encoder, + tokenizer, + transformer3d, + args, + config, + accelerator, + weight_dtype, + global_step, + ) + + logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]} + progress_bar.set_postfix(**logs) + + if global_step >= args.max_train_steps: + break + + if args.validation_prompts is not None and epoch % args.validation_epochs == 0: + log_validation( + vae, + text_encoder, + tokenizer, + transformer3d, + args, + config, + accelerator, + weight_dtype, + global_step, + ) + + # Create the pipeline using the trained modules and save it. + accelerator.wait_for_everyone() + if accelerator.is_main_process: + transformer3d = unwrap_model(transformer3d) + + if args.use_deepspeed or args.use_fsdp or accelerator.is_main_process: + gc.collect() + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") + accelerator.save_state(save_path) + logger.info(f"Saved state to {save_path}") + + accelerator.end_training() + + +if __name__ == "__main__": + main() diff --git a/scripts/wan2.1_causal_forcing/train_causal_consistency_distill.sh b/scripts/wan2.1_causal_forcing/train_causal_consistency_distill.sh new file mode 100644 index 00000000..b7cfd39f --- /dev/null +++ b/scripts/wan2.1_causal_forcing/train_causal_consistency_distill.sh @@ -0,0 +1,54 @@ +export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/" +export DATASET_NAME="datasets/internal_datasets/" +export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export STAGE1_CKPT="output_dir_wan2.1_causal_forcing_ar_diffusion/checkpoint-8000/diffusion_pytorch_model.safetensors" +# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. +# export NCCL_IB_DISABLE=1 +# export NCCL_P2P_DISABLE=1 +NCCL_DEBUG=INFO + +# Causal-Forcing Stage 2: Causal Consistency Distillation (CCD). +# - --transformer_path / --teacher_transformer_path point to the Stage 1 AR-diffusion ckpt. +# - num_frame_per_block=3 -> chunkwise; set to 1 for the framewise variant. +# - discrete_cd_N=48 mirrors the official `causal_cd_chunkwise.yaml`. +accelerate launch --mixed_precision="bf16" scripts/wan2.1_causal_forcing/train_causal_consistency_distill.py \ + --config_path="config/wan2.1/wan_civitai.yaml" \ + --pretrained_model_name_or_path=$MODEL_NAME \ + --train_data_dir=$DATASET_NAME \ + --train_data_meta=$DATASET_META_NAME \ + --transformer_path=$STAGE1_CKPT \ + --teacher_transformer_path=$STAGE1_CKPT \ + --image_sample_size=640 \ + --video_sample_size=640 \ + --token_sample_size=640 \ + --fix_sample_size 480 832 \ + --video_sample_stride=2 \ + --video_sample_n_frames=81 \ + --train_batch_size=1 \ + --gradient_accumulation_steps=1 \ + --dataloader_num_workers=8 \ + --num_train_epochs=100 \ + --checkpointing_steps=200 \ + --learning_rate=2.0e-06 \ + --lr_scheduler="constant_with_warmup" \ + --lr_warmup_steps=100 \ + --seed=42 \ + --output_dir="output_dir_wan2.1_causal_forcing_ccd" \ + --gradient_checkpointing \ + --mixed_precision="bf16" \ + --adam_weight_decay=0.0 \ + --adam_beta1=0.0 \ + --adam_beta2=0.999 \ + --adam_epsilon=1e-10 \ + --vae_mini_batch=1 \ + --max_grad_norm=10.0 \ + --random_hw_adapt \ + --training_with_video_token_length \ + --enable_bucket \ + --num_frame_per_block=3 \ + --shift=5.0 \ + --discrete_cd_N=48 \ + --guidance_scale=3.0 \ + --ema_weight=0.99 \ + --ema_start_step=200 \ + --trainable_modules "." \ No newline at end of file diff --git a/scripts/wan2.1_causal_forcing/train_causal_dmd.py b/scripts/wan2.1_causal_forcing/train_causal_dmd.py new file mode 100644 index 00000000..752d3369 --- /dev/null +++ b/scripts/wan2.1_causal_forcing/train_causal_dmd.py @@ -0,0 +1,3063 @@ +"""Modified from https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image.py +""" +#!/usr/bin/env python +# coding=utf-8 +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and + +import argparse +import contextlib +import gc +import json +import logging +import math +import os +import pickle +import shutil +import sys + +import accelerate +import diffusers +import numpy as np +import torch +import torch.distributed as dist +import torch.nn.functional as F +import torch.utils.checkpoint +import torchvision.transforms.functional as TF +import transformers +from accelerate import Accelerator, FullyShardedDataParallelPlugin +from accelerate.logging import get_logger +from accelerate.state import AcceleratorState +from accelerate.utils import ProjectConfiguration, set_seed +from diffusers import DDIMScheduler, FlowMatchEulerDiscreteScheduler +from diffusers.optimization import get_scheduler +from diffusers.training_utils import (EMAModel, + compute_density_for_timestep_sampling, + compute_loss_weighting_for_sd3) +from diffusers.utils import check_min_version, is_wandb_available +from diffusers.utils.torch_utils import is_compiled_module +from einops import rearrange +from omegaconf import OmegaConf +from packaging import version +from PIL import Image +from torch.distributed.fsdp.fully_sharded_data_parallel import ( + FullOptimStateDictConfig, FullStateDictConfig, ShardedOptimStateDictConfig, + ShardedStateDictConfig) +from torch.utils.data import BatchSampler, Dataset, RandomSampler +from torch.utils.tensorboard import SummaryWriter +from torchvision import transforms +from tqdm.auto import tqdm +from transformers import AutoTokenizer +from transformers.utils import ContextManagers + +import datasets + +current_file_path = os.path.abspath(__file__) +project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dirname(current_file_path)), os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))] +for project_root in project_roots: + sys.path.insert(0, project_root) if project_root not in sys.path else None + +from videox_fun.data import (ASPECT_RATIO_512, ASPECT_RATIO_RANDOM_CROP_512, + ASPECT_RATIO_RANDOM_CROP_PROB, + AspectRatioBatchImageVideoSampler, + ImageVideoDataset, ImageVideoSampler, + RandomSampler, TextDataset, get_closest_ratio, + get_random_mask) +from videox_fun.models import (AutoencoderKLWan, CLIPModel, WanT5EncoderModel, + WanTransformer3DModel, + WanTransformer3DModel_SelfForcing) +from videox_fun.pipeline import (WanI2VPipeline, WanPipeline, + WanSelfForcingPipeline) +from videox_fun.utils.discrete_sampler import DiscreteSampling +from videox_fun.utils.utils import (calculate_dimensions, get_image_latent, + get_image_to_video_latent, + save_videos_grid) + +if is_wandb_available(): + import wandb + + +def initialize_kv_cache_for_training(batch_size, num_frames, frame_seq_length, num_layers, num_heads, head_dim, dtype, device): + """Initialize KV cache for block-by-block training""" + kv_cache_size = num_frames * frame_seq_length + kv_cache = [] + + for _ in range(num_layers): + kv_cache.append({ + "k": torch.zeros([batch_size, kv_cache_size, num_heads, head_dim], dtype=dtype, device=device), + "v": torch.zeros([batch_size, kv_cache_size, num_heads, head_dim], dtype=dtype, device=device), + "global_end_index": torch.tensor([0], dtype=torch.long, device=device), + "local_end_index": torch.tensor([0], dtype=torch.long, device=device) + }) + + return kv_cache + + +def initialize_crossattn_cache_for_training(batch_size, text_len, num_layers, num_heads, head_dim, dtype, device): + """Initialize cross-attention cache for block-by-block training""" + crossattn_cache = [] + + for _ in range(num_layers): + crossattn_cache.append({ + "k": torch.zeros([batch_size, text_len, num_heads, head_dim], dtype=dtype, device=device), + "v": torch.zeros([batch_size, text_len, num_heads, head_dim], dtype=dtype, device=device), + "is_init": False + }) + + return crossattn_cache + + +def filter_kwargs(cls, kwargs): + import inspect + sig = inspect.signature(cls.__init__) + valid_params = set(sig.parameters.keys()) - {'self', 'cls'} + filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_params} + return filtered_kwargs + +def get_random_downsample_ratio(sample_size, image_ratio=[], + all_choices=False, rng=None): + def _create_special_list(length): + if length == 1: + return [1.0] + if length >= 2: + first_element = 0.75 + remaining_sum = 1.0 - first_element + other_elements_value = remaining_sum / (length - 1) + special_list = [first_element] + [other_elements_value] * (length - 1) + return special_list + + if sample_size >= 1536: + number_list = [1, 1.25, 1.5, 2, 2.5, 3] + image_ratio + elif sample_size >= 1024: + number_list = [1, 1.25, 1.5, 2] + image_ratio + elif sample_size >= 768: + number_list = [1, 1.25, 1.5] + image_ratio + elif sample_size >= 512: + number_list = [1] + image_ratio + else: + number_list = [1] + + if all_choices: + return number_list + + number_list_prob = np.array(_create_special_list(len(number_list))) + if rng is None: + return np.random.choice(number_list, p = number_list_prob) + else: + return rng.choice(number_list, p = number_list_prob) + +def resize_mask(mask, latent, process_first_frame_only=True): + latent_size = latent.size() + batch_size, channels, num_frames, height, width = mask.shape + + if process_first_frame_only: + target_size = list(latent_size[2:]) + target_size[0] = 1 + first_frame_resized = F.interpolate( + mask[:, :, 0:1, :, :], + size=target_size, + mode='trilinear', + align_corners=False + ) + + target_size = list(latent_size[2:]) + target_size[0] = target_size[0] - 1 + if target_size[0] != 0: + remaining_frames_resized = F.interpolate( + mask[:, :, 1:, :, :], + size=target_size, + mode='trilinear', + align_corners=False + ) + resized_mask = torch.cat([first_frame_resized, remaining_frames_resized], dim=2) + else: + resized_mask = first_frame_resized + else: + target_size = list(latent_size[2:]) + resized_mask = F.interpolate( + mask, + size=target_size, + mode='trilinear', + align_corners=False + ) + return resized_mask + +# Will error if the minimal version of diffusers is not installed. Remove at your own risks. +check_min_version("0.18.0.dev0") + +logger = get_logger(__name__, log_level="INFO") + +def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, transformer3d, args, config, accelerator, weight_dtype, global_step): + try: + is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine' + if is_deepspeed: + origin_config = transformer3d.config + transformer3d.config = accelerator.unwrap_model(transformer3d).config + with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + logger.info("Running validation... ") + scheduler = FlowMatchEulerDiscreteScheduler( + **filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs'])) + ) + + if args.train_mode != "normal": + raise NotImplementedError(f"Validation for train_mode '{args.train_mode}' is not yet supported with WanSelfForcingPipeline. Only T2V (train_mode='normal') is currently supported.") + else: + pipeline = WanSelfForcingPipeline( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d, + scheduler=scheduler, + ) + pipeline = pipeline.to(accelerator.device) + + if args.seed is None: + generator = None + else: + rank_seed = args.seed + accelerator.process_index + generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed) + logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}") + + for i in range(len(args.validation_prompts)): + if args.train_mode != "normal": + raise NotImplementedError(f"Validation for train_mode '{args.train_mode}' is not yet supported with WanSelfForcingPipeline. Only T2V (train_mode='normal') is currently supported.") + else: + if args.fix_sample_size is not None: + height, width = args.fix_sample_size + else: + height, width = args.video_sample_size, args.video_sample_size + sample = pipeline( + args.validation_prompts[i], + num_frames = args.video_sample_n_frames, + negative_prompt = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走", + height = height, + width = width, + generator = generator, + guidance_scale = 1.0, + num_inference_steps = len(args.denoising_step_indices_list), + num_frame_per_block = args.num_frame_per_block, + independent_first_frame = args.independent_first_frame, + context_noise = args.context_noise, + ).videos + os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True) + save_videos_grid( + sample, + os.path.join( + args.output_dir, + f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.mp4" + ) + ) + + del pipeline + gc.collect() + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype) + if not args.enable_text_encoder_in_dataloader: + text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype) + if is_deepspeed: + transformer3d.config = origin_config + except Exception as e: + gc.collect() + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + print(f"Eval error on rank {accelerator.process_index} with info {e}") + vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype) + if not args.enable_text_encoder_in_dataloader: + text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype) + +def linear_decay(initial_value, final_value, total_steps, current_step): + if current_step >= total_steps: + return final_value + current_step = max(0, current_step) + step_size = (final_value - initial_value) / total_steps + current_value = initial_value + step_size * current_step + return current_value + +def generate_timestep_with_lognorm(low, high, shape, device="cpu", generator=None): + u = torch.normal(mean=0.0, std=1.0, size=shape, device=device, generator=generator) + t = 1 / (1 + torch.exp(-u)) * (high - low) + low + return torch.clip(t.to(torch.int32), low, high - 1) + +def parse_args(): + parser = argparse.ArgumentParser(description="Simple example of a training script.") + parser.add_argument( + "--input_perturbation", type=float, default=0, help="The scale of input perturbation. Recommended 0.1." + ) + parser.add_argument( + "--pretrained_model_name_or_path", + type=str, + default=None, + required=True, + help="Path to pretrained model or model identifier from huggingface.co/models.", + ) + parser.add_argument( + "--revision", + type=str, + default=None, + required=False, + help="Revision of pretrained model identifier from huggingface.co/models.", + ) + parser.add_argument( + "--variant", + type=str, + default=None, + help="Variant of the model files of the pretrained model identifier from huggingface.co/models, 'e.g.' fp16", + ) + parser.add_argument( + "--train_data_dir", + type=str, + default=None, + help=( + "A folder containing the training data. " + ), + ) + parser.add_argument( + "--train_data_meta", + type=str, + default=None, + help=( + "A csv containing the training data. " + ), + ) + parser.add_argument( + "--max_train_samples", + type=int, + default=None, + help=( + "For debugging purposes or quicker training, truncate the number of training examples to this " + "value if set." + ), + ) + parser.add_argument( + "--validation_prompts", + type=str, + default=None, + nargs="+", + help=("A set of prompts evaluated every `--validation_epochs` and logged to `--report_to`."), + ) + parser.add_argument( + "--validation_paths", + type=str, + default=None, + nargs="+", + help=("A set of control videos evaluated every `--validation_epochs` and logged to `--report_to`."), + ) + parser.add_argument( + "--negative_prompt", + type=str, + default="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走", + help=("The negative prompt of cfg distill"), + ) + parser.add_argument( + "--output_dir", + type=str, + default="sd-model-finetuned", + help="The output directory where the model predictions and checkpoints will be written.", + ) + parser.add_argument( + "--cache_dir", + type=str, + default=None, + help="The directory where the downloaded models and datasets will be stored.", + ) + parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.") + parser.add_argument( + "--random_flip", + action="store_true", + help="whether to randomly flip images horizontally", + ) + parser.add_argument( + "--use_came", + action="store_true", + help="whether to use came", + ) + parser.add_argument( + "--multi_stream", + action="store_true", + help="whether to use cuda multi-stream", + ) + parser.add_argument( + "--train_batch_size", type=int, default=16, help="Batch size (per device) for the training dataloader." + ) + parser.add_argument( + "--vae_mini_batch", type=int, default=32, help="mini batch size for vae." + ) + parser.add_argument("--num_train_epochs", type=int, default=100) + parser.add_argument( + "--max_train_steps", + type=int, + default=None, + help="Total number of training steps to perform. If provided, overrides num_train_epochs.", + ) + parser.add_argument( + "--gradient_accumulation_steps", + type=int, + default=1, + help="Number of updates steps to accumulate before performing a backward/update pass.", + ) + parser.add_argument( + "--gradient_checkpointing", + action="store_true", + help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.", + ) + parser.add_argument( + "--learning_rate", + type=float, + default=1e-4, + help="Initial learning rate (after the potential warmup period) to use.", + ) + parser.add_argument( + "--learning_rate_critic", + type=float, + default=1e-4, + help="Initial learning rate (after the potential warmup period) to use.", + ) + parser.add_argument( + "--scale_lr", + action="store_true", + default=False, + help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.", + ) + parser.add_argument( + "--lr_scheduler", + type=str, + default="constant", + help=( + 'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",' + ' "constant", "constant_with_warmup"]' + ), + ) + parser.add_argument( + "--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler." + ) + parser.add_argument( + "--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes." + ) + parser.add_argument( + "--allow_tf32", + action="store_true", + help=( + "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see" + " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices" + ), + ) + parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model for the generator.") + parser.add_argument( + "--dataloader_num_workers", + type=int, + default=0, + help=( + "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process." + ), + ) + parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.") + parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.") + parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.") + parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer") + parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.") + parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.") + parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.") + parser.add_argument( + "--prediction_type", + type=str, + default=None, + help="The prediction_type that shall be used for training. Choose between 'epsilon' or 'v_prediction' or leave `None`. If left to `None` the default prediction type of the scheduler: `noise_scheduler.config.prediciton_type` is chosen.", + ) + parser.add_argument( + "--hub_model_id", + type=str, + default=None, + help="The name of the repository to keep in sync with the local `output_dir`.", + ) + parser.add_argument( + "--logging_dir", + type=str, + default="logs", + help=( + "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to" + " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***." + ), + ) + parser.add_argument( + "--report_model_info", action="store_true", help="Whether or not to report more info about model (such as norm, grad)." + ) + parser.add_argument( + "--mixed_precision", + type=str, + default=None, + choices=["no", "fp16", "bf16"], + help=( + "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >=" + " 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the" + " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config." + ), + ) + parser.add_argument( + "--report_to", + type=str, + default="tensorboard", + help=( + 'The integration to report the results and logs to. Supported platforms are `"tensorboard"`' + ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.' + ), + ) + parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank") + parser.add_argument( + "--checkpointing_steps", + type=int, + default=500, + help=( + "Save a checkpoint of the training state every X updates. These checkpoints are only suitable for resuming" + " training using `--resume_from_checkpoint`." + ), + ) + parser.add_argument( + "--checkpoints_total_limit", + type=int, + default=None, + help=("Max number of checkpoints to store."), + ) + parser.add_argument( + "--resume_from_checkpoint", + type=str, + default=None, + help=( + "Whether training should be resumed from a previous checkpoint. Use a path saved by" + ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.' + ), + ) + parser.add_argument("--noise_offset", type=float, default=0, help="The scale of noise offset.") + parser.add_argument( + "--validation_epochs", + type=int, + default=5, + help="Run validation every X epochs.", + ) + parser.add_argument( + "--validation_steps", + type=int, + default=2000, + help="Run validation every X steps.", + ) + parser.add_argument( + "--tracker_project_name", + type=str, + default="text2image-fine-tune", + help=( + "The `project_name` argument passed to Accelerator.init_trackers for" + " more information see https://huggingface.co/docs/accelerate/v0.17.0/en/package_reference/accelerator#accelerate.Accelerator" + ), + ) + + parser.add_argument( + "--snr_loss", action="store_true", help="Whether or not to use snr_loss." + ) + parser.add_argument( + "--uniform_sampling", action="store_true", help="Whether or not to use uniform_sampling." + ) + parser.add_argument( + "--enable_text_encoder_in_dataloader", action="store_true", help="Whether or not to use text encoder in dataloader." + ) + parser.add_argument( + "--enable_bucket", action="store_true", help="Whether enable bucket sample in datasets." + ) + parser.add_argument( + "--random_ratio_crop", action="store_true", help="Whether enable random ratio crop sample in datasets." + ) + parser.add_argument( + "--random_frame_crop", action="store_true", help="Whether enable random frame crop sample in datasets." + ) + parser.add_argument( + "--random_hw_adapt", action="store_true", help="Whether enable random adapt height and width in datasets." + ) + parser.add_argument( + "--training_with_video_token_length", action="store_true", help="The training stage of the model in training.", + ) + parser.add_argument( + "--auto_tile_batch_size", action="store_true", help="Whether to auto tile batch size.", + ) + parser.add_argument( + "--motion_sub_loss", action="store_true", help="Whether enable motion sub loss." + ) + parser.add_argument( + "--motion_sub_loss_ratio", type=float, default=0.25, help="The ratio of motion sub loss." + ) + parser.add_argument( + "--train_sampling_steps", + type=int, + default=1000, + help="Run train_sampling_steps.", + ) + parser.add_argument( + "--keep_all_node_same_token_length", + action="store_true", + help="Reference of the length token.", + ) + parser.add_argument( + "--token_sample_size", + type=int, + default=512, + help="Sample size of the token.", + ) + parser.add_argument( + "--video_sample_size", + type=int, + default=512, + help="Sample size of the video.", + ) + parser.add_argument( + "--image_sample_size", + type=int, + default=512, + help="Sample size of the image.", + ) + parser.add_argument( + "--fix_sample_size", + nargs=2, type=int, default=None, + help="Fix Sample size [height, width] when using bucket and collate_fn." + ) + parser.add_argument( + "--video_sample_stride", + type=int, + default=4, + help="Sample stride of the video.", + ) + parser.add_argument( + "--video_sample_n_frames", + type=int, + default=17, + help="Num frame of video.", + ) + parser.add_argument( + "--video_repeat", + type=int, + default=0, + help="Num of repeat video.", + ) + parser.add_argument( + "--config_path", + type=str, + default=None, + help=( + "The config of the model in training." + ), + ) + parser.add_argument( + "--transformer_path", + type=str, + default=None, + help=("If you want to load the weight from other transformers, input its path."), + ) + parser.add_argument( + "--ode_transformer_path", + type=str, + default=None, + help=("If you want to load the ode-trained weight into generator transformer3d, input its path."), + ) + parser.add_argument( + "--real_score_pretrained_model_name_or_path", + type=str, + default=None, + help=( + "Path to a non-causal pretrained model used as the DMD real_score teacher. " + "For CF Stage 3 this should point at Wan2.1-T2V-14B. If unset, falls back to " + "`--pretrained_model_name_or_path` (1.3B) — useful only for sanity testing." + ), + ) + parser.add_argument( + "--vae_path", + type=str, + default=None, + help=("If you want to load the weight from other vaes, input its path."), + ) + + parser.add_argument( + '--trainable_modules', + nargs='+', + help='Enter a list of trainable modules' + ) + parser.add_argument( + '--trainable_modules_low_learning_rate', + nargs='+', + default=[], + help='Enter a list of trainable modules with lower learning rate' + ) + parser.add_argument( + '--tokenizer_max_length', + type=int, + default=512, + help='Max length of tokenizer' + ) + parser.add_argument( + "--use_deepspeed", action="store_true", help="Whether or not to use deepspeed." + ) + parser.add_argument( + "--use_fsdp", action="store_true", help="Whether or not to use fsdp." + ) + parser.add_argument( + "--low_vram", action="store_true", help="Whether enable low_vram mode." + ) + parser.add_argument( + "--train_mode", + type=str, + default="normal", + help=( + 'The format of training data. Support `"normal"`' + ' (default), `"i2v"`.' + ), + ) + parser.add_argument( + "--gen_update_interval", + type=int, + default=5, + help="The ratio to update transformer3d.", + ) + parser.add_argument( + "--fake_guidance_scale", + type=float, + default=0.0, + help="The cfg scale for fake iscore.", + ) + parser.add_argument( + "--real_guidance_scale", + type=float, + default=4.5, + help="The cfg scale for real score.", + ) + parser.add_argument( + '--denoising_step_indices_list', + nargs='+', + type=int, + default=[1000, 500], + help="Denoising step indices (in train_sampling_steps space). CF Stage 3 DMD default = 2-step [1000, 500].", + ) + parser.add_argument( + "--num_frame_per_block", + type=int, + default=3, + help="Number of frames per block for Self-Forcing causal training" + ) + parser.add_argument( + "--independent_first_frame", + action="store_true", + help="Whether first frame is independent ([1, N, N, ...] pattern)" + ) + parser.add_argument( + "--use_kv_cache_training", + action="store_true", + help="Use KV cache block-by-block training (matches original Self-Forcing)" + ) + parser.add_argument( + "--score_num_frames", + type=int, + default=21, + help="Number of latent frames for score computation window (default: 21, matching base model). " + "fake_score/real_score always receive this many frames." + ) + parser.add_argument( + "--min_length_prob_bias", + type=float, + default=0.0, + help="Probability bias for sampling the minimum length (score_num_frames). " + "0.0 = uniform sampling (default), 0.5 = 50%% prob for min length, " + "remaining prob distributed equally among longer lengths. " + "Use this to increase 21-frame training ratio." + ) + parser.add_argument( + "--context_noise", + type=int, + default=0, + help="Context noise level for KV cache update (matches training config)" + ) + parser.add_argument( + "--use_teacher_forcing", + action="store_true", + help="Enable teacher forcing training (pass clean_x to transformer)" + ) + parser.add_argument( + "--teacher_forcing_prob", + type=float, + default=1.0, + help="Probability of applying teacher forcing per step (1.0 = always)" + ) + + args = parser.parse_args() + env_local_rank = int(os.environ.get("LOCAL_RANK", -1)) + if env_local_rank != -1 and env_local_rank != args.local_rank: + args.local_rank = env_local_rank + + return args + + +def main(): + args = parse_args() + + if args.report_to == "wandb" and args.hub_token is not None: + raise ValueError( + "You cannot use both --report_to=wandb and --hub_token due to a security risk of exposing your token." + " Please use `huggingface-cli login` to authenticate with the Hub." + ) + + logging_dir = os.path.join(args.output_dir, args.logging_dir) + + config = OmegaConf.load(args.config_path) + accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir) + + accelerator = Accelerator( + gradient_accumulation_steps=args.gradient_accumulation_steps, + mixed_precision=args.mixed_precision, + log_with=args.report_to, + project_config=accelerator_project_config, + ) + accelerator_fake_score_transformer3d = Accelerator( + gradient_accumulation_steps=args.gradient_accumulation_steps, + mixed_precision=args.mixed_precision, + log_with=args.report_to, + project_config=accelerator_project_config, + ) + + deepspeed_plugin = accelerator.state.deepspeed_plugin if hasattr(accelerator.state, "deepspeed_plugin") else None + fsdp_plugin = accelerator.state.fsdp_plugin if hasattr(accelerator.state, "fsdp_plugin") else None + if deepspeed_plugin is not None: + zero_stage = int(deepspeed_plugin.zero_stage) + fsdp_stage = 0 + print(f"Using DeepSpeed Zero stage: {zero_stage}") + + args.use_deepspeed = True + if zero_stage == 3: + print(f"Auto set save_state to True because zero_stage == 3") + args.save_state = True + elif fsdp_plugin is not None: + from torch.distributed.fsdp import ShardingStrategy + zero_stage = 0 + if fsdp_plugin.sharding_strategy is ShardingStrategy.FULL_SHARD: + fsdp_stage = 3 + elif fsdp_plugin.sharding_strategy is None: # The fsdp_plugin.sharding_strategy is None in FSDP 2. + fsdp_stage = 3 + elif fsdp_plugin.sharding_strategy is ShardingStrategy.SHARD_GRAD_OP: + fsdp_stage = 2 + else: + fsdp_stage = 0 + print(f"Using FSDP stage: {fsdp_stage}") + + args.use_fsdp = True + if fsdp_stage == 3: + print(f"Auto set save_state to True because fsdp_stage == 3") + args.save_state = True + else: + zero_stage = 0 + fsdp_stage = 0 + print("DeepSpeed is not enabled.") + + if accelerator.is_main_process: + writer = SummaryWriter(log_dir=logging_dir) + + # Make one log on every process with the configuration for debugging. + logging.basicConfig( + format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", + datefmt="%m/%d/%Y %H:%M:%S", + level=logging.INFO, + ) + logger.info(accelerator.state, main_process_only=False) + if accelerator.is_local_main_process: + datasets.utils.logging.set_verbosity_warning() + transformers.utils.logging.set_verbosity_warning() + diffusers.utils.logging.set_verbosity_info() + else: + datasets.utils.logging.set_verbosity_error() + transformers.utils.logging.set_verbosity_error() + diffusers.utils.logging.set_verbosity_error() + + # If passed along, set the training seed now. + if args.seed is not None: + set_seed(args.seed) + rng = np.random.default_rng(np.random.PCG64(args.seed + accelerator.process_index)) + torch_rng = torch.Generator(accelerator.device).manual_seed(args.seed + accelerator.process_index) + else: + rng = None + torch_rng = None + index_rng = np.random.default_rng(np.random.PCG64(43)) + print(f"Init rng with seed {args.seed + accelerator.process_index}. Process_index is {accelerator.process_index}") + + # Handle the repository creation + if accelerator.is_main_process: + if args.output_dir is not None: + os.makedirs(args.output_dir, exist_ok=True) + + # For mixed precision training we cast all non-trainable weigths (vae, non-lora text_encoder and non-lora transformer3d) to half-precision + # as these weights are only used for inference, keeping weights in full precision is not required. + weight_dtype = torch.float32 + if accelerator.mixed_precision == "fp16": + weight_dtype = torch.float16 + args.mixed_precision = accelerator.mixed_precision + elif accelerator.mixed_precision == "bf16": + weight_dtype = torch.bfloat16 + args.mixed_precision = accelerator.mixed_precision + + # Load scheduler, tokenizer and models. + noise_scheduler = FlowMatchEulerDiscreteScheduler( + **filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs'])) + ) + + # Get Tokenizer + tokenizer = AutoTokenizer.from_pretrained( + os.path.join(args.pretrained_model_name_or_path, config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer')), + ) + + def deepspeed_zero_init_disabled_context_manager(): + """ + returns either a context list that includes one that will disable zero.Init or an empty context list + """ + deepspeed_plugin = AcceleratorState().deepspeed_plugin if accelerate.state.is_initialized() else None + if deepspeed_plugin is None: + return [] + + return [deepspeed_plugin.zero3_init_context_manager(enable=False)] + + # Currently Accelerate doesn't know how to handle multiple models under Deepspeed ZeRO stage 3. + # For this to work properly all models must be run through `accelerate.prepare`. But accelerate + # will try to assign the same optimizer with the same weights to all models during + # `deepspeed.initialize`, which of course doesn't work. + # + # For now the following workaround will partially support Deepspeed ZeRO-3, by excluding the 2 + # frozen models from being partitioned during `zero.Init` which gets called during + # `from_pretrained` So CLIPTextModel and AutoencoderKL will not enjoy the parameter sharding + # across multiple gpus and only UNet2DConditionModel will get ZeRO sharded. + with ContextManagers(deepspeed_zero_init_disabled_context_manager()): + # Get Text encoder + text_encoder = WanT5EncoderModel.from_pretrained( + os.path.join(args.pretrained_model_name_or_path, config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')), + additional_kwargs=OmegaConf.to_container(config['text_encoder_kwargs']), + low_cpu_mem_usage=True, + torch_dtype=weight_dtype, + ) + text_encoder = text_encoder.eval() + # Get Vae + vae = AutoencoderKLWan.from_pretrained( + os.path.join(args.pretrained_model_name_or_path, config['vae_kwargs'].get('vae_subpath', 'vae')), + additional_kwargs=OmegaConf.to_container(config['vae_kwargs']), + ) + vae.eval() + # Get Clip Image Encoder + if args.train_mode != "normal": + clip_image_encoder = CLIPModel.from_pretrained( + os.path.join(args.pretrained_model_name_or_path, config['image_encoder_kwargs'].get('image_encoder_subpath', 'image_encoder')), + ) + clip_image_encoder = clip_image_encoder.eval() + else: + clip_image_encoder = None + + # Get Transformer. + # Keep all three transformers in fp32 — the same lesson learned in Stage 1/2: + # accelerate's mixed_precision="bf16" autocasts forward to bf16 while keeping + # master weights & Adam moments in fp32. If params live in bf16, every update + # (LR*grad ~ 1e-5 for LR=2e-6) falls below bf16 mantissa precision and is + # rounded to zero. CF Stage 3 FSDP also keeps fp32 master weights (no + # compute_dtype set on the MixedPrecision policy). + generator_transformer3d = WanTransformer3DModel_SelfForcing.from_pretrained( + os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')), + transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']), + low_cpu_mem_usage=True, + torch_dtype=torch.float32, + ) + # CF Stage 3 uses a 14B non-causal teacher as real_score; pass an explicit + # `--real_score_pretrained_model_name_or_path` to point at Wan2.1-T2V-14B. + real_score_base = args.real_score_pretrained_model_name_or_path or args.pretrained_model_name_or_path + real_score_transformer3d = WanTransformer3DModel.from_pretrained( + os.path.join(real_score_base, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')), + transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']), + low_cpu_mem_usage=True, + torch_dtype=torch.float32, + ) + fake_score_transformer3d = WanTransformer3DModel.from_pretrained( + os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')), + transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']), + low_cpu_mem_usage=True, + torch_dtype=torch.float32, + ) + + # Freeze vae and text_encoder and set generator_transformer3d to trainable + vae.requires_grad_(False) + text_encoder.requires_grad_(False) + generator_transformer3d.requires_grad_(False) + real_score_transformer3d.requires_grad_(False) + fake_score_transformer3d.requires_grad_(False) + if args.train_mode != "normal": + clip_image_encoder.requires_grad_(False) + + # CF Stage 3 init + if args.transformer_path is not None: + print(f"From checkpoint: {args.transformer_path}") + if args.transformer_path.endswith("safetensors"): + from safetensors.torch import load_file, safe_open + state_dict = load_file(args.transformer_path) + else: + state_dict = torch.load(args.transformer_path, map_location="cpu") + state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict + state_dict = state_dict["generator_ema"] if "generator_ema" in state_dict else state_dict + if any(k.startswith("model.") for k in state_dict.keys()): + state_dict = {k.replace("model.", "", 1) if k.startswith("model.") else k: v for k, v in state_dict.items()} + + m, u = generator_transformer3d.load_state_dict(state_dict, strict=False) + m, u = real_score_transformer3d.load_state_dict(state_dict, strict=False) + m, u = fake_score_transformer3d.load_state_dict(state_dict, strict=False) + print(f"missing keys: {len(m)}, unexpected keys: {len(u)}") + assert len(u) == 0 + + if args.ode_transformer_path is not None: + print(f"From checkpoint: {args.ode_transformer_path}") + if args.ode_transformer_path.endswith("safetensors"): + from safetensors.torch import load_file, safe_open + state_dict = load_file(args.ode_transformer_path) + else: + state_dict = torch.load(args.ode_transformer_path, map_location="cpu") + state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict + state_dict = state_dict["generator_ema"] if "generator_ema" in state_dict else state_dict + state_dict = state_dict["generator"] if "generator" in state_dict else state_dict + if any(k.startswith("model.") for k in state_dict.keys()): + state_dict = {k.replace("model.", "", 1) if k.startswith("model.") else k: v for k, v in state_dict.items()} + + m, u = generator_transformer3d.load_state_dict(state_dict, strict=False) + print(f"missing keys: {len(m)}, unexpected keys: {len(u)}") + assert len(u) == 0 + + if args.vae_path is not None: + print(f"From checkpoint: {args.vae_path}") + if args.vae_path.endswith("safetensors"): + from safetensors.torch import load_file, safe_open + state_dict = load_file(args.vae_path) + else: + state_dict = torch.load(args.vae_path, map_location="cpu") + state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict + + m, u = vae.load_state_dict(state_dict, strict=False) + print(f"missing keys: {len(m)}, unexpected keys: {len(u)}") + assert len(u) == 0 + + # A good trainable modules is showed below now. + # For 3D Patch: trainable_modules = ['ff.net', 'pos_embed', 'attn2', 'proj_out', 'timepositionalencoding', 'h_position', 'w_position'] + # For 2D Patch: trainable_modules = ['ff.net', 'attn2', 'timepositionalencoding', 'h_position', 'w_position'] + generator_transformer3d.train() + fake_score_transformer3d.train() + if accelerator.is_main_process: + accelerator.print( + f"Trainable modules '{args.trainable_modules}'." + ) + for name, param in generator_transformer3d.named_parameters(): + for trainable_module_name in args.trainable_modules + args.trainable_modules_low_learning_rate: + if trainable_module_name in name: + param.requires_grad = True + break + for name, param in fake_score_transformer3d.named_parameters(): + for trainable_module_name in args.trainable_modules + args.trainable_modules_low_learning_rate: + if trainable_module_name in name: + param.requires_grad = True + break + + # Create EMA for the generator_transformer3d. + if args.use_ema: + if zero_stage == 3: + raise NotImplementedError("FSDP does not support EMA.") + + ema_generator_transformer3d = EMAModel( + generator_transformer3d.parameters(), + model_cls=WanTransformer3DModel_SelfForcing, + model_config=generator_transformer3d.config, + ) + + # `accelerate` 0.16.0 will have better support for customized saving + if version.parse(accelerate.__version__) >= version.parse("0.16.0"): + # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format + if fsdp_stage != 0 or zero_stage == 3: + def save_model_hook(models, weights, output_dir): + accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True) + if accelerator.is_main_process: + from safetensors.torch import save_file + + safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors") + accelerate_state_dict = {k: v.to(dtype=weight_dtype) for k, v in accelerate_state_dict.items()} + save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"}) + + with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file: + pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file) + + def load_model_hook(models, input_dir): + pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl") + if os.path.exists(pkl_path): + with open(pkl_path, 'rb') as file: + loaded_number, _ = pickle.load(file) + batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0) + print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.") + else: + # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format + def save_model_hook(models, weights, output_dir): + if accelerator.is_main_process: + if args.use_ema and models[0] is generator_transformer3d: + ema_generator_transformer3d.save_pretrained(os.path.join(output_dir, "transformer_ema")) + + models[0].save_pretrained(os.path.join(output_dir, "transformer")) + if not args.use_deepspeed: + weights.pop() + + with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file: + pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file) + + def load_model_hook(models, input_dir): + if args.use_ema and models[0] is generator_transformer3d: + ema_path = os.path.join(input_dir, "transformer_ema") + _, ema_kwargs = WanTransformer3DModel_SelfForcing.load_config(ema_path, return_unused_kwargs=True) + load_model = WanTransformer3DModel_SelfForcing.from_pretrained( + input_dir, subfolder="transformer_ema", + transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']), + ) + load_model = EMAModel(load_model.parameters(), model_cls=WanTransformer3DModel_SelfForcing, model_config=load_model.config) + load_model.load_state_dict(ema_kwargs) + + ema_generator_transformer3d.load_state_dict(load_model.state_dict()) + ema_generator_transformer3d.to(accelerator.device) + del load_model + + for i in range(len(models)): + # pop models so that they are not loaded again + model = models.pop() + + # load diffusers style into model + load_model = WanTransformer3DModel.from_pretrained( + input_dir, subfolder="transformer" + ) + model.register_to_config(**load_model.config) + + model.load_state_dict(load_model.state_dict()) + del load_model + + pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl") + if os.path.exists(pkl_path): + with open(pkl_path, 'rb') as file: + loaded_number, _ = pickle.load(file) + batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0) + print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.") + + accelerator.register_save_state_pre_hook(save_model_hook) + accelerator.register_load_state_pre_hook(load_model_hook) + accelerator_fake_score_transformer3d.register_save_state_pre_hook(save_model_hook) + accelerator_fake_score_transformer3d.register_load_state_pre_hook(load_model_hook) + + if args.gradient_checkpointing: + generator_transformer3d.enable_gradient_checkpointing() + fake_score_transformer3d.enable_gradient_checkpointing() + + # Enable TF32 for faster training on Ampere GPUs, + # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices + if args.allow_tf32: + torch.backends.cuda.matmul.allow_tf32 = True + + if args.scale_lr: + args.learning_rate = ( + args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes + ) + + # Initialize the optimizer + if args.use_8bit_adam: + try: + import bitsandbytes as bnb + except ImportError: + raise ImportError( + "Please install bitsandbytes to use 8-bit Adam. You can do so by running `pip install bitsandbytes`" + ) + + optimizer_cls = bnb.optim.AdamW8bit + elif args.use_came: + try: + from came_pytorch import CAME + except Exception: + raise ImportError( + "Please install came_pytorch to use CAME. You can do so by running `pip install came_pytorch`" + ) + + optimizer_cls = CAME + else: + optimizer_cls = torch.optim.AdamW + + trainable_params = list(filter(lambda p: p.requires_grad, generator_transformer3d.parameters())) + trainable_params_optim = [ + {'params': [], 'lr': args.learning_rate}, + {'params': [], 'lr': args.learning_rate / 2}, + ] + in_already = [] + for name, param in generator_transformer3d.named_parameters(): + high_lr_flag = False + if name in in_already: + continue + for trainable_module_name in args.trainable_modules: + if trainable_module_name in name: + in_already.append(name) + high_lr_flag = True + trainable_params_optim[0]['params'].append(param) + if accelerator.is_main_process: + print(f"Set {name} to lr : {args.learning_rate}") + break + if high_lr_flag: + continue + for trainable_module_name in args.trainable_modules_low_learning_rate: + if trainable_module_name in name: + in_already.append(name) + trainable_params_optim[1]['params'].append(param) + if accelerator.is_main_process: + print(f"Set {name} to lr : {args.learning_rate / 2}") + break + + fake_trainable_params = list(filter(lambda p: p.requires_grad, fake_score_transformer3d.parameters())) + fake_trainable_params_optim = [ + {'params': [], 'lr': args.learning_rate}, + {'params': [], 'lr': args.learning_rate / 2}, + ] + in_already = [] + for name, param in fake_score_transformer3d.named_parameters(): + high_lr_flag = False + if name in in_already: + continue + for trainable_module_name in args.trainable_modules: + if trainable_module_name in name: + in_already.append(name) + high_lr_flag = True + fake_trainable_params_optim[0]['params'].append(param) + if accelerator.is_main_process: + print(f"Set {name} to lr : {args.learning_rate}") + break + if high_lr_flag: + continue + for trainable_module_name in args.trainable_modules_low_learning_rate: + if trainable_module_name in name: + in_already.append(name) + fake_trainable_params_optim[1]['params'].append(param) + if accelerator.is_main_process: + print(f"Set {name} to lr : {args.learning_rate / 2}") + break + + if args.use_came: + optimizer = optimizer_cls( + trainable_params_optim, + lr=args.learning_rate, + # weight_decay=args.adam_weight_decay, + betas=(0.9, 0.999, 0.9999), + eps=(1e-30, 1e-16) + ) + critic_optimizer = optimizer_cls( + fake_trainable_params_optim, + lr=args.learning_rate_critic, + # weight_decay=args.adam_weight_decay, + betas=(0.9, 0.999, 0.9999), + eps=(1e-30, 1e-16) + ) + else: + optimizer = optimizer_cls( + trainable_params_optim, + lr=args.learning_rate, + betas=(args.adam_beta1, args.adam_beta2), + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + ) + critic_optimizer = optimizer_cls( + fake_trainable_params_optim, + lr=args.learning_rate_critic, + betas=(args.adam_beta1, args.adam_beta2), + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + ) + + # Get the training dataset + sample_n_frames_bucket_interval = vae.config.temporal_compression_ratio + + if args.fix_sample_size is not None and args.enable_bucket: + args.video_sample_size = max(max(args.fix_sample_size), args.video_sample_size) + args.image_sample_size = max(max(args.fix_sample_size), args.image_sample_size) + args.training_with_video_token_length = False + args.random_hw_adapt = False + + # Get the dataset + if args.train_mode != "normal" or args.use_teacher_forcing: + train_dataset = ImageVideoDataset( + args.train_data_meta, args.train_data_dir, + video_sample_size=args.video_sample_size, video_sample_stride=args.video_sample_stride, video_sample_n_frames=args.video_sample_n_frames, + video_repeat=args.video_repeat, + image_sample_size=args.image_sample_size, + enable_bucket=args.enable_bucket, enable_inpaint=True if args.train_mode != "normal" else False, + ) + else: + train_dataset = TextDataset( + args.train_data_meta + ) + + def get_length_to_frame_num(token_length): + if args.image_sample_size > args.video_sample_size: + sample_sizes = list(range(args.video_sample_size, args.image_sample_size + 1, 128)) + + if sample_sizes[-1] != args.image_sample_size: + sample_sizes.append(args.image_sample_size) + else: + sample_sizes = [args.image_sample_size] + + length_to_frame_num = { + sample_size: min(token_length / sample_size / sample_size, args.video_sample_n_frames) // sample_n_frames_bucket_interval * sample_n_frames_bucket_interval + 1 for sample_size in sample_sizes + } + + return length_to_frame_num + + if (args.enable_bucket and args.train_mode != "normal") or args.use_teacher_forcing: + aspect_ratio_sample_size = {key : [x / 512 * args.video_sample_size for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()} + batch_sampler_generator = torch.Generator().manual_seed(args.seed) + batch_sampler = AspectRatioBatchImageVideoSampler( + sampler=RandomSampler(train_dataset, generator=batch_sampler_generator), dataset=train_dataset.dataset, + batch_size=args.train_batch_size, train_folder = args.train_data_dir, drop_last=True, + aspect_ratios=aspect_ratio_sample_size, + ) + + def collate_fn(examples): + # Get token length + target_token_length = args.video_sample_n_frames * args.token_sample_size * args.token_sample_size + length_to_frame_num = get_length_to_frame_num(target_token_length) + + # Create new output + new_examples = {} + new_examples["target_token_length"] = target_token_length + new_examples["pixel_values"] = [] + new_examples["text"] = [] + # Used in Inpaint mode + if args.train_mode != "normal": + new_examples["mask_pixel_values"] = [] + new_examples["mask"] = [] + new_examples["clip_pixel_values"] = [] + + # Get downsample ratio in image and videos + pixel_value = examples[0]["pixel_values"] + data_type = examples[0]["data_type"] + f, h, w, c = np.shape(pixel_value) + if data_type == 'image': + random_downsample_ratio = 1 if not args.random_hw_adapt else get_random_downsample_ratio(args.image_sample_size, image_ratio=[args.image_sample_size / args.video_sample_size], rng=rng) + + aspect_ratio_sample_size = {key : [x / 512 * args.image_sample_size / random_downsample_ratio for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()} + aspect_ratio_random_crop_sample_size = {key : [x / 512 * args.image_sample_size / random_downsample_ratio for x in ASPECT_RATIO_RANDOM_CROP_512[key]] for key in ASPECT_RATIO_RANDOM_CROP_512.keys()} + + batch_video_length = args.video_sample_n_frames + sample_n_frames_bucket_interval + else: + if args.random_hw_adapt: + if args.training_with_video_token_length: + local_min_size = np.min(np.array([np.mean(np.array([np.shape(example["pixel_values"])[1], np.shape(example["pixel_values"])[2]])) for example in examples])) + # The video will be resized to a lower resolution than its own. + choice_list = [length for length in list(length_to_frame_num.keys()) if length < local_min_size * 1.25] + if len(choice_list) == 0: + choice_list = list(length_to_frame_num.keys()) + if rng is None: + local_video_sample_size = np.random.choice(choice_list) + else: + local_video_sample_size = rng.choice(choice_list) + batch_video_length = length_to_frame_num[local_video_sample_size] + random_downsample_ratio = args.video_sample_size / local_video_sample_size + else: + random_downsample_ratio = get_random_downsample_ratio( + args.video_sample_size, rng=rng) + batch_video_length = args.video_sample_n_frames + sample_n_frames_bucket_interval + else: + random_downsample_ratio = 1 + batch_video_length = args.video_sample_n_frames + sample_n_frames_bucket_interval + + aspect_ratio_sample_size = {key : [x / 512 * args.video_sample_size / random_downsample_ratio for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()} + aspect_ratio_random_crop_sample_size = {key : [x / 512 * args.video_sample_size / random_downsample_ratio for x in ASPECT_RATIO_RANDOM_CROP_512[key]] for key in ASPECT_RATIO_RANDOM_CROP_512.keys()} + + if args.fix_sample_size is not None: + fix_sample_size = [int(x / 16) * 16 for x in args.fix_sample_size] + elif args.random_ratio_crop: + if rng is None: + random_sample_size = aspect_ratio_random_crop_sample_size[ + np.random.choice(list(aspect_ratio_random_crop_sample_size.keys()), p = ASPECT_RATIO_RANDOM_CROP_PROB) + ] + else: + random_sample_size = aspect_ratio_random_crop_sample_size[ + rng.choice(list(aspect_ratio_random_crop_sample_size.keys()), p = ASPECT_RATIO_RANDOM_CROP_PROB) + ] + random_sample_size = [int(x / 16) * 16 for x in random_sample_size] + else: + closest_size, closest_ratio = get_closest_ratio(h, w, ratios=aspect_ratio_sample_size) + closest_size = [int(x / 16) * 16 for x in closest_size] + + min_example_length = min( + [example["pixel_values"].shape[0] for example in examples] + ) + batch_video_length = int(min(batch_video_length, min_example_length)) + + # Magvae needs the number of frames to be 4n + 1. + batch_video_length = (batch_video_length - 1) // sample_n_frames_bucket_interval * sample_n_frames_bucket_interval + 1 + + # KV cache training requires latent frames divisible by num_frame_per_block + if args.use_kv_cache_training: + k = (batch_video_length - 1) // sample_n_frames_bucket_interval + if args.independent_first_frame: + # latent_frames - 1 = k must be divisible by num_frame_per_block + k = (k // args.num_frame_per_block) * args.num_frame_per_block + else: + # latent_frames = k + 1 must be divisible by num_frame_per_block + k = ((k + 1) // args.num_frame_per_block) * args.num_frame_per_block - 1 + batch_video_length = k * sample_n_frames_bucket_interval + 1 + + if batch_video_length <= 0: + batch_video_length = 1 + + for example in examples: + if args.fix_sample_size is not None: + # To 0~1 + pixel_values = torch.from_numpy(example["pixel_values"]).permute(0, 3, 1, 2).contiguous() + pixel_values = pixel_values / 255. + + # Get adapt hw for resize + fix_sample_size = list(map(lambda x: int(x), fix_sample_size)) + transform = transforms.Compose([ + transforms.Resize(fix_sample_size, interpolation=transforms.InterpolationMode.BILINEAR), # Image.BICUBIC + transforms.CenterCrop(fix_sample_size), + transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True), + ]) + elif args.random_ratio_crop: + # To 0~1 + pixel_values = torch.from_numpy(example["pixel_values"]).permute(0, 3, 1, 2).contiguous() + pixel_values = pixel_values / 255. + + # Get adapt hw for resize + b, c, h, w = pixel_values.size() + th, tw = random_sample_size + if th / tw > h / w: + nh = int(th) + nw = int(w / h * nh) + else: + nw = int(tw) + nh = int(h / w * nw) + + transform = transforms.Compose([ + transforms.Resize([nh, nw]), + transforms.CenterCrop([int(x) for x in random_sample_size]), + transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True), + ]) + else: + # To 0~1 + pixel_values = torch.from_numpy(example["pixel_values"]).permute(0, 3, 1, 2).contiguous() + pixel_values = pixel_values / 255. + + # Get adapt hw for resize + closest_size = list(map(lambda x: int(x), closest_size)) + if closest_size[0] / h > closest_size[1] / w: + resize_size = closest_size[0], int(w * closest_size[0] / h) + else: + resize_size = int(h * closest_size[1] / w), closest_size[1] + + transform = transforms.Compose([ + transforms.Resize(resize_size, interpolation=transforms.InterpolationMode.BILINEAR), # Image.BICUBIC + transforms.CenterCrop(closest_size), + transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True), + ]) + + new_examples["pixel_values"].append(transform(pixel_values)[:batch_video_length]) + new_examples["text"].append(example["text"]) + + if args.train_mode != "normal": + mask = get_random_mask(new_examples["pixel_values"][-1].size(), image_start_only=True) + mask_pixel_values = new_examples["pixel_values"][-1] * (1 - mask) + # Wan 2.1 use 0 for masked pixels + # + torch.ones_like(new_examples["pixel_values"][-1]) * -1 * mask + new_examples["mask_pixel_values"].append(mask_pixel_values) + new_examples["mask"].append(mask) + + clip_pixel_values = new_examples["pixel_values"][-1][0].permute(1, 2, 0).contiguous() + clip_pixel_values = (clip_pixel_values * 0.5 + 0.5) * 255 + new_examples["clip_pixel_values"].append(clip_pixel_values) + + # Limit the number of frames to the same + new_examples["pixel_values"] = torch.stack([example for example in new_examples["pixel_values"]]) + if args.train_mode != "normal": + new_examples["mask_pixel_values"] = torch.stack([example for example in new_examples["mask_pixel_values"]]) + new_examples["mask"] = torch.stack([example for example in new_examples["mask"]]) + new_examples["clip_pixel_values"] = torch.stack([example for example in new_examples["clip_pixel_values"]]) + + # Encode prompts when enable_text_encoder_in_dataloader=True + if args.enable_text_encoder_in_dataloader: + prompt_ids = tokenizer( + new_examples['text'], + max_length=args.tokenizer_max_length, + padding="max_length", + add_special_tokens=True, + truncation=True, + return_tensors="pt" + ) + text_input_ids = prompt_ids.input_ids + prompt_attention_mask = prompt_ids.attention_mask + + seq_lens = prompt_attention_mask.gt(0).sum(dim=1).long() + prompt_embeds = text_encoder(text_input_ids.to("cpu"), attention_mask=prompt_attention_mask.to("cpu"))[0] + prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] + + new_examples['encoder_attention_mask'] = prompt_ids.attention_mask + new_examples['encoder_hidden_states'] = prompt_embeds + + neg_txt = [ + "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" for text in batch['text'] + ] + neg_prompt_ids = tokenizer( + neg_txt, + max_length=args.tokenizer_max_length, + padding="max_length", + add_special_tokens=True, + truncation=True, + return_tensors="pt" + ) + neg_text_input_ids = neg_prompt_ids.input_ids + neg_prompt_attention_mask = neg_prompt_ids.attention_mask + + neg_seq_lens = neg_prompt_attention_mask.gt(0).sum(dim=1).long() + neg_prompt_embeds = text_encoder(neg_text_input_ids.to("cpu"), attention_mask=neg_prompt_attention_mask.to("cpu"))[0] + neg_prompt_embeds = [u[:v] for u, v in zip(neg_prompt_embeds, neg_seq_lens)] + + new_examples['neg_encoder_attention_mask'] = neg_prompt_ids.attention_mask + new_examples['neg_encoder_hidden_states'] = neg_prompt_embeds + + return new_examples + + # DataLoaders creation: + train_dataloader = torch.utils.data.DataLoader( + train_dataset, + batch_sampler=batch_sampler, + collate_fn=collate_fn, + persistent_workers=True if args.dataloader_num_workers != 0 else False, + num_workers=args.dataloader_num_workers, + ) + elif args.train_mode == "normal": + def collate_fn(examples): + new_examples = {} + new_examples["text"] = [] + for example in examples: + new_examples["text"].append(example["text"]) + + # Encode prompts when enable_text_encoder_in_dataloader=True + if args.enable_text_encoder_in_dataloader: + prompt_ids = tokenizer( + new_examples['text'], + max_length=args.tokenizer_max_length, + padding="max_length", + add_special_tokens=True, + truncation=True, + return_tensors="pt" + ) + text_input_ids = prompt_ids.input_ids + prompt_attention_mask = prompt_ids.attention_mask + + seq_lens = prompt_attention_mask.gt(0).sum(dim=1).long() + prompt_embeds = text_encoder(text_input_ids.to("cpu"), attention_mask=prompt_attention_mask.to("cpu"))[0] + prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] + + new_examples['encoder_attention_mask'] = prompt_ids.attention_mask + new_examples['encoder_hidden_states'] = prompt_embeds + + neg_txt = [ + "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" for text in batch['text'] + ] + neg_prompt_ids = tokenizer( + neg_txt, + max_length=args.tokenizer_max_length, + padding="max_length", + add_special_tokens=True, + truncation=True, + return_tensors="pt" + ) + neg_text_input_ids = neg_prompt_ids.input_ids + neg_prompt_attention_mask = neg_prompt_ids.attention_mask + + neg_seq_lens = neg_prompt_attention_mask.gt(0).sum(dim=1).long() + neg_prompt_embeds = text_encoder(neg_text_input_ids.to("cpu"), attention_mask=neg_prompt_attention_mask.to("cpu"))[0] + neg_prompt_embeds = [u[:v] for u, v in zip(neg_prompt_embeds, neg_seq_lens)] + + new_examples['neg_encoder_attention_mask'] = neg_prompt_ids.attention_mask + new_examples['neg_encoder_hidden_states'] = neg_prompt_embeds + + return new_examples + + batch_sampler_generator = torch.Generator().manual_seed(args.seed) + batch_sampler = BatchSampler(RandomSampler(train_dataset, generator=batch_sampler_generator), batch_size=args.train_batch_size, drop_last=True) + + # DataLoaders creation: + train_dataloader = torch.utils.data.DataLoader( + train_dataset, + batch_sampler=batch_sampler, + collate_fn=collate_fn, + persistent_workers=True if args.dataloader_num_workers != 0 else False, + num_workers=args.dataloader_num_workers, + ) + else: + # DataLoaders creation: + batch_sampler_generator = torch.Generator().manual_seed(args.seed) + batch_sampler = ImageVideoSampler(RandomSampler(train_dataset, generator=batch_sampler_generator), train_dataset, args.train_batch_size) + train_dataloader = torch.utils.data.DataLoader( + train_dataset, + batch_sampler=batch_sampler, + persistent_workers=True if args.dataloader_num_workers != 0 else False, + num_workers=args.dataloader_num_workers, + ) + + # Scheduler and math around the number of training steps. + overrode_max_train_steps = False + num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) + if args.max_train_steps is None: + args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch + overrode_max_train_steps = True + + lr_scheduler = get_scheduler( + args.lr_scheduler, + optimizer=optimizer, + num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes, + num_training_steps=args.max_train_steps * accelerator.num_processes, + ) + fake_score_lr_scheduler = get_scheduler( + args.lr_scheduler, + optimizer=optimizer, + num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes, + num_training_steps=args.max_train_steps * accelerator.num_processes, + ) + + # Prepare everything with our `accelerator`. + generator_transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( + generator_transformer3d, optimizer, train_dataloader, lr_scheduler + ) + fake_score_transformer3d, critic_optimizer, fake_score_lr_scheduler= accelerator_fake_score_transformer3d.prepare( + fake_score_transformer3d, critic_optimizer, fake_score_lr_scheduler + ) + if fsdp_stage != 0 or zero_stage != 0: + from functools import partial + + from videox_fun.dist import set_multi_gpus_devices, shard_model + shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype) + real_score_transformer3d = shard_fn(real_score_transformer3d) + if fsdp_stage != 0 or zero_stage != 0: + from functools import partial + + from videox_fun.dist import set_multi_gpus_devices, shard_model + shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype) + text_encoder = shard_fn(text_encoder) + + # Move text_encode and vae to gpu and cast to weight_dtype + vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype) + real_score_transformer3d.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype) + if not args.enable_text_encoder_in_dataloader: + text_encoder.to(accelerator.device if not args.low_vram else "cpu") + if args.train_mode != "normal": + clip_image_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype) + + if args.use_ema: + ema_generator_transformer3d.to(accelerator.device) + + # We need to recalculate our total training steps as the size of the training dataloader may have changed. + num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) + if overrode_max_train_steps: + args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch + # Afterwards we recalculate our number of training epochs + args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch) + + # We need to initialize the trackers we use, and also store our configuration. + # The trackers initializes automatically on the main process. + if accelerator.is_main_process: + tracker_config = dict(vars(args)) + keys_to_pop = [k for k, v in tracker_config.items() if isinstance(v, list)] + for k in keys_to_pop: + tracker_config.pop(k) + print(f"Removed tracker_config['{k}']") + accelerator.init_trackers(args.tracker_project_name, tracker_config) + + # Function for unwrapping if model was compiled with `torch.compile`. + def unwrap_model(model): + model = accelerator.unwrap_model(model) + model = model._orig_mod if is_compiled_module(model) else model + return model + + # Train! + total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps + + logger.info("***** Running training *****") + logger.info(f" Num examples = {len(train_dataset)}") + logger.info(f" Num Epochs = {args.num_train_epochs}") + logger.info(f" Instantaneous batch size per device = {args.train_batch_size}") + logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}") + logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}") + logger.info(f" Total optimization steps = {args.max_train_steps}") + global_step = 0 + first_epoch = 0 + + # Potentially load in the weights and states from a previous save + if args.resume_from_checkpoint: + if args.resume_from_checkpoint != "latest": + path = os.path.basename(args.resume_from_checkpoint) + else: + # Get the most recent checkpoint + dirs = os.listdir(args.output_dir) + dirs = [d for d in dirs if d.startswith("checkpoint")] + dirs = sorted(dirs, key=lambda x: int(x.split("-")[1])) + path = dirs[-1] if len(dirs) > 0 else None + + if path is None: + accelerator.print( + f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run." + ) + args.resume_from_checkpoint = None + initial_global_step = 0 + else: + global_step = int(path.split("-")[1]) + + initial_global_step = global_step + + pkl_path = os.path.join(os.path.join(args.output_dir, path), "sampler_pos_start.pkl") + if os.path.exists(pkl_path): + with open(pkl_path, 'rb') as file: + _, first_epoch = pickle.load(file) + else: + first_epoch = global_step // num_update_steps_per_epoch + print(f"Load pkl from {pkl_path}. Get first_epoch = {first_epoch}.") + + accelerator.print(f"Resuming from checkpoint {path}") + fake_score_path = os.path.join(path, "fake_score") + accelerator.load_state(os.path.join(args.output_dir, path)) + accelerator_fake_score_transformer3d.load_state(os.path.join(args.output_dir, fake_score_path)) + else: + initial_global_step = 0 + + progress_bar = tqdm( + range(0, args.max_train_steps), + initial=initial_global_step, + desc="Steps", + # Only show the progress bar once on each machine. + disable=not accelerator.is_local_main_process, + ) + + if args.multi_stream and args.train_mode != "normal": + # create extra cuda streams to speedup inpaint vae computation + vae_stream_1 = torch.cuda.Stream() + vae_stream_2 = torch.cuda.Stream() + else: + vae_stream_1 = None + vae_stream_2 = None + + idx_sampling = DiscreteSampling(args.train_sampling_steps, uniform_sampling=args.uniform_sampling) + + def randomize_denoising_step_indices( + denoising_step_indices_list, + train_sampling_steps, + torch_rng, + accelerator, + jitter_ratio=0.3, + ): + indices = list(denoising_step_indices_list) + n = len(indices) + + if n <= 2: + low = indices[1] + high = indices[0] - 1 + random_tail = torch.randint(low, high + 1, (1,)).item() + + result = torch.tensor([indices[0], random_tail]) + else: + result = [0] * n + result[0] = indices[0] + result[-1] = indices[-1] + + for i in range(1, n - 1): + gap_upper = indices[i - 1] - indices[i] + gap_lower = indices[i] - indices[i + 1] + + max_jitter = int(min(gap_upper, gap_lower) * jitter_ratio) + + if max_jitter > 0: + jitter = torch.randint( + -max_jitter, max_jitter + 1, (1,) + ).item() + else: + jitter = 0 + + result[i] = indices[i] + jitter + + for i in range(1, n): + if result[i] >= result[i - 1]: + result[i] = result[i - 1] - 1 + + result = [max(1, min(train_sampling_steps, x)) for x in result] + result = torch.tensor(result) + + if dist.is_initialized(): + result = result.to(accelerator.device) + dist.broadcast(result, src=0) + result = result.cpu() + return result + + for epoch in range(first_epoch, args.num_train_epochs): + train_dmd_loss = 0.0 + train_denoising_loss = 0.0 + batch_sampler.sampler.generator = torch.Generator().manual_seed(args.seed + epoch) + for step, batch in enumerate(train_dataloader): + # Data batch sanity check + if args.train_mode != "normal" and epoch == first_epoch and step == 0: + pixel_values, texts = batch['pixel_values'].cpu(), batch['text'] + pixel_values = rearrange(pixel_values, "b f c h w -> b c f h w") + os.makedirs(os.path.join(args.output_dir, "sanity_check"), exist_ok=True) + for idx, (pixel_value, text) in enumerate(zip(pixel_values, texts)): + pixel_value = pixel_value[None, ...] + gif_name = '-'.join(text.replace('/', '').split()[:10]) if not text == '' else f'{global_step}-{idx}' + save_videos_grid(pixel_value, f"{args.output_dir}/sanity_check/{gif_name[:10]}.mp4", rescale=True) + + clip_pixel_values, mask_pixel_values, texts = batch['clip_pixel_values'].cpu(), batch['mask_pixel_values'].cpu(), batch['text'] + mask_pixel_values = rearrange(mask_pixel_values, "b f c h w -> b c f h w") + for idx, (clip_pixel_value, pixel_value, text) in enumerate(zip(clip_pixel_values, mask_pixel_values, texts)): + pixel_value = pixel_value[None, ...] + Image.fromarray(np.uint8(clip_pixel_value)).save(f"{args.output_dir}/sanity_check/clip_{gif_name[:10] if not text == '' else f'{global_step}-{idx}'}.png") + save_videos_grid(pixel_value, f"{args.output_dir}/sanity_check/mask_{gif_name[:10] if not text == '' else f'{global_step}-{idx}'}.mp4", rescale=True) + + with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + if args.train_mode != "normal" or args.use_teacher_forcing: + # Convert images to latent space + pixel_values = batch["pixel_values"].to(weight_dtype) + + # Increase the batch size when the length of the latent sequence of the current sample is small + if args.auto_tile_batch_size and args.training_with_video_token_length and zero_stage != 3: + if args.video_sample_n_frames * args.token_sample_size * args.token_sample_size // 16 >= pixel_values.size()[1] * pixel_values.size()[3] * pixel_values.size()[4]: + pixel_values = torch.tile(pixel_values, (4, 1, 1, 1, 1)) + if args.enable_text_encoder_in_dataloader: + batch['encoder_hidden_states'] = torch.tile(batch['encoder_hidden_states'], (4, 1, 1)) + batch['encoder_attention_mask'] = torch.tile(batch['encoder_attention_mask'], (4, 1)) + batch['neg_encoder_hidden_states'] = torch.tile(batch['neg_encoder_hidden_states'], (4, 1, 1)) + batch['neg_encoder_attention_mask'] = torch.tile(batch['neg_encoder_attention_mask'], (4, 1)) + else: + batch['text'] = batch['text'] * 4 + elif args.video_sample_n_frames * args.token_sample_size * args.token_sample_size // 4 >= pixel_values.size()[1] * pixel_values.size()[3] * pixel_values.size()[4]: + pixel_values = torch.tile(pixel_values, (2, 1, 1, 1, 1)) + if args.enable_text_encoder_in_dataloader: + batch['encoder_hidden_states'] = torch.tile(batch['encoder_hidden_states'], (2, 1, 1)) + batch['encoder_attention_mask'] = torch.tile(batch['encoder_attention_mask'], (2, 1)) + batch['neg_encoder_hidden_states'] = torch.tile(batch['neg_encoder_hidden_states'], (2, 1, 1)) + batch['neg_encoder_attention_mask'] = torch.tile(batch['neg_encoder_attention_mask'], (2, 1)) + else: + batch['text'] = batch['text'] * 2 + if args.train_mode != "normal": + clip_pixel_values = batch["clip_pixel_values"].to(weight_dtype) + mask_pixel_values = batch["mask_pixel_values"].to(weight_dtype) + mask = batch["mask"].to(weight_dtype) + # Increase the batch size when the length of the latent sequence of the current sample is small + if args.auto_tile_batch_size and args.training_with_video_token_length and zero_stage != 3: + if args.video_sample_n_frames * args.token_sample_size * args.token_sample_size // 16 >= pixel_values.size()[1] * pixel_values.size()[3] * pixel_values.size()[4]: + clip_pixel_values = torch.tile(clip_pixel_values, (4, 1, 1, 1)) + mask_pixel_values = torch.tile(mask_pixel_values, (4, 1, 1, 1, 1)) + mask = torch.tile(mask, (4, 1, 1, 1, 1)) + elif args.video_sample_n_frames * args.token_sample_size * args.token_sample_size // 4 >= pixel_values.size()[1] * pixel_values.size()[3] * pixel_values.size()[4]: + clip_pixel_values = torch.tile(clip_pixel_values, (2, 1, 1, 1)) + mask_pixel_values = torch.tile(mask_pixel_values, (2, 1, 1, 1, 1)) + mask = torch.tile(mask, (2, 1, 1, 1, 1)) + + if args.random_frame_crop: + def _create_special_list(length): + if length == 1: + return [1.0] + if length >= 2: + last_element = 0.90 + remaining_sum = 1.0 - last_element + other_elements_value = remaining_sum / (length - 1) + special_list = [other_elements_value] * (length - 1) + [last_element] + return special_list + select_frames = [_tmp for _tmp in list(range(sample_n_frames_bucket_interval + 1, args.video_sample_n_frames + sample_n_frames_bucket_interval, sample_n_frames_bucket_interval))] + select_frames_prob = np.array(_create_special_list(len(select_frames))) + + if len(select_frames) != 0: + if rng is None: + temp_n_frames = np.random.choice(select_frames, p = select_frames_prob) + else: + temp_n_frames = rng.choice(select_frames, p = select_frames_prob) + else: + temp_n_frames = 1 + + # Magvae needs the number of frames to be 4n + 1. + temp_n_frames = (temp_n_frames - 1) // sample_n_frames_bucket_interval + 1 + + pixel_values = pixel_values[:, :temp_n_frames, :, :] + mask_pixel_values = mask_pixel_values[:, :temp_n_frames, :, :] + mask = mask[:, :temp_n_frames, :, :] + + # Keep all node same token length to accelerate the traning when resolution grows. + if args.keep_all_node_same_token_length: + if args.token_sample_size > 256: + numbers_list = list(range(256, args.token_sample_size + 1, 128)) + + if numbers_list[-1] != args.token_sample_size: + numbers_list.append(args.token_sample_size) + else: + numbers_list = [256] + numbers_list = [_number * _number * args.video_sample_n_frames for _number in numbers_list] + + actual_token_length = index_rng.choice(numbers_list) + actual_video_length = (min( + actual_token_length / pixel_values.size()[-1] / pixel_values.size()[-2], args.video_sample_n_frames + ) - 1) // sample_n_frames_bucket_interval * sample_n_frames_bucket_interval + 1 + actual_video_length = int(max(actual_video_length, 1)) + + # Magvae needs the number of frames to be 4n + 1. + actual_video_length = (actual_video_length - 1) // sample_n_frames_bucket_interval + 1 + + pixel_values = pixel_values[:, :actual_video_length, :, :] + mask_pixel_values = mask_pixel_values[:, :actual_video_length, :, :] + mask = mask[:, :actual_video_length, :, :] + + if args.low_vram: + torch.cuda.empty_cache() + vae.to(accelerator.device) + if args.train_mode != "normal": + clip_image_encoder.to(accelerator.device) + real_score_transformer3d = real_score_transformer3d.to("cpu") + if not args.enable_text_encoder_in_dataloader: + text_encoder.to("cpu") + + with torch.no_grad(): + # This way is quicker when batch grows up + def _batch_encode_vae(pixel_values): + pixel_values = rearrange(pixel_values, "b f c h w -> b c f h w") + bs = args.vae_mini_batch + new_pixel_values = [] + for i in range(0, pixel_values.shape[0], bs): + pixel_values_bs = pixel_values[i : i + bs] + pixel_values_bs = vae.encode(pixel_values_bs)[0] + pixel_values_bs = pixel_values_bs.sample() + new_pixel_values.append(pixel_values_bs) + return torch.cat(new_pixel_values, dim = 0) + if args.use_teacher_forcing: + clean_latents = _batch_encode_vae(pixel_values) + else: + clean_latents = None + + if args.train_mode != "normal": + # Encode inpaint latents. + mask_latents = _batch_encode_vae(mask_pixel_values) + if vae_stream_2 is not None: + torch.cuda.current_stream().wait_stream(vae_stream_2) + + mask = rearrange(mask, "b f c h w -> b c f h w") + mask = torch.concat( + [ + torch.repeat_interleave(mask[:, :, 0:1], repeats=4, dim=2), + mask[:, :, 1:] + ], dim=2 + ) + mask = mask.view(mask.shape[0], mask.shape[2] // 4, 4, mask.shape[3], mask.shape[4]) + mask = mask.transpose(1, 2) + mask = resize_mask(1 - mask, mask_latents) + + inpaint_latents = torch.concat([mask, mask_latents], dim=1) + + clip_context = [] + for clip_pixel_value in clip_pixel_values: + clip_image = Image.fromarray(np.uint8(clip_pixel_value.float().cpu().numpy())) + clip_image = TF.to_tensor(clip_image).sub_(0.5).div_(0.5).to(clip_image_encoder.device, weight_dtype) + _clip_context = clip_image_encoder([clip_image[:, None, :, :]]) + clip_context.append(_clip_context) + clip_context = torch.cat(clip_context) + + if args.use_teacher_forcing: + target_shape = clean_latents.size() + else: + target_shape = mask_latents.size() + else: + text = batch['text'] + if args.fix_sample_size is not None: + local_sample_size = [int(x / 16) * 16 for x in args.fix_sample_size] + num_frames = args.video_sample_n_frames + else: + if args.random_hw_adapt and args.training_with_video_token_length: + # Get token length + target_token_length = args.video_sample_n_frames * args.token_sample_size * args.token_sample_size + length_to_frame_num = get_length_to_frame_num(target_token_length) + + if rng is None: + local_length = np.random.choice(list(length_to_frame_num.keys())) + else: + local_length = rng.choice(list(length_to_frame_num.keys())) + num_frames = length_to_frame_num[local_length] + + aspect_ratio_sample_size = {key : [x / 512 * local_length for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()} + if rng is None: + aspect_ratio_key = np.random.choice(list(aspect_ratio_sample_size.keys())) + else: + aspect_ratio_key = rng.choice(list(aspect_ratio_sample_size.keys())) + local_sample_size = aspect_ratio_sample_size[aspect_ratio_key] + else: + num_frames = args.video_sample_n_frames + + aspect_ratio_sample_size = {key : [x / 512 * args.video_sample_size for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()} + if rng is None: + aspect_ratio_key = np.random.choice(list(aspect_ratio_sample_size.keys())) + else: + aspect_ratio_key = rng.choice(list(aspect_ratio_sample_size.keys())) + local_sample_size = aspect_ratio_sample_size[aspect_ratio_key] + local_sample_size = [int(x / 16) * 16 for x in local_sample_size] + + # Compute latent frame count + latent_num_frames = int((num_frames - 1) // vae.temporal_compression_ratio + 1) + + # Align latent_num_frames to num_frame_per_block for KV cache training + if args.use_kv_cache_training: + if args.independent_first_frame: + # latent_frames - 1 must be divisible by num_frame_per_block + k = latent_num_frames - 1 + k = (k // args.num_frame_per_block) * args.num_frame_per_block + latent_num_frames = k + 1 + else: + # latent_frames must be divisible by num_frame_per_block + latent_num_frames = (latent_num_frames // args.num_frame_per_block) * args.num_frame_per_block + latent_num_frames = max(latent_num_frames, args.num_frame_per_block) + + target_shape = ( + len(text), + vae.latent_channels, + latent_num_frames, + int(local_sample_size[0] // vae.spatial_compression_ratio), + int(local_sample_size[1] // vae.spatial_compression_ratio), + ) + clean_latents = None + + if args.low_vram: + vae.to('cpu') + real_score_transformer3d = real_score_transformer3d.to("cpu") + if args.train_mode != "normal": + clip_image_encoder.to('cpu') + torch.cuda.empty_cache() + if not args.enable_text_encoder_in_dataloader: + text_encoder.to(accelerator.device) + + if args.enable_text_encoder_in_dataloader: + prompt_embeds = batch['encoder_hidden_states'].to(device=accelerator.device) + neg_prompt_embeds = batch['neg_encoder_hidden_states'].to(device=accelerator.device) + else: + with torch.no_grad(): + prompt_ids = tokenizer( + batch['text'], + padding="max_length", + max_length=args.tokenizer_max_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt" + ) + text_input_ids = prompt_ids.input_ids + prompt_attention_mask = prompt_ids.attention_mask + + seq_lens = prompt_attention_mask.gt(0).sum(dim=1).long() + prompt_embeds = text_encoder(text_input_ids.to(accelerator.device), attention_mask=prompt_attention_mask.to(accelerator.device))[0] + prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] + + neg_txt = [ + "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" for text in batch['text'] + ] + neg_prompt_ids = tokenizer( + neg_txt, + padding="max_length", + max_length=args.tokenizer_max_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt" + ) + neg_text_input_ids = neg_prompt_ids.input_ids + neg_prompt_attention_mask = neg_prompt_ids.attention_mask + + neg_seq_lens = neg_prompt_attention_mask.gt(0).sum(dim=1).long() + neg_prompt_embeds = text_encoder(neg_text_input_ids.to(accelerator.device), attention_mask=neg_prompt_attention_mask.to(accelerator.device))[0] + neg_prompt_embeds = [u[:v] for u, v in zip(neg_prompt_embeds, neg_seq_lens)] + + if args.low_vram: + generator_transformer3d = generator_transformer3d.to(accelerator.device) + real_score_transformer3d = real_score_transformer3d.to(accelerator.device) + fake_score_transformer3d = fake_score_transformer3d.to(accelerator.device) + if not args.enable_text_encoder_in_dataloader: + text_encoder.to('cpu') + torch.cuda.empty_cache() + + with accelerator.accumulate(generator_transformer3d): + def get_sigmas(timesteps, n_dim=4, dtype=torch.float32): + sigmas = noise_scheduler.sigmas.to(device=accelerator.device, dtype=dtype) + schedule_timesteps = noise_scheduler.timesteps.to(accelerator.device) + timesteps = timesteps.to(accelerator.device) + + step_indices = [ + torch.argmin(torch.abs(schedule_timesteps - t)).item() + for t in timesteps + ] + step_indices = torch.tensor(step_indices, device=accelerator.device) + sigma = sigmas[step_indices].flatten() + + while len(sigma.shape) < n_dim: + sigma = sigma.unsqueeze(-1) + return sigma + + def add_noise(latents, noise, timesteps): + sigmas = get_sigmas(timesteps, n_dim=latents.ndim, dtype=latents.dtype) + return (1.0 - sigmas) * latents + sigmas * noise + + def generate_and_sync_list(num_denoising_steps, device): + indices = torch.randint(low=0, high=num_denoising_steps, size=(1,), generator=torch_rng, device=device) + if dist.is_initialized(): + dist.broadcast(indices, src=0) + return indices.tolist() + + def convert_flow_pred_to_x0( + scheduler, + flow_pred: torch.Tensor, + xt: torch.Tensor, + timestep: torch.Tensor + ) -> torch.Tensor: + """ + Convert flow matching's prediction to x0 prediction. + Supports both 4D [B, C, H, W] and 5D [B, C, F, H, W] inputs. + """ + original_dtype = flow_pred.dtype + device = flow_pred.device + + flow_pred = flow_pred.double() + xt = xt.double() + timesteps = scheduler.timesteps.to(device).double() + sigmas = scheduler.sigmas.to(device).double() + timestep = timestep.to(device).double() + + timestep_id = torch.argmin((timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) + sigma_t = sigmas[timestep_id] + + ndim = flow_pred.ndim + if ndim == 4: + sigma_t = sigma_t.view(-1, 1, 1, 1) + elif ndim == 5: + sigma_t = sigma_t.view(-1, 1, 1, 1, 1) + else: + raise ValueError(f"Expected 4D or 5D input, got {ndim}D tensor.") + + x0_pred = xt - sigma_t * flow_pred + return x0_pred.to(original_dtype) + + # Create discrete denoising steps (per-step, with optional randomization) + if getattr(args, 'randomize_step_indices', False): + random_indices = randomize_denoising_step_indices( + args.denoising_step_indices_list, + args.train_sampling_steps, + torch_rng, + accelerator, + jitter_ratio=getattr(args, 'index_jitter_ratio', 0.30), + ) + else: + random_indices = torch.tensor(args.denoising_step_indices_list) + + denoising_step_list = noise_scheduler.timesteps[args.train_sampling_steps - random_indices] + + # --- Main Training Logic --- + bsz, channel, num_frames, height, width = target_shape + if step % args.gen_update_interval == 0: + if args.use_kv_cache_training: + # Calculate frame_seq_length + patch_h, patch_w = accelerator.unwrap_model(generator_transformer3d).config.patch_size[1:] + frame_seq_length = (target_shape[3] * target_shape[4]) // (patch_h * patch_w) + + # Determine block structure with variable-length support + if not args.independent_first_frame: + assert num_frames % args.num_frame_per_block == 0 + max_num_blocks = num_frames // args.num_frame_per_block + assert args.score_num_frames % args.num_frame_per_block == 0 + min_num_blocks = args.score_num_frames // args.num_frame_per_block + else: + assert (num_frames - 1) % args.num_frame_per_block == 0 + max_num_blocks = (num_frames - 1) // args.num_frame_per_block + if args.score_num_frames > 1: + assert (args.score_num_frames - 1) % args.num_frame_per_block == 0 + min_num_blocks = (args.score_num_frames - 1) // args.num_frame_per_block + else: + min_num_blocks = 0 + + # Random sample number of blocks (Self-Forcing variable-length training) + if args.min_length_prob_bias > 0.0 and max_num_blocks > min_num_blocks: + # Weighted sampling: give min_num_blocks a higher probability + num_options = max_num_blocks - min_num_blocks + 1 + bias = min(args.min_length_prob_bias, 0.99) + remaining_prob = (1.0 - bias) / (num_options - 1) + probs = [remaining_prob] * num_options + probs[0] = bias # min_num_blocks gets the bias + probs_tensor = torch.tensor(probs, device=accelerator.device) + block_indices = torch.multinomial(probs_tensor, 1, generator=torch_rng) + num_generated_blocks = (min_num_blocks + block_indices).item() + else: + num_generated_blocks = torch.randint( + min_num_blocks, max_num_blocks + 1, (1,), + generator=torch_rng, device=accelerator.device + ).item() + if dist.is_initialized(): + _sync = torch.tensor([num_generated_blocks], device=accelerator.device) + dist.broadcast(_sync, src=0) + num_generated_blocks = _sync.item() + + all_num_frames = [args.num_frame_per_block] * num_generated_blocks + if args.independent_first_frame: + all_num_frames = [1] + all_num_frames + + num_generated_frames = sum(all_num_frames) + + # Initialize KV cache + num_layers = generator_transformer3d.config.num_layers + num_heads = generator_transformer3d.config.num_heads + head_dim = generator_transformer3d.config.dim // num_heads + text_len = 512 # T5 sequence length + + kv_cache = initialize_kv_cache_for_training( + batch_size=bsz, + num_frames=num_frames, + frame_seq_length=frame_seq_length, + num_layers=num_layers, + num_heads=num_heads, + head_dim=head_dim, + dtype=weight_dtype, + device=accelerator.device + ) + + crossattn_cache = initialize_crossattn_cache_for_training( + batch_size=bsz, + text_len=text_len, + num_layers=num_layers, + num_heads=num_heads, + head_dim=head_dim, + dtype=weight_dtype, + device=accelerator.device + ) + + # Block-by-block generation + generator_noise = torch.randn(target_shape, device=accelerator.device, generator=torch_rng, dtype=weight_dtype) + current_start_frame = 0 + num_input_frames = 0 # T2V mode + + # Use actual batch size from generator_noise (may differ due to SP) + actual_bsz = generator_noise.shape[0] + output_pred = torch.zeros_like(generator_noise) + + # Decide whether to use teacher forcing for this video (once per video, not per block) + use_teacher_forcing_step = ( + args.use_teacher_forcing and + torch.rand(1, generator=torch_rng, device=accelerator.device).item() < args.teacher_forcing_prob + ) + + # Same exit step across all blocks (matches original Self-Forcing default) + num_denoising_steps = len(denoising_step_list) + final_step_index = generate_and_sync_list(num_denoising_steps, device=accelerator.device)[0] + + # Only blocks in the last score_num_frames get gradient at exit step + # (matches Self-Forcing: start_gradient_frame_index = num_output_frames - 21) + start_gradient_frame_index = num_generated_frames - args.score_num_frames + + for block_idx, current_num_frames in enumerate(all_num_frames): + # Extract noise for current block + start_idx = current_start_frame - num_input_frames + end_idx = start_idx + current_num_frames + noisy_input = generator_noise[:, :, start_idx:end_idx] + + for index, current_timestep in enumerate(denoising_step_list): + is_final_step = (index == final_step_index) + timestep = torch.full( + [bsz, current_num_frames], + current_timestep, + device=noisy_input.device, + dtype=torch.int64 + ) + + # Only enable gradient for final step AND block within score window + if not is_final_step or current_start_frame < start_gradient_frame_index: + context_manager = torch.no_grad() + else: + context_manager = contextlib.nullcontext() + + with context_manager: + # Convert noisy_input to list format + noisy_input_list = [noisy_input[i] for i in range(bsz)] + + # Use full seq_len (consistent with inference code) + full_seq_len = frame_seq_length * num_frames + + generator_pred_block = generator_transformer3d( + x=noisy_input_list, + context=prompt_embeds, + t=timestep, + seq_len=full_seq_len, + kv_cache=kv_cache, + crossattn_cache=crossattn_cache, + current_start=current_start_frame * frame_seq_length, + cache_start=None, + y=inpaint_latents if args.train_mode != "normal" else None, + clip_fea=clip_context if args.train_mode != "normal" else None, + ) + + # Stack list output to tensor: [B, C, F, H, W] + if isinstance(generator_pred_block, list): + generator_pred_block = torch.stack(generator_pred_block, dim=0) + + # Flatten timestep for convert_flow_pred_to_x0: [B, F] -> [B*F] + generator_pred_block = convert_flow_pred_to_x0( + scheduler=noise_scheduler, + flow_pred=generator_pred_block, + xt=noisy_input, + timestep=timestep[:, 0] + ) + + if is_final_step: + break + + # Add noise for next step + next_timestep = denoising_step_list[index + 1] * torch.ones( + bsz, dtype=torch.long, device=noisy_input.device + ) + noisy_input = add_noise( + generator_pred_block, + torch.randn(generator_pred_block.shape, dtype=generator_pred_block.dtype, device=generator_pred_block.device, generator=torch_rng), + next_timestep + ) + + # Record output + output_pred[:, :, current_start_frame:current_start_frame + current_num_frames] = generator_pred_block + + # Update KV cache with clean context (consistent with inference: feed denoised_pred directly) + if block_idx < len(all_num_frames) - 1: + context_timestep = torch.ones([bsz, current_num_frames], device=accelerator.device, dtype=torch.int64) * args.context_noise + + # Use clean latents for teacher forcing, otherwise use denoised prediction directly + if use_teacher_forcing_step and clean_latents is not None: + context_input = clean_latents[:, :, start_idx:end_idx] + else: + context_input = generator_pred_block + + context_input_list = [context_input[i] for i in range(bsz)] + + # Use full seq_len (consistent with inference code) + full_seq_len = frame_seq_length * num_frames + + with torch.no_grad(): + generator_transformer3d( + x=context_input_list, + context=prompt_embeds, + t=context_timestep, + seq_len=full_seq_len, + kv_cache=kv_cache, + crossattn_cache=crossattn_cache, + current_start=current_start_frame * frame_seq_length, + cache_start=None, + y=inpaint_latents if args.train_mode != "normal" else None, + clip_fea=clip_context if args.train_mode != "normal" else None, + ) + + current_start_frame += current_num_frames + + # Final output — slice generated frames (may be < num_frames for variable-length) + generator_pred_full = output_pred[:, :, :num_generated_frames] + + # Gradient mask: first block gets no gradient when generating > min frames + # (matches Self-Forcing reference: model/base.py L182-L190) + min_num_frames_score = args.score_num_frames + need_gradient_mask = (num_generated_frames != min_num_frames_score) + gradient_mask = None + if need_gradient_mask: + gradient_mask = torch.ones_like(generator_pred_full, dtype=torch.bool) + if args.independent_first_frame: + gradient_mask[:, :, :1] = False + else: + gradient_mask[:, :, :args.num_frame_per_block] = False + + # Slice for score computation: last score_num_frames frames + if num_generated_frames > args.score_num_frames: + # Re-encode boundary for cleaner score input + generator_pred_for_score, score_num_frames, _ = slice_for_score( + generator_pred_full, vae, weight_dtype, + score_num_frames=args.score_num_frames, + independent_first_frame=args.independent_first_frame, + ) + else: + generator_pred_for_score = generator_pred_full + score_num_frames = num_generated_frames + + # Compute score_mask for DMD loss (matches Self-Forcing: dmd.py L199-204) + score_mask = None + if gradient_mask is not None: + mask_offset = num_generated_frames - score_num_frames + score_mask = gradient_mask[:, :, mask_offset:mask_offset + score_num_frames] + + # generator_pred = the sliced version for DMD loss + generator_pred = generator_pred_for_score + seq_len = frame_seq_length * score_num_frames # Score always on fixed window + + else: + # === Block mask training (flex attention, no KV cache) === + # Block mask training: use flex attention to process entire video at once + # Note: for long videos, use KV cache mode instead + score_mask = None # Block mask mode: no gradient mask needed + if num_frames > args.score_num_frames: + raise ValueError( + f"Block mask mode does not support variable-length training " + f"(video produces {num_frames} latent frames > score_num_frames={args.score_num_frames}). " + f"Use --use_kv_cache_training for long video training." + ) + + patch_h_bm, patch_w_bm = accelerator.unwrap_model(generator_transformer3d).config.patch_size[1:] + frame_seqlen_bm = (height * width) // (patch_h_bm * patch_w_bm) + + # Standard backward simulation training + generator_noise = torch.randn(target_shape, device=accelerator.device, generator=torch_rng, dtype=weight_dtype) + num_denoising_steps = len(denoising_step_list) + final_step_index = generate_and_sync_list(num_denoising_steps, device=generator_noise.device)[0] + + # Precompute seq_len once (same for all steps) + seq_len = frame_seqlen_bm * num_frames + + # Decide whether to use teacher forcing for this step + use_teacher_forcing_step = ( + args.use_teacher_forcing and + torch.rand(1, generator=torch_rng, device=accelerator.device).item() < args.teacher_forcing_prob + ) + + # Create appropriate block mask based on teacher forcing decision + if use_teacher_forcing_step and clean_latents is not None: + # Teacher forcing: clean + noisy sequence mask + accelerator.unwrap_model(generator_transformer3d).create_teacher_forcing_mask( + device=accelerator.device, + num_frames=num_frames, + frame_seqlen=frame_seqlen_bm, + num_frame_per_block=args.num_frame_per_block, + ) + # Prepare clean_x and aug_t for teacher forcing + clean_x = [clean_latents[i] for i in range(clean_latents.size(0))] + aug_t = torch.zeros(bsz, device=accelerator.device, dtype=torch.int64) + else: + # Standard causal mask + accelerator.unwrap_model(generator_transformer3d).create_block_mask_for_training( + num_frames=num_frames, + frame_seqlen=frame_seqlen_bm, + num_frame_per_block=args.num_frame_per_block, + independent_first_frame=args.independent_first_frame, + device=accelerator.device + ) + clean_x = None + aug_t = None + + for index, current_timestep in enumerate(denoising_step_list): + is_final_step = (index == final_step_index) + timestep = torch.full( + generator_noise.shape[:1], + current_timestep, + device=generator_noise.device, + dtype=torch.int64 + ) + + with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + context_manager = torch.no_grad() if not is_final_step else contextlib.nullcontext() + + with context_manager: + # Convert to list format for transformer + generator_noise_list = [generator_noise[i] for i in range(bsz)] + clean_x_list = [clean_latents[i] for i in range(bsz)] if clean_x is not None else None + + # Use block_mask for causal training (一次性处理整个视频) + generator_pred = generator_transformer3d( + x=generator_noise_list, + context=prompt_embeds, + t=timestep, + seq_len=seq_len, + y=inpaint_latents if args.train_mode != "normal" else None, + clip_fea=clip_context if args.train_mode != "normal" else None, + clean_x=clean_x_list, + aug_t=aug_t, + ) + generator_pred = convert_flow_pred_to_x0( + scheduler=noise_scheduler, + flow_pred=generator_pred, + xt=generator_noise, + timestep=timestep + ) + + if is_final_step: + break + + next_timestep = denoising_step_list[index + 1] * torch.ones( + generator_noise.shape[:1], dtype=torch.long, device=generator_noise.device + ) + generator_noise = add_noise( + generator_pred, + torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng), + next_timestep + ) + + # Common code for both KV cache and block mask training + indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu() + generator_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device) + generator_denoised_input = add_noise( + generator_pred, + torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng), + generator_timestep + ).detach().to(accelerator.device, dtype=weight_dtype) + + # Compute fake score + with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device), torch.no_grad(): + fake_score_main_cond = fake_score_transformer3d( + x=generator_denoised_input, + context=prompt_embeds, + t=generator_timestep, + seq_len=seq_len, + y=inpaint_latents if args.train_mode != "normal" else None, + clip_fea=clip_context if args.train_mode != "normal" else None, + ) + fake_score_main_cond = convert_flow_pred_to_x0( + scheduler=noise_scheduler, + flow_pred=fake_score_main_cond, + xt=generator_denoised_input, + timestep=generator_timestep + ) + + if args.fake_guidance_scale != 0.0: + fake_score_main_uncond = fake_score_transformer3d( + x=generator_denoised_input, + context=neg_prompt_embeds, + t=generator_timestep, + seq_len=seq_len, + y=inpaint_latents if args.train_mode != "normal" else None, + clip_fea=clip_context if args.train_mode != "normal" else None, + ) + fake_score_main_uncond = convert_flow_pred_to_x0( + scheduler=noise_scheduler, + flow_pred=fake_score_main_uncond, + xt=generator_denoised_input, + timestep=generator_timestep + ) + fake_score_main = fake_score_main_uncond + ( + fake_score_main_cond - fake_score_main_uncond + ) * args.fake_guidance_scale + else: + fake_score_main = fake_score_main_cond + + # Compute real score + real_score_main_cond = real_score_transformer3d( + x=generator_denoised_input, + context=prompt_embeds, + t=generator_timestep, + seq_len=seq_len, + y=inpaint_latents if args.train_mode != "normal" else None, + clip_fea=clip_context if args.train_mode != "normal" else None, + ) + real_score_main_cond = convert_flow_pred_to_x0( + scheduler=noise_scheduler, + flow_pred=real_score_main_cond, + xt=generator_denoised_input, + timestep=generator_timestep + ) + + real_score_main_uncond = real_score_transformer3d( + x=generator_denoised_input, + context=neg_prompt_embeds, + t=generator_timestep, + seq_len=seq_len, + y=inpaint_latents if args.train_mode != "normal" else None, + clip_fea=clip_context if args.train_mode != "normal" else None, + ) + real_score_main_uncond = convert_flow_pred_to_x0( + scheduler=noise_scheduler, + flow_pred=real_score_main_uncond, + xt=generator_denoised_input, + timestep=generator_timestep + ) + + real_score_main = real_score_main_uncond + ( + real_score_main_cond - real_score_main_uncond + ) * args.real_guidance_scale + + # DMD loss + fake_to_real_grad = fake_score_main - real_score_main + generator_to_real_norm = generator_pred - real_score_main + normalizer = torch.abs(generator_to_real_norm).mean(dim=[1, 2, 3, 4], keepdim=True) + fake_to_real_grad = fake_to_real_grad / normalizer + fake_to_real_grad = torch.nan_to_num(fake_to_real_grad) + + # Apply gradient mask: only compute loss on unmasked elements + # (matches Self-Forcing dmd.py: F.mse_loss(x[mask], target[mask])) + if score_mask is not None: + dmd_loss = 0.5 * F.mse_loss( + generator_pred.double()[score_mask], + (generator_pred.double() - fake_to_real_grad.double()).detach()[score_mask], + reduction="mean" + ) + else: + dmd_loss = 0.5 * F.mse_loss( + generator_pred.double(), + (generator_pred.double() - fake_to_real_grad.double()).detach(), + reduction="mean" + ) + + avg_dmd_loss = accelerator.gather(dmd_loss.repeat(args.train_batch_size)).mean() + train_dmd_loss += avg_dmd_loss.item() / args.gradient_accumulation_steps + + if args.low_vram: + real_score_transformer3d = real_score_transformer3d.to("cpu") + fake_score_transformer3d = fake_score_transformer3d.to("cpu") + torch.cuda.empty_cache() + + accelerator.backward(dmd_loss) + if accelerator.sync_gradients: + accelerator.clip_grad_norm_(trainable_params, args.max_grad_norm) + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + + if args.low_vram: + fake_score_transformer3d = fake_score_transformer3d.to(accelerator.device) + torch.cuda.empty_cache() + + with accelerator_fake_score_transformer3d.accumulate(fake_score_transformer3d): + # --- Fake Critic Denoising Loss --- + + if args.use_kv_cache_training: + # KV cache mode: block-by-block generation + fake_score_critic_noise = torch.randn(target_shape, device=accelerator.device, generator=torch_rng, dtype=weight_dtype) + + # Calculate frame_seq_length + frame_seq_length = (target_shape[3] * target_shape[4]) // (patch_h * patch_w) + + # Determine block structure (variable-length, mirrors generator branch) + if not args.independent_first_frame: + max_num_blocks_critic = num_frames // args.num_frame_per_block + min_num_blocks_critic = args.score_num_frames // args.num_frame_per_block + else: + max_num_blocks_critic = (num_frames - 1) // args.num_frame_per_block + if args.score_num_frames > 1: + min_num_blocks_critic = (args.score_num_frames - 1) // args.num_frame_per_block + else: + min_num_blocks_critic = 0 + + # Random sample number of blocks (mirrors generator's variable-length training) + if args.min_length_prob_bias > 0.0 and max_num_blocks_critic > min_num_blocks_critic: + num_options = max_num_blocks_critic - min_num_blocks_critic + 1 + bias = min(args.min_length_prob_bias, 0.99) + remaining_prob = (1.0 - bias) / (num_options - 1) + probs = [remaining_prob] * num_options + probs[0] = bias # min_num_blocks_critic gets the bias + probs_tensor = torch.tensor(probs, device=accelerator.device) + block_indices = torch.multinomial(probs_tensor, 1, generator=torch_rng) + num_generated_blocks_critic = (min_num_blocks_critic + block_indices).item() + else: + num_generated_blocks_critic = torch.randint( + min_num_blocks_critic, max_num_blocks_critic + 1, (1,), + generator=torch_rng, device=accelerator.device + ).item() + if dist.is_initialized(): + _sync = torch.tensor([num_generated_blocks_critic], device=accelerator.device) + dist.broadcast(_sync, src=0) + num_generated_blocks_critic = _sync.item() + + all_num_frames = [args.num_frame_per_block] * num_generated_blocks_critic + if args.independent_first_frame: + all_num_frames = [1] + all_num_frames + num_generated_frames_critic = sum(all_num_frames) + + # Initialize KV cache + num_layers = generator_transformer3d.config.num_layers + num_heads = generator_transformer3d.config.num_heads + head_dim = generator_transformer3d.config.dim // num_heads + text_len = 512 + + critic_kv_cache = initialize_kv_cache_for_training( + batch_size=bsz, + num_frames=num_frames, + frame_seq_length=frame_seq_length, + num_layers=num_layers, + num_heads=num_heads, + head_dim=head_dim, + dtype=weight_dtype, + device=accelerator.device + ) + + critic_crossattn_cache = initialize_crossattn_cache_for_training( + batch_size=bsz, + text_len=text_len, + num_layers=num_layers, + num_heads=num_heads, + head_dim=head_dim, + dtype=weight_dtype, + device=accelerator.device + ) + + current_start_frame = 0 + num_input_frames = 0 + output_pred = torch.zeros_like(fake_score_critic_noise) + + # Decide whether to use teacher forcing for this video + use_teacher_forcing_step = ( + args.use_teacher_forcing and + torch.rand(1, generator=torch_rng, device=accelerator.device).item() < args.teacher_forcing_prob + ) + + # Same exit step across all blocks (matches original Self-Forcing default) + num_denoising_steps = len(denoising_step_list) + final_step_index = generate_and_sync_list(num_denoising_steps, device=accelerator.device)[0] + + for block_idx, current_num_frames in enumerate(all_num_frames): + start_idx = current_start_frame - num_input_frames + end_idx = start_idx + current_num_frames + noisy_input = fake_score_critic_noise[:, :, start_idx:end_idx] + + for index, current_timestep in enumerate(denoising_step_list): + is_final_step = (index == final_step_index) + timestep = torch.full( + [bsz, current_num_frames], + current_timestep, + device=noisy_input.device, + dtype=torch.int64 + ) + + context_manager = torch.no_grad() + + with context_manager: + noisy_input_list = [noisy_input[i] for i in range(bsz)] + + # Use full seq_len (consistent with inference code) + full_seq_len = frame_seq_length * num_frames + + fake_score_denoised_pred_block = generator_transformer3d( + x=noisy_input_list, + context=prompt_embeds, + t=timestep, + seq_len=full_seq_len, + kv_cache=critic_kv_cache, + crossattn_cache=critic_crossattn_cache, + current_start=current_start_frame * frame_seq_length, + cache_start=None, + y=inpaint_latents if args.train_mode != "normal" else None, + clip_fea=clip_context if args.train_mode != "normal" else None, + ) + + # Stack list output to tensor: [B, C, F, H, W] + if isinstance(fake_score_denoised_pred_block, list): + fake_score_denoised_pred_block = torch.stack(fake_score_denoised_pred_block, dim=0) + + fake_score_denoised_pred_block = convert_flow_pred_to_x0( + scheduler=noise_scheduler, + flow_pred=fake_score_denoised_pred_block, + xt=noisy_input, + timestep=timestep[:, 0] + ) + + if is_final_step: + break + + next_timestep = denoising_step_list[index + 1] * torch.ones( + bsz, dtype=torch.long, device=noisy_input.device + ) + noisy_input = add_noise( + fake_score_denoised_pred_block, + torch.randn(fake_score_denoised_pred_block.shape, dtype=fake_score_denoised_pred_block.dtype, device=fake_score_denoised_pred_block.device, generator=torch_rng), + next_timestep + ) + + output_pred[:, :, current_start_frame:current_start_frame + current_num_frames] = fake_score_denoised_pred_block + + # Update KV cache with clean context (consistent with inference: feed denoised_pred directly) + if block_idx < len(all_num_frames) - 1: + context_timestep = torch.ones([bsz, current_num_frames], device=accelerator.device, dtype=torch.int64) * args.context_noise + + # Use clean latents for teacher forcing, otherwise use denoised prediction directly + if use_teacher_forcing_step and clean_latents is not None: + context_input = clean_latents[:, :, start_idx:end_idx] + else: + context_input = fake_score_denoised_pred_block + + context_input_list = [context_input[i] for i in range(bsz)] + + # Use full seq_len (consistent with inference code) + full_seq_len = frame_seq_length * num_frames + + with torch.no_grad(): + generator_transformer3d( + x=context_input_list, + context=prompt_embeds, + t=context_timestep, + seq_len=full_seq_len, + kv_cache=critic_kv_cache, + crossattn_cache=critic_crossattn_cache, + current_start=current_start_frame * frame_seq_length, + cache_start=None, + y=inpaint_latents if args.train_mode != "normal" else None, + clip_fea=clip_context if args.train_mode != "normal" else None, + ) + + current_start_frame += current_num_frames + + fake_score_denoised_pred_full = output_pred[:, :, :num_generated_frames_critic] + + # Slice for critic score: last score_num_frames frames + if num_generated_frames_critic > args.score_num_frames: + fake_score_denoised_pred, critic_score_num_frames, _ = slice_for_score( + fake_score_denoised_pred_full, vae, weight_dtype, + score_num_frames=args.score_num_frames, + independent_first_frame=args.independent_first_frame, + ) + else: + fake_score_denoised_pred = fake_score_denoised_pred_full + critic_score_num_frames = num_generated_frames_critic + + seq_len = frame_seq_length * critic_score_num_frames + + else: + with torch.no_grad(): + # Block mask mode: use flex attention to process entire video at once + + patch_h_bm, patch_w_bm = accelerator.unwrap_model(generator_transformer3d).config.patch_size[1:] + frame_seqlen_bm = (height * width) // (patch_h_bm * patch_w_bm) + seq_len = frame_seqlen_bm * num_frames + + fake_score_critic_noise = torch.randn(target_shape, device=accelerator.device, generator=torch_rng, dtype=weight_dtype) + num_denoising_steps = len(denoising_step_list) + final_step_index = generate_and_sync_list(num_denoising_steps, device=fake_score_critic_noise.device)[0] + + # Decide whether to use teacher forcing for this step + use_teacher_forcing_step = ( + args.use_teacher_forcing and + torch.rand(1, generator=torch_rng, device=accelerator.device).item() < args.teacher_forcing_prob + ) + + # Create appropriate block mask based on teacher forcing decision + if use_teacher_forcing_step and clean_latents is not None: + # Teacher forcing: clean + noisy sequence mask + accelerator.unwrap_model(generator_transformer3d).create_teacher_forcing_mask( + device=accelerator.device, + num_frames=num_frames, + frame_seqlen=frame_seqlen_bm, + num_frame_per_block=args.num_frame_per_block, + ) + clean_x = [clean_latents[i] for i in range(clean_latents.size(0))] + aug_t = torch.zeros(bsz, device=accelerator.device, dtype=torch.int64) + else: + # Standard causal mask + accelerator.unwrap_model(generator_transformer3d).create_block_mask_for_training( + num_frames=num_frames, + frame_seqlen=frame_seqlen_bm, + num_frame_per_block=args.num_frame_per_block, + independent_first_frame=args.independent_first_frame, + device=accelerator.device + ) + clean_x = None + aug_t = None + + for index, current_timestep in enumerate(denoising_step_list): + is_final_step = (index == final_step_index) + timestep = torch.full( + fake_score_critic_noise.shape[:1], + current_timestep, + device=fake_score_critic_noise.device, + dtype=torch.int64 + ) + + + with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + # Convert to list format for transformer + fake_score_critic_noise_list = [fake_score_critic_noise[i] for i in range(bsz)] + clean_x_list = [clean_latents[i] for i in range(bsz)] if clean_x is not None else None + + fake_score_denoised_pred = generator_transformer3d( + x=fake_score_critic_noise_list, + context=prompt_embeds, + t=timestep, + seq_len=seq_len, + y=inpaint_latents if args.train_mode != "normal" else None, + clip_fea=clip_context if args.train_mode != "normal" else None, + clean_x=clean_x_list, + aug_t=aug_t, + ) + fake_score_denoised_pred = convert_flow_pred_to_x0( + scheduler=noise_scheduler, + flow_pred=fake_score_denoised_pred, + xt=fake_score_critic_noise, + timestep=timestep + ) + + if is_final_step: + break + + next_timestep = denoising_step_list[index + 1] * torch.ones( + fake_score_critic_noise.shape[:1], + dtype=torch.long, + device=fake_score_critic_noise.device + ) + + fake_score_critic_noise = add_noise( + fake_score_denoised_pred, + torch.randn(fake_score_denoised_pred.shape, dtype=fake_score_denoised_pred.dtype, device=fake_score_denoised_pred.device, generator=torch_rng), + next_timestep + ) + + indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu() + critic_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device) + critic_noise = torch.randn(fake_score_denoised_pred.shape, dtype=fake_score_denoised_pred.dtype, device=fake_score_denoised_pred.device, generator=torch_rng) + + fake_score_denoised_input = add_noise( + fake_score_denoised_pred, + critic_noise, + critic_timestep + ) + + with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + fake_score_denoised_output = fake_score_transformer3d( + x=fake_score_denoised_input, + context=prompt_embeds, + t=critic_timestep, + seq_len=seq_len, + y=inpaint_latents if args.train_mode != "normal" else None, + clip_fea=clip_context if args.train_mode != "normal" else None, + ) + + def custom_mse_loss(noise_pred, target, weighting=None, threshold=50): + noise_pred = noise_pred.float() + target = target.float() + diff = noise_pred - target + mse_loss = F.mse_loss(noise_pred, target, reduction='none') + mask = (diff.abs() <= threshold).float() + masked_loss = mse_loss * mask + if weighting is not None: + masked_loss = masked_loss * weighting + final_loss = masked_loss.mean() + return final_loss + + denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred) + avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean() + train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps + + accelerator_fake_score_transformer3d.backward(denoising_loss) + if accelerator_fake_score_transformer3d.sync_gradients: + accelerator_fake_score_transformer3d.clip_grad_norm_(fake_trainable_params, args.max_grad_norm) + critic_optimizer.step() + fake_score_lr_scheduler.step() + critic_optimizer.zero_grad() + + if args.low_vram: + fake_score_transformer3d = fake_score_transformer3d.to(accelerator.device) + generator_transformer3d = generator_transformer3d.to(accelerator.device) + + # Checks if the accelerator has performed an optimization step behind the scenes + if accelerator.sync_gradients: + + if args.use_ema: + ema_generator_transformer3d.step(generator_transformer3d.parameters()) + progress_bar.update(1) + global_step += 1 + accelerator.log({"train_denoising_loss": train_denoising_loss, "train_dmd_loss": train_dmd_loss}, step=global_step) + train_dmd_loss = 0.0 + train_denoising_loss = 0.0 + + if global_step % args.checkpointing_steps == 0: + if args.use_deepspeed or args.use_fsdp or accelerator.is_main_process: + # _before_ saving state, check if this save would set us over the `checkpoints_total_limit` + if args.checkpoints_total_limit is not None: + checkpoints = os.listdir(args.output_dir) + checkpoints = [d for d in checkpoints if d.startswith("checkpoint")] + checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1])) + + # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints + if len(checkpoints) >= args.checkpoints_total_limit: + num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1 + removing_checkpoints = checkpoints[0:num_to_remove] + + logger.info( + f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints" + ) + logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}") + + for removing_checkpoint in removing_checkpoints: + removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint) + shutil.rmtree(removing_checkpoint) + + gc.collect() + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") + fake_score_save_path = os.path.join(save_path, "fake_score") + accelerator.save_state(save_path) + accelerator_fake_score_transformer3d.save_state(fake_score_save_path) + logger.info(f"Saved state to {save_path}") + + if args.validation_prompts is not None and global_step % args.validation_steps == 0: + if args.use_ema: + # Store the generator parameters temporarily and load the EMA parameters to perform inference. + ema_generator_transformer3d.store(generator_transformer3d.parameters()) + ema_generator_transformer3d.copy_to(generator_transformer3d.parameters()) + log_validation( + vae, + text_encoder, + tokenizer, + clip_image_encoder, + generator_transformer3d, + args, + config, + accelerator, + weight_dtype, + global_step, + ) + if args.use_ema: + # Switch back to the original generator parameters. + ema_generator_transformer3d.restore(generator_transformer3d.parameters()) + + logs = {"denoising_loss": denoising_loss.detach().item(), "dmd_loss": dmd_loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]} + progress_bar.set_postfix(**logs) + + if global_step >= args.max_train_steps: + break + + if args.validation_prompts is not None and epoch % args.validation_epochs == 0: + if args.use_ema: + # Store the generator parameters temporarily and load the EMA parameters to perform inference. + ema_generator_transformer3d.store(generator_transformer3d.parameters()) + ema_generator_transformer3d.copy_to(generator_transformer3d.parameters()) + log_validation( + vae, + text_encoder, + tokenizer, + clip_image_encoder, + generator_transformer3d, + args, + config, + accelerator, + weight_dtype, + global_step, + ) + if args.use_ema: + # Switch back to the original generator parameters. + ema_generator_transformer3d.restore(generator_transformer3d.parameters()) + + # Create the pipeline using the trained modules and save it. + accelerator.wait_for_everyone() + if args.use_ema: + # Copy EMA weights to generator on ALL processes before save_state (collective op). + ema_generator_transformer3d.copy_to(generator_transformer3d.parameters()) + if accelerator.is_main_process: + generator_transformer3d = unwrap_model(generator_transformer3d) + + if args.use_deepspeed or args.use_fsdp or accelerator.is_main_process: + gc.collect() + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") + fake_score_save_path = os.path.join(save_path, "fake_score") + accelerator.save_state(save_path) + accelerator_fake_score_transformer3d.save_state(fake_score_save_path) + if args.use_ema and accelerator.is_main_process: + ema_generator_transformer3d.save_pretrained(os.path.join(save_path, "transformer_ema")) + logger.info(f"Saved state to {save_path}") + + accelerator.end_training() + + +if __name__ == "__main__": + main() diff --git a/scripts/wan2.1_causal_forcing/train_causal_dmd.sh b/scripts/wan2.1_causal_forcing/train_causal_dmd.sh new file mode 100755 index 00000000..b9ff9e99 --- /dev/null +++ b/scripts/wan2.1_causal_forcing/train_causal_dmd.sh @@ -0,0 +1,64 @@ +export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/" +export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B" +export DATASET_NAME="datasets/internal_datasets/" +export DATASET_META_NAME="datasets/internal_datasets/metadata.json" +export STAGE2_CKPT="output_dir_wan2.1_causal_forcing_ccd/checkpoint-5000/transformer/diffusion_pytorch_model.safetensors" +# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. +# export NCCL_IB_DISABLE=1 +# export NCCL_P2P_DISABLE=1 +NCCL_DEBUG=INFO + +# Causal-Forcing Stage 3: Distribution Matching Distillation (DMD), frame-wise 2-step variant. +# - --real_score_pretrained_model_name_or_path -> Wan2.1-T2V-14B non-causal teacher (DMD real_score). +# - --transformer_path -> Stage 2 CCD ckpt (generator/critic init). +# - num_frame_per_block=1 -> frame-wise; --denoising_step_indices_list 1000 500 -> 2-step DMD. +# - train_mode="normal" (default) -> TextDataset (prompt-only); generation shape from video_sample_* / fix_sample_size. +# - Note: DMD parser has no --shift; the flow scheduler shift is fixed inside the training loop. +accelerate launch --mixed_precision="bf16" --use_fsdp \ + --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \ + --fsdp_transformer_layer_cls_to_wrap=CasualWanAttentionBlock \ + --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT \ + --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False \ + scripts/wan2.1_causal_forcing/train_causal_dmd.py \ + --config_path="config/wan2.1/wan_civitai.yaml" \ + --pretrained_model_name_or_path=$MODEL_NAME \ + --real_score_pretrained_model_name_or_path=$REAL_SCORE_MODEL_NAME \ + --train_data_dir=$DATASET_NAME \ + --train_data_meta=$DATASET_META_NAME \ + --ode_transformer_path=$STAGE2_CKPT \ + --image_sample_size=640 \ + --video_sample_size=640 \ + --token_sample_size=640 \ + --fix_sample_size 480 832 \ + --video_sample_stride=2 \ + --video_sample_n_frames=81 \ + --train_batch_size=1 \ + --gradient_accumulation_steps=1 \ + --dataloader_num_workers=8 \ + --num_train_epochs=100 \ + --checkpointing_steps=200 \ + --learning_rate=2.0e-06 \ + --learning_rate_critic=2.0e-06 \ + --lr_scheduler="constant_with_warmup" \ + --lr_warmup_steps=100 \ + --seed=42 \ + --output_dir="output_dir_wan2.1_causal_forcing_dmd" \ + --gradient_checkpointing \ + --mixed_precision="bf16" \ + --adam_weight_decay=0.0 \ + --adam_beta1=0.0 \ + --adam_beta2=0.999 \ + --adam_epsilon=1e-10 \ + --vae_mini_batch=1 \ + --max_grad_norm=10.0 \ + --random_hw_adapt \ + --training_with_video_token_length \ + --enable_bucket \ + --num_frame_per_block=1 \ + --use_kv_cache_training \ + --denoising_step_indices_list 1000 500 \ + --real_guidance_scale=3.0 \ + --fake_guidance_scale=0.0 \ + --gen_update_interval=5 \ + --resume_from_checkpoint="latest" \ + --trainable_modules "." diff --git a/videox_fun/pipeline/__init__.py b/videox_fun/pipeline/__init__.py index fef32ce4..6c58ae50 100755 --- a/videox_fun/pipeline/__init__.py +++ b/videox_fun/pipeline/__init__.py @@ -43,34 +43,4 @@ WanI2VPipeline = WanFunInpaintPipeline Wan2_2FunPipeline = Wan2_2Pipeline -Wan2_2I2VPipeline = Wan2_2FunInpaintPipeline - -import importlib.util - -if importlib.util.find_spec("paifuser") is not None: - # --------------------------------------------------------------- # - # Sparse Attention - # --------------------------------------------------------------- # - from paifuser.ops import sparse_reset - - # Wan2.1 - WanFunInpaintPipeline.__call__ = sparse_reset(WanFunInpaintPipeline.__call__) - WanFunPipeline.__call__ = sparse_reset(WanFunPipeline.__call__) - WanFunControlPipeline.__call__ = sparse_reset(WanFunControlPipeline.__call__) - WanI2VPipeline.__call__ = sparse_reset(WanI2VPipeline.__call__) - WanPipeline.__call__ = sparse_reset(WanPipeline.__call__) - WanVacePipeline.__call__ = sparse_reset(WanVacePipeline.__call__) - - # Phantom - WanFunPhantomPipeline.__call__ = sparse_reset(WanFunPhantomPipeline.__call__) - - # Wan2.2 - Wan2_2FunInpaintPipeline.__call__ = sparse_reset(Wan2_2FunInpaintPipeline.__call__) - Wan2_2FunPipeline.__call__ = sparse_reset(Wan2_2FunPipeline.__call__) - Wan2_2FunControlPipeline.__call__ = sparse_reset(Wan2_2FunControlPipeline.__call__) - Wan2_2Pipeline.__call__ = sparse_reset(Wan2_2Pipeline.__call__) - Wan2_2I2VPipeline.__call__ = sparse_reset(Wan2_2I2VPipeline.__call__) - Wan2_2TI2VPipeline.__call__ = sparse_reset(Wan2_2TI2VPipeline.__call__) - Wan2_2S2VPipeline.__call__ = sparse_reset(Wan2_2S2VPipeline.__call__) - Wan2_2VaceFunPipeline.__call__ = sparse_reset(Wan2_2VaceFunPipeline.__call__) - Wan2_2AnimatePipeline.__call__ = sparse_reset(Wan2_2AnimatePipeline.__call__) \ No newline at end of file +Wan2_2I2VPipeline = Wan2_2FunInpaintPipeline \ No newline at end of file diff --git a/videox_fun/pipeline/pipeline_wan_self_forcing.py b/videox_fun/pipeline/pipeline_wan_self_forcing.py index d3a7bf8d..52d2a995 100644 --- a/videox_fun/pipeline/pipeline_wan_self_forcing.py +++ b/videox_fun/pipeline/pipeline_wan_self_forcing.py @@ -23,9 +23,17 @@ def stochastic_sampling_timesteps(num_inference_steps, shift, device, num_timesteps=1000): - """Official FlashHead timestep schedule with shift transform.""" + """Official FlashHead timestep schedule with shift transform. + + Hardcoded schedules match the CF reference for distilled few-step inference: + - 4-step: [1000, 750, 500, 250] — Stage 2 CCD (`causal_cd_framewise.yaml`) + - 2-step: [1000, 500] — Stage 3 DMD (`causal_forcing_dmd_framewise_2step.yaml`) + Anything else falls back to linspace. + """ if num_inference_steps == 4: timesteps = [1000, 750, 500, 250] + elif num_inference_steps == 2: + timesteps = [1000, 500] else: timesteps = np.linspace(num_timesteps, 1, num_inference_steps, dtype=np.float32).tolist() timesteps = torch.tensor(timesteps + [0.0], dtype=torch.float32, device=device) @@ -778,19 +786,20 @@ def __call__( denoised_pred.shape, dtype=denoised_pred.dtype, device=device, generator=generator ) else: - # Get current sigma for x0 conversion - sigma_t = self.scheduler.sigmas[step_idx] - - # Convert to x0: x0 = x_t - sigma_t * flow_pred - denoised_pred = noisy_input - sigma_t * flow_pred - - if step_idx < len(denoise_timesteps) - 1: - # Not the last step: add noise for next timestep - next_sigma = self.scheduler.sigmas[step_idx + 1] - local_noise = torch.randn(denoised_pred.shape, device=denoised_pred.device, dtype=denoised_pred.dtype, generator=generator) - noisy_input = (1 - next_sigma) * denoised_pred + next_sigma * local_noise - else: - noisy_input = denoised_pred + # Delegate to the scheduler's own step() so each sampler + # (Flow Euler, UniPC, DPM++) applies its correct + # multi-step formula. The previous "predict-x0 + resample + # fresh noise" hybrid did NOT match CF's UniPC reference + # (CF's CausalDiffusionInferencePipeline calls + # sample_scheduler.step(...)). For framewise AR rollouts + # this mismatch caused subtle drift on top of CF's own + # output for the same checkpoint. + t = denoise_timesteps[step_idx] + step_out = self.scheduler.step( + flow_pred, t, noisy_input, return_dict=False + ) + noisy_input = step_out[0] + denoised_pred = noisy_input progress_bar.update()