nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[feat] World model training using third person games (#1443)

Browse Source 
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Mingjia Huo
2026-07-10 05:09:07 +00:00
committed by GitHub
co-authored by Claude Fable 5
parent af4a77553c
commit 4c08ffce49
37 changed files with 4618 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
assets/eval/worldmodel_synthetic_flow_calibration.json
+22
View File
@@ -0,0 +1,22 @@
{
"alpha_yaw": 0.08734091699186919,
"alpha_pitch": 0.08169667696275307,
"alpha_turn": 5.724587470723463e-17,
"beta_fwd": 0.02842768078408099,
"beta_strafe": 0.022531015077067108,
"focal_length": 457.0,
"frame_shape": [
352,
640
],
"calibrated_from": [
"1_wasd_only",
"camera",
"camera4hold_alpha1",
"fully_random",
"wasdonly_alpha1",
"wasd4holdrandview_simple_1key1mouse1"
],
"residual_rms": 15.890399609478676,
"n_equations": 4125232
}
docs/training/train_infra.md
+46
View File
@@ -161,6 +161,18 @@ training:
decay_interval_steps: 0
```
`training.data.data_path` can also mix multiple preprocessed datasets by using a mapping from dataset path to repeat count:
```yaml
training:
data:
data_path:
data/zeldam2-clean: 1
data/multi3d_games: 2
```
The repeat count duplicates that dataset's parquet file list before shuffling/sampling, so the example above trains with roughly twice as much `multi3d_games` exposure as `zeldam2-clean`. Paths are just suggested locations; use any local path that contains a FastVideo preprocessed parquet dataset.
See [Training Trackers](trackers.md) to configure Weights & Biases or SwanLab,
including SwanLab installation and authentication.
@@ -326,6 +338,40 @@ Self-Forcing inherits all DMD2 parameters, plus:
| `enable_gradient_in_rollout` | `true` | Enable backprop through rollout |
| `start_gradient_frame` | `0` | Frame index where gradients begin |
### Streaming Long Tuning
`StreamingLongTuningMethod` extends Self-Forcing for LongLive-style rollouts. It
keeps a streaming state, generates overlapping chunks, and trains only the new
frames while preserving context from earlier chunks.
For the MatrixGame2/Zelda world-model example, self-forcing and long tuning are
separate runs: first train or load the 1k-step self-forcing checkpoint using
`examples/train/scenario/worldmodel/zelda/self_forcing_causal_i2v.yaml`,
then run
`examples/train/scenario/worldmodel/zelda/streaming_long_tuning_causal_i2v.yaml`
from that checkpoint for the 3k-step streaming long-tuning stage.
```yaml
method:
_target_: fastvideo.train.methods.distribution_matching.streaming_long_tuning.StreamingLongTuningMethod
streaming_chunk_size: 9
streaming_max_length: 39
streaming_fixed_overlap_latents: 3
streaming_reencode_overlap_anchor: true
streaming_anchor_inject_k: 1
streaming_require_full_blocks: true
multi_phased_distill_schedule:
- stage: streaming_long
start_step: 0
end_step: 3000
num_latent_t: 39
streaming_training: true
```
See
`examples/train/scenario/worldmodel/zelda/streaming_long_tuning_causal_i2v.yaml`
for a complete MatrixGame2/Zelda configuration.
---
## Callbacks
examples/train/configs/distribution_matching/wan/self_forcing_causal_t2v.yaml
-3
View File
@@ -100,9 +100,6 @@ callbacks:
sampling_steps: [4]
sampling_timesteps: [1000, 750, 500, 250]
num_frames: 81
# Validation/inference uses standard CFG in both clean and Self-Forcing,
# so this directly matches Self-Forcing guidance_scale=3.0.
guidance_scale: 3.0
pipeline:
flow_shift: 5
examples/train/scenario/worldmodel/README.md
+96 -9
View File
@@ -1,32 +1,119 @@
# World-Model: Matrix-Game 2.0 I2V
Three training scenarios for the Matrix-Game 2.0 I2V world model on the
new YAML-driven trainer (`fastvideo/train/entrypoint/train.py`).
Training scenarios for the Matrix-Game 2.0 I2V world model on Solaris (Minecraft)
data and Zelda data, using the new YAML-driven trainer
(`fastvideo/train/entrypoint/train.py`).
## Solaris Configs
| Config | Method | Student | Notes |
|---|---|---|---|
| `finetune_i2v.yaml` | `FineTuneMethod` | `MatrixGame2Model` (bidirectional) | Multi-step SFT from `mg_bidirectional_Solaris`. |
| `dfsft_causal_i2v.yaml` | `DiffusionForcingSFTMethod` | `MatrixGame2CausalModel` | Diffusion-Forcing SFT with chunkwise timesteps. |
| `self_forcing_causal_i2v.yaml` | `SelfForcingMethod` | `MatrixGame2CausalModel` | DMD/Self-Forcing distillation; teacher = bidirectional, critic = bidirectional. |
| `solaris/finetune_i2v.yaml` | `FineTuneMethod` | `MatrixGame2Model` (bidirectional) | Multi-step SFT from `mg_bidirectional_Solaris`. |
| `solaris/dfsft_causal_i2v.yaml` | `DiffusionForcingSFTMethod` | `MatrixGame2CausalModel` | Diffusion-Forcing SFT with chunkwise timesteps. |
| `solaris/self_forcing_causal_i2v.yaml` | `SelfForcingMethod` | `MatrixGame2CausalModel` | Matrix-Game 2.0 DMD/Self-Forcing distillation; teacher = bidirectional, critic = bidirectional. |
## Zelda Configs
| Config | Method | Student | Notes |
|---|---|---|---|
| `zelda/finetune_i2v.yaml` | `FineTuneMethod` | `MatrixGame2Model` (bidirectional) | Zelda bidirectional I2V finetuning from `FastVideo/Matrix-Game-2.0-Base-Diffusers`. Uses 33-frame clips and Zelda validation with action overlays. |
| `zelda/dfsft_causal_i2v.yaml` | `DiffusionForcingSFTMethod` | `MatrixGame2CausalModel` | Zelda causal Diffusion-Forcing SFT from `mignonjia/mg_bidirectional_zelda`. Uses the same Zelda data, resolution, optimizer, and validation defaults as the Zelda finetune config. |
| `zelda/self_forcing_causal_i2v.yaml` | `SelfForcingMethod` | `MatrixGame2CausalModel` | Zelda DMD/Self-Forcing distillation; student init = `mignonjia/mg_causal_zelda`, teacher = bidirectional (`mignonjia/mg_bidirectional_zelda`), critic = bidirectional. |
| `zelda/streaming_long_tuning_causal_i2v.yaml` | `StreamingLongTuningMethod` | `MatrixGame2CausalModel` | LongLive-style streaming long tuning from the 1k-step Zelda self-forcing checkpoint. |
Zelda world-model distillation is a two-run workflow: first run
`zelda/self_forcing_causal_i2v.yaml` to train or load the 1k-step
self-forcing checkpoint (`mignonjia/mg_sf_distilled_zelda_1k_steps`), then run
`zelda/streaming_long_tuning_causal_i2v.yaml` for the 3k-step streaming
long-tuning stage. The long-tuning YAML starts from that 1k-step checkpoint; it
does not run the short self-forcing stage inside the same config.
## Zelda Training Data
The Zelda training configs use `data/zeldam2-clean` as a suggested local path.
Download the dataset from Hugging Face before running those configs:
```bash
python scripts/huggingface/download_hf.py \
--repo_id mignonjia/zeldam2-clean \
--local_dir data/zeldam2-clean \
--repo_type dataset
```
You can store the dataset elsewhere; update `training.data.data_path` in the
YAML to point at that location.
## Multi3D Training Data
`zelda/finetune_i2v.yaml` includes an optional, commented-out Multi3D entry.
Enable it only when you want to mix Zelda with multi-game data from
`data/multi3d_games`. You can store this dataset anywhere; before enabling it,
update the matching commented `training.data.data_path` key in the YAML to the
correct location.
To mix datasets in a training YAML, set `training.data.data_path` to a
path-to-repeat-count mapping. For example, `zelda/finetune_i2v.yaml` can use
`data/zeldam2-clean: 1` and `# data/multi3d_games: 10`; uncommenting the
Multi3D entry repeats the multi-game parquet list ten times before training
samples are shuffled.
## World Model Validation Data
The Zelda validation configs expect a small public validation bundle under
`data/zelda_validation_data`.
Download it from Hugging Face before running the Zelda scenarios:
```bash
python scripts/huggingface/download_hf.py \
--repo_id mignonjia/zelda_validation_data \
--local_dir data/zelda_validation_data \
--repo_type dataset
```
The bundle contains `validation_zelda.json`, `images/`, and `actions/`.
The Zelda configs point
`callbacks.validation.dataset_file` at
`data/zelda_validation_data/validation_zelda.json`.
## Usage
### Solaris
```bash
bash examples/train/run.sh \
examples/train/scenario/worldmodel/finetune_i2v.yaml
examples/train/scenario/worldmodel/solaris/finetune_i2v.yaml
bash examples/train/run.sh \
examples/train/scenario/worldmodel/dfsft_causal_i2v.yaml
examples/train/scenario/worldmodel/solaris/dfsft_causal_i2v.yaml
bash examples/train/run.sh \
examples/train/scenario/worldmodel/self_forcing_causal_i2v.yaml
examples/train/scenario/worldmodel/solaris/self_forcing_causal_i2v.yaml
```
### Zelda
```bash
# Finetuning / DFSFT
bash examples/train/run.sh \
examples/train/scenario/worldmodel/zelda/finetune_i2v.yaml
bash examples/train/run.sh \
examples/train/scenario/worldmodel/zelda/dfsft_causal_i2v.yaml
# Distillation / long tuning
bash examples/train/run.sh \
examples/train/scenario/worldmodel/zelda/self_forcing_causal_i2v.yaml
bash examples/train/run.sh \
examples/train/scenario/worldmodel/zelda/streaming_long_tuning_causal_i2v.yaml
```
Override any field on the command line:
```bash
bash examples/train/run.sh \
examples/train/scenario/worldmodel/dfsft_causal_i2v.yaml \
examples/train/scenario/worldmodel/solaris/dfsft_causal_i2v.yaml \
--training.distributed.num_gpus 8 \
--training.optimizer.learning_rate 1e-5
```
examples/train/scenario/worldmodel/dfsft_causal_i2v.yaml → examples/train/scenario/worldmodel/solaris/dfsft_causal_i2v.yaml
View File
examples/train/scenario/worldmodel/finetune_i2v.yaml → examples/train/scenario/worldmodel/solaris/finetune_i2v.yaml
View File
examples/train/scenario/worldmodel/self_forcing_causal_i2v.yaml → examples/train/scenario/worldmodel/solaris/self_forcing_causal_i2v.yaml
+1 -1
View File
@@ -97,4 +97,4 @@ callbacks:
guidance_scale: 6.0
pipeline:
flow_shift: 5
flow_shift: 5
examples/train/scenario/worldmodel/zelda/dfsft_causal_i2v.yaml
+94
View File
@@ -0,0 +1,94 @@
# Diffusion-Forcing SFT: Zelda world model I2V Causal
models:
student:
_target_: fastvideo.train.models.matrixgame2.matrixgame2_causal.MatrixGame2CausalModel
init_from: mignonjia/mg_bidirectional_zelda
trainable: true
method:
_target_: fastvideo.train.methods.fine_tuning.dfsft.DiffusionForcingSFTMethod
chunk_size: 3
training:
distributed:
num_gpus: 8
sp_size: 1
tp_size: 1
hsdp_replicate_dim: 1
hsdp_shard_dim: 8
data:
data_path:
data/zeldam2-clean: 1
dataloader_num_workers: 1
train_batch_size: 1
training_cfg_rate: 0.0
seed: 42
num_latent_t: 9
num_height: 480
num_width: 832
num_frames: 33
optimizer:
learning_rate: 2.0e-5
betas: [0.9, 0.95]
weight_decay: 0.0
lr_scheduler: constant
lr_warmup_steps: 0
loop:
max_train_steps: 60000
gradient_accumulation_steps: 1
checkpoint:
output_dir: outputs/matrixgame_finetune/checkpoints/zelda_causal_dfsft
training_state_checkpointing_steps: 5000
checkpoints_total_limit: 3
tracker:
entity: hapo-exp
project_name: mg_1.3b_zelda
run_name: zelda_causal_dfsft
model:
enable_gradient_checkpointing_type: full
dit_precision: fp32
callbacks:
grad_clip:
max_grad_norm: 1.0
validation:
pipeline_target: fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.MatrixGame2CausalDMDPipeline
dataset_file: data/zelda_validation_data/validation_zelda.json
every_steps: 200
sampling_steps: [40]
sampling_timesteps: [1000, 975, 950, 925, 900, 875, 850, 825, 800, 775,
750, 725, 700, 675, 650, 625, 600, 575, 550, 525,
500, 475, 450, 425, 400, 375, 350, 325, 300, 275,
250, 225, 200, 175, 150, 125, 100, 75, 50, 25]
num_frames: 33
overlay_actions: true
guidance_scale: 6.0
metrics:
enabled: true
names:
- vbench.imaging_quality
- vbench.aesthetic_quality
- vbench.temporal_flickering
- vbench.motion_smoothness
- vbench.subject_consistency
- vbench.background_consistency
- vbench.dynamic_degree
- optical_flow.synthetic_optical_flow
calibration_path: assets/eval/worldmodel_synthetic_flow_calibration.json
skip_missing_deps: true
strict: false
unload_after_validation: true
pipeline:
flow_shift: 5
dit_config:
local_attn_size: 6
sink_size: 1
examples/train/scenario/worldmodel/zelda/finetune_i2v.yaml
+88
View File
@@ -0,0 +1,88 @@
# Matrix-Game 2.0 Zelda + multi-game I2V finetune.
models:
student:
_target_: fastvideo.train.models.matrixgame2.matrixgame2.MatrixGame2Model
init_from: FastVideo/Matrix-Game-2.0-Base-Diffusers
trainable: true
method:
_target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
training:
distributed:
num_gpus: 8
sp_size: 1
tp_size: 1
hsdp_replicate_dim: 1
hsdp_shard_dim: 8
data:
data_path:
data/zeldam2-clean: 1
# data/multi3d_games: 10
dataloader_num_workers: 1
train_batch_size: 1
training_cfg_rate: 0.0 # unused for MatrixGame2 I2V; no text_embedding CFG dropout
seed: 42
num_latent_t: 9
num_height: 480
num_width: 832
num_frames: 33
optimizer:
learning_rate: 2.0e-5
betas: [0.9, 0.95]
weight_decay: 0.0
lr_scheduler: constant
lr_warmup_steps: 0
loop:
max_train_steps: 60000
gradient_accumulation_steps: 1
checkpoint:
output_dir: outputs/matrixgame_finetune/checkpoints/zelda_with_mg_init
training_state_checkpointing_steps: 5000
checkpoints_total_limit: 3
tracker:
project_name: mg_1.3b_zelda
run_name: zelda_with_mg_init
model:
enable_gradient_checkpointing_type: full
dit_precision: fp32
callbacks:
grad_clip:
_target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
max_grad_norm: 1.0
validation:
_target_: fastvideo.train.callbacks.validation.ValidationCallback
pipeline_target: fastvideo.pipelines.basic.matrixgame2.matrixgame2_i2v_pipeline.MatrixGame2I2VPipeline
dataset_file: data/zelda_validation_data/validation_zelda.json
every_steps: 200
sampling_steps: [40]
num_frames: 33
overlay_actions: true
guidance_scale: 6.0
metrics:
enabled: true
names:
- vbench.imaging_quality
- vbench.aesthetic_quality
- vbench.temporal_flickering
- vbench.motion_smoothness
- vbench.subject_consistency
- vbench.background_consistency
- vbench.dynamic_degree
- optical_flow.synthetic_optical_flow
calibration_path: assets/eval/worldmodel_synthetic_flow_calibration.json
skip_missing_deps: true
strict: false
unload_after_validation: true
pipeline:
flow_shift: 5
examples/train/scenario/worldmodel/zelda/self_forcing_causal_i2v.yaml
+120
View File
@@ -0,0 +1,120 @@
# Self-forcing distillation: Zelda world model I2V Causal
models:
student:
_target_: fastvideo.train.models.matrixgame2.matrixgame2_causal.MatrixGame2CausalModel
init_from: mignonjia/mg_causal_zelda
trainable: true
teacher:
_target_: fastvideo.train.models.matrixgame2.matrixgame2.MatrixGame2Model
init_from: mignonjia/mg_bidirectional_zelda
trainable: false
disable_custom_init_weights: true
critic:
_target_: fastvideo.train.models.matrixgame2.matrixgame2.MatrixGame2Model
init_from: mignonjia/mg_bidirectional_zelda
trainable: true
disable_custom_init_weights: true
method:
_target_: fastvideo.train.methods.distribution_matching.self_forcing.SelfForcingMethod
rollout_mode: simulate
generator_update_interval: 5
dmd_denoising_steps: [1000, 750, 500, 250]
warp_denoising_step: true
min_timestep_ratio: 0.02
max_timestep_ratio: 0.98
chunk_size: 3
student_sample_type: sde
same_step_across_blocks: true
last_step_only: false
context_noise: 0.0
enable_gradient_in_rollout: true
start_gradient_frame: 0
# Critic optimizer
fake_score_learning_rate: 3.0e-7
fake_score_betas: [0.9, 0.95]
fake_score_lr_scheduler: constant
training:
distributed:
num_gpus: 8
sp_size: 1
tp_size: 1
hsdp_replicate_dim: 1
hsdp_shard_dim: 8
data:
data_path: data/zeldam2-clean
dataloader_num_workers: 4
train_batch_size: 1
training_cfg_rate: 0.0
seed: 1001
num_latent_t: 9
num_height: 480
num_width: 832
num_frames: 33
optimizer:
learning_rate: 3.0e-6
betas: [0.9, 0.95]
weight_decay: 0.0
lr_scheduler: constant
lr_warmup_steps: 0
loop:
max_train_steps: 1000
gradient_accumulation_steps: 1
checkpoint:
output_dir: outputs/zelda_causal_self_forcing
training_state_checkpointing_steps: 1000
checkpoints_total_limit: 1
tracker:
project_name: wangame_sf
run_name: mg2_self_forcing_9_latents
model:
enable_gradient_checkpointing_type: full
dit_precision: fp32
callbacks:
ema:
decay: 0.99
start_iter: 200
grad_clip:
max_grad_norm: 1.0
validation:
pipeline_target: fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.MatrixGame2CausalDMDPipeline
dataset_file: data/zelda_validation_data/validation_zelda.json
every_steps: 100
sampling_steps: [4]
sampling_timesteps: [1000, 750, 500, 250]
num_frames: 153
overlay_actions: true
keyboard_value_scale: 1.0
metrics:
enabled: true
names:
- vbench.imaging_quality
- vbench.aesthetic_quality
- vbench.temporal_flickering
- vbench.motion_smoothness
- vbench.subject_consistency
- vbench.background_consistency
- vbench.dynamic_degree
- optical_flow.synthetic_optical_flow
calibration_path: assets/eval/worldmodel_synthetic_flow_calibration.json
skip_missing_deps: true
strict: false
unload_after_validation: true
pipeline:
flow_shift: 5
dit_config:
local_attn_size: 6
sink_size: 1
examples/train/scenario/worldmodel/zelda/streaming_long_tuning_causal_i2v.yaml
+140
View File
@@ -0,0 +1,140 @@
# MatrixGame2 I2V LongLive-style streaming distillation: Zelda world model I2V Causal
# Student init from self forcing after 1k steps
models:
student:
_target_: fastvideo.train.models.matrixgame2.matrixgame2_causal.MatrixGame2CausalModel
init_from: mignonjia/mg_sf_distilled_zelda_1k_steps
trainable: true
# transformer_override_safetensor: outputs/matrixgame_dmd/checkpoints/mg_zelda_sf_m2/checkpoint-1000_weight_only/ema/generator_ema.safetensors
teacher:
_target_: fastvideo.train.models.matrixgame2.matrixgame2.MatrixGame2Model
init_from: mignonjia/mg_bidirectional_zelda
trainable: false
disable_custom_init_weights: true
critic:
_target_: fastvideo.train.models.matrixgame2.matrixgame2.MatrixGame2Model
init_from: mignonjia/mg_bidirectional_zelda
trainable: true
disable_custom_init_weights: true
method:
_target_: fastvideo.train.methods.distribution_matching.streaming_long_tuning.StreamingLongTuningMethod
rollout_mode: simulate
generator_update_interval: 5
dmd_denoising_steps: [1000, 750, 500, 250]
warp_denoising_step: true
min_timestep_ratio: 0.02
max_timestep_ratio: 0.98
chunk_size: 3
student_sample_type: sde
same_step_across_blocks: true
last_step_only: false
context_noise: 0.0
enable_gradient_in_rollout: true
start_gradient_frame: 0
streaming_training: true
streaming_chunk_size: 9
streaming_max_length: 39
streaming_fixed_overlap_latents: 3
streaming_reencode_overlap_anchor: true
streaming_anchor_inject_k: 1
streaming_require_full_blocks: true
multi_phased_distill_schedule:
- stage: streaming_long
start_step: 0
end_step: 3000
num_latent_t: 39
streaming_training: true
streaming_chunk_size: 9
streaming_max_length: 39
streaming_fixed_overlap_latents: 3
# Critic optimizer
fake_score_learning_rate: 3.0e-7
fake_score_betas: [0.9, 0.95]
fake_score_lr_scheduler: constant
training:
distributed:
num_gpus: 8
sp_size: 1
tp_size: 1
hsdp_replicate_dim: 1
hsdp_shard_dim: 8
data:
data_path: data/zeldam2-clean
dataloader_num_workers: 4
train_batch_size: 1
training_cfg_rate: 0.0
seed: 1001
num_latent_t: 39
num_height: 480
num_width: 832
num_frames: 153
optimizer:
learning_rate: 3.0e-6
betas: [0.9, 0.95]
weight_decay: 0.0
lr_scheduler: constant
lr_warmup_steps: 0
loop:
max_train_steps: 3000
gradient_accumulation_steps: 1
checkpoint:
output_dir: outputs/zelda_causal_long_tuning
training_state_checkpointing_steps: 1000
checkpoints_total_limit: 2
tracker:
project_name: wangame_sf
run_name: mg2_39only_streaming_long
model:
enable_gradient_checkpointing_type: full
dit_precision: fp32
callbacks:
ema:
decay: 0.99
start_iter: 200
grad_clip:
max_grad_norm: 1.0
validation:
pipeline_target: fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.MatrixGame2CausalDMDPipeline
dataset_file: data/zelda_validation_data/validation_zelda.json
every_steps: 100
sampling_steps: [4]
sampling_timesteps: [1000, 750, 500, 250]
num_frames: 153
overlay_actions: true
keyboard_value_scale: 1.0
metrics:
enabled: true
names:
- vbench.imaging_quality
- vbench.aesthetic_quality
- vbench.temporal_flickering
- vbench.motion_smoothness
- vbench.subject_consistency
- vbench.background_consistency
- vbench.dynamic_degree
- optical_flow.synthetic_optical_flow
calibration_path: assets/eval/worldmodel_synthetic_flow_calibration.json
skip_missing_deps: true
strict: false
unload_after_validation: true
pipeline:
flow_shift: 5
dit_config:
local_attn_size: 6
sink_size: 1
fastvideo/dataset/parquet_dataset_map_style.py
+52 -4
View File
@@ -2,6 +2,7 @@
import os
import pickle
import random
from collections.abc import Sequence
from typing import Any
import pyarrow as pa
@@ -94,8 +95,56 @@ class DP_SP_BatchSampler(Sampler[list[int]]):
return len(self.sp_group_local_indices) // self.batch_size
def get_parquet_files_and_length(path: str):
dataset_root = os.path.realpath(os.path.expanduser(path))
def _parse_data_path_specs(path: str | Sequence[str] | dict[str, int]) -> list[tuple[str, int]]:
"""Parse one or more dataset roots with old-framework repeat counts."""
if isinstance(path, dict):
return [
(str(root), int(repeat))
for root, repeat in path.items()
if int(repeat) > 0
]
if isinstance(path, Sequence) and not isinstance(path, str):
return [(str(root), 1) for root in path]
specs: list[tuple[str, int]] = []
for part in str(path).split(","):
part = part.strip()
if not part:
continue
if ":" in part:
dir_path, count_str = part.rsplit(":", 1)
count = int(count_str)
else:
dir_path, count = part, 1
if count > 0:
specs.append((dir_path.strip(), count))
return specs
def get_parquet_files_and_length(path: str | Sequence[str] | dict[str, int]):
specs = _parse_data_path_specs(path)
if len(specs) != 1 or specs[0][1] != 1:
all_file_names: list[str] = []
all_lengths: list[int] = []
for root, repeat in specs:
file_names, lengths = get_parquet_files_and_length(str(root))
for _ in range(repeat):
all_file_names.extend(file_names)
all_lengths.extend(lengths)
if not all_file_names:
raise FileNotFoundError(
"No parquet files found under dataset paths: "
f"{path}. "
"Please verify these paths point to preprocessed parquet data."
)
file_lengths = sorted(
zip(all_file_names, all_lengths, strict=True),
key=lambda x: x[0],
)
file_names_sorted, lengths_sorted = zip(*file_lengths, strict=True)
return file_names_sorted, lengths_sorted
dataset_root = os.path.realpath(os.path.expanduser(specs[0][0]))
# Check if cached info exists
cache_dir = os.path.join(dataset_root, "map_style_cache")
cache_file = os.path.join(cache_dir, "file_info.pkl")
@@ -268,7 +317,7 @@ class LatentsParquetMapStyleDataset(Dataset):
def __init__(
self,
path: str,
path: str | Sequence[str] | dict[str, int],
batch_size: int,
parquet_schema: pa.Schema,
cfg_rate: float = 0.0,
@@ -358,7 +407,6 @@ class LatentsParquetMapStyleDataset(Dataset):
def __len__(self):
return sum(self.lengths)
# ────────────────────────────────────────────────────────────────────────────
# 3. Loader helper – everything else stays just like your original trainer
# ────────────────────────────────────────────────────────────────────────────
fastvideo/eval/README.md
+38 -2
View File
@@ -122,6 +122,44 @@ to run a subset of the Evaluator's registered metrics on this batch
(useful for scoring different corpora with different metric subsets in
successive calls without burning model loads).
### Training-time validation metrics
The modular trainer can score validation videos during training through
`callbacks.validation.metrics`. Each distributed rank that writes local
validation videos also runs its local metrics on `cuda:<local_rank>`;
rank 0 only merges scalar summaries and logs artifacts.
```yaml
callbacks:
validation:
pipeline_target: fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.MatrixGame2CausalDMDPipeline
dataset_file: examples/distill/MatrixGame2.0/validation.json
sampling_steps: [4]
metrics:
enabled: true
names:
- vbench.imaging_quality
- vbench.aesthetic_quality
- optical_flow.synthetic_optical_flow
- common.fvd
skip_missing_deps: true
strict: false
unload_after_validation: true
# `synthetic_optical_flow` also needs actions in the validation
# manifest (`action_path`) plus a repo-local calibration file.
calibration_path: assets/eval/worldmodel_synthetic_flow_calibration.json
```
Metric summaries are written under
`<output_dir>/eval/step_<step>/inference_steps_<n>_rank_<rank>.json` and
scalar means are logged with the `metrics/validation/...` prefix.
Install `fastvideo[eval]` for optical-flow dependencies such as
`ptlflow`; VBench metrics use the pinned submodule under
`fastvideo/third_party/eval/vbench`. FVD uses `ref_video` entries from
the validation manifest when present, or the standard
`FASTVIDEO_FVD_REF_FEATURES` / eval-cache reference feature path.
### CLI
```bash
@@ -350,5 +388,3 @@ sweep several baselines into a table, see
- **MIND** metrics. Depend on a separate `vipe` upstream submodule.
- **VBench-2.0**. Sibling vbench2 package; needs its own port.
- **Training-time eval callback** (`EvalCallback`) and the
`RolloutEvaluator` helper.
fastvideo/eval/metrics/optical_flow/synthetic_optical_flow/_thirdperson.py
+1 -1
View File
@@ -63,7 +63,7 @@ class ThirdPersonCalibration:
beta_fwd: float
beta_strafe: float
focal_length: float
r_z: float
r_z: float = 0.0
r_y: float = 0.0
init_pitch: float = 0.0
notes: str = ""
fastvideo/models/dits/causal_wanvideo.py
+2 -1
View File
@@ -665,7 +665,8 @@ class CausalWanTransformer3DModel(BaseDiT):
else:
causal_kwargs = {
"kv_cache": kv_cache[block_index],
"crossattn_cache": crossattn_cache[block_index],
"crossattn_cache": (crossattn_cache[block_index]
if crossattn_cache is not None else None),
"current_start": current_start,
"cache_start": cache_start,
"block_mask": self.block_mask,
fastvideo/models/dits/matrixgame2/action_module.py
+20 -7
View File
@@ -16,6 +16,8 @@ from fastvideo.layers.rotary_embedding import (
)
from fastvideo.platforms import AttentionBackendEnum
from .utils import retain_kv_with_sink
DISABLE_COMPILE = False
flex_attention = torch.compile(
@@ -100,6 +102,7 @@ def _update_kv_cache_and_attend(
start_frame: int,
num_frame_per_block: int,
local_attn_size: int,
sink_size: int,
*,
use_k_for_num_tokens: bool = False,
store_first_only: bool = False,
@@ -133,7 +136,6 @@ def _update_kv_cache_and_attend(
k.shape[1] if use_k_for_num_tokens else q.shape[1]
) == num_frame_per_block
sink_size = 0
max_attention_size = local_attn_size
sink_tokens = sink_size * 1
kv_cache_size = kv_cache["k"].shape[1]
@@ -206,16 +208,22 @@ def _update_kv_cache_and_attend(
# Store new k, v in cache
if store_first_only:
kv_cache["k"][:, local_start_index:local_end_index] = k[:1]
kv_cache["v"][:, local_start_index:local_end_index] = v[:1]
kv_cache["k"][:, local_start_index:local_end_index] = k[
: kv_cache["k"].shape[0]
]
kv_cache["v"][:, local_start_index:local_end_index] = v[
: kv_cache["v"].shape[0]
]
else:
kv_cache["k"][:, local_start_index:local_end_index] = k
kv_cache["v"][:, local_start_index:local_end_index] = v
# Retrieve from cache and perform attention
cache_start = max(0, local_end_index - max_attention_size)
cached_k = kv_cache["k"][:, cache_start:local_end_index]
cached_v = kv_cache["v"][:, cache_start:local_end_index]
cached_k, cached_v = retain_kv_with_sink(
kv_cache["k"][:, :local_end_index],
kv_cache["v"][:, :local_end_index],
min(local_end_index, max_attention_size),
sink_tokens,
)
if repeat_factor is not None:
cached_k = cached_k.repeat(repeat_factor, 1, 1, 1)
@@ -261,6 +269,7 @@ class ActionModule(nn.Module):
enable_mouse=True,
enable_keyboard=True,
local_attn_size=6,
sink_size=0,
blocks: list | None = None,
):
super().__init__()
@@ -274,6 +283,7 @@ class ActionModule(nn.Module):
)
blocks = blocks if blocks is not None else []
self.local_attn_size = local_attn_size
self.sink_size = sink_size
self.enable_mouse = enable_mouse
self.enable_keyboard = enable_keyboard
@@ -565,6 +575,7 @@ class ActionModule(nn.Module):
start_frame,
num_frame_per_block,
self.local_attn_size,
self.sink_size,
use_k_for_num_tokens=False,
)
else:
@@ -689,6 +700,7 @@ class ActionModule(nn.Module):
start_frame,
num_frame_per_block,
self.local_attn_size,
self.sink_size,
use_k_for_num_tokens=True,
store_first_only=True,
repeat_factor=S,
@@ -719,6 +731,7 @@ class ActionModule(nn.Module):
start_frame,
num_frame_per_block,
self.local_attn_size,
self.sink_size,
use_k_for_num_tokens=True,
)
else:
fastvideo/models/dits/matrixgame2/causal_model.py
+89 -12
View File
@@ -46,6 +46,7 @@ from fastvideo.platforms import AttentionBackendEnum, current_platform
from .action_module import ActionModule
from .model import MatrixGame2CrossAttention
from .utils import retain_kv_with_sink
logger = init_logger(__name__)
@@ -345,9 +346,12 @@ class CausalMatrixGame2SelfAttention(nn.Module):
kv_cache["k"][:, local_start_index:local_end_index] = stored_key
kv_cache["v"][:, local_start_index:local_end_index] = v
kv_start = max(0, local_end_index - max_attention_size)
k_for_attn = kv_cache["k"][:, kv_start:local_end_index]
v_for_attn = kv_cache["v"][:, kv_start:local_end_index]
k_for_attn, v_for_attn = retain_kv_with_sink(
kv_cache["k"][:, :local_end_index],
kv_cache["v"][:, :local_end_index],
min(local_end_index, max_attention_size),
sink_tokens,
)
if relativistic:
window_len, query_lo, _ = relativistic_window_offsets(
@@ -418,6 +422,7 @@ class CausalMatrixGame2TransformerBlock(nn.Module):
self.hidden_dim = dim
self.num_attention_heads = num_heads
self.local_attn_size = local_attn_size
self.sink_size = sink_size
dim_head = dim // num_heads
if qk_norm == "rms_norm":
self.norm_q = RMSNorm(dim_head, eps=eps)
@@ -474,6 +479,7 @@ class CausalMatrixGame2TransformerBlock(nn.Module):
),
patch_size=action_config["patch_size"],
local_attn_size=local_attn_size,
sink_size=sink_size,
qk_norm=action_config["qk_norm"],
qkv_bias=action_config["qkv_bias"],
vae_time_compression_ratio=action_config[
@@ -647,6 +653,8 @@ class CausalMatrixGame2WanModel(BaseDiT):
self.num_channels_latents = config.num_channels_latents
self.patch_size = config.patch_size
# Long tuning controls the causal window from YAML; consume it here
# like Wan so train wrappers do not need runtime patching.
arch_cfg = getattr(config, "arch_config", None)
self.local_attn_size = (
getattr(
@@ -662,6 +670,12 @@ class CausalMatrixGame2WanModel(BaseDiT):
if arch_cfg
else getattr(config, "sink_size", 0)
)
if self.sink_size < 0:
raise ValueError("sink_size must be non-negative")
if self.local_attn_size != -1 and self.sink_size >= self.local_attn_size:
raise ValueError(
"sink_size must be smaller than local_attn_size for "
"MatrixGame2 causal attention")
self.rope_cache_policy = (
getattr(arch_cfg, "rope_cache_policy",
getattr(config, "rope_cache_policy", "absolute"))
@@ -749,6 +763,52 @@ class CausalMatrixGame2WanModel(BaseDiT):
self.__post_init__()
def set_causal_attention_window(
self,
*,
local_attn_size: int | None = None,
sink_size: int | None = None,
) -> None:
local_attn_size = (
self.local_attn_size if local_attn_size is None else int(local_attn_size)
)
sink_size = self.sink_size if sink_size is None else int(sink_size)
if sink_size < 0:
raise ValueError("sink_size must be non-negative")
if local_attn_size != -1 and sink_size >= local_attn_size:
raise ValueError(
"sink_size must be smaller than local_attn_size for "
"MatrixGame2 causal attention")
self.local_attn_size = local_attn_size
self.sink_size = sink_size
self.block_mask = None
self.block_mask_keyboard = None
self.block_mask_mouse = None
for block in self.blocks:
if hasattr(block, "local_attn_size"):
block.local_attn_size = local_attn_size
if hasattr(block, "sink_size"):
block.sink_size = sink_size
attn1 = getattr(block, "attn1", None)
if attn1 is not None:
if hasattr(attn1, "local_attn_size"):
attn1.local_attn_size = local_attn_size
if hasattr(attn1, "sink_size"):
attn1.sink_size = sink_size
action_model = getattr(block, "action_model", None)
if action_model is not None:
if hasattr(action_model, "local_attn_size"):
action_model.local_attn_size = local_attn_size
if hasattr(action_model, "sink_size"):
action_model.sink_size = sink_size
def set_sink_size(self, sink_size: int) -> None:
self.set_causal_attention_window(sink_size=sink_size)
@staticmethod
def _prepare_blockwise_causal_attn_mask(
device: torch.device | str,
@@ -756,6 +816,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
frame_seqlen: int = 880,
num_frame_per_block: int = 1,
local_attn_size: int = -1,
sink_size: int = 0,
) -> BlockMask:
total_length = num_frames * frame_seqlen
padded_length = math.ceil(total_length / 128) * 128 - total_length
@@ -782,7 +843,10 @@ class CausalMatrixGame2WanModel(BaseDiT):
else:
return (
(kv_idx < ends[q_idx])
& (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))
& (
(kv_idx < sink_size * frame_seqlen)
| (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))
)
) | (q_idx == kv_idx)
block_mask = create_block_mask(
@@ -796,8 +860,9 @@ class CausalMatrixGame2WanModel(BaseDiT):
)
if not dist.is_initialized() or dist.get_rank() == 0:
print(
f" cache a block wise causal mask with block size of {num_frame_per_block} frames"
logger.info(
"cache a block wise causal mask with block size of %s frames",
num_frame_per_block,
)
return block_mask
@@ -809,6 +874,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
frame_seqlen: int = 880,
num_frame_per_block: int = 1,
local_attn_size: int = -1,
sink_size: int = 0,
) -> BlockMask:
total_length2 = num_frames * frame_seqlen
padded_length2 = math.ceil(total_length2 / 32) * 32 - total_length2
@@ -834,7 +900,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
else:
return (
(kv_idx < ends2[q_idx])
& (kv_idx >= (ends2[q_idx] - local_attn_size))
& ((kv_idx < sink_size) | (kv_idx >= (ends2[q_idx] - local_attn_size)))
) | (q_idx == kv_idx)
block_mask2 = create_block_mask(
@@ -848,8 +914,9 @@ class CausalMatrixGame2WanModel(BaseDiT):
)
if not dist.is_initialized() or dist.get_rank() == 0:
print(
f" cache a block wise causal mask for keyboard with block size of {num_frame_per_block} frames"
logger.info(
"cache a block wise causal mask for keyboard with block size of %s frames",
num_frame_per_block,
)
return block_mask2
@@ -861,6 +928,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
frame_seqlen: int = 1,
num_frame_per_block: int = 1,
local_attn_size: int = -1,
sink_size: int = 0,
) -> BlockMask:
total_length2 = num_frames * frame_seqlen
padded_length2 = math.ceil(total_length2 / 32) * 32 - total_length2
@@ -886,7 +954,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
else:
return (
(kv_idx < ends2[q_idx])
& (kv_idx >= (ends2[q_idx] - local_attn_size))
& ((kv_idx < sink_size) | (kv_idx >= (ends2[q_idx] - local_attn_size)))
) | (q_idx == kv_idx)
block_mask2 = create_block_mask(
@@ -900,8 +968,9 @@ class CausalMatrixGame2WanModel(BaseDiT):
)
if not dist.is_initialized() or dist.get_rank() == 0:
print(
f" cache a block wise causal mask for action with block size of {num_frame_per_block} frames"
logger.info(
"cache a block wise causal mask for action with block size of %s frames",
num_frame_per_block,
)
return block_mask2
@@ -1063,6 +1132,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
frame_seqlen=post_patch_height * post_patch_width,
num_frame_per_block=self.num_frame_per_block,
local_attn_size=self.local_attn_size,
sink_size=self.sink_size,
)
if self.use_rope_keyboard:
block_mask_keyboard = self._prepare_blockwise_causal_attn_mask_action(
@@ -1071,6 +1141,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
frame_seqlen=1,
num_frame_per_block=self.num_frame_per_block,
local_attn_size=self.local_attn_size,
sink_size=self.sink_size,
)
else:
block_mask_keyboard = self._prepare_blockwise_causal_attn_mask_keyboard(
@@ -1079,6 +1150,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
frame_seqlen=post_patch_height * post_patch_width,
num_frame_per_block=self.num_frame_per_block,
local_attn_size=self.local_attn_size,
sink_size=self.sink_size,
)
block_mask_mouse = self._prepare_blockwise_causal_attn_mask_action(
device=hidden_states.device,
@@ -1086,6 +1158,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
frame_seqlen=1,
num_frame_per_block=self.num_frame_per_block,
local_attn_size=self.local_attn_size,
sink_size=self.sink_size,
)
if kv_cache is None:
kv_cache = [None] * len(self.blocks)
@@ -1239,6 +1312,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
frame_seqlen=post_patch_height * post_patch_width,
num_frame_per_block=self.num_frame_per_block,
local_attn_size=self.local_attn_size,
sink_size=self.sink_size,
)
if self.block_mask_keyboard is None:
if self.use_rope_keyboard:
@@ -1248,6 +1322,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
frame_seqlen=1,
num_frame_per_block=self.num_frame_per_block,
local_attn_size=self.local_attn_size,
sink_size=self.sink_size,
)
else:
self.block_mask_keyboard = self._prepare_blockwise_causal_attn_mask_keyboard(
@@ -1256,6 +1331,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
frame_seqlen=post_patch_height * post_patch_width,
num_frame_per_block=self.num_frame_per_block,
local_attn_size=self.local_attn_size,
sink_size=self.sink_size,
)
if self.block_mask_mouse is None:
self.block_mask_mouse = self._prepare_blockwise_causal_attn_mask_action(
@@ -1264,6 +1340,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
frame_seqlen=1,
num_frame_per_block=self.num_frame_per_block,
local_attn_size=self.local_attn_size,
sink_size=self.sink_size,
)
hidden_states = self.patch_embedding(hidden_states)
fastvideo/models/dits/matrixgame2/model.py
+1 -1
View File
@@ -452,7 +452,7 @@ class MatrixGame2WanModel(BaseDiT):
and len(encoder_hidden_states_image) > 0
):
encoder_hidden_states_image = encoder_hidden_states_image[0]
else:
elif not isinstance(encoder_hidden_states_image, torch.Tensor):
encoder_hidden_states_image = None
batch_size, num_channels, num_frames, height, width = (
fastvideo/models/dits/matrixgame2/utils.py
+256
View File
@@ -4,12 +4,18 @@ import asyncio
import os
import random
import numpy as np
import torch
from PIL import Image
from fastvideo.distributed.parallel_state import get_local_torch_device
from fastvideo.utils import logger
try:
import cv2
except ImportError:
cv2 = None
CAM_VALUE = 0.1
CAMERA_MAP = {
@@ -34,6 +40,256 @@ KEYBOARD_MAP_7 = { # templerun_distilled_model: still/w/s/left/right/a/d
"d": [0, 0, 0, 0, 0, 0, 1], # d
}
KEYBOARD_MAP = KEYBOARD_MAP_4 # Default for backward compatibility
SOLARIS_MOVEMENT_KEY_INDICES = {
"W": 11,
"S": 12,
"A": 13,
"D": 14,
}
def retain_kv_with_sink(
k: torch.Tensor,
v: torch.Tensor,
target_len: int,
sink_size: int,
) -> tuple[torch.Tensor, torch.Tensor]:
if k.shape[1] <= target_len or sink_size <= 0:
return k[:, -target_len:], v[:, -target_len:]
sink_len = min(int(sink_size), target_len, k.shape[1])
tail_len = target_len - sink_len
sink_k = k[:, :sink_len]
sink_v = v[:, :sink_len]
if tail_len <= 0:
return sink_k, sink_v
tail_k = k[:, sink_len:][:, -tail_len:]
tail_v = v[:, sink_len:][:, -tail_len:]
return torch.cat([sink_k, tail_k], dim=1), torch.cat([sink_v, tail_v], dim=1)
def _require_cv2() -> bool:
if cv2 is None:
logger.warning(
"OpenCV is not available; skipping MatrixGame2 validation overlay."
)
return False
return True
def _matrixgame2_overlay_keys_from_keyboard(
keyboard_frame: np.ndarray | torch.Tensor,
) -> dict[str, bool]:
vec = np.asarray(keyboard_frame, dtype=np.float32).reshape(-1)
dim = vec.shape[0]
if dim >= 23:
return {
key: bool(vec[idx] > 0.5)
for key, idx in SOLARIS_MOVEMENT_KEY_INDICES.items()
}
if dim == 7:
return {
"W": bool(vec[1] > 0.5),
"S": bool(vec[2] > 0.5),
"A": bool(vec[5] > 0.5),
"D": bool(vec[6] > 0.5),
}
if dim >= 6:
return {
"W": bool(vec[0] > 0.5),
"S": bool(vec[1] > 0.5),
"A": bool(vec[2] > 0.5),
"D": bool(vec[3] > 0.5),
}
if dim == 4:
return {
"W": bool(vec[0] > 0.5),
"S": bool(vec[1] > 0.5),
"A": bool(vec[2] > 0.5),
"D": bool(vec[3] > 0.5),
}
if dim == 2:
return {
"W": bool(vec[0] > 0.5),
"S": bool(vec[1] > 0.5),
"A": False,
"D": False,
}
return {"W": False, "S": False, "A": False, "D": False}
def draw_rounded_rectangle(
image: np.ndarray,
top_left: tuple[int, int],
bottom_right: tuple[int, int],
color: tuple[int, int, int],
radius: int = 10,
alpha: float = 0.5,
) -> None:
if not _require_cv2():
return
overlay = image.copy()
x1, y1 = top_left
x2, y2 = bottom_right
cv2.rectangle(overlay, (x1 + radius, y1), (x2 - radius, y2), color, -1)
cv2.rectangle(overlay, (x1, y1 + radius), (x2, y2 - radius), color, -1)
cv2.ellipse(overlay, (x1 + radius, y1 + radius), (radius, radius), 180, 0, 90, color, -1)
cv2.ellipse(overlay, (x2 - radius, y1 + radius), (radius, radius), 270, 0, 90, color, -1)
cv2.ellipse(overlay, (x1 + radius, y2 - radius), (radius, radius), 90, 0, 90, color, -1)
cv2.ellipse(overlay, (x2 - radius, y2 - radius), (radius, radius), 0, 0, 90, color, -1)
cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0, image)
def draw_keys_on_frame(
frame: np.ndarray,
keys: dict[str, bool],
key_size: tuple[int, int] = (30, 30),
top_margin: int = 15,
) -> None:
if not _require_cv2():
return
left_margin = 15
gap = 3
key_positions = {
"W": (left_margin + key_size[0] + gap, top_margin),
"A": (left_margin, top_margin + key_size[1] + gap),
"S": (left_margin + key_size[0] + gap, top_margin + key_size[1] + gap),
"D": (
left_margin + (key_size[0] + gap) * 2,
top_margin + key_size[1] + gap,
),
}
for key, (x, y) in key_positions.items():
pressed = keys.get(key, False)
color = (0, 255, 0) if pressed else (200, 200, 200)
alpha = 0.8 if pressed else 0.5
draw_rounded_rectangle(
frame,
(x, y),
(x + key_size[0], y + key_size[1]),
color,
radius=5,
alpha=alpha,
)
text_size = cv2.getTextSize(
key, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1
)[0]
text_x = x + (key_size[0] - text_size[0]) // 2
text_y = y + (key_size[1] + text_size[1]) // 2
cv2.putText(
frame,
key,
(text_x, text_y),
cv2.FONT_HERSHEY_SIMPLEX,
0.5,
(0, 0, 0),
1,
)
def draw_mouse_on_frame(
frame: np.ndarray,
yaw: float,
pitch: float,
top_margin: int = 15,
) -> None:
if not _require_cv2():
return
_, width, _ = frame.shape
right_margin = 15
crosshair_radius = 25
crosshair_x = width - right_margin - crosshair_radius
crosshair_y = top_margin + crosshair_radius
dx = int(yaw * crosshair_radius * 8)
dy = int(-pitch * crosshair_radius * 8)
max_arrow = crosshair_radius - 5
dx = max(-max_arrow, min(max_arrow, dx))
dy = max(-max_arrow, min(max_arrow, dy))
cv2.circle(frame, (crosshair_x, crosshair_y), crosshair_radius, (50, 50, 50), -1)
cv2.circle(frame, (crosshair_x, crosshair_y), crosshair_radius, (200, 200, 200), 1)
cv2.line(
frame,
(crosshair_x - crosshair_radius + 5, crosshair_y),
(crosshair_x + crosshair_radius - 5, crosshair_y),
(100, 100, 100),
1,
)
cv2.line(
frame,
(crosshair_x, crosshair_y - crosshair_radius + 5),
(crosshair_x, crosshair_y + crosshair_radius - 5),
(100, 100, 100),
1,
)
if abs(dx) > 1 or abs(dy) > 1:
cv2.arrowedLine(
frame,
(crosshair_x, crosshair_y),
(crosshair_x + dx, crosshair_y + dy),
(0, 255, 0),
2,
tipLength=0.3,
)
def overlay_validation_actions_on_frames(
frames: list[np.ndarray],
keyboard_cond: np.ndarray | torch.Tensor | None = None,
mouse_cond: np.ndarray | torch.Tensor | None = None,
) -> list[np.ndarray]:
if (keyboard_cond is None and mouse_cond is None) or not _require_cv2():
return frames
if keyboard_cond is not None and torch.is_tensor(keyboard_cond):
keyboard_cond = keyboard_cond.detach().cpu().float().numpy()
if mouse_cond is not None and torch.is_tensor(mouse_cond):
mouse_cond = mouse_cond.detach().cpu().float().numpy()
if keyboard_cond is not None:
keyboard_cond = np.asarray(keyboard_cond, dtype=np.float32)
if keyboard_cond.ndim == 3 and keyboard_cond.shape[0] == 1:
keyboard_cond = keyboard_cond[0]
if mouse_cond is not None:
mouse_cond = np.asarray(mouse_cond, dtype=np.float32)
if mouse_cond.ndim == 3 and mouse_cond.shape[0] == 1:
mouse_cond = mouse_cond[0]
processed_frames: list[np.ndarray] = []
for frame_idx, frame in enumerate(frames):
frame = np.ascontiguousarray(frame.copy())
if keyboard_cond is not None and frame_idx < len(keyboard_cond):
draw_keys_on_frame(
frame,
_matrixgame2_overlay_keys_from_keyboard(keyboard_cond[frame_idx]),
)
if mouse_cond is not None and frame_idx < len(mouse_cond):
# MG action convention: mouse[0]=pitch, mouse[1]=yaw.
pitch = float(mouse_cond[frame_idx, 0])
yaw = float(mouse_cond[frame_idx, 1])
draw_mouse_on_frame(frame, yaw=yaw, pitch=pitch)
processed_frames.append(frame)
return processed_frames
def scale_keyboard_condition(
keyboard_condition: torch.Tensor,
scale: float,
) -> torch.Tensor:
if float(scale) == 1.0:
return keyboard_condition
return keyboard_condition * float(scale)
def expand_action_to_frames(action: dict, num_frames: int) -> tuple[torch.Tensor, torch.Tensor]:
result = {}
fastvideo/registry.py
+6
View File
@@ -718,6 +718,12 @@ def _register_configs() -> None:
"FastVideo/Matrix-Game-2.0-Base-Diffusers",
"FastVideo/Matrix-Game-2.0-GTA-Diffusers",
"FastVideo/Matrix-Game-2.0-TempleRun-Diffusers",
# Zelda World Model
"mignonjia/mg_longtuning_distilled_zelda", # distilled using streaming long tuning which init from mg_sf_distilled_zelda_1k_steps
"mignonjia/mg_sf_distilled_zelda_1k_steps", # distilled using self forcing for 1k steps
"mignonjia/mg_sf_distilled_zelda", # distilled using self forcing for 3k steps
"mignonjia/mg_causal_zelda",
"mignonjia/mg_bidirectional_zelda",
],
model_detectors=[
lambda path: any(token in path.lower() for token in (
fastvideo/tests/dataset/test_parquet_dataset_map_style.py
+117
View File
@@ -0,0 +1,117 @@
# SPDX-License-Identifier: Apache-2.0
"""Unit tests for parquet map-style dataset path parsing."""
from __future__ import annotations
import pickle
import pytest
from fastvideo.dataset import parquet_dataset_map_style as parquet_dataset
from fastvideo.dataset.parquet_dataset_map_style import _parse_data_path_specs
def test_parse_data_path_specs_accepts_old_repeat_string() -> None:
# Dataset parsing keeps compatibility with the old "path:repeat" string
# form used by existing training configs.
assert _parse_data_path_specs("data/path1:2,data/path2:1") == [
("data/path1", 2),
("data/path2", 1),
]
def test_parse_data_path_specs_accepts_yaml_mapping() -> None:
# New YAML mapping form should reach the dataset layer as path -> repeat
# and parse to the same internal spec representation.
assert _parse_data_path_specs({
"data/path1": 1,
"data/path2": 2,
}) == [
("data/path1", 1),
("data/path2", 2),
]
def test_parse_data_path_specs_accepts_path_list() -> None:
assert _parse_data_path_specs(["data/a", "data/b"]) == [
("data/a", 1),
("data/b", 1),
]
def test_parse_data_path_specs_drops_non_positive_repeats() -> None:
assert _parse_data_path_specs({"data/a": 0, "data/b": -1, "data/c": 2}) == [
("data/c", 2),
]
assert _parse_data_path_specs("data/a:0,data/b:2") == [("data/b", 2)]
def test_parse_data_path_specs_rejects_malformed_repeats() -> None:
# Repeat counts must parse as ints; anything else is an explicit error
# rather than a silently mis-weighted dataset.
with pytest.raises(ValueError):
_parse_data_path_specs("data/a:abc")
with pytest.raises(ValueError):
_parse_data_path_specs("data/a:1.5")
with pytest.raises(ValueError):
_parse_data_path_specs({"data/a": "abc"}) # type: ignore[dict-item]
class _DummyWorldGroup:
def barrier(self) -> None:
return None
def _write_root_cache(dataset_root, filename: str, length: int) -> str:
cache_dir = dataset_root / "map_style_cache"
cache_dir.mkdir(parents=True)
parquet_file = dataset_root / filename
parquet_file.touch()
with (cache_dir / "file_info.pkl").open("wb") as f:
pickle.dump(((str(parquet_file),), (length,)), f)
return str(parquet_file)
def test_get_parquet_files_and_length_repeats_single_path(tmp_path, monkeypatch) -> None:
# get_parquet_files_and_length applies repeat counts after reading the
# per-root parquet cache, so a repeated root duplicates both names and rows.
dataset_root = tmp_path / "dataset"
parquet_file = _write_root_cache(dataset_root, "sample.parquet", 7)
monkeypatch.setattr(parquet_dataset, "get_world_rank", lambda: 0)
monkeypatch.setattr(parquet_dataset, "get_world_group", _DummyWorldGroup)
file_names, lengths = parquet_dataset.get_parquet_files_and_length({
str(dataset_root): 2,
})
assert file_names == (parquet_file, parquet_file)
assert lengths == (7, 7)
def test_get_parquet_files_and_length_mixes_roots_and_resorts(tmp_path, monkeypatch) -> None:
# Multiple roots are expanded per repeat count and then globally re-sorted
# by filename, so the mix order is independent of the mapping order.
root_a = tmp_path / "dataset_a"
root_b = tmp_path / "dataset_b"
file_a = _write_root_cache(root_a, "a.parquet", 5)
file_b = _write_root_cache(root_b, "b.parquet", 9)
monkeypatch.setattr(parquet_dataset, "get_world_rank", lambda: 0)
monkeypatch.setattr(parquet_dataset, "get_world_group", _DummyWorldGroup)
file_names, lengths = parquet_dataset.get_parquet_files_and_length({
str(root_b): 1,
str(root_a): 2,
})
assert file_names == (file_a, file_a, file_b)
assert lengths == (5, 5, 9)
def test_get_parquet_files_and_length_raises_when_all_repeats_dropped() -> None:
# Zero/negative repeats are dropped at parse time; if that leaves nothing
# to read, the mix branch fails loudly instead of yielding an empty dataset.
with pytest.raises(FileNotFoundError):
parquet_dataset.get_parquet_files_and_length({"data/a": 0})
fastvideo/tests/train/callbacks/test_validation.py
+347 -1
View File
@@ -14,11 +14,21 @@ distributed init and is exercised by Phase 2/3 tests.
"""
from __future__ import annotations
from types import SimpleNamespace
import numpy as np
import pytest
import torch
from fastvideo.train.callbacks.callback import CallbackDict
from fastvideo.train.callbacks.ema import EMACallback
from fastvideo.train.callbacks.validation import ValidationCallback
from fastvideo.train.callbacks.validation import (
DEFAULT_VALIDATION_VBENCH_METRICS,
SYNTHETIC_OPTICAL_FLOW_LOG_KEYS,
SYNTHETIC_OPTICAL_FLOW_METRIC,
ValidationCallback,
_ValidationMetricStats,
)
# ---------------------------------------------------------------------------
@@ -37,6 +47,7 @@ def _make_callback(
num_frames: int | None = None,
sampling_timesteps: list[int] | None = None,
output_dir: str | None = None,
overlay_actions: bool = False,
) -> ValidationCallback:
return ValidationCallback(
pipeline_target=_PIPE_TARGET,
@@ -47,6 +58,7 @@ def _make_callback(
num_frames=num_frames,
sampling_timesteps=sampling_timesteps,
output_dir=output_dir,
overlay_actions=overlay_actions,
)
@@ -67,12 +79,14 @@ class TestConstructor:
assert cb.num_frames is None
assert cb.sampling_timesteps is None
assert cb.output_dir is None
assert cb.overlay_actions is False
assert cb.offload_training_state is False
assert cb.unload_pipeline_after_validation is False
# Lazy fields not yet populated.
assert cb._pipeline is None
assert cb._sampling_param is None
assert cb.validation_random_generator is None
assert cb.metrics_config.enabled is False
def test_string_inputs_are_coerced(self) -> None:
# YAML often produces strings for numeric fields; the
@@ -85,6 +99,7 @@ class TestConstructor:
guidance_scale="4.5", # type: ignore[arg-type]
num_frames="77", # type: ignore[arg-type]
sampling_timesteps=["1000", "500"],
overlay_actions=1, # type: ignore[arg-type]
offload_training_state="1", # type: ignore[arg-type]
unload_pipeline_after_validation="false", # type: ignore[arg-type]
)
@@ -93,6 +108,7 @@ class TestConstructor:
assert cb.guidance_scale == 4.5
assert cb.num_frames == 77
assert cb.sampling_timesteps == [1000, 500]
assert cb.overlay_actions is True
assert cb.offload_training_state is True
assert cb.unload_pipeline_after_validation is False
@@ -109,6 +125,45 @@ class TestConstructor:
"another": "value",
}
def test_metrics_true_uses_default_vbench_subset(self) -> None:
cb = ValidationCallback(
pipeline_target=_PIPE_TARGET,
dataset_file="x.json",
metrics=True,
)
assert cb.metrics_config.enabled is True
assert cb.metrics_config.names == DEFAULT_VALIDATION_VBENCH_METRICS
assert "metrics" not in cb.pipeline_kwargs
def test_metrics_mapping_is_coerced(self) -> None:
cb = ValidationCallback(
pipeline_target=_PIPE_TARGET,
dataset_file="x.json",
metrics={
"names": "vbench.aesthetic_quality",
"device": "cpu",
"calibration_path": "/tmp/calibration.json",
"mouse_pitch_sign": "-1",
"skip_missing_deps": False,
"strict": True,
"unload_after_validation": False,
"loader_threads": "2",
"prefetch_factor": "3",
"log_prefix": "custom/validation",
},
)
assert cb.metrics_config.enabled is True
assert cb.metrics_config.names == ["vbench.aesthetic_quality"]
assert cb.metrics_config.device == "cpu"
assert cb.metrics_config.calibration_path == "/tmp/calibration.json"
assert cb.metrics_config.mouse_pitch_sign == -1
assert cb.metrics_config.skip_missing_deps is False
assert cb.metrics_config.strict is True
assert cb.metrics_config.unload_after_validation is False
assert cb.metrics_config.loader_threads == 2
assert cb.metrics_config.prefetch_factor == 3
assert cb.metrics_config.log_prefix == "custom/validation"
# ---------------------------------------------------------------------------
# B. on_validation_begin gating
@@ -238,3 +293,294 @@ class TestStateDict:
{"validation_rng": torch.tensor([1, 2, 3], dtype=torch.uint8)}
)
assert cb.validation_random_generator is None
# ---------------------------------------------------------------------------
# E. metric result aggregation
# ---------------------------------------------------------------------------
class TestMetricAggregation:
def test_accumulates_scores_and_scalar_details(self) -> None:
stats = _ValidationMetricStats()
row: dict = {"path": "sample.mp4"}
ValidationCallback._accumulate_metric_results(
stats,
row,
{
"vbench.aesthetic_quality": SimpleNamespace(
name="vbench.aesthetic_quality",
score=0.5,
details={"ignored": [1, 2, 3]},
),
SYNTHETIC_OPTICAL_FLOW_METRIC: SimpleNamespace(
name=SYNTHETIC_OPTICAL_FLOW_METRIC,
score=1.5,
details={
**{key: float(i) for i, key in enumerate(SYNTHETIC_OPTICAL_FLOW_LOG_KEYS)},
"pixel_epe_mean_std": 99.0,
"pixel_epe_mean_max": 100.0,
"pixel_epe_mean_auc": 101.0,
},
),
},
)
assert stats.sums["vbench.aesthetic_quality"] == 0.5
assert stats.counts["vbench.aesthetic_quality"] == 1.0
for i, key in enumerate(SYNTHETIC_OPTICAL_FLOW_LOG_KEYS):
metric_key = f"{SYNTHETIC_OPTICAL_FLOW_METRIC}.{key}"
assert stats.sums[metric_key] == float(i)
assert stats.counts[metric_key] == 1.0
assert SYNTHETIC_OPTICAL_FLOW_METRIC not in stats.sums
assert f"{SYNTHETIC_OPTICAL_FLOW_METRIC}.pixel_epe_mean_std" not in stats.sums
assert f"{SYNTHETIC_OPTICAL_FLOW_METRIC}.pixel_epe_mean_max" not in stats.sums
assert f"{SYNTHETIC_OPTICAL_FLOW_METRIC}.pixel_epe_mean_auc" not in stats.sums
assert "vbench.aesthetic_quality.ignored" not in stats.sums
assert row["vbench.aesthetic_quality"] == 0.5
def test_merges_metric_stats(self) -> None:
dst = _ValidationMetricStats(
sums={"a": 1.0},
counts={"a": 1.0},
per_video=[{"path": "a.mp4"}],
errors=["first"],
)
src = _ValidationMetricStats(
sums={"a": 2.0, "b": 3.0},
counts={"a": 2.0, "b": 1.0},
per_video=[{"path": "b.mp4"}],
errors=["second"],
)
ValidationCallback._merge_metric_stats(dst, src)
assert dst.sums == {"a": 3.0, "b": 3.0}
assert dst.counts == {"a": 3.0, "b": 1.0}
assert dst.per_video == [{"path": "a.mp4"}, {"path": "b.mp4"}]
assert dst.errors == ["first", "second"]
def test_metric_log_name_replaces_dots(self) -> None:
cb = ValidationCallback(
pipeline_target=_PIPE_TARGET,
dataset_file="x.json",
metrics={"names": ["vbench.aesthetic_quality"]},
)
assert (
cb._metric_log_name("vbench.aesthetic_quality")
== "metrics/validation/vbench/aesthetic_quality"
)
def test_metric_device_uses_local_rank(
self,
monkeypatch: pytest.MonkeyPatch,
) -> None:
cb = ValidationCallback(
pipeline_target=_PIPE_TARGET,
dataset_file="x.json",
metrics={"names": ["vbench.aesthetic_quality"]},
)
cb.world_group = SimpleNamespace(local_rank=5)
monkeypatch.setattr(
torch.cuda,
"is_available",
lambda: True,
)
assert cb._metric_device() == "cuda:5"
def test_metric_extras_include_actions_and_calibration(self) -> None:
cb = ValidationCallback(
pipeline_target=_PIPE_TARGET,
dataset_file="x.json",
metrics={
"names": ["optical_flow.synthetic_optical_flow"],
"calibration_path": "/tmp/calibration.json",
"mouse_pitch_sign": -1,
},
)
actions = {
"keyboard": np.zeros((3, 6)),
"mouse": np.zeros((3, 2)),
}
extras = cb._validation_metric_extras(
video_filenames=["a.mp4", "b.mp4"],
actions=[actions, None],
mouse_pitch_signs=[None, 1],
)
assert extras[0]["actions"] is actions
assert extras[0]["calibration"] == "/tmp/calibration.json"
assert extras[0]["mouse_pitch_sign"] == -1
assert "actions" not in extras[1]
assert extras[1]["mouse_pitch_sign"] == 1
def test_validation_actions_prefers_loaded_conditions(self) -> None:
keyboard = np.zeros((4, 6))
mouse = np.zeros((4, 2))
actions = ValidationCallback._validation_actions(
{
"keyboard_cond": keyboard,
"mouse_cond": mouse,
}
)
assert actions is not None
assert np.array_equal(actions["keyboard"], keyboard)
assert np.array_equal(actions["mouse"], mouse)
def test_validation_actions_loads_action_path(
self,
tmp_path,
) -> None:
path = tmp_path / "actions.npy"
np.save(
path,
{
"keyboard": np.ones((4, 6)),
"mouse": np.ones((4, 2)),
},
allow_pickle=True,
)
actions = ValidationCallback._validation_actions({"action_path": str(path)})
assert actions is not None
assert actions["keyboard"].shape == (4, 6)
assert actions["mouse"].shape == (4, 2)
def test_available_paths_requires_all_paths(self) -> None:
assert ValidationCallback._available_paths(["a.mp4", "b.mp4"]) == ["a.mp4", "b.mp4"]
assert ValidationCallback._available_paths(["a.mp4", None]) is None
# ---------------------------------------------------------------------------
# F. action overlay plumbing
# ---------------------------------------------------------------------------
class TestActionOverlay:
def test_save_validation_videos_supports_overlay_suffix(
self,
tmp_path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
cb = _make_callback()
cb.global_rank = 3
calls: list[tuple[str, int]] = []
def fake_mimsave(
fname: str,
video: list[np.ndarray],
*,
fps: int,
) -> None:
calls.append((fname, fps))
monkeypatch.setattr(
"fastvideo.train.callbacks.validation.imageio.mimsave",
fake_mimsave,
)
saved = cb._save_validation_videos(
[[np.zeros((2, 2, 3), dtype=np.uint8)]],
output_dir=str(tmp_path),
step=7,
num_inference_steps=4,
fps=25,
suffix="_overlay",
)
assert saved.filenames == [
str(tmp_path / "validation_step_7_inference_steps_4_rank_3_video_0_overlay.mp4")
]
assert saved.indices == [0]
assert calls == [(saved.filenames[0], 25)]
def test_save_validation_videos_skips_failed_artifact(
self,
tmp_path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
cb = _make_callback()
cb.global_rank = 0
calls: list[str] = []
def fake_mimsave(
fname: str,
video: list[np.ndarray],
*,
fps: int,
) -> None:
del video, fps
calls.append(fname)
if len(calls) == 1:
raise OSError("ffmpeg failed")
monkeypatch.setattr(
"fastvideo.train.callbacks.validation.imageio.mimsave",
fake_mimsave,
)
saved = cb._save_validation_videos(
[
[np.zeros((2, 2, 3), dtype=np.uint8)],
[np.ones((2, 2, 3), dtype=np.uint8)],
],
output_dir=str(tmp_path),
step=8,
num_inference_steps=5,
fps=24,
)
assert saved.indices == [1]
assert saved.filenames == [
str(tmp_path / "validation_step_8_inference_steps_5_rank_0_video_1.mp4")
]
assert len(calls) == 2
def test_post_process_validation_frames_uses_student_hook(self) -> None:
cb = _make_callback(overlay_actions=True)
frames = [np.zeros((2, 2, 3), dtype=np.uint8)]
action = {
"keyboard": np.ones((1, 4), dtype=np.float32),
"mouse": np.zeros((1, 2), dtype=np.float32),
}
class Student:
def post_process_validation_frames(
self,
received_frames: list[np.ndarray],
*,
action: dict,
) -> list[np.ndarray]:
assert received_frames is frames
assert action["keyboard"].shape == (1, 4)
return [np.full((2, 2, 3), 255, dtype=np.uint8)]
cb.method = SimpleNamespace(student=Student())
overlay = cb._post_process_validation_frames(
frames,
action=action,
)
assert overlay is not None
assert int(overlay[0][0, 0, 0]) == 255
def test_post_process_validation_frames_without_hook_is_noop(self) -> None:
cb = _make_callback(overlay_actions=True)
cb.method = SimpleNamespace(student=object())
overlay = cb._post_process_validation_frames(
[np.zeros((2, 2, 3), dtype=np.uint8)],
action=None,
)
assert overlay is None
fastvideo/tests/train/fixtures/wan_causal_t2v_streaming_long_min.yaml
+109
View File
@@ -0,0 +1,109 @@
# SPDX-License-Identifier: Apache-2.0
# Minimum config to run single training steps of StreamingLongTuningMethod
# on WanCausalModel for the per-method smoke test. Uses the real Wan 2.1 1.3B
# checkpoint for the causal student and the bidirectional teacher/critic
# (mirroring the MatrixGame/LongLive production layout at 1.3B scale), with
# tiny synthetic latents.
#
# Deviations from the production zelda config, chosen to keep the smoke test
# transformer-only:
# - cfg_uncond.text: zero — avoids loading the UMT5 encoder for the
# negative-prompt embedding on the teacher CFG branch.
# - streaming_reencode_overlap_anchor: false — avoids loading the VAE for the
# overlap-anchor re-encode.
models:
student:
_target_: fastvideo.train.models.wan.WanCausalModel
init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
trainable: true
teacher:
_target_: fastvideo.train.models.wan.WanModel
init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
trainable: false
critic:
_target_: fastvideo.train.models.wan.WanModel
init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
trainable: true
method:
_target_: fastvideo.train.methods.distribution_matching.streaming_long_tuning.StreamingLongTuningMethod
rollout_mode: simulate
generator_update_interval: 1
dmd_denoising_steps: [1000, 750, 500, 250]
warp_denoising_step: true
min_timestep_ratio: 0.02
max_timestep_ratio: 0.98
chunk_size: 3
student_sample_type: sde
same_step_across_blocks: true
last_step_only: false
context_noise: 0.0
enable_gradient_in_rollout: true
start_gradient_frame: 0
cfg_uncond:
text: zero
streaming_training: true
streaming_chunk_size: 3
streaming_max_length: 6
# With chunk_size 3 and fixed overlap 1, the second streaming chunk carries
# 2 new frames + 1 overlap latent, exercising the overlap/chunk-mask path.
streaming_fixed_overlap_latents: 1
streaming_reencode_overlap_anchor: false
multi_phased_distill_schedule:
- stage: self_forcing
start_step: 0
end_step: 1
num_latent_t: 3
- stage: streaming_long
start_step: 1
num_latent_t: 6
streaming_training: true
streaming_chunk_size: 3
streaming_max_length: 6
streaming_fixed_overlap_latents: 1
# Critic optimizer
fake_score_learning_rate: 3.0e-7
fake_score_betas: [0.9, 0.95]
fake_score_lr_scheduler: constant
training:
dit_precision: bf16
# Required for grad-enabled rollouts: switches the student's streaming KV
# cache to the checkpoint-safe variant so cache writes don't invalidate
# tensors saved for backward.
model:
enable_gradient_checkpointing_type: full
distributed:
num_gpus: 1
sp_size: 1
tp_size: 1
data:
seed: 42
train_batch_size: 1
training_cfg_rate: 0.0
num_latent_t: 6
num_height: 64
num_width: 64
num_frames: 21
optimizer:
learning_rate: 2.0e-6
betas: [0.9, 0.999]
weight_decay: 0.01
lr_scheduler: constant
lr_warmup_steps: 0
loop:
max_train_steps: 4
gradient_accumulation_steps: 1
pipeline: {}
fastvideo/tests/train/methods/grad_norm_refs.json
+3
View File
@@ -6,6 +6,9 @@
"test_matrixgame2_finetune": {
"H200": 4.8247
},
"test_streaming_long_tuning": {
"H200": 0.0049
},
"test_wan_causal_dfsft": {
"GB200": 2.9781,
"L40S": 3.2562
fastvideo/tests/train/methods/test_streaming_long_tuning.py
+362
View File
@@ -0,0 +1,362 @@
# SPDX-License-Identifier: Apache-2.0
"""Per-method GPU smoke test: ``WanCausalModel`` + ``StreamingLongTuningMethod``.
Mirrors ``test_wan_causal_cd.py``. LongLive-style streaming long tuning runs a
short self-forcing stage first, then switches to a persistent streaming
sequence trained chunk-by-chunk with the causal student KV cache. The GPU test
walks a single method instance through both stages: the short-stage step, the
first streaming chunk, an overlapped chunk, the chunk that fills the sequence,
and the post-reset chunk, asserting finite losses, student/critic gradient
flow, and the streaming-length bookkeeping at every step.
The pure-tensor helpers (stage selection, schedule validation, first-frame
padding, latent truncation, chunk sizing) are covered by CPU-only tests at the
bottom — they need no GPU, no mocks, and run wherever this file is collected.
"""
from __future__ import annotations
import os
os.environ.setdefault("MASTER_ADDR", "localhost")
os.environ.setdefault("MASTER_PORT", "29523")
from pathlib import Path
from types import SimpleNamespace
import pytest
import torch
from fastvideo.train.methods.distribution_matching.streaming_long_tuning import (
DistillStage,
StreamingLongTuningMethod,
parse_multi_phased_distill_schedule,
select_distill_stage,
)
from fastvideo.train.models.wan import WanCausalModel, WanModel
from fastvideo.train.utils.config import load_run_config
from .grad_norm_regression import check_grad_norm_regression
_FIXTURE = str(
Path(__file__).resolve().parent.parent / "fixtures"
/ "wan_causal_t2v_streaming_long_min.yaml")
def _build_synthetic_batch(
device: torch.device,
dtype: torch.dtype,
) -> dict[str, torch.Tensor]:
batch_size = 1
return {
"text_embedding":
torch.randn(batch_size, 16, 4096, device=device, dtype=dtype),
"text_attention_mask":
torch.ones(batch_size, 16, device=device),
"vae_latent":
torch.randn(batch_size, 16, 6, 8, 8, device=device, dtype=dtype),
}
def _assert_layer0_grads(model, *, role: str) -> None:
blocks = model.transformer.blocks
assert blocks is not None and len(blocks) > 0
trainable = [p for p in blocks[0].parameters() if p.requires_grad]
assert len(trainable) > 0, f"{role} layer 0 has no trainable parameters"
for i, p in enumerate(trainable):
assert p.grad is not None, f"{role} layer 0 param[{i}] has None grad"
assert torch.isfinite(p.grad).all().item(), (
f"{role} layer 0 param[{i}] grad contains NaN/Inf")
assert any(
p.grad.detach().float().norm().item() > 0.0 for p in trainable), (
f"all {role} layer-0 grads are exactly zero; the loss did not "
"reach the first transformer block")
def _assert_finite_losses(loss_map: dict[str, torch.Tensor]) -> None:
for name in ("total_loss", "generator_loss", "fake_score_loss"):
loss = loss_map[name]
assert torch.is_tensor(loss), f"{name} must be a torch.Tensor"
assert torch.isfinite(loss).item(), (
f"{name} is not finite: {loss.item()}")
@pytest.mark.usefixtures("distributed_setup")
def test_streaming_long_tuning_multi_stage_train_steps(
monkeypatch: pytest.MonkeyPatch) -> None:
if not torch.cuda.is_available():
pytest.skip("requires CUDA")
cfg = load_run_config(_FIXTURE)
device = torch.device("cuda:0")
dtype = torch.bfloat16
monkeypatch.setattr(
"fastvideo.train.utils.dataloader."
"build_parquet_t2v_train_dataloader",
lambda *args, **kwargs: None,
)
student = WanCausalModel(
init_from=cfg.models["student"]["init_from"],
training_config=cfg.training,
trainable=True,
)
student.transformer = student.transformer.to(device=device, dtype=dtype)
teacher = WanModel(
init_from=cfg.models["teacher"]["init_from"],
training_config=cfg.training,
trainable=False,
)
teacher.transformer = teacher.transformer.to(device=device, dtype=dtype)
critic = WanModel(
init_from=cfg.models["critic"]["init_from"],
training_config=cfg.training,
trainable=True,
)
critic.transformer = critic.transformer.to(device=device, dtype=dtype)
method = StreamingLongTuningMethod(
cfg=cfg,
role_models={
"student": student,
"teacher": teacher,
"critic": critic,
},
)
method.on_train_start()
batch = _build_synthetic_batch(device, dtype)
# Iteration 0 — short self-forcing stage (stage 0, num_latent_t=3).
loss_map, outputs, metrics = method.single_train_step(dict(batch), 0)
_assert_finite_losses(loss_map)
assert metrics["distill_stage_index"] == 0.0
assert metrics["active_num_latent_t"] == 3.0
assert metrics["update_student"] == 1.0
method.backward(loss_map, outputs, grad_accum_rounds=1)
_assert_layer0_grads(student, role="student")
_assert_layer0_grads(critic, role="critic")
assert all(not p.requires_grad
for p in teacher.transformer.parameters()), (
"teacher must stay frozen")
# Iteration 1 — first streaming chunk: 3 fresh latents, no overlap.
method.optimizers_zero_grad(1)
loss_map, outputs, metrics = method.single_train_step(dict(batch), 1)
_assert_finite_losses(loss_map)
assert metrics["distill_stage_index"] == 1.0
assert metrics["streaming_current_length"] == 3.0
assert metrics["streaming_max_length"] == 6.0
method.backward(loss_map, outputs, grad_accum_rounds=1)
_assert_layer0_grads(student, role="student")
_assert_layer0_grads(critic, role="critic")
vis = outputs["dmd_latent_vis_dict"]
assert int(vis["streaming_chunk_mask"].sum().item()) == 3
# Iteration 2 — overlapped chunk: 2 new latents + 1 overlap anchor.
method.optimizers_zero_grad(2)
loss_map, outputs, metrics = method.single_train_step(dict(batch), 2)
_assert_finite_losses(loss_map)
assert metrics["streaming_current_length"] == 5.0
method.backward(loss_map, outputs, grad_accum_rounds=1)
mask = outputs["dmd_latent_vis_dict"]["streaming_chunk_mask"]
assert mask.dtype == torch.bool
assert int(mask.sum().item()) == 2, (
"overlap latents must be masked out of the training chunk")
# Iteration 3 — final chunk fills the sequence (1 new latent).
method.optimizers_zero_grad(3)
loss_map, outputs, metrics = method.single_train_step(dict(batch), 3)
_assert_finite_losses(loss_map)
assert metrics["streaming_current_length"] == 6.0
method.backward(loss_map, outputs, grad_accum_rounds=1)
# Iteration 4 — the exhausted sequence resets and a fresh chunk starts.
method.optimizers_zero_grad(4)
loss_map, outputs, metrics = method.single_train_step(dict(batch), 4)
_assert_finite_losses(loss_map)
assert metrics["streaming_current_length"] == 3.0, (
"streaming state must reset once the sequence reaches max length")
method.backward(loss_map, outputs, grad_accum_rounds=1)
_assert_layer0_grads(student, role="student")
# Grad-norm regression on the post-reset chunk (grads zeroed at the top
# of this iteration, so the reference captures exactly one step). Keep
# this last: under FASTVIDEO_GRADNORM_UPDATE=1 it records and skips.
check_grad_norm_regression(
"test_streaming_long_tuning",
student.transformer,
)
# ---------------------------------------------------------------------------
# CPU-only helper tests (no GPU, no mocks)
# ---------------------------------------------------------------------------
def _method_without_init() -> StreamingLongTuningMethod:
# Bare instance for exercising stateless helpers; no attributes beyond
# the ones each test sets are touched.
return StreamingLongTuningMethod.__new__(StreamingLongTuningMethod)
def _make_stage(**overrides) -> DistillStage:
defaults = dict(
name="streaming_long",
start_step=0,
end_step=None,
num_latent_t=6,
streaming_training=True,
streaming_chunk_size=3,
streaming_max_length=6,
)
defaults.update(overrides)
return DistillStage(**defaults)
class TestStageSelection:
def test_boundaries_between_stages(self) -> None:
short = _make_stage(
name="self_forcing",
start_step=0,
end_step=10,
streaming_training=False,
streaming_chunk_size=None,
streaming_max_length=None,
)
long = _make_stage(start_step=10)
stages = [short, long]
assert select_distill_stage(stages, 0) is short
assert select_distill_stage(stages, 9) is short
assert select_distill_stage(stages, 10) is long
assert select_distill_stage(stages, 10_000) is long
def test_past_last_closed_stage_falls_back_to_last(self) -> None:
short = _make_stage(
name="self_forcing",
start_step=0,
end_step=10,
streaming_training=False,
streaming_chunk_size=None,
streaming_max_length=None,
)
assert select_distill_stage([short], 99) is short
class TestScheduleParsing:
def test_rejects_overlapping_stages(self) -> None:
with pytest.raises(ValueError, match="ordered and non-overlapping"):
parse_multi_phased_distill_schedule(
[
{"stage": "self_forcing", "start_step": 0, "end_step": 10},
{"stage": "streaming_long", "start_step": 5},
],
default_num_latent_t=6,
default_streaming_chunk_size=3,
)
def test_rejects_fixed_overlap_not_smaller_than_chunk(self) -> None:
with pytest.raises(ValueError, match="smaller than"):
parse_multi_phased_distill_schedule(
[{
"stage": "streaming_long",
"streaming_chunk_size": 3,
"streaming_fixed_overlap_latents": 3,
}],
default_num_latent_t=6,
)
class TestPadFirstFrameLatent:
def test_pads_short_condition_with_zeros(self) -> None:
method = _method_without_init()
first_frame = torch.ones(1, 4, 1, 2, 2)
batch = method._with_padded_first_frame_latent(
{"first_frame_latent": first_frame},
target_latent_frames=3,
)
padded = batch["first_frame_latent"]
assert padded.shape == (1, 4, 3, 2, 2)
assert torch.equal(padded[:, :, :1], first_frame)
assert padded[:, :, 1:].abs().sum().item() == 0.0
def test_passthrough_without_condition_or_when_long_enough(self) -> None:
method = _method_without_init()
raw: dict = {"other": 1}
assert method._with_padded_first_frame_latent(
raw, target_latent_frames=3) is raw
full = {"first_frame_latent": torch.zeros(1, 4, 3, 2, 2)}
assert method._with_padded_first_frame_latent(
full, target_latent_frames=3) is full
def test_rejects_non_5d_condition(self) -> None:
method = _method_without_init()
with pytest.raises(ValueError, match=r"\[B, C, T, H, W\]"):
method._with_padded_first_frame_latent(
{"first_frame_latent": torch.zeros(4, 1, 2, 2)},
target_latent_frames=3,
)
class TestTruncateBatchLatents:
def test_slices_every_temporal_field(self) -> None:
method = _method_without_init()
# No _build_attention_metadata on the student → metadata rebuild is
# a no-op, so only the tensor slicing is exercised.
method.student = object()
batch = SimpleNamespace(
latents=torch.randn(1, 6, 2, 4, 4),
noise_latents=torch.randn(1, 6, 2, 4, 4),
noisy_model_input=torch.randn(1, 6, 2, 4, 4),
noise=torch.randn(1, 6, 2, 4, 4),
timesteps=torch.zeros(1, 6),
sigmas=torch.linspace(0, 1, 6),
raw_latent_shape=(1, 6, 2, 4, 4),
attn_metadata=None,
attn_metadata_vsa=None,
)
method._truncate_batch_latents(batch, 3)
assert batch.latents.shape == (1, 3, 2, 4, 4)
assert batch.noise_latents.shape == (1, 3, 2, 4, 4)
assert batch.noisy_model_input.shape == (1, 3, 2, 4, 4)
assert batch.noise.shape == (1, 3, 2, 4, 4)
assert batch.timesteps.shape == (1, 3)
assert batch.sigmas.shape == (3, )
assert batch.raw_latent_shape == (1, 2, 3, 4, 4)
class TestStageChunkSize:
def test_prefers_stage_then_config_then_rollout_block(self) -> None:
method = _method_without_init()
method.method_config = {"streaming_chunk_size": 5}
method._chunk_size = 7
assert method._stage_chunk_size(_make_stage()) == 3
assert method._stage_chunk_size(
_make_stage(streaming_chunk_size=None)) == 5
method.method_config = {}
assert method._stage_chunk_size(
_make_stage(streaming_chunk_size=None)) == 7
def test_rejects_non_positive_chunk_size(self) -> None:
method = _method_without_init()
method.method_config = {}
method._chunk_size = 0
with pytest.raises(ValueError, match="must be positive"):
method._stage_chunk_size(_make_stage(streaming_chunk_size=None))
fastvideo/tests/train/trainer/__init__.py
View File
fastvideo/tests/train/trainer/test_validation.py
+139
View File
@@ -0,0 +1,139 @@
# SPDX-License-Identifier: Apache-2.0
"""CPU-only integration tests for Trainer validation hook dispatch."""
from __future__ import annotations
from types import SimpleNamespace
from typing import Any
import torch
from fastvideo.train.callbacks.validation import ValidationCallback
from fastvideo.train.trainer import Trainer
from fastvideo.train.utils.training_config import TrainingConfig
class _RecordingValidationCallback(ValidationCallback):
def __init__(self, **kwargs: Any) -> None:
super().__init__(**kwargs)
self.run_calls: list[int] = []
def _run_validation(self, method: Any, step: int) -> None: # type: ignore[override]
del method
self.run_calls.append(step)
class _DummyTracker:
def __init__(self) -> None:
self.logs: list[tuple[dict[str, float], int]] = []
self.finished = False
def log(self, metrics: dict[str, float], step: int) -> None:
self.logs.append((metrics, step))
def finish(self) -> None:
self.finished = True
class _DummyMethod:
def __init__(self) -> None:
self.weight = torch.nn.Parameter(torch.tensor(1.0))
self.train_start_calls = 0
self.zero_grad_steps: list[int] = []
self.optimizer_steps: list[int] = []
self.backward_calls = 0
self.tracker = None
def set_tracker(self, tracker: Any) -> None:
self.tracker = tracker
def on_train_start(self) -> None:
self.train_start_calls += 1
def manages_optimization(self) -> bool:
return False
def single_train_step(
self,
batch: dict[str, Any],
iteration: int,
) -> tuple[dict[str, torch.Tensor], dict[str, Any], dict[str, float]]:
assert batch["sample"] == "x"
loss = self.weight * 0.0 + 1.0
return {"total_loss": loss}, {}, {"iteration_seen": float(iteration)}
def backward(
self,
loss_map: dict[str, torch.Tensor],
outputs: dict[str, Any],
*,
grad_accum_rounds: int,
) -> None:
del outputs
self.backward_calls += 1
(loss_map["total_loss"] / grad_accum_rounds).backward()
def optimizers_schedulers_step(self, iteration: int) -> None:
self.optimizer_steps.append(iteration)
def optimizers_zero_grad(self, iteration: int) -> None:
self.zero_grad_steps.append(iteration)
self.weight.grad = None
def test_trainer_runs_validation_callback_during_training(
monkeypatch,
) -> None:
tracker = _DummyTracker()
group = SimpleNamespace(rank=0, local_rank=0, rank_in_group=0, world_size=1)
monkeypatch.setattr("fastvideo.train.trainer.get_world_group", lambda: group)
monkeypatch.setattr("fastvideo.train.trainer.get_sp_group", lambda: group)
monkeypatch.setattr(
"fastvideo.train.callbacks.validation.get_world_group",
lambda: group,
)
monkeypatch.setattr(
"fastvideo.train.callbacks.validation.get_sp_group",
lambda: group,
)
monkeypatch.setattr(
"fastvideo.train.trainer.build_tracker",
lambda *args, **kwargs: tracker,
)
cfg = TrainingConfig()
cfg.tracker.project_name = ""
cfg.loop.gradient_accumulation_steps = 1
callback_configs = {
"validation": {
"_target_": f"{__name__}._RecordingValidationCallback",
"pipeline_target": "unused.pipeline.Target",
"dataset_file": "unused.json",
"every_steps": 2,
}
}
trainer = Trainer(
cfg,
callback_configs=callback_configs,
)
method = _DummyMethod()
trainer.run(
method,
dataloader=[{"sample": "x"}],
max_steps=3,
)
validation = trainer.callbacks._callbacks["validation"]
assert isinstance(validation, _RecordingValidationCallback)
assert validation.run_calls == [0, 2]
assert method.train_start_calls == 1
assert method.backward_calls == 3
assert method.zero_grad_steps == [0, 1, 2, 3]
assert method.optimizer_steps == [1, 2, 3]
assert [step for _, step in tracker.logs] == [1, 2, 3]
assert tracker.finished is True
fastvideo/tests/train/utils/test_config.py
+43
View File
@@ -155,6 +155,7 @@ def test_full_yaml_populates_all_training_fields(tmp_path: Path) -> None:
assert t.distributed.pin_cpu_memory is True
assert t.data.train_batch_size == 2
assert t.data.data_path == "/some/path"
assert t.data.num_frames == 33
assert t.data.seed == 42
@@ -244,6 +245,48 @@ def test_betas_parses_list_and_string_forms(
assert cfg.training.optimizer.betas == expected
def test_data_path_mapping_parses_repeat_counts(tmp_path: Path) -> None:
# Config loading should preserve structured multi-dataset paths so the
# dataset layer can interpret repeat counts later.
data = _minimal_yaml()
data["training"] = {
"data": {
"data_path": {
"data/path1": 1,
"data/path2": 2,
}
}
}
cfg = load_run_config(_write_yaml(tmp_path, data))
assert cfg.training.data.data_path == {
"data/path1": 1,
"data/path2": 2,
}
def test_dotted_override_replaces_mapping_data_path(tmp_path: Path) -> None:
# A dict-valued data_path is a single leaf for overrides: a scalar
# --training.data.data_path replaces the whole mapping.
data = _minimal_yaml()
data["training"] = {
"data": {
"data_path": {
"data/path1": 1,
"data/path2": 2,
}
}
}
cfg = load_run_config(
_write_yaml(tmp_path, data),
overrides=["--training.data.data_path=data/only"],
)
assert cfg.training.data.data_path == "data/only"
def test_dotted_overrides_apply_with_type_coercion(tmp_path: Path) -> None:
path = _write_yaml(tmp_path, _minimal_yaml())
overrides = [
fastvideo/train/callbacks/validation.py
+760 -49
View File
@@ -9,8 +9,11 @@ section. The pipeline class is resolved from
from __future__ import annotations
import contextlib
import gc
import json
import os
from dataclasses import dataclass
from copy import deepcopy
from dataclasses import dataclass, field
from typing import Any, TYPE_CHECKING
import imageio
@@ -46,6 +49,64 @@ logger = init_logger(__name__)
class _ValidationStepResult:
videos: list[list[np.ndarray]]
captions: list[str]
overlay_videos: list[list[np.ndarray]] = field(default_factory=list)
overlay_captions: list[str] = field(default_factory=list)
ref_videos: list[str | None] = field(default_factory=list)
actions: list[dict[str, Any] | None] = field(default_factory=list)
mouse_pitch_signs: list[int | None] = field(default_factory=list)
@dataclass(slots=True)
class _ValidationMetricStats:
sums: dict[str, float] = field(default_factory=dict)
counts: dict[str, float] = field(default_factory=dict)
per_video: list[dict[str, Any]] = field(default_factory=list)
errors: list[str] = field(default_factory=list)
@dataclass(slots=True)
class _SavedValidationVideos:
filenames: list[str] = field(default_factory=list)
indices: list[int] = field(default_factory=list)
@dataclass(slots=True)
class _ValidationMetricsConfig:
enabled: bool = False
names: list[str] = field(default_factory=list)
device: str = "cuda"
calibration_path: str | None = None
mouse_pitch_sign: int = 1
skip_missing_deps: bool = True
strict: bool = False
unload_after_validation: bool = True
loader_threads: int = 1
prefetch_factor: int = 2
log_prefix: str = "metrics/validation"
DEFAULT_VALIDATION_VBENCH_METRICS = [
"vbench.imaging_quality",
"vbench.aesthetic_quality",
"vbench.temporal_flickering",
"vbench.motion_smoothness",
"vbench.subject_consistency",
"vbench.background_consistency",
"vbench.dynamic_degree",
]
SYNTHETIC_OPTICAL_FLOW_METRIC = "optical_flow.synthetic_optical_flow"
SYNTHETIC_OPTICAL_FLOW_LOG_KEYS = (
"mf_epe_mean",
"mf_angle_err_mean",
"mf_cosine_mean",
"mf_mag_ratio_mean",
"pixel_epe_mean_mean",
"px_angle_rmse_mean",
"fl_all_mean",
"foe_dist_mean",
"flow_kl_2d_mean",
)
class ValidationCallback(Callback):
@@ -68,6 +129,8 @@ class ValidationCallback(Callback):
num_frames: int | None = None,
output_dir: str | None = None,
sampling_timesteps: list[int] | None = None,
overlay_actions: bool = False,
keyboard_value_scale: float = 1.0,
offload_training_state: bool = False,
unload_pipeline_after_validation: bool = False,
**pipeline_kwargs: Any,
@@ -80,6 +143,11 @@ class ValidationCallback(Callback):
self.num_frames = (int(num_frames) if num_frames is not None else None)
self.output_dir = (str(output_dir) if output_dir is not None else None)
self.sampling_timesteps = ([int(s) for s in sampling_timesteps] if sampling_timesteps is not None else None)
self.overlay_actions = self._coerce_bool(overlay_actions)
# Validation-only action amplification for world model; training keeps raw action values.
self.keyboard_value_scale = float(keyboard_value_scale)
metrics_config = pipeline_kwargs.pop("metrics", None)
self.metrics_config = self._parse_metrics_config(metrics_config)
self.offload_training_state = self._coerce_bool(offload_training_state)
self.unload_pipeline_after_validation = self._coerce_bool(unload_pipeline_after_validation)
self.pipeline_kwargs = dict(pipeline_kwargs)
@@ -88,10 +156,56 @@ class ValidationCallback(Callback):
self._pipeline: Any | None = None
self._pipeline_key: tuple[Any, ...] | None = None
self._sampling_param: SamplingParam | None = None
self._metric_evaluator: Any | None = None
self.tracker: Any = DummyTracker()
self.validation_random_generator: (torch.Generator | None) = None
self.seed: int = 0
@staticmethod
def _parse_metrics_config(config: Any) -> _ValidationMetricsConfig:
if config is None or config is False:
return _ValidationMetricsConfig(enabled=False)
if config is True:
return _ValidationMetricsConfig(
enabled=True,
names=list(DEFAULT_VALIDATION_VBENCH_METRICS),
)
if isinstance(config, str):
return _ValidationMetricsConfig(
enabled=True,
names=[config],
)
if isinstance(config, list | tuple):
return _ValidationMetricsConfig(
enabled=bool(config),
names=[str(name) for name in config],
)
if not isinstance(config, dict):
raise TypeError("callbacks.validation.metrics must be a bool, string, list, or mapping")
enabled = bool(config.get("enabled", True))
raw_names = config.get("names", config.get("metrics", None))
if raw_names is None:
names = list(DEFAULT_VALIDATION_VBENCH_METRICS)
elif isinstance(raw_names, str):
names = [raw_names]
else:
names = [str(name) for name in raw_names]
return _ValidationMetricsConfig(
enabled=enabled and bool(names),
names=names,
device=str(config.get("device", "cuda")),
calibration_path=(str(config["calibration_path"]) if config.get("calibration_path") is not None else None),
mouse_pitch_sign=int(config.get("mouse_pitch_sign", 1)),
skip_missing_deps=bool(config.get("skip_missing_deps", True)),
strict=bool(config.get("strict", False)),
unload_after_validation=bool(config.get("unload_after_validation", True)),
loader_threads=int(config.get("loader_threads", 1)),
prefetch_factor=int(config.get("prefetch_factor", 2)),
log_prefix=str(config.get("log_prefix", "metrics/validation")),
)
@staticmethod
def _coerce_bool(value: Any) -> bool:
if isinstance(value, str):
@@ -354,6 +468,7 @@ class ValidationCallback(Callback):
try:
transformer.eval()
num_sp_groups = (self.world_group.world_size // self.sp_group.world_size)
sp = self._get_sampling_param()
for num_inference_steps in self.sampling_steps:
result = self._run_validation_for_steps(
@@ -364,72 +479,536 @@ class ValidationCallback(Callback):
if self.rank_in_sp_group != 0:
continue
os.makedirs(
output_dir,
exist_ok=True,
)
local_videos = self._save_validation_videos(
result.videos,
output_dir=output_dir,
step=step,
num_inference_steps=num_inference_steps,
fps=sp.fps,
)
local_overlay_videos = self._save_validation_videos(
result.overlay_videos,
output_dir=output_dir,
step=step,
num_inference_steps=num_inference_steps,
fps=sp.fps,
suffix="_overlay",
)
local_video_filenames = local_videos.filenames
local_overlay_video_filenames = local_overlay_videos.filenames
local_captions = self._select_by_indices(
result.captions,
local_videos.indices,
)
local_overlay_captions = self._select_by_indices(
result.overlay_captions,
local_overlay_videos.indices,
)
local_ref_videos = self._select_by_indices(
result.ref_videos,
local_videos.indices,
)
local_actions = self._select_by_indices(
result.actions,
local_videos.indices,
)
local_mouse_pitch_signs = self._select_by_indices(
result.mouse_pitch_signs,
local_videos.indices,
)
local_metric_stats = self._evaluate_validation_metrics(
video_filenames=local_video_filenames,
captions=local_captions,
ref_videos=local_ref_videos,
actions=local_actions,
mouse_pitch_signs=local_mouse_pitch_signs,
fps=sp.fps,
output_dir=output_dir,
step=step,
num_inference_steps=num_inference_steps,
)
if self.global_rank == 0:
all_videos = list(result.videos)
all_captions = list(result.captions)
all_video_filenames = list(local_video_filenames)
all_overlay_video_filenames = list(local_overlay_video_filenames)
all_captions = list(local_captions)
all_overlay_captions = list(local_overlay_captions)
all_metric_stats = local_metric_stats
for sp_idx in range(1, num_sp_groups):
src = (sp_idx * self.sp_world_size)
recv_v = (self.world_group.recv_object(src=src))
recv_c = (self.world_group.recv_object(src=src))
all_videos.extend(recv_v)
recv_ov = (self.world_group.recv_object(src=src))
recv_oc = (self.world_group.recv_object(src=src))
recv_m = (self.world_group.recv_object(src=src))
all_video_filenames.extend(recv_v)
all_overlay_video_filenames.extend(recv_ov)
all_captions.extend(recv_c)
all_overlay_captions.extend(recv_oc)
self._merge_metric_stats(
all_metric_stats,
recv_m,
)
os.makedirs(
output_dir,
exist_ok=True,
self._log_validation_metrics(
all_metric_stats,
step=step,
)
video_filenames: list[str] = []
sp = self._get_sampling_param()
for i, video in enumerate(all_videos):
fname = os.path.join(
output_dir,
f"validation_step_{step}"
f"_inference_steps_"
f"{num_inference_steps}"
f"_video_{i}.mp4",
)
imageio.mimsave(
fname,
video,
self._log_validation_video_artifacts(
all_video_filenames,
all_captions,
key=f"validation_videos_{num_inference_steps}_steps",
step=step,
fps=sp.fps,
)
if all_overlay_video_filenames:
self._log_validation_video_artifacts(
all_overlay_video_filenames,
all_overlay_captions,
key=(f"validation_videos_{num_inference_steps}"
f"_steps_overlay"),
step=step,
fps=sp.fps,
)
video_filenames.append(fname)
video_logs = []
for fname, cap in zip(
video_filenames,
all_captions,
strict=True,
):
art = self.tracker.video(
fname,
caption=cap,
fps=sp.fps,
)
if art is not None:
video_logs.append(art)
if video_logs:
logs = {
f"validation_videos_"
f"{num_inference_steps}"
f"_steps": video_logs
}
self.tracker.log_artifacts(
logs,
step,
)
else:
self.world_group.send_object(
result.videos,
local_video_filenames,
dst=0,
)
self.world_group.send_object(
result.captions,
local_captions,
dst=0,
)
self.world_group.send_object(
local_overlay_video_filenames,
dst=0,
)
self.world_group.send_object(
local_overlay_captions,
dst=0,
)
self.world_group.send_object(
local_metric_stats,
dst=0,
)
finally:
if was_training:
transformer.train()
self._release_metric_evaluator()
def _save_validation_videos(
self,
videos: list[list[np.ndarray]],
*,
output_dir: str,
step: int,
num_inference_steps: int,
fps: int,
suffix: str = "",
) -> _SavedValidationVideos:
saved = _SavedValidationVideos()
for i, video in enumerate(videos):
fname = os.path.join(
output_dir,
f"validation_step_{step}"
f"_inference_steps_{num_inference_steps}"
f"_rank_{self.global_rank}"
f"_video_{i}{suffix}.mp4",
)
# Validation video encoding can fail due transient ffmpeg or
# filesystem errors; skip the artifact so training can continue.
try:
imageio.mimsave(
fname,
video,
fps=fps,
)
except Exception as exc:
logger.exception(
"Failed to save validation video %s on rank %s; skipping artifact: %s",
fname,
self.global_rank,
exc,
)
with contextlib.suppress(OSError):
os.remove(fname)
continue
saved.filenames.append(fname)
saved.indices.append(i)
return saved
@staticmethod
def _select_by_indices(
values: list[Any],
indices: list[int],
) -> list[Any]:
return [values[i] for i in indices if i < len(values)]
def _log_validation_video_artifacts(
self,
video_filenames: list[str],
captions: list[str],
*,
key: str,
step: int,
fps: int,
) -> None:
video_logs = []
for fname, cap in zip(
video_filenames,
captions,
strict=True,
):
art = self.tracker.video(
fname,
caption=cap,
fps=fps,
)
if art is not None:
video_logs.append(art)
if video_logs:
self.tracker.log_artifacts(
{key: video_logs},
step,
)
# ----------------------------------------------------------
# Metric evaluation
# ----------------------------------------------------------
def _metric_device(self) -> str:
device = self.metrics_config.device
if device == "cuda" and torch.cuda.is_available():
local_rank = int(getattr(self.world_group, "local_rank", 0))
return f"cuda:{local_rank}"
return device
def _get_metric_evaluator(self) -> Any:
if self._metric_evaluator is not None:
return self._metric_evaluator
from fastvideo.eval import Evaluator
cfg = self.metrics_config
self._metric_evaluator = Evaluator(
metrics=cfg.names,
device=self._metric_device(),
num_gpus=1,
loader_threads=cfg.loader_threads,
prefetch_factor=cfg.prefetch_factor,
pre_upload=False,
skip_missing_deps=cfg.skip_missing_deps,
)
logger.info(
"Initialized validation metrics: %s",
", ".join(self._metric_evaluator.metric_names),
)
return self._metric_evaluator
def _release_metric_evaluator(self) -> None:
if self._metric_evaluator is None:
return
if self.metrics_config.unload_after_validation:
self._metric_evaluator.unload()
self._metric_evaluator = None
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
def _evaluate_validation_metrics(
self,
*,
video_filenames: list[str],
captions: list[str],
ref_videos: list[str | None],
actions: list[dict[str, Any] | None],
mouse_pitch_signs: list[int | None],
fps: int,
output_dir: str,
step: int,
num_inference_steps: int,
) -> _ValidationMetricStats:
stats = _ValidationMetricStats()
if not self.metrics_config.enabled or not video_filenames:
return stats
try:
evaluator = self._get_metric_evaluator()
from fastvideo.eval import samples_from
samples = samples_from(
video=video_filenames,
reference=self._available_paths(ref_videos),
text_prompts=captions,
fps=float(fps),
extras=self._validation_metric_extras(
video_filenames=video_filenames,
actions=actions,
mouse_pitch_signs=mouse_pitch_signs,
),
)
results = evaluator.evaluate(samples=samples)
for filename, metric_results in zip(
video_filenames,
results,
strict=True,
):
row: dict[str, Any] = {"path": filename}
self._accumulate_metric_results(
stats,
row,
metric_results,
)
stats.per_video.append(row)
for metric_name, metric_result in results.corpus.items():
row = {"path": "<corpus>"}
self._accumulate_metric_results(
stats,
row,
{metric_name: metric_result},
)
stats.per_video.append(row)
except Exception as exc:
message = ("Validation metric evaluation failed on rank "
f"{self.global_rank}: {exc}")
logger.exception(message)
stats.errors.append(message)
if self.metrics_config.strict:
raise
finally:
if self._metric_evaluator is not None:
self._metric_evaluator.release_cuda_memory()
self._write_metric_summary(
stats,
output_dir=output_dir,
step=step,
num_inference_steps=num_inference_steps,
)
return stats
@staticmethod
def _accumulate_metric_results(
stats: _ValidationMetricStats,
row: dict[str, Any],
metric_results: dict[str, Any],
) -> None:
for metric_name, metric_result in metric_results.items():
details = getattr(metric_result, "details", {}) or {}
if metric_name == SYNTHETIC_OPTICAL_FLOW_METRIC:
if not isinstance(details, dict):
continue
for detail_name in SYNTHETIC_OPTICAL_FLOW_LOG_KEYS:
key = f"{metric_name}.{detail_name}"
ValidationCallback._accumulate_scalar(
stats,
row,
key,
details.get(detail_name),
)
continue
score = getattr(metric_result, "score", None)
ValidationCallback._accumulate_scalar(
stats,
row,
metric_name,
score,
)
if not isinstance(details, dict):
continue
for detail_name, value in details.items():
key = f"{metric_name}.{detail_name}"
ValidationCallback._accumulate_scalar(
stats,
row,
key,
value,
)
@staticmethod
def _accumulate_scalar(
stats: _ValidationMetricStats,
row: dict[str, Any],
key: str,
value: Any,
) -> None:
if not isinstance(value, float | int | np.floating | np.integer):
return
value_float = float(value)
if not np.isfinite(value_float):
return
row[key] = value_float
stats.sums[key] = stats.sums.get(key, 0.0) + value_float
stats.counts[key] = stats.counts.get(key, 0.0) + 1.0
@staticmethod
def _merge_metric_stats(
dst: _ValidationMetricStats,
src: _ValidationMetricStats,
) -> None:
for key, value in src.sums.items():
dst.sums[key] = dst.sums.get(key, 0.0) + float(value)
for key, value in src.counts.items():
dst.counts[key] = dst.counts.get(key, 0.0) + float(value)
dst.per_video.extend(src.per_video)
dst.errors.extend(src.errors)
def _log_validation_metrics(
self,
stats: _ValidationMetricStats,
*,
step: int,
) -> None:
if stats.errors:
message = "Validation metric evaluation failed:\n" + "\n".join(stats.errors)
if self.metrics_config.strict:
raise RuntimeError(message)
logger.warning(message)
logs: dict[str, float] = {}
for key, metric_sum in stats.sums.items():
metric_count = stats.counts.get(key, 0.0)
if metric_count <= 0:
continue
value = float(metric_sum / metric_count)
if not np.isfinite(value):
continue
logs[self._metric_log_name(key)] = value
if logs:
self.tracker.log(
logs,
step,
)
def _metric_log_name(
self,
key: str,
) -> str:
return f"{self.metrics_config.log_prefix}/{key.replace('.', '/')}"
def _write_metric_summary(
self,
stats: _ValidationMetricStats,
*,
output_dir: str,
step: int,
num_inference_steps: int,
) -> None:
if not self.metrics_config.enabled:
return
summary_dir = os.path.join(
output_dir,
"eval",
f"step_{step}",
)
os.makedirs(
summary_dir,
exist_ok=True,
)
metrics = {
key: float(stats.sums[key] / stats.counts[key])
for key in sorted(stats.sums) if stats.counts.get(key, 0.0) > 0
}
payload = {
"step": int(step),
"num_inference_steps": int(num_inference_steps),
"rank": int(self.global_rank),
"metrics": metrics,
"per_video": stats.per_video,
"errors": stats.errors,
}
path = os.path.join(
summary_dir,
f"inference_steps_{num_inference_steps}_rank_{self.global_rank}.json",
)
with open(
path,
"w",
encoding="utf-8",
) as f:
json.dump(
payload,
f,
indent=2,
)
def _validation_metric_extras(
self,
*,
video_filenames: list[str],
actions: list[dict[str, Any] | None],
mouse_pitch_signs: list[int | None],
) -> list[dict[str, Any]]:
extras: list[dict[str, Any]] = []
for i, _filename in enumerate(video_filenames):
extra: dict[str, Any] = {}
action = actions[i] if i < len(actions) else None
if action is not None:
extra["actions"] = action
if self.metrics_config.calibration_path is not None:
extra["calibration"] = self.metrics_config.calibration_path
mouse_pitch_sign = mouse_pitch_signs[i] if i < len(mouse_pitch_signs) else None
extra["mouse_pitch_sign"] = int(mouse_pitch_sign or self.metrics_config.mouse_pitch_sign)
extras.append(extra)
return extras
@staticmethod
def _available_paths(paths: list[str | None]) -> list[str] | None:
if not paths or any(path is None for path in paths):
return None
return [str(path) for path in paths]
@staticmethod
def _validation_actions(validation_batch: dict[str, Any]) -> dict[str, Any] | None:
keyboard = validation_batch.get("keyboard_cond")
mouse = validation_batch.get("mouse_cond")
if keyboard is not None and mouse is not None:
return {
"keyboard": np.asarray(keyboard),
"mouse": np.asarray(mouse),
}
action_path = validation_batch.get("action_path")
if not isinstance(action_path, str) or not os.path.isfile(action_path):
return None
try:
actions_obj = np.load(action_path, allow_pickle=True)
if isinstance(actions_obj, np.ndarray) and actions_obj.dtype == object:
actions_obj = actions_obj.item()
except Exception as exc:
logger.warning(
"Failed to load validation action file %s for metrics: %s",
action_path,
exc,
)
return None
if not isinstance(actions_obj, dict):
return None
keyboard = actions_obj.get("keyboard")
mouse = actions_obj.get("mouse")
if keyboard is None or mouse is None:
return None
return {
"keyboard": np.asarray(keyboard),
"mouse": np.asarray(mouse),
}
def _validation_mouse_pitch_sign(
self,
validation_batch: dict[str, Any],
) -> int:
if validation_batch.get("mouse_pitch_sign") is not None:
return int(validation_batch["mouse_pitch_sign"])
flipped = validation_batch.get("mouse_pitch_flipped")
if isinstance(flipped, str):
flipped = flipped.lower() in {"1", "true", "yes", "y"}
if bool(flipped):
return -1
return int(self.metrics_config.mouse_pitch_sign)
# ----------------------------------------------------------
# Pipeline management
@@ -437,9 +1016,51 @@ class ValidationCallback(Callback):
def _get_sampling_param(self) -> SamplingParam:
if self._sampling_param is None:
self._sampling_param = (SamplingParam.from_pretrained(self.training_config.model_path))
self._sampling_param = (SamplingParam.from_pretrained(self._pipeline_model_path()))
return self._sampling_param
def _pipeline_model_path(self) -> str:
tc = self.training_config
pipeline_config = getattr(tc, "pipeline_config", None)
model_path = getattr(pipeline_config, "_fastvideo_train_model_path", None)
return str(model_path or tc.model_path)
def _validation_pipeline_config(self, transformer: torch.nn.Module) -> Any:
tc = self.training_config
pipeline_config = deepcopy(tc.pipeline_config)
self._sync_runtime_dit_arch_config(
pipeline_config,
transformer,
)
return pipeline_config
@staticmethod
def _sync_runtime_dit_arch_config(
pipeline_config: Any,
transformer: torch.nn.Module,
) -> None:
dit_config = getattr(
pipeline_config,
"dit_config",
None,
)
arch_config = getattr(
dit_config,
"arch_config",
None,
)
if arch_config is None:
return
for name in ("local_attn_size", "sink_size"):
if not hasattr(arch_config, name) or not hasattr(transformer, name):
continue
setattr(
arch_config,
name,
getattr(transformer, name),
)
def _get_pipeline(
self,
*,
@@ -479,9 +1100,18 @@ class ValidationCallback(Callback):
kwargs.update(self.pipeline_kwargs)
self._pipeline = PipelineCls.from_pretrained(
tc.model_path,
self._pipeline_model_path(),
**kwargs,
)
if tc.pipeline_config is not None:
self._pipeline.fastvideo_args.pipeline_config = self._validation_pipeline_config(transformer)
arch_config = self._pipeline.fastvideo_args.pipeline_config.dit_config.arch_config
logger.info(
"Validation pipeline runtime config: local_attn_size=%s sink_size=%s boundary_ratio=%s",
getattr(arch_config, "local_attn_size", None),
getattr(arch_config, "sink_size", None),
getattr(self._pipeline.fastvideo_args.pipeline_config.dit_config, "boundary_ratio", None),
)
self._pipeline_key = key
return self._pipeline
@@ -552,8 +1182,47 @@ class ValidationCallback(Callback):
if ("image" in validation_batch and validation_batch["image"] is not None):
batch.pil_image = validation_batch["image"]
self._attach_action_conditions(
batch,
validation_batch,
sampling_param.num_frames,
)
return batch
def _attach_action_conditions(
self,
batch: ForwardBatch,
validation_batch: dict[str, Any],
num_frames: int,
) -> None:
for name in ("keyboard_cond", "mouse_cond"):
value = validation_batch.get(name)
if value is None:
continue
array = np.asarray(value)
if array.size == 0:
continue
array = array[:num_frames]
tensor = torch.as_tensor(
array,
dtype=torch.bfloat16,
)
student = getattr(self.method, "student", None)
prepare_action = getattr(student, "prepare_validation_action_condition", None)
if prepare_action is not None:
tensor = prepare_action(
tensor,
name=name,
keyboard_value_scale=self.keyboard_value_scale,
)
tensor = tensor.unsqueeze(0)
setattr(
batch,
name,
tensor,
)
# ----------------------------------------------------------
# Validation loop
# ----------------------------------------------------------
@@ -579,6 +1248,10 @@ class ValidationCallback(Callback):
tc,
model_path=tc.model_path,
)
self._sync_runtime_dit_arch_config(
inference_args.pipeline_config,
transformer,
)
# Propagate sampling_timesteps to pipeline_config so
# causal/DMD denoising stages can read them.
@@ -586,7 +1259,12 @@ class ValidationCallback(Callback):
inference_args.pipeline_config.dmd_denoising_steps = ([int(s) for s in self.sampling_timesteps])
videos: list[list[np.ndarray]] = []
overlay_videos: list[list[np.ndarray]] = []
captions: list[str] = []
overlay_captions: list[str] = []
ref_videos: list[str | None] = []
actions: list[dict[str, Any] | None] = []
mouse_pitch_signs: list[int | None] = []
for validation_batch in dataloader:
batch = self._prepare_validation_batch(
@@ -597,6 +1275,11 @@ class ValidationCallback(Callback):
assert (batch.prompt is not None and isinstance(batch.prompt, str))
captions.append(batch.prompt)
ref_video = validation_batch.get("ref_video")
ref_videos.append(ref_video if isinstance(ref_video, str) else None)
action = self._validation_actions(validation_batch)
actions.append(action)
mouse_pitch_signs.append(self._validation_mouse_pitch_sign(validation_batch))
with torch.no_grad():
output_batch = pipeline.forward(
@@ -621,10 +1304,38 @@ class ValidationCallback(Callback):
x = (x.transpose(0, 1).transpose(1, 2).squeeze(-1))
frames.append((x * 255).numpy().astype(np.uint8))
videos.append(frames)
if self.overlay_actions:
overlay_frames = self._post_process_validation_frames(
frames,
action=action,
)
if overlay_frames is not None:
overlay_videos.append(overlay_frames)
overlay_captions.append(batch.prompt)
return _ValidationStepResult(
videos=videos,
captions=captions,
overlay_videos=overlay_videos,
overlay_captions=overlay_captions,
ref_videos=ref_videos,
actions=actions,
mouse_pitch_signs=mouse_pitch_signs,
)
def _post_process_validation_frames(
self,
frames: list[np.ndarray],
*,
action: dict[str, Any] | None,
) -> list[np.ndarray] | None:
student = getattr(self.method, "student", None)
post_process = getattr(student, "post_process_validation_frames", None)
if post_process is None:
return None
return post_process(
frames,
action=action,
)
# ----------------------------------------------------------
fastvideo/train/methods/distribution_matching/__init__.py
+3
View File
@@ -3,8 +3,11 @@
from fastvideo.train.methods.distribution_matching.dmd2 import DMD2Method
from fastvideo.train.methods.distribution_matching.self_forcing import (
SelfForcingMethod, )
from fastvideo.train.methods.distribution_matching.streaming_long_tuning import (
StreamingLongTuningMethod, )
__all__ = [
"DMD2Method",
"SelfForcingMethod",
"StreamingLongTuningMethod",
]
fastvideo/train/methods/distribution_matching/streaming_long_tuning.py
+1471
View File
File diff suppressed because it is too large Load Diff
fastvideo/train/models/matrixgame2/matrixgame2.py
+166 -5
View File
@@ -13,6 +13,7 @@ from fastvideo.distributed import (
get_sp_group,
get_world_group,
)
from fastvideo.models.dits.matrixgame2.utils import scale_keyboard_condition
from fastvideo.pipelines import TrainingBatch
from fastvideo.training.training_utils import normalize_dit_input
@@ -84,12 +85,14 @@ class MatrixGame2Model(WanModel):
raw_batch,
key="keyboard_cond",
expected_frames=self._expected_action_frames(tc.data.num_latent_t),
expected_dim=self._expected_action_dim("keyboard"),
dtype=dtype,
)
mouse_cond = self._get_optional_action(
raw_batch,
key="mouse_cond",
expected_frames=self._expected_action_frames(tc.data.num_latent_t),
expected_dim=self._expected_action_dim("mouse"),
dtype=dtype,
)
@@ -172,23 +175,64 @@ class MatrixGame2Model(WanModel):
raise ValueError("text_dict cannot be None for Matrix-Game 2.0")
hidden_states = noise_input.permute(0, 2, 1, 3, 4)
if hidden_states.shape[1] == 16:
cond_latents = text_dict.get("image_latents")
# Streaming long tuning scores a window at a time; use the
# window-specific first-image latent when the method provides it.
cond_latents = self._get_score_window_image_latents(
text_dict,
batch_info=text_dict.get("_batch_info"),
num_latents=hidden_states.shape[2],
)
if cond_latents is None:
raise RuntimeError("Matrix-Game 2.0 requires image_latents in conditional_dict "
"when noise_input has 16 channels")
num_t = hidden_states.shape[2]
cond_latents = cond_latents[:, :, :num_t]
hidden_states = torch.cat([hidden_states, cond_latents], dim=1)
# For streaming long tuning, align actions to the same score window as
# the latent/image condition; otherwise this falls back to the prefix.
keyboard_cond, mouse_cond = self._get_score_window_actions(
text_dict,
batch_info=text_dict.get("_batch_info"),
num_latents=hidden_states.shape[2],
)
return {
"hidden_states": hidden_states,
"encoder_hidden_states": None,
"timestep": timestep.to(device=self.device, dtype=torch.bfloat16),
"encoder_hidden_states_image": text_dict["encoder_hidden_states_image"],
"keyboard_cond": text_dict["keyboard_cond"],
"mouse_cond": text_dict["mouse_cond"],
"keyboard_cond": keyboard_cond,
"mouse_cond": mouse_cond,
"return_dict": False,
}
def predict_noise(
self,
noisy_latents: torch.Tensor,
timestep: torch.Tensor,
batch: TrainingBatch,
*,
conditional: bool,
cfg_uncond: dict[str, Any] | None = None,
attn_kind: Literal["dense", "vsa"] = "dense",
) -> torch.Tensor:
# Streaming long tuning re-anchors each score window on its first
# image/latent, so pass window metadata through existing dict plumbing.
info = getattr(batch, "matrixgame_streaming_chunk_info", None)
for text_dict in (batch.conditional_dict, batch.unconditional_dict):
if isinstance(text_dict, dict):
text_dict["_batch_info"] = info
try:
return super().predict_noise(
noisy_latents,
timestep,
batch,
conditional=conditional,
cfg_uncond=cfg_uncond,
attn_kind=attn_kind,
)
finally:
for text_dict in (batch.conditional_dict, batch.unconditional_dict):
if isinstance(text_dict, dict):
text_dict.pop("_batch_info", None)
def _build_matrixgame_cond_concat(
self,
image_latents: torch.Tensor,
@@ -235,21 +279,138 @@ class MatrixGame2Model(WanModel):
).transpose(1, 2)
return torch.cat([mask_lat_size, image_latents], dim=1)
def _get_score_window_image_latents(
self,
text_dict: dict[str, Any],
*,
batch_info: dict[str, Any] | None,
num_latents: int,
) -> torch.Tensor | None:
# Prefer a window-specific first-image latent when long tuning has
# recomputed one; otherwise keep the original MG first-frame condition.
if batch_info is not None:
condition_latents = batch_info.get("condition_image_latents")
if condition_latents is not None:
if condition_latents.shape[2] != num_latents:
raise ValueError("condition_image_latents temporal length must match "
"noise_input: "
f"{condition_latents.shape[2]} vs {num_latents}")
return condition_latents
image_latents = text_dict.get("image_latents")
if image_latents is None:
return None
return image_latents[:, :, :num_latents]
def _get_score_window_actions(
self,
text_dict: dict[str, Any],
*,
batch_info: dict[str, Any] | None,
num_latents: int,
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
frame_start = 0
if batch_info is not None:
frame_start = int(batch_info.get("frame_start", 0) or 0)
frame_end = frame_start + int(num_latents)
action_frame_start = frame_start * self._temporal_compression_ratio()
action_frame_end = ((frame_end - 1) * self._temporal_compression_ratio()) + 1
def _slice_action(action: torch.Tensor | None) -> torch.Tensor | None:
if action is None:
return None
if action.shape[1] < action_frame_end:
raise ValueError("Matrix-Game 2.0 score action tensor is shorter than "
"the requested streaming window: "
f"got={action.shape[1]}, required>={action_frame_end}")
return action[:, action_frame_start:action_frame_end]
return (
_slice_action(text_dict.get("keyboard_cond")),
_slice_action(text_dict.get("mouse_cond")),
)
def _get_optional_action(
self,
raw_batch: dict[str, Any],
*,
key: str,
expected_frames: int,
expected_dim: int,
dtype: torch.dtype,
) -> torch.Tensor | None:
value = raw_batch.get(key)
if value is None or value.numel() == 0:
return None
# Actions must follow the same first-image window as the latent
# condition; allow only padded zero channels during format alignment.
value = self._align_action_feature_dim(
value,
name=key,
expected_dim=expected_dim,
)
if value.shape[1] < expected_frames:
raise ValueError(f"{key} has {value.shape[1]} frames but requires at least {expected_frames}")
return value[:, :expected_frames].to(device=self.device, dtype=dtype)
def prepare_validation_action_condition(
self,
action: torch.Tensor,
*,
name: str,
keyboard_value_scale: float = 1.0,
) -> torch.Tensor:
channel = name.removesuffix("_cond")
action = self._align_action_feature_dim(
action,
name=name,
expected_dim=self._expected_action_dim(channel),
)
if name == "keyboard_cond":
action = scale_keyboard_condition(action, float(keyboard_value_scale))
return action
def post_process_validation_frames(
self,
frames: list[Any],
*,
action: dict[str, Any] | None = None,
) -> list[Any] | None:
if action is None:
return None
from fastvideo.models.dits.matrixgame2.utils import (
overlay_validation_actions_on_frames, )
return overlay_validation_actions_on_frames(
frames,
keyboard_cond=action.get("keyboard"),
mouse_cond=action.get("mouse"),
)
def _expected_action_dim(self, channel: str) -> int:
action_config = getattr(self.transformer, "action_config", {}) or {}
return int(action_config.get(f"{channel}_dim_in", 0) or 0)
def _align_action_feature_dim(
self,
action: torch.Tensor,
*,
name: str,
expected_dim: int,
) -> torch.Tensor:
if expected_dim <= 0:
return action
actual_dim = int(action.shape[-1])
if actual_dim == expected_dim:
return action
if actual_dim > expected_dim:
extra = action[..., expected_dim:]
if bool(torch.count_nonzero(extra).item() == 0):
return action[..., :expected_dim]
raise ValueError(f"{name} feature dim mismatch: got={actual_dim}, "
f"expected={expected_dim}, and trailing channels are not all zero.")
raise ValueError(f"{name} feature dim mismatch: got={actual_dim}, expected={expected_dim}")
def _expected_action_frames(self, num_latent_t: int) -> int:
return (num_latent_t - 1) * self._temporal_compression_ratio() + 1
fastvideo/train/models/wan/wan_causal.py
+7
View File
@@ -129,6 +129,13 @@ class WanCausalModel(WanModel, CausalModelBase):
if (self._should_snapshot_streaming_cache() and torch.is_grad_enabled()):
kv_cache = self._snapshot_kv_cache_indices(kv_cache)
# The cross-attn cache is stateful inside the checkpointed block
# forward (is_init flip + k/v writes); a checkpoint recompute that
# sees the mutated dict skips the k/v projections and trips
# torch's saved-tensor-count check. Grad-enabled forwards
# recompute cross k/v instead, matching the legacy
# gradient_checkpointing branch of _forward_inference.
crossattn_cache = None
model_kwargs: dict[str, Any] = {
"kv_cache": kv_cache,
fastvideo/train/utils/config.py
+11 -1
View File
@@ -343,6 +343,15 @@ def _build_training_config(
if init_from is not None:
model_path = str(init_from)
raw_data_path = da.get("data_path", "") or ""
data_path: str | list[str] | dict[str, int]
if isinstance(raw_data_path, dict):
data_path = {str(path): int(repeat) for path, repeat in raw_data_path.items()}
elif isinstance(raw_data_path, list | tuple):
data_path = [str(path) for path in raw_data_path]
else:
data_path = str(raw_data_path)
preprocessed_data_type = str(da.get("preprocessed_data_type", "t2v") or "t2v").strip().lower()
if preprocessed_data_type not in {"t2v", "text_only"}:
raise ValueError("training.data.preprocessed_data_type must be one of "
@@ -359,7 +368,7 @@ def _build_training_config(
pin_cpu_memory=bool(d.get("pin_cpu_memory", False)),
),
data=DataConfig(
data_path=str(da.get("data_path", "") or ""),
data_path=data_path,
preprocessed_data_type=preprocessed_data_type,
train_batch_size=int(da.get("train_batch_size", 1) or 1),
dataloader_num_workers=int(da.get("dataloader_num_workers", 0) or 0),
@@ -392,6 +401,7 @@ def _build_training_config(
),
tracker=TrackerConfig(
trackers=list(tr.get("trackers", []) or []),
entity=str(tr.get("entity", "") or ""),
project_name=str(tr.get("project_name", "fastvideo") or "fastvideo"),
run_name=str(tr.get("run_name", "") or ""),
),
fastvideo/train/utils/tracking.py
+2
View File
@@ -39,6 +39,7 @@ def build_tracker(
tracker_log_dir = os.path.join(tracker_log_dir, "tracker")
tracker_config_dict = config if trackers else None
tracker_entity = tracker_config.entity or None
tracker_run_name = tracker_config.run_name or None
project = (tracker_config.project_name or "fastvideo")
@@ -47,5 +48,6 @@ def build_tracker(
experiment_name=project,
config=tracker_config_dict,
log_dir=tracker_log_dir,
entity=tracker_entity,
run_name=tracker_run_name,
)
fastvideo/train/utils/training_config.py
+2 -1
View File
@@ -22,7 +22,7 @@ class DistributedConfig:
@dataclass(slots=True)
class DataConfig:
data_path: str = ""
data_path: str | list[str] | dict[str, int] = ""
preprocessed_data_type: str = "t2v"
train_batch_size: int = 1
dataloader_num_workers: int = 0
@@ -63,6 +63,7 @@ class CheckpointConfig:
@dataclass(slots=True)
class TrackerConfig:
trackers: list[str] = field(default_factory=list)
entity: str = ""
project_name: str = "fastvideo"
run_name: str = ""
fastvideo/training/trackers.py
+4
View File
@@ -305,6 +305,7 @@ class WandbTracker(BaseTracker):
log_dir: str,
*,
config: dict[str, Any] | None = None,
entity: str | None = None,
run_name: str | None = None,
) -> None:
super().__init__()
@@ -315,6 +316,7 @@ class WandbTracker(BaseTracker):
self._wandb = wandb
self._run = wandb.init(
entity=entity,
project=experiment_name,
dir=log_dir,
config=(_sanitize_wandb_config(config) if config is not None else None),
@@ -515,6 +517,7 @@ def initialize_trackers(
experiment_name: str,
config: dict[str, Any] | None,
log_dir: str,
entity: str | None = None,
run_name: str | None = None,
) -> BaseTracker:
"""Create tracker instances based on ``trackers`` configuration."""
@@ -538,6 +541,7 @@ def initialize_trackers(
experiment_name,
os.path.abspath(log_dir),
config=config,
entity=entity,
run_name=run_name,
))
elif tracker_name == Trackers.SWANLAB.value: