Merge remote-tracking branch 'origin/main' into gpu_memory_mode

comfyui/z_image/nodes.py

@@ -97,12 +97,6 @@ class LoadZImageTransformerModel:
 
         model_path = folder_paths.get_full_path("diffusion_models", model_name)
         transformer_state_dict = load_torch_file(model_path, safe_load=True)
-        # Some ComfyUI-native checkpoints store every diffusion-model tensor under a
-        # `model.diffusion_model.` prefix (e.g. CivitAI Z-Image fine-tunes). Strip it so
-        # the conversion gate below and `convert_state_dict` receive bare keys.
-        _prefix = "model.diffusion_model."
-        if transformer_state_dict and all(k.startswith(_prefix) for k in transformer_state_dict):
-            transformer_state_dict = {k[len(_prefix):]: v for k, v in transformer_state_dict.items()}
         def convert_state_dict(old_state_dict):
             new_state_dict = {}
             for key, value in old_state_dict.items():

scripts/taehv/README_TRAIN.md

@@ -128,7 +128,7 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun
 
 ```bash
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
-export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata.json"
 ```
 
 **Absolute Path**:
@@ -173,7 +173,7 @@ wget https://github.com/madebyollin/taehv/raw/main/taew2_2.safetensors
 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-TI2V-5B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
-export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata.json"
 # Optional: warm-start from released TAE weights instead of training from scratch.
 export TAE_PATH="models/Diffusion_Transformer/taew2_2.safetensors"
 NCCL_DEBUG=INFO
@@ -203,6 +203,7 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv.py \
   --adam_epsilon=1e-08 \
   --vae_mini_batch=1 \
   --max_grad_norm=1.0 \
+  --random_hw_adapt \
   --enable_bucket \
   --trainable_modules "." \
   --resume_from_checkpoint=latest
@@ -236,11 +237,10 @@ export TAE_PATH="taew2_1.safetensors"
 | `--use_taehv_sequential` | Run the TAE in O(1)-memory sequential mode instead of parallel | - |
 | `--latent_loss_weight` | Weight of the latent MSE (TAE latent vs. full VAE latent) relative to pixel L1 | 1.0 |
 | `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Demo/` |
-| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` |
+| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Demo/metadata.json` |
 | `--train_batch_size` | Batch size (per device) | 1 |
 | `--video_sample_size` | Training resolution | 512 |
 | `--video_sample_stride` | Video sample stride | 1 |
-| `--video_sample_strides` | Space-separated stride candidates; one is drawn at random per clip (overrides `--video_sample_stride`) | `None` |
 | `--video_sample_n_frames` | Number of frames to sample. **Must be 4k+1** (33, 49, 81, ...) | 33 |
 | `--vae_mini_batch` | Mini batch size for teacher VAE encoding | 1 |
 | `--gradient_accumulation_steps` | Gradient accumulation steps | 1 |
@@ -253,26 +253,22 @@ export TAE_PATH="taew2_1.safetensors"
 | `--lr_warmup_steps` | Learning rate warmup steps | 100 |
 | `--use_8bit_adam` / `--use_came` | Alternative optimizers | - |
 | `--use_ema` | Keep an EMA copy of the TAE (used for validation and final save) | - |
-| `--ema_decay` | EMA decay; also disables the EMA warmup ramp. Defaults to the `EMAModel` default (0.9999) | `None` |
-| `--seed` | Random seed; defaults to 42 so a bare CLI run stays reproducible | 42 |
+| `--seed` | Random seed | 42 |
 | `--output_dir` | Output directory | `output_dir_taehv_w2.2` |
 | `--mixed_precision` | Mixed precision: `fp16/bf16` | `bf16` |
 | `--max_grad_norm` | Maximum gradient norm for clipping | 1.0 |
 | `--enable_bucket` | Enable bucket training without cropping, groups by resolution | - |
-| `--fix_sample_size` | Fixed `[height, width]` canvas overriding the square `--video_sample_size` (aspect-preserving cover crop) | `None` |
+| `--random_hw_adapt` | Randomly scale videos to a range of resolutions | - |
 | `--low_vram` | Keep the teacher VAE on CPU and move it to GPU only when encoding | - |
 | `--trainable_modules` | Trainable modules (`"."` means all modules) | `"."` |
 | `--trainable_modules_low_learning_rate` | Trainable modules with lr/2 | `[]` |
 | `--resume_from_checkpoint` | Resume training from checkpoint, use `"latest"` to auto-select | `latest` |
 | `--validation_steps` / `--validation_epochs` | Run validation every N steps / epochs | 2000 / 5 |
 | `--validation_paths` | Video paths for validation reconstruction comparison | `"asset/1.mp4"` |
-| `--validation_n_frames` | Validation clip length (4k+1); defaults to `--video_sample_n_frames` | `None` |
 
 **Sample Size Configuration Guide**:
-- `video_sample_size` sets a square training canvas; pair it with `--enable_bucket` to train on native aspect ratios.
-- `--fix_sample_size 720 1280` pins a fixed non-square canvas (aspect-preserving cover resize + center crop); with `--enable_bucket` it is applied per batch and `video_sample_size` is bumped to `max(fix_sample_size)`.
-- `--video_sample_n_frames` must satisfy `4k+1` (e.g. 33, 49, 81) because both the full VAE and the TAE are causal 4x temporal compressors.
-- `--video_sample_strides 1 2 3` draws a temporal stride at random for every clip (multi-scale temporal augmentation); the frame count requested shrinks as the stride grows, so the time span a clip covers stays roughly the same.
+- `video_sample_size` represents the training resolution; when `random_hw_adapt` is enabled, it represents the minimum resolution.
+- `video_sample_n_frames` must satisfy `4k+1` (e.g. 33, 49, 81) because both the full VAE and the TAE are causal 4x temporal compressors.
 
 ### 3.4 Training Validation
 
@@ -283,18 +279,15 @@ You can configure validation parameters to periodically reconstruct test videos
 | `--validation_steps` | Run validation every N steps | 2000 |
 | `--validation_epochs` | Run validation every N epochs | 5 |
 | `--validation_paths` | Validation video paths | `"asset/1.mp4"` |
-| `--validation_n_frames` | Validation clip length, **4k+1**; omit to reuse `--video_sample_n_frames` | 81 |
 
 ```bash
   --validation_paths "asset/1.mp4" \
   --validation_steps=2000 \
-  --validation_epochs=5 \
-  --validation_n_frames=81
+  --validation_epochs=5
 ```
 
 **Notes**:
 - Validation saves three videos per sample into `output_dir/sample/`: `*_input.mp4` (resized input), `*_taehv.mp4` (TAE reconstruction), `*_fullvae.mp4` (full VAE reconstruction for reference).
-- Validation clips are truncated to `min(video length, --validation_n_frames)` and then aligned down to `4k+1`, so a longer setting never fails on short assets.
 - When `--use_ema` is enabled, validation runs with the EMA weights.
 
 ### 3.5 Training Tips
@@ -304,7 +297,7 @@ You can configure validation parameters to periodically reconstruct test videos
 - **Decoder-only distillation**: Add `--freeze_tae_encoder` to only improve decoding quality.
 - **Memory**: The teacher full VAE is the main memory consumer. Use `--low_vram` to keep it on CPU between encoding steps, or reduce `--video_sample_n_frames` / `--video_sample_size`.
 - **Sequential TAE**: `--use_taehv_sequential` trades speed for O(1) activation memory w.r.t. video length.
-- **EMA**: `--use_ema` is recommended for the final deliverable; the final saved `taehv` directory contains the EMA weights. Pass `--ema_decay` (e.g. `0.999`, as `train_taehv_super.sh` does) when short runs need the average to catch up faster than the 0.9999 default.
+- **EMA**: `--use_ema` is recommended for the final deliverable; the final saved `taehv` directory contains the EMA weights.
 
 ### 3.6 Multi-Node Distributed Training
 

scripts/taehv/README_TRAIN_SUPER.md

@@ -1,434 +0,0 @@
-# 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.

scripts/taehv/README_TRAIN_SUPER_zh-CN.md

@@ -1,419 +0,0 @@
-# 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。

scripts/taehv/README_TRAIN_zh-CN.md

@@ -128,7 +128,7 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun
 
 ```bash
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
-export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata.json"
 ```
 
 **绝对路径**：
@@ -173,7 +173,7 @@ wget https://github.com/madebyollin/taehv/raw/main/taew2_2.safetensors
 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-TI2V-5B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
-export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata.json"
 # 可选：用官方发布的 TAE 权重热启动，而不是从头训练。
 export TAE_PATH="models/Diffusion_Transformer/taew2_2.safetensors"
 NCCL_DEBUG=INFO
@@ -203,6 +203,7 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv.py \
   --adam_epsilon=1e-08 \
   --vae_mini_batch=1 \
   --max_grad_norm=1.0 \
+  --random_hw_adapt \
   --enable_bucket \
   --trainable_modules "." \
   --resume_from_checkpoint=latest
@@ -236,11 +237,10 @@ export TAE_PATH="taew2_1.safetensors"
 | `--use_taehv_sequential` | 以 O(1) 显存的串行模式运行 TAE，而不是并行模式 | - |
 | `--latent_loss_weight` | latent MSE（TAE latent 对齐完整 VAE latent）相对 pixel L1 的权重 | 1.0 |
 | `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Demo/` |
-| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` |
+| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Videos-Demo/metadata.json` |
 | `--train_batch_size` | 每卡批次大小 | 1 |
 | `--video_sample_size` | 训练分辨率 | 512 |
 | `--video_sample_stride` | 视频采样步幅 | 1 |
-| `--video_sample_strides` | 空格分隔的候选步幅，每个片段随机抽取一个（覆盖 `--video_sample_stride`） | `None` |
 | `--video_sample_n_frames` | 采样帧数，**必须为 4k+1**（33、49、81……） | 33 |
 | `--vae_mini_batch` | teacher VAE 编码时的迷你批次大小 | 1 |
 | `--gradient_accumulation_steps` | 梯度累积步数 | 1 |
@@ -253,26 +253,22 @@ export TAE_PATH="taew2_1.safetensors"
 | `--lr_warmup_steps` | 学习率预热步数 | 100 |
 | `--use_8bit_adam` / `--use_came` | 可选优化器 | - |
 | `--use_ema` | 维护 TAE 的 EMA 副本（用于验证与最终保存） | - |
-| `--ema_decay` | EMA 衰减系数，同时关闭 EMA warmup；不传则用 `EMAModel` 默认值（0.9999） | `None` |
-| `--seed` | 随机种子；默认 42，直接跑命令行也可复现 | 42 |
+| `--seed` | 随机种子 | 42 |
 | `--output_dir` | 输出目录 | `output_dir_taehv_w2.2` |
 | `--mixed_precision` | 混合精度：`fp16/bf16` | `bf16` |
 | `--max_grad_norm` | 梯度裁剪阈值 | 1.0 |
 | `--enable_bucket` | 启用分桶训练，不裁剪视频，按分辨率分组训练 | - |
-| `--fix_sample_size` | 固定画布 `[height, width]`，覆盖方形 `--video_sample_size`（等比缩放后中心裁剪） | `None` |
+| `--random_hw_adapt` | 自动缩放视频到一定范围内的随机尺寸 | - |
 | `--low_vram` | 将 teacher VAE 放在 CPU，仅在编码时搬上 GPU | - |
 | `--trainable_modules` | 可训练模块（`"."` 表示所有模块） | `"."` |
 | `--trainable_modules_low_learning_rate` | 以 lr/2 训练的模块 | `[]` |
 | `--resume_from_checkpoint` | 恢复训练路径，使用 `"latest"` 自动选择最新 checkpoint | `latest` |
 | `--validation_steps` / `--validation_epochs` | 每 N 步 / 每 N 个 epoch 执行一次验证 | 2000 / 5 |
 | `--validation_paths` | 用于重建对比验证的视频路径 | `"asset/1.mp4"` |
-| `--validation_n_frames` | 验证片段帧数（4k+1），不传则复用 `--video_sample_n_frames` | `None` |
 
 **Sample Size 配置指南**：
-- `video_sample_size` 设置方形训练画布；配合 `--enable_bucket` 可按原始宽高比训练。
-- `--fix_sample_size 720 1280` 固定非方形画布（等比缩放 + 中心裁剪）；配合 `--enable_bucket` 时按批生效，`video_sample_size` 会提升到 `max(fix_sample_size)`。
+- `video_sample_size` 表示训练分辨率；当启用 `random_hw_adapt` 时，表示分辨率的最小值。
 - `video_sample_n_frames` 必须满足 `4k+1`（如 33、49、81），因为完整 VAE 与 TAE 都是因果 4 倍时间压缩。
-- `--video_sample_strides 1 2 3` 会为每个片段随机抽取一个时间步幅（多尺度时间增广）；步幅越大请求的帧数越少，因此单个片段覆盖的时间跨度基本不变。
 
 ### 3.4 训练验证
 
@@ -283,18 +279,15 @@ export TAE_PATH="taew2_1.safetensors"
 | `--validation_steps` | 每 N 步执行一次验证 | 2000 |
 | `--validation_epochs` | 每 N 个 epoch 执行一次验证 | 5 |
 | `--validation_paths` | 验证视频路径 | `"asset/1.mp4"` |
-| `--validation_n_frames` | 验证片段帧数，**必须为 4k+1**；不传则复用 `--video_sample_n_frames` | 81 |
 
 ```bash
   --validation_paths "asset/1.mp4" \
   --validation_steps=2000 \
-  --validation_epochs=5 \
-  --validation_n_frames=81
+  --validation_epochs=5
 ```
 
 **注意事项**：
 - 验证会为每个样本在 `output_dir/sample/` 保存三个视频：`*_input.mp4`（缩放后的输入）、`*_taehv.mp4`（TAE 重建）、`*_fullvae.mp4`（完整 VAE 重建，供参考对比）。
-- 验证片段会先截取 `min(视频长度, --validation_n_frames)` 帧、再向下对齐到 `4k+1`，因此把该值设大不会在短视频素材上失败。
 - 启用 `--use_ema` 时，验证使用 EMA 权重。
 
 ### 3.5 训练技巧
@@ -304,7 +297,7 @@ export TAE_PATH="taew2_1.safetensors"
 - **Decoder-only 蒸馏**：加上 `--freeze_tae_encoder` 只提升解码质量。
 - **显存**：teacher 完整 VAE 是主要的显存消耗者。使用 `--low_vram` 可以将其放在 CPU、仅在编码时搬上 GPU；也可以降低 `--video_sample_n_frames` / `--video_sample_size`。
 - **串行 TAE**：`--use_taehv_sequential` 以速度换显存，激活显存与视频长度无关（O(1)）。
-- **EMA**：推荐开启 `--use_ema` 作为最终交付物；最终保存的 `taehv` 目录里是 EMA 权重。短训练希望均值更快跟上时，可传 `--ema_decay`（如 `train_taehv_super.sh` 用的 `0.999`）替代默认的 0.9999。
+- **EMA**：推荐开启 `--use_ema` 作为最终交付物；最终保存的 `taehv` 目录里是 EMA 权重。
 
 ### 3.6 多机分布式训练
 

scripts/taehv/train_taehv.py

@@ -17,10 +17,6 @@ The teacher VAE family is selected by --config_path:
     config/wan2.1/wan_civitai.yaml                -> AutoencoderKLWan    (16ch, taew2_1)
     config/wan2.2/wan_civitai_5b.yaml             -> AutoencoderKLWan3_8 (48ch, taew2_2)
 (any wan2.2 *_2.2vae.yaml also works, they share the same vae_kwargs)
-
-The training loop mirrors scripts/wan2.1_fun/train.py (derived from
-https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image.py):
-same shared argument surface, dataset/sampler, EMA and checkpoint hooks.
 """
 #!/usr/bin/env python
 # coding=utf-8
@@ -38,8 +34,8 @@ same shared argument surface, dataset/sampler, EMA and checkpoint hooks.
 # See the License for the specific language governing permissions and
 
 import argparse
-import copy
 import gc
+import inspect
 import logging
 import math
 import os
@@ -65,6 +61,7 @@ from einops import rearrange
 from omegaconf import OmegaConf
 from packaging import version
 from PIL import Image
+from torch.utils.data import RandomSampler
 from torch.utils.tensorboard import SummaryWriter
 from torchvision import transforms
 from transformers.utils import ContextManagers
@@ -94,6 +91,13 @@ check_min_version("0.18.0.dev0")
 logger = get_logger(__name__, log_level="INFO")
 
 
+def filter_kwargs(cls, kwargs):
+    """Filter kwargs to only include parameters accepted by the class __init__."""
+    sig = inspect.signature(cls.__init__)
+    valid_keys = set(sig.parameters.keys()) - {"self"}
+    return {k: v for k, v in kwargs.items() if k in valid_keys}
+
+
 def linear_decay(initial_value, final_value, total_steps, current_step):
     if current_step >= total_steps:
         return final_value
@@ -105,13 +109,8 @@ def linear_decay(initial_value, final_value, total_steps, current_step):
 
 def log_validation(vae, tae, args, accelerator, weight_dtype, global_step):
     """Validation: encode video -> decode with TAE / full VAE, save comparison videos."""
-    # The guard sits *inside* the contexts on purpose: `Accelerator.autocast` enters and
-    # exits AMP without a try/finally, so an exception escaping the block strands autocast
-    # enabled and leaves its cast-weight cache holding tensors built under no_grad. The
-    # next training step would then produce a graph-free loss and `accelerator.backward`
-    # would die with "element 0 of tensors does not require grad".
-    with torch.no_grad(), accelerator.autocast():
-        try:
+    try:
+        with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype):
             logger.info("Running validation...")
 
             if args.validation_paths is None or len(args.validation_paths) == 0:
@@ -121,10 +120,6 @@ def log_validation(vae, tae, args, accelerator, weight_dtype, global_step):
             from decord import VideoReader
 
             unwrapped_tae = accelerator.unwrap_model(tae)
-            # Under --low_vram the teacher VAE waits on CPU between encoding steps.
-            teacher_on_gpu = bool(args.low_vram)
-            if teacher_on_gpu:
-                vae.to(accelerator.device, dtype=weight_dtype)
 
             for i, video_path in enumerate(args.validation_paths):
                 if not os.path.exists(video_path):
@@ -133,7 +128,7 @@ def log_validation(vae, tae, args, accelerator, weight_dtype, global_step):
 
                 # Load video frames
                 vr = VideoReader(video_path)
-                num_frames = min(len(vr), args.validation_n_frames)
+                num_frames = min(len(vr), args.video_sample_n_frames)
                 # Align to temporal compression ratio (4k+1)
                 temporal_ratio = vae.config.temporal_compression_ratio
                 num_frames = (num_frames - 1) // temporal_ratio * temporal_ratio + 1
@@ -149,8 +144,8 @@ def log_validation(vae, tae, args, accelerator, weight_dtype, global_step):
 
                 # Resize to target size aligned to the spatial compression ratio
                 spatial_ratio = int(vae.config.spatial_compression_ratio)
-                h_target = int(args.video_sample_height / spatial_ratio) * spatial_ratio
-                w_target = int(args.video_sample_width / spatial_ratio) * spatial_ratio
+                h_target = int(args.video_sample_size / spatial_ratio) * spatial_ratio
+                w_target = h_target  # square for simplicity in validation
                 pixel_values = F.interpolate(
                     pixel_values, size=(h_target, w_target), mode='bilinear', align_corners=False
                 )
@@ -166,33 +161,29 @@ def log_validation(vae, tae, args, accelerator, weight_dtype, global_step):
                 z_full = vae.encode(pixel_values).latent_dist.mode()
                 decoded_full = vae.decode(z_full, return_dict=False)[0]
 
-                # Save videos; save_videos_grid goes through numpy, which has no bf16.
+                # Save videos
                 os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
                 save_videos_grid(
-                    pixel_values.float(),
+                    pixel_values,
                     os.path.join(args.output_dir, f"sample/step{global_step}_val{i}_input.mp4"),
                     rescale=True, fps=24
                 )
                 save_videos_grid(
-                    decoded_tae.float(),
+                    decoded_tae,
                     os.path.join(args.output_dir, f"sample/step{global_step}_val{i}_taehv.mp4"),
                     rescale=True, fps=24
                 )
                 save_videos_grid(
-                    decoded_full.float(),
+                    decoded_full,
                     os.path.join(args.output_dir, f"sample/step{global_step}_val{i}_fullvae.mp4"),
                     rescale=True, fps=24
                 )
                 logger.info(f"Saved validation video {i} at step {global_step}")
 
-            if teacher_on_gpu:
-                vae.to("cpu")
-                torch.cuda.empty_cache()
-
-        except Exception as e:
-            gc.collect()
-            torch.cuda.empty_cache()
-            print(f"Eval error on rank {accelerator.process_index} with info {e}")
+    except Exception as e:
+        gc.collect()
+        torch.cuda.empty_cache()
+        print(f"Eval error on rank {accelerator.process_index} with info {e}")
 
 
 def parse_args():
@@ -227,14 +218,14 @@ def parse_args():
     parser.add_argument(
         "--train_data_dir",
         type=str,
-        default="datasets/X-Fun-Videos-Demo",
+        default=None,
         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."),
+        default=None,
+        help=("A json/csv containing the training data meta."),
     )
     parser.add_argument(
         "--max_train_samples",
@@ -264,9 +255,7 @@ def parse_args():
         default=None,
         help="The directory where the downloaded models and datasets will be stored.",
     )
-    # The default must be a concrete int: the batch-sampler generators and `worker_init_fn`
-    # below call `manual_seed(args.seed)` / `args.seed + process_index` without a None guard.
-    parser.add_argument("--seed", type=int, default=42, help="A seed for reproducible training.")
+    parser.add_argument("--seed", type=int, default=None, 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."
     )
@@ -329,12 +318,6 @@ def parse_args():
         ),
     )
     parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model.")
-    parser.add_argument(
-        "--ema_decay",
-        type=float,
-        default=None,
-        help="EMA decay for the tracked parameters. Defaults to the EMAModel default (0.9999).",
-    )
     parser.add_argument(
         "--non_ema_revision",
         type=str,
@@ -406,12 +389,6 @@ def parse_args():
         default=2000,
         help="Run validation every X steps.",
     )
-    parser.add_argument(
-        "--validation_n_frames",
-        type=int,
-        default=None,
-        help="Validation clip length, 4k+1. Defaults to `--video_sample_n_frames`.",
-    )
     parser.add_argument(
         "--tracker_project_name",
         type=str,
@@ -421,40 +398,36 @@ def parse_args():
     parser.add_argument(
         "--enable_bucket", action="store_true", help="Whether enable bucket sample in datasets."
     )
+    parser.add_argument(
+        "--random_ratio_crop", action="store_true", help="Whether enable random ratio crop sample in datasets."
+    )
+    parser.add_argument(
+        "--random_hw_adapt", action="store_true", help="Whether enable random adapt height and width in datasets."
+    )
     parser.add_argument(
         "--video_sample_size",
         type=int,
         default=512,
         help="Training video sample size.",
     )
-    parser.add_argument(
-        "--fix_sample_size",
-        nargs=2,
-        type=int,
-        default=None,
-        help="Fix Sample size [height, width] to override `--video_sample_size` with a fixed "
-             "non-square shape. Under `--enable_bucket` the scalar `--video_sample_size` is bumped to "
-             "at least max(fix_sample_size), so the bucket ratios still cover the canvas.",
-    )
     parser.add_argument(
         "--video_sample_stride",
         type=int,
         default=1,
         help="Sample stride of the video.",
     )
-    parser.add_argument(
-        "--video_sample_strides",
-        nargs="+",
-        type=int,
-        default=None,
-        help="Optional per-video random sampling strides; one is drawn for every clip.",
-    )
     parser.add_argument(
         "--video_sample_n_frames",
         type=int,
         default=33,
         help="Num frame of video. Must be 4k+1 to match the causal temporal compression.",
     )
+    parser.add_argument(
+        "--video_repeat",
+        type=int,
+        default=0,
+        help="Num of repeat video.",
+    )
     parser.add_argument(
         "--tae_path",
         type=str,
@@ -539,45 +512,10 @@ def parse_args():
     if args.non_ema_revision is None:
         args.non_ema_revision = args.revision
 
-    # Derive the effective canvas once, so every downstream use agrees.
-    if args.fix_sample_size is not None:
-        args.video_sample_height, args.video_sample_width = args.fix_sample_size
-        # The bucket sampler groups by the scalar sample size; keep it >= the fixed canvas.
-        if args.enable_bucket:
-            args.video_sample_size = max(max(args.fix_sample_size), args.video_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 args.validation_n_frames is None:
-        args.validation_n_frames = args.video_sample_n_frames
     # Both the full VAE and the TAE are causal 4x temporal compressors, so the
     # clip length must be 4k+1 for the decode output to match the input.
-    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.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) < 1:
-        parser.error("Epochs, batch sizes, checkpoint interval and stride must be positive")
-    if min(args.dataloader_num_workers, args.validation_steps, args.validation_epochs, args.lr_warmup_steps) < 0:
-        parser.error("Workers, validation intervals 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 args.ema_decay is not None and not 0 <= args.ema_decay < 1:
-        parser.error("ema_decay must be in [0, 1)")
-    if not 0 <= args.adam_beta1 < 1 or not 0 <= args.adam_beta2 < 1:
-        parser.error("Invalid Adam betas")
-    if args.adam_weight_decay < 0:
-        parser.error("Adam weight decay cannot be negative")
-    if any(not math.isfinite(value) or value <= 0 for value in
-           (args.learning_rate, args.adam_epsilon, args.max_grad_norm)):
-        parser.error("Learning rate, epsilon and gradient norm must be finite and positive")
-    if args.latent_loss_weight < 0:
-        parser.error("latent_loss_weight cannot be negative")
+    assert (args.video_sample_n_frames - 1) % 4 == 0, \
+        f"video_sample_n_frames must be 4k+1, got {args.video_sample_n_frames}"
 
     return args
 
@@ -714,9 +652,9 @@ def main():
             arch_variant=args.tae_arch_variant,
         )
     tae.parallel = not args.use_taehv_sequential
-    # Trainable parameters (and the Adam moments built below) stay in FP32 like the
-    # other training scripts; mixed precision comes from `accelerator.autocast()`.
-    tae.float()
+    # The TAE is tiny (~20MB); keep it fully in weight_dtype so forward works
+    # regardless of whether autocast is active.
+    tae.to(weight_dtype)
 
     # Set trainable parameters
     tae.requires_grad_(False)
@@ -749,17 +687,10 @@ def main():
             # every local shard of the copy pairs 1:1 with the live shard.
             ema_tae = FSDPEMA(ema_module, source=tae, accelerator=accelerator, fsdp_plugin=fsdp_plugin)
         else:
-            ema_module = ema_module.float()
+            ema_module = ema_module.to(weight_dtype)
             if args.tae_path is not None:
                 ema_module.load_state_dict(tae.state_dict())
-            # An explicit decay also disables the warmup ramp, mirroring train_taehv_super.py.
-            ema_kwargs = {} if args.ema_decay is None else {
-                "decay": args.ema_decay, "min_decay": args.ema_decay, "use_ema_warmup": False,
-            }
-            ema_tae = EMAModel(
-                ema_module.parameters(), model_cls=AutoencoderTinyWan, model_config=ema_module.config,
-                **ema_kwargs,
-            )
+            ema_tae = EMAModel(ema_module.parameters(), model_cls=AutoencoderTinyWan, model_config=ema_module.config)
 
     # ==================== Save/Load Hooks ====================
     if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
@@ -891,8 +822,7 @@ def main():
 
     train_dataset = VideoDataset(
         args.train_data_meta, args.train_data_dir,
-        sample_size=(args.video_sample_height, args.video_sample_width),
-        sample_stride=args.video_sample_strides or args.video_sample_stride,
+        sample_size=args.video_sample_size, sample_stride=args.video_sample_stride,
         sample_n_frames=args.video_sample_n_frames,
         enable_bucket=args.enable_bucket, enable_inpaint=False,
     )
@@ -914,9 +844,6 @@ def main():
             aspect_ratios=aspect_ratio_sample_size,
         )
 
-        # A fixed canvas overrides the closest bucket size, so every batch is uniform.
-        fix_size = None if args.fix_sample_size is None else [args.video_sample_height, args.video_sample_width]
-
         def collate_fn(examples):
             new_examples = {}
             new_examples["pixel_values"] = []
@@ -924,13 +851,9 @@ def main():
             pixel_value = examples[0]["pixel_values"]
             f, h, w, c = np.shape(pixel_value)
 
-            if fix_size is None:
-                closest_size, closest_ratio = get_closest_ratio(h, w, ratios=aspect_ratio_sample_size)
-                spatial_ratio = int(vae.config.spatial_compression_ratio)
-                closest_size = [int(x / spatial_ratio) * spatial_ratio for x in closest_size]
-                target_size = closest_size
-            else:
-                target_size = fix_size
+            closest_size, closest_ratio = get_closest_ratio(h, w, ratios=aspect_ratio_sample_size)
+            spatial_ratio = int(vae.config.spatial_compression_ratio)
+            closest_size = [int(x / spatial_ratio) * spatial_ratio for x in closest_size]
 
             min_example_length = min([example["pixel_values"].shape[0] for example in examples])
             batch_video_length = int(min(args.video_sample_n_frames + sample_n_frames_bucket_interval, min_example_length))
@@ -942,19 +865,14 @@ def main():
                 pixel_values = torch.from_numpy(example["pixel_values"]).permute(0, 3, 1, 2).contiguous()
                 pixel_values = pixel_values / 255.0
 
-                if fix_size is None:
-                    if closest_size[0] / h > closest_size[1] / w:
-                        resize_size = closest_size[0], int(w * closest_size[0] / h)
-                    else:
-                        resize_size = int(h * closest_size[1] / w), closest_size[1]
+                if closest_size[0] / h > closest_size[1] / w:
+                    resize_size = closest_size[0], int(w * closest_size[0] / h)
                 else:
-                    # Aspect-preserving cover resize, same as scripts/taehv/train_taehv_super.py.
-                    scale = max(fix_size[0] / h, fix_size[1] / w)
-                    resize_size = math.ceil(h * scale), math.ceil(w * scale)
+                    resize_size = int(h * closest_size[1] / w), closest_size[1]
 
                 transform = transforms.Compose([
                     transforms.Resize(resize_size, interpolation=transforms.InterpolationMode.BILINEAR),
-                    transforms.CenterCrop(target_size),
+                    transforms.CenterCrop(closest_size),
                     transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
                 ])
                 pixel_values = transform(pixel_values)[:batch_video_length]
@@ -982,11 +900,6 @@ def main():
             worker_init_fn=worker_init_fn(args.seed + accelerator.process_index)
         )
 
-    # Snapshot the sampler's empty bucket layout so every epoch can restore it. The shape is
-    # sampler-specific -- ImageVideoSampler holds flat lists while AspectRatioBatchImageVideoSampler
-    # holds one list per aspect ratio -- so the initial value is copied rather than rebuilt.
-    initial_bucket = copy.deepcopy(batch_sampler.bucket)
-
     # ==================== LR Scheduler ====================
     overrode_max_train_steps = False
     num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
@@ -1076,18 +989,8 @@ def main():
     )
 
     for epoch in range(first_epoch, args.num_train_epochs):
-        if global_step >= args.max_train_steps:
-            break
         train_loss = 0.0
-        train_dataloader.set_epoch(epoch)
         batch_sampler.sampler.generator = torch.Generator().manual_seed(args.seed + epoch)
-        # Indices that never filled a batch stay behind in `bucket` and would leak into the next
-        # epoch, silently changing the batch composition, so the layout is restored here.
-        # `_pos_start` is deliberately NOT reset: RandomSampler.__iter__ already zeroes it at the
-        # end of a full pass, and this script's resume relies on load_model_hook writing the saved
-        # position back into it. train_taehv_super.py can reset it unconditionally only because it
-        # pairs that with `resume_step` + accelerator.skip_first_batches, which this script lacks.
-        batch_sampler.bucket = copy.deepcopy(initial_bucket)
         for step, batch in enumerate(train_dataloader):
             with accelerator.accumulate(tae):
                 pixel_values = batch["pixel_values"].to(weight_dtype)  # [B, F, C, H, W]
@@ -1112,11 +1015,10 @@ def main():
                     vae.to('cpu')
                     torch.cuda.empty_cache()
 
-                # 2. Student forward: TAE encode -> decode (FP32 weights, autocast activations)
+                # 2. Student forward: TAE encode -> decode
                 # Use the unwrapped model since encode/decode are not exposed by DDP.
-                with accelerator.autocast():
-                    z_tae = unwrap_model(tae).encode(pixel_values).latent_dist.mode()
-                    decoded = unwrap_model(tae).decode(z_tae).sample
+                z_tae = unwrap_model(tae).encode(pixel_values).latent_dist.mode()
+                decoded = unwrap_model(tae).decode(z_tae).sample
                 # Trim to the input temporal span in case the TAE padded extra frames.
                 decoded = decoded[:, :, :pixel_values.shape[2]]
 
@@ -1227,12 +1129,6 @@ def main():
                 log_validation(vae, tae, args, accelerator, weight_dtype, global_step)
                 if args.use_ema:
                     ema_tae.restore(tae.parameters())
-        # Epoch-level break, mirroring train_taehv_super.py: the `break` inside the step loop only
-        # ends that epoch, so without this the remaining epochs would each run one extra step past
-        # max_train_steps and write a spurious checkpoint. Placed after the epoch validation so the
-        # final epoch still gets validated before exiting.
-        if global_step >= args.max_train_steps:
-            break
 
     # Close the bar before the end-of-run checkpoint: tqdm keeps redrawing a live bar whenever
     # something else writes to the console. PauseAwareTqdm.close() rebases the closing line onto

scripts/taehv/train_taehv.sh

@@ -1,6 +1,6 @@
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-TI2V-5B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
-export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata.json"
 export TAE_PATH="models/Diffusion_Transformer/taew2_2.safetensors"
 # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
 # export NCCL_IB_DISABLE=1
@@ -9,36 +9,30 @@ NCCL_DEBUG=INFO
 
 accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv.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=49 \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --tae_path=$TAE_PATH \
+  --train_data_dir=$DATASET_NAME \
+  --train_data_meta=$DATASET_META_NAME \
+  --video_sample_size=512 \
+  --video_sample_stride=1 \
+  --video_sample_n_frames=33 \
   --train_batch_size=1 \
   --gradient_accumulation_steps=1 \
-  --dataloader_num_workers=8 \
-  --max_train_steps=10000 \
-  --checkpointing_steps=1000 \
-  --validation_steps=1000 \
-  --validation_paths "asset/inpaint_video.mp4" \
+  --dataloader_num_workers=4 \
+  --num_train_epochs=100 \
+  --checkpointing_steps=500 \
   --learning_rate=1e-04 \
   --lr_scheduler="constant_with_warmup" \
-  --lr_warmup_steps=200 \
+  --lr_warmup_steps=100 \
   --latent_loss_weight=1.0 \
-  --use_ema \
-  --ema_decay=0.999 \
   --seed=42 \
   --output_dir="output_dir_taehv_w2.2" \
   --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 \
+  --random_hw_adapt \
   --enable_bucket \
   --trainable_modules "." \
   --resume_from_checkpoint=latest

scripts/taehv/train_taehv_super.py

@@ -1,950 +0,0 @@
-"""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="Run the TAE in sequential (O(1)-memory) mode instead of parallel.",
-    )
-    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="Fix Sample size [height, width] to override `--video_sample_size` with a fixed "
-             "non-square shape.",
-    )
-    parser.add_argument(
-        "--video_sample_strides",
-        nargs="+",
-        type=int,
-        default=None,
-        help="Optional per-video random sampling strides; one is drawn for every clip.",
-    )
-    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()

scripts/taehv/train_taehv_super.sh

@@ -1,50 +0,0 @@
-#!/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

videox_fun/data/__init__.py

@@ -19,4 +19,4 @@ from .utils import (VIDEO_READER_TIMEOUT, Camera, VideoReader_contextmanager,
                     prepare_lingbot_dit_cond_dict,
                     prepare_lingbot_dit_cond_dict_from_c2ws, process_pose_file,
                     process_pose_params, ray_condition, resize_frame,
-                    resize_image_with_target_area, sample_clip_start)
+                    resize_image_with_target_area)

videox_fun/data/dataset_image_video.py

@@ -30,7 +30,7 @@ except ImportError:
 from .utils import (VIDEO_READER_TIMEOUT, VideoReader_contextmanager,
                     get_random_mask, get_video_reader_batch, padding_image,
                     process_pose_file, process_pose_params, resize_frame,
-                    resize_image_with_target_area, sample_clip_start)
+                    resize_image_with_target_area)
 
 
 class ImageVideoSampler(BatchSampler):
@@ -99,7 +99,6 @@ class ImageVideoDataset(Dataset):
         enable_inpaint=False,
         inpaint_mask_fill_value=0,
         return_file_name=False,
-        random_start=True,
     ):
         # Loading annotations from files
         print(f"loading annotations from {ann_path} ...")
@@ -137,8 +136,6 @@ class ImageVideoDataset(Dataset):
 
         self.video_length_drop_start = video_length_drop_start
         self.video_length_drop_end = video_length_drop_end
-        # Whether the clip start is drawn randomly (training) or pinned (validation/eval); see `sample_clip_start`.
-        self.random_start = random_start
 
         # Video params: resize, center crop, normalize to [-1, 1]
         self.video_sample_stride    = video_sample_stride
@@ -191,13 +188,10 @@ class ImageVideoDataset(Dataset):
                         f"Video too short: sampled {min_sample_n_frames} frames < required {min_video_sample_n_frames}."
                     )
 
-                # Select contiguous clip; the start is random unless this is an eval dataset
+                # 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.video_sample_stride + 1)
-                start_idx = sample_clip_start(
-                    video_length, clip_length,
-                    drop_start=self.video_length_drop_start, random_start=self.random_start,
-                )
+                start_idx   = random.randint(int(self.video_length_drop_start * video_length), video_length - clip_length) if video_length != clip_length else 0
                 batch_index = np.linspace(start_idx, start_idx + clip_length - 1, min_sample_n_frames, dtype=int)
 
                 try:
@@ -369,11 +363,7 @@ class LingbotImageVideoDataset(ImageVideoDataset):
 
                 video_length = int(self.video_length_drop_end * len(video_reader))
                 clip_length = min(video_length, (min_sample_n_frames - 1) * self.video_sample_stride + 1)
-                # The start is random unless this is an eval dataset
-                start_idx = sample_clip_start(
-                    video_length, clip_length,
-                    drop_start=self.video_length_drop_start, random_start=self.random_start,
-                )
+                start_idx = random.randint(int(self.video_length_drop_start * video_length), video_length - clip_length) if video_length != clip_length else 0
                 batch_index = np.linspace(start_idx, start_idx + clip_length - 1, min_sample_n_frames, dtype=int)
 
                 try:
@@ -513,7 +503,6 @@ class ImageVideoControlDataset(Dataset):
         enable_subject_info=False,
         padding_subject_info=True,
         return_file_name=False,
-        random_start=True,
     ):
         # Loading annotations from files
         print(f"loading annotations from {ann_path} ...")
@@ -554,8 +543,6 @@ class ImageVideoControlDataset(Dataset):
 
         self.video_length_drop_start = video_length_drop_start
         self.video_length_drop_end = video_length_drop_end
-        # Whether the clip start is drawn randomly (training) or pinned (validation/eval); see `sample_clip_start`.
-        self.random_start = random_start
 
         # Video params: resize, center crop, normalize to [-1, 1]
         self.video_sample_stride    = video_sample_stride
@@ -616,13 +603,10 @@ class ImageVideoControlDataset(Dataset):
                         f"Video too short: sampled {min_sample_n_frames} frames < required {min_video_sample_n_frames}."
                     )
 
-                # Select contiguous clip; the start is random unless this is an eval dataset
+                # 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.video_sample_stride + 1)
-                start_idx = sample_clip_start(
-                    video_length, clip_length,
-                    drop_start=self.video_length_drop_start, random_start=self.random_start,
-                )
+                start_idx   = random.randint(int(self.video_length_drop_start * video_length), video_length - clip_length) if video_length != clip_length else 0
                 batch_index = np.linspace(start_idx, start_idx + clip_length - 1, min_sample_n_frames, dtype=int)
 
                 try:

videox_fun/data/dataset_video.py

@@ -34,8 +34,7 @@ except ImportError:
     from .utils import AVVideoReader as VideoReader
 
 from .utils import (VIDEO_READER_TIMEOUT, VideoReader_contextmanager,
-                    get_random_mask, get_video_reader_batch, resize_frame,
-                    sample_clip_start)
+                    get_random_mask, get_video_reader_batch, resize_frame)
 
 
 def load_audio(path, sr, mono=True, native_sr=False, res_type=None):
@@ -88,7 +87,6 @@ class WebVid10M(Dataset):
         enable_bucket=False, 
         enable_inpaint=False, 
         is_image=False,
-        random_start=True,
     ):
         print(f"loading annotations from {csv_path} ...")
         with open(csv_path, 'r') as csvfile:
@@ -102,7 +100,6 @@ class WebVid10M(Dataset):
         self.enable_bucket   = enable_bucket
         self.enable_inpaint  = enable_inpaint
         self.is_image        = is_image
-        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([
@@ -121,10 +118,10 @@ class WebVid10M(Dataset):
         
         if not self.is_image:
             clip_length = min(video_length, (self.sample_n_frames - 1) * self.sample_stride + 1)
-            start_idx   = sample_clip_start(video_length, clip_length, random_start=self.random_start)
+            start_idx   = random.randint(0, video_length - clip_length)
             batch_index = np.linspace(start_idx, start_idx + clip_length - 1, self.sample_n_frames, dtype=int)
         else:
-            batch_index = [sample_clip_start(video_length, 1, random_start=self.random_start)]
+            batch_index = [random.randint(0, video_length - 1)]
 
         if not self.enable_bucket:
             pixel_values = torch.from_numpy(video_reader.get_batch(batch_index).asnumpy()).permute(0, 3, 1, 2).contiguous()
@@ -176,7 +173,6 @@ 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} ...")
@@ -193,7 +189,6 @@ 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(
@@ -217,23 +212,17 @@ 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) - 1)
-                    // sample_stride) + 1
+                int(len(video_reader) * (self.video_length_drop_end - self.video_length_drop_start) // self.sample_stride)
             )
             if min_sample_n_frames == 0:
                 raise ValueError(f"No Frames in video.")
 
-            # Select contiguous clip; the start is random unless this is an eval dataset
+            # 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) * sample_stride + 1)
-            start_idx = sample_clip_start(
-                video_length, clip_length,
-                drop_start=self.video_length_drop_start, random_start=self.random_start,
-            )
+            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
             batch_index = np.linspace(start_idx, start_idx + clip_length - 1, min_sample_n_frames, dtype=int)
 
             try:
@@ -314,7 +303,6 @@ class VideoSpeechDataset(Dataset):
         audio_stereo=False,
         audio_span_includes_last_frame=False,
         enable_ref2va=False,
-        random_start=True,
     ):
         # Loading annotations from files
         print(f"loading annotations from {ann_path} ...")
@@ -335,8 +323,6 @@ class VideoSpeechDataset(Dataset):
         # Video params: resize, center crop, normalize to [-1, 1]
         self.video_sample_stride = video_sample_stride
         self.video_sample_n_frames = video_sample_n_frames
-        # Whether the clip start is drawn randomly (training) or pinned (validation/eval); see `sample_clip_start`.
-        self.random_start = random_start
         # Fewest sampled frames a clip has to yield to be usable. A model whose VAE only encodes certain frame
         # counts (MiniMax-H3 needs `17n + 5`, so at least 5) raises this so that clips which cannot fill one chunk
         # are skipped by `__getitem__`'s retry instead of reaching the collate as a short batch. The default of 1 is
@@ -451,11 +437,10 @@ class VideoSpeechDataset(Dataset):
                     "run needs; skipping."
                 )
 
-            # Select the starting frame; pinned when random_start=False so eval clips are reproducible
+            # Randomly select the starting frame
             frame_span = (actual_n_frames - 1) * frame_step
-            start_frame = sample_clip_start(
-                total_frames, int(math.ceil(frame_span)) + 1, random_start=self.random_start
-            )
+            max_start = total_frames - 1 - int(math.ceil(frame_span))
+            start_frame = random.randint(0, max_start) if max_start > 0 else 0
             frame_indices = [
                 min(total_frames - 1, int(round(start_frame + index * frame_step)))
                 for index in range(actual_n_frames)
@@ -760,7 +745,6 @@ class VideoSpeechControlDataset(Dataset):
         audio_native_sr_resample=False,
         audio_stereo=False,
         audio_span_includes_last_frame=False,
-        random_start=True,
     ):
         # Loading annotations from files
         print(f"loading annotations from {ann_path} ...")
@@ -781,8 +765,6 @@ class VideoSpeechControlDataset(Dataset):
         # Video params: resize, center crop, normalize to [-1, 1]
         self.video_sample_stride = video_sample_stride
         self.video_sample_n_frames = video_sample_n_frames
-        # Whether the clip start is drawn randomly (training) or pinned (validation/eval); see `sample_clip_start`.
-        self.random_start = random_start
         # Fewest sampled frames a clip has to yield to be usable; see `VideoSpeechDataset` for the rationale.
         self.min_video_sample_n_frames = max(1, int(min_video_sample_n_frames))
         # Frame rate the sampled clip has to land on; see `VideoSpeechDataset` for the rationale.
@@ -889,11 +871,10 @@ class VideoSpeechControlDataset(Dataset):
                     "run needs; skipping."
                 )
 
-            # Select the starting frame; pinned when random_start=False so eval clips are reproducible
+            # Randomly select the starting frame
             frame_span = (actual_n_frames - 1) * frame_step
-            start_frame = sample_clip_start(
-                total_frames, int(math.ceil(frame_span)) + 1, random_start=self.random_start
-            )
+            max_start = total_frames - 1 - int(math.ceil(frame_span))
+            start_frame = random.randint(0, max_start) if max_start > 0 else 0
             frame_indices = [
                 min(total_frames - 1, int(round(start_frame + index * frame_step)))
                 for index in range(actual_n_frames)
@@ -1159,7 +1140,6 @@ class VideoAnimateDataset(Dataset):
         video_length_drop_start=0.1, 
         video_length_drop_end=0.9,
         return_file_name=False,
-        random_start=True,
     ):
         # Loading annotations from files
         print(f"loading annotations from {ann_path} ...")
@@ -1195,8 +1175,6 @@ class VideoAnimateDataset(Dataset):
 
         self.video_length_drop_start = video_length_drop_start
         self.video_length_drop_end = video_length_drop_end
-        # Whether the clip start is drawn randomly (training) or pinned (validation/eval); see `sample_clip_start`.
-        self.random_start = random_start
 
         # Video params: resize, center crop, normalize to [-1, 1]
         self.video_sample_stride = video_sample_stride
@@ -1232,13 +1210,10 @@ class VideoAnimateDataset(Dataset):
             if min_sample_n_frames == 0:
                 raise ValueError(f"No Frames in video.")
 
-            # Select contiguous clip; the start is random unless this is an eval dataset
+            # 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.video_sample_stride + 1)
-            start_idx = sample_clip_start(
-                video_length, clip_length,
-                drop_start=self.video_length_drop_start, random_start=self.random_start,
-            )
+            start_idx = random.randint(int(self.video_length_drop_start * video_length), video_length - clip_length) if video_length != clip_length else 0
             batch_index = np.linspace(start_idx, start_idx + clip_length - 1, min_sample_n_frames, dtype=int)
 
             try:

videox_fun/data/utils.py

@@ -82,28 +82,6 @@ class _AVBatchFrames:
     def asnumpy(self):
         return np.stack(self._frames)
 
-
-def sample_clip_start(video_length, clip_length, drop_start=0.0, random_start=True):
-    """Pick the first frame index of a contiguous clip.
-
-    `video_length` is the usable tail boundary (already scaled by `video_length_drop_end`) and
-    `clip_length` the span the clip covers, so `video_length - clip_length` is the last valid start.
-    With `random_start` the start is drawn uniformly over every valid position, which is the temporal
-    augmentation training sets want; with `random_start=False` it is pinned to the
-    `video_length_drop_start` boundary, so validation/eval datasets rebuild the very same clip on
-    every pass and their metrics stay comparable across steps. The head boundary is clamped to the
-    last valid start, hence a clip that already fills the usable range can never run past the video.
-
-    Note this only decides *where* a clip starts; how many frames it yields stays the caller's
-    business, since the dataset classes differ in that formula.
-    """
-    last_valid_start = max(video_length - clip_length, 0)
-    start_idx = min(int(drop_start * video_length), last_valid_start)
-    if random_start and last_valid_start > start_idx:
-        start_idx = random.randint(start_idx, last_valid_start)
-    return start_idx
-
-
 def get_random_mask(shape, image_start_only=False):
     f, c, h, w = shape
     mask = torch.zeros((f, 1, h, w), dtype=torch.uint8)

videox_fun/models/wan_tae.py

@@ -329,11 +329,6 @@ class AutoencoderTinyWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
     - taew2_2*.pth / .safetensors: Wan 2.2 5B (48ch latent, patch_size=2)
     """
 
-    # The TAE has no checkpointed block path, but the flag is declared so that
-    # `enable_gradient_checkpointing()` (a shared training-script call) is a
-    # no-op instead of raising, mirroring AutoencoderKLWan.
-    _supports_gradient_checkpointing = True
-
     # per-channel statistics of the full-size Wan VAE latent spaces
     WAN2_1_MEAN = [-0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508,
                    0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921]
@@ -454,10 +449,6 @@ 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):