From d5d4e76ae7f7bdb9407827fef51e5f15a9cb6035 Mon Sep 17 00:00:00 2001 From: "huangkunzhe.hkz" Date: Tue, 20 Oct 2026 16:50:34 +0800 Subject: [PATCH] add the training script for taew2_2_super --- scripts/taehv/README_TRAIN_SUPER.md | 434 ++++++++++ scripts/taehv/README_TRAIN_SUPER_zh-CN.md | 419 ++++++++++ scripts/taehv/train_taehv_super.py | 949 ++++++++++++++++++++++ scripts/taehv/train_taehv_super.sh | 50 ++ videox_fun/data/dataset_video.py | 13 +- videox_fun/models/wan_tae.py | 4 + 6 files changed, 1866 insertions(+), 3 deletions(-) create mode 100644 scripts/taehv/README_TRAIN_SUPER.md create mode 100644 scripts/taehv/README_TRAIN_SUPER_zh-CN.md create mode 100644 scripts/taehv/train_taehv_super.py create mode 100644 scripts/taehv/train_taehv_super.sh diff --git a/scripts/taehv/README_TRAIN_SUPER.md b/scripts/taehv/README_TRAIN_SUPER.md new file mode 100644 index 0000000..26b6057 --- /dev/null +++ b/scripts/taehv/README_TRAIN_SUPER.md @@ -0,0 +1,434 @@ +# TAE Super (Tiny AutoEncoder) Training Guide + +This document provides a complete workflow for training / fine-tuning the Wan2.2 +Super decoder (`AutoencoderTinyWan`), including environment setup, data +preparation, training, and using the trained checkpoint for inference. + +> **Note**: Training uses the frozen official Wan2.2 encoder and updates only the +> Super decoder's approximately 19M parameters. The tiny encoder is frozen and is +> not used in the training forward pass. +> +> ```text +> x (video, [-1, 1]) +> -> official Wan2.2 encoder (frozen) : z = vae.encode(x).latent_dist.mode() +> -> Super decoder : x_hat = tae(z).sample +> loss = 0.1 L1 + 0.1 LPIPS + 0.02 DINO + 1.0 grid +> ``` +> +> The losses compare reconstruction and input after mapping RGB to [0,1]. The +> decoder learns to consume the official Wan latent space directly. + +The script supports the Wan2.2 5B Super model: + +| Family | Latent | Official full VAE | `--config_path` | Models | +|--------|--------|-------------------|-----------------|--------| +| taew2_2_super | 48ch, patch_size=2 | `AutoencoderKLWan3_8` | `config/wan2.2/wan_civitai_5b.yaml` | Wan2.2 TI2V-5B / Fun-2.2VAE | + +**Differences from [train_taehv.py](train_taehv.py)**: + +| Item | `train_taehv.py` | `train_taehv_super.py` | +|------|-----------------|-----------------------| +| Reconstruction path | Tiny encoder → tiny decoder; official encoder supplies latent targets | Frozen official encoder → Super decoder | +| Trainable parameters | Tiny encoder and decoder by default, controlled by `--trainable_modules` | Super decoder only; official and tiny encoders are frozen | +| Losses | Pixel L1 + latent MSE | Pixel L1 + LPIPS + DINO + grid | + +--- + +## Table of Contents + +- [1. Environment Setup](#1-environment-setup) +- [2. Data Preparation](#2-data-preparation) + - [2.1 Quick Test Dataset](#21-quick-test-dataset) + - [2.2 Dataset Structure](#22-dataset-structure) + - [2.3 metadata.json Format](#23-metadatajson-format) + - [2.4 Relative vs Absolute Path Usage](#24-relative-vs-absolute-path-usage) +- [3. TAE Super Training](#3-tae-super-training) + - [3.1 Download Pretrained Model](#31-download-pretrained-model) + - [3.2 Quick Start](#32-quick-start) + - [3.3 Training Parameter Reference](#33-training-parameter-reference) + - [3.4 Training Validation](#34-training-validation) + - [3.5 Training Tips](#35-training-tips) + - [3.6 Further Tuning](#36-further-tuning) +- [4. Inference Testing](#4-inference-testing) + - [4.1 Checkpoint Layout](#41-checkpoint-layout) + - [4.2 Use the Trained TAE in Predict Scripts](#42-use-the-trained-tae-in-predict-scripts) +- [5. Additional Resources](#5-additional-resources) + +--- + +## 1. Environment Setup + +Run the commands from the VideoX-Fun repository root. Install dependencies into +your Python environment with pip. + +**Option 1: Using requirements.txt** + +```bash +pip install -r requirements.txt +pip install lpips torchvision av modelscope opencv-python +``` + +**Option 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 lpips torchvision av modelscope opencv-python +``` + +LPIPS's AlexNet backbone and DINOv2 download automatically when missing and reuse +the PyTorch cache. LPIPS's calibration weights ship with the package. For offline +DINOv2 loading, use `--dino_repo` and `--dino_weights`. + +> The script supports a single device and ordinary DDP through Accelerate. The +> official Wan VAE and the frozen LPIPS/DINOv2 feature networks remain on the GPU. + +--- + +## 2. Data Preparation + +### 2.1 Quick Test Dataset + +We provide a test dataset containing several training samples. + +```bash +modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo +``` + +### 2.2 Dataset Structure + +``` +📦 datasets/ +├── 📂 my_dataset/ +│ ├── 📂 train/ +│ │ ├── 📄 video001.mp4 +│ │ ├── 📄 video002.mp4 +│ │ └── 📄 ... +│ └── 📄 metadata.json +``` + +### 2.3 metadata.json Format + +**Relative Path Format** (example format): +```json +[ + { + "file_path": "train/video001.mp4", + "text": "A beautiful sunset over the ocean, golden hour lighting", + "type": "video", + "width": 1024, + "height": 1024 + } +] +``` + +**Absolute Path Format**: +```json +[ + { + "file_path": "/mnt/data/videos/sunset.mp4", + "text": "A beautiful sunset over the ocean", + "type": "video", + "width": 1024, + "height": 1024 + } +] +``` + +**Key Field Descriptions**: + +- `file_path`: Video path (relative or absolute path) +- `text`: Video description (not used by the TAE loss, kept for meta format compatibility) +- `type`: Data type, should be `"video"` +- `width` / `height`: Video width and height (recommended metadata). The Super + training entry uses a fixed resize/crop canvas. + +### 2.4 Relative vs Absolute Path Usage + +**Relative Path**: + +```bash +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" +``` + +**Absolute Path**: + +```bash +export DATASET_NAME="" +export DATASET_META_NAME="/mnt/data/metadata.json" +``` + +> 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. + +--- + +## 3. TAE Super Training + +### 3.1 Download Pretrained Model + +The training script needs the official Wan2.2 VAE weights, which ship inside the +model directory. Download the released Super weights for fine-tuning: + +```bash +mkdir -p models/Diffusion_Transformer + +modelscope download --model Wan-AI/Wan2.2-TI2V-5B --local_dir models/Diffusion_Transformer/Wan2.2-TI2V-5B + +wget -O models/Diffusion_Transformer/taew2_2_super.pth \ + https://github.com/madebyollin/taehv/raw/main/taew2_2_super.pth +``` + +LPIPS and DINOv2 weights are loaded separately by the loss implementation. + +### 3.2 Quick Start + +**Wan2.2 TI2V-5B / Fun-2.2VAE (taew2_2_super) Example**: + +The launcher uses a 720x1280 canvas, 25 training frames, 81 validation frames, +10,000 updates, and EMA. Set the GPU/process count with `accelerate config`; +four GPUs with batch size 1 and accumulation 1 give an effective batch of 4. + +```bash +export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-TI2V-5B" +export DATASET_NAME="${DATASET_NAME:-datasets/X-Fun-Videos-Demo}" +export DATASET_META_NAME="${DATASET_META_NAME:-$DATASET_NAME/metadata_add_width_height.json}" +export TAE_PATH="models/Diffusion_Transformer/taew2_2_super.pth" +NCCL_DEBUG=INFO + +accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \ + --config_path="config/wan2.2/wan_civitai_5b.yaml" \ + --pretrained_model_name_or_path="$MODEL_NAME" \ + --tae_path="$TAE_PATH" \ + --train_data_dir="$DATASET_NAME" \ + --train_data_meta="$DATASET_META_NAME" \ + --fix_sample_size 720 1280 \ + --video_sample_strides 1 2 3 \ + --video_sample_n_frames=25 \ + --validation_n_frames=81 \ + --train_batch_size=1 \ + --gradient_accumulation_steps=1 \ + --dataloader_num_workers=8 \ + --max_train_steps=10000 \ + --checkpointing_steps=1000 \ + --validation_steps=1000 \ + --validation_max_samples=0 \ + --learning_rate=5e-06 \ + --lr_scheduler="cosine" \ + --lr_warmup_steps=200 \ + --lr_min_ratio=0.2 \ + --l1_loss_weight=0.1 \ + --lpips_loss_weight=0.1 \ + --dino_loss_weight=0.02 \ + --grid_loss_weight=1.0 \ + --grid_period=16 \ + --feature_batch_size=13 \ + --use_ema \ + --ema_decay=0.999 \ + --selection_lpips_ratio=1.02 \ + --selection_grid_ratio=1.05 \ + --seed=42 \ + --output_dir="output_dir_taehv_w2.2_super" \ + --mixed_precision="bf16" \ + --adam_beta1=0.9 \ + --adam_beta2=0.99 \ + --adam_weight_decay=1e-2 \ + --adam_epsilon=1e-08 \ + --vae_mini_batch=1 \ + --max_grad_norm=1.0 +``` + +These settings are also provided in `scripts/taehv/train_taehv_super.sh`. Edit its +training arguments and run `bash scripts/taehv/train_taehv_super.sh`. + +**Train from scratch**: + +```bash +accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \ + --pretrained_model_name_or_path="$MODEL_NAME" \ + --train_data_dir="$DATASET_NAME" \ + --train_data_meta="$DATASET_META_NAME" \ + --from_scratch \ + --fix_sample_size 720 1280 \ + --video_sample_n_frames=25 \ + --video_sample_strides 1 2 3 \ + --max_train_steps=10000 \ + --learning_rate=5e-6 --lr_warmup_steps=200 --lr_min_ratio=0.2 \ + --use_ema --ema_decay=0.999 \ + --dataloader_num_workers=8 --checkpointing_steps=1000 --validation_steps=1000 \ + --output_dir="output_dir_taehv_w2.2_super_scratch" +``` + +Omit `--tae_path` when using `--from_scratch`. Only the decoder is trained; the +saved tiny encoder remains randomly initialized in a scratch run. + +**Resume interrupted training**: + +```bash +accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \ + --pretrained_model_name_or_path="$MODEL_NAME" \ + --train_data_dir="$DATASET_NAME" \ + --train_data_meta="$DATASET_META_NAME" \ + --fix_sample_size 720 1280 --video_sample_n_frames=25 \ + --video_sample_strides 1 2 3 --max_train_steps=10000 \ + --learning_rate=5e-6 --lr_warmup_steps=200 --lr_min_ratio=0.2 \ + --use_ema --ema_decay=0.999 \ + --dataloader_num_workers=8 --checkpointing_steps=1000 --validation_steps=1000 \ + --output_dir="output_dir_taehv_w2.2_super" \ + --resume_from_checkpoint=latest +``` + +Restore with the same settings used to create the checkpoint, including validation +metadata when supplied. `--max_train_steps` is the total update budget. Keep the +original schedule horizon. The checkpoint restores optimizer, scheduler, EMA, RNG, +step and batch cursor; changing workers or sampling settings can change the resumed +sample sequence. Use a new output directory for a new run. + +### 3.3 Training Parameter Reference + +Values below describe the launcher configuration; they are not all Python CLI defaults. + +| Parameter | Description | Example Value | +|-----------|-------------|---------------| +| `--config_path` | Official Wan VAE component configuration | `config/wan2.2/wan_civitai_5b.yaml` | +| `--pretrained_model_name_or_path` | Model directory containing the official VAE weights | `models/Diffusion_Transformer/Wan2.2-TI2V-5B` | +| `--tae_path` | Super weights, file or Diffusers directory; omit to initialize randomly | `models/Diffusion_Transformer/taew2_2_super.pth` | +| `--from_scratch` | Randomly initialize the Super decoder; excludes weight initialization and resume | - | +| `--train_data_dir` / `--train_data_meta` | Video root and JSON metadata array | `datasets/X-Fun-Videos-Demo` | +| `--fix_sample_size` | Fixed `[height, width]`, overriding square `--video_sample_size` | `720 1280` | +| `--video_sample_strides` | Choose one sampling stride per training clip | `1 2 3` | +| `--video_sample_n_frames` | Training clip length, must be 4k+1 | 25 | +| `--train_batch_size` / `--gradient_accumulation_steps` | Batch per device / gradient accumulation | 1 / 1 | +| `--vae_mini_batch` | Official VAE encoding mini-batch size | 1 | +| `--dataloader_num_workers` | DataLoader workers | 8 | +| `--max_train_steps` | Total optimizer updates; overrides `--num_train_epochs` | 10000 | +| `--checkpointing_steps` / `--checkpoints_total_limit` | Save interval / checkpoint count limit | 1000 / `None` | +| `--learning_rate` / `--lr_scheduler` | Peak learning rate / scheduler | `5e-6` / `cosine` | +| `--lr_warmup_steps` / `--lr_min_ratio` | Warmup updates / final-to-peak ratio for cosine decay | 200 / 0.2 | +| `--l1_loss_weight` / `--lpips_loss_weight` | Pixel L1 / LPIPS loss weights | 0.1 / 0.1 | +| `--dino_loss_weight` / `--grid_loss_weight` | DINO / periodic grid loss weights | 0.02 / 1.0 | +| `--feature_batch_size` / `--feature_checkpointing` | Perceptual frame batch / activation checkpointing | 13 / off | +| `--use_ema` / `--ema_decay` | Maintain decoder EMA for validation and selection | enabled / 0.999 | +| `--seed` / `--mixed_precision` | Random seed / mixed precision | 42 / `bf16` | +| `--output_dir` | Output directory | `output_dir_taehv_w2.2_super` | +| `--resume_from_checkpoint` | Restore a checkpoint directory or select the latest | `latest` | + +**Sample Size Configuration Guide**: + +- `--fix_sample_size 720 1280` resizes to cover the canvas and center crops while preserving aspect ratio. Without it, `--video_sample_size` selects a square canvas. +- A 25-frame clip at stride 3 needs at least 73 source frames; 81-frame validation needs at least 81. Short inputs shorten the batch to the shortest 4k+1 clip length, recorded as `frames` in the training log. +- Training uses random temporal starts. Validation starts at frame 0 with stride 1. Data loading retains the original random replacement/retry behavior on decode failures. + +**Loss Configuration Guide**: + +- The frozen official encoder supplies latents; only the Super decoder is optimized. +- LPIPS uses AlexNet; DINOv2 uses intermediate ViT-B/14 patch features. Frozen feature networks still pass reconstruction gradients to the decoder. +- Grid loss measures residual means by spatial phase with each frame/channel's DC removed. Inspect local crops alongside this metric. + +### 3.4 Training Validation + +Provide a separate JSON array through `--validation_data_meta` in the training +command, or add that argument to the launcher. Relative video paths resolve under +`--train_data_dir`, so use absolute video paths when validation data has another root. + +| Parameter | Description | Recommended Value | +|-----------|-------------|-------------------| +| `--validation_data_meta` | Held-out video metadata array | Your validation index | +| `--validation_steps` | Evaluate every N optimizer updates | 1000 | +| `--validation_n_frames` | Validation clip length, 4k+1 | 81 | +| `--validation_max_samples` | Validation clip cap; zero uses the complete index | 0 | + +Validation runs at initialization, periodically, and at the final step. It reports +PSNR, SSIM, LPIPS and grid, and saves the first sample's input/reconstruction preview +as `sample/step-N.mp4`. Metrics use floating-point output before video encoding, +averaging frames within a clip and then clips equally. + +With EMA enabled, validation and best-weight selection use EMA; otherwise they +use raw weights. Selection maximizes PSNR subject to LPIPS/grid ratios relative +to the first validation, controlled by `--selection_lpips_ratio` and +`--selection_grid_ratio`. The initialization is also a candidate. Best inference +weights go to `best/taehv`; resumable checkpoints retain raw training weights. + +For a short execution check, use the official demo: + +```bash +accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \ + --pretrained_model_name_or_path="$MODEL_NAME" --tae_path="$TAE_PATH" \ + --train_data_dir="datasets/X-Fun-Videos-Demo" \ + --train_data_meta="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" \ + --validation_data_meta="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" \ + --fix_sample_size 64 96 --video_sample_n_frames=5 --validation_n_frames=5 \ + --video_sample_strides 1 --max_train_steps=2 --lr_warmup_steps=0 \ + --checkpointing_steps=1 --validation_steps=1 --validation_max_samples=1 \ + --dataloader_num_workers=0 --use_ema --ema_decay=0.999 \ + --output_dir="output_dir_taehv_w2.2_super_smoke" +``` + +Sharing the demo across training/validation checks execution only. Use independent +video sources for actual experiments, and use native-resolution data to assess +720p quality. + +### 3.5 Training Tips + +- **Warm-start vs. from scratch**: Fine-tune from `taew2_2_super.pth` or a previous Super directory; use `--from_scratch` for a new decoder. Initialization, dataset and budget all affect the result. +- **Loss balance**: Start with the four-loss settings in the launcher and inspect both reconstruction metrics and native-resolution crops. +- **Memory**: Lower the training frame count, canvas size or feature batch size, or enable `--feature_checkpointing` if needed. Perceptual networks are frozen but their prediction branches still need activation memory. +- **EMA**: The launcher enables EMA. `best/taehv` holds selected inference weights; `checkpoint-N/taehv` holds raw weights for training recovery. +- **Distributed training**: Use Accelerate's single-device or ordinary DDP configuration. Global batch equals device count times per-device batch times accumulation. + +### 3.6 Further Tuning + +- Try longer temporal context and check both the initial frames and later frames, accounting for the extra compute. +- Explore pixel MSE or official-decoder output supervision, balancing pixel accuracy, perceptual quality and teacher compute cost. +- Adjust LPIPS/grid weights with controlled comparisons, checking fine textures, text and local grid artifacts alongside average metrics. + +These are possible extensions; the current training entry uses the four-loss recipe above. + +--- + +## 4. Inference Testing + +### 4.1 Checkpoint Layout + +Each resumable checkpoint is written as `output_dir/checkpoint-{step}/`: + +```text +output_dir_taehv_w2.2_super/ + args-step-N.yaml # resolved invocation record + train.jsonl # total loss, four terms, LR and actual frame count + validation.jsonl # metrics and selection decisions, when validation is enabled + best/taehv/ # selected inference weights, when validation is enabled + checkpoint-1000/ + taehv/ # raw Super weights and config.json + optimizer.bin + custom_checkpoint_*.pkl # scheduler and optional EMA state + random_states_*.pkl + trainer_state.json + sample/ # validation videos + logs/ # TensorBoard logs +``` + +The `taehv` directories use Diffusers `save_pretrained` format and load through +`AutoencoderTinyWan.from_pretrained`. `args-step-N.yaml` records the invocation; +training arguments are configured directly in the shell launcher. + +### 4.2 Use the Trained TAE in Predict Scripts + +Point `tae_path` in a compatible Wan2.2 5B TAE predict script to the selected +`best/taehv` directory: + +| Script | Family | +|--------|--------| +| `examples/wan2.2/predict_ti2v_tae.py` | taew2_2 | +| `examples/wan2.2_fun/predict_t2v_2.2vae_tae.py` | taew2_2 | +| `examples/wan2.2_fun/predict_i2v_2.2vae_tae.py` | taew2_2 | + +```python +# e.g. in examples/wan2.2_fun/predict_t2v_2.2vae_tae.py +tae_path = "output_dir_taehv_w2.2_super/best/taehv" +``` + +The trained decoder consumes official Wan2.2 5B latents. A scratch-trained +checkpoint has no trained tiny encoder; retain the official Wan encoder for +image/video conditioning when using such a checkpoint. diff --git a/scripts/taehv/README_TRAIN_SUPER_zh-CN.md b/scripts/taehv/README_TRAIN_SUPER_zh-CN.md new file mode 100644 index 0000000..ddb8a97 --- /dev/null +++ b/scripts/taehv/README_TRAIN_SUPER_zh-CN.md @@ -0,0 +1,419 @@ +# TAE Super（Tiny AutoEncoder）训练指南 + +本文档提供 Wan2.2 Super decoder（`AutoencoderTinyWan`）训练/微调的完整流程， +包括环境配置、数据准备、训练，以及将训练好的 checkpoint 用于推理。 + +> **说明**：训练使用冻结的官方 Wan2.2 encoder，只更新 Super decoder 约 19M 参数。 +> Tiny encoder 冻结，不参与训练前向。 +> +> ```text +> x (video, [-1, 1]) +> -> official Wan2.2 encoder (frozen) : z = vae.encode(x).latent_dist.mode() +> -> Super decoder : x_hat = tae(z).sample +> loss = 0.1 L1 + 0.1 LPIPS + 0.02 DINO + 1.0 grid +> ``` +> +> 损失在重建与输入 RGB 映射到 [0,1] 后计算，decoder 直接学习解码官方 Wan latent。 + +脚本支持 Wan2.2 5B Super 模型： + +| 家族 | Latent | 官方完整 VAE | `--config_path` | 适用模型 | +|------|--------|--------------|-----------------|----------| +| taew2_2_super | 48ch, patch_size=2 | `AutoencoderKLWan3_8` | `config/wan2.2/wan_civitai_5b.yaml` | Wan2.2 TI2V-5B / Fun-2.2VAE | + +**与 [train_taehv.py](train_taehv.py) 的差异**： + +| 对比项 | `train_taehv.py` | `train_taehv_super.py` | +|--------|-----------------|-----------------------| +| 重建路径 | 小 encoder → 小 decoder，官方 encoder 提供 latent 监督 | 冻结官方 encoder → Super decoder | +| 更新参数 | 默认训练小 encoder 和 decoder，由 `--trainable_modules` 控制 | 仅 Super decoder；官方和小 encoder 均冻结 | +| 损失 | 像素 L1 + latent MSE | 像素 L1 + LPIPS + DINO + grid | + +--- + +## 目录 + +- [一、环境配置](#一环境配置) +- [二、数据准备](#二数据准备) + - [2.1 快速测试数据集](#21-快速测试数据集) + - [2.2 数据集结构](#22-数据集结构) + - [2.3 metadata.json 格式](#23-metadatajson-格式) + - [2.4 相对路径与绝对路径使用方案](#24-相对路径与绝对路径使用方案) +- [三、TAE Super 训练](#三tae-super-训练) + - [3.1 下载预训练模型](#31-下载预训练模型) + - [3.2 快速开始](#32-快速开始) + - [3.3 训练常用参数解析](#33-训练常用参数解析) + - [3.4 训练验证](#34-训练验证) + - [3.5 训练技巧](#35-训练技巧) + - [3.6 后续调优方向](#36-后续调优方向) +- [四、推理测试](#四推理测试) + - [4.1 Checkpoint 目录结构](#41-checkpoint-目录结构) + - [4.2 在 Predict 脚本中使用训练好的 TAE](#42-在-predict-脚本中使用训练好的-tae) +- [五、更多资源](#五更多资源) + +--- + +## 一、环境配置 + +以下命令在 VideoX-Fun 仓库根目录执行，使用 pip 将依赖安装到当前 Python 环境。 + +**方式 1：使用 requirements.txt** + +```bash +pip install -r requirements.txt +pip install lpips torchvision av modelscope opencv-python +``` + +**方式 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 lpips torchvision av modelscope opencv-python +``` + +LPIPS 的 AlexNet 主干和 DINOv2 在缺少缓存时自动下载，后续复用 PyTorch 缓存； +LPIPS 校准权重随安装包提供。离线加载 DINOv2 时，可使用 `--dino_repo` 和 `--dino_weights`。 + +> 脚本通过 Accelerate 支持单设备与普通 DDP。官方 Wan VAE 和冻结的 LPIPS/DINOv2 +> 特征网络均保留在 GPU 上。 + +--- + +## 二、数据准备 + +### 2.1 快速测试数据集 + +我们提供了一个测试的数据集，其中包含若干训练数据。 + +```bash +modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo +``` + +### 2.2 数据集结构 + +``` +📦 datasets/ +├── 📂 my_dataset/ +│ ├── 📂 train/ +│ │ ├── 📄 video001.mp4 +│ │ ├── 📄 video002.mp4 +│ │ └── 📄 ... +│ └── 📄 metadata.json +``` + +### 2.3 metadata.json 格式 + +**相对路径格式**（示例格式）： +```json +[ + { + "file_path": "train/video001.mp4", + "text": "A beautiful sunset over the ocean, golden hour lighting", + "type": "video", + "width": 1024, + "height": 1024 + } +] +``` + +**绝对路径格式**： +```json +[ + { + "file_path": "/mnt/data/videos/sunset.mp4", + "text": "A beautiful sunset over the ocean", + "type": "video", + "width": 1024, + "height": 1024 + } +] +``` + +**关键字段说明**： + +- `file_path`：视频路径（相对或绝对路径） +- `text`：视频描述（TAE 损失不使用，仅为 meta 格式兼容而保留） +- `type`：数据类型，固定为 `"video"` +- `width` / `height`：视频宽高（推荐保留）。Super 训练入口使用固定的缩放/裁剪画布。 + +### 2.4 相对路径与绝对路径使用方案 + +**相对路径**： + +```bash +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" +``` + +**绝对路径**： + +```bash +export DATASET_NAME="" +export DATASET_META_NAME="/mnt/data/metadata.json" +``` + +> 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。 + +--- + +## 三、TAE Super 训练 + +### 3.1 下载预训练模型 + +训练脚本需要随模型目录提供的官方 Wan2.2 VAE 权重。微调时同时下载发布版 Super 权重： + +```bash +mkdir -p models/Diffusion_Transformer + +modelscope download --model Wan-AI/Wan2.2-TI2V-5B --local_dir models/Diffusion_Transformer/Wan2.2-TI2V-5B + +wget -O models/Diffusion_Transformer/taew2_2_super.pth \ + https://github.com/madebyollin/taehv/raw/main/taew2_2_super.pth +``` + +LPIPS 与 DINOv2 权重由损失实现单独加载。 + +### 3.2 快速开始 + +**Wan2.2 TI2V-5B / Fun-2.2VAE（taew2_2_super）训练示例**： + +启动器采用 720x1280 画布、训练 25 帧、验证 81 帧、10,000 步和 EMA。 +使用 `accelerate config` 配置 GPU/进程数；四卡、每卡 batch=1、累积=1 时有效 batch 为 4。 + +```bash +export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-TI2V-5B" +export DATASET_NAME="${DATASET_NAME:-datasets/X-Fun-Videos-Demo}" +export DATASET_META_NAME="${DATASET_META_NAME:-$DATASET_NAME/metadata_add_width_height.json}" +export TAE_PATH="models/Diffusion_Transformer/taew2_2_super.pth" +NCCL_DEBUG=INFO + +accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \ + --config_path="config/wan2.2/wan_civitai_5b.yaml" \ + --pretrained_model_name_or_path="$MODEL_NAME" \ + --tae_path="$TAE_PATH" \ + --train_data_dir="$DATASET_NAME" \ + --train_data_meta="$DATASET_META_NAME" \ + --fix_sample_size 720 1280 \ + --video_sample_strides 1 2 3 \ + --video_sample_n_frames=25 \ + --validation_n_frames=81 \ + --train_batch_size=1 \ + --gradient_accumulation_steps=1 \ + --dataloader_num_workers=8 \ + --max_train_steps=10000 \ + --checkpointing_steps=1000 \ + --validation_steps=1000 \ + --validation_max_samples=0 \ + --learning_rate=5e-06 \ + --lr_scheduler="cosine" \ + --lr_warmup_steps=200 \ + --lr_min_ratio=0.2 \ + --l1_loss_weight=0.1 \ + --lpips_loss_weight=0.1 \ + --dino_loss_weight=0.02 \ + --grid_loss_weight=1.0 \ + --grid_period=16 \ + --feature_batch_size=13 \ + --use_ema \ + --ema_decay=0.999 \ + --selection_lpips_ratio=1.02 \ + --selection_grid_ratio=1.05 \ + --seed=42 \ + --output_dir="output_dir_taehv_w2.2_super" \ + --mixed_precision="bf16" \ + --adam_beta1=0.9 \ + --adam_beta2=0.99 \ + --adam_weight_decay=1e-2 \ + --adam_epsilon=1e-08 \ + --vae_mini_batch=1 \ + --max_grad_norm=1.0 +``` + +相同参数已写入 `scripts/taehv/train_taehv_super.sh`，也可修改其中的训练参数后， +运行 `bash scripts/taehv/train_taehv_super.sh`。 + +**从头训练**： + +```bash +accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \ + --pretrained_model_name_or_path="$MODEL_NAME" \ + --train_data_dir="$DATASET_NAME" \ + --train_data_meta="$DATASET_META_NAME" \ + --from_scratch \ + --fix_sample_size 720 1280 \ + --video_sample_n_frames=25 \ + --video_sample_strides 1 2 3 \ + --max_train_steps=10000 \ + --learning_rate=5e-6 --lr_warmup_steps=200 --lr_min_ratio=0.2 \ + --use_ema --ema_decay=0.999 \ + --dataloader_num_workers=8 --checkpointing_steps=1000 --validation_steps=1000 \ + --output_dir="output_dir_taehv_w2.2_super_scratch" +``` + +使用 `--from_scratch` 时不要传 `--tae_path`。只训练 decoder，保存的 tiny encoder +仍是随机初始化权重。 + +**中断后恢复训练**： + +```bash +accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \ + --pretrained_model_name_or_path="$MODEL_NAME" \ + --train_data_dir="$DATASET_NAME" \ + --train_data_meta="$DATASET_META_NAME" \ + --fix_sample_size 720 1280 --video_sample_n_frames=25 \ + --video_sample_strides 1 2 3 --max_train_steps=10000 \ + --learning_rate=5e-6 --lr_warmup_steps=200 --lr_min_ratio=0.2 \ + --use_ema --ema_decay=0.999 \ + --dataloader_num_workers=8 --checkpointing_steps=1000 --validation_steps=1000 \ + --output_dir="output_dir_taehv_w2.2_super" \ + --resume_from_checkpoint=latest +``` + +恢复时保持创建 checkpoint 时的配置，使用过的验证索引也需继续提供。 +`--max_train_steps` 表示总更新预算，应保留原 schedule 预算。Checkpoint 恢复 +optimizer、scheduler、EMA、RNG、step 和 batch 游标；改变 workers 或采样设置可能改变 +恢复后的样本顺序。新训练应使用新的输出目录。 + +### 3.3 训练常用参数解析 + +下表列出启动器的配置值，不代表所有参数的 Python CLI 默认值。 + +| 参数 | 说明 | 示例值 | +|------|------|--------| +| `--config_path` | 官方 Wan VAE 组件配置 | `config/wan2.2/wan_civitai_5b.yaml` | +| `--pretrained_model_name_or_path` | 包含官方 VAE 权重的模型目录 | `models/Diffusion_Transformer/Wan2.2-TI2V-5B` | +| `--tae_path` | Super 权重文件或 Diffusers 目录；不传则随机初始化 | `models/Diffusion_Transformer/taew2_2_super.pth` | +| `--from_scratch` | 随机初始化 Super decoder；不能与加载权重或恢复组合 | - | +| `--train_data_dir` / `--train_data_meta` | 视频根目录与 JSON 元数据数组 | `datasets/X-Fun-Videos-Demo` | +| `--fix_sample_size` | 固定 `[高, 宽]`，覆盖正方形 `--video_sample_size` | `720 1280` | +| `--video_sample_strides` | 每个训练片段随机选择一个采样 stride | `1 2 3` | +| `--video_sample_n_frames` | 训练片段帧数，必须为 4k+1 | 25 | +| `--train_batch_size` / `--gradient_accumulation_steps` | 每卡 batch / 梯度累积 | 1 / 1 | +| `--vae_mini_batch` | 官方 VAE 编码 mini-batch | 1 | +| `--dataloader_num_workers` | DataLoader workers | 8 | +| `--max_train_steps` | 总优化器更新次数，覆盖 `--num_train_epochs` | 10000 | +| `--checkpointing_steps` / `--checkpoints_total_limit` | 保存间隔 / checkpoint 数量上限 | 1000 / `None` | +| `--learning_rate` / `--lr_scheduler` | 峰值学习率 / 调度器 | `5e-6` / `cosine` | +| `--lr_warmup_steps` / `--lr_min_ratio` | 预热步数 / cosine 终点与峰值之比 | 200 / 0.2 | +| `--l1_loss_weight` / `--lpips_loss_weight` | 像素 L1 / LPIPS 损失权重 | 0.1 / 0.1 | +| `--dino_loss_weight` / `--grid_loss_weight` | DINO / 周期格纹损失权重 | 0.02 / 1.0 | +| `--feature_batch_size` / `--feature_checkpointing` | 感知网络帧 batch / 激活 checkpointing | 13 / 关闭 | +| `--use_ema` / `--ema_decay` | 维护用于验证和选模的 decoder EMA | 开启 / 0.999 | +| `--seed` / `--mixed_precision` | 随机种子 / 混合精度 | 42 / `bf16` | +| `--output_dir` | 输出目录 | `output_dir_taehv_w2.2_super` | +| `--resume_from_checkpoint` | 恢复 checkpoint 目录或选择最新断点 | `latest` | + +**Sample Size 配置指南**： + +- `--fix_sample_size 720 1280` 等比例缩放覆盖画布后中心裁剪；未指定时，`--video_sample_size` 控制正方形画布。 +- 25 帧、stride=3 至少需要 73 个源视频帧；81 帧验证至少需要 81 帧。短视频会使 batch 按最短视频截为 4k+1 帧，实际长度记录在训练日志的 `frames` 字段。 +- 训练随机选择时间起点，验证固定从第 0 帧开始、stride=1。解码失败时沿用原有随机换视频重试逻辑。 + +**损失配置指南**： + +- 冻结官方 encoder 提供 latent，只优化 Super decoder。 +- LPIPS 使用 AlexNet，DINOv2 使用 ViT-B/14 中间 patch 特征；特征网络冻结，但仍将重建分支的梯度传回 decoder。 +- Grid loss 按空间相位统计残差均值，并去掉每帧每通道的 DC；应结合局部原生裁剪判断伪影。 + +### 3.4 训练验证 + +在训练命令中通过 `--validation_data_meta` 提供独立 JSON 数组，或将该参数加入启动器。 +相对视频路径基于 `--train_data_dir` 解析；验证数据位于其他根目录时，可使用绝对视频路径。 + +| 参数 | 说明 | 推荐值 | +|------|------|--------| +| `--validation_data_meta` | 独立验证集元数据数组 | 自己的验证索引 | +| `--validation_steps` | 每 N 次优化器更新执行验证 | 1000 | +| `--validation_n_frames` | 验证片段帧数，4k+1 | 81 | +| `--validation_max_samples` | 验证片段数量上限，0 表示全部 | 0 | + +初始化、周期节点和最后一步均执行验证，报告 PSNR、SSIM、LPIPS 和 grid，并将首个 +样本的输入/重建预览保存为 `sample/step-N.mp4`。指标使用视频编码前的浮点输出， +先对片段内帧平均，再对片段等权平均。 + +开启 EMA 时验证和最佳权重选择使用 EMA，否则使用 raw。选模在相对首次验证的 +LPIPS/grid 比例约束下最大化 PSNR，由 `--selection_lpips_ratio`、`--selection_grid_ratio` +控制。初始权重也是候选。最佳推理权重写入 `best/taehv`，可恢复 checkpoint 保留 raw +训练权重。 + +可使用官方 Demo 进行短流程检查： + +```bash +accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \ + --pretrained_model_name_or_path="$MODEL_NAME" --tae_path="$TAE_PATH" \ + --train_data_dir="datasets/X-Fun-Videos-Demo" \ + --train_data_meta="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" \ + --validation_data_meta="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" \ + --fix_sample_size 64 96 --video_sample_n_frames=5 --validation_n_frames=5 \ + --video_sample_strides 1 --max_train_steps=2 --lr_warmup_steps=0 \ + --checkpointing_steps=1 --validation_steps=1 --validation_max_samples=1 \ + --dataloader_num_workers=0 --use_ema --ema_decay=0.999 \ + --output_dir="output_dir_taehv_w2.2_super_smoke" +``` + +短测中训练和验证共用 Demo 只用于检查执行；正式实验应按视频来源隔离，评估 720p +质量时应使用原生分辨率数据。 + +### 3.5 训练技巧 + +- **热启动与从头训练**：从 `taew2_2_super.pth` 或已有 Super 目录微调；使用 `--from_scratch` 训练新 decoder。初始化、数据和预算都会影响结果。 +- **损失平衡**：先使用启动器中的四项损失配置，同时检查重建指标和原生分辨率裁剪。 +- **显存**：按需降低训练帧数、画布尺寸或特征 batch，或开启 `--feature_checkpointing`。感知网络虽然冻结，预测分支仍需要激活显存。 +- **EMA**：启动器默认开启 EMA。`best/taehv` 为选出的推理权重，`checkpoint-N/taehv` 为恢复训练使用的 raw 权重。 +- **分布式训练**：使用 Accelerate 单设备或普通 DDP 配置。有效 batch 为设备数、每卡 batch 和梯度累积的乘积。 + +### 3.6 后续调优方向 + +- 尝试更长的时间上下文，同时检查视频起始段和后续帧，并计入额外计算开销。 +- 探索像素 MSE 或官方 decoder 输出监督，平衡像素精度、感知质量和 teacher 计算成本。 +- 通过受控对照调整 LPIPS/grid 权重，结合细纹理、文字和局部格纹检查平均指标。 + +以上为可探索的扩展方向，当前训练入口使用前述四项损失配置。 + +--- + +## 四、推理测试 + +### 4.1 Checkpoint 目录结构 + +可恢复 checkpoint 写入 `output_dir/checkpoint-{step}/`： + +```text +output_dir_taehv_w2.2_super/ + args-step-N.yaml # 实际运行参数记录 + train.jsonl # 总损失、四项损失、LR 和实际帧数 + validation.jsonl # 提供验证索引时保存指标与选模结果 + best/taehv/ # 提供验证索引时保存最佳推理权重 + checkpoint-1000/ + taehv/ # raw Super 权重与 config.json + optimizer.bin + custom_checkpoint_*.pkl # scheduler 和可选 EMA 状态 + random_states_*.pkl + trainer_state.json + sample/ # 验证视频 + logs/ # TensorBoard 日志 +``` + +`taehv` 目录使用 Diffusers `save_pretrained` 格式，可由 +`AutoencoderTinyWan.from_pretrained` 加载。`args-step-N.yaml` 记录本次调用参数， +训练参数直接在 shell 启动器中配置。 + +### 4.2 在 Predict 脚本中使用训练好的 TAE + +将兼容的 Wan2.2 5B TAE predict 脚本中的 `tae_path` 指向选出的 `best/taehv` 目录： + +| 脚本 | 家族 | +|------|------| +| `examples/wan2.2/predict_ti2v_tae.py` | taew2_2 | +| `examples/wan2.2_fun/predict_t2v_2.2vae_tae.py` | taew2_2 | +| `examples/wan2.2_fun/predict_i2v_2.2vae_tae.py` | taew2_2 | + +```python +# 例如 examples/wan2.2_fun/predict_t2v_2.2vae_tae.py 中 +tae_path = "output_dir_taehv_w2.2_super/best/taehv" +``` + +训练好的 decoder 接收官方 Wan2.2 5B latent。从头训练的 checkpoint 没有训练好的 +小 encoder；此类 checkpoint 用于图像/视频条件生成时，应保留官方 Wan encoder。 diff --git a/scripts/taehv/train_taehv_super.py b/scripts/taehv/train_taehv_super.py new file mode 100644 index 0000000..72e8aa6 --- /dev/null +++ b/scripts/taehv/train_taehv_super.py @@ -0,0 +1,949 @@ +"""Super decoder training, copied and adapted from train_taehv.py. + +Frozen official Wan2.2 encoder -> trainable Super decoder, using the four +TinyVAE losses and the stage 4a B training recipe. +""" +#!/usr/bin/env python +# coding=utf-8 +# Copyright 2025 The VideoX-Fun 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 + +# ruff: noqa: E402 -- local imports follow the repository sys.path bootstrap. +import argparse +import json +import logging +import math +import os +import shutil +import sys + +import diffusers +import torch +from accelerate import Accelerator +from accelerate.logging import get_logger +from accelerate.utils import ProjectConfiguration, set_seed +from diffusers.optimization import get_scheduler +from diffusers.training_utils import EMAModel +from diffusers.utils import check_min_version +from omegaconf import OmegaConf +from torch import nn +from torch.nn import functional as F +from torch.utils.checkpoint import checkpoint +from torchvision import transforms + +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 ImageVideoSampler, RandomSampler, VideoDataset +from videox_fun.models import AutoencoderKLWan3_8, AutoencoderTinyWan +from videox_fun.utils.tqdm_bar import PauseAwareTqdm +from videox_fun.utils.utils import save_videos_grid + +check_min_version("0.18.0.dev0") +logger = get_logger(__name__, log_level="INFO") + + +# TinyVAE reconstruction losses: RGB [0,1], BCTHW, FP32 reductions. +def grid_loss_per_frame(pred, target, period=16): + """Reference: ZYT-World: A Real-Time Controllable World Model for + Closed-Loop Autonomous-Driving Simulation (https://arxiv.org/abs/2609.21712). + + Penalize NCHW residual phase means (16x16 by default) after per-frame/channel + DC removal to suppress periodic grid artifacts. + """ + n, c, h, w = pred.shape + h, w = h // period * period, w // period * period + if min(h, w) == 0: + raise ValueError("Images must be at least one grid period in size") + residual = pred.float()[..., :h, :w] - target.float()[..., :h, :w] + phase = residual.reshape(n, c, h // period, period, w // period, period).mean((2, 4)) + return (phase - phase.mean((-2, -1), keepdim=True)).abs().mean((1, 2, 3)) + +class TinyVAEReconstructionLoss(nn.Module): + def __init__( + self, + weights, + dino_repo=None, + dino_weights=None, + feature_batch_size=13, + feature_checkpointing=False, + grid_period=16, + lpips_model=None, + dino_model=None, + ): + super().__init__() + self.weights = { + k: float(weights[k]) for k in ("l1", "lpips", "dino", "grid") + } + if any(not math.isfinite(v) or v < 0 for v in self.weights.values()): + raise ValueError("Loss weights must be finite and nonnegative") + if not any(self.weights.values()) or min(feature_batch_size, grid_period) < 1: + raise ValueError("Need a positive objective, feature batch size and grid period") + if lpips_model is None: + import lpips + + lpips_model = lpips.LPIPS(net="alex", version="0.1", verbose=False) + self.lpips = lpips_model.eval().requires_grad_(False) + if self.weights["dino"] and dino_model is None: + dino_model = torch.hub.load( + dino_repo or "facebookresearch/dinov2:7764ea0f912e53c92e82eb78a2a1631e92725fc8", + "dinov2_vitb14", + source="local" if dino_repo else "github", + pretrained=not bool(dino_weights), + trust_repo=True, + skip_validation=True, + ) + if dino_weights: + dino_model.load_state_dict( + torch.load(dino_weights, map_location="cpu", weights_only=True), strict=True + ) + self.dino = dino_model.eval().requires_grad_(False) if self.weights["dino"] else None + self.feature_batch_size = feature_batch_size + self.feature_checkpointing = feature_checkpointing + self.grid_period = grid_period + self.register_buffer("mean", torch.tensor([0.485, 0.456, 0.406])[None, :, None, None]) + self.register_buffer("std", torch.tensor([0.229, 0.224, 0.225])[None, :, None, None]) + self.eval() + + def train(self, mode=True): + return super().train(False) + + def dino_features(self, x): + x = F.interpolate(x.float(), (224, 560), mode="bilinear", align_corners=False) + return self.dino.get_intermediate_layers( + (x - self.mean) / self.std, + n=(2, 5, 8, 11), + reshape=False, + return_class_token=False, + norm=True, + ) + + def feature_terms(self, pred, target): + lp = ( + self.lpips(pred * 2 - 1, target * 2 - 1).float().flatten() + if self.weights["lpips"] + else pred.new_zeros(len(pred)) + ) + semantic = pred.new_zeros(len(pred)) + if self.dino is not None: + with torch.no_grad(): + reference = self.dino_features(target) + features = self.dino_features(pred) + semantic = torch.stack( + [ + (fp.float() - ft.float()).square().sum((-2, -1)) + / ft.float().square().sum((-2, -1)).clamp_min(1e-8) + for fp, ft in zip(features, reference, strict=True) + ] + ).mean(0) + return lp, semantic + + def forward(self, pred, target): + if pred.ndim != 5 or pred.shape != target.shape or pred.shape[1] != 3: + raise ValueError("Expected equal BCTHW RGB tensors in [0,1]") + p = pred.permute(0, 2, 1, 3, 4).flatten(0, 1).float() + x = target.permute(0, 2, 1, 3, 4).flatten(0, 1).float() + terms = { + "l1": (p - x).abs().mean(), + "grid": grid_loss_per_frame(p, x, self.grid_period).mean(), + } + chunks = [] + for start in range(0, len(p), self.feature_batch_size): + inputs = ( + p[start : start + self.feature_batch_size], + x[start : start + self.feature_batch_size], + ) + if self.feature_checkpointing and torch.is_grad_enabled() and p.requires_grad: + chunks.append(checkpoint(self.feature_terms, *inputs, use_reentrant=False)) + else: + chunks.append(self.feature_terms(*inputs)) + terms.update( + lpips=torch.cat([v[0] for v in chunks]).mean(), + dino=torch.cat([v[1] for v in chunks]).mean(), + ) + return sum(self.weights[k] * v for k, v in terms.items()), terms + + @torch.no_grad() + def metrics(self, pred, target): + """Unweighted frame means; RGB Gaussian 11x11 valid SSIM, no DINO in eval.""" + p = pred.permute(0, 2, 1, 3, 4).flatten(0, 1).float() + x = target.permute(0, 2, 1, 3, 4).flatten(0, 1).float() + rows = [] + for start in range(0, len(p), self.feature_batch_size): + a, b = ( + p[start : start + self.feature_batch_size], + x[start : start + self.feature_batch_size], + ) + lp = self.lpips(a * 2 - 1, b * 2 - 1).float().flatten() + with torch.autocast(p.device.type, enabled=False): + q = torch.arange(11, device=p.device, dtype=torch.float32) - 5 + kernel = torch.exp(-q.square() / (2 * 1.5**2)) + kernel /= kernel.sum() + window = (kernel[:, None] * kernel[None, :])[None, None].expand(3, 1, 11, 11) + mu_a, mu_b = (F.conv2d(v, window, groups=3) for v in (a, b)) + va = (F.conv2d(a.square(), window, groups=3) - mu_a.square()).clamp_min(0) + vb = (F.conv2d(b.square(), window, groups=3) - mu_b.square()).clamp_min(0) + cov = F.conv2d(a * b, window, groups=3) - mu_a * mu_b + ss = ((2 * mu_a * mu_b + 0.01**2) * (2 * cov + 0.03**2)) / ( + (mu_a.square() + mu_b.square() + 0.01**2) * (va + vb + 0.03**2) + ) + psnr = -10 * (a - b).square().mean((1, 2, 3)).clamp_min(1e-12).log10() + grid = grid_loss_per_frame(a, b, self.grid_period) + rows.append(torch.stack((psnr, ss.mean((1, 2, 3)), lp, grid), dim=1)) + return dict(zip(("psnr", "ssim", "lpips", "grid"), torch.cat(rows).mean(0).tolist())) + +@torch.no_grad() +def log_validation(vae, tae, args, accelerator, weight_dtype, global_step, + criterion, validation_dataloader, selection_state, ema_tae=None): + """Validate the Super decoder using official latents, then select inference weights.""" + accelerator.wait_for_everyone() + if accelerator.is_main_process: + logger.info("Running validation...") + unwrapped_tae = accelerator.unwrap_model(tae) + trainable_params = list(unwrapped_tae.model.decoder.parameters()) + was_training = unwrapped_tae.training + if ema_tae is not None: + ema_tae.store(trainable_params) + ema_tae.copy_to(trainable_params) + unwrapped_tae.eval() + rows = [] + try: + for i, batch in enumerate(validation_dataloader): + if args.validation_max_samples and i >= args.validation_max_samples: + break + pixel_values = batch["pixel_values"].to(accelerator.device) + with accelerator.autocast(): + gt_latents = vae.encode(pixel_values.to(weight_dtype)).latent_dist.mode() + decoded = unwrapped_tae(gt_latents).sample[:, :, :pixel_values.shape[2]] + decoded = decoded.float().add(1).mul(0.5) + target = pixel_values.float().add(1).mul(0.5) + rows.append(criterion.metrics(decoded, target)) + if i == 0: + save_videos_grid( + torch.cat([target, decoded]), + os.path.join(args.output_dir, f"sample/step-{global_step}.mp4"), n_rows=2, + ) + if not rows: + raise ValueError("Validation set is empty") + metrics = {key: sum(row[key] for row in rows) / len(rows) for key in rows[0]} + if selection_state["reference"] is None: + selection_state["reference"] = metrics + reference = selection_state["reference"] + eligible = (metrics["lpips"] <= reference["lpips"] * args.selection_lpips_ratio and + metrics["grid"] <= reference["grid"] * args.selection_grid_ratio) + best = eligible and (selection_state["best_psnr"] is None or metrics["psnr"] > selection_state["best_psnr"]) + if best: + selection_state["best_psnr"] = metrics["psnr"] + unwrapped_tae.save_pretrained(os.path.join(args.output_dir, "best", "taehv")) + row = {"step": global_step, "weights": "ema" if ema_tae is not None else "raw", + "eligible": eligible, "best": best, **metrics} + with open(os.path.join(args.output_dir, "validation.jsonl"), "a") as file: + file.write(json.dumps(row) + "\n") + accelerator.log({f"validation/{key}": value for key, value in metrics.items()}, step=global_step) + logger.info(row) + finally: + if ema_tae is not None: + ema_tae.restore(trainable_params) + unwrapped_tae.train(was_training) + accelerator.wait_for_everyone() + + + +def parse_args(): + parser = argparse.ArgumentParser(description="Training script for the Wan2.2 Super decoder.") + parser.add_argument( + "--config_path", + type=str, + default="config/wan2.2/wan_civitai_5b.yaml", + help="Path to the Wan2.2 model config yaml; vae_kwargs configures the frozen official VAE.", + ) + parser.add_argument( + "--pretrained_model_name_or_path", + type=str, + default=None, + required=True, + help="Path to pretrained model directory (contains the full VAE weights).", + ) + parser.add_argument( + "--train_data_dir", + type=str, + default="datasets/X-Fun-Videos-Demo", + help=("A folder containing the training data."), + ) + parser.add_argument( + "--train_data_meta", + type=str, + default="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json", + help=("A JSON array containing the training video metadata."), + ) + parser.add_argument( + "--output_dir", + type=str, + default="output_dir_taehv_super", + help="The output directory where the model predictions and checkpoints will be written.", + ) + parser.add_argument("--seed", type=int, default=42, help="A seed for reproducible training.") + parser.add_argument( + "--train_batch_size", type=int, default=1, help="Batch size (per device) for the training dataloader." + ) + parser.add_argument( + "--vae_mini_batch", type=int, default=1, help="Mini batch size for VAE encoding." + ) + 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( + "--learning_rate", + type=float, + default=5e-6, + help="Initial learning rate (after the potential warmup period) to use.", + ) + parser.add_argument( + "--lr_scheduler", + type=str, + default="cosine", + 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=200, help="Number of steps for the warmup in the lr scheduler." + ) + parser.add_argument( + "--lr_min_ratio", type=float, default=0.2, + help="Final LR / peak LR for the cosine scheduler (stage 4a: 1e-6 / 5e-6).", + ) + 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." + ), + ) + parser.add_argument("--use_ema", action=argparse.BooleanOptionalAction, default=False, help="Whether to use EMA model.") + parser.add_argument( + "--dataloader_num_workers", + type=int, + default=0, + help="Number of subprocesses to use for data loading.", + ) + 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.99, 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 log directory.", + ) + parser.add_argument( + "--mixed_precision", + type=str, + default="bf16", + choices=["no", "fp16", "bf16"], + help="Whether to use mixed precision. Choose between fp16 and bf16.", + ) + parser.add_argument( + "--report_to", + type=str, + default="tensorboard", + help='The integration to report the results and logs to.', + ) + parser.add_argument( + "--checkpointing_steps", + type=int, + default=500, + help="Save a checkpoint of the training state every X updates.", + ) + 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.", + ) + parser.add_argument( + "--validation_steps", + type=int, + default=2000, + help="Run validation every X steps.", + ) + parser.add_argument( + "--tracker_project_name", + type=str, + default="taehv-train", + help="The `project_name` argument passed to Accelerator.init_trackers.", + ) + parser.add_argument( + "--video_sample_size", + type=int, + default=512, + help="Training video sample size.", + ) + parser.add_argument( + "--video_sample_stride", + type=int, + default=1, + help="Sample stride of the video.", + ) + parser.add_argument( + "--video_sample_n_frames", + type=int, + default=25, + help="Num frame of video. Must be 4k+1 to match the causal temporal compression.", + ) + parser.add_argument( + "--tae_path", + type=str, + default=None, + help="Optional Super weights to warm-start from (taew2_2_super.pth/.safetensors, a directory, " + "or a previous checkpoint-xxx/taehv dir). If not given, train from scratch.", + ) + parser.add_argument( + "--use_taehv_sequential", action="store_true", + help="Use the sequential TAE decoder instead of the parallel path.", + ) + parser.add_argument( + "--from_scratch", + action="store_true", + help="Initialize a new Super decoder without student weights.", + ) + parser.add_argument( + "--fix_sample_size", + nargs=2, + type=int, + default=None, + help="Fixed [height, width] overriding video_sample_size.", + ) + parser.add_argument( + "--video_sample_strides", + nargs="+", + type=int, + default=None, + help="Optional per-video random sampling strides.", + ) + parser.add_argument( + "--validation_data_meta", + type=str, + default=None, + help="Held-out validation JSON index.", + ) + parser.add_argument( + "--validation_n_frames", + type=int, + default=81, + help="Validation clip length, 4k+1.", + ) + parser.add_argument( + "--validation_max_samples", + type=int, + default=0, + help="Validation sample cap; zero evaluates the complete index.", + ) + parser.add_argument( + "--dino_repo", + type=str, + default=None, + help="Optional local pinned DINOv2 source directory.", + ) + parser.add_argument( + "--dino_weights", + type=str, + default=None, + help="Optional local ViT-B/14 weights.", + ) + parser.add_argument( + "--feature_batch_size", + type=int, + default=13, + help="Frame batch size for perceptual networks.", + ) + parser.add_argument( + "--feature_checkpointing", + action=argparse.BooleanOptionalAction, + default=False, + help="Checkpoint perceptual feature computation.", + ) + parser.add_argument( + "--grid_period", + type=int, + default=16, + help="Spatial residual period for grid loss.", + ) + parser.add_argument( + "--l1_loss_weight", + type=float, + default=0.1, + help="RGB L1 reconstruction loss weight.", + ) + parser.add_argument( + "--lpips_loss_weight", + type=float, + default=0.1, + help="LPIPS reconstruction loss weight.", + ) + parser.add_argument( + "--dino_loss_weight", + type=float, + default=0.02, + help="DINOv2 patch feature loss weight.", + ) + parser.add_argument( + "--grid_loss_weight", + type=float, + default=1.0, + help="Periodic spatial residual loss weight.", + ) + parser.add_argument( + "--ema_decay", + type=float, + default=0.999, + help="EMA decay for decoder parameters.", + ) + parser.add_argument( + "--selection_lpips_ratio", + type=float, + default=1.02, + help="Maximum LPIPS relative to the first validation.", + ) + parser.add_argument( + "--selection_grid_ratio", + type=float, + default=1.05, + help="Maximum grid loss relative to the first validation.", + ) + + args = parser.parse_args() + + if args.fix_sample_size is not None: + args.video_sample_height, args.video_sample_width = args.fix_sample_size + else: + args.video_sample_height = args.video_sample_width = args.video_sample_size + if any(size < 32 or size % 16 for size in (args.video_sample_height, args.video_sample_width)): + parser.error("Video dimensions must be >=32 and divisible by 16") + if any(frames < 1 or (frames - 1) % 4 for frames in (args.video_sample_n_frames, args.validation_n_frames)): + parser.error("Video frame counts must be 4k+1") + if args.from_scratch and (args.tae_path or args.resume_from_checkpoint): + parser.error("--from_scratch cannot be combined with weight initialization or resume") + if args.max_train_steps is not None and args.max_train_steps < 1: + parser.error("max_train_steps must be positive") + if min(args.num_train_epochs, args.train_batch_size, args.vae_mini_batch, + args.gradient_accumulation_steps, args.checkpointing_steps, args.video_sample_stride, + args.feature_batch_size, args.grid_period) < 1: + parser.error("Epochs, batch sizes, checkpoint interval, stride and grid period must be positive") + if min(args.dataloader_num_workers, args.validation_steps, args.validation_max_samples, args.lr_warmup_steps) < 0: + parser.error("Workers, validation intervals/sample cap and warmup cannot be negative") + if args.video_sample_strides is not None and (not args.video_sample_strides or min(args.video_sample_strides) < 1): + parser.error("video_sample_strides must be positive") + if args.checkpoints_total_limit is not None and args.checkpoints_total_limit < 1: + parser.error("checkpoints_total_limit must be positive") + if not 0 <= args.ema_decay < 1 or not 0 <= args.adam_beta1 < 1 or not 0 <= args.adam_beta2 < 1: + parser.error("Invalid EMA decay or Adam betas") + if args.adam_weight_decay < 0: + parser.error("Adam weight decay cannot be negative") + if not 0 <= args.lr_min_ratio <= 1: + parser.error("lr_min_ratio must be between 0 and 1") + if any(not math.isfinite(value) or value <= 0 for value in + (args.learning_rate, args.adam_epsilon, args.max_grad_norm, + args.selection_lpips_ratio, args.selection_grid_ratio)): + parser.error("Learning rate, epsilon, gradient norm and selection ratios must be finite and positive") + loss_weights = (args.l1_loss_weight, args.lpips_loss_weight, args.dino_loss_weight, + args.grid_loss_weight) + if not any(loss_weights) or any(not math.isfinite(value) or value < 0 for value in loss_weights): + parser.error("Loss weights must be finite and nonnegative, with at least one positive term") + return args + + + + +def main(): + args = parse_args() + logging_dir = os.path.join(args.output_dir, args.logging_dir) + 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=None if args.report_to == "none" else args.report_to, + project_config=accelerator_project_config, + ) + if accelerator.distributed_type.name not in ("NO", "MULTI_GPU", "MULTI_CPU"): + raise ValueError("This decoder trainer supports single device and ordinary DDP") + + 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: + diffusers.utils.logging.set_verbosity_info() + else: + diffusers.utils.logging.set_verbosity_error() + set_seed(args.seed) + if args.allow_tf32: + torch.backends.cuda.matmul.allow_tf32 = True + if accelerator.is_main_process: + os.makedirs(args.output_dir, exist_ok=True) + if not args.resume_from_checkpoint and ( + os.path.exists(os.path.join(args.output_dir, "train.jsonl")) or + any(name.startswith("checkpoint-") for name in os.listdir(args.output_dir)) + ): + raise FileExistsError("Use a new output_dir or --resume_from_checkpoint") + + # For mixed precision training, keep trainable parameters and Adam moments FP32. + weight_dtype = torch.float32 + if accelerator.mixed_precision == "fp16": + weight_dtype = torch.float16 + elif accelerator.mixed_precision == "bf16": + weight_dtype = torch.bfloat16 + + # ==================== Load Config ==================== + config = OmegaConf.load(args.config_path) + if config['vae_kwargs'].get('vae_type') != 'AutoencoderKLWan3_8': + raise ValueError("Super training requires the official 48-channel Wan2.2 VAE") + + # ==================== Load Models ==================== + vae = AutoencoderKLWan3_8.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() + vae.requires_grad_(False) + vae.to(accelerator.device, dtype=weight_dtype) + + tae_load_kwargs = {"latent_channels": 48, "patch_size": 2, "arch_variant": "super"} + if args.tae_path is not None and not args.resume_from_checkpoint: + logger.info(f"Warm-starting Super TAE from checkpoint: {args.tae_path}") + tae = AutoencoderTinyWan.from_pretrained( + args.tae_path, additional_kwargs=tae_load_kwargs, + ) + else: + logger.info("Initializing Super TAE (training state is restored below when resuming)") + tae = AutoencoderTinyWan(**tae_load_kwargs) + if tae.config.arch_variant != "super" or tae.config.latent_channels != 48 or tae.config.patch_size != 2: + raise ValueError("Expected a 48-channel, patch-size-2 Super checkpoint") + tae.parallel = not args.use_taehv_sequential + tae.float() + tae.requires_grad_(False) + tae.model.decoder.requires_grad_(True) + tae.train() + + trainable_params = list(filter(lambda p: p.requires_grad, tae.parameters())) + total_params = sum(p.numel() for p in tae.parameters()) + trainable_count = sum(p.numel() for p in trainable_params) + logger.info(f"Total params: {total_params / 1e6:.2f}M, Trainable params: {trainable_count / 1e6:.2f}M") + ema_tae = None + if args.use_ema: + ema_tae = EMAModel( + trainable_params, decay=args.ema_decay, min_decay=args.ema_decay, use_ema_warmup=False, + ) + + # ==================== Save/Load Hooks ==================== + def save_model_hook(models, weights, output_dir): + if accelerator.is_main_process: + accelerator.unwrap_model(models[0]).save_pretrained(os.path.join(output_dir, "taehv")) + trainer_state = {"step": global_step, "epoch": epoch, "batch": resume_step, **selection_state} + with open(os.path.join(output_dir, "trainer_state.json"), "w") as file: + json.dump(trainer_state, file, indent=2) + # The model is saved in Diffusers format by this hook, as in train_taehv.py. + weights.pop() + + def load_model_hook(models, input_dir): + while models: + model = accelerator.unwrap_model(models.pop()) + load_model = AutoencoderTinyWan.from_pretrained(input_dir, subfolder="taehv") + model.register_to_config(**load_model.config) + model.load_state_dict(load_model.state_dict(), strict=True) + del load_model + + accelerator.register_save_state_pre_hook(save_model_hook) + accelerator.register_load_state_pre_hook(load_model_hook) + + # ==================== Optimizer ==================== + optimizer = torch.optim.AdamW( + trainable_params, lr=args.learning_rate, + betas=(args.adam_beta1, args.adam_beta2), + weight_decay=args.adam_weight_decay, eps=args.adam_epsilon, + ) + + # ==================== Dataset ==================== + train_dataset = VideoDataset( + args.train_data_meta, args.train_data_dir, + sample_size=args.video_sample_size, + sample_stride=args.video_sample_strides or args.video_sample_stride, + sample_n_frames=args.video_sample_n_frames, + enable_bucket=True, enable_inpaint=False, text_drop_ratio=0, + ) + if not len(train_dataset): + raise ValueError("Training set is empty") + + def collate_fn(examples): + """Fixed canvas, aspect-preserving resize/crop; align the shortest clip to 4k+1.""" + frames = (min(len(e["pixel_values"]) for e in examples) - 1) // 4 * 4 + 1 + height, width = args.video_sample_height, args.video_sample_width + videos = [] + for example in examples: + x = torch.from_numpy(example["pixel_values"][:frames]).permute(0, 3, 1, 2).float() + h, w = x.shape[-2:] + scale = max(height / h, width / w) + resize_size = (math.ceil(h * scale), math.ceil(w * scale)) + x = transforms.functional.resize(x, resize_size, antialias=True) + x = transforms.functional.center_crop(x, (height, width)) + videos.append(x.div_(127.5).sub_(1)) + return {"pixel_values": torch.stack(videos).permute(0, 2, 1, 3, 4).contiguous()} + + 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, drop_last=True, + ) + 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, + ) + validation_dataloader = None + if args.validation_data_meta is not None: + validation_dataset = VideoDataset( + args.validation_data_meta, args.train_data_dir, + sample_n_frames=args.validation_n_frames, sample_stride=1, + enable_bucket=True, enable_inpaint=False, text_drop_ratio=0, + random_start=False, + ) + validation_dataloader = torch.utils.data.DataLoader( + validation_dataset, batch_size=1, collate_fn=collate_fn, + ) + + # ==================== LR Scheduler ==================== + tae, optimizer, train_dataloader = accelerator.prepare(tae, optimizer, train_dataloader) + num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) + if num_update_steps_per_epoch == 0: + raise ValueError("Training data does not contain a complete batch") + if args.max_train_steps is None: + args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch + args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch) + lr_scheduler = get_scheduler( + args.lr_scheduler, optimizer=optimizer.optimizer, + num_warmup_steps=args.lr_warmup_steps, + num_training_steps=args.max_train_steps, + ) + if args.lr_scheduler == "cosine": + # Keep Diffusers warmup/cosine and add stage 4a's nonzero final LR. + cosine_schedule = lr_scheduler.lr_lambdas[0] + + def lr_lambda(step): + factor = cosine_schedule(step) + if step < args.lr_warmup_steps: + return factor + return args.lr_min_ratio + (1 - args.lr_min_ratio) * factor + + lr_scheduler.lr_lambdas = [lr_lambda] * len(lr_scheduler.lr_lambdas) + # Register explicitly so it advances once per optimizer update, at any world size. + accelerator.register_for_checkpointing(lr_scheduler) + if args.use_ema: + ema_tae.to(accelerator.device) + accelerator.register_for_checkpointing(ema_tae) + + # ==================== TinyVAE Losses ==================== + loss_weights = { + "l1": args.l1_loss_weight, + "lpips": args.lpips_loss_weight, + "dino": args.dino_loss_weight, + "grid": args.grid_loss_weight, + } + criterion = TinyVAEReconstructionLoss( + loss_weights, + dino_repo=args.dino_repo, dino_weights=args.dino_weights, + feature_batch_size=args.feature_batch_size, + feature_checkpointing=args.feature_checkpointing, grid_period=args.grid_period, + ).to(accelerator.device) + + if accelerator.is_main_process: + tracker_config = dict(vars(args)) + keys_to_pop = [key for key, value in tracker_config.items() if isinstance(value, list)] + for key in keys_to_pop: + tracker_config.pop(key) + accelerator.init_trackers(args.tracker_project_name, tracker_config) + + # ==================== Training ==================== + total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps + logger.info("***** Running Super decoder 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}") + logger.info(f" Loss weights = {loss_weights}") + global_step = 0 + first_epoch = 0 + resume_step = 0 + selection_state = {"reference": None, "best_psnr": None} + + # Resume from checkpoint + if args.resume_from_checkpoint: + if args.resume_from_checkpoint != "latest": + path = args.resume_from_checkpoint + else: + dirs = [name for name in os.listdir(args.output_dir) + if name.startswith("checkpoint-") and name.split("-")[-1].isdigit()] + dirs = sorted(dirs, key=lambda name: int(name.split("-")[-1])) + if not dirs: + raise FileNotFoundError(f"No checkpoints in {args.output_dir}") + path = os.path.join(args.output_dir, dirs[-1]) + with open(os.path.join(path, "trainer_state.json")) as file: + trainer_state = json.load(file) + global_step = trainer_state["step"] + first_epoch = trainer_state["epoch"] + resume_step = trainer_state["batch"] + selection_state["reference"] = trainer_state["reference"] + selection_state["best_psnr"] = trainer_state["best_psnr"] + accelerator.print(f"Resuming from checkpoint {path}") + accelerator.load_state(path) + if args.use_ema: + ema_tae.to(accelerator.device) + initial_global_step = global_step + epoch = first_epoch + if accelerator.is_main_process: + OmegaConf.save(OmegaConf.create(vars(args)), os.path.join(args.output_dir, f"args-step-{global_step}.yaml")) + if validation_dataloader is not None and selection_state["reference"] is None: + log_validation(vae, tae, args, accelerator, weight_dtype, global_step, + criterion, validation_dataloader, selection_state, ema_tae) + + progress_bar = PauseAwareTqdm( + range(0, args.max_train_steps), initial=initial_global_step, desc="Steps", + disable=not accelerator.is_local_main_process, + ) + optimizer.zero_grad() + for epoch in range(first_epoch, args.num_train_epochs): + if global_step >= args.max_train_steps: + break + train_loss = torch.zeros(1 + len(loss_weights), device=accelerator.device) + micro_steps = 0 + train_dataloader.set_epoch(epoch) + batch_sampler.sampler.generator = torch.Generator().manual_seed(args.seed + epoch) + batch_sampler.sampler._pos_start = 0 + batch_sampler.bucket = {"image": [], "video": []} + epoch_dataloader = accelerator.skip_first_batches(train_dataloader, resume_step) if resume_step else train_dataloader + for step, batch in enumerate(epoch_dataloader, start=resume_step): + with accelerator.accumulate(tae): + pixel_values = batch["pixel_values"].to(accelerator.device) # [B, C, F, H, W], FP32 target + bsz = pixel_values.shape[0] + + # 1. Official latents, frozen. + with torch.no_grad(), accelerator.autocast(): + gt_latents_list = [] + for i in range(0, bsz, args.vae_mini_batch): + encoded = vae.encode(pixel_values[i:i + args.vae_mini_batch].to(weight_dtype)).latent_dist.mode() + gt_latents_list.append(encoded) + gt_latents = torch.cat(gt_latents_list, dim=0).detach() + + # 2. Student forward: official latent -> Super decoder, through DDP. + with accelerator.autocast(): + decoded = tae(gt_latents).sample[:, :, :pixel_values.shape[2]] + decoded = decoded.float().add(1).mul(0.5) + target = pixel_values.float().add(1).mul(0.5) + loss, loss_terms = criterion(decoded, target) + if not torch.isfinite(loss): + raise FloatingPointError(f"Nonfinite loss at step {global_step}") + + # 3. Gather losses for logging and backpropagate. + train_loss += torch.stack([loss.detach(), *[loss_terms[key].detach() for key in loss_weights]]) + micro_steps += 1 + accelerator.backward(loss) + if accelerator.sync_gradients: + accelerator.clip_grad_norm_(trainable_params, args.max_grad_norm) + optimizer.step() + optimizer.zero_grad() + if accelerator.sync_gradients and not accelerator.optimizer_step_was_skipped: + lr_scheduler.step() + if args.use_ema: + ema_tae.step(trainable_params) + + resume_step = step + 1 + if accelerator.sync_gradients and not accelerator.optimizer_step_was_skipped: + progress_bar.update(1) + global_step += 1 + averaged_loss = accelerator.reduce(train_loss / micro_steps, reduction="mean").tolist() + logs = dict(zip(("loss", *loss_weights), averaged_loss)) + logs.update(step=global_step, frames=pixel_values.shape[2], lr=lr_scheduler.get_last_lr()[0]) + accelerator.log(logs, step=global_step) + if accelerator.is_main_process: + with open(os.path.join(args.output_dir, "train.jsonl"), "a") as file: + file.write(json.dumps(logs) + "\n") + train_loss.zero_() + micro_steps = 0 + progress_bar.set_postfix(step_loss=logs["loss"], lr=logs["lr"]) + + # Validation + if validation_dataloader is not None and args.validation_steps and global_step % args.validation_steps == 0: + with progress_bar.paused(): + log_validation(vae, tae, args, accelerator, weight_dtype, global_step, + criterion, validation_dataloader, selection_state, ema_tae) + + # Checkpointing + if global_step % args.checkpointing_steps == 0: + if accelerator.is_main_process and args.checkpoints_total_limit is not None: + checkpoints = [name for name in os.listdir(args.output_dir) + if name.startswith("checkpoint-") and name.split("-")[-1].isdigit()] + checkpoints = sorted(checkpoints, key=lambda name: int(name.split("-")[-1])) + num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1 + for removing_checkpoint in checkpoints[:max(num_to_remove, 0)]: + shutil.rmtree(os.path.join(args.output_dir, removing_checkpoint)) + accelerator.wait_for_everyone() + save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") + with progress_bar.paused(): + accelerator.save_state(save_path) + logger.info(f"Saved state to {save_path}") + if global_step >= args.max_train_steps: + break + if global_step >= args.max_train_steps: + break + resume_step = 0 + progress_bar.close() + + # Final validation and save. Restore raw parameters after EMA validation. + if validation_dataloader is not None and (not args.validation_steps or global_step % args.validation_steps): + log_validation(vae, tae, args, accelerator, weight_dtype, global_step, + criterion, validation_dataloader, selection_state, ema_tae) + accelerator.wait_for_everyone() + 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/taehv/train_taehv_super.sh b/scripts/taehv/train_taehv_super.sh new file mode 100644 index 0000000..c332fcf --- /dev/null +++ b/scripts/taehv/train_taehv_super.sh @@ -0,0 +1,50 @@ +#!/usr/bin/env bash +export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-TI2V-5B" +export DATASET_NAME="${DATASET_NAME:-datasets/X-Fun-Videos-Demo}" +export DATASET_META_NAME="${DATASET_META_NAME:-$DATASET_NAME/metadata_add_width_height.json}" +export TAE_PATH="models/Diffusion_Transformer/taew2_2_super.pth" +# 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 + +accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \ + --config_path="config/wan2.2/wan_civitai_5b.yaml" \ + --pretrained_model_name_or_path="$MODEL_NAME" \ + --tae_path="$TAE_PATH" \ + --train_data_dir="$DATASET_NAME" \ + --train_data_meta="$DATASET_META_NAME" \ + --fix_sample_size 720 1280 \ + --video_sample_strides 1 2 3 \ + --video_sample_n_frames=25 \ + --validation_n_frames=81 \ + --train_batch_size=1 \ + --gradient_accumulation_steps=1 \ + --dataloader_num_workers=8 \ + --max_train_steps=10000 \ + --checkpointing_steps=1000 \ + --validation_steps=1000 \ + --validation_max_samples=0 \ + --learning_rate=5e-06 \ + --lr_scheduler="cosine" \ + --lr_warmup_steps=200 \ + --lr_min_ratio=0.2 \ + --l1_loss_weight=0.1 \ + --lpips_loss_weight=0.1 \ + --dino_loss_weight=0.02 \ + --grid_loss_weight=1.0 \ + --grid_period=16 \ + --feature_batch_size=13 \ + --use_ema \ + --ema_decay=0.999 \ + --selection_lpips_ratio=1.02 \ + --selection_grid_ratio=1.05 \ + --seed=42 \ + --output_dir="output_dir_taehv_w2.2_super" \ + --mixed_precision="bf16" \ + --adam_beta1=0.9 \ + --adam_beta2=0.99 \ + --adam_weight_decay=1e-2 \ + --adam_epsilon=1e-08 \ + --vae_mini_batch=1 \ + --max_grad_norm=1.0 diff --git a/videox_fun/data/dataset_video.py b/videox_fun/data/dataset_video.py index 400a3da..736ff6e 100644 --- a/videox_fun/data/dataset_video.py +++ b/videox_fun/data/dataset_video.py @@ -173,6 +173,7 @@ class VideoDataset(Dataset): video_length_drop_start=0.0, video_length_drop_end=1.0, text_drop_ratio=0.1, + random_start=True, ): # Loading annotations from files print(f"loading annotations from {ann_path} ...") @@ -189,6 +190,7 @@ class VideoDataset(Dataset): self.video_length_drop_start = video_length_drop_start self.video_length_drop_end = video_length_drop_end self.text_drop_ratio = text_drop_ratio + self.random_start = random_start sample_size = tuple(sample_size) if not isinstance(sample_size, int) else (sample_size, sample_size) self.pixel_transforms = transforms.Compose( @@ -212,17 +214,22 @@ class VideoDataset(Dataset): with VideoReader_contextmanager(video_dir, num_threads=2) as video_reader: # Calculate frame sampling range with length dropout + sample_stride = (random.choice(self.sample_stride) + if isinstance(self.sample_stride, (tuple, list)) else self.sample_stride) min_sample_n_frames = min( self.sample_n_frames, - int(len(video_reader) * (self.video_length_drop_end - self.video_length_drop_start) // self.sample_stride) + int((len(video_reader) * (self.video_length_drop_end - self.video_length_drop_start) - 1) + // sample_stride) + 1 ) if min_sample_n_frames == 0: raise ValueError(f"No Frames in video.") # Select contiguous clip with random start position video_length = int(self.video_length_drop_end * len(video_reader)) - clip_length = min(video_length, (min_sample_n_frames - 1) * self.sample_stride + 1) - start_idx = random.randint(int(self.video_length_drop_start * video_length), video_length - clip_length) if video_length != clip_length else 0 + clip_length = min(video_length, (min_sample_n_frames - 1) * sample_stride + 1) + start_idx = int(self.video_length_drop_start * video_length) + if self.random_start and video_length != clip_length: + start_idx = random.randint(start_idx, video_length - clip_length) batch_index = np.linspace(start_idx, start_idx + clip_length - 1, min_sample_n_frames, dtype=int) try: diff --git a/videox_fun/models/wan_tae.py b/videox_fun/models/wan_tae.py index 6f8a9bd..c0de53f 100644 --- a/videox_fun/models/wan_tae.py +++ b/videox_fun/models/wan_tae.py @@ -449,6 +449,10 @@ class AutoencoderTinyWan(ModelMixin, ConfigMixin, FromOriginalModelMixin): return (decoded,) return DecoderOutput(sample=decoded) + def forward(self, z: torch.Tensor, return_dict: bool = True): + # Use forward so decoder training runs through DDP hooks. + return self.decode(z, return_dict=return_dict) + @classmethod def from_pretrained(cls, pretrained_model_path, additional_kwargs={}, **kwargs): def filter_kwargs(cls, kwargs):