From 4c08ffce497bb7fe732454135d1fb7bf6733fbb2 Mon Sep 17 00:00:00 2001 From: Mingjia Huo Date: Thu, 9 Jul 2026 22:09:07 -0700 Subject: [PATCH] [feat] World model training using third person games (#1443) Co-authored-by: Claude Fable 5 --- ...worldmodel_synthetic_flow_calibration.json | 22 + docs/training/train_infra.md | 46 + .../wan/self_forcing_causal_t2v.yaml | 3 - examples/train/scenario/worldmodel/README.md | 105 +- .../{ => solaris}/dfsft_causal_i2v.yaml | 0 .../{ => solaris}/finetune_i2v.yaml | 0 .../self_forcing_causal_i2v.yaml | 2 +- .../worldmodel/zelda/dfsft_causal_i2v.yaml | 94 ++ .../worldmodel/zelda/finetune_i2v.yaml | 88 + .../zelda/self_forcing_causal_i2v.yaml | 120 ++ .../streaming_long_tuning_causal_i2v.yaml | 140 ++ .../dataset/parquet_dataset_map_style.py | 56 +- fastvideo/eval/README.md | 40 +- .../synthetic_optical_flow/_thirdperson.py | 2 +- fastvideo/models/dits/causal_wanvideo.py | 3 +- .../models/dits/matrixgame2/action_module.py | 27 +- .../models/dits/matrixgame2/causal_model.py | 101 +- fastvideo/models/dits/matrixgame2/model.py | 2 +- fastvideo/models/dits/matrixgame2/utils.py | 256 +++ fastvideo/registry.py | 6 + .../dataset/test_parquet_dataset_map_style.py | 117 ++ .../tests/train/callbacks/test_validation.py | 348 +++- .../wan_causal_t2v_streaming_long_min.yaml | 109 ++ .../tests/train/methods/grad_norm_refs.json | 3 + .../methods/test_streaming_long_tuning.py | 362 ++++ fastvideo/tests/train/trainer/__init__.py | 0 .../tests/train/trainer/test_validation.py | 139 ++ fastvideo/tests/train/utils/test_config.py | 43 + fastvideo/train/callbacks/validation.py | 809 ++++++++- .../methods/distribution_matching/__init__.py | 3 + .../streaming_long_tuning.py | 1471 +++++++++++++++++ .../train/models/matrixgame2/matrixgame2.py | 171 +- fastvideo/train/models/wan/wan_causal.py | 7 + fastvideo/train/utils/config.py | 12 +- fastvideo/train/utils/tracking.py | 2 + fastvideo/train/utils/training_config.py | 3 +- fastvideo/training/trackers.py | 4 + 37 files changed, 4618 insertions(+), 98 deletions(-) create mode 100644 assets/eval/worldmodel_synthetic_flow_calibration.json rename examples/train/scenario/worldmodel/{ => solaris}/dfsft_causal_i2v.yaml (100%) rename examples/train/scenario/worldmodel/{ => solaris}/finetune_i2v.yaml (100%) rename examples/train/scenario/worldmodel/{ => solaris}/self_forcing_causal_i2v.yaml (99%) create mode 100644 examples/train/scenario/worldmodel/zelda/dfsft_causal_i2v.yaml create mode 100644 examples/train/scenario/worldmodel/zelda/finetune_i2v.yaml create mode 100644 examples/train/scenario/worldmodel/zelda/self_forcing_causal_i2v.yaml create mode 100644 examples/train/scenario/worldmodel/zelda/streaming_long_tuning_causal_i2v.yaml create mode 100644 fastvideo/tests/dataset/test_parquet_dataset_map_style.py create mode 100644 fastvideo/tests/train/fixtures/wan_causal_t2v_streaming_long_min.yaml create mode 100644 fastvideo/tests/train/methods/test_streaming_long_tuning.py create mode 100644 fastvideo/tests/train/trainer/__init__.py create mode 100644 fastvideo/tests/train/trainer/test_validation.py create mode 100644 fastvideo/train/methods/distribution_matching/streaming_long_tuning.py diff --git a/assets/eval/worldmodel_synthetic_flow_calibration.json b/assets/eval/worldmodel_synthetic_flow_calibration.json new file mode 100644 index 000000000..c5505048f --- /dev/null +++ b/assets/eval/worldmodel_synthetic_flow_calibration.json @@ -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 +} diff --git a/docs/training/train_infra.md b/docs/training/train_infra.md index 12582f807..8af5c5f3a 100644 --- a/docs/training/train_infra.md +++ b/docs/training/train_infra.md @@ -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 diff --git a/examples/train/configs/distribution_matching/wan/self_forcing_causal_t2v.yaml b/examples/train/configs/distribution_matching/wan/self_forcing_causal_t2v.yaml index 0c111e662..b4d5971cd 100644 --- a/examples/train/configs/distribution_matching/wan/self_forcing_causal_t2v.yaml +++ b/examples/train/configs/distribution_matching/wan/self_forcing_causal_t2v.yaml @@ -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 diff --git a/examples/train/scenario/worldmodel/README.md b/examples/train/scenario/worldmodel/README.md index 14ed7037e..cdb82d594 100644 --- a/examples/train/scenario/worldmodel/README.md +++ b/examples/train/scenario/worldmodel/README.md @@ -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 ``` diff --git a/examples/train/scenario/worldmodel/dfsft_causal_i2v.yaml b/examples/train/scenario/worldmodel/solaris/dfsft_causal_i2v.yaml similarity index 100% rename from examples/train/scenario/worldmodel/dfsft_causal_i2v.yaml rename to examples/train/scenario/worldmodel/solaris/dfsft_causal_i2v.yaml diff --git a/examples/train/scenario/worldmodel/finetune_i2v.yaml b/examples/train/scenario/worldmodel/solaris/finetune_i2v.yaml similarity index 100% rename from examples/train/scenario/worldmodel/finetune_i2v.yaml rename to examples/train/scenario/worldmodel/solaris/finetune_i2v.yaml diff --git a/examples/train/scenario/worldmodel/self_forcing_causal_i2v.yaml b/examples/train/scenario/worldmodel/solaris/self_forcing_causal_i2v.yaml similarity index 99% rename from examples/train/scenario/worldmodel/self_forcing_causal_i2v.yaml rename to examples/train/scenario/worldmodel/solaris/self_forcing_causal_i2v.yaml index 50757867d..9f8d8796f 100644 --- a/examples/train/scenario/worldmodel/self_forcing_causal_i2v.yaml +++ b/examples/train/scenario/worldmodel/solaris/self_forcing_causal_i2v.yaml @@ -97,4 +97,4 @@ callbacks: guidance_scale: 6.0 pipeline: - flow_shift: 5 + flow_shift: 5 \ No newline at end of file diff --git a/examples/train/scenario/worldmodel/zelda/dfsft_causal_i2v.yaml b/examples/train/scenario/worldmodel/zelda/dfsft_causal_i2v.yaml new file mode 100644 index 000000000..e3fe4d9de --- /dev/null +++ b/examples/train/scenario/worldmodel/zelda/dfsft_causal_i2v.yaml @@ -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 diff --git a/examples/train/scenario/worldmodel/zelda/finetune_i2v.yaml b/examples/train/scenario/worldmodel/zelda/finetune_i2v.yaml new file mode 100644 index 000000000..7b643ce53 --- /dev/null +++ b/examples/train/scenario/worldmodel/zelda/finetune_i2v.yaml @@ -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 diff --git a/examples/train/scenario/worldmodel/zelda/self_forcing_causal_i2v.yaml b/examples/train/scenario/worldmodel/zelda/self_forcing_causal_i2v.yaml new file mode 100644 index 000000000..1ad114eb2 --- /dev/null +++ b/examples/train/scenario/worldmodel/zelda/self_forcing_causal_i2v.yaml @@ -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 diff --git a/examples/train/scenario/worldmodel/zelda/streaming_long_tuning_causal_i2v.yaml b/examples/train/scenario/worldmodel/zelda/streaming_long_tuning_causal_i2v.yaml new file mode 100644 index 000000000..4b3424794 --- /dev/null +++ b/examples/train/scenario/worldmodel/zelda/streaming_long_tuning_causal_i2v.yaml @@ -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 diff --git a/fastvideo/dataset/parquet_dataset_map_style.py b/fastvideo/dataset/parquet_dataset_map_style.py index c14eee93e..be41ab46e 100644 --- a/fastvideo/dataset/parquet_dataset_map_style.py +++ b/fastvideo/dataset/parquet_dataset_map_style.py @@ -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 # ──────────────────────────────────────────────────────────────────────────── diff --git a/fastvideo/eval/README.md b/fastvideo/eval/README.md index 574e0f91f..1433b0ef1 100644 --- a/fastvideo/eval/README.md +++ b/fastvideo/eval/README.md @@ -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:`; +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 +`/eval/step_/inference_steps__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. diff --git a/fastvideo/eval/metrics/optical_flow/synthetic_optical_flow/_thirdperson.py b/fastvideo/eval/metrics/optical_flow/synthetic_optical_flow/_thirdperson.py index fd5402e62..4f372784a 100644 --- a/fastvideo/eval/metrics/optical_flow/synthetic_optical_flow/_thirdperson.py +++ b/fastvideo/eval/metrics/optical_flow/synthetic_optical_flow/_thirdperson.py @@ -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 = "" diff --git a/fastvideo/models/dits/causal_wanvideo.py b/fastvideo/models/dits/causal_wanvideo.py index 2adc618f8..88eced7ce 100644 --- a/fastvideo/models/dits/causal_wanvideo.py +++ b/fastvideo/models/dits/causal_wanvideo.py @@ -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, diff --git a/fastvideo/models/dits/matrixgame2/action_module.py b/fastvideo/models/dits/matrixgame2/action_module.py index 4732b7ba3..b7fc40511 100644 --- a/fastvideo/models/dits/matrixgame2/action_module.py +++ b/fastvideo/models/dits/matrixgame2/action_module.py @@ -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: diff --git a/fastvideo/models/dits/matrixgame2/causal_model.py b/fastvideo/models/dits/matrixgame2/causal_model.py index 682d1cfa7..1b98e2910 100644 --- a/fastvideo/models/dits/matrixgame2/causal_model.py +++ b/fastvideo/models/dits/matrixgame2/causal_model.py @@ -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) diff --git a/fastvideo/models/dits/matrixgame2/model.py b/fastvideo/models/dits/matrixgame2/model.py index a759c9512..2bf90b93b 100644 --- a/fastvideo/models/dits/matrixgame2/model.py +++ b/fastvideo/models/dits/matrixgame2/model.py @@ -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 = ( diff --git a/fastvideo/models/dits/matrixgame2/utils.py b/fastvideo/models/dits/matrixgame2/utils.py index 4dd937699..f51f7251e 100644 --- a/fastvideo/models/dits/matrixgame2/utils.py +++ b/fastvideo/models/dits/matrixgame2/utils.py @@ -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 = {} diff --git a/fastvideo/registry.py b/fastvideo/registry.py index 29a0b9aa1..4296cf642 100644 --- a/fastvideo/registry.py +++ b/fastvideo/registry.py @@ -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 ( diff --git a/fastvideo/tests/dataset/test_parquet_dataset_map_style.py b/fastvideo/tests/dataset/test_parquet_dataset_map_style.py new file mode 100644 index 000000000..04dd55766 --- /dev/null +++ b/fastvideo/tests/dataset/test_parquet_dataset_map_style.py @@ -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}) diff --git a/fastvideo/tests/train/callbacks/test_validation.py b/fastvideo/tests/train/callbacks/test_validation.py index 4b2af70e2..c70f03c90 100644 --- a/fastvideo/tests/train/callbacks/test_validation.py +++ b/fastvideo/tests/train/callbacks/test_validation.py @@ -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 diff --git a/fastvideo/tests/train/fixtures/wan_causal_t2v_streaming_long_min.yaml b/fastvideo/tests/train/fixtures/wan_causal_t2v_streaming_long_min.yaml new file mode 100644 index 000000000..7e59ff81d --- /dev/null +++ b/fastvideo/tests/train/fixtures/wan_causal_t2v_streaming_long_min.yaml @@ -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: {} diff --git a/fastvideo/tests/train/methods/grad_norm_refs.json b/fastvideo/tests/train/methods/grad_norm_refs.json index e4775534e..a026a5f84 100644 --- a/fastvideo/tests/train/methods/grad_norm_refs.json +++ b/fastvideo/tests/train/methods/grad_norm_refs.json @@ -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 diff --git a/fastvideo/tests/train/methods/test_streaming_long_tuning.py b/fastvideo/tests/train/methods/test_streaming_long_tuning.py new file mode 100644 index 000000000..8a9c04190 --- /dev/null +++ b/fastvideo/tests/train/methods/test_streaming_long_tuning.py @@ -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)) diff --git a/fastvideo/tests/train/trainer/__init__.py b/fastvideo/tests/train/trainer/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/fastvideo/tests/train/trainer/test_validation.py b/fastvideo/tests/train/trainer/test_validation.py new file mode 100644 index 000000000..c2e2f0d6a --- /dev/null +++ b/fastvideo/tests/train/trainer/test_validation.py @@ -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 diff --git a/fastvideo/tests/train/utils/test_config.py b/fastvideo/tests/train/utils/test_config.py index f23b5ced8..a80cf0517 100644 --- a/fastvideo/tests/train/utils/test_config.py +++ b/fastvideo/tests/train/utils/test_config.py @@ -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 = [ diff --git a/fastvideo/train/callbacks/validation.py b/fastvideo/train/callbacks/validation.py index f557aef8f..47bd28a32 100644 --- a/fastvideo/train/callbacks/validation.py +++ b/fastvideo/train/callbacks/validation.py @@ -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": ""} + 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, ) # ---------------------------------------------------------- diff --git a/fastvideo/train/methods/distribution_matching/__init__.py b/fastvideo/train/methods/distribution_matching/__init__.py index 4edb43cf7..015f8e124 100644 --- a/fastvideo/train/methods/distribution_matching/__init__.py +++ b/fastvideo/train/methods/distribution_matching/__init__.py @@ -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", ] diff --git a/fastvideo/train/methods/distribution_matching/streaming_long_tuning.py b/fastvideo/train/methods/distribution_matching/streaming_long_tuning.py new file mode 100644 index 000000000..f884a2dc0 --- /dev/null +++ b/fastvideo/train/methods/distribution_matching/streaming_long_tuning.py @@ -0,0 +1,1471 @@ +# SPDX-License-Identifier: Apache-2.0 +"""LongLive-style multi-stage self-forcing for streaming rollouts.""" + +from __future__ import annotations + +import copy +from dataclasses import dataclass +from typing import Any + +import torch +import torch.distributed as dist + +from fastvideo.logger import init_logger +from fastvideo.models.utils import pred_noise_to_pred_video +from fastvideo.train.methods.base import LogScalar +from fastvideo.train.methods.distribution_matching.self_forcing import ( + SelfForcingMethod, ) +from fastvideo.train.models.base import CausalModelBase +from fastvideo.train.utils.config import get_optional_float, get_optional_int + +logger = init_logger(__name__) + + +@dataclass(frozen=True, slots=True) +class DistillStage: + """Resolved multi-stage distillation stage.""" + + name: str + start_step: int + end_step: int | None + num_latent_t: int + streaming_training: bool + streaming_chunk_size: int | None = None + streaming_max_length: int | None = None + streaming_min_new_frame: int | None = None + streaming_fixed_overlap_latents: int | None = None + + +@dataclass(slots=True) +class _StreamingState: + stage: DistillStage + batch: Any + current_length: int = 0 + previous_latents: torch.Tensor | None = None + cache_tag: str = "streaming_long" + + +@dataclass(frozen=True, slots=True) +class _StreamingChunkInfo: + chunk_start: int + chunk_end: int + train_start: int + train_end: int + new_frames: int + overlap: int + current_length: int + max_length: int + + +def _as_bool(raw: Any, *, where: str) -> bool: + if isinstance(raw, bool): + return raw + raise ValueError(f"Expected bool at {where}, got {type(raw).__name__}") + + +def _as_int(raw: Any, *, where: str) -> int: + if isinstance(raw, bool): + raise ValueError(f"Expected int at {where}, got bool") + if isinstance(raw, int): + return int(raw) + if isinstance(raw, float) and raw.is_integer(): + return int(raw) + if isinstance(raw, str) and raw.strip(): + return int(raw) + raise ValueError(f"Expected int at {where}, got {type(raw).__name__}") + + +def parse_multi_phased_distill_schedule( + raw: Any, + *, + default_num_latent_t: int, + default_streaming_chunk_size: int | None = None, + default_streaming_max_length: int | None = None, +) -> list[DistillStage]: + """Parse a compact/list multi-stage schedule. + + Preferred YAML form: + + ``[{stage: self_forcing, end_step: 700, num_latent_t: 21}, ...]`` + + Compact string form accepts ``"700:21,3000:240"`` and treats the + first stage as non-streaming self-forcing and later stages as streaming. + """ + + if raw is None or raw == "": + max_length = default_streaming_max_length or default_num_latent_t + return [ + DistillStage( + name="streaming_long", + start_step=0, + end_step=None, + num_latent_t=int(max_length), + streaming_training=True, + streaming_chunk_size=default_streaming_chunk_size, + streaming_max_length=int(max_length), + ) + ] + + stages: list[DistillStage] = [] + previous_end = 0 + + if isinstance(raw, str): + parts = [part.strip() for part in raw.split(",") if part.strip()] + for idx, part in enumerate(parts): + fields = [field.strip() for field in part.split(":")] + end_step: int | None + if len(fields) == 2: + start_step = previous_end + end_step = _as_int( + fields[0], + where="multi_phased_distill_schedule.end_step", + ) + num_latent_t = _as_int( + fields[1], + where="multi_phased_distill_schedule.num_latent_t", + ) + elif len(fields) == 3: + start_step = _as_int( + fields[0], + where="multi_phased_distill_schedule.start_step", + ) + end_step = _as_int( + fields[1], + where="multi_phased_distill_schedule.end_step", + ) + num_latent_t = _as_int( + fields[2], + where="multi_phased_distill_schedule.num_latent_t", + ) + else: + raise ValueError("multi_phased_distill_schedule string entries must be " + "'end_step:num_latent_t' or 'start:end:num_latent_t', " + f"got {part!r}") + streaming = idx > 0 + stages.append( + DistillStage( + name="streaming_long" if streaming else "self_forcing", + start_step=start_step, + end_step=end_step, + num_latent_t=num_latent_t, + streaming_training=streaming, + streaming_chunk_size=(default_streaming_chunk_size if streaming else None), + streaming_max_length=num_latent_t if streaming else None, + )) + previous_end = end_step + elif isinstance(raw, list | tuple): + for idx, stage_raw in enumerate(raw): + if not isinstance(stage_raw, dict): + raise ValueError("multi_phased_distill_schedule list entries must be dicts") + name = str(stage_raw.get("stage", "") or stage_raw.get("name", "")).strip() + streaming_raw = stage_raw.get("streaming_training", None) + if streaming_raw is None: + streaming = name in {"streaming_long", "long", "streaming"} + else: + streaming = _as_bool( + streaming_raw, + where=f"multi_phased_distill_schedule[{idx}].streaming_training", + ) + if not name: + name = "streaming_long" if streaming else "self_forcing" + + start_step = _as_int( + stage_raw.get("start_step", previous_end), + where=f"multi_phased_distill_schedule[{idx}].start_step", + ) + end_raw = stage_raw.get("end_step", None) + end_step = (None if end_raw is None else _as_int( + end_raw, + where=f"multi_phased_distill_schedule[{idx}].end_step", + )) + num_latent_t = _as_int( + stage_raw.get( + "num_latent_t", + stage_raw.get( + "streaming_max_length", + stage_raw.get("max_length", default_num_latent_t), + ), + ), + where=f"multi_phased_distill_schedule[{idx}].num_latent_t", + ) + streaming_max = stage_raw.get("streaming_max_length", None) + chunk_size = stage_raw.get("streaming_chunk_size", None) + min_new = stage_raw.get("streaming_min_new_frame", None) + fixed_overlap = stage_raw.get("streaming_fixed_overlap_latents", None) + stages.append( + DistillStage( + name=name, + start_step=start_step, + end_step=end_step, + num_latent_t=num_latent_t, + streaming_training=streaming, + streaming_chunk_size=(None if chunk_size is None else _as_int( + chunk_size, + where=("multi_phased_distill_schedule" + f"[{idx}].streaming_chunk_size"), + )), + streaming_max_length=(None if streaming_max is None else _as_int( + streaming_max, + where=("multi_phased_distill_schedule" + f"[{idx}].streaming_max_length"), + )), + streaming_min_new_frame=(None if min_new is None else _as_int( + min_new, + where=("multi_phased_distill_schedule" + f"[{idx}].streaming_min_new_frame"), + )), + streaming_fixed_overlap_latents=(None if fixed_overlap is None else _as_int( + fixed_overlap, + where=("multi_phased_distill_schedule" + f"[{idx}].streaming_fixed_overlap_latents"), + )), + )) + if end_step is not None: + previous_end = end_step + else: + raise ValueError("multi_phased_distill_schedule must be a list, string, or empty") + + if not stages: + raise ValueError("multi_phased_distill_schedule produced no stages") + + previous_end = 0 + for stage in stages: + if stage.start_step < previous_end: + raise ValueError("Distillation stages must be ordered and non-overlapping") + if stage.end_step is not None and stage.end_step <= stage.start_step: + raise ValueError("Distillation stage end_step must be > start_step") + if stage.num_latent_t <= 0: + raise ValueError("Distillation stage num_latent_t must be positive") + if stage.streaming_training: + max_length = stage.streaming_max_length or stage.num_latent_t + chunk_size = stage.streaming_chunk_size or default_streaming_chunk_size + if max_length <= 0: + raise ValueError("streaming_max_length must be positive") + if chunk_size is None or chunk_size <= 0: + raise ValueError("streaming_chunk_size must be positive") + if stage.streaming_fixed_overlap_latents is not None: + if stage.streaming_fixed_overlap_latents < 0: + raise ValueError("streaming_fixed_overlap_latents must be non-negative") + if stage.streaming_fixed_overlap_latents >= chunk_size: + raise ValueError("streaming_fixed_overlap_latents must be smaller than " + "streaming_chunk_size") + if stage.end_step is not None: + previous_end = stage.end_step + return stages + + +def select_distill_stage( + stages: list[DistillStage], + iteration: int, +) -> DistillStage: + """Select the active stage for ``iteration``.""" + + step = int(iteration) + for stage in stages: + if stage.end_step is None: + if step >= stage.start_step: + return stage + elif stage.start_step <= step < stage.end_step: + return stage + return stages[-1] + + +class StreamingLongTuningMethod(SelfForcingMethod): + """Two-stage MatrixGame/LongLive-style self-forcing method. + + Stage 1 uses the existing full self-forcing rollout over a short latent + horizon. Stage 2 uses a persistent streaming sequence and trains on chunks + generated with the causal student cache. + """ + + def __init__( + self, + *, + cfg: Any, + role_models: dict[str, Any], + ) -> None: + pipeline_config = getattr(cfg.training, "pipeline_config", None) + vae_config = getattr(pipeline_config, "vae_config", None) + if vae_config is not None: + vae_config.load_encoder = True + vae_config.load_decoder = True + + super().__init__(cfg=cfg, role_models=role_models) + if not isinstance(self.student, CausalModelBase): + raise ValueError("StreamingLongTuningMethod requires a causal student") + + mcfg = self.method_config + default_chunk = get_optional_int( + mcfg, + "streaming_chunk_size", + where="method.streaming_chunk_size", + ) + default_max = get_optional_int( + mcfg, + "streaming_max_length", + where="method.streaming_max_length", + ) + default_num_latent_t = int(self.training_config.data.num_latent_t) + self._distill_stages = parse_multi_phased_distill_schedule( + mcfg.get("multi_phased_distill_schedule", None), + default_num_latent_t=default_num_latent_t, + default_streaming_chunk_size=default_chunk, + default_streaming_max_length=default_max, + ) + self._streaming_state: _StreamingState | None = None + self._streaming_anchor_inject_k = int( + get_optional_int( + mcfg, + "streaming_anchor_inject_k", + where="method.streaming_anchor_inject_k", + ) or 1) + if self._streaming_anchor_inject_k < 1: + raise ValueError("method.streaming_anchor_inject_k must be >= 1") + reencode_anchor_raw = mcfg.get("streaming_reencode_overlap_anchor", True) + if reencode_anchor_raw is None: + reencode_anchor_raw = True + self._streaming_reencode_overlap_anchor = _as_bool( + reencode_anchor_raw, + where="method.streaming_reencode_overlap_anchor", + ) + require_full_blocks_raw = mcfg.get("streaming_require_full_blocks", False) + if require_full_blocks_raw is None: + require_full_blocks_raw = False + self._streaming_require_full_blocks = _as_bool( + require_full_blocks_raw, + where="method.streaming_require_full_blocks", + ) + real_score_cfg_raw = mcfg.get("real_score_cfg", True) + if real_score_cfg_raw is None: + real_score_cfg_raw = True + self._real_score_cfg = _as_bool( + real_score_cfg_raw, + where="method.real_score_cfg", + ) + self._score_min_timestep = self._score_timestep_from_ratio("min_timestep_ratio") + self._score_max_timestep = self._score_timestep_from_ratio("max_timestep_ratio") + + max_stage_latents = max(s.num_latent_t for s in self._distill_stages) + if max_stage_latents > default_num_latent_t: + raise ValueError("training.data.num_latent_t must be at least the largest " + "multi-stage horizon so prepared conditioning is not trimmed " + "before streaming: " + f"num_latent_t={default_num_latent_t}, " + f"max_stage_latents={max_stage_latents}") + + def single_train_step( + self, + batch: dict[str, Any], + iteration: int, + ) -> tuple[ + dict[str, torch.Tensor], + dict[str, Any], + dict[str, LogScalar], + ]: + stage = select_distill_stage(self._distill_stages, iteration) + if not stage.streaming_training: + self._streaming_state = None + return self._single_train_step_short(batch, iteration, stage) + return self._single_train_step_streaming(batch, iteration, stage) + + def _single_train_step_short( + self, + batch: dict[str, Any], + iteration: int, + stage: DistillStage, + ) -> tuple[ + dict[str, torch.Tensor], + dict[str, Any], + dict[str, LogScalar], + ]: + batch = self._with_padded_first_frame_latent( + batch, + target_latent_frames=int(self.training_config.data.num_latent_t), + ) + training_batch = self.student.prepare_batch( + batch, + generator=self.cuda_generator, + latents_source="zeros", + ) + self._truncate_batch_latents(training_batch, stage.num_latent_t) + return self._losses_for_batch( + training_batch, + iteration, + stage=stage, + chunk_mask=None, + ) + + def _single_train_step_streaming( + self, + batch: dict[str, Any], + iteration: int, + stage: DistillStage, + ) -> tuple[ + dict[str, torch.Tensor], + dict[str, Any], + dict[str, LogScalar], + ]: + state = self._ensure_streaming_state(batch, stage) + update_student = self._should_update_student(iteration) + + generator_loss = torch.zeros( + (), + device=state.batch.latents.device, + dtype=state.batch.latents.dtype, + ) + student_ctx = None + generator_outputs: dict[str, Any] = {} + pred_x0, chunk_mask, chunk_info = self._generate_streaming_chunk( + state, + with_grad=update_student, + ) + if update_student: + student_ctx = ( + state.batch.timesteps, + state.batch.attn_metadata_vsa, + ) + generator_loss = self._dmd_loss_masked( + pred_x0, + state.batch, + chunk_mask=chunk_mask, + ) + generator_outputs["dmd_latent_vis_dict"] = { + "generator_pred_video": pred_x0.detach(), + "streaming_chunk_mask": chunk_mask.detach(), + } + self.student.backward( + generator_loss, + student_ctx, + grad_accum_rounds=self._gradient_accumulation_steps(), + ) + generator_loss = generator_loss.detach() + student_ctx = None + + fake_score_loss, critic_ctx, critic_outputs = (self._critic_flow_matching_loss_for_x0( + pred_x0.detach(), + state.batch, + chunk_mask=chunk_mask, + )) + self._log_train_chunk( + iteration=iteration, + stage=stage, + update_student=update_student, + chunk_info=chunk_info, + ) + + total_loss = generator_loss + fake_score_loss + outputs: dict[str, Any] = dict(generator_outputs) + outputs.update(critic_outputs) + outputs["_fv_backward"] = { + "update_student": False, + "student_ctx": student_ctx, + "critic_ctx": critic_ctx, + } + metrics: dict[str, LogScalar] = { + "update_student": float(update_student), + "distill_stage_index": float(self._distill_stages.index(stage)), + "streaming_current_length": float(state.current_length), + "streaming_max_length": float(self._stage_max_length(stage)), + } + return ( + { + "total_loss": total_loss, + "generator_loss": generator_loss, + "fake_score_loss": fake_score_loss, + }, + outputs, + metrics, + ) + + def _losses_for_batch( + self, + training_batch: Any, + iteration: int, + *, + stage: DistillStage, + chunk_mask: torch.Tensor | None, + ) -> tuple[ + dict[str, torch.Tensor], + dict[str, Any], + dict[str, LogScalar], + ]: + update_student = self._should_update_student(iteration) + + generator_loss = torch.zeros( + (), + device=training_batch.latents.device, + dtype=training_batch.latents.dtype, + ) + student_ctx = None + if update_student: + generator_pred_x0 = self._student_rollout( + training_batch, + with_grad=True, + ) + student_ctx = ( + training_batch.timesteps, + training_batch.attn_metadata_vsa, + ) + generator_loss = self._dmd_loss_masked( + generator_pred_x0, + training_batch, + chunk_mask=chunk_mask, + ) + self.student.backward( + generator_loss, + student_ctx, + grad_accum_rounds=self._gradient_accumulation_steps(), + ) + generator_loss = generator_loss.detach() + student_ctx = None + + with torch.no_grad(): + generator_pred_x0 = self._student_rollout( + training_batch, + with_grad=False, + ) + + fake_score_loss, critic_ctx, critic_outputs = (self._critic_flow_matching_loss_for_x0( + generator_pred_x0, + training_batch, + chunk_mask=chunk_mask, + )) + + total_loss = generator_loss + fake_score_loss + outputs: dict[str, Any] = dict(critic_outputs) + outputs["_fv_backward"] = { + "update_student": False, + "student_ctx": student_ctx, + "critic_ctx": critic_ctx, + } + metrics: dict[str, LogScalar] = { + "update_student": float(update_student), + "distill_stage_index": float(self._distill_stages.index(stage)), + "active_num_latent_t": float(stage.num_latent_t), + } + self._log_train_chunk( + iteration=iteration, + stage=stage, + update_student=update_student, + chunk_info=_StreamingChunkInfo( + chunk_start=0, + chunk_end=int(stage.num_latent_t) - 1, + train_start=0, + train_end=int(stage.num_latent_t) - 1, + new_frames=int(stage.num_latent_t), + overlap=0, + current_length=int(stage.num_latent_t), + max_length=int(stage.num_latent_t), + ), + ) + return ( + { + "total_loss": total_loss, + "generator_loss": generator_loss, + "fake_score_loss": fake_score_loss, + }, + outputs, + metrics, + ) + + def _gradient_accumulation_steps(self) -> int: + loop_config = getattr(self.training_config, "loop", None) + return max(1, int(getattr(loop_config, "gradient_accumulation_steps", 1) or 1)) + + def _student_rollout_streaming( + self, + batch: Any, + *, + with_grad: bool, + ) -> torch.Tensor: + # Short-stage rollout should respect the truncated batch length above. + return super()._student_rollout_streaming(batch, with_grad=with_grad) + + def _ensure_streaming_state( + self, + raw_batch: dict[str, Any], + stage: DistillStage, + ) -> _StreamingState: + state = self._streaming_state + if (state is None or state.stage != stage or not self._can_generate_more(state)): + if state is not None: + self.student.clear_caches(cache_tag=state.cache_tag) + raw_batch = self._with_padded_first_frame_latent( + raw_batch, + target_latent_frames=int(self.training_config.data.num_latent_t), + ) + training_batch = self.student.prepare_batch( + raw_batch, + generator=self.cuda_generator, + latents_source="zeros", + ) + self._truncate_batch_latents( + training_batch, + self._stage_max_length(stage), + ) + cache_tag = f"streaming_long_{id(training_batch)}" + self.student.clear_caches(cache_tag=cache_tag) + state = _StreamingState( + stage=stage, + batch=training_batch, + current_length=0, + previous_latents=None, + cache_tag=cache_tag, + ) + self._streaming_state = state + return state + + def _with_padded_first_frame_latent( + self, + raw_batch: dict[str, Any], + *, + target_latent_frames: int, + ) -> dict[str, Any]: + first_frame_latent = raw_batch.get("first_frame_latent") + if first_frame_latent is None: + return raw_batch + if first_frame_latent.ndim != 5: + raise ValueError("first_frame_latent must have shape [B, C, T, H, W], " + f"got {tuple(first_frame_latent.shape)}") + target_latent_frames = int(target_latent_frames) + if first_frame_latent.shape[2] >= target_latent_frames: + return raw_batch + + # Long tuning passes first-frame-only MatrixGame I2V condition as a full-length latent tensor. + pad_frames = target_latent_frames - int(first_frame_latent.shape[2]) + pad = torch.zeros( + first_frame_latent.shape[0], + first_frame_latent.shape[1], + pad_frames, + first_frame_latent.shape[3], + first_frame_latent.shape[4], + device=first_frame_latent.device, + dtype=first_frame_latent.dtype, + ) + batch = dict(raw_batch) + batch["first_frame_latent"] = torch.cat([first_frame_latent, pad], dim=2) + return batch + + def _can_generate_more(self, state: _StreamingState) -> bool: + max_length = self._stage_max_length(state.stage) + return state.current_length < max_length + + def _generate_streaming_chunk( + self, + state: _StreamingState, + *, + with_grad: bool, + ) -> tuple[torch.Tensor, torch.Tensor, _StreamingChunkInfo]: + assert isinstance(self.student, CausalModelBase) + + batch = state.batch + latents = batch.latents + if latents is None: + raise RuntimeError("TrainingBatch.latents is required") + + device = latents.device + dtype = latents.dtype + batch_size = int(latents.shape[0]) + max_length = self._stage_max_length(state.stage) + chunk_size = self._stage_chunk_size(state.stage) + current = int(state.current_length) + remaining = max_length - current + if remaining <= 0: + raise RuntimeError("Streaming sequence has no remaining frames") + + if state.previous_latents is None: + new_frames = min(chunk_size, remaining) + overlap = 0 + else: + new_frames = self._select_new_frame_count( + state.stage, + remaining=remaining, + device=device, + ) + overlap = min( + chunk_size - new_frames, + int(state.previous_latents.shape[1]), + ) + + if new_frames <= 0: + raise RuntimeError("Streaming new frame count must be positive") + if current + new_frames > max_length: + raise RuntimeError("Streaming chunk would exceed max length") + + rollout_block_size = int(self._chunk_size) + if rollout_block_size <= 0: + raise ValueError("method.chunk_size must be positive") + if self._streaming_require_full_blocks and new_frames % rollout_block_size != 0: + raise ValueError("Streaming new-frame count must be divisible by " + "method.chunk_size: " + f"new_frames={new_frames}, " + f"chunk_size={rollout_block_size}") + + denoising_steps = self._get_denoising_step_list(device) + num_steps = int(denoising_steps.numel()) + noise_full = torch.randn( + ( + batch_size, + new_frames, + int(latents.shape[2]), + int(latents.shape[3]), + int(latents.shape[4]), + ), + device=device, + dtype=dtype, + generator=self.cuda_generator, + ) + + block_lengths: list[int] = [] + remaining_block_frames = int(new_frames) + while remaining_block_frames > 0: + current_block = min(rollout_block_size, remaining_block_frames) + block_lengths.append(current_block) + remaining_block_frames -= current_block + + exit_indices = self._sample_exit_indices( + num_blocks=len(block_lengths), + num_steps=num_steps, + device=device, + ) + block0_exit_idx = int(exit_indices[0]) + denoised_from, denoised_to = self._get_self_forcing_timestep_bounds( + denoising_steps, + block0_exit_idx, + ) + + denoised_blocks: list[torch.Tensor] = [] + local_start = 0 + for block_idx, block_frames in enumerate(block_lengths): + block_start = current + local_start + noisy_block = noise_full[:, local_start:local_start + block_frames] + exit_idx = int(exit_indices[block_idx]) + pred_x0_block: torch.Tensor | None = None + + for step_idx, current_timestep in enumerate(denoising_steps): + timestep_block = current_timestep * torch.ones( + (batch_size, block_frames), + device=device, + dtype=torch.float32, + ) + exit_flag = step_idx == exit_idx + enable_grad = (bool(with_grad) and bool(self._enable_gradient_in_rollout) and torch.is_grad_enabled()) + + if not exit_flag: + with torch.no_grad(): + pred_noise = self.student.predict_noise_streaming( + noisy_block, + timestep_block, + batch, + conditional=True, + cache_tag=state.cache_tag, + store_kv=False, + cur_start_frame=block_start, + cfg_uncond=self._cfg_uncond, + attn_kind="vsa", + ) + if pred_noise is None: + raise RuntimeError("predict_noise_streaming returned None") + pred_x0_block = pred_noise_to_pred_video( + pred_noise=pred_noise.flatten(0, 1), + noise_input_latent=noisy_block.flatten(0, 1), + timestep=timestep_block, + scheduler=self._sf_scheduler, + ).unflatten(0, pred_noise.shape[:2]) + + if step_idx + 1 >= num_steps: + break + next_timestep = denoising_steps[step_idx + 1] + if self._student_sample_type == "sde": + noisy_block = self._sf_add_noise( + pred_x0_block, + torch.randn( + pred_x0_block.shape, + device=device, + dtype=pred_x0_block.dtype, + generator=self.cuda_generator, + ), + next_timestep * torch.ones( + (batch_size, block_frames), + device=device, + dtype=torch.float32, + ), + ) + else: + next_timestep_block = next_timestep * torch.ones( + (batch_size, block_frames), + device=device, + dtype=torch.float32, + ) + sigma_cur = self._timestep_to_sigma(timestep_block).view( + batch_size, + block_frames, + 1, + 1, + 1, + ) + sigma_next = self._timestep_to_sigma(next_timestep_block).view( + batch_size, + block_frames, + 1, + 1, + 1, + ) + eps = (noisy_block - (1 - sigma_cur) * pred_x0_block) / sigma_cur.clamp_min(1e-8) + noisy_block = (1 - sigma_next) * pred_x0_block + sigma_next * eps + continue + + with torch.set_grad_enabled(enable_grad): + pred_noise = self.student.predict_noise_streaming( + noisy_block, + timestep_block, + batch, + conditional=True, + cache_tag=state.cache_tag, + store_kv=False, + cur_start_frame=block_start, + cfg_uncond=self._cfg_uncond, + attn_kind="vsa", + ) + if pred_noise is None: + raise RuntimeError("predict_noise_streaming returned None") + pred_x0_block = pred_noise_to_pred_video( + pred_noise=pred_noise.flatten(0, 1), + noise_input_latent=noisy_block.flatten(0, 1), + timestep=timestep_block, + scheduler=self._sf_scheduler, + ).unflatten(0, pred_noise.shape[:2]) + break + + if pred_x0_block is None: + raise RuntimeError("Streaming rollout produced no block") + denoised_blocks.append(pred_x0_block) + + with torch.no_grad(): + context_timestep = torch.ones( + (batch_size, block_frames), + device=device, + dtype=torch.float32, + ) * float(self._context_noise) + context_latents = self._sf_add_noise( + pred_x0_block.detach(), + torch.randn( + pred_x0_block.shape, + device=device, + dtype=pred_x0_block.dtype, + generator=self.cuda_generator, + ), + context_timestep, + ) + + _ = self.student.predict_noise_streaming( + context_latents, + context_timestep, + batch, + conditional=True, + cache_tag=state.cache_tag, + store_kv=True, + cur_start_frame=block_start, + cfg_uncond=self._cfg_uncond, + attn_kind="vsa", + ) + + local_start += block_frames + + if not denoised_blocks: + raise RuntimeError("Streaming rollout produced no chunk") + pred_x0_chunk = torch.cat(denoised_blocks, dim=1) + batch.denoised_timestep_from = denoised_from + batch.denoised_timestep_to = denoised_to + + if overlap > 0 and state.previous_latents is not None: + overlap_latents = state.previous_latents[:, -overlap:].detach() + if self._streaming_reencode_overlap_anchor: + anchor_source = torch.cat( + [state.previous_latents.detach(), pred_x0_chunk], + dim=1, + ) + full_chunk = self._process_first_frame_anchor( + anchor_source, + target_latents=chunk_size, + ) + else: + full_chunk = torch.cat([overlap_latents, pred_x0_chunk], dim=1) + chunk_mask = torch.zeros( + (batch_size, chunk_size, 1, 1, 1), + device=device, + dtype=torch.bool, + ) + chunk_mask[:, overlap:] = True + else: + full_chunk = pred_x0_chunk + chunk_mask = torch.ones( + (batch_size, new_frames, 1, 1, 1), + device=device, + dtype=torch.bool, + ) + + state.current_length = current + new_frames + temporal_compression_ratio = int(self.training_config.pipeline_config.vae_config.arch_config. + temporal_compression_ratio # type: ignore[union-attr] + ) + chunk_info = _StreamingChunkInfo( + chunk_start=current - overlap, + chunk_end=current + new_frames - 1, + train_start=current, + train_end=current + new_frames - 1, + new_frames=new_frames, + overlap=overlap, + current_length=state.current_length, + max_length=max_length, + ) + score_frame_start = max(0, current - overlap) + batch.matrixgame_streaming_chunk_info = { + "frame_start": score_frame_start, + "frame_end": score_frame_start + int(full_chunk.shape[1]), + "new_frame_start": int(overlap), + "new_frames": int(new_frames), + "gradient_mask": chunk_mask, + "condition_image_latents": None, + "condition_image_embeds": None, + "generator_chunk_start": int(current), + "generator_chunk_end": int(current + new_frames), + "generator_action_frame_end": int((current + new_frames - 1) * temporal_compression_ratio + 1), + "window_label": f"window_{score_frame_start}_{score_frame_start + int(full_chunk.shape[1])}", + } + state.previous_latents = full_chunk.detach()[:, -chunk_size:] + + if not dist.is_initialized() or dist.get_rank() == 0: + logger.info( + "[StreamingLong] stage=%s start=%s new=%s overlap=%s " + "blocks=%s current_length=%s/%s", + state.stage.name, + current, + new_frames, + overlap, + block_lengths, + state.current_length, + max_length, + ) + + return full_chunk, chunk_mask, chunk_info + + def _get_self_forcing_timestep_bounds( + self, + step_list: torch.Tensor, + exit_idx: int, + ) -> tuple[int | None, int | None]: + scheduler_timesteps = self._sf_scheduler.timesteps.to( + device=step_list.device, + dtype=torch.float32, + ) + exit_idx = min(int(exit_idx), int(step_list.numel()) - 1) + denoising_timestep = step_list[exit_idx].to(dtype=scheduler_timesteps.dtype) + denoised_timestep_from = int(self.student.num_train_timesteps) - int( + torch.argmin( + (scheduler_timesteps - denoising_timestep).abs(), + dim=0, + ).item()) + if exit_idx == int(step_list.numel()) - 1: + return denoised_timestep_from, 0 + + next_denoising_timestep = step_list[exit_idx + 1].to(dtype=scheduler_timesteps.dtype) + denoised_timestep_to = int(self.student.num_train_timesteps) - int( + torch.argmin( + (scheduler_timesteps - next_denoising_timestep).abs(), + dim=0, + ).item()) + return denoised_timestep_from, denoised_timestep_to + + def _process_first_frame_anchor( + self, + frames: torch.Tensor, + *, + target_latents: int, + ) -> torch.Tensor: + if frames.shape[1] <= 1: + return frames + + process_latents = min(int(target_latents), int(frames.shape[1])) + if process_latents <= 1: + return frames[:, -process_latents:] + + vae = getattr(self.student, "vae", None) + if vae is None: + raise RuntimeError("Streaming first-frame anchor requires student.vae") + self._require_vae_encoder(vae) + + inject_k = min(int(self._streaming_anchor_inject_k), process_latents) + if inject_k == 1: + with torch.no_grad(): + latents_to_decode = frames[:, :-(process_latents - 1), ...] + decode_latents = self._denormalize_wan_latents( + latents_to_decode, + vae=vae, + ) + pixels = vae.decode(decode_latents) + last_frame_pixel = pixels[:, :, -1:, :, :].float() + image_latent = self._encode_wan_latent( + vae, + last_frame_pixel, + like=frames, + ) + + remaining = frames[:, -(process_latents - 1):, ...] + return torch.cat([image_latent, remaining], dim=1) + + out_slice = frames[:, -process_latents:, ...] + with torch.no_grad(): + first_k = out_slice[:, :inject_k, ...] + decode_latents = self._denormalize_wan_latents( + first_k, + vae=vae, + ) + pixels = vae.decode(decode_latents).float() + new_first_k = self._encode_wan_latent( + vae, + pixels, + like=frames, + ) + + remaining = out_slice[:, inject_k:, ...] + return torch.cat([new_first_k, remaining], dim=1) + + def _denormalize_wan_latents( + self, + frames: torch.Tensor, + *, + vae: Any, + ) -> torch.Tensor: + latents = frames.permute(0, 2, 1, 3, 4).float() + cfg = getattr(vae, "config", None) + if (cfg is not None and hasattr(cfg, "latents_mean") and hasattr(cfg, "latents_std")): + latents_mean = torch.tensor( + cfg.latents_mean, + device=latents.device, + dtype=latents.dtype, + ).view(1, -1, 1, 1, 1) + latents_std = torch.tensor( + cfg.latents_std, + device=latents.device, + dtype=latents.dtype, + ).view(1, -1, 1, 1, 1) + return latents * latents_std + latents_mean + + if hasattr(vae, "scaling_factor"): + scaling_factor = self._vae_tensor_attr( + vae, + "scaling_factor", + like=latents, + default=1.0, + ) + latents = latents / scaling_factor + shift_factor = self._vae_tensor_attr( + vae, + "shift_factor", + like=latents, + default=0.0, + ) + return latents + shift_factor + + def _normalize_wan_latents( + self, + latents: torch.Tensor, + *, + vae: Any, + ) -> torch.Tensor: + cfg = getattr(vae, "config", None) + if (cfg is not None and hasattr(cfg, "latents_mean") and hasattr(cfg, "latents_std")): + latents_mean = torch.tensor( + cfg.latents_mean, + device=latents.device, + dtype=latents.dtype, + ).view(1, -1, 1, 1, 1) + latents_std = torch.tensor( + cfg.latents_std, + device=latents.device, + dtype=latents.dtype, + ).view(1, -1, 1, 1, 1) + return (latents - latents_mean) / latents_std + + shift_factor = self._vae_tensor_attr( + vae, + "shift_factor", + like=latents, + default=0.0, + ) + latents = latents - shift_factor + if hasattr(vae, "scaling_factor"): + scaling_factor = self._vae_tensor_attr( + vae, + "scaling_factor", + like=latents, + default=1.0, + ) + latents = latents * scaling_factor + return latents + + def _encode_wan_latent( + self, + vae: Any, + pixels: torch.Tensor, + *, + like: torch.Tensor, + ) -> torch.Tensor: + encoded = vae.encode(pixels) + if isinstance(encoded, torch.Tensor): + raw_latent = encoded + elif hasattr(encoded, "mode"): + raw_latent = encoded.mode() + else: + latent_dist = getattr(encoded, "latent_dist", None) + if latent_dist is not None: + raw_latent = (latent_dist.mode() if hasattr(latent_dist, "mode") else latent_dist.mean) + else: + mean = getattr(encoded, "mean", None) + if not isinstance(mean, torch.Tensor): + raise RuntimeError("Unsupported VAE encode output") + raw_latent = mean + + raw_latent = raw_latent.to(device=like.device, dtype=torch.float32) + normalized = self._normalize_wan_latents( + raw_latent, + vae=vae, + ) + return normalized.permute(0, 2, 1, 3, 4).to(dtype=like.dtype) + + def _require_vae_encoder(self, vae: Any) -> None: + target = getattr(vae, "module", vae) + encoder = getattr(target, "encoder", None) + if encoder is None: + raise RuntimeError("Streaming first-frame anchor requires a VAE " + "loaded with encoder weights") + if (getattr(target, "use_feature_cache", False) and not hasattr(target, "_enc_feat_map")): + config = getattr(target, "config", None) + if config is not None: + config.load_encoder = True + clear_cache = getattr(target, "clear_cache", None) + if callable(clear_cache): + clear_cache() + if (getattr(target, "use_feature_cache", False) and not hasattr(target, "_enc_feat_map")): + raise RuntimeError("Streaming first-frame anchor requires VAE " + "encoder feature-cache buffers") + + def _vae_tensor_attr( + self, + vae: Any, + name: str, + *, + like: torch.Tensor, + default: float, + ) -> torch.Tensor: + value = getattr(vae, name, None) + if value is None and hasattr(vae, "module"): + value = getattr(vae.module, name, None) + if value is None: + value = default + if isinstance(value, torch.Tensor): + return value.to(device=like.device, dtype=like.dtype) + return torch.tensor(value, device=like.device, dtype=like.dtype) + + def _log_train_chunk( + self, + *, + iteration: int, + stage: DistillStage, + update_student: bool, + chunk_info: _StreamingChunkInfo, + ) -> None: + if dist.is_initialized() and dist.get_rank() != 0: + return + + generator_state = "on" if update_student else "off" + train_target = "generator+critic" if update_student else "critic" + logger.info( + "[StreamingLong] step=%s stage=%s chunk_latents=%s-%s " + "train_latents=%s-%s new=%s overlap=%s current_length=%s/%s " + "train=%s generator=%s critic=on", + iteration, + stage.name, + chunk_info.chunk_start, + chunk_info.chunk_end, + chunk_info.train_start, + chunk_info.train_end, + chunk_info.new_frames, + chunk_info.overlap, + chunk_info.current_length, + chunk_info.max_length, + train_target, + generator_state, + ) + + def _select_new_frame_count( + self, + stage: DistillStage, + *, + remaining: int, + device: torch.device, + ) -> int: + chunk_size = self._stage_chunk_size(stage) + fixed_overlap = stage.streaming_fixed_overlap_latents + if fixed_overlap is None: + fixed_overlap = get_optional_int( + self.method_config, + "streaming_fixed_overlap_latents", + where="method.streaming_fixed_overlap_latents", + ) + if fixed_overlap is not None: + return min(chunk_size - int(fixed_overlap), remaining) + + min_new = stage.streaming_min_new_frame + if min_new is None: + min_new = get_optional_int( + self.method_config, + "streaming_min_new_frame", + where="method.streaming_min_new_frame", + ) + if min_new is None: + min_new = chunk_size + min_new = max(1, int(min_new)) + + max_new = min(chunk_size, remaining) + if max_new <= min_new: + return max_new + + # Match LongLive's step-by-num_frame_per_block selection. With the + # official 21/18 setting this chooses 18 until the final short chunk. + block = int(self._chunk_size) + choices = list(range(min_new, max_new, block)) + if not choices: + choices = [max_new] + + if not dist.is_initialized() or dist.get_rank() == 0: + idx = torch.randint( + low=0, + high=len(choices), + size=(1, ), + device=device, + dtype=torch.long, + generator=self.cuda_generator, + ) + else: + idx = torch.empty((1, ), device=device, dtype=torch.long) + if dist.is_initialized(): + dist.broadcast(idx, src=0) + return int(choices[int(idx.item())]) + + def _stage_max_length(self, stage: DistillStage) -> int: + return int(stage.streaming_max_length or stage.num_latent_t) + + def _stage_chunk_size(self, stage: DistillStage) -> int: + chunk_size = stage.streaming_chunk_size + if chunk_size is None: + chunk_size = get_optional_int( + self.method_config, + "streaming_chunk_size", + where="method.streaming_chunk_size", + ) + if chunk_size is None: + chunk_size = self._chunk_size + if int(chunk_size) <= 0: + raise ValueError("streaming_chunk_size must be positive") + return int(chunk_size) + + def _truncate_batch_latents(self, batch: Any, num_latent_t: int) -> None: + num_latent_t = int(num_latent_t) + if batch.latents is not None: + batch.latents = batch.latents[:, :num_latent_t] + batch.raw_latent_shape = ( + int(batch.latents.shape[0]), + int(batch.latents.shape[2]), + int(batch.latents.shape[1]), + int(batch.latents.shape[3]), + int(batch.latents.shape[4]), + ) + if batch.noise_latents is not None: + batch.noise_latents = batch.noise_latents[:, :num_latent_t] + if batch.noisy_model_input is not None: + batch.noisy_model_input = batch.noisy_model_input[:, :num_latent_t] + if batch.noise is not None: + batch.noise = batch.noise[:, :num_latent_t] + if batch.timesteps is not None and batch.timesteps.ndim == 2: + batch.timesteps = batch.timesteps[:, :num_latent_t] + if batch.sigmas is not None and batch.sigmas.ndim >= 1: + if batch.sigmas.ndim == 1: + batch.sigmas = batch.sigmas[:num_latent_t] + else: + batch.sigmas = batch.sigmas[:, :num_latent_t] + self._rebuild_attention_metadata(batch) + + def _rebuild_attention_metadata(self, batch: Any) -> None: + build_metadata = getattr(self.student, "_build_attention_metadata", None) + if not callable(build_metadata): + return + batch.attn_metadata = None + batch.attn_metadata_vsa = None + build_metadata(batch) + batch.attn_metadata_vsa = copy.copy(batch.attn_metadata) + if batch.attn_metadata is not None: + batch.attn_metadata.VSA_sparsity = 0.0 + + def _critic_flow_matching_loss_for_x0( + self, + generator_pred_x0: torch.Tensor, + batch: Any, + *, + chunk_mask: torch.Tensor | None, + ) -> tuple[torch.Tensor, Any, dict[str, Any]]: + device = generator_pred_x0.device + batch_size = int(generator_pred_x0.shape[0]) + fake_score_timestep = self._sample_score_timestep( + batch_size, + device=device, + ) + + noise = torch.randn( + generator_pred_x0.shape, + device=device, + dtype=generator_pred_x0.dtype, + generator=self.cuda_generator, + ) + noisy_x0 = self._sf_add_noise(generator_pred_x0, noise, fake_score_timestep) + + pred_noise = self.critic.predict_noise( + noisy_x0, + fake_score_timestep, + batch, + conditional=True, + cfg_uncond=self._cfg_uncond, + attn_kind="dense", + ) + target = noise - generator_pred_x0 + flow_matching_loss = self._masked_mean( + (pred_noise - target)**2, + chunk_mask, + ) + + batch.fake_score_latent_vis_dict = { + "generator_pred_video": generator_pred_x0, + "fake_score_timestep": fake_score_timestep, + } + outputs = {"fake_score_latent_vis_dict": batch.fake_score_latent_vis_dict} + return ( + flow_matching_loss, + (batch.timesteps, batch.attn_metadata), + outputs, + ) + + def _dmd_loss_masked( + self, + generator_pred_x0: torch.Tensor, + batch: Any, + *, + chunk_mask: torch.Tensor | None, + ) -> torch.Tensor: + guidance_scale = get_optional_float( + self.method_config, + "real_score_guidance_scale", + where="method.real_score_guidance_scale", + ) + if guidance_scale is None: + guidance_scale = 1.0 + device = generator_pred_x0.device + batch_size = int(generator_pred_x0.shape[0]) + + with torch.no_grad(): + timestep = self._sample_score_timestep( + batch_size, + device=device, + ) + noise = torch.randn( + generator_pred_x0.shape, + device=device, + dtype=generator_pred_x0.dtype, + generator=self.cuda_generator, + ) + noisy_latents = self._sf_add_noise(generator_pred_x0, noise, timestep) + + faker_x0 = self._predict_x0_with_scheduler( + self.critic, + noisy_latents, + timestep, + batch, + conditional=True, + attn_kind="dense", + ) + real_cond_x0 = self._predict_x0_with_scheduler( + self.teacher, + noisy_latents, + timestep, + batch, + conditional=True, + attn_kind="dense", + ) + if self._real_score_cfg: + real_uncond_x0 = self._predict_x0_with_scheduler( + self.teacher, + noisy_latents, + timestep, + batch, + conditional=False, + attn_kind="dense", + ) + real_cfg_x0 = real_uncond_x0 + (real_cond_x0 - real_uncond_x0) * guidance_scale + else: + real_cfg_x0 = real_cond_x0 + denom = torch.abs(generator_pred_x0.float() - real_cfg_x0.float(), ).mean(dim=[1, 2, 3, 4], + keepdim=True).clamp_min(1e-6) + grad = (faker_x0 - real_cfg_x0) / denom + grad = torch.nan_to_num(grad) + + loss = 0.5 * (generator_pred_x0.float() - (generator_pred_x0.float() - grad.float()).detach())**2 + return self._masked_mean(loss, chunk_mask) + + def _score_timestep_from_ratio(self, key: str) -> int | None: + ratio = get_optional_float( + self.method_config, + key, + where=f"method.{key}", + ) + if ratio is None: + return None + if not 0.0 <= ratio <= 1.0: + raise ValueError(f"method.{key} must be in [0, 1], got {ratio}") + return int(round(float(ratio) * int(self.student.num_train_timesteps))) + + def _sample_score_timestep( + self, + batch_size: int, + *, + device: torch.device, + ) -> torch.Tensor: + timestep = torch.randint( + 0, + int(self.student.num_train_timesteps), + [batch_size], + device=device, + dtype=torch.long, + generator=self.cuda_generator, + ) + timestep = self.student.shift_and_clamp_timestep(timestep) + min_timestep = self._score_min_timestep + max_timestep = self._score_max_timestep + if min_timestep is not None or max_timestep is not None: + timestep = timestep.clamp( + min=0 if min_timestep is None else int(min_timestep), + max=(int(self.student.num_train_timesteps) if max_timestep is None else int(max_timestep)), + ) + return timestep + + def _masked_mean( + self, + values: torch.Tensor, + mask: torch.Tensor | None, + ) -> torch.Tensor: + if mask is None: + return values.mean() + mask_f = mask.to(device=values.device, dtype=values.dtype) + while mask_f.ndim < values.ndim: + mask_f = mask_f.unsqueeze(-1) + denom = mask_f.expand_as(values).sum().clamp_min(1.0) + return (values * mask_f).sum() / denom diff --git a/fastvideo/train/models/matrixgame2/matrixgame2.py b/fastvideo/train/models/matrixgame2/matrixgame2.py index e4df5ac34..54f127f34 100644 --- a/fastvideo/train/models/matrixgame2/matrixgame2.py +++ b/fastvideo/train/models/matrixgame2/matrixgame2.py @@ -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 diff --git a/fastvideo/train/models/wan/wan_causal.py b/fastvideo/train/models/wan/wan_causal.py index c7cfb78a7..63893668c 100644 --- a/fastvideo/train/models/wan/wan_causal.py +++ b/fastvideo/train/models/wan/wan_causal.py @@ -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, diff --git a/fastvideo/train/utils/config.py b/fastvideo/train/utils/config.py index df5fc71a5..3bebe61e0 100644 --- a/fastvideo/train/utils/config.py +++ b/fastvideo/train/utils/config.py @@ -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 ""), ), diff --git a/fastvideo/train/utils/tracking.py b/fastvideo/train/utils/tracking.py index 7ad28a2e4..91e06889b 100644 --- a/fastvideo/train/utils/tracking.py +++ b/fastvideo/train/utils/tracking.py @@ -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, ) diff --git a/fastvideo/train/utils/training_config.py b/fastvideo/train/utils/training_config.py index 0916e62e2..25a546cef 100644 --- a/fastvideo/train/utils/training_config.py +++ b/fastvideo/train/utils/training_config.py @@ -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 = "" diff --git a/fastvideo/training/trackers.py b/fastvideo/training/trackers.py index 6683befd6..aa108696a 100644 --- a/fastvideo/training/trackers.py +++ b/fastvideo/training/trackers.py @@ -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: