add the training script for taew2_2_super
This commit is contained in:
huangkunzhe.hkz
@@ -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.
|
||||
@@ -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。
|
||||
@@ -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()
|
||||
@@ -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
|
||||
@@ -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:
|
||||
|
||||
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user