From ce4fa8ef2d9d9d30b8e1063df9a6e97ae099663e Mon Sep 17 00:00:00 2001 From: H1yori233 Date: Wed, 22 Jul 2026 22:40:39 -0700 Subject: [PATCH] [feat]: add bidirectional and causal WanTrack support --- docs/training/wantrack.md | 188 ++++++ .../wantrack/bidirectional_i2v.yaml | 82 +++ .../fine_tuning/wantrack/causal_i2v.yaml | 88 +++ fastvideo/configs/models/dits/__init__.py | 3 +- fastvideo/configs/models/dits/trackwan.py | 61 ++ fastvideo/dataset/dataloader/schema.py | 42 +- fastvideo/dataset/preprocessing_datasets.py | 16 +- fastvideo/models/dits/causal_wanvideo.py | 45 +- fastvideo/models/dits/trackwan/__init__.py | 17 + fastvideo/models/dits/trackwan/model.py | 212 +++++++ .../models/dits/trackwan/track_encoder.py | 186 ++++++ fastvideo/models/dits/wanvideo.py | 39 +- fastvideo/models/registry.py | 8 + fastvideo/pipelines/pipeline_batch_info.py | 5 + .../preprocess/preprocess_pipeline_base.py | 8 + .../preprocess_pipeline_i2v_track.py | 214 +++++++ .../pipelines/preprocess/v1_preprocess.py | 7 +- fastvideo/train/callbacks/track_validation.py | 305 +++++++++ fastvideo/train/models/wan/wan.py | 2 + fastvideo/train/models/wan/wan_causal.py | 1 + fastvideo/train/models/wantrack/__init__.py | 11 + fastvideo/train/models/wantrack/inference.py | 336 ++++++++++ fastvideo/train/models/wantrack/wantrack.py | 595 ++++++++++++++++++ .../train/models/wantrack/wantrack_causal.py | 11 + fastvideo/train/utils/config.py | 17 +- mkdocs.yml | 1 + .../wan_fun_control_to_trackwan.py | 412 ++++++++++++ 27 files changed, 2868 insertions(+), 44 deletions(-) create mode 100644 docs/training/wantrack.md create mode 100644 examples/train/configs/fine_tuning/wantrack/bidirectional_i2v.yaml create mode 100644 examples/train/configs/fine_tuning/wantrack/causal_i2v.yaml create mode 100644 fastvideo/configs/models/dits/trackwan.py create mode 100644 fastvideo/models/dits/trackwan/__init__.py create mode 100644 fastvideo/models/dits/trackwan/model.py create mode 100644 fastvideo/models/dits/trackwan/track_encoder.py create mode 100644 fastvideo/pipelines/preprocess/preprocess_pipeline_i2v_track.py create mode 100644 fastvideo/train/callbacks/track_validation.py create mode 100644 fastvideo/train/models/wantrack/__init__.py create mode 100644 fastvideo/train/models/wantrack/inference.py create mode 100644 fastvideo/train/models/wantrack/wantrack.py create mode 100644 fastvideo/train/models/wantrack/wantrack_causal.py create mode 100644 scripts/checkpoint_conversion/wan_fun_control_to_trackwan.py diff --git a/docs/training/wantrack.md b/docs/training/wantrack.md new file mode 100644 index 000000000..b181c0f5f --- /dev/null +++ b/docs/training/wantrack.md @@ -0,0 +1,188 @@ +# WanTrack training + +WanTrack extends Wan I2V with sparse point tracks. The same model wrapper and +conditioning path are used for preprocessing, training, validation, and +standalone inference. Both bidirectional and causal checkpoints are supported. + +## Build an initialization checkpoint + +WanTrack uses the pretrained control slot from Wan2.1-Fun Control as its track +slot. Convert a VideoX-Fun control checkpoint together with a diffusers-format +Wan2.1-Fun InP base: + +```bash +python scripts/checkpoint_conversion/wan_fun_control_to_trackwan.py \ + --inp-base models/Wan2.1-Fun-1.3B-InP-Diffusers \ + --control-ckpt models/Wan2.1-Fun-1.3B-Control/diffusion_pytorch_model.safetensors \ + --out models/wantrack-control-init +``` + +For the causal transformer, add `--causal` and use a separate output directory: + +```bash +python scripts/checkpoint_conversion/wan_fun_control_to_trackwan.py \ + --inp-base models/Wan2.1-Fun-1.3B-InP-Diffusers \ + --control-ckpt models/Wan2.1-Fun-1.3B-Control/diffusion_pytorch_model.safetensors \ + --out models/wantrack-control-causal-init \ + --causal +``` + +`--causal` selects `CausalTrackWanTransformer3DModel` in the generated +transformer config. Without it, the converter selects +`TrackWanTransformer3DModel`. + +The converted patch input has 52 channels: + +| Channels | Meaning | +|----------|---------| +| `0:16` | Noisy video latent | +| `16:20` | I2V mask | +| `20:36` | First-frame latent | +| `36:52` | Track map | + +The first three entries form the usual 36-channel Wan I2V input. A +`TrackEncoder` rasterizes sparse points on the VAE grid and appends the +16-channel track map. + +## Prepare point-track data + +Each video needs an `.npz` sidecar referenced by `points_path` in its metadata. +The archive must contain: + +- `tracks`: floating-point source-pixel coordinates with shape `[T, N, 2]`. + The final dimension is `(x, y)`. +- `visibility`: a boolean or numeric visibility mask with shape `[T, N]`. + +`T` is the source-video timeline and `N` is the number of point tracks. The +preprocessor applies the same temporal sample and center-crop/resize transform +to the sidecar as it applies to the video. Coordinates should therefore be in +the original video's pixel space, not normalized or pre-cropped. Points outside +the retained crop are marked invisible. + +For example, a merged-dataset annotation can contain: + +```json +[ + { + "path": "videos/clip_0001.mp4", + "points_path": "tracks/clip_0001.npz", + "cap": ["A cyclist follows a winding road."], + "resolution": {"width": 1920, "height": 1080}, + "fps": 24.0, + "duration": 5.0, + "num_frames": 120 + } +] +``` + +The paths are relative to the dataset root named in the merge file: + +```text +data/wantrack/raw,data/wantrack/metadata.json +``` + +All examples combined into one training batch must have a stackable point +dimension. Keeping `N` fixed across the dataset is the simplest option. + +Run the I2V-track preprocessor with matching pixel and latent lengths. Wan's +temporal compression is four, so 81 pixel frames produce 21 latent frames: + +```bash +torchrun --nproc_per_node=1 fastvideo/pipelines/preprocess/v1_preprocess.py \ + --model_path models/wantrack-control-init \ + --data_merge_path data/wantrack/data_merge.txt \ + --output_dir data/wantrack/preprocessed \ + --preprocess_task i2v_track \ + --num_frames 81 \ + --num_latent_t 21 \ + --max_height 480 \ + --max_width 832 \ + --train_fps 16 \ + --preprocess_video_batch_size 1 +``` + +The trainer reads the resulting +`data/wantrack/preprocessed/combined_parquet_dataset` directory with +`preprocessed_data_type: i2v_track`. + +## Train the bidirectional model + +The bidirectional training wrapper is +`fastvideo.train.models.wantrack.WanTrackModel`; it loads +`TrackWanTransformer3DModel`. + +```bash +torchrun --nproc_per_node=8 -m fastvideo.train.entrypoint.train \ + --config examples/train/configs/fine_tuning/wantrack/bidirectional_i2v.yaml +``` + +The example optionally subsamples points and applies temporal track masking. +Track IDs are sampled once for a training sample and then reused by every +denoising call and its conditional/unconditional branches. Re-sampling IDs per +call would change the point embeddings even when the coordinates are +unchanged. + +## Train the causal model + +The causal wrapper is +`fastvideo.train.models.wantrack.WanTrackCausalModel`; it loads +`CausalTrackWanTransformer3DModel`. The example uses +`TeacherForcingSFTMethod`, so clean history and the noisy current chunk receive +the same I2V and track conditioning: + +```bash +torchrun --nproc_per_node=8 -m fastvideo.train.entrypoint.train \ + --config examples/train/configs/fine_tuning/wantrack/causal_i2v.yaml +``` + +Track encoding is causal at the VAE temporal boundary. The implementation +left-pads the complete source track sequence at its global beginning, encodes +it, and then slices the latent track map at the current latent `start_frame`. +It must not independently pad and encode each later chunk: doing so would +reset temporal context and produce a different feature for the same point +history. + +The same stable `track_ids` are reused across all chunks, denoising steps, and +CFG branches. Causal sampling follows the RobotWM streaming contract: +`predict_noise_streaming()` owns the KV caches, each denoised block is committed +once as clean context, and caches are cleared at the sample boundary. + +## Validation and inference + +Both example configs use +`fastvideo.train.callbacks.track_validation.TrackValidationCallback`. It loads +fixed samples from the preprocessed WanTrack parquet, builds conditions through +the student's normal `prepare_batch()` path, generates videos, overlays the +active tracks, and logs them to the configured tracker. Set +`callbacks.track_validation.val_data_path` to use a dedicated validation +parquet; otherwise it samples from the training parquet. + +The callback and standalone callers share two helpers: + +```python +from fastvideo.train.models.wantrack.inference import ( + prepare_wantrack_batch, + sample_wantrack, +) + +batch = prepare_wantrack_batch( + model, + raw_batch, + seed=1000, + latents_source="zeros", +) +latents = sample_wantrack( + model, + batch, + num_inference_steps=30, + seed=1000, + text_guidance_scale=3.0, + motion_guidance_scale=1.5, +) +video = model.decode_latents(latents) +``` + +`sample_wantrack()` denoises the complete clip for `WanTrackModel`. For +`WanTrackCausalModel`, it uses the existing `CausalModelBase` streaming API and +the transformer's configured block size; it does not modify or fork the common +Wan causal denoising stage. diff --git a/examples/train/configs/fine_tuning/wantrack/bidirectional_i2v.yaml b/examples/train/configs/fine_tuning/wantrack/bidirectional_i2v.yaml new file mode 100644 index 000000000..9937fecce --- /dev/null +++ b/examples/train/configs/fine_tuning/wantrack/bidirectional_i2v.yaml @@ -0,0 +1,82 @@ +# WanTrack bidirectional I2V fine-tuning. +# +# Build models/wantrack-control-init with +# scripts/checkpoint_conversion/wan_fun_control_to_trackwan.py, then preprocess +# the dataset with --preprocess_task i2v_track. + +models: + student: + _target_: fastvideo.train.models.wantrack.WanTrackModel + init_from: models/wantrack-control-init + trainable: true + track_augmentation: + enabled: true + min_points: 1000 + max_points: 2500 + temporal_mask_probability: 0.2 + temporal_mask_chunk_size: 8 + motion_dropout_probability: 0.0 + text_dropout_probability: 0.0 + +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/wantrack/preprocessed/combined_parquet_dataset + preprocessed_data_type: i2v_track + dataloader_num_workers: 4 + train_batch_size: 1 + training_cfg_rate: 0.0 + seed: 1000 + num_latent_t: 21 + num_height: 480 + num_width: 832 + num_frames: 81 + + 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: 4000 + gradient_accumulation_steps: 1 + + checkpoint: + output_dir: outputs/wantrack_bidirectional_i2v + training_state_checkpointing_steps: 1000 + checkpoints_total_limit: 3 + resume_from_checkpoint: null + + tracker: + project_name: fastvideo + run_name: wantrack_bidirectional_i2v + + model: + enable_gradient_checkpointing_type: full + +callbacks: + grad_clip: + _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback + max_grad_norm: 1.0 + track_validation: + _target_: fastvideo.train.callbacks.track_validation.TrackValidationCallback + every_steps: 250 + num_val_samples: 2 + num_inference_steps: 30 + guidance_scale: 3.0 + motion_guidance_scale: 1.5 + validate_at_start: false + +pipeline: + flow_shift: 5 diff --git a/examples/train/configs/fine_tuning/wantrack/causal_i2v.yaml b/examples/train/configs/fine_tuning/wantrack/causal_i2v.yaml new file mode 100644 index 000000000..61c9d16f6 --- /dev/null +++ b/examples/train/configs/fine_tuning/wantrack/causal_i2v.yaml @@ -0,0 +1,88 @@ +# WanTrack causal I2V teacher-forcing fine-tuning. +# +# Build models/wantrack-control-causal-init with the checkpoint converter's +# --causal flag. Track IDs and the full encoded track map are shared by every +# latent chunk. + +models: + student: + _target_: fastvideo.train.models.wantrack.WanTrackCausalModel + init_from: models/wantrack-control-causal-init + trainable: true + track_augmentation: + enabled: true + min_points: 1000 + max_points: 2500 + temporal_mask_probability: 0.2 + temporal_mask_chunk_size: 8 + motion_dropout_probability: 0.0 + text_dropout_probability: 0.0 + +method: + _target_: fastvideo.train.methods.fine_tuning.tfsft.TeacherForcingSFTMethod + 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/wantrack/preprocessed/combined_parquet_dataset + preprocessed_data_type: i2v_track + dataloader_num_workers: 4 + train_batch_size: 1 + training_cfg_rate: 0.0 + seed: 1000 + num_latent_t: 21 + num_height: 480 + num_width: 832 + num_frames: 81 + + 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: 4000 + gradient_accumulation_steps: 1 + + checkpoint: + output_dir: outputs/wantrack_causal_i2v + training_state_checkpointing_steps: 1000 + checkpoints_total_limit: 3 + resume_from_checkpoint: null + + tracker: + project_name: fastvideo + run_name: wantrack_causal_i2v + + model: + enable_gradient_checkpointing_type: full + +callbacks: + grad_clip: + _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback + max_grad_norm: 1.0 + track_validation: + _target_: fastvideo.train.callbacks.track_validation.TrackValidationCallback + every_steps: 250 + num_val_samples: 2 + num_inference_steps: 30 + guidance_scale: 3.0 + motion_guidance_scale: 1.5 + validate_at_start: false + +pipeline: + flow_shift: 5 + dit_config: + local_attn_size: 6 + sink_size: 1 + rope_cache_policy: relativistic + causal_train_attention: triton diff --git a/fastvideo/configs/models/dits/__init__.py b/fastvideo/configs/models/dits/__init__.py index 1c2f97a1a..2014fdc5f 100644 --- a/fastvideo/configs/models/dits/__init__.py +++ b/fastvideo/configs/models/dits/__init__.py @@ -11,6 +11,7 @@ from fastvideo.configs.models.dits.longcat import LongCatVideoConfig from fastvideo.configs.models.dits.ltx2 import LTX2VideoConfig from fastvideo.configs.models.dits.magi_human import MagiHumanVideoConfig from fastvideo.configs.models.dits.stable_audio import StableAudioConfig +from fastvideo.configs.models.dits.trackwan import CausalTrackWanVideoConfig, TrackWanVideoConfig from fastvideo.configs.models.dits.wanvideo import WanVideoConfig from fastvideo.configs.models.dits.zimage import ZImageDiTConfig from fastvideo.configs.models.dits.hyworld import HYWorldConfig @@ -23,5 +24,5 @@ __all__ = [ "DreamXWorldARConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "FluxDiTConfig", "Flux2Config", "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig", "MagiHumanVideoConfig", "StableAudioConfig", "GlmImageDiTConfig", "LingBotWorld2CausalFastVideoConfig", "LingBotVideoConfig", - "ZImageDiTConfig" + "ZImageDiTConfig", "TrackWanVideoConfig", "CausalTrackWanVideoConfig" ] diff --git a/fastvideo/configs/models/dits/trackwan.py b/fastvideo/configs/models/dits/trackwan.py new file mode 100644 index 000000000..d6f786beb --- /dev/null +++ b/fastvideo/configs/models/dits/trackwan.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Configuration for track-conditioned Wan transformers.""" + +from dataclasses import dataclass, field + +from fastvideo.configs.models.dits.base import DiTArchConfig +from fastvideo.configs.models.dits.wanvideo import WanVideoArchConfig, WanVideoConfig + + +def _default_track_config() -> dict[str, int | bool]: + return { + "id_dim": 64, + "track_channels": 16, + "vae_spatial_compression": 8, + "vae_temporal_compression": 4, + "max_track_id": 100_000, + "zero_init_head": False, + "use_bias": False, + } + + +@dataclass +class TrackWanVideoArchConfig(WanVideoArchConfig): + """Wan I2V input widened with a latent-aligned point-track map. + + Channel order is fixed and checkpoint-visible: + noisy latent (16), I2V mask (4), first-frame latent (16), track map (16). + """ + + in_channels: int = 52 + out_channels: int = 16 + image_dim: int = 1280 + added_kv_proj_dim: int | None = 5120 + track_config: dict[str, int | bool] = field(default_factory=_default_track_config) + + def __post_init__(self) -> None: + super().__post_init__() + track_channels = int(self.track_config.get("track_channels", 0)) + expected_in_channels = self.num_channels_latents + 20 + track_channels + if track_channels <= 0: + raise ValueError("track_config.track_channels must be positive") + if self.in_channels != expected_in_channels: + raise ValueError("TrackWan in_channels must equal latent channels + 20 I2V " + f"channels + track channels; got {self.in_channels}, expected " + f"{expected_in_channels}") + + +@dataclass +class TrackWanVideoConfig(WanVideoConfig): + arch_config: DiTArchConfig = field(default_factory=TrackWanVideoArchConfig) + prefix: str = "Wan" + + +@dataclass +class CausalTrackWanVideoArchConfig(TrackWanVideoArchConfig): + """Explicit config type for the causal TrackWan architecture.""" + + +@dataclass +class CausalTrackWanVideoConfig(TrackWanVideoConfig): + arch_config: DiTArchConfig = field(default_factory=CausalTrackWanVideoArchConfig) diff --git a/fastvideo/dataset/dataloader/schema.py b/fastvideo/dataset/dataloader/schema.py index 9d8e5e4e4..5a711cb95 100644 --- a/fastvideo/dataset/dataloader/schema.py +++ b/fastvideo/dataset/dataloader/schema.py @@ -50,6 +50,46 @@ pyarrow_schema_i2v = pa.schema([ pa.field("fps", pa.float64()), ]) +pyarrow_schema_i2v_track = pa.schema([ + pa.field("id", pa.string()), + pa.field("vae_latent_bytes", pa.binary()), + pa.field("vae_latent_shape", pa.list_(pa.int64())), + pa.field("vae_latent_dtype", pa.string()), + pa.field("text_embedding_bytes", pa.binary()), + pa.field("text_embedding_shape", pa.list_(pa.int64())), + pa.field("text_embedding_dtype", pa.string()), + pa.field("first_frame_latent_bytes", pa.binary()), + pa.field("first_frame_latent_shape", pa.list_(pa.int64())), + pa.field("first_frame_latent_dtype", pa.string()), + pa.field("clip_feature_bytes", pa.binary()), + pa.field("clip_feature_shape", pa.list_(pa.int64())), + pa.field("clip_feature_dtype", pa.string()), + # MotionStream point tracks, normalized to [0, 1]. + pa.field("track_points_bytes", pa.binary()), + pa.field("track_points_shape", pa.list_(pa.int64())), # [T, N, 2] + pa.field("track_points_dtype", pa.string()), + pa.field("track_visibility_bytes", pa.binary()), + pa.field("track_visibility_shape", pa.list_(pa.int64())), # [T, N] + pa.field("track_visibility_dtype", pa.string()), + # Optional preprocessing signals used by object-covering track sampling. + # New preprocessing writes dense fallback values when segmentation or + # informativeness weights are unavailable. + pa.field("object_ids_bytes", pa.binary()), + pa.field("object_ids_shape", pa.list_(pa.int64())), # [N] + pa.field("object_ids_dtype", pa.string()), + pa.field("track_weights_bytes", pa.binary()), + pa.field("track_weights_shape", pa.list_(pa.int64())), # [N] + pa.field("track_weights_dtype", pa.string()), + pa.field("file_name", pa.string()), + pa.field("caption", pa.string()), + pa.field("media_type", pa.string()), + pa.field("width", pa.int64()), + pa.field("height", pa.int64()), + pa.field("num_frames", pa.int64()), + pa.field("duration_sec", pa.float64()), + pa.field("fps", pa.float64()), +]) + pyarrow_schema_t2v = pa.schema([ pa.field("id", pa.string()), @@ -189,4 +229,4 @@ pyarrow_schema_matrixgame2_ode_trajectory = pa.schema([ pa.field("file_name", pa.string()), pa.field("caption", pa.string()), pa.field("media_type", pa.string()), # 'image' or 'video' -]) \ No newline at end of file +]) diff --git a/fastvideo/dataset/preprocessing_datasets.py b/fastvideo/dataset/preprocessing_datasets.py index 9d1ec44e7..ba9c6f446 100644 --- a/fastvideo/dataset/preprocessing_datasets.py +++ b/fastvideo/dataset/preprocessing_datasets.py @@ -41,6 +41,7 @@ class PreprocessBatch: sample_frame_index: list[int] | None = None sample_num_frames: int | None = None action_path: str | None = None + points_path: str | None = None # Processed data pixel_values: torch.Tensor | None = None @@ -471,6 +472,8 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, item["path"] = opj(folder, item["path"]) if "action_path" in item and item["action_path"]: item["action_path"] = opj(folder, item["action_path"]) + if "points_path" in item and item["points_path"]: + item["points_path"] = opj(folder, item["points_path"]) return data_items @@ -492,7 +495,8 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, resolution=item.get("resolution"), fps=item.get("fps"), duration=item.get("duration"), - action_path=item.get("action_path")) + action_path=item.get("action_path"), + points_path=item.get("points_path")) # Apply filtering stages if not self._apply_filter_stages(batch, filter_counts): @@ -570,6 +574,16 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, # Add action_path if batch.action_path: result["action_path"] = batch.action_path + if batch.points_path: + result["points_path"] = batch.points_path + if batch.sample_frame_index is not None: + result["sample_frame_index"] = torch.tensor(batch.sample_frame_index, + dtype=torch.long) + if batch.resolution is not None: + if "width" in batch.resolution: + result["source_width"] = int(batch.resolution["width"]) + if "height" in batch.resolution: + result["source_height"] = int(batch.resolution["height"]) return result diff --git a/fastvideo/models/dits/causal_wanvideo.py b/fastvideo/models/dits/causal_wanvideo.py index 62310ed17..d2129047b 100644 --- a/fastvideo/models/dits/causal_wanvideo.py +++ b/fastvideo/models/dits/causal_wanvideo.py @@ -40,7 +40,7 @@ from fastvideo.models.dits._causal_train_attention import ( ) from fastvideo.models.dits._relative_rope import relativistic_window_offsets from fastvideo.models.dits.base import BaseDiT -from fastvideo.models.dits.wanvideo import WanT2VCrossAttention, WanTimeTextImageEmbedding +from fastvideo.models.dits.wanvideo import WanI2VCrossAttention, WanT2VCrossAttention, WanTimeTextImageEmbedding from fastvideo.platforms import AttentionBackendEnum, current_platform logger = init_logger(__name__) @@ -311,12 +311,25 @@ class CausalWanTransformerBlock(nn.Module): elementwise_affine=True, dtype=torch.float32) - # 2. Cross-attention - # Only T2V for now - self.attn2 = WanT2VCrossAttention(dim, - num_heads, - qk_norm=qk_norm, - eps=eps) + # 2. Cross-attention. I2V checkpoints prepend CLIP image tokens and + # carry a second pair of projections, which must be preserved during + # causal rollouts. + if added_kv_proj_dim is not None: + self.attn2 = WanI2VCrossAttention( + dim, + num_heads, + qk_norm=qk_norm, + eps=eps, + prefix=f"{prefix}.attn2", + ) + else: + self.attn2 = WanT2VCrossAttention( + dim, + num_heads, + qk_norm=qk_norm, + eps=eps, + prefix=f"{prefix}.attn2", + ) self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift( dim, norm_type="layer", @@ -757,11 +770,10 @@ class CausalWanTransformer3DModel(BaseDiT): orig_dtype = hidden_states.dtype if not isinstance(encoder_hidden_states, torch.Tensor): encoder_hidden_states = encoder_hidden_states[0] - if isinstance(encoder_hidden_states_image, - list) and len(encoder_hidden_states_image) > 0: - encoder_hidden_states_image = encoder_hidden_states_image[0] - else: - encoder_hidden_states_image = None + if isinstance(encoder_hidden_states_image, list): + encoder_hidden_states_image = ( + encoder_hidden_states_image[0] + if encoder_hidden_states_image else None) batch_size, num_channels, num_frames, height, width = hidden_states.shape p_t, p_h, p_w = self.patch_size @@ -876,11 +888,10 @@ class CausalWanTransformer3DModel(BaseDiT): teacher_forcing = clean_x is not None if not isinstance(encoder_hidden_states, torch.Tensor): encoder_hidden_states = encoder_hidden_states[0] - if isinstance(encoder_hidden_states_image, - list) and len(encoder_hidden_states_image) > 0: - encoder_hidden_states_image = encoder_hidden_states_image[0] - else: - encoder_hidden_states_image = None + if isinstance(encoder_hidden_states_image, list): + encoder_hidden_states_image = ( + encoder_hidden_states_image[0] + if encoder_hidden_states_image else None) batch_size, num_channels, num_frames, height, width = hidden_states.shape p_t, p_h, p_w = self.patch_size diff --git a/fastvideo/models/dits/trackwan/__init__.py b/fastvideo/models/dits/trackwan/__init__.py new file mode 100644 index 000000000..3b8211441 --- /dev/null +++ b/fastvideo/models/dits/trackwan/__init__.py @@ -0,0 +1,17 @@ +from fastvideo.models.dits.trackwan.model import ( + CausalTrackWanTransformer3DModel, + TrackWanTransformer3DModel, +) +from fastvideo.models.dits.trackwan.track_encoder import TrackEncoder + +__all__ = [ + "TrackEncoder", + "TrackWanTransformer3DModel", + "CausalTrackWanTransformer3DModel", +] + +# Entry points for model registry discovery. +EntryClass = [ + TrackWanTransformer3DModel, + CausalTrackWanTransformer3DModel, +] diff --git a/fastvideo/models/dits/trackwan/model.py b/fastvideo/models/dits/trackwan/model.py new file mode 100644 index 000000000..624b1cd12 --- /dev/null +++ b/fastvideo/models/dits/trackwan/model.py @@ -0,0 +1,212 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Bidirectional and causal track-conditioned Wan transformers.""" + +from __future__ import annotations + +from typing import Any + +import torch + +from fastvideo.configs.models.dits.trackwan import TrackWanVideoConfig +from fastvideo.models.dits.causal_wanvideo import CausalWanTransformer3DModel +from fastvideo.models.dits.trackwan.track_encoder import TrackEncoder +from fastvideo.models.dits.wanvideo import WanTransformer3DModel + + +class _TrackConditioningMixin: + supports_track_input = True + track_channels: int + track_encoder: TrackEncoder + + def _init_track_conditioning( + self, + config: TrackWanVideoConfig, + hf_config: dict[str, Any], + ) -> None: + track_config = dict(getattr(config, "track_config", None) or {}) + # Checkpoint config is authoritative for checkpoint-visible TrackEncoder + # shapes and legacy bias tensors. Defaults only fill missing fields. + if isinstance(hf_config, dict): + track_config.update(dict(hf_config.get("track_config", {}) or {})) + + self.track_channels = int( + track_config.get("track_channels", self.in_channels - 36)) + base_channels = self.in_channels - self.track_channels + expected_base_channels = self.num_channels_latents + 20 + if self.track_channels <= 0: + raise ValueError("TrackWan requires positive track_channels") + if base_channels != expected_base_channels: + raise ValueError( + "TrackWan expects noisy latent + 20 I2V channels before " + f"track conditioning; got {base_channels}, expected " + f"{expected_base_channels}") + + self.track_encoder = TrackEncoder( + id_dim=int(track_config.get("id_dim", 64)), + track_channels=self.track_channels, + vae_spatial_compression=int( + track_config.get("vae_spatial_compression", 8)), + vae_temporal_compression=int( + track_config.get("vae_temporal_compression", 4)), + max_track_id=int(track_config.get("max_track_id", 100_000)), + zero_init=bool(track_config.get("zero_init_head", False)), + use_bias=bool(track_config.get("use_bias", False)), + ) + + def _append_track_conditioning( + self, + hidden_states: torch.Tensor, + *, + track_points: torch.Tensor | None, + track_visibility: torch.Tensor | None, + track_ids: torch.Tensor | None, + start_frame: int, + ) -> torch.Tensor: + expected_channels = self.in_channels - self.track_channels + if hidden_states.ndim != 5: + raise ValueError( + "TrackWan hidden_states must be [B, C, T, H, W], got " + f"{tuple(hidden_states.shape)}") + if hidden_states.shape[1] != expected_channels: + raise ValueError( + "TrackWan received an unexpected pre-track channel count: " + f"{hidden_states.shape[1]} (expected {expected_channels})") + if (track_points is None) != (track_visibility is None): + raise ValueError( + "track_points and track_visibility must be provided together") + + batch_size, _, latent_t, latent_h, latent_w = hidden_states.shape + if track_points is None: + track_map = hidden_states.new_zeros( + batch_size, + self.track_channels, + latent_t, + latent_h, + latent_w, + ) + else: + if track_points.shape[0] != batch_size: + raise ValueError( + "track batch size must match hidden_states batch size") + temporal_ratio = self.track_encoder.vae_temporal_compression + full_latent_t = ((track_points.shape[1] - 1) // + temporal_ratio + 1) + end_frame = int(start_frame) + latent_t + if start_frame < 0 or end_frame > full_latent_t: + raise ValueError( + "track sequence does not cover the requested latent " + f"window [{start_frame}, {end_frame}); available " + f"latent frames: {full_latent_t}") + + # Encode from the global beginning so the causal temporal conv is + # left-padded exactly once, then take the requested latent chunk. + full_track_map = self.track_encoder( + track_points, + track_visibility, + full_latent_t, + latent_h, + latent_w, + track_ids=track_ids, + ) + track_map = full_track_map[:, :, start_frame:end_frame] + track_map = track_map.to(dtype=hidden_states.dtype) + + return torch.cat([hidden_states, track_map], dim=1) + + +class TrackWanTransformer3DModel( + _TrackConditioningMixin, WanTransformer3DModel): + """Bidirectional Wan I2V transformer conditioned on sparse point tracks.""" + + def __init__( + self, + config: TrackWanVideoConfig, + hf_config: dict[str, Any], + ) -> None: + super().__init__(config=config, hf_config=hf_config) + self._init_track_conditioning(config, hf_config) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + timestep: torch.LongTensor, + encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] + | None = None, + guidance: torch.Tensor | None = None, + r_timestep: torch.Tensor | None = None, + track_points: torch.Tensor | None = None, + track_visibility: torch.Tensor | None = None, + track_ids: torch.Tensor | None = None, + **kwargs: Any, + ) -> torch.Tensor: + hidden_states = self._append_track_conditioning( + hidden_states, + track_points=track_points, + track_visibility=track_visibility, + track_ids=track_ids, + start_frame=0, + ) + return super().forward( + hidden_states, + encoder_hidden_states, + timestep, + encoder_hidden_states_image=encoder_hidden_states_image, + guidance=guidance, + r_timestep=r_timestep, + **kwargs, + ) + + +class CausalTrackWanTransformer3DModel( + _TrackConditioningMixin, CausalWanTransformer3DModel): + """Causal Wan I2V transformer with globally aligned track chunks.""" + + def __init__( + self, + config: TrackWanVideoConfig, + hf_config: dict[str, Any], + ) -> None: + super().__init__(config=config, hf_config=hf_config) + self._init_track_conditioning(config, hf_config) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + timestep: torch.LongTensor, + encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] + | None = None, + start_frame: int = 0, + clean_x: torch.Tensor | None = None, + aug_t: torch.Tensor | None = None, + track_points: torch.Tensor | None = None, + track_visibility: torch.Tensor | None = None, + track_ids: torch.Tensor | None = None, + **kwargs: Any, + ) -> torch.Tensor: + hidden_states = self._append_track_conditioning( + hidden_states, + track_points=track_points, + track_visibility=track_visibility, + track_ids=track_ids, + start_frame=start_frame, + ) + if clean_x is not None: + clean_x = self._append_track_conditioning( + clean_x, + track_points=track_points, + track_visibility=track_visibility, + track_ids=track_ids, + start_frame=start_frame, + ) + return super().forward( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep, + encoder_hidden_states_image=encoder_hidden_states_image, + start_frame=start_frame, + clean_x=clean_x, + aug_t=aug_t, + **kwargs, + ) diff --git a/fastvideo/models/dits/trackwan/track_encoder.py b/fastvideo/models/dits/trackwan/track_encoder.py new file mode 100644 index 000000000..fce79816b --- /dev/null +++ b/fastvideo/models/dits/trackwan/track_encoder.py @@ -0,0 +1,186 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Sparse point tracks to a dense, latent-aligned conditioning map.""" + +from __future__ import annotations + +import math + +import torch +import torch.nn.functional as F +from torch import nn + + +def sinusoidal_embedding( + ids: torch.Tensor, + dim: int, + max_period: int = 10_000, +) -> torch.Tensor: + """Map integer identity labels to sinusoidal embeddings.""" + half = dim // 2 + frequencies = torch.exp( + -math.log(max_period) * + torch.arange(half, device=ids.device, dtype=torch.float32) / + max(half, 1)) + angles = ids.float().unsqueeze(-1) * frequencies + embedding = torch.cat([torch.cos(angles), torch.sin(angles)], dim=-1) + if dim % 2: + embedding = F.pad(embedding, (0, 1)) + return embedding + + +class TrackEncoder(nn.Module): + """Encode ``[B, T, N, 2]`` normalized tracks on the Wan latent grid.""" + + def __init__( + self, + id_dim: int, + track_channels: int, + vae_spatial_compression: int = 8, + vae_temporal_compression: int = 4, + max_track_id: int = 100_000, + zero_init: bool = False, + use_bias: bool = False, + ) -> None: + super().__init__() + if id_dim <= 0 or track_channels <= 0: + raise ValueError("id_dim and track_channels must be positive") + if vae_spatial_compression <= 0 or vae_temporal_compression <= 0: + raise ValueError("VAE compression ratios must be positive") + if max_track_id <= 0: + raise ValueError("max_track_id must be positive") + + self.id_dim = int(id_dim) + self.track_channels = int(track_channels) + self.vae_spatial_compression = int(vae_spatial_compression) + self.vae_temporal_compression = int(vae_temporal_compression) + self.max_track_id = int(max_track_id) + self.use_bias = bool(use_bias) + + # Bias-free convolutions keep empty cells exactly zero. Left-padding + # below gives the same temporal mapping as the causal Wan VAE: + # T_latent = (T_pixel - 1) // ratio + 1. + self.temporal_conv = nn.Conv3d( + self.id_dim, + self.track_channels, + kernel_size=(self.vae_temporal_compression, 1, 1), + stride=(self.vae_temporal_compression, 1, 1), + bias=self.use_bias, + ) + self.proj = nn.Conv3d( + self.track_channels, + self.track_channels, + kernel_size=1, + bias=self.use_bias, + ) + if self.temporal_conv.bias is not None: + nn.init.zeros_(self.temporal_conv.bias) + if self.proj.bias is not None: + nn.init.zeros_(self.proj.bias) + if zero_init: + nn.init.zeros_(self.proj.weight) + + def sample_ids( + self, + batch_size: int, + num_tracks: int, + device: torch.device, + generator: torch.Generator | None = None, + ) -> torch.Tensor: + """Sample IDs once at a batch/session boundary.""" + if num_tracks > self.max_track_id: + raise ValueError( + f"num_tracks ({num_tracks}) exceeds max_track_id " + f"({self.max_track_id}); unique IDs are unavailable") + return torch.stack([ + torch.randperm( + self.max_track_id, + device=device, + generator=generator, + )[:num_tracks] for _ in range(batch_size) + ]) + + def forward( + self, + coords: torch.Tensor, + visibility: torch.Tensor, + latent_t: int, + latent_h: int, + latent_w: int, + track_ids: torch.Tensor | None = None, + ) -> torch.Tensor: + if coords.ndim != 4 or coords.shape[-1] != 2: + raise ValueError( + "track coords must have shape [B, T, N, 2], got " + f"{tuple(coords.shape)}") + if visibility.shape != coords.shape[:-1]: + raise ValueError( + "track visibility must have shape [B, T, N] matching " + f"coords, got {tuple(visibility.shape)}") + if latent_t <= 0 or latent_h <= 0 or latent_w <= 0: + raise ValueError("latent dimensions must be positive") + + batch_size, _, num_tracks, _ = coords.shape + device = coords.device + if num_tracks > self.max_track_id: + raise ValueError( + f"num_tracks ({num_tracks}) exceeds max_track_id " + f"({self.max_track_id})") + if track_ids is None: + # A deterministic fallback is required for repeated denoising and + # causal chunks. Training samples random IDs once in its wrapper. + track_ids = torch.arange( + num_tracks, + device=device, + ).unsqueeze(0).expand(batch_size, num_tracks) + elif track_ids.shape != (batch_size, num_tracks): + raise ValueError( + "track_ids must have shape [B, N], got " + f"{tuple(track_ids.shape)}") + else: + track_ids = track_ids.to(device=device) + + identities = sinusoidal_embedding(track_ids, self.id_dim) + coords_float = coords.float() + x = (coords_float[..., 0].clamp(0.0, 1.0) * + (latent_w - 1)).round().long() + y = (coords_float[..., 1].clamp(0.0, 1.0) * + (latent_h - 1)).round().long() + visible = visibility.to(device=device, dtype=torch.float32) + contribution = visible.unsqueeze(-1) * identities.unsqueeze(1) + + cell_index = y * latent_w + x + dense = torch.zeros( + ( + batch_size, + coords.shape[1], + latent_h * latent_w, + self.id_dim, + ), + device=device, + dtype=torch.float32, + ) + dense.scatter_add_( + 2, + cell_index.unsqueeze(-1).expand(-1, -1, -1, self.id_dim), + contribution, + ) + dense = dense.view( + batch_size, + coords.shape[1], + latent_h, + latent_w, + self.id_dim, + ).permute(0, 4, 1, 2, 3).contiguous() + dense = dense.to(dtype=self.temporal_conv.weight.dtype) + dense = F.pad( + dense, + (0, 0, 0, 0, self.vae_temporal_compression - 1, 0), + ) + dense = self.temporal_conv(dense) + if dense.shape[2:] != (latent_t, latent_h, latent_w): + dense = F.interpolate( + dense, + size=(latent_t, latent_h, latent_w), + mode="nearest", + ) + return self.proj(dense) diff --git a/fastvideo/models/dits/wanvideo.py b/fastvideo/models/dits/wanvideo.py index 6b925c5f0..5d8db08c0 100644 --- a/fastvideo/models/dits/wanvideo.py +++ b/fastvideo/models/dits/wanvideo.py @@ -258,7 +258,7 @@ class WanI2VCrossAttention(WanSelfAttention): self.norm_added_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() self.norm_added_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() - def forward(self, x, context, context_lens): + def forward(self, x, context, context_lens, crossattn_cache=None): r""" Args: x(Tensor): Shape [B, L1, C] @@ -269,13 +269,29 @@ class WanI2VCrossAttention(WanSelfAttention): context = context[:, 257:] b, n, d = x.size(0), self.num_heads, self.head_dim - # compute query, key, value + # Query depends on the current latent tokens. Text/image keys and + # values only depend on conditioning and can be reused by causal + # streaming rollouts. q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d) - k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d) - v = self.to_v(context)[0].view(b, -1, n, d) - k_img = self.norm_added_k(self.add_k_proj(context_img)[0]).view( - b, -1, n, d) - v_img = self.add_v_proj(context_img)[0].view(b, -1, n, d) + if crossattn_cache is not None and crossattn_cache.get("is_init", False): + k = crossattn_cache["k"] + v = crossattn_cache["v"] + k_img = crossattn_cache["k_img"] + v_img = crossattn_cache["v_img"] + else: + k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d) + v = self.to_v(context)[0].view(b, -1, n, d) + k_img = self.norm_added_k(self.add_k_proj(context_img)[0]).view( + b, -1, n, d) + v_img = self.add_v_proj(context_img)[0].view(b, -1, n, d) + if crossattn_cache is not None: + crossattn_cache.update({ + "is_init": True, + "k": k, + "v": v, + "k_img": k_img, + "v_img": v_img, + }) img_x = self.attn(q, k_img, v_img) # compute attention x = self.attn(q, k, v) if k.size(1) > 0 else torch.zeros_like(q) @@ -694,11 +710,10 @@ class WanTransformer3DModel(BaseDiT): orig_dtype = hidden_states.dtype if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor): encoder_hidden_states = encoder_hidden_states[0] - if isinstance(encoder_hidden_states_image, - list) and len(encoder_hidden_states_image) > 0: - encoder_hidden_states_image = encoder_hidden_states_image[0] - else: - encoder_hidden_states_image = None + if isinstance(encoder_hidden_states_image, list): + encoder_hidden_states_image = ( + encoder_hidden_states_image[0] + if encoder_hidden_states_image else None) batch_size, num_channels, num_frames, height, width = hidden_states.shape p_t, p_h, p_w = self.patch_size diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index eb3308c18..4c62583e1 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -32,9 +32,13 @@ _TEXT_TO_VIDEO_DIT_MODELS = { "HYWorldTransformer3DModel": ("dits", "hyworld", "HYWorldTransformer3DModel"), "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"), + "TrackWanTransformer3DModel": + ("dits", "trackwan", "TrackWanTransformer3DModel"), "DreamXWorldTransformer3DModel": ("dits", "dreamx_world", "DreamXWorldTransformer3DModel"), "DreamXWorldARTransformer3DModel": ("dits", "dreamx_world_ar", "DreamXWorldARTransformer3DModel"), "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"), + "CausalTrackWanTransformer3DModel": + ("dits", "trackwan", "CausalTrackWanTransformer3DModel"), "CosmosTransformer3DModel": ("dits", "cosmos", "CosmosTransformer3DModel"), "Cosmos25Transformer3DModel": ("dits", "cosmos2_5", "Cosmos25Transformer3DModel"), "LongCatVideoTransformer3DModel": ("dits", "longcat_video_dit", "LongCatVideoTransformer3DModel"), # Wrapper (Phase 1) @@ -55,9 +59,13 @@ _TEXT_TO_VIDEO_DIT_MODELS = { _IMAGE_TO_VIDEO_DIT_MODELS = { # "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoDiT"), "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"), + "TrackWanTransformer3DModel": + ("dits", "trackwan", "TrackWanTransformer3DModel"), "DreamXWorldTransformer3DModel": ("dits", "dreamx_world", "DreamXWorldTransformer3DModel"), "DreamXWorldARTransformer3DModel": ("dits", "dreamx_world_ar", "DreamXWorldARTransformer3DModel"), "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"), + "CausalTrackWanTransformer3DModel": + ("dits", "trackwan", "CausalTrackWanTransformer3DModel"), "LingBotWorld2CausalFastTransformer3DModel": ( "dits", "lingbotworld2", diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index b8faf4983..c1ec0f823 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -301,6 +301,11 @@ class TrainingBatch: image_latents: torch.Tensor | None = None infos: list[dict[str, Any]] | None = None mask_lat_size: torch.Tensor | None = None + # WanTrack point trajectories. IDs are sampled once per batch/session + # and reused by all denoising steps and causal chunks. + track_points: torch.Tensor | None = None + track_visibility: torch.Tensor | None = None + track_ids: torch.Tensor | None = None # ODE trajectory supervision trajectory_latents: torch.Tensor | None = None diff --git a/fastvideo/pipelines/preprocess/preprocess_pipeline_base.py b/fastvideo/pipelines/preprocess/preprocess_pipeline_base.py index 9795faded..ef17bbfa6 100644 --- a/fastvideo/pipelines/preprocess/preprocess_pipeline_base.py +++ b/fastvideo/pipelines/preprocess/preprocess_pipeline_base.py @@ -286,6 +286,14 @@ class BasePreprocessPipeline(ComposedPipelineBase): if "action_path" in data: valid_data["action_path"] = [data["action_path"][i] for i in valid_indices] + if "points_path" in data: + valid_data["points_path"] = [data["points_path"][i] for i in valid_indices] + if "sample_frame_index" in data: + valid_data["sample_frame_index"] = [data["sample_frame_index"][i] for i in valid_indices] + if "source_width" in data: + valid_data["source_width"] = [data["source_width"][i] for i in valid_indices] + if "source_height" in data: + valid_data["source_height"] = [data["source_height"][i] for i in valid_indices] # VAE with torch.autocast("cuda", dtype=torch.float32): diff --git a/fastvideo/pipelines/preprocess/preprocess_pipeline_i2v_track.py b/fastvideo/pipelines/preprocess/preprocess_pipeline_i2v_track.py new file mode 100644 index 000000000..e13451bd1 --- /dev/null +++ b/fastvideo/pipelines/preprocess/preprocess_pipeline_i2v_track.py @@ -0,0 +1,214 @@ +# SPDX-License-Identifier: Apache-2.0 +"""I2V preprocessing with point-track conditioning.""" + +from typing import Any + +import numpy as np + +from fastvideo.dataset.dataloader.schema import pyarrow_schema_i2v_track +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.pipelines.preprocess.preprocess_pipeline_i2v import ( + PreprocessPipeline_I2V, ) + + +class PreprocessPipeline_I2V_Track(PreprocessPipeline_I2V): + """Store I2V features and spatially aligned MotionStream point tracks.""" + + def get_pyarrow_schema(self): + return pyarrow_schema_i2v_track + + @staticmethod + def _get_source_size( + sidecar: Any, + valid_data: dict[str, Any], + idx: int, + points_path: str, + ) -> tuple[float, float]: + if "width" in sidecar and "height" in sidecar: + width = float(np.asarray(sidecar["width"]).item()) + height = float(np.asarray(sidecar["height"]).item()) + elif "source_width" in valid_data and "source_height" in valid_data: + width = float(valid_data["source_width"][idx]) + height = float(valid_data["source_height"][idx]) + else: + raise ValueError(f"{points_path}: source width and height are required to align " + "pixel-space tracks. Store them in the sidecar or manifest resolution.") + + if width <= 0 or height <= 0: + raise ValueError(f"{points_path}: invalid source size ({width}, {height})") + return width, height + + @staticmethod + def _normalize_after_center_crop( + tracks: np.ndarray, + visibility: np.ndarray, + *, + source_width: float, + source_height: float, + target_width: int, + target_height: int, + ) -> tuple[np.ndarray, np.ndarray]: + """Apply ``CenterCropResizeVideo`` geometry and normalize coordinates.""" + if target_width <= 0 or target_height <= 0: + raise ValueError(f"Invalid target size ({target_width}, {target_height})") + target_ratio = target_height / target_width + if source_height / source_width > target_ratio: + crop_height = int(source_width * target_ratio) + crop_width = int(source_width) + else: + crop_height = int(source_height) + crop_width = int(source_height / target_ratio) + if crop_width <= 0 or crop_height <= 0: + raise ValueError(f"Center crop is empty for source size " + f"({source_width}, {source_height}) and target size " + f"({target_width}, {target_height})") + + crop_top = int(round((source_height - crop_height) / 2.0)) + crop_left = int(round((source_width - crop_width) / 2.0)) + + x = tracks[..., 0] + y = tracks[..., 1] + finite = np.isfinite(x) & np.isfinite(y) + in_crop = (finite + & (x >= crop_left) + & (x < crop_left + crop_width) + & (y >= crop_top) + & (y < crop_top + crop_height)) + + normalized = np.empty_like(tracks, dtype=np.float32) + normalized[..., 0] = (x - crop_left) / crop_width + normalized[..., 1] = (y - crop_top) / crop_height + normalized = np.nan_to_num(normalized, nan=0.0, posinf=1.0, neginf=0.0) + np.clip(normalized, 0.0, 1.0, out=normalized) + + aligned_visibility = np.asarray(visibility, dtype=np.float32) * in_crop + aligned_visibility = np.nan_to_num(aligned_visibility, nan=0.0, posinf=1.0, neginf=0.0) + np.clip(aligned_visibility, 0.0, 1.0, out=aligned_visibility) + return ( + np.ascontiguousarray(normalized, dtype=np.float32), + np.ascontiguousarray(aligned_visibility, dtype=np.float32), + ) + + def get_extra_features( + self, + valid_data: dict[str, Any], + fastvideo_args: FastVideoArgs, + ) -> dict[str, Any]: + features = super().get_extra_features(valid_data, fastvideo_args) + + points_paths = valid_data.get("points_path") + frame_indices_batch = valid_data.get("sample_frame_index") + if not points_paths or frame_indices_batch is None: + raise ValueError("i2v_track preprocessing requires points_path and sampled frame " + "indices for every video.") + non_video_paths = [path for path in valid_data["path"] if not str(path).lower().endswith(".mp4")] + if non_video_paths: + raise ValueError("i2v_track preprocessing currently supports videos only; " + f"received {non_video_paths[0]!r}") + + target_height = int(valid_data["pixel_values"].shape[-2]) + target_width = int(valid_data["pixel_values"].shape[-1]) + expected_frames = int(valid_data["pixel_values"].shape[2]) + track_points: list[np.ndarray] = [] + track_visibility: list[np.ndarray] = [] + object_ids: list[np.ndarray] = [] + track_weights: list[np.ndarray] = [] + + samples = zip(points_paths, frame_indices_batch, strict=True) + for idx, (points_path, frame_indices) in enumerate(samples): + indices = np.asarray(frame_indices, dtype=np.int64) + if indices.ndim != 1 or indices.size != expected_frames: + raise ValueError(f"{points_path}: expected {expected_frames} sampled frame " + f"indices, got shape {indices.shape}") + if indices.size == 0 or indices.min() < 0: + raise ValueError(f"{points_path}: sampled frame indices must be non-negative") + + with np.load(points_path) as sidecar: + if "tracks" not in sidecar or "visibility" not in sidecar: + raise ValueError(f"{points_path}: sidecar must contain 'tracks' and 'visibility'") + tracks = np.asarray(sidecar["tracks"], dtype=np.float32) + visibility = np.asarray(sidecar["visibility"], dtype=np.float32) + if tracks.ndim != 3 or tracks.shape[-1] != 2: + raise ValueError(f"{points_path}: tracks must have shape [T, N, 2], got " + f"{tracks.shape}") + if visibility.shape != tracks.shape[:2]: + raise ValueError(f"{points_path}: visibility shape {visibility.shape} does " + f"not match tracks {tracks.shape[:2]}") + if indices.max() >= tracks.shape[0]: + raise ValueError(f"{points_path}: sampled frame {indices.max()} exceeds " + f"{tracks.shape[0]} track frames") + + source_width, source_height = self._get_source_size(sidecar, valid_data, idx, points_path) + num_tracks = tracks.shape[1] + sample_object_ids = np.asarray( + sidecar["object_ids"] if "object_ids" in sidecar else np.full(num_tracks, -1, dtype=np.float32), + dtype=np.float32, + ) + sample_track_weights = np.asarray( + sidecar["track_weights"] if "track_weights" in sidecar else np.zeros(num_tracks, dtype=np.float32), + dtype=np.float32, + ) + if sample_object_ids.shape != (num_tracks, ): + raise ValueError(f"{points_path}: object_ids must have shape " + f"[{num_tracks}], got {sample_object_ids.shape}") + if sample_track_weights.shape != (num_tracks, ): + raise ValueError(f"{points_path}: track_weights must have shape " + f"[{num_tracks}], got {sample_track_weights.shape}") + tracks = tracks[indices] + visibility = visibility[indices] + + tracks, visibility = self._normalize_after_center_crop( + tracks, + visibility, + source_width=source_width, + source_height=source_height, + target_width=target_width, + target_height=target_height, + ) + track_points.append(tracks) + track_visibility.append(visibility) + object_ids.append(np.ascontiguousarray(sample_object_ids, dtype=np.float32)) + track_weights.append(np.ascontiguousarray(sample_track_weights, dtype=np.float32)) + + features["track_points"] = track_points + features["track_visibility"] = track_visibility + features["object_ids"] = object_ids + features["track_weights"] = track_weights + return features + + def create_record( + self, + video_name: str, + vae_latent: np.ndarray, + text_embedding: np.ndarray, + valid_data: dict[str, Any], + idx: int, + extra_features: dict[str, Any] | None = None, + ) -> dict[str, Any]: + record = super().create_record( + video_name=video_name, + vae_latent=vae_latent, + text_embedding=text_embedding, + valid_data=valid_data, + idx=idx, + extra_features=extra_features, + ) + + for name in ( + "track_points", + "track_visibility", + "object_ids", + "track_weights", + ): + if extra_features is None or name not in extra_features: + raise ValueError(f"Missing required WanTrack feature {name!r} for {video_name}") + array = np.ascontiguousarray(extra_features[name], dtype=np.float32) + record[f"{name}_bytes"] = array.tobytes() + record[f"{name}_shape"] = list(array.shape) + record[f"{name}_dtype"] = str(array.dtype) + record.pop(name, None) + + return record + + +EntryClass = PreprocessPipeline_I2V_Track diff --git a/fastvideo/pipelines/preprocess/v1_preprocess.py b/fastvideo/pipelines/preprocess/v1_preprocess.py index 2c8bc6f79..3d74bb7fd 100644 --- a/fastvideo/pipelines/preprocess/v1_preprocess.py +++ b/fastvideo/pipelines/preprocess/v1_preprocess.py @@ -8,6 +8,7 @@ from fastvideo.distributed import (get_world_size, maybe_init_distributed_enviro from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger from fastvideo.pipelines.preprocess.preprocess_pipeline_i2v import (PreprocessPipeline_I2V) +from fastvideo.pipelines.preprocess.preprocess_pipeline_i2v_track import (PreprocessPipeline_I2V_Track) from fastvideo.pipelines.preprocess.preprocess_pipeline_ode_trajectory import (PreprocessPipeline_ODE_Trajectory) from fastvideo.pipelines.preprocess.preprocess_pipeline_t2v import (PreprocessPipeline_T2V) from fastvideo.pipelines.preprocess.preprocess_pipeline_text import (PreprocessPipeline_Text) @@ -52,6 +53,8 @@ def main(args) -> None: PreprocessPipeline = PreprocessPipeline_T2V elif args.preprocess_task == "i2v": PreprocessPipeline = PreprocessPipeline_I2V + elif args.preprocess_task == "i2v_track": + PreprocessPipeline = PreprocessPipeline_I2V_Track elif args.preprocess_task == "text_only": PreprocessPipeline = PreprocessPipeline_Text elif args.preprocess_task == "ode_trajectory": @@ -72,7 +75,8 @@ def main(args) -> None: PreprocessPipeline = PreprocessPipeline_MatrixGame2_ODE_Trajectory else: raise ValueError(f"Invalid preprocess task: {args.preprocess_task}. " - f"Valid options: t2v, i2v, ode_trajectory, text_only, matrixgame2, matrixgame2_ode_trajectory") + f"Valid options: t2v, i2v, i2v_track, ode_trajectory, text_only, " + f"matrixgame2, matrixgame2_ode_trajectory") logger.info("Preprocess task: %s using %s", args.preprocess_task, PreprocessPipeline.__name__) @@ -115,6 +119,7 @@ if __name__ == "__main__": choices=[ "t2v", "i2v", + "i2v_track", "text_only", "ode_trajectory", "matrixgame2", diff --git a/fastvideo/train/callbacks/track_validation.py b/fastvideo/train/callbacks/track_validation.py new file mode 100644 index 000000000..17cfa8db3 --- /dev/null +++ b/fastvideo/train/callbacks/track_validation.py @@ -0,0 +1,305 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Track-conditioned validation shared by bidirectional and causal WanTrack.""" + +from __future__ import annotations + +import colorsys +import glob +import os +from typing import Any + +import imageio.v2 as imageio +import numpy as np +import pyarrow.parquet as pq +import torch +from PIL import Image, ImageDraw + +from fastvideo.distributed import get_world_group +from fastvideo.logger import init_logger +from fastvideo.train.callbacks.callback import Callback +from fastvideo.train.models.wantrack.inference import ( + prepare_wantrack_batch, + sample_wantrack, +) +from fastvideo.training.trackers import DummyTracker + +logger = init_logger(__name__) + + +def _track_colors(count: int) -> np.ndarray: + colors = np.empty((count, 3), dtype=np.uint8) + for index in range(count): + hue = index / max(count, 1) + red, green, blue = colorsys.hsv_to_rgb(hue, 0.9, 1.0) + colors[index] = ( + int(red * 255), + int(green * 255), + int(blue * 255), + ) + return colors + + +def _overlay_tracks( + frames: np.ndarray, + track_points: torch.Tensor, + track_visibility: torch.Tensor, + *, + stride: int, + tail: int, + radius: int, + visibility_threshold: float, +) -> list[np.ndarray]: + points = track_points[0].float().cpu().numpy() + visibility = track_visibility[0].float().cpu().numpy() + frame_count = min( + int(frames.shape[0]), + int(points.shape[0]), + int(visibility.shape[0]), + ) + frames = frames[:frame_count] + points = points[:frame_count, ::stride] + visibility = visibility[:frame_count, ::stride] + colors = _track_colors(int(points.shape[1])) + + height, width = int(frames.shape[1]), int(frames.shape[2]) + points = points.copy() + points[..., 0] *= width + points[..., 1] *= height + + output: list[np.ndarray] = [] + for frame_index, frame in enumerate(frames): + image = Image.fromarray(frame).convert("RGB") + draw = ImageDraw.Draw(image) + start = max(0, frame_index - tail) + for track_index, color_value in enumerate(colors): + color = tuple(int(value) for value in color_value) + trail: list[tuple[float, float]] = [] + for history_index in range(start, frame_index + 1): + if visibility[history_index, track_index] < visibility_threshold: + trail = [] + continue + x, y = points[history_index, track_index] + if 0 <= x < width and 0 <= y < height: + trail.append((float(x), float(y))) + if len(trail) >= 2: + draw.line(trail, fill=color, width=1) + if visibility[frame_index, track_index] >= visibility_threshold: + x, y = points[frame_index, track_index] + if 0 <= x < width and 0 <= y < height: + draw.ellipse( + ( + float(x - radius), + float(y - radius), + float(x + radius), + float(y + radius), + ), + fill=color, + ) + output.append(np.asarray(image)) + return output + + +class TrackValidationCallback(Callback): + """Generate fixed track-conditioned samples through the live student.""" + + def __init__( + self, + *, + every_steps: int = 250, + val_data_path: str | None = None, + num_val_samples: int = 1, + num_inference_steps: int = 30, + guidance_scale: float = 1.0, + motion_guidance_scale: float = 1.0, + output_dir: str | None = None, + fps: int = 24, + grid_stride: int = 3, + tail: int = 12, + radius: int = 2, + visibility_threshold: float = 0.5, + seed: int = 1000, + validate_at_start: bool = False, + ) -> None: + self.every_steps = int(every_steps) + self.val_data_path = (str(val_data_path) if val_data_path is not None else None) + self.num_val_samples = int(num_val_samples) + self.num_inference_steps = int(num_inference_steps) + self.guidance_scale = float(guidance_scale) + self.motion_guidance_scale = float(motion_guidance_scale) + self.output_dir = (str(output_dir) if output_dir is not None else None) + self.fps = int(fps) + self.grid_stride = int(grid_stride) + self.tail = int(tail) + self.radius = int(radius) + self.visibility_threshold = float(visibility_threshold) + self.seed = int(seed) + self.validate_at_start = bool(validate_at_start) + + if self.every_steps <= 0: + raise ValueError("every_steps must be positive") + if self.num_val_samples <= 0: + raise ValueError("num_val_samples must be positive") + if self.num_inference_steps <= 0: + raise ValueError("num_inference_steps must be positive") + if self.grid_stride <= 0: + raise ValueError("grid_stride must be positive") + + self.tracker: Any = DummyTracker() + self._samples: list[dict[str, Any]] = [] + self._is_main = False + self._did_start_validation = False + + def on_train_start(self, method: Any, iteration: int = 0) -> None: + del iteration + self.tracker = getattr(method, "tracker", None) or DummyTracker() + self._is_main = int(get_world_group().rank) == 0 + try: + self._samples = self._load_samples() + logger.info( + "WanTrack validation loaded %d fixed sample(s)", + len(self._samples), + ) + except Exception as error: # noqa: BLE001 + logger.warning( + "WanTrack validation setup failed; disabling callback: %s", + error, + ) + self._samples = [] + + def on_validation_begin(self, method: Any, iteration: int = 0) -> None: + if not self._samples: + return + run_at_start = (self.validate_at_start and not self._did_start_validation and iteration == 0) + run_periodic = iteration > 0 and iteration % self.every_steps == 0 + if not run_at_start and not run_periodic: + return + self._did_start_validation = True + try: + self._run(method, iteration) + except Exception as error: # noqa: BLE001 + logger.warning( + "WanTrack validation failed at step %d: %s", + iteration, + error, + ) + + def _load_samples(self) -> list[dict[str, Any]]: + from fastvideo.dataset.dataloader.schema import ( + pyarrow_schema_i2v_track, ) + from fastvideo.dataset.utils import ( + collate_rows_from_parquet_schema, ) + + data_path = self.val_data_path or str(self.training_config.data.data_path) + files = sorted(glob.glob( + os.path.join(data_path, "**", "*.parquet"), + recursive=True, + )) + if not files: + raise FileNotFoundError(f"No WanTrack validation parquet under {data_path}") + + rows: list[dict[str, Any]] = [] + for path in files: + table = pq.read_table( + path, + columns=pyarrow_schema_i2v_track.names, + ) + rows.extend(table.to_pylist()) + if len(rows) >= self.num_val_samples: + break + + text_len = int(self.training_config.pipeline_config.text_encoder_configs[0].arch_config.text_len) + collated = collate_rows_from_parquet_schema( + rows[:self.num_val_samples], + pyarrow_schema_i2v_track, + text_padding_length=text_len, + cfg_rate=0.0, + seed=self.seed, + ) + + infos = collated.get("info_list") or [{} for _ in rows] + samples: list[dict[str, Any]] = [] + batch_size = int(collated["text_embedding"].shape[0]) + for index in range(batch_size): + sample = { + name: value[index:index + 1].clone() + for name, value in collated.items() if torch.is_tensor(value) + } + sample["info_list"] = [infos[index]] + samples.append(sample) + return samples + + @torch.no_grad() + def _run(self, method: Any, iteration: int) -> None: + student = method.student + transformer = student.transformer + was_training = bool(transformer.training) + transformer.eval() + try: + artifacts: list[Any] = [] + for sample_index, raw_sample in enumerate(self._samples): + seed = self.seed + sample_index + batch = prepare_wantrack_batch( + student, + raw_sample, + seed=seed, + latents_source="zeros", + ) + sampled = sample_wantrack( + student, + batch, + num_inference_steps=self.num_inference_steps, + seed=seed, + text_guidance_scale=self.guidance_scale, + motion_guidance_scale=self.motion_guidance_scale, + ) + decoded = student.decode_latents(sampled) + + if not self._is_main: + continue + output_dir = self.output_dir or os.path.join( + self.training_config.checkpoint.output_dir, + "track_validation", + ) + os.makedirs(output_dir, exist_ok=True) + output_path = os.path.join( + output_dir, + f"step{iteration:06d}_sample{sample_index}.mp4", + ) + video = (decoded[0].clamp(0, 1).float().cpu().numpy() * 255.0).astype(np.uint8) + frames = video.transpose(1, 2, 3, 0) + conditional = batch.conditional_dict + if conditional is None: + raise RuntimeError("WanTrack validation batch has no conditioning") + frames_with_tracks = _overlay_tracks( + frames, + conditional["track_points"], + conditional["track_visibility"], + stride=self.grid_stride, + tail=self.tail, + radius=self.radius, + visibility_threshold=self.visibility_threshold, + ) + imageio.mimsave( + output_path, + frames_with_tracks, + fps=self.fps, + macro_block_size=1, + ) + info = raw_sample.get("info_list") or [{}] + caption = str(info[0].get("caption", "")) + artifact = self.tracker.video( + output_path, + caption=(f"WanTrack step {iteration}: " + f"{caption[:120]}"), + ) + if artifact is not None: + artifacts.append(artifact) + + if self._is_main and artifacts: + self.tracker.log_artifacts( + {"track_validation/generated": artifacts}, + iteration, + ) + finally: + if was_training: + transformer.train() diff --git a/fastvideo/train/models/wan/wan.py b/fastvideo/train/models/wan/wan.py index 7737639b3..26327f7b0 100644 --- a/fastvideo/train/models/wan/wan.py +++ b/fastvideo/train/models/wan/wan.py @@ -592,7 +592,9 @@ class WanModel(ModelBase): text_dict: dict[str, torch.Tensor] | None, clean_x: torch.Tensor | None = None, aug_t: torch.Tensor | None = None, + start_frame: int = 0, ) -> dict[str, Any]: + del start_frame if text_dict is None: raise ValueError("text_dict cannot be None for " "Wan distillation") diff --git a/fastvideo/train/models/wan/wan_causal.py b/fastvideo/train/models/wan/wan_causal.py index 1f8f50012..7d39ad3fd 100644 --- a/fastvideo/train/models/wan/wan_causal.py +++ b/fastvideo/train/models/wan/wan_causal.py @@ -172,6 +172,7 @@ class WanCausalModel(WanModel, CausalModelBase): noisy_latents, timestep_full, text_dict, + start_frame=cur_start_frame, )) input_kwargs["timestep"] = (timestep_full.to( device=self.device, diff --git a/fastvideo/train/models/wantrack/__init__.py b/fastvideo/train/models/wantrack/__init__.py new file mode 100644 index 000000000..77c320136 --- /dev/null +++ b/fastvideo/train/models/wantrack/__init__.py @@ -0,0 +1,11 @@ +from fastvideo.train.models.wantrack.wantrack import ( + TrackAugmentationConfig, + WanTrackModel, +) +from fastvideo.train.models.wantrack.wantrack_causal import WanTrackCausalModel + +__all__ = [ + "TrackAugmentationConfig", + "WanTrackModel", + "WanTrackCausalModel", +] diff --git a/fastvideo/train/models/wantrack/inference.py b/fastvideo/train/models/wantrack/inference.py new file mode 100644 index 000000000..54c5ebbda --- /dev/null +++ b/fastvideo/train/models/wantrack/inference.py @@ -0,0 +1,336 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Shared bidirectional and causal WanTrack sampling helpers.""" + +from __future__ import annotations + +from dataclasses import replace +from typing import Any, Literal + +import torch + +from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler, ) +from fastvideo.pipelines import TrainingBatch +from fastvideo.train.models.base import CausalModelBase, ModelBase + +_Branch = Literal["full", "no_text", "no_motion"] + + +def prepare_wantrack_batch( + model: ModelBase, + raw_batch: dict[str, Any], + *, + seed: int, + latents_source: Literal["data", "zeros"] = "zeros", +) -> TrainingBatch: + """Build inference conditions through the same path used by training.""" + augmentation = getattr(model, "track_augmentation", None) + if augmentation is None: + raise TypeError("WanTrack inference requires a WanTrack training model") + + generator = torch.Generator(device=model.device).manual_seed(int(seed)) + model.track_augmentation = replace( + augmentation, + track_dropout_probability=0.0, + temporal_mask_probability=0.0, + motion_dropout_probability=0.0, + text_dropout_probability=0.0, + ) + try: + batch = model.prepare_batch( + raw_batch, + generator=generator, + latents_source=latents_source, + ) + finally: + model.track_augmentation = augmentation + + # Streaming inference owns its cache/mask geometry. Training attention + # metadata must not leak into validation or standalone sampling. + batch.attn_metadata = None + batch.attn_metadata_vsa = None + return batch + + +def _branch_args(branch: _Branch) -> tuple[bool, dict[str, Any] | None]: + if branch == "full": + return True, None + if branch == "no_text": + return False, { + "text": "zero", + "track": "keep", + "on_missing": "ignore", + } + if branch == "no_motion": + return False, { + "text": "keep", + "track": "drop", + "on_missing": "ignore", + } + raise ValueError(f"Unknown WanTrack CFG branch: {branch!r}") + + +def _predict( + model: ModelBase, + latents: torch.Tensor, + timestep: torch.Tensor, + batch: TrainingBatch, + *, + branch: _Branch, + cache_tag: str, + start_frame: int, + store_kv: bool, +) -> torch.Tensor | None: + conditional, cfg_uncond = _branch_args(branch) + if isinstance(model, CausalModelBase): + return model.predict_noise_streaming( + latents, + timestep, + batch, + conditional=conditional, + cache_tag=cache_tag, + store_kv=store_kv, + cur_start_frame=start_frame, + cfg_uncond=cfg_uncond, + attn_kind="dense", + ) + if store_kv: + return None + return model.predict_noise( + latents, + timestep, + batch, + conditional=conditional, + cfg_uncond=cfg_uncond, + attn_kind="dense", + ) + + +def _guided_prediction( + model: ModelBase, + latents: torch.Tensor, + timestep: torch.Tensor, + batch: TrainingBatch, + *, + start_frame: int, + text_guidance_scale: float, + motion_guidance_scale: float, + motion_cfg: bool, +) -> tuple[torch.Tensor, tuple[_Branch, ...]]: + full = _predict( + model, + latents, + timestep, + batch, + branch="full", + cache_tag="wantrack_full", + start_frame=start_frame, + store_kv=False, + ) + if full is None: + raise RuntimeError("WanTrack prediction unexpectedly returned None") + + text_scale = float(text_guidance_scale) + motion_scale = float(motion_guidance_scale) + if text_scale == 1.0 and motion_scale == 1.0: + return full, ("full", ) + + no_text = _predict( + model, + latents, + timestep, + batch, + branch="no_text", + cache_tag="wantrack_no_text", + start_frame=start_frame, + store_kv=False, + ) + if no_text is None: + raise RuntimeError("WanTrack no-text prediction returned None") + + if not motion_cfg: + return no_text + text_scale * (full - no_text), ( + "full", + "no_text", + ) + + no_motion = _predict( + model, + latents, + timestep, + batch, + branch="no_motion", + cache_tag="wantrack_no_motion", + start_frame=start_frame, + store_kv=False, + ) + if no_motion is None: + raise RuntimeError("WanTrack no-motion prediction returned None") + + denominator = text_scale + motion_scale + alpha = text_scale / denominator if denominator > 0 else 0.5 + base = alpha * no_text + (1.0 - alpha) * no_motion + guided = (base + text_scale * (full - no_text) + motion_scale * (full - no_motion)) + return guided, ("full", "no_text", "no_motion") + + +def _store_causal_context( + model: CausalModelBase, + latents: torch.Tensor, + batch: TrainingBatch, + *, + start_frame: int, + branches: tuple[_Branch, ...], +) -> None: + timestep = torch.zeros( + latents.shape[:2], + device=latents.device, + dtype=torch.float32, + ) + batch.timesteps = timestep + for branch in branches: + _predict( + model, + latents, + timestep, + batch, + branch=branch, + cache_tag=f"wantrack_{branch}", + start_frame=start_frame, + store_kv=True, + ) + + +def _sample_block( + model: ModelBase, + latents: torch.Tensor, + batch: TrainingBatch, + *, + scheduler: FlowMatchEulerDiscreteScheduler, + num_inference_steps: int, + start_frame: int, + text_guidance_scale: float, + motion_guidance_scale: float, + motion_cfg: bool, +) -> tuple[torch.Tensor, tuple[_Branch, ...]]: + scheduler.set_timesteps(int(num_inference_steps), device=latents.device) + branches: tuple[_Branch, ...] = ("full", ) + for current_timestep in scheduler.timesteps: + timestep = torch.full( + latents.shape[:2], + float(current_timestep.item()), + device=latents.device, + dtype=torch.float32, + ) + batch.timesteps = timestep + prediction, branches = _guided_prediction( + model, + latents, + timestep, + batch, + start_frame=start_frame, + text_guidance_scale=text_guidance_scale, + motion_guidance_scale=motion_guidance_scale, + motion_cfg=motion_cfg, + ) + latents = scheduler.step( + prediction.float(), + current_timestep, + latents.float(), + return_dict=False, + )[0] + return latents, branches + + +@torch.no_grad() +def sample_wantrack( + model: ModelBase, + batch: TrainingBatch, + *, + num_inference_steps: int = 30, + seed: int = 0, + text_guidance_scale: float = 1.0, + motion_guidance_scale: float = 1.0, + motion_cfg: bool = True, + chunk_size: int | None = None, +) -> torch.Tensor: + """Generate normalized latents in ``[B, T, C, H, W]`` layout. + + Bidirectional models denoise the complete clip. Causal models reuse the + same streaming prediction and cache API as RobotWM. + """ + if batch.latents is None or batch.latents.ndim != 5: + raise ValueError("WanTrack inference requires [B, T, C, H, W] latents") + if num_inference_steps <= 0: + raise ValueError("num_inference_steps must be positive") + + generator = torch.Generator(device="cpu").manual_seed(int(seed)) + latents = torch.randn( + tuple(batch.latents.shape), + generator=generator, + dtype=torch.float32, + ).to(device=batch.latents.device, dtype=batch.latents.dtype) + scheduler = FlowMatchEulerDiscreteScheduler(shift=float(getattr(model, "timestep_shift", 5.0)), ) + + if not isinstance(model, CausalModelBase): + sampled, _ = _sample_block( + model, + latents, + batch, + scheduler=scheduler, + num_inference_steps=num_inference_steps, + start_frame=0, + text_guidance_scale=text_guidance_scale, + motion_guidance_scale=motion_guidance_scale, + motion_cfg=motion_cfg, + ) + return sampled + + if chunk_size is None: + transformer = getattr(model, "transformer", None) + chunk_size = int( + getattr( + transformer, + "num_frame_per_block", + getattr(transformer.config.arch_config, "num_frames_per_block", 3), + )) + chunk_size = int(chunk_size) + if chunk_size <= 0: + raise ValueError("chunk_size must be positive") + num_frames = int(latents.shape[1]) + if num_frames % chunk_size != 0: + raise ValueError("Causal WanTrack inference requires latent frames " + f"divisible by chunk_size; got {num_frames} and " + f"{chunk_size}") + + for branch in ("full", "no_text", "no_motion"): + model.clear_caches(cache_tag=f"wantrack_{branch}") + + sampled_blocks: list[torch.Tensor] = [] + try: + for start_frame in range(0, num_frames, chunk_size): + block = latents[:, start_frame:start_frame + chunk_size] + block, branches = _sample_block( + model, + block, + batch, + scheduler=scheduler, + num_inference_steps=num_inference_steps, + start_frame=start_frame, + text_guidance_scale=text_guidance_scale, + motion_guidance_scale=motion_guidance_scale, + motion_cfg=motion_cfg, + ) + sampled_blocks.append(block) + _store_causal_context( + model, + block, + batch, + start_frame=start_frame, + branches=branches, + ) + finally: + for branch in ("full", "no_text", "no_motion"): + model.clear_caches(cache_tag=f"wantrack_{branch}") + + return torch.cat(sampled_blocks, dim=1) diff --git a/fastvideo/train/models/wantrack/wantrack.py b/fastvideo/train/models/wantrack/wantrack.py new file mode 100644 index 000000000..ea2f18328 --- /dev/null +++ b/fastvideo/train/models/wantrack/wantrack.py @@ -0,0 +1,595 @@ +# SPDX-License-Identifier: Apache-2.0 +"""WanTrack training integration shared by bidirectional and causal models.""" + +from __future__ import annotations + +import copy +from collections.abc import Mapping +from dataclasses import dataclass, fields +from typing import Any, Literal + +import torch + +from fastvideo.dataset.dataloader.schema import pyarrow_schema_i2v_track +from fastvideo.distributed import get_sp_group, get_world_group +from fastvideo.pipelines import TrainingBatch +from fastvideo.train.models.wan.wan import WanModel +from fastvideo.train.utils.dataloader import build_parquet_t2v_train_dataloader +from fastvideo.train.utils.moduleloader import load_module_from_path +from fastvideo.training.training_utils import normalize_dit_input + + +@dataclass(slots=True) +class TrackAugmentationConfig: + """Configuration-driven MotionStream-style training augmentation.""" + + enabled: bool = True + min_points: int = 1000 + max_points: int = 2500 + sparse_object_sampling: bool = False + extra_points: int = 0 + extra_point_sampling: Literal["random", "weighted"] = "random" + track_dropout_probability: float = 0.0 + temporal_mask_probability: float = 0.2 + temporal_mask_chunk_size: int = 8 + motion_dropout_probability: float = 0.0 + text_dropout_probability: float = 0.0 + + def __post_init__(self) -> None: + if self.min_points < 0 or self.max_points < 0: + raise ValueError("min_points and max_points must be non-negative") + if self.max_points and self.min_points > self.max_points: + raise ValueError("min_points cannot exceed max_points") + if self.extra_points < 0: + raise ValueError("extra_points must be non-negative") + if self.extra_point_sampling not in {"random", "weighted"}: + raise ValueError("extra_point_sampling must be 'random' or " + "'weighted'") + if self.temporal_mask_chunk_size <= 0: + raise ValueError("temporal_mask_chunk_size must be positive") + for name in ( + "track_dropout_probability", + "temporal_mask_probability", + "motion_dropout_probability", + "text_dropout_probability", + ): + value = float(getattr(self, name)) + if not 0.0 <= value <= 1.0: + raise ValueError(f"{name} must be in [0, 1], got {value}") + + @classmethod + def from_value( + cls, + value: TrackAugmentationConfig | Mapping[str, Any] | None, + ) -> TrackAugmentationConfig: + if value is None: + return cls() + if isinstance(value, cls): + return value + if not isinstance(value, Mapping): + raise TypeError("track_augmentation must be a mapping or " + "TrackAugmentationConfig") + valid_fields = {item.name for item in fields(cls)} + unknown = set(value) - valid_fields + if unknown: + raise ValueError("Unknown track_augmentation fields: " + f"{sorted(unknown)}") + return cls(**dict(value)) + + +class WanTrackModel(WanModel): + """Bidirectional WanTrack model and shared causal conditioning.""" + + _transformer_cls_name = "TrackWanTransformer3DModel" + + def __init__( + self, + *, + track_augmentation: TrackAugmentationConfig + | Mapping[str, Any] + | None = None, + freeze_track_encoder: bool = False, + **kwargs: Any, + ) -> None: + self.track_augmentation = TrackAugmentationConfig.from_value(track_augmentation) + self.freeze_track_encoder = bool(freeze_track_encoder) + super().__init__(**kwargs) + if self.freeze_track_encoder: + track_encoder = getattr(self.transformer, "track_encoder", None) + if track_encoder is None: + raise ValueError("freeze_track_encoder=true requires a " + "transformer.track_encoder") + track_encoder.requires_grad_(False) + + def init_preprocessors(self, training_config: Any) -> None: + self.vae = load_module_from_path( + model_path=str(training_config.model_path), + module_type="vae", + training_config=training_config, + ) + self.world_group = get_world_group() + self.sp_group = get_sp_group() + self._init_timestep_mechanics() + + preprocessed_data_type = str(getattr( + training_config.data, + "preprocessed_data_type", + "", + )).strip().lower() + if preprocessed_data_type != "i2v_track": + raise ValueError("WanTrack requires data.preprocessed_data_type='i2v_track', " + f"got {preprocessed_data_type!r}") + + text_len = (training_config.pipeline_config.text_encoder_configs[0].arch_config.text_len) + self.dataloader = build_parquet_t2v_train_dataloader( + training_config.data, + text_len=int(text_len), + parquet_schema=pyarrow_schema_i2v_track, + ) + self.start_step = 0 + + def prepare_batch( + self, + raw_batch: dict[str, Any], + *, + generator: torch.Generator, + latents_source: Literal["data", "zeros"] = "data", + ) -> TrainingBatch: + if self._requires_negative_conditioning: + self.ensure_negative_conditioning() + assert self.training_config is not None + training_config = self.training_config + dtype = self._get_training_dtype() + device = self.device + + required = ( + "first_frame_latent", + "clip_feature", + "track_points", + "track_visibility", + "text_embedding", + "text_attention_mask", + ) + missing = [name for name in required if name not in raw_batch] + if missing: + raise ValueError(f"WanTrack batch is missing required fields: {missing}") + + latent_t = int(training_config.data.num_latent_t) + if latents_source == "data": + if "vae_latent" not in raw_batch: + raise ValueError("WanTrack requires vae_latent when " + "latents_source='data'") + latents = raw_batch["vae_latent"][:, :, :latent_t].to(device=device, dtype=dtype) + elif latents_source == "zeros": + batch_size = int(raw_batch["text_embedding"].shape[0]) + vae_config = (training_config.pipeline_config.vae_config.arch_config) + latents = torch.zeros( + batch_size, + int(vae_config.z_dim), + latent_t, + int(training_config.data.num_height) // int(vae_config.spatial_compression_ratio), + int(training_config.data.num_width) // int(vae_config.spatial_compression_ratio), + device=device, + dtype=dtype, + ) + else: + raise ValueError(f"Unknown latents_source: {latents_source!r}") + first_frame_latent = raw_batch["first_frame_latent"][:, :, :latent_t].to(device=device, dtype=dtype) + image_embeds = raw_batch["clip_feature"] + if not torch.is_tensor(image_embeds) or image_embeds.numel() == 0: + raise ValueError("WanTrack requires non-empty clip_feature") + image_embeds = image_embeds.to(device=device, dtype=dtype) + + expected_track_frames = ((latent_t - 1) * self._temporal_compression_ratio() + 1) + track_points = raw_batch["track_points"] + track_visibility = raw_batch["track_visibility"] + if track_points.ndim != 4 or track_points.shape[-1] != 2: + raise ValueError("track_points must have shape [B, T, N, 2], got " + f"{tuple(track_points.shape)}") + if track_visibility.shape != track_points.shape[:-1]: + raise ValueError("track_visibility must have shape [B, T, N], got " + f"{tuple(track_visibility.shape)}") + if track_points.shape[1] < expected_track_frames: + raise ValueError("track sequence is shorter than the requested latent clip: " + f"{track_points.shape[1]} < {expected_track_frames}") + track_points = track_points[:, :expected_track_frames].to(device=device, dtype=torch.float32) + track_visibility = track_visibility[:, :expected_track_frames].to(device=device, dtype=torch.float32) + num_tracks = int(track_points.shape[2]) + + object_ids = raw_batch.get("object_ids") + if (torch.is_tensor(object_ids) and object_ids.numel() > 0 + and object_ids.shape == (track_points.shape[0], num_tracks)): + object_ids = object_ids.to(device=device, dtype=torch.long) + else: + object_ids = None + + track_weights = raw_batch.get("track_weights") + if (torch.is_tensor(track_weights) and track_weights.numel() > 0 + and track_weights.shape == (track_points.shape[0], num_tracks)): + track_weights = track_weights.to(device=device, dtype=torch.float32) + else: + track_weights = None + + encoder_hidden_states = raw_batch["text_embedding"].to(device=device, dtype=dtype) + encoder_hidden_states, track_visibility = self._augment_conditions( + encoder_hidden_states, + track_visibility, + object_ids=object_ids, + track_weights=track_weights, + generator=generator, + ) + batch_size = int(track_points.shape[0]) + max_track_id = int(getattr( + getattr(self.transformer, "track_encoder", None), + "max_track_id", + 100_000, + )) + if num_tracks > max_track_id: + raise ValueError(f"num_tracks ({num_tracks}) exceeds max_track_id " + f"({max_track_id})") + track_ids = torch.stack([ + torch.randperm( + max_track_id, + generator=generator, + device=device, + )[:num_tracks] for _ in range(batch_size) + ]) + + training_batch = TrainingBatch( + latents=latents, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=raw_batch["text_attention_mask"].to(device=device, dtype=dtype), + image_embeds=image_embeds, + image_latents=first_frame_latent, + track_points=track_points, + track_visibility=track_visibility, + track_ids=track_ids, + infos=raw_batch.get("info_list"), + ) + training_batch.latents = normalize_dit_input( + "wan", + training_batch.latents, + self.vae, + ) + training_batch = self._prepare_dit_inputs(training_batch, generator) + training_batch = self._build_attention_metadata(training_batch) + training_batch.attn_metadata_vsa = copy.copy(training_batch.attn_metadata) + if training_batch.attn_metadata is not None: + training_batch.attn_metadata.VSA_sparsity = 0.0 # type: ignore[attr-defined] + return training_batch + + def _augment_conditions( + self, + encoder_hidden_states: torch.Tensor, + track_visibility: torch.Tensor, + *, + object_ids: torch.Tensor | None, + track_weights: torch.Tensor | None, + generator: torch.Generator, + ) -> tuple[torch.Tensor, torch.Tensor]: + config = self.track_augmentation + if not config.enabled: + return encoder_hidden_states, track_visibility + + visibility = track_visibility.clone() + batch_size, num_frames, num_tracks = visibility.shape + device = visibility.device + + # Keep tensors rectangular for the collator; sampling drops tracks by + # zeroing their visibility. Sparse mode preserves one visible query per + # segmented object, then adds background/context queries. + if config.sparse_object_sampling: + if object_ids is None: + raise ValueError("sparse_object_sampling requires object_ids " + "in the WanTrack parquet") + for batch_index in range(batch_size): + valid = visibility[batch_index, 0] > 0.5 + keep = torch.zeros( + num_tracks, + device=device, + dtype=torch.bool, + ) + sample_object_ids = object_ids[batch_index] + for object_id in torch.unique(sample_object_ids): + if int(object_id.item()) < 0: + continue + candidates = torch.nonzero( + (sample_object_ids == object_id) & valid, + as_tuple=True, + )[0] + if candidates.numel() == 0: + candidates = torch.nonzero( + sample_object_ids == object_id, + as_tuple=True, + )[0] + if candidates.numel() > 0: + choice = torch.randint( + int(candidates.numel()), + (1, ), + generator=generator, + device=device, + ) + keep[candidates[choice]] = True + + pool = valid & ~keep + extra_count = min( + config.extra_points, + int(pool.sum().item()), + ) + if extra_count > 0: + weights = pool.to(dtype=torch.float32) + if (config.extra_point_sampling == "weighted" and track_weights is not None): + weights *= (0.05 + track_weights[batch_index].clamp_min(0)) + selected = torch.multinomial( + weights, + extra_count, + replacement=False, + generator=generator, + ) + keep[selected] = True + visibility[batch_index] *= keep.to(dtype=visibility.dtype).unsqueeze(0) + elif config.max_points > 0 and config.min_points < num_tracks: + high = min(config.max_points, num_tracks) + low = min(config.min_points, high) + for batch_index in range(batch_size): + keep_count = int(torch.randint( + low, + high + 1, + (1, ), + generator=generator, + device=device, + ).item()) + permutation = torch.randperm( + num_tracks, + generator=generator, + device=device, + ) + keep = torch.zeros( + num_tracks, + device=device, + dtype=visibility.dtype, + ) + keep[permutation[:keep_count]] = 1 + visibility[batch_index] *= keep.unsqueeze(0) + + if config.track_dropout_probability > 0: + for batch_index in range(batch_size): + kept = visibility[batch_index].amax(dim=0) > 0.5 + kept_indices = torch.nonzero(kept, as_tuple=True)[0] + if kept_indices.numel() <= 1: + continue + drop = (torch.rand( + num_tracks, + generator=generator, + device=device, + ) < config.track_dropout_probability) & kept + if int((kept & ~drop).sum().item()) == 0: + spare = kept_indices[torch.randint( + int(kept_indices.numel()), + (1, ), + generator=generator, + device=device, + )] + drop[spare] = False + visibility[batch_index, :, drop] = 0 + + chunk_size = config.temporal_mask_chunk_size + num_chunks = (num_frames + chunk_size - 1) // chunk_size + if config.temporal_mask_probability > 0: + drop_chunks = torch.rand( + batch_size, + num_chunks, + generator=generator, + device=device, + ) < config.temporal_mask_probability + for chunk_index in range(num_chunks): + start = chunk_index * chunk_size + end = min(start + chunk_size, num_frames) + visibility[:, start:end] *= (~drop_chunks[:, chunk_index]).to(dtype=visibility.dtype).view(-1, 1, 1) + + if config.motion_dropout_probability > 0: + drop_motion = torch.rand( + batch_size, + generator=generator, + device=device, + ) < config.motion_dropout_probability + visibility[drop_motion] = 0 + + if config.text_dropout_probability > 0: + drop_text = torch.rand( + batch_size, + generator=generator, + device=device, + ) < config.text_dropout_probability + encoder_hidden_states = encoder_hidden_states.clone() + encoder_hidden_states[drop_text] = 0 + + return encoder_hidden_states, visibility + + def _prepare_dit_inputs( + self, + training_batch: TrainingBatch, + generator: torch.Generator, + ) -> TrainingBatch: + training_batch = super()._prepare_dit_inputs( + training_batch, + generator, + ) + if training_batch.image_latents is None: + raise RuntimeError("WanTrack requires first_frame_latent") + if training_batch.image_embeds is None: + raise RuntimeError("WanTrack requires clip_feature") + + image_condition = self._build_i2v_condition(training_batch.image_latents) + training_batch.image_latents = image_condition + assert training_batch.noisy_model_input is not None + training_batch.noisy_model_input = torch.cat( + [training_batch.noisy_model_input, image_condition], + dim=1, + ) + assert training_batch.conditional_dict is not None + extras = { + "image_latents": image_condition, + "encoder_hidden_states_image": training_batch.image_embeds, + "track_points": training_batch.track_points, + "track_visibility": training_batch.track_visibility, + "track_ids": training_batch.track_ids, + } + training_batch.conditional_dict.update(extras) + if training_batch.unconditional_dict is None: + training_batch.unconditional_dict = dict(training_batch.conditional_dict) + else: + training_batch.unconditional_dict.update(extras) + return training_batch + + def _build_distill_input_kwargs( + self, + noise_input: torch.Tensor, + timestep: torch.Tensor, + text_dict: dict[str, Any] | None, + clean_x: torch.Tensor | None = None, + aug_t: torch.Tensor | None = None, + start_frame: int = 0, + ) -> dict[str, Any]: + if text_dict is None: + raise ValueError("text_dict cannot be None for WanTrack") + hidden_states = self._append_i2v_for_window( + noise_input, + text_dict, + start_frame=start_frame, + ) + kwargs: dict[str, Any] = { + "hidden_states": hidden_states, + "encoder_hidden_states": text_dict["encoder_hidden_states"], + "encoder_attention_mask": text_dict["encoder_attention_mask"], + "encoder_hidden_states_image": text_dict.get("encoder_hidden_states_image"), + "track_points": text_dict.get("track_points"), + "track_visibility": text_dict.get("track_visibility"), + "track_ids": text_dict.get("track_ids"), + "timestep": timestep, + "return_dict": False, + } + if clean_x is not None: + kwargs["clean_x"] = self._append_i2v_for_window( + clean_x, + text_dict, + start_frame=start_frame, + ) + kwargs["aug_t"] = aug_t + return kwargs + + def _append_i2v_for_window( + self, + latents: torch.Tensor, + text_dict: dict[str, Any], + *, + start_frame: int, + ) -> torch.Tensor: + hidden_states = latents.permute(0, 2, 1, 3, 4) + if hidden_states.shape[1] == 36: + return hidden_states + if hidden_states.shape[1] != 16: + raise ValueError("WanTrack expects a 16-channel latent or a 36-channel " + f"latent+I2V input, got {hidden_states.shape[1]}") + image_condition = text_dict.get("image_latents") + if image_condition is None: + raise RuntimeError("WanTrack requires image_latents for a 16-channel input") + start = int(start_frame) + end = start + hidden_states.shape[2] + if start < 0 or end > image_condition.shape[2]: + raise ValueError("I2V condition does not cover the requested latent window " + f"[{start}, {end})") + return torch.cat( + [hidden_states, image_condition[:, :, start:end]], + dim=1, + ) + + def _get_uncond_text_dict( + self, + batch: TrainingBatch, + *, + cfg_uncond: dict[str, Any] | None, + ) -> dict[str, Any]: + sanitized_cfg = None + track_policy = "keep" + if cfg_uncond is not None: + sanitized_cfg = dict(cfg_uncond) + policies = [] + for channel in ("track", "motion"): + value = sanitized_cfg.pop(channel, None) + if value is not None: + if not isinstance(value, str): + raise ValueError(f"cfg_uncond.{channel} must be a string") + policies.append(value.strip().lower()) + if policies: + if len(set(policies)) != 1: + raise ValueError("cfg_uncond.track and cfg_uncond.motion conflict") + track_policy = policies[0] + + result = dict(super()._get_uncond_text_dict( + batch, + cfg_uncond=sanitized_cfg, + )) + conditional = batch.conditional_dict + if conditional is None: + raise RuntimeError("Missing conditional_dict in TrainingBatch") + for key in ( + "image_latents", + "encoder_hidden_states_image", + "track_points", + "track_visibility", + "track_ids", + ): + result.setdefault(key, conditional.get(key)) + + if track_policy not in {"keep", "zero", "drop"}: + raise ValueError("cfg_uncond.track/motion must be one of {keep, zero, drop}, " + f"got {track_policy!r}") + if track_policy in {"zero", "drop"}: + result["track_points"] = None + result["track_visibility"] = None + result["track_ids"] = None + return result + + def _build_i2v_condition( + self, + image_latents: torch.Tensor, + ) -> torch.Tensor: + if image_latents.ndim != 5: + raise ValueError("first_frame_latent must be [B, C, T, H, W], got " + f"{tuple(image_latents.shape)}") + if image_latents.shape[1] == 20: + return image_latents + if image_latents.shape[1] != 16: + raise ValueError("first_frame_latent must have 16 or 20 channels, got " + f"{image_latents.shape[1]}") + + ratio = self._temporal_compression_ratio() + batch_size, _, latent_t, height, width = image_latents.shape + pixel_frames = (latent_t - 1) * ratio + 1 + mask = torch.ones( + batch_size, + 1, + pixel_frames, + height, + width, + device=image_latents.device, + dtype=image_latents.dtype, + ) + mask[:, :, 1:] = 0 + mask = torch.cat([ + mask[:, :, :1].repeat_interleave(ratio, dim=2), + mask[:, :, 1:], + ], dim=2) + mask = mask.view( + batch_size, + -1, + ratio, + height, + width, + ).transpose(1, 2) + return torch.cat([mask, image_latents], dim=1) + + def _temporal_compression_ratio(self) -> int: + assert self.training_config is not None + return int(self.training_config.pipeline_config.vae_config.arch_config.temporal_compression_ratio) diff --git a/fastvideo/train/models/wantrack/wantrack_causal.py b/fastvideo/train/models/wantrack/wantrack_causal.py new file mode 100644 index 000000000..40cf39de9 --- /dev/null +++ b/fastvideo/train/models/wantrack/wantrack_causal.py @@ -0,0 +1,11 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Causal WanTrack training model.""" + +from fastvideo.train.models.wan.wan_causal import WanCausalModel +from fastvideo.train.models.wantrack.wantrack import WanTrackModel + + +class WanTrackCausalModel(WanTrackModel, WanCausalModel): + """WanTrack with block-causal attention and streaming cache support.""" + + _transformer_cls_name = "CausalTrackWanTransformer3DModel" diff --git a/fastvideo/train/utils/config.py b/fastvideo/train/utils/config.py index a81314fb3..113609b4c 100644 --- a/fastvideo/train/utils/config.py +++ b/fastvideo/train/utils/config.py @@ -54,8 +54,7 @@ class RunConfig: if dataclasses.is_dataclass(obj) and not isinstance(obj, type): return { f.name: _safe_asdict(getattr(obj, f.name)) - for f in dataclasses.fields(obj) - if not callable(getattr(obj, f.name)) + for f in dataclasses.fields(obj) if not callable(getattr(obj, f.name)) } if isinstance(obj, dict): return {k: _safe_asdict(v) for k, v in obj.items()} @@ -349,12 +348,10 @@ def _build_training_config( 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 " - "{'t2v', 'text_only'}, got " - f"{preprocessed_data_type!r}" - ) + if preprocessed_data_type not in {"t2v", "i2v_track", "text_only"}: + raise ValueError("training.data.preprocessed_data_type must be one of " + "{'t2v', 'i2v_track', 'text_only'}, got " + f"{preprocessed_data_type!r}") return TrainingConfig( distributed=DistributedConfig( @@ -424,9 +421,7 @@ def _build_training_config( ) -def _parse_cli_overrides( - overrides: list[str], -) -> dict[str, Any]: +def _parse_cli_overrides(overrides: list[str], ) -> dict[str, Any]: """Parse ``--dotted.key value`` CLI overrides. Returns a flat dict mapping dotted keys to parsed diff --git a/mkdocs.yml b/mkdocs.yml index 3cc6a92a0..2b88c1382 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -172,6 +172,7 @@ nav: - Experiment Tracking: training/trackers.md - Data Preprocessing: training/data_preprocess.md - Fine-tuning: training/finetune.md + - WanTrack Training: training/wantrack.md - Attn-QAT Training: training/attn_qat.md - Examples: - Training Index: training/examples/examples_training_index.md diff --git a/scripts/checkpoint_conversion/wan_fun_control_to_trackwan.py b/scripts/checkpoint_conversion/wan_fun_control_to_trackwan.py new file mode 100644 index 000000000..837e5bf5a --- /dev/null +++ b/scripts/checkpoint_conversion/wan_fun_control_to_trackwan.py @@ -0,0 +1,412 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""Build a WanTrack initialization checkpoint from Wan2.1-Fun Control weights. + +The VideoX-Fun Control DiT receives 48 input channels in this order: + + noisy latent (16) | spatial control (16) | reference latent (16) + +WanTrack receives 52 input channels: + + noisy latent (16) | I2V mask (4) | first-frame latent (16) | track map (16) + +This converter preserves the pretrained control pathway by moving the Fun +Control channels into WanTrack's track slot. The four new I2V mask channels are +initialized to zero. The remainder of the model is renamed to the Diffusers +key layout and checked against a compatible Fun-InP Diffusers checkpoint. +""" + +from __future__ import annotations + +import argparse +import json +import re +import shutil +from pathlib import Path + +import torch +from safetensors.torch import load_file, save_file +from torch import nn + + +SOURCE_IN_CHANNELS = 48 +TARGET_IN_CHANNELS = 52 +LATENT_CHANNELS = 16 +TRACK_CHANNELS = 16 +VAE_SPATIAL_COMPRESSION = 8 +VAE_TEMPORAL_COMPRESSION = 4 + +# VideoX-Fun Wan names to the Diffusers/FastVideo Wan state-dict surface. +# patch_embedding is handled separately because its input channels are remapped. +PARAM_NAME_MAP: tuple[tuple[re.Pattern[str], str], ...] = tuple( + (re.compile(pattern), replacement) + for pattern, replacement in { + r"^text_embedding\.0\.(.*)$": r"condition_embedder.text_embedder.linear_1.\1", + r"^text_embedding\.2\.(.*)$": r"condition_embedder.text_embedder.linear_2.\1", + r"^time_embedding\.0\.(.*)$": r"condition_embedder.time_embedder.linear_1.\1", + r"^time_embedding\.2\.(.*)$": r"condition_embedder.time_embedder.linear_2.\1", + r"^time_projection\.1\.(.*)$": r"condition_embedder.time_proj.\1", + r"^img_emb\.proj\.0\.(.*)$": r"condition_embedder.image_embedder.norm1.\1", + r"^img_emb\.proj\.1\.(.*)$": r"condition_embedder.image_embedder.ff.net.0.proj.\1", + r"^img_emb\.proj\.3\.(.*)$": r"condition_embedder.image_embedder.ff.net.2.\1", + r"^img_emb\.proj\.4\.(.*)$": r"condition_embedder.image_embedder.norm2.\1", + r"^head\.modulation$": r"scale_shift_table", + r"^head\.head\.(.*)$": r"proj_out.\1", + r"^blocks\.(\d+)\.self_attn\.q\.(.*)$": r"blocks.\1.attn1.to_q.\2", + r"^blocks\.(\d+)\.self_attn\.k\.(.*)$": r"blocks.\1.attn1.to_k.\2", + r"^blocks\.(\d+)\.self_attn\.v\.(.*)$": r"blocks.\1.attn1.to_v.\2", + r"^blocks\.(\d+)\.self_attn\.o\.(.*)$": r"blocks.\1.attn1.to_out.0.\2", + r"^blocks\.(\d+)\.self_attn\.norm_q\.(.*)$": r"blocks.\1.attn1.norm_q.\2", + r"^blocks\.(\d+)\.self_attn\.norm_k\.(.*)$": r"blocks.\1.attn1.norm_k.\2", + r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$": r"blocks.\1.attn2.to_q.\2", + r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$": r"blocks.\1.attn2.to_k.\2", + r"^blocks\.(\d+)\.cross_attn\.k_img\.(.*)$": r"blocks.\1.attn2.add_k_proj.\2", + r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$": r"blocks.\1.attn2.to_v.\2", + r"^blocks\.(\d+)\.cross_attn\.v_img\.(.*)$": r"blocks.\1.attn2.add_v_proj.\2", + r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$": r"blocks.\1.attn2.to_out.0.\2", + r"^blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$": r"blocks.\1.attn2.norm_q.\2", + r"^blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$": r"blocks.\1.attn2.norm_k.\2", + r"^blocks\.(\d+)\.cross_attn\.norm_k_img\.(.*)$": r"blocks.\1.attn2.norm_added_k.\2", + r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.net.0.proj.\2", + r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.net.2.\2", + r"^blocks\.(\d+)\.modulation$": r"blocks.\1.scale_shift_table", + r"^blocks\.(\d+)\.norm3\.(.*)$": r"blocks.\1.norm2.\2", + }.items() +) + + +def load_sharded_safetensors(directory: Path) -> dict[str, torch.Tensor]: + """Load every safetensors shard in a directory and reject duplicate keys.""" + shard_paths = sorted(directory.glob("*.safetensors")) + if not shard_paths: + raise FileNotFoundError(f"No safetensors files found under {directory}") + + state_dict: dict[str, torch.Tensor] = {} + for shard_path in shard_paths: + shard = load_file(str(shard_path), device="cpu") + duplicate_keys = state_dict.keys() & shard.keys() + if duplicate_keys: + examples = ", ".join(sorted(duplicate_keys)[:8]) + raise RuntimeError(f"Duplicate parameter keys across base shards: {examples}") + state_dict.update(shard) + return state_dict + + +def rename_control_state_dict(source: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Rename all VideoX-Fun parameters and reject unknown or colliding keys.""" + converted: dict[str, torch.Tensor] = {} + for source_key, value in source.items(): + if source_key.startswith("patch_embedding."): + target_key = source_key + else: + target_key = "" + for pattern, replacement in PARAM_NAME_MAP: + if pattern.match(source_key): + target_key = pattern.sub(replacement, source_key) + break + if not target_key: + raise ValueError(f"No target mapping for source parameter: {source_key}") + + if target_key in converted: + raise RuntimeError(f"Multiple source parameters map to target parameter: {target_key}") + converted[target_key] = value + return converted + + +def synthesize_image_query_norms( + converted: dict[str, torch.Tensor], + reference: dict[str, torch.Tensor], +) -> None: + """Create Diffusers image-query norm tensors absent from VideoX-Fun.""" + expected_keys = sorted(key for key in reference if key.endswith("attn2.norm_added_q.weight")) + for query_key in expected_keys: + key_norm_key = query_key.replace("norm_added_q", "norm_added_k") + if key_norm_key not in converted: + raise RuntimeError(f"Cannot synthesize {query_key}: converted state lacks {key_norm_key}") + if query_key in converted: + raise RuntimeError(f"Refusing to overwrite converted parameter: {query_key}") + converted[query_key] = torch.zeros_like(converted[key_norm_key]) + + +def remap_patch_embedding(source_weight: torch.Tensor) -> torch.Tensor: + """Remap Fun Control's 48-channel patch projection into WanTrack's 52 channels.""" + if source_weight.ndim != 5: + raise ValueError(f"Expected a 5D patch embedding weight, got shape {tuple(source_weight.shape)}") + if source_weight.shape[1] != SOURCE_IN_CHANNELS: + raise ValueError( + f"Expected patch_embedding.weight input channels={SOURCE_IN_CHANNELS}, " + f"got {source_weight.shape[1]}" + ) + + target_shape = (source_weight.shape[0], TARGET_IN_CHANNELS, *source_weight.shape[2:]) + target_weight = source_weight.new_zeros(target_shape) + target_weight[:, 0:16] = source_weight[:, 0:16] + target_weight[:, 20:36] = source_weight[:, 32:48] + target_weight[:, 36:52] = source_weight[:, 16:32] + + if torch.count_nonzero(target_weight[:, 16:20]).item() != 0: + raise RuntimeError("WanTrack I2V mask channels must be initialized to zero") + return target_weight + + +def validate_converted_body( + converted: dict[str, torch.Tensor], + reference: dict[str, torch.Tensor], +) -> None: + """Require the converted Wan body to match the Fun-InP key and shape surface.""" + converted_keys = set(converted) + reference_keys = set(reference) + missing = sorted(reference_keys - converted_keys) + unexpected = sorted(converted_keys - reference_keys) + if missing or unexpected: + raise RuntimeError( + "Converted transformer keys do not match the Fun-InP reference. " + f"missing={missing[:8]} ({len(missing)} total), " + f"unexpected={unexpected[:8]} ({len(unexpected)} total)" + ) + + patch_key = "patch_embedding.weight" + if patch_key not in converted: + raise RuntimeError(f"Converted transformer lacks required parameter: {patch_key}") + reference_patch = reference[patch_key] + converted_patch = converted[patch_key] + if reference_patch.ndim != converted_patch.ndim: + raise RuntimeError( + f"{patch_key} rank mismatch: reference={reference_patch.ndim}, converted={converted_patch.ndim}" + ) + if reference_patch.shape[0] != converted_patch.shape[0] or reference_patch.shape[2:] != converted_patch.shape[2:]: + raise RuntimeError( + f"{patch_key} non-input dimensions differ: " + f"reference={tuple(reference_patch.shape)}, converted={tuple(converted_patch.shape)}" + ) + if converted_patch.shape[1] != TARGET_IN_CHANNELS: + raise RuntimeError(f"{patch_key} must have {TARGET_IN_CHANNELS} input channels") + + shape_mismatches = [] + for key in sorted(reference_keys - {patch_key}): + if converted[key].shape != reference[key].shape: + shape_mismatches.append((key, tuple(reference[key].shape), tuple(converted[key].shape))) + if shape_mismatches: + raise RuntimeError( + "Converted transformer shapes do not match the Fun-InP reference: " + f"{shape_mismatches[:8]} ({len(shape_mismatches)} total)" + ) + + +def build_track_encoder_state( + id_dim: int, + dtype: torch.dtype, + seed: int, + zero_init_head: bool, +) -> dict[str, torch.Tensor]: + """Initialize the two bias-free TrackEncoder convolution weights reproducibly.""" + if id_dim <= 0: + raise ValueError(f"--id-dim must be positive, got {id_dim}") + + with torch.random.fork_rng(devices=[]): + torch.manual_seed(seed) + temporal_conv = nn.Conv3d( + id_dim, + TRACK_CHANNELS, + kernel_size=(VAE_TEMPORAL_COMPRESSION, 1, 1), + stride=(VAE_TEMPORAL_COMPRESSION, 1, 1), + bias=False, + ) + projection = nn.Conv3d(TRACK_CHANNELS, TRACK_CHANNELS, kernel_size=1, bias=False) + if zero_init_head: + nn.init.zeros_(projection.weight) + + state = { + "track_encoder.temporal_conv.weight": temporal_conv.weight.detach().to(dtype=dtype).contiguous(), + "track_encoder.proj.weight": projection.weight.detach().to(dtype=dtype).contiguous(), + } + expected_shapes = { + "track_encoder.temporal_conv.weight": ( + TRACK_CHANNELS, + id_dim, + VAE_TEMPORAL_COMPRESSION, + 1, + 1, + ), + "track_encoder.proj.weight": (TRACK_CHANNELS, TRACK_CHANNELS, 1, 1, 1), + } + for key, expected_shape in expected_shapes.items(): + if state[key].shape != expected_shape: + raise RuntimeError(f"Unexpected initialized shape for {key}: {tuple(state[key].shape)}") + return state + + +def copy_passthrough_components(source_dir: Path, output_dir: Path) -> None: + """Copy every base bundle entry except the transformer directory.""" + for source_path in sorted(source_dir.iterdir(), key=lambda path: path.name): + if source_path.name == "transformer": + continue + target_path = output_dir / source_path.name + if source_path.is_dir(): + shutil.copytree(source_path, target_path, symlinks=True) + else: + shutil.copy2(source_path, target_path, follow_symlinks=False) + + +def update_model_index(output_dir: Path, transformer_class: str) -> None: + """Point the copied bundle index at the selected FastVideo transformer.""" + model_index_path = output_dir / "model_index.json" + if not model_index_path.is_file(): + raise FileNotFoundError(f"Fun-InP base does not provide {model_index_path.name}") + + model_index = json.loads(model_index_path.read_text(encoding="utf-8")) + transformer_entry = model_index.get("transformer") + if not isinstance(transformer_entry, list) or len(transformer_entry) != 2: + raise ValueError(f"Unsupported transformer entry in {model_index_path}: {transformer_entry!r}") + transformer_entry[1] = transformer_class + model_index_path.write_text(json.dumps(model_index, indent=2) + "\n", encoding="utf-8") + + +def write_transformer_config( + base_dir: Path, + output_dir: Path, + transformer_class: str, + id_dim: int, + max_track_id: int, + zero_init_head: bool, +) -> None: + """Write the widened, loader-compatible WanTrack transformer config.""" + source_config_path = base_dir / "transformer" / "config.json" + if not source_config_path.is_file(): + raise FileNotFoundError(f"Missing Fun-InP transformer config: {source_config_path}") + + config = json.loads(source_config_path.read_text(encoding="utf-8")) + config.update( + { + "_class_name": transformer_class, + "in_channels": TARGET_IN_CHANNELS, + "out_channels": LATENT_CHANNELS, + "track_config": { + "id_dim": id_dim, + "track_channels": TRACK_CHANNELS, + "vae_spatial_compression": VAE_SPATIAL_COMPRESSION, + "vae_temporal_compression": VAE_TEMPORAL_COMPRESSION, + "max_track_id": max_track_id, + "zero_init_head": zero_init_head, + "use_bias": False, + }, + } + ) + config_path = output_dir / "transformer" / "config.json" + config_path.write_text(json.dumps(config, indent=2) + "\n", encoding="utf-8") + + +def validate_paths(base_dir: Path, control_checkpoint: Path, output_dir: Path) -> None: + """Validate input paths and refuse ambiguous or destructive output layouts.""" + if not base_dir.is_dir(): + raise NotADirectoryError(f"Fun-InP base directory does not exist: {base_dir}") + if not control_checkpoint.is_file(): + raise FileNotFoundError(f"Fun Control checkpoint does not exist: {control_checkpoint}") + if control_checkpoint.suffix != ".safetensors": + raise ValueError(f"Fun Control checkpoint must be a .safetensors file: {control_checkpoint}") + + base_resolved = base_dir.resolve() + output_resolved = output_dir.resolve() + if output_resolved == base_resolved: + raise ValueError("Output directory must differ from the Fun-InP base directory") + if output_resolved.is_relative_to(base_resolved) or base_resolved.is_relative_to(output_resolved): + raise ValueError("Output directory and Fun-InP base directory must not contain one another") + if output_dir.exists() and not output_dir.is_dir(): + raise NotADirectoryError(f"Output path exists and is not a directory: {output_dir}") + if output_dir.exists() and any(output_dir.iterdir()): + raise FileExistsError(f"Refusing to overwrite non-empty output directory: {output_dir}") + + +def parse_args() -> argparse.Namespace: + """Parse command-line arguments.""" + parser = argparse.ArgumentParser( + description="Convert VideoX-Fun Wan Control weights into a FastVideo WanTrack bundle.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument( + "--inp-base", + required=True, + type=Path, + help="Compatible Fun-InP Diffusers bundle used for config, key validation, and passthrough components.", + ) + parser.add_argument( + "--control-ckpt", + required=True, + type=Path, + help="VideoX-Fun Wan Control diffusion_pytorch_model.safetensors.", + ) + parser.add_argument("--out", required=True, type=Path, help="Output directory for the converted WanTrack bundle.") + parser.add_argument("--id-dim", type=int, default=128, help="Sinusoidal track-ID embedding dimension.") + parser.add_argument( + "--max-track-id", + type=int, + default=100_000, + help="Exclusive upper bound for sampled track IDs.", + ) + parser.add_argument("--seed", type=int, default=0, help="Seed used to initialize TrackEncoder weights.") + parser.add_argument( + "--zero-init-head", + action=argparse.BooleanOptionalAction, + default=True, + help="Zero-initialize TrackEncoder's final projection.", + ) + parser.add_argument( + "--causal", + action="store_true", + help="Emit CausalTrackWanTransformer3DModel instead of the bidirectional transformer class.", + ) + return parser.parse_args() + + +def main() -> None: + """Convert the checkpoint and create a complete WanTrack bundle.""" + args = parse_args() + if args.max_track_id <= 0: + raise ValueError(f"--max-track-id must be positive, got {args.max_track_id}") + validate_paths(args.inp_base, args.control_ckpt, args.out) + + reference_state = load_sharded_safetensors(args.inp_base / "transformer") + control_state = load_file(str(args.control_ckpt), device="cpu") + converted_state = rename_control_state_dict(control_state) + synthesize_image_query_norms(converted_state, reference_state) + + patch_key = "patch_embedding.weight" + if patch_key not in converted_state: + raise RuntimeError(f"Fun Control checkpoint lacks required parameter: {patch_key}") + converted_state[patch_key] = remap_patch_embedding(converted_state[patch_key]) + validate_converted_body(converted_state, reference_state) + + converted_state.update( + build_track_encoder_state( + id_dim=args.id_dim, + dtype=converted_state[patch_key].dtype, + seed=args.seed, + zero_init_head=args.zero_init_head, + ) + ) + + transformer_class = "CausalTrackWanTransformer3DModel" if args.causal else "TrackWanTransformer3DModel" + args.out.mkdir(parents=True, exist_ok=True) + copy_passthrough_components(args.inp_base, args.out) + (args.out / "transformer").mkdir() + update_model_index(args.out, transformer_class) + write_transformer_config( + base_dir=args.inp_base, + output_dir=args.out, + transformer_class=transformer_class, + id_dim=args.id_dim, + max_track_id=args.max_track_id, + zero_init_head=args.zero_init_head, + ) + + output_weights = args.out / "transformer" / "diffusion_pytorch_model.safetensors" + contiguous_state = {key: value.contiguous() for key, value in converted_state.items()} + save_file(contiguous_state, str(output_weights), metadata={"format": "pt"}) + print( + f"Converted {len(reference_state)} Wan parameters and 2 TrackEncoder parameters " + f"to {output_weights} ({transformer_class})." + ) + + +if __name__ == "__main__": + main()