Refactor Wan Model Training & Add Wan-VACE Training Support#352

Open

ninatu wants to merge 9 commits intomainfrom

ninatu/wan_training

Commits on Mar 11, 2026

Update wan configs for training

ninatu
and
martinarroyo
committed
Wan training: Resolve training mode bug with dropout and layer_forward

ninatu
and
martinarroyo
committed
Wan training: use learning rate from config

ninatu
and
martinarroyo
committed
Fix: Ensure prepare_sample_fn is used for 'tfrecord' dataset type

ninatu
and
martinarroyo
committed
Wan training: Set default dropout to 0.0 in Wan configs

ninatu
and
martinarroyo
committed
Wan 2.1 training: Resolve checkpoint loading issues with larger TPU slices and different topologies

ninatu
and
martinarroyo
committed
Wan training: Fix WAN training timestep sampling with continuous sampling and introduce disable_training_weights, add max_grad_norm and max_abs_grad logging.

ninatu
and
martinarroyo
committed
Abstract common WAN training components into BaseWanTrainer

ninatu
and
martinarroyo
committed
Add WAN-VACE training functionality

ninatu
and
martinarroyo
committed