diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml index 5d847a94d..dcd14e131 100644 --- a/docs/source/_toctree.yml +++ b/docs/source/_toctree.yml @@ -69,6 +69,8 @@ title: VLA-JEPA - local: eo1 title: EO-1 + - local: fastwam + title: FastWAM - local: groot title: NVIDIA GR00T N1.5 - local: xvla diff --git a/docs/source/fastwam.mdx b/docs/source/fastwam.mdx new file mode 100644 index 000000000..18b4775f8 --- /dev/null +++ b/docs/source/fastwam.mdx @@ -0,0 +1,167 @@ +# FastWAM + +FastWAM is a World Action Model policy for robot control. The LeRobot integration exposes FastWAM through the standard policy API so it can be configured with `policy.type=fastwam`, trained with `lerobot-train`, and loaded through the LeRobot pretrained policy interface. + +## Model Overview + +FastWAM keeps video modeling during training, but uses direct action prediction at inference time instead of iteratively generating future observations. This LeRobot policy wraps the FastWAM action model, adapts LeRobot batches to FastWAM training samples, and provides the standard processor pipeline for normalization and action postprocessing. + +The implementation initializes the visual world-model components from `Wan-AI/Wan2.2-TI2V-5B` by default and predicts action chunks with shape `[batch, action_horizon, action_dim]`. + +### What the LeRobot Integration Covers + +- Standard `policy.type=fastwam` configuration through LeRobot +- Image, state, action, and language-task batch adaptation +- Action chunk inference through `select_action` and `predict_action_chunk` +- Checkpoint save/load through the LeRobot policy APIs +- Configurable LIBERO gripper action postprocessing + +## Installation Requirements + +Install LeRobot from source, then install FastWAM dependencies: + +```bash +pip install -e ".[fastwam]" +``` + +This installs the FastWAM policy extra from `pyproject.toml`: `transformers`, +`diffusers`, `ftfy`, and `regex`, plus LeRobot's base dependencies. + +For LIBERO evaluation, install the benchmark dependencies too: + +```bash +pip install -e ".[fastwam,libero]" +``` + +This installs both extras. In addition to the FastWAM dependencies above, the +`libero` extra installs LeRobot dataset dependencies, `hf-libero` on Linux, and +`scipy`. + +FastWAM uses the Wan2.2 TI2V backbone. The default model id is: + +```python +policy.model_id=Wan-AI/Wan2.2-TI2V-5B +``` + +## Data Requirements + +FastWAM expects a LeRobot dataset with: + +- one or more visual observations whose widths concatenate to `policy.image_size[1]` +- `observation.state` when `policy.proprio_dim` is not `None` +- `action` +- a language task instruction through the dataset task field, or precomputed `context` and `context_mask` tensors + +The default visual setup is one image feature named `observation.images.image` with shape `(3, 224, 448)`. If the dataset uses two cameras, configure `policy.input_features` so their heights match `224` and their widths sum to `448`. + +## Usage + +Create a new FastWAM policy with: + +```bash +lerobot-train \ + --dataset.repo_id=your-org/your-dataset \ + --policy.type=fastwam \ + --policy.action_dim=7 \ + --policy.proprio_dim=8 \ + --policy.action_horizon=32 \ + --policy.n_action_steps=10 \ + --policy.image_size='[224,448]' \ + --output_dir=./outputs/fastwam_training \ + --job_name=fastwam_training \ + --steps=300000 \ + --batch_size=8 \ + --policy.device=cuda +``` + +Evaluate an existing LeRobot-format checkpoint on LIBERO-10 with: + +```bash +lerobot-eval \ + --policy.path=ZibinDong/fastwam_libero_uncond_2cam224 \ + --policy.device=cuda \ + --policy.torch_dtype=float32 \ + --policy.n_action_steps=10 \ + --env.type=libero \ + --env.task=libero_10 \ + --env.observation_height=224 \ + --env.observation_width=224 \ + --eval.batch_size=1 \ + --eval.n_episodes=50 \ + --seed=0 \ + --env.episode_length=600 +``` + +For `libero_goal`, `libero_spatial`, and `libero_object`, use +`--env.episode_length=300`. + +For real-robot rollout, use the same checkpoint path: + +```bash +lerobot-rollout \ + --robot.type=so101_follower \ + --robot.port=/dev/ttyACM0 \ + --policy.path=your-org/fastwam-real-robot +``` + +## Configuration Notes + +### Image Features + +`policy.image_size` is the size of the concatenated FastWAM image tensor as `(height, width)`. Each configured image feature must have shape `(3, height, camera_width)`, and all camera widths must sum to the configured width. + +### Action Chunking + +`policy.action_horizon` controls the number of future actions supervised during training and predicted during inference. `policy.n_action_steps` controls how many actions are consumed before the policy predicts a fresh chunk. `policy.n_action_steps` must be less than or equal to `policy.action_horizon`. + +### Wan Components + +FastWAM loads the Wan VAE, video DiT, text encoder, and tokenizer from the configured Wan model directory or Hugging Face Hub model id. LeRobot-format FastWAM checkpoints saved by `save_pretrained` also copy the local Wan component files needed by `from_pretrained`. + +### Attention Backend + +FastWAM's DiT uses PyTorch's `scaled_dot_product_attention` (SDPA) for all attention. It does **not** use FlashAttention: its Mixture-of-Transformers (MoT) routing needs arbitrary boolean `[query, key]` attention masks, which the FlashAttention varlen API cannot express. Installing the `flash-attn` package therefore has no effect on the FastWAM path. (Note that SDPA itself may still select PyTorch's own flash / memory-efficient / math kernel internally — this is unrelated to the `flash-attn` package.) + +### LIBERO Action Toggle + +FastWAM LIBERO checkpoints use `policy.toggle_action_dimensions=[-1]` by +default to match the gripper action convention used by the original FastWAM +evaluation pipeline: + +```bash +--policy.toggle_action_dimensions='[-1]' +``` + +## Results + +Evaluated on LIBERO with [`ZibinDong/fastwam_libero_uncond_2cam224`](https://huggingface.co/ZibinDong/fastwam_libero_uncond_2cam224): + +| Suite | Success rate | n_episodes | +| -------------- | -----------: | ---------: | +| libero_spatial | 97.6% | 500 | +| libero_object | 99.0% | 500 | +| libero_goal | 95.0% | 500 | +| libero_10 | 94.0% | 500 | +| **average** | **96.4%** | 2000 | + +Reproduce: `lerobot-eval --policy.path=ZibinDong/fastwam_libero_uncond_2cam224 --policy.device=cuda --policy.torch_dtype=float32 --policy.n_action_steps=10 --env.type=libero --env.task=libero_spatial --env.observation_height=256 --env.observation_width=256 --eval.batch_size=1 --eval.n_episodes=50 --seed=0 --env.episode_length=300` (1x H20 140 GB). + +## References + +- [Fast-WAM paper](https://arxiv.org/abs/2603.16666) +- [Fast-WAM project page](https://yuantianyuan01.github.io/FastWAM/) +- [Fast-WAM code](https://github.com/yuantianyuan01/FastWAM) +- [Released upstream checkpoints](https://huggingface.co/yuanty/fastwam) +- [Wan2.2 TI2V 5B](https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B) + +## Citation + +```bibtex +@article{yuan2026fastwam, + title = {Fast-WAM: Do World Action Models Need Test-time Future Imagination?}, + author = {Tianyuan Yuan and Zibin Dong and Yicheng Liu and Hang Zhao}, + journal = {arXiv preprint arXiv:2603.16666}, + year = {2026}, + url = {https://arxiv.org/abs/2603.16666} +} +``` diff --git a/docs/source/policy_fastwam_README.md b/docs/source/policy_fastwam_README.md new file mode 100644 index 000000000..6af0eaa79 --- /dev/null +++ b/docs/source/policy_fastwam_README.md @@ -0,0 +1,56 @@ +## Research Paper + +Paper: https://arxiv.org/abs/2603.16666 + +## Repository + +Code: https://github.com/yuantianyuan01/FastWAM + +Project page: https://yuantianyuan01.github.io/FastWAM/ + +## Citation + +```bibtex +@article{yuan2026fastwam, + title = {Fast-WAM: Do World Action Models Need Test-time Future Imagination?}, + author = {Tianyuan Yuan and Zibin Dong and Yicheng Liu and Hang Zhao}, + journal = {arXiv preprint arXiv:2603.16666}, + year = {2026}, + url = {https://arxiv.org/abs/2603.16666} +} +``` + +## Additional Resources + +Base video model: https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B + +Released upstream checkpoints: https://huggingface.co/yuanty/fastwam + +## Results + +Evaluated on LIBERO with [`ZibinDong/fastwam_libero_uncond_2cam224`](https://huggingface.co/ZibinDong/fastwam_libero_uncond_2cam224): + +| Suite | Success rate | n_episodes | +| -------------- | -----------: | ---------: | +| libero_spatial | 97.6% | 500 | +| libero_object | 99.0% | 500 | +| libero_goal | 95.0% | 500 | +| libero_10 | 94.0% | 500 | +| **average** | **96.4%** | 2000 | + +Reproduce: `lerobot-eval --policy.path=ZibinDong/fastwam_libero_uncond_2cam224 --policy.device=cuda --policy.torch_dtype=float32 --policy.n_action_steps=10 --env.type=libero --env.task=libero_spatial --env.observation_height=256 --env.observation_width=256 --eval.batch_size=1 --eval.n_episodes=50 --seed=0 --env.episode_length=300`. + +For LIBERO-10, use `--env.task=libero_10 --env.episode_length=600`: + +```bash +lerobot-eval \ + --policy.path=ZibinDong/fastwam_libero_uncond_2cam224 \ + --policy.device=cuda \ + --policy.torch_dtype=float32 \ + --policy.n_action_steps=10 \ + --env.type=libero \ + --env.task=libero_10 --env.observation_height=256 --env.observation_width=256 \ + --eval.batch_size=1 \ + --eval.n_episodes=50 \ + --seed=0 --env.episode_length=600 +``` diff --git a/pyproject.toml b/pyproject.toml index 28a8948b7..08ba7fc45 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -229,6 +229,10 @@ robometer = ["lerobot[transformers-dep]", "lerobot[qwen-vl-utils-dep]", "lerobot topreward = ["lerobot[transformers-dep]"] xvla = ["lerobot[transformers-dep]"] eo1 = ["lerobot[transformers-dep]", "lerobot[qwen-vl-utils-dep]"] +fastwam = [ + "lerobot[transformers-dep]", + "lerobot[diffusers-dep]", +] hilserl = ["lerobot[transformers-dep]", "lerobot[dataset]", "gym-hil>=0.1.14,<0.2.0", "lerobot[grpcio-dep]", "lerobot[placo-dep]"] vla_jepa = ["lerobot[transformers-dep]", "lerobot[diffusers-dep]", "lerobot[qwen-vl-utils-dep]"] @@ -308,6 +312,7 @@ all = [ "lerobot[pi]", "lerobot[molmoact2]", "lerobot[smolvla]", + "lerobot[fastwam]", # "lerobot[groot]", TODO(Steven): Gr00t requires specific installation instructions for flash-attn "lerobot[xvla]", "lerobot[hilserl]", @@ -444,7 +449,8 @@ default.extend-ignore-identifiers-re = [ "is_compileable", "ROBOTIS", "OT_VALUE", - "VanderBilt" + "VanderBilt", + "seperated_timestep", ] # TODO: Uncomment when ready to use diff --git a/src/lerobot/policies/__init__.py b/src/lerobot/policies/__init__.py index 68d23c9ca..4daa6abc5 100644 --- a/src/lerobot/policies/__init__.py +++ b/src/lerobot/policies/__init__.py @@ -18,6 +18,7 @@ from .act.configuration_act import ACTConfig as ACTConfig from .diffusion.configuration_diffusion import DiffusionConfig as DiffusionConfig from .eo1.configuration_eo1 import EO1Config as EO1Config from .factory import get_policy_class, make_policy, make_policy_config, make_pre_post_processors +from .fastwam.configuration_fastwam import FastWAMConfig as FastWAMConfig from .gaussian_actor.configuration_gaussian_actor import GaussianActorConfig as GaussianActorConfig from .groot.configuration_groot import GrootConfig as GrootConfig from .molmoact2.configuration_molmoact2 import MolmoAct2Config as MolmoAct2Config @@ -42,6 +43,7 @@ __all__ = [ "ACTConfig", "DiffusionConfig", "EO1Config", + "FastWAMConfig", "GaussianActorConfig", "GrootConfig", "MolmoAct2Config", diff --git a/src/lerobot/policies/factory.py b/src/lerobot/policies/factory.py index b82eaeb72..f5acb0170 100644 --- a/src/lerobot/policies/factory.py +++ b/src/lerobot/policies/factory.py @@ -47,6 +47,7 @@ from lerobot.utils.feature_utils import dataset_to_policy_features from .act.configuration_act import ACTConfig from .diffusion.configuration_diffusion import DiffusionConfig from .eo1.configuration_eo1 import EO1Config +from .fastwam.configuration_fastwam import FastWAMConfig from .gaussian_actor.configuration_gaussian_actor import GaussianActorConfig from .groot.configuration_groot import GrootConfig from .molmoact2.configuration_molmoact2 import MolmoAct2Config @@ -162,6 +163,10 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]: from .vla_jepa.modeling_vla_jepa import VLAJEPAPolicy return VLAJEPAPolicy + elif name == "fastwam": + from .fastwam.modeling_fastwam import FastWAMPolicy + + return FastWAMPolicy else: try: return _get_policy_cls_from_policy_name(name=name) @@ -218,6 +223,8 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig: return MolmoAct2Config(**kwargs) elif policy_type == "vla_jepa": return VLAJEPAConfig(**kwargs) + elif policy_type == "fastwam": + return FastWAMConfig(**kwargs) else: try: config_cls = PreTrainedConfig.get_choice_class(policy_type) @@ -451,6 +458,14 @@ def make_pre_post_processors( dataset_stats=kwargs.get("dataset_stats"), ) + elif isinstance(policy_cfg, FastWAMConfig): + from .fastwam.processor_fastwam import make_fastwam_pre_post_processors + + processors = make_fastwam_pre_post_processors( + config=policy_cfg, + dataset_stats=kwargs.get("dataset_stats"), + ) + else: try: processors = _make_processors_from_policy_config( diff --git a/src/lerobot/policies/fastwam/README.md b/src/lerobot/policies/fastwam/README.md new file mode 120000 index 000000000..d78b9ef36 --- /dev/null +++ b/src/lerobot/policies/fastwam/README.md @@ -0,0 +1 @@ +../../../../docs/source/policy_fastwam_README.md \ No newline at end of file diff --git a/src/lerobot/policies/fastwam/__init__.py b/src/lerobot/policies/fastwam/__init__.py new file mode 100644 index 000000000..8488e7b78 --- /dev/null +++ b/src/lerobot/policies/fastwam/__init__.py @@ -0,0 +1,23 @@ +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .configuration_fastwam import FastWAMConfig +from .modeling_fastwam import FastWAMPolicy +from .processor_fastwam import make_fastwam_pre_post_processors + +__all__ = [ + "FastWAMConfig", + "FastWAMPolicy", + "make_fastwam_pre_post_processors", +] diff --git a/src/lerobot/policies/fastwam/configuration_fastwam.py b/src/lerobot/policies/fastwam/configuration_fastwam.py new file mode 100644 index 000000000..a3ef4f602 --- /dev/null +++ b/src/lerobot/policies/fastwam/configuration_fastwam.py @@ -0,0 +1,399 @@ +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from lerobot.configs import ( + FeatureType, + NormalizationMode, + PolicyFeature, + PreTrainedConfig, +) +from lerobot.optim import AdamWConfig +from lerobot.utils.constants import ACTION, OBS_STATE + +WAN22_MODEL_ID = "Wan-AI/Wan2.2-TI2V-5B" +WAN22_DIFFUSERS_MODEL_ID = "Wan-AI/Wan2.2-TI2V-5B-Diffusers" +FASTWAM_BASE_MODEL_ID = "lerobot/fastwam_base" +WAN_T5_TOKENIZER_ID = "google/umt5-xxl" + + +_FASTWAM_VIDEO_BASE_COMPAT_KEYS = ( + "patch_size", + "in_dim", + "hidden_dim", + "ffn_dim", + "freq_dim", + "text_dim", + "out_dim", + "num_heads", + "attn_head_dim", + "num_layers", +) + +_FASTWAM_ACTION_BASE_COMPAT_KEYS = ( + "hidden_dim", + "ffn_dim", + "num_heads", + "attn_head_dim", + "num_layers", + "text_dim", + "freq_dim", +) + + +def default_video_dit_config(action_dim: int) -> dict[str, Any]: + return { + "patch_size": [1, 2, 2], + "in_dim": 48, + "hidden_dim": 3072, + "ffn_dim": 14336, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 48, + "num_heads": 24, + "attn_head_dim": 128, + "num_layers": 30, + "eps": 1.0e-6, + "seperated_timestep": True, + "use_gradient_checkpointing": False, + "video_attention_mask_mode": "first_frame_causal", + "action_conditioned": False, + "action_dim": action_dim, + "action_group_causal_mask_mode": "group_diagonal", + "fp32_attention": True, + } + + +def default_action_dit_config(action_dim: int) -> dict[str, Any]: + return { + "action_dim": action_dim, + "hidden_dim": 1024, + "ffn_dim": 4096, + "num_heads": 24, + "attn_head_dim": 128, + "num_layers": 30, + "text_dim": 4096, + "freq_dim": 256, + "eps": 1.0e-6, + "use_gradient_checkpointing": False, + "fp32_attention": True, + } + + +def _coerce_enum(enum_cls: type, value: Any) -> Any: + if isinstance(value, enum_cls): + return value + try: + return enum_cls(value) + except (TypeError, ValueError) as exc: + member = getattr(enum_cls, str(value), None) + if member is None: + raise ValueError(f"Cannot coerce {value!r} into {enum_cls.__name__}.") from exc + return member + + +def _coerce_policy_features(features: dict[str, Any] | None) -> dict[str, PolicyFeature] | None: + if features is None: + return None + coerced = {} + for name, feature in features.items(): + if isinstance(feature, PolicyFeature): + coerced[name] = feature + continue + coerced[name] = PolicyFeature( + type=_coerce_enum(FeatureType, feature["type"]), + shape=tuple(feature["shape"]), + ) + return coerced + + +def _is_local_model_id(value: str) -> bool: + path = Path(value).expanduser() + return path.is_absolute() or value.startswith(("./", "../", "~")) or path.exists() + + +def _validate_wan_model_id(value: str, field_name: str) -> str: + if value == WAN22_MODEL_ID or _is_local_model_id(value): + return value + raise ValueError(f"`{field_name}` must be `{WAN22_MODEL_ID}` or an explicit local path, got `{value}`.") + + +def is_fastwam_base_compatible_config(config: FastWAMConfig) -> bool: + """Return whether `fastwam_base` partial weights can initialize this config.""" + + default_video_config = default_video_dit_config(config.action_dim) + default_action_config = default_action_dit_config(config.action_dim) + return all( + config.video_dit_config.get(key) == default_video_config.get(key) + for key in _FASTWAM_VIDEO_BASE_COMPAT_KEYS + ) and all( + config.action_dit_config.get(key) == default_action_config.get(key) + for key in _FASTWAM_ACTION_BASE_COMPAT_KEYS + ) + + +@PreTrainedConfig.register_subclass("fastwam") +@dataclass +class FastWAMConfig(PreTrainedConfig): + """Configuration for the FastWAM LeRobot policy. + + Args: + action_dim (int): Number of scalar action channels per timestep. + proprio_dim (int | None): Number of proprioception channels used as an + extra text-context token. `None` disables proprio conditioning. + action_horizon (int): Number of actions predicted by one policy call. + num_video_frames (int): Raw video sampling window (in dataset frames). The + model actually operates on `model_video_frames` frames after subsampling + by `action_video_freq_ratio`. + action_video_freq_ratio (int): Actions are sampled at this multiple of the + video frame rate. Video frames are taken every `action_video_freq_ratio`-th + raw frame, so the model sees `(num_video_frames - 1) // ratio + 1` frames + spanning the same time window as `action_horizon` actions (ratio actions + per video frame). + image_size (tuple[int, int]): Concatenated image size as `(height, width)`. + context_len (int): Maximum text embedding token length. + video_dit_config (dict[str, Any] | None): Wan video expert config. + action_dit_config (dict[str, Any] | None): Action expert config. + use_gradient_checkpointing (bool): Enable activation checkpointing in both DiT + experts (trades compute for memory; propagated into the DiT configs). + freeze_video_expert (bool): Freeze the ~5B Wan video expert + (`model.video_expert`) so only the action expert + proprio encoder train. + Cuts the AdamW optimizer footprint substantially; the video expert keeps its + pretrained weights. (If enabled, also set `loss.lambda_video=0` to skip the + now-gradient-free video loss compute.) + """ + + n_obs_steps: int = 1 + action_dim: int = 7 + proprio_dim: int | None = 8 + action_horizon: int = 32 + n_action_steps: int = 32 + num_video_frames: int = 33 + action_video_freq_ratio: int = 4 + image_size: tuple[int, int] = (224, 448) + context_len: int = 128 + model_id: str = WAN22_MODEL_ID + tokenizer_model_id: str = WAN_T5_TOKENIZER_ID + text_encoder_model_id: str = WAN22_DIFFUSERS_MODEL_ID + base_model_id: str | None = FASTWAM_BASE_MODEL_ID + tokenizer_max_len: int = 128 + load_text_encoder: bool = True + mot_checkpoint_mixed_attn: bool = False + torch_dtype: str = "bfloat16" + prompt_template: str = ( + "A video recorded from a robot's point of view executing the following instruction: {task}" + ) + num_inference_steps: int = 10 + inference_seed: int | None = 42 + rand_device: str = "cpu" + text_cfg_scale: float = 1.0 + negative_prompt: str = "" + sigma_shift: float | None = None + tiled: bool = False + fp32_attention: bool = True + use_gradient_checkpointing: bool = False + freeze_video_expert: bool = False + toggle_action_dimensions: list[int] = field(default_factory=list) + video_scheduler: dict[str, float | int] = field( + default_factory=lambda: {"train_shift": 5.0, "infer_shift": 5.0, "num_train_timesteps": 1000} + ) + action_scheduler: dict[str, float | int] = field( + default_factory=lambda: {"train_shift": 5.0, "infer_shift": 5.0, "num_train_timesteps": 1000} + ) + loss: dict[str, float] = field(default_factory=lambda: {"lambda_video": 1.0, "lambda_action": 1.0}) + video_dit_config: dict[str, Any] | None = None + action_dit_config: dict[str, Any] | None = None + normalization_mapping: dict[str, NormalizationMode] = field( + default_factory=lambda: { + "VISUAL": NormalizationMode.IDENTITY, + "STATE": NormalizationMode.MEAN_STD, + "ACTION": NormalizationMode.MEAN_STD, + } + ) + input_features: dict[str, PolicyFeature] | None = None + output_features: dict[str, PolicyFeature] | None = None + optimizer_lr: float = 1.0e-4 + optimizer_weight_decay: float = 1.0e-2 + + def __post_init__(self) -> None: + super().__post_init__() + self.image_size = tuple(self.image_size) + self.model_id = _validate_wan_model_id(self.model_id, "model_id") + self.input_features = _coerce_policy_features(self.input_features) + self.output_features = _coerce_policy_features(self.output_features) + self.toggle_action_dimensions = [int(dim) for dim in self.toggle_action_dimensions] + self.video_dit_config = self.video_dit_config or default_video_dit_config(self.action_dim) + self.action_dit_config = self.action_dit_config or default_action_dit_config(self.action_dim) + self.video_dit_config["fp32_attention"] = bool(self.fp32_attention) + self.action_dit_config["fp32_attention"] = bool(self.fp32_attention) + self.video_dit_config["use_gradient_checkpointing"] = bool(self.use_gradient_checkpointing) + self.action_dit_config["use_gradient_checkpointing"] = bool(self.use_gradient_checkpointing) + if self.input_features is None: + height, width = self.image_size + self.input_features = { + "observation.images.image": PolicyFeature( + type=FeatureType.VISUAL, + shape=(3, height, width), + ) + } + if self.proprio_dim is not None: + self.input_features[OBS_STATE] = PolicyFeature( + type=FeatureType.STATE, + shape=(self.proprio_dim,), + ) + if self.output_features is None: + self.output_features = {ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(self.action_dim,))} + self.validate_features() + if self.pretrained_path or self.use_peft or not self.base_model_id: + return + if not is_fastwam_base_compatible_config(self): + return + self.pretrained_path = Path(self.base_model_id) + self._auto_pretrained_path = True + + def _save_pretrained(self, save_directory: Path) -> None: + if not getattr(self, "_auto_pretrained_path", False): + super()._save_pretrained(save_directory) + return + + pretrained_path = self.pretrained_path + self.pretrained_path = None + try: + super()._save_pretrained(save_directory) + finally: + self.pretrained_path = pretrained_path + + def get_optimizer_preset(self) -> AdamWConfig: + return AdamWConfig(lr=self.optimizer_lr, weight_decay=self.optimizer_weight_decay) + + def get_scheduler_preset(self) -> None: + return None + + def set_dataset_feature_metadata(self, dataset_features: dict[str, Any]) -> None: + """Rebuild visual input features from the dataset's real camera keys. + + FastWAM's `__post_init__` installs a synthetic single-image default + (`observation.images.image` at full `image_size` width). For datasets + with one or more separately-named cameras (e.g. `observation.images.top`, + `observation.images.wrist`), this hook — invoked by `make_policy` once the + dataset metadata is known — replaces that default with the actual camera + keys, each declared at the policy's native per-camera resolution + (`image_size[0]` x `image_size[1] // num_cameras`). The accompanying + resize step in `make_fastwam_pre_post_processors` resizes raw frames to + match, so heterogeneous source resolutions (e.g. 480x640) are supported. + """ + image_keys = sorted( + key + for key, feature in dataset_features.items() + if key.startswith("observation.images.") and feature.get("dtype") in ("video", "image") + ) + if not image_keys: + return + height, total_width = self.image_size + per_cam_width = total_width // len(image_keys) + new_inputs: dict[str, PolicyFeature] = { + key: PolicyFeature(type=FeatureType.VISUAL, shape=(3, height, per_cam_width)) + for key in image_keys + } + if self.proprio_dim is not None and OBS_STATE in dataset_features: + new_inputs[OBS_STATE] = PolicyFeature(type=FeatureType.STATE, shape=(self.proprio_dim,)) + self.input_features = new_inputs + self.validate_features() + + def validate_features(self) -> None: + if self.action_dim <= 0: + raise ValueError(f"`action_dim` must be positive, got {self.action_dim}.") + if self.action_horizon <= 0: + raise ValueError(f"`action_horizon` must be positive, got {self.action_horizon}.") + if self.n_action_steps > self.action_horizon: + raise ValueError("`n_action_steps` cannot exceed `action_horizon`.") + if self.action_video_freq_ratio <= 0: + raise ValueError( + f"`action_video_freq_ratio` must be positive, got {self.action_video_freq_ratio}." + ) + # Video frames are subsampled by action_video_freq_ratio; the resulting model frame + # count must satisfy T % 4 == 1 for the VAE temporal tokenization (mirrors the + # original FastWAM dataset asserts). + if (self.num_video_frames - 1) % self.action_video_freq_ratio != 0: + raise ValueError( + f"`num_video_frames - 1` ({self.num_video_frames - 1}) must be divisible by " + f"`action_video_freq_ratio` ({self.action_video_freq_ratio})." + ) + if ((self.num_video_frames - 1) // self.action_video_freq_ratio) % 4 != 0: + raise ValueError( + f"Subsampled video transitions ({(self.num_video_frames - 1) // self.action_video_freq_ratio}) " + "must be divisible by 4 for VAE tokenization (i.e. model_video_frames % 4 == 1)." + ) + if self.action_horizon % ((self.num_video_frames - 1) // self.action_video_freq_ratio) != 0: + raise ValueError( + f"`action_horizon` ({self.action_horizon}) must be divisible by the number of " + f"video transitions ({(self.num_video_frames - 1) // self.action_video_freq_ratio})." + ) + if not self.image_features: + raise ValueError("FastWAM requires at least one image feature.") + if self.action_feature is None: + raise ValueError("FastWAM requires `action` in output_features.") + action_shape = tuple(self.action_feature.shape) + if action_shape != (self.action_dim,): + raise ValueError( + f"FastWAM action feature shape must be ({self.action_dim},), got {action_shape}." + ) + if self.proprio_dim is not None: + state_feature = self.robot_state_feature + if state_feature is None: + raise ValueError("FastWAM requires `observation.state` when `proprio_dim` is set.") + state_shape = tuple(state_feature.shape) + if state_shape != (self.proprio_dim,): + raise ValueError( + f"FastWAM state feature shape must be ({self.proprio_dim},), got {state_shape}." + ) + height, width = self.image_size + image_width_sum = 0 + for name, feature in self.image_features.items(): + shape = tuple(feature.shape) + if len(shape) != 3 or shape[0] != 3: + raise ValueError(f"FastWAM image feature `{name}` must have shape (3, H, W), got {shape}.") + if shape[1] != height: + raise ValueError(f"FastWAM image feature `{name}` height must be {height}, got {shape[1]}.") + image_width_sum += shape[2] + if image_width_sum != width: + raise ValueError(f"FastWAM image feature widths must sum to {width}, got {image_width_sum}.") + + @property + def model_video_frames(self) -> int: + """Number of video frames the model actually operates on, after subsampling the + raw `num_video_frames` window by `action_video_freq_ratio` (e.g. 33 -> 9).""" + return (self.num_video_frames - 1) // self.action_video_freq_ratio + 1 + + @property + def observation_delta_indices(self) -> list[int]: + # Load the video frames the model is supervised on: the future window subsampled by + # action_video_freq_ratio (e.g. [0, 4, 8, ..., 32] -> 9 frames). Each video frame is + # thus `action_video_freq_ratio` actions apart, while actions load at the full rate + # (`action_delta_indices` = range(action_horizon)). Returning None would load only the + # current frame, making the video target a static repeat (degenerate supervision). + return list(range(0, self.num_video_frames, self.action_video_freq_ratio)) + + @property + def action_delta_indices(self) -> list[int]: + return list(range(self.action_horizon)) + + @property + def reward_delta_indices(self) -> None: + return None diff --git a/src/lerobot/policies/fastwam/modeling_fastwam.py b/src/lerobot/policies/fastwam/modeling_fastwam.py new file mode 100644 index 000000000..10671e717 --- /dev/null +++ b/src/lerobot/policies/fastwam/modeling_fastwam.py @@ -0,0 +1,440 @@ +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import logging +from collections import deque +from typing import Any + +import torch +from torch import Tensor + +from lerobot.policies.pretrained import PreTrainedPolicy +from lerobot.utils.constants import OBS_STATE +from lerobot.utils.import_utils import require_package + +from .configuration_fastwam import FastWAMConfig +from .wan import ( + ActionDiT, + FastWAM, + MoT, + WanVideoDiT, + build_wan_tokenizer, + load_pretrained_wan_text_encoder, + load_pretrained_wan_vae, +) + + +class FastWAMPolicy(PreTrainedPolicy): + """LeRobot policy wrapper for FastWAM. + + Attention backend: FastWAM's DiT uses ``torch.nn.functional.scaled_dot_product_attention`` + (SDPA) for all attention. It does not use FlashAttention, because MoT routing requires + arbitrary boolean ``[query, key]`` masks that the FlashAttention varlen API cannot express; + installing ``flash-attn`` has no effect on the FastWAM path. (SDPA may still dispatch to + PyTorch's own flash/mem-efficient/math kernel internally, unrelated to the ``flash-attn`` package.) + + Args: + config (FastWAMConfig): FastWAM policy configuration. + dataset_stats (dict[str, dict[str, Tensor]] | None): Optional LeRobot + dataset statistics passed by the training/evaluation stack. + """ + + config_class = FastWAMConfig + name = "fastwam" + + def __init__( + self, + config: FastWAMConfig, + dataset_stats: dict[str, dict[str, Tensor]] | None = None, + **kwargs: Any, + ): + # FastWAM's Wan2.2 backbone needs transformers (UMT5 text encoder/tokenizer) and + # diffusers (Wan VAE), both behind the `fastwam` extra. Fail fast with an actionable + # message in base installs rather than deep in Wan component construction. + require_package("transformers", extra="fastwam") + require_package("diffusers", extra="fastwam") + # `make_policy`/`from_pretrained` forward extra kwargs (e.g. `dataset_meta`); the + # dataset feature metadata is already applied to `config` by make_policy upstream, + # so we accept and ignore them, matching the other LeRobot policies. + super().__init__(config, dataset_stats) + config.validate_features() + self.config = config + self.dataset_stats = dataset_stats + self.model = self._build_core_model(config) + if config.freeze_video_expert and getattr(self.model, "video_expert", None) is not None: + # Freeze the ~5B Wan video expert; get_optim_params filters on requires_grad, + # so its params drop out of the optimizer (and DDP skips them). + self.model.video_expert.requires_grad_(False) + # The transformer blocks are re-parented onto the MoTLayers (single FSDP owner), so + # `video_expert.requires_grad_` no longer reaches them — freeze them via the layers. + mot = getattr(self.model, "mot", None) + if mot is not None and getattr(mot, "layers", None) is not None: + for layer in mot.layers: + if "video" in layer.blocks: + layer.blocks["video"].requires_grad_(False) + self.reset() + + @classmethod + def _load_as_safetensor(cls, model, model_file: str, map_location: str, strict: bool): + """Shape-aware load that supports cross-embodiment fine-tuning. + + `safetensors.load_model(strict=False)` ignores missing/unexpected keys but + still raises on a shape mismatch for a shared key. When fine-tuning from a + checkpoint trained on a different embodiment (e.g. the LIBERO 7-DoF / 8-dim + checkpoint adapted to a 6-DoF / 6-dim arm), the action encoder/head and + proprio encoder legitimately differ in shape. With `strict=False` we drop + only those shape-mismatched tensors — leaving them at their freshly + initialized values — and load every compatible tensor. With `strict=True` + the standard exact-match loader is used. + """ + from safetensors import safe_open + + model_state_dict = model.state_dict() + mismatched = [] + with safe_open(model_file, framework="pt") as f: + checkpoint_keys = list(f.keys()) + for key in checkpoint_keys: + if key in model_state_dict and tuple(model_state_dict[key].shape) != tuple( + f.get_slice(key).get_shape() + ): + mismatched.append(key) + + if not mismatched: + return super()._load_as_safetensor(model, model_file, map_location, strict) + if strict: + raise RuntimeError( + f"FastWAM: {len(mismatched)} checkpoint tensors have a shape mismatch under " + f"strict=True: {mismatched}" + ) + + from safetensors.torch import load_file + + logging.warning( + "FastWAM cross-embodiment load: reinitializing %d shape-mismatched tensor(s), keeping " + "every compatible weight: %s", + len(mismatched), + mismatched, + ) + state_dict = load_file(model_file, device="cpu") + for key in mismatched: + state_dict.pop(key, None) + model.load_state_dict(state_dict, strict=False) + if map_location and map_location != "cpu": + model.to(map_location) + return model + + def get_optim_params(self) -> list[Tensor]: + # Return the trainable tensors directly (a single param group). The optimizer + # builder wraps these in a param group; returning a bare {"params": [...]} dict + # instead would make `list(...)` yield the key string "params". + params = ( + list(self.model.dit.parameters()) if hasattr(self.model, "dit") else list(self.model.parameters()) + ) + proprio_encoder = getattr(self.model, "proprio_encoder", None) + if proprio_encoder is not None: + params.extend(list(proprio_encoder.parameters())) + return [p for p in params if p.requires_grad] + + def reset(self) -> None: + self._action_queue: deque[Tensor] = deque([], maxlen=self.config.n_action_steps) + + def _batch_to_training_sample(self, batch: dict[str, Tensor]) -> dict[str, Tensor]: + """Adapt a standard LeRobot batch to the FastWAM-native sample that + `FastWAM.build_inputs` consumes (`video`, `action`, `context`/`context_mask`, + per-frame `proprio`). + + The LeRobot training loop passes raw `observation.images.*`, a single-step + `observation.state` `[B, D]`, `action`, and a language `task` string. We do + only the translation `build_inputs` can't: stack the camera frames into a + video, encode the prompt with the (frozen) text encoder (mirroring inference, + so language-conditioned datasets need no precomputed context), and give proprio + the per-frame axis `build_inputs` indexes. All shape/presence validation is + left to `build_inputs`, the single authority on the contract. + """ + sample = dict(batch) + if "video" not in sample: + sample["video"] = _stack_video_from_images(batch, self.config) + if "context" not in sample or "context_mask" not in sample: + prompt = _prompt_from_batch(batch=batch, config=self.config) + if prompt is None: + raise KeyError( + "FastWAM training requires a `task`/`prompt` to encode text context, " + "or precomputed `context`/`context_mask` in the batch." + ) + sample["context"], sample["context_mask"] = self.model.encode_prompt(prompt) + if self.config.proprio_dim is not None and "proprio" not in sample: + state = sample.get(OBS_STATE) + if state is not None: + # LeRobot gives a single-step state [B, D]; build_inputs expects + # per-frame [B, T, D] and uses frame 0, so add a T=1 axis. + sample["proprio"] = state.unsqueeze(1) if state.ndim == 2 else state + return sample + + def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict[str, Any]]: + """Compute FastWAM training loss for a LeRobot batch. + + Args: + batch (dict[str, Tensor]): Batch containing FastWAM-ready keys + (`video`, `action`, `context`, `context_mask`) or LeRobot keys + that can be adapted (`observation.images.*`, `observation.state`, + `action`, `action_is_pad`). + + Returns: + tuple[Tensor, dict[str, Any]]: The scalar loss to backprop, and a dict of + logging metrics (e.g. `loss_video`, `loss_action`) — the `(loss, output_dict)` + contract the LeRobot training loop expects. + """ + + sample = self._batch_to_training_sample(batch) + loss, metrics = self.model.training_loss(sample) + return loss, dict(metrics or {}) + + @torch.no_grad() + def predict_action_chunk(self, batch: dict[str, Tensor], **_: Any) -> Tensor: + """Predict a chunk of actions from the current FastWAM observation. + + Args: + batch (dict[str, Tensor]): Inference batch with `input_image` or + image observation keys, plus `context/context_mask` or `prompt`. + + Returns: + Tensor: Action chunk with shape `[B, action_horizon, action_dim]`. + """ + + self.eval() + infer_kwargs = _batch_to_infer_kwargs(batch=batch, config=self.config) + batch_size = _infer_kwargs_batch_size(infer_kwargs) + if batch_size == 1: + action = _action_from_model_output(self.model.infer_action(**infer_kwargs)) + else: + action = torch.cat( + [ + _action_from_model_output( + self.model.infer_action( + **_slice_infer_kwargs(infer_kwargs, index=i, batch_size=batch_size) + ) + ) + for i in range(batch_size) + ], + dim=0, + ) + return action.to(device=batch_device(batch), dtype=torch.float32) + + @torch.no_grad() + def select_action(self, batch: dict[str, Tensor], **kwargs: Any) -> Tensor: + self.eval() + if len(self._action_queue) == 0: + actions = self.predict_action_chunk(batch, **kwargs)[:, : self.config.n_action_steps] + self._action_queue.extend(actions.transpose(0, 1)) + return self._action_queue.popleft() + + def _build_core_model(self, config: FastWAMConfig) -> FastWAM: + """Build the FastWAM core for training / inference. + + Only the trainable parts (the MoT DiT and the proprio encoder) are + materialized empty here and then filled from the policy's + `model.safetensors` by the base `from_pretrained`. The *frozen* Wan2.2 VAE + and UMT5 text encoder are loaded with their real weights from the + `Wan-AI/Wan2.2-TI2V-5B-Diffusers` repo (cached in the HF cache, shared + across checkpoints) and are intentionally excluded from `model.safetensors` + — see `FastWAM.__init__`. The tokenizer comes from `google/umt5-xxl`. + """ + dtype = _dtype_from_name(config.torch_dtype) + device = config.device + video_expert = WanVideoDiT(**config.video_dit_config).to(device=device, dtype=dtype) + action_expert = ActionDiT(**config.action_dit_config).to(device=device, dtype=dtype) + mot = MoT( + mixtures={"video": video_expert, "action": action_expert}, + mot_checkpoint_mixed_attn=config.mot_checkpoint_mixed_attn, + ) + text_encoder = ( + load_pretrained_wan_text_encoder( + model_id=config.text_encoder_model_id, torch_dtype=dtype, device=device + ) + if config.load_text_encoder + else None + ) + return FastWAM( + video_expert=video_expert, + action_expert=action_expert, + mot=mot, + vae=load_pretrained_wan_vae(torch_dtype=dtype, device=device), + text_encoder=text_encoder, + tokenizer=build_wan_tokenizer( + model_id=config.tokenizer_model_id, tokenizer_max_len=config.tokenizer_max_len + ), + text_dim=int(config.video_dit_config["text_dim"]), + proprio_dim=config.proprio_dim, + device=device, + torch_dtype=dtype, + video_train_shift=float(config.video_scheduler["train_shift"]), + video_infer_shift=float(config.video_scheduler["infer_shift"]), + video_num_train_timesteps=int(config.video_scheduler["num_train_timesteps"]), + action_train_shift=float(config.action_scheduler["train_shift"]), + action_infer_shift=float(config.action_scheduler["infer_shift"]), + action_num_train_timesteps=int(config.action_scheduler["num_train_timesteps"]), + loss_lambda_video=float(config.loss["lambda_video"]), + loss_lambda_action=float(config.loss["lambda_action"]), + ) + + +def _scalar(value: Any) -> Any: + """Unwrap a 0-/1-element tensor (e.g. from DataLoader collation) to a Python scalar.""" + return value.item() if isinstance(value, Tensor) else value + + +def _batch_to_infer_kwargs(batch: dict[str, Tensor], config: FastWAMConfig) -> dict[str, Any]: + return { + "prompt": _prompt_from_batch(batch=batch, config=config), + "input_image": _input_image_from_batch(batch, config), + "action_horizon": config.action_horizon, + "proprio": batch.get("proprio", batch.get(OBS_STATE)), + "context": batch.get("context"), + "context_mask": batch.get("context_mask"), + "negative_prompt": batch.get("negative_prompt", config.negative_prompt), + "text_cfg_scale": float(_scalar(batch.get("text_cfg_scale", config.text_cfg_scale))), + "num_inference_steps": int(_scalar(batch.get("num_inference_steps", config.num_inference_steps))), + "sigma_shift": batch.get("sigma_shift", config.sigma_shift), + "seed": batch.get("seed", config.inference_seed), + "rand_device": batch.get("rand_device", config.rand_device), + "tiled": bool(batch.get("tiled", config.tiled)), + } + + +def _prompt_from_batch(batch: dict[str, Tensor], config: FastWAMConfig) -> Any: + prompt = batch.get("prompt") + if prompt is not None: + return prompt + + task = batch.get("task") + if task is None: + return None + if isinstance(task, str): + return config.prompt_template.format(task=task) + if isinstance(task, (list, tuple)): + return [config.prompt_template.format(task=str(item)) for item in task] + return config.prompt_template.format(task=str(task)) + + +def _action_from_model_output(output: Any) -> Tensor: + action = output["action"] if isinstance(output, dict) else output + if action.ndim == 2: + action = action.unsqueeze(0) + return action + + +def _infer_kwargs_batch_size(infer_kwargs: dict[str, Any]) -> int: + image = infer_kwargs["input_image"] + if not isinstance(image, Tensor): + raise TypeError(f"`input_image` must be a tensor, got {type(image).__name__}.") + if image.ndim == 3: + return 1 + if image.ndim == 4: + return int(image.shape[0]) + raise ValueError(f"`input_image` must be [B,C,H,W] or [C,H,W], got {tuple(image.shape)}.") + + +def _slice_infer_kwargs(infer_kwargs: dict[str, Any], *, index: int, batch_size: int) -> dict[str, Any]: + return { + key: _slice_infer_value(value, index=index, batch_size=batch_size) + for key, value in infer_kwargs.items() + } + + +def _slice_infer_value(value: Any, *, index: int, batch_size: int) -> Any: + if isinstance(value, Tensor) and value.ndim > 0 and value.shape[0] == batch_size: + return value[index : index + 1] + if isinstance(value, (list, tuple)) and len(value) == batch_size: + return value[index] + return value + + +def _dtype_from_name(name: str) -> torch.dtype: + dtype_map = {"float32": torch.float32, "float16": torch.float16, "bfloat16": torch.bfloat16} + if name not in dtype_map: + raise ValueError(f"Unsupported torch dtype `{name}`.") + return dtype_map[name] + + +def batch_device(batch: dict[str, Any]) -> torch.device: + for value in batch.values(): + if isinstance(value, Tensor): + return value.device + return torch.device("cpu") + + +def _resize_frames(frames: Tensor, size: tuple[int, int]) -> Tensor: + """Resize a frame tensor to `size` (H, W), tolerating a leading temporal/batch stack. + + `interpolate` only accepts a single leading batch dim (`[N, C, H, W]`), but FastWAM camera + tensors arrive as `[B, C, H, W]` (live eval) or `[B, T, C, H, W]` (temporal stack), so flatten + any leading dims into the batch, resize, then restore. A no-op when already at `size`. + """ + if tuple(frames.shape[-2:]) == size: + return frames + lead = frames.shape[:-3] + flat = frames.reshape(-1, *frames.shape[-3:]) + flat = torch.nn.functional.interpolate( + flat, size=size, mode="bilinear", align_corners=False, antialias=True + ) + return flat.reshape(*lead, *flat.shape[-3:]) + + +def _stack_video_from_images(batch: dict[str, Tensor], config: FastWAMConfig) -> Tensor: + # Exclude the `*_is_pad` companion tensors that delta-timestamp loading adds alongside + # each camera (shape [B, T]); they share the `observation.images.` prefix but are not frames. + image_keys = sorted(k for k in batch if k.startswith("observation.images.") and not k.endswith("_is_pad")) + if not image_keys: + raise KeyError("FastWAM batch must contain `video` or `observation.images.*` keys.") + per_cam = (int(config.image_size[0]), int(config.image_size[1]) // len(image_keys)) + images = [_resize_frames(batch[key], per_cam) for key in image_keys] + # Cameras concatenate along width (last dim) in both the single-frame and temporal case. + image = torch.cat(images, dim=-1) if len(images) > 1 else images[0] + if image.ndim == 4: + # [B, C, H, W]: a single frame (e.g. the live eval observation) -> repeat across time. + image = image.unsqueeze(2).repeat(1, 1, config.model_video_frames, 1, 1) + elif image.ndim == 5: + # [B, T, C, H, W]: temporal stack from delta-timestamp loading -> [B, C, T, H, W]. + image = image.permute(0, 2, 1, 3, 4) + else: + raise ValueError(f"Expected image batch [B,C,H,W] or temporal [B,T,C,H,W], got {tuple(image.shape)}.") + return image + + +def _input_image_from_batch(batch: dict[str, Tensor], config: FastWAMConfig) -> Tensor: + if "input_image" in batch: + return _prepare_infer_image(batch["input_image"], config) + video = batch.get("video") + if video is None: + video = _stack_video_from_images(batch, config) + if video.ndim == 5: + return _prepare_infer_image(video[:, :, 0], config) + if video.ndim == 4: + return _prepare_infer_image(video, config) + raise ValueError(f"Cannot build input image from tensor with shape {tuple(video.shape)}.") + + +def _prepare_infer_image(image: Tensor, config: FastWAMConfig) -> Tensor: + if image.ndim == 3: + image = image.unsqueeze(0) + if image.ndim != 4: + raise ValueError(f"Expected image tensor [B,C,H,W] or [C,H,W], got {tuple(image.shape)}.") + + # Resize to the full configured resolution (no-op when the video path already produced it, but + # also covers a directly-supplied `input_image`). The model owns its input resolution — see + # `_stack_video_from_images` — so we resize rather than assert on a mismatch. + target_h, target_w = int(config.image_size[0]), int(config.image_size[1]) + return _resize_frames(image, (target_h, target_w)) diff --git a/src/lerobot/policies/fastwam/processor_fastwam.py b/src/lerobot/policies/fastwam/processor_fastwam.py new file mode 100644 index 000000000..31f3b9277 --- /dev/null +++ b/src/lerobot/policies/fastwam/processor_fastwam.py @@ -0,0 +1,142 @@ +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import torch + +from lerobot.configs import PipelineFeatureType, PolicyFeature +from lerobot.processor import ( + ActionProcessorStep, + AddBatchDimensionProcessorStep, + DeviceProcessorStep, + NormalizerProcessorStep, + PolicyAction, + PolicyProcessorPipeline, + ProcessorStepRegistry, + RenameObservationsProcessorStep, + UnnormalizerProcessorStep, + policy_action_to_transition, + transition_to_policy_action, +) +from lerobot.utils.constants import ( + POLICY_POSTPROCESSOR_DEFAULT_NAME, + POLICY_PREPROCESSOR_DEFAULT_NAME, +) + +from .configuration_fastwam import FastWAMConfig + + +@dataclass +@ProcessorStepRegistry.register(name="fastwam_action_toggle_processor") +class FastWAMActionToggleProcessorStep(ActionProcessorStep): + """Apply FastWAM LIBERO toggle semantics to configured action dimensions.""" + + toggle_dimensions: list[int] + + def action(self, action: PolicyAction) -> PolicyAction: + if not self.toggle_dimensions: + return action + processed_action = action.clone() + action_dim = int(processed_action.shape[-1]) + for dim in self.toggle_dimensions: + resolved_dim = dim if dim >= 0 else action_dim + dim + if resolved_dim < 0 or resolved_dim >= action_dim: + raise ValueError( + f"FastWAM action toggle dimension {dim} is out of bounds for action dim {action_dim}." + ) + value = processed_action[..., resolved_dim] + value = value * 2.0 - 1.0 + processed_action[..., resolved_dim] = torch.sign(-value) + return processed_action + + def get_config(self) -> dict[str, Any]: + return {"toggle_dimensions": self.toggle_dimensions} + + def transform_features( + self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]] + ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]: + return features + + +def make_fastwam_pre_post_processors( + config: FastWAMConfig, + dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None, +) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]: + """Create LeRobot pre- and post-processing pipelines for FastWAM. + + Args: + config (FastWAMConfig): Policy configuration controlling device and + normalization feature metadata. + dataset_stats (dict[str, dict[str, torch.Tensor]] | None): Optional + LeRobot dataset statistics used by normalization processors. + + Returns: + tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]: Input and + output processor pipelines discoverable by LeRobot. + """ + + # NOTE: no visual normalization here. VISUAL is IDENTITY (see configuration_fastwam.normalization_mapping) + # — images pass through in [0, 1] and the model maps them to the Wan VAE's [-1, 1] at the encode + # boundary. This is deliberate: `lerobot_train.py` overrides the normalizer stats with + # `dataset.meta.stats` when fine-tuning, and a real dataset's per-channel image std is the tiny + # frame-to-frame brightness variance, which would blow images far outside [-1,1] and saturate them. + # STATE/ACTION still normalize with dataset stats below. + normalization_stats: dict[str, dict[str, Any]] = dict(dataset_stats or {}) + + # NOTE: no resize step here. The model is the single authority on input resolution: it resizes + # each camera to the per-camera target (image_size split across cameras) in + # `_stack_video_from_images` / `_prepare_infer_image`, on every path (train forward, rollout and + # eval select_action). A preprocessor resize step would be both redundant (the model re-resizes + # anyway) and unsafe across fine-tuning: its `resize_size` would be inherited from the base + # checkpoint's camera geometry, not this dataset's, making the concatenation N_cameras x too wide. + + input_steps = [ + RenameObservationsProcessorStep(rename_map={}), + AddBatchDimensionProcessorStep(), + DeviceProcessorStep(device=config.device), + NormalizerProcessorStep( + features={**config.input_features, **config.output_features}, + norm_map=config.normalization_mapping, + stats=normalization_stats, + device=config.device, + ), + ] + output_steps = [ + UnnormalizerProcessorStep( + features=config.output_features, + norm_map=config.normalization_mapping, + stats=normalization_stats, + ), + ] + if config.toggle_action_dimensions: + output_steps.append( + FastWAMActionToggleProcessorStep(toggle_dimensions=config.toggle_action_dimensions) + ) + output_steps.append(DeviceProcessorStep(device="cpu")) + return ( + PolicyProcessorPipeline[dict[str, Any], dict[str, Any]]( + steps=input_steps, + name=POLICY_PREPROCESSOR_DEFAULT_NAME, + ), + PolicyProcessorPipeline[PolicyAction, PolicyAction]( + steps=output_steps, + name=POLICY_POSTPROCESSOR_DEFAULT_NAME, + to_transition=policy_action_to_transition, + to_output=transition_to_policy_action, + ), + ) diff --git a/src/lerobot/policies/fastwam/wan/README.md b/src/lerobot/policies/fastwam/wan/README.md new file mode 100644 index 000000000..7b7d61033 --- /dev/null +++ b/src/lerobot/policies/fastwam/wan/README.md @@ -0,0 +1,34 @@ +# FastWAM `wan` package + +This package holds FastWAM's model implementation. It mixes a small **vendored +subset of the official Wan2.2 source tree** with FastWAM's own code, kept flat in +a single directory. + +## Vendored from Wan2.2 + +- Upstream repository: https://github.com/Wan-Video/Wan2.2 +- Upstream commit: `42bf4cfaa384bc21833865abc2f9e6c0e67233dc` +- License: Apache-2.0, matching the license in `LICENSE.txt` from the upstream repository + +Copied files: + +- `model.py` (was `wan/modules/model.py`), trimmed: the flash-attention path + (the vendored `attention.py` and the block/model `forward`s) was removed. + FastWAM's DiT uses SDPA instead (see `video_dit.py`). +- `get_sampling_sigmas` in `video_dit.py` (was `wan/utils/fm_solvers.py`), inlined + next to its only caller. + +This subset only backs FastWAM's **custom MoT video DiT**. The Wan2.2 VAE, +UMT5 text encoder, and tokenizer are no longer vendored - they come from +`diffusers.AutoencoderKLWan`, `transformers.UMT5EncoderModel`, and +`transformers.AutoTokenizer` (see `components.py` and `adapters.py`). + +## FastWAM's own code + +- `video_dit.py` builds on `model` (`sinusoidal_embedding_1d`, `rope_params`, + `rope_apply`, …) and computes attention with SDPA (`fastwam_masked_attention`). Its + `WanContinuousFlowMatchScheduler` uses `get_sampling_sigmas` for Wan-compatible + inference timesteps. +- `components.py` / `adapters.py` load the VAE, text encoder, tokenizer, and the + custom DiT weights. +- `modular.py` defines the FastWAM model (`ActionDiT`, `MoT`, `FastWAM`, …). diff --git a/src/lerobot/policies/fastwam/wan/__init__.py b/src/lerobot/policies/fastwam/wan/__init__.py new file mode 100644 index 000000000..11c0aaf6f --- /dev/null +++ b/src/lerobot/policies/fastwam/wan/__init__.py @@ -0,0 +1,33 @@ +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .adapters import WanVideoVAE38 +from .components import ( + build_wan_tokenizer, + load_pretrained_wan_text_encoder, + load_pretrained_wan_vae, +) +from .modular import ActionDiT, FastWAM, MoT +from .video_dit import WanVideoDiT + +__all__ = [ + "ActionDiT", + "FastWAM", + "MoT", + "WanVideoDiT", + "WanVideoVAE38", + "build_wan_tokenizer", + "load_pretrained_wan_text_encoder", + "load_pretrained_wan_vae", +] diff --git a/src/lerobot/policies/fastwam/wan/adapters.py b/src/lerobot/policies/fastwam/wan/adapters.py new file mode 100644 index 000000000..a2e8a1068 --- /dev/null +++ b/src/lerobot/policies/fastwam/wan/adapters.py @@ -0,0 +1,108 @@ +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from diffusers import AutoencoderKLWan + + +class WanVideoVAE38(torch.nn.Module): + """FastWAM VAE contract over `diffusers.AutoencoderKLWan` (Wan2.2-TI2V-5B). + + 16x spatial / 4x temporal compression, 48 latent channels. diffusers' + `AutoencoderKLWan` returns *raw* latents (it does not apply `latents_mean`/ + `latents_std`), so `encode`/`decode` here apply the same standardization the + Wan reference uses — `(latents - mean) / std` — done in fp32 for stability. + `encode` uses the deterministic posterior mode, matching the original VAE + which returned the latent mean `mu`. + """ + + upsampling_factor = 16 + temporal_downsample_factor = 4 + z_dim = 48 + + def __init__( + self, + dtype: torch.dtype = torch.float32, + device: str | torch.device = "cuda", + *, + pretrained: AutoencoderKLWan, + ) -> None: + super().__init__() + # The Wan2.2 VAE is a fixed pretrained model — it is never trained from scratch, + # so a real `AutoencoderKLWan` (with weights) must always be supplied (loaded from + # the diffusers repo by `load_pretrained_wan_vae`). No random/offline build path. + self.vae = pretrained.to(device=device, dtype=dtype) + + # Read the standardization stats from the VAE's own config (diffusers populates + # these from vae/config.json) — single source of truth, no local copy. diffusers' + # encode/decode return *raw* latents, so we apply (latent - mean) / std ourselves. + # Non-persistent: kept out of state_dict. + self.register_buffer( + "latents_mean", + torch.tensor(self.vae.config.latents_mean).view(1, self.z_dim, 1, 1, 1), + persistent=False, + ) + self.register_buffer( + "latents_std", + torch.tensor(self.vae.config.latents_std).view(1, self.z_dim, 1, 1, 1), + persistent=False, + ) + + def _device_dtype(self) -> tuple[torch.device, torch.dtype]: + param = next(self.vae.parameters()) + return param.device, param.dtype + + def encode( + self, + videos: list[torch.Tensor] | torch.Tensor, + device: str | torch.device | None = None, + tiled: bool = False, + tile_size: tuple[int, int] = (34, 34), + tile_stride: tuple[int, int] = (18, 16), + ) -> torch.Tensor: + del device, tile_size, tile_stride + if tiled: + raise NotImplementedError("Tiled Wan2.2 VAE encoding is not supported by the FastWAM adapter.") + if isinstance(videos, (list, tuple)): + videos = torch.stack(list(videos)) + dev, dtype = self._device_dtype() + mu = self.vae.encode(videos.to(device=dev, dtype=dtype)).latent_dist.mode().float() + mean = self.latents_mean.float().to(mu.device) + std = self.latents_std.float().to(mu.device) + return (mu - mean) / std + + def decode( + self, + hidden_states: list[torch.Tensor] | torch.Tensor, + device: str | torch.device | None = None, + tiled: bool = False, + tile_size: tuple[int, int] = (34, 34), + tile_stride: tuple[int, int] = (18, 16), + ) -> torch.Tensor: + del device, tile_size, tile_stride + if tiled: + raise NotImplementedError("Tiled Wan2.2 VAE decoding is not supported by the FastWAM adapter.") + if isinstance(hidden_states, (list, tuple)): + hidden_states = torch.stack(list(hidden_states)) + dev, dtype = self._device_dtype() + z = hidden_states.float() + z = z * self.latents_std.float().to(z.device) + self.latents_mean.float().to(z.device) + out = self.vae.decode(z.to(device=dev, dtype=dtype)).sample + return out.float().clamp_(-1.0, 1.0) diff --git a/src/lerobot/policies/fastwam/wan/components.py b/src/lerobot/policies/fastwam/wan/components.py new file mode 100644 index 000000000..59a9b610e --- /dev/null +++ b/src/lerobot/policies/fastwam/wan/components.py @@ -0,0 +1,175 @@ +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import logging +from collections.abc import Sequence +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import torch +from huggingface_hub import snapshot_download +from safetensors.torch import load_file + +from lerobot.utils.import_utils import _diffusers_available, _transformers_available, require_package + +if TYPE_CHECKING or _transformers_available: + from transformers import AutoTokenizer, UMT5EncoderModel +else: + AutoTokenizer = None + UMT5EncoderModel = None + +if TYPE_CHECKING or _diffusers_available: + from diffusers import AutoencoderKLWan +else: + AutoencoderKLWan = None + +from .adapters import WanVideoVAE38 +from .video_dit import WanVideoDiT + +logger = logging.getLogger(__name__) + +# The custom MoT video DiT still ships in the original (non-diffusers) Wan2.2 +# repo as sharded `diffusion_pytorch_model*.safetensors`; the VAE and UMT5 text +# encoder come from the diffusers conversion. Tokenizer is the stock UMT5 one. +WAN_DIT_PATTERN = "diffusion_pytorch_model*.safetensors" +WAN_T5_TOKENIZER = "google/umt5-xxl" +WAN22_DIFFUSERS_MODEL_ID = "Wan-AI/Wan2.2-TI2V-5B-Diffusers" + + +class WanTextEncoder(torch.nn.Module): + """FastWAM text-encoder contract over `transformers.UMT5EncoderModel`. + + Exposes `.dim` (hidden size) and `forward(ids, mask) -> [B, L, dim]`, matching + the call in `FastWAM.encode_prompt`. + """ + + def __init__( + self, + dtype: torch.dtype = torch.bfloat16, + device: str | torch.device = "cuda", + *, + pretrained: torch.nn.Module, + ) -> None: + super().__init__() + # UMT5-XXL is a fixed pretrained encoder — never trained from scratch, so a real + # `UMT5EncoderModel` (with weights) must always be supplied (loaded from the + # diffusers repo by `load_pretrained_wan_text_encoder`). No random/offline build. + self.model = pretrained.to(device=device, dtype=dtype) + self.dim = int(self.model.config.d_model) + + def forward(self, ids: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + return self.model(input_ids=ids, attention_mask=mask.long()).last_hidden_state + + +class WanTokenizer: + """UMT5 tokenizer wrapper returning `(input_ids, attention_mask)` like the + FastWAM call site expects.""" + + def __init__(self, name: str = WAN_T5_TOKENIZER, seq_len: int = 512) -> None: + require_package("transformers", extra="fastwam") + self.tokenizer = AutoTokenizer.from_pretrained(name) + self.seq_len = int(seq_len) + + def __call__( + self, + sequence: str | Sequence[str], + return_mask: bool = False, + add_special_tokens: bool = True, + **_: Any, + ): + if isinstance(sequence, str): + sequence = [sequence] + out = self.tokenizer( + list(sequence), + padding="max_length", + truncation=True, + max_length=self.seq_len, + add_special_tokens=add_special_tokens, + return_tensors="pt", + ) + if return_mask: + return out.input_ids, out.attention_mask + return out.input_ids + + +def build_wan_tokenizer(*, model_id: str = WAN_T5_TOKENIZER, tokenizer_max_len: int) -> WanTokenizer: + return WanTokenizer(name=model_id, seq_len=int(tokenizer_max_len)) + + +def load_pretrained_wan_vae(*, torch_dtype: torch.dtype, device: str) -> WanVideoVAE38: + """Load real Wan2.2 VAE weights from the diffusers repo (offline base creation).""" + require_package("diffusers", extra="fastwam") + vae = AutoencoderKLWan.from_pretrained(WAN22_DIFFUSERS_MODEL_ID, subfolder="vae", torch_dtype=torch_dtype) + return WanVideoVAE38(dtype=torch_dtype, device=device, pretrained=vae) + + +def load_pretrained_wan_text_encoder( + *, + model_id: str = WAN22_DIFFUSERS_MODEL_ID, + subfolder: str | None = "text_encoder", + torch_dtype: torch.dtype, + device: str, +) -> WanTextEncoder: + """Load UMT5-XXL encoder weights (defaults to the Wan2.2 diffusers repo). + + Must stay compatible with the tokenizer (see `build_wan_tokenizer`): the encoder's + embedding table is indexed by the tokenizer's vocabulary. + """ + require_package("transformers", extra="fastwam") + encoder = UMT5EncoderModel.from_pretrained(model_id, subfolder=subfolder, torch_dtype=torch_dtype) + return WanTextEncoder(dtype=torch_dtype, device=device, pretrained=encoder) + + +def resolve_wan_dit_paths( + model_id_or_path: str | Path, + *, + cache_dir: str | Path | None = None, + local_files_only: bool = False, + revision: str | None = None, +) -> list[Path]: + """Resolve the custom MoT DiT shards from the original Wan2.2 repo or a local dir.""" + path = Path(model_id_or_path).expanduser() + if path.is_dir(): + return sorted(path.glob(WAN_DIT_PATTERN)) + + snapshot_path = snapshot_download( + repo_id=str(model_id_or_path), + revision=revision, + cache_dir=cache_dir, + local_files_only=local_files_only, + allow_patterns=[WAN_DIT_PATTERN], + ) + return sorted(Path(snapshot_path).glob(WAN_DIT_PATTERN)) + + +def load_wan_video_dit( + paths: list[str | Path], + *, + dit_config: dict[str, Any], + torch_dtype: torch.dtype, + device: str, +) -> WanVideoDiT: + model = WanVideoDiT(**dit_config) + state_dict = _read_wan_dit_safetensors(paths) + model.load_state_dict(state_dict, strict=False) + return model.to(device=device, dtype=torch_dtype) + + +def _read_wan_dit_safetensors(paths: list[str | Path]) -> dict[str, torch.Tensor]: + state_dict = {} + for path in paths: + state_dict.update(load_file(str(path), device="cpu")) + return state_dict diff --git a/src/lerobot/policies/fastwam/wan/model.py b/src/lerobot/policies/fastwam/wan/model.py new file mode 100644 index 000000000..329d1c48a --- /dev/null +++ b/src/lerobot/policies/fastwam/wan/model.py @@ -0,0 +1,341 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +import math + +import torch +import torch.nn as nn + + +def sinusoidal_embedding_1d(dim, position): + # preprocess + if dim % 2 != 0: + raise ValueError(f"dim must be even, got {dim}.") + half = dim // 2 + position = position.type(torch.float64) + + # calculation + sinusoid = torch.outer(position, torch.pow(10000, -torch.arange(half).to(position).div(half))) + x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1) + return x + + +@torch.amp.autocast("cuda", enabled=False) +def rope_params(max_seq_len, dim, theta=10000): + if dim % 2 != 0: + raise ValueError(f"dim must be even, got {dim}.") + freqs = torch.outer( + torch.arange(max_seq_len), 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float64).div(dim)) + ) + freqs = torch.polar(torch.ones_like(freqs), freqs) + return freqs + + +@torch.amp.autocast("cuda", enabled=False) +def rope_apply(x, grid_sizes, freqs): + n, c = x.size(2), x.size(3) // 2 + + # split freqs + freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) + + # loop over samples + output = [] + for i, (f, h, w) in enumerate(grid_sizes.tolist()): + seq_len = f * h * w + + # precompute multipliers + x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape(seq_len, n, -1, 2)) + freqs_i = torch.cat( + [ + freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1), + freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), + freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1), + ], + dim=-1, + ).reshape(seq_len, 1, -1) + + # apply rotary embedding + x_i = torch.view_as_real(x_i * freqs_i).flatten(2) + x_i = torch.cat([x_i, x[i, seq_len:]]) + + # append to collection + output.append(x_i) + return torch.stack(output).float() + + +class WanRMSNorm(nn.Module): + def __init__(self, dim, eps=1e-5): + super().__init__() + self.dim = dim + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + r""" + Args: + x(Tensor): Shape [B, L, C] + """ + return self._norm(x.float()).type_as(x) * self.weight + + def _norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) + + +class WanLayerNorm(nn.LayerNorm): + def __init__(self, dim, eps=1e-6, elementwise_affine=False): + super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps) + + def forward(self, x): + r""" + Args: + x(Tensor): Shape [B, L, C] + """ + return super().forward(x.float()).type_as(x) + + +class WanSelfAttention(nn.Module): + def __init__(self, dim, num_heads, qk_norm=True, eps=1e-6): + if dim % num_heads != 0: + raise ValueError(f"dim ({dim}) must be divisible by num_heads ({num_heads}).") + super().__init__() + self.num_heads = num_heads + self.head_dim = dim // num_heads + + # layers + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + + # NOTE: FastWAM never runs the upstream Wan attention forward. FastWAMAttentionBlock + # reuses only the q/k/v/o/norm submodules defined above and computes attention via + # `fastwam_masked_attention` (SDPA). The original flash-attention forward was removed, + # which also collapsed the former WanCrossAttention subclass into this class (it only + # differed by its forward): self- and cross-attention now share the same projection module. + + +class WanAttentionBlock(nn.Module): + def __init__(self, dim, ffn_dim, num_heads, qk_norm=True, cross_attn_norm=False, eps=1e-6): + super().__init__() + self.dim = dim + self.ffn_dim = ffn_dim + self.num_heads = num_heads + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + + # layers + self.norm1 = WanLayerNorm(dim, eps) + self.self_attn = WanSelfAttention(dim, num_heads, qk_norm, eps) + self.norm3 = WanLayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() + self.cross_attn = WanSelfAttention(dim, num_heads, qk_norm, eps) + self.norm2 = WanLayerNorm(dim, eps) + self.ffn = nn.Sequential( + nn.Linear(dim, ffn_dim), nn.GELU(approximate="tanh"), nn.Linear(ffn_dim, dim) + ) + + # modulation + self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + # NOTE: The upstream Wan block forward (self-attention + cross-attention + FFN via + # flash-attention) was removed. FastWAM subclasses this block as FastWAMAttentionBlock + # and overrides forward to use SDPA with explicit boolean masks; only __init__ (the + # norm/attention/ffn submodules) is reused here. + + +class Head(nn.Module): + def __init__(self, dim, out_dim, patch_size, eps=1e-6): + super().__init__() + self.dim = dim + self.out_dim = out_dim + self.patch_size = patch_size + self.eps = eps + + # layers + out_dim = math.prod(patch_size) * out_dim + self.norm = WanLayerNorm(dim, eps) + self.head = nn.Linear(dim, out_dim) + + # modulation + self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + + def forward(self, x, e): + r""" + Args: + x(Tensor): Shape [B, L1, C] + e(Tensor): Shape [B, L1, C] + """ + with torch.amp.autocast("cuda", dtype=torch.float32): + e = (self.modulation.unsqueeze(0) + e.unsqueeze(2)).chunk(2, dim=2) + x = self.head(self.norm(x) * (1 + e[1].squeeze(2)) + e[0].squeeze(2)) + return x + + +class WanModel(nn.Module): + r""" + Wan diffusion backbone supporting both text-to-video and image-to-video. + """ + + def __init__( + self, + model_type="t2v", + patch_size=(1, 2, 2), + text_len=512, + in_dim=16, + dim=2048, + ffn_dim=8192, + freq_dim=256, + text_dim=4096, + out_dim=16, + num_heads=16, + num_layers=32, + qk_norm=True, + cross_attn_norm=True, + eps=1e-6, + ): + r""" + Initialize the diffusion model backbone. + + Args: + model_type (`str`, *optional*, defaults to 't2v'): + Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video) + patch_size (`tuple`, *optional*, defaults to (1, 2, 2)): + 3D patch dimensions for video embedding (t_patch, h_patch, w_patch) + text_len (`int`, *optional*, defaults to 512): + Fixed length for text embeddings + in_dim (`int`, *optional*, defaults to 16): + Input video channels (C_in) + dim (`int`, *optional*, defaults to 2048): + Hidden dimension of the transformer + ffn_dim (`int`, *optional*, defaults to 8192): + Intermediate dimension in feed-forward network + freq_dim (`int`, *optional*, defaults to 256): + Dimension for sinusoidal time embeddings + text_dim (`int`, *optional*, defaults to 4096): + Input dimension for text embeddings + out_dim (`int`, *optional*, defaults to 16): + Output video channels (C_out) + num_heads (`int`, *optional*, defaults to 16): + Number of attention heads + num_layers (`int`, *optional*, defaults to 32): + Number of transformer blocks + qk_norm (`bool`, *optional*, defaults to True): + Enable query/key normalization + cross_attn_norm (`bool`, *optional*, defaults to False): + Enable cross-attention normalization + eps (`float`, *optional*, defaults to 1e-6): + Epsilon value for normalization layers + """ + + super().__init__() + + if model_type not in ["t2v", "i2v", "ti2v", "s2v"]: + raise ValueError(f"model_type must be one of ['t2v', 'i2v', 'ti2v', 's2v'], got {model_type!r}.") + self.model_type = model_type + + self.patch_size = patch_size + self.text_len = text_len + self.in_dim = in_dim + self.dim = dim + self.ffn_dim = ffn_dim + self.freq_dim = freq_dim + self.text_dim = text_dim + self.out_dim = out_dim + self.num_heads = num_heads + self.num_layers = num_layers + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + + # embeddings + self.patch_embedding = nn.Conv3d(in_dim, dim, kernel_size=patch_size, stride=patch_size) + self.text_embedding = nn.Sequential( + nn.Linear(text_dim, dim), nn.GELU(approximate="tanh"), nn.Linear(dim, dim) + ) + + self.time_embedding = nn.Sequential(nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6)) + + # blocks + self.blocks = nn.ModuleList( + [ + WanAttentionBlock(dim, ffn_dim, num_heads, qk_norm, cross_attn_norm, eps) + for _ in range(num_layers) + ] + ) + + # head + self.head = Head(dim, out_dim, patch_size, eps) + + # buffers (don't use register_buffer otherwise dtype will be changed in to()) + if (dim % num_heads) != 0 or (dim // num_heads) % 2 != 0: + raise ValueError( + f"dim ({dim}) must be divisible by num_heads ({num_heads}) with an even head dim." + ) + d = dim // num_heads + self.freqs = torch.cat( + [ + rope_params(1024, d - 4 * (d // 6)), + rope_params(1024, 2 * (d // 6)), + rope_params(1024, 2 * (d // 6)), + ], + dim=1, + ) + + # initialize weights + self.init_weights() + + # NOTE: The upstream Wan diffusion forward (flash-attention based) was removed. + # FastWAM's WanVideoDiT subclasses this model, rebuilds `self.blocks` with + # FastWAMAttentionBlock, and provides its own SDPA-based forward. Only the + # constructor (embeddings, blocks, head, rope buffers) and the helpers below + # (unpatchify / init_weights) are reused. WanModel is never run directly. + + def unpatchify(self, x, grid_sizes): + r""" + Reconstruct video tensors from patch embeddings. + + Args: + x (List[Tensor]): + List of patchified features, each with shape [L, C_out * prod(patch_size)] + grid_sizes (Tensor): + Original spatial-temporal grid dimensions before patching, + shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches) + + Returns: + List[Tensor]: + Reconstructed video tensors with shape [C_out, F, H / 8, W / 8] + """ + + c = self.out_dim + out = [] + for u, v in zip(x, grid_sizes.tolist(), strict=False): + u = u[: math.prod(v)].view(*v, *self.patch_size, c) + u = torch.einsum("fhwpqrc->cfphqwr", u) + u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size, strict=False)]) + out.append(u) + return out + + def init_weights(self): + r""" + Initialize model parameters using Xavier initialization. + """ + + # basic init + for m in self.modules(): + if isinstance(m, nn.Linear): + nn.init.xavier_uniform_(m.weight) + if m.bias is not None: + nn.init.zeros_(m.bias) + + # init embeddings + nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1)) + for m in self.text_embedding.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=0.02) + for m in self.time_embedding.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=0.02) + + # init output layer + nn.init.zeros_(self.head.head.weight) diff --git a/src/lerobot/policies/fastwam/wan/modular.py b/src/lerobot/policies/fastwam/wan/modular.py new file mode 100644 index 000000000..fac96776b --- /dev/null +++ b/src/lerobot/policies/fastwam/wan/modular.py @@ -0,0 +1,1912 @@ +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import logging +import re +from collections.abc import Sequence +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as functional +from PIL import Image + +from .components import ( + WAN22_DIFFUSERS_MODEL_ID, + WAN_T5_TOKENIZER, + build_wan_tokenizer, + load_pretrained_wan_text_encoder, + load_pretrained_wan_vae, + load_wan_video_dit, + resolve_wan_dit_paths, +) +from .video_dit import ( + FastWAMAttentionBlock, + WanContinuousFlowMatchScheduler, + fastwam_masked_attention, + gradient_checkpoint_forward, + modulate, + precompute_freqs_cis, + sinusoidal_embedding_1d, +) + +logger = logging.getLogger(__name__) + + +def _apply_block_norm(block, name: str, x: torch.Tensor) -> torch.Tensor: + apply_norm = getattr(block, f"apply_{name}", None) + if apply_norm is not None: + return apply_norm(x) + return getattr(block, name)(x) + + +class ActionHead(nn.Module): + def __init__(self, hidden_dim: int, out_dim: int, eps: float): + super().__init__() + self.norm = nn.LayerNorm(hidden_dim, eps=eps, elementwise_affine=False) + self.proj = nn.Linear(hidden_dim, out_dim) + self.modulation = nn.Parameter(torch.randn(1, 2, hidden_dim) / hidden_dim**0.5) + + def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + shift, scale = (self.modulation.to(dtype=t.dtype, device=t.device) + t.unsqueeze(1)).chunk(2, dim=1) + shift = shift.squeeze(1) + scale = scale.squeeze(1) + return self.proj(self.norm(x) * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)) + + +class ActionDiT(nn.Module): + def __init__( + self, + hidden_dim: int, + action_dim: int, + ffn_dim: int, + text_dim: int, + freq_dim: int, + eps: float, + num_heads: int, + attn_head_dim: int, + num_layers: int, + use_gradient_checkpointing: bool = False, + fp32_attention: bool = True, + ): + super().__init__() + self.hidden_dim = hidden_dim + self.action_dim = action_dim + self.ffn_dim = ffn_dim + self.text_dim = text_dim + self.freq_dim = freq_dim + self.num_heads = num_heads + self.attn_head_dim = attn_head_dim + + if num_heads <= 0: + raise ValueError(f"`num_heads` must be > 0, got {num_heads}") + if attn_head_dim <= 0: + raise ValueError(f"`attn_head_dim` must be > 0, got {attn_head_dim}") + if attn_head_dim % 2 != 0: + raise ValueError(f"`attn_head_dim` must be even for RoPE, got {attn_head_dim}") + + self.action_encoder = nn.Linear(action_dim, hidden_dim) + self.text_embedding = nn.Sequential( + nn.Linear(text_dim, hidden_dim), + nn.GELU(approximate="tanh"), + nn.Linear(hidden_dim, hidden_dim), + ) + self.time_embedding = nn.Sequential( + nn.Linear(freq_dim, hidden_dim), + nn.SiLU(), + nn.Linear(hidden_dim, hidden_dim), + ) + self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(hidden_dim, hidden_dim * 6)) + self.blocks = nn.ModuleList( + [ + FastWAMAttentionBlock( + hidden_dim=hidden_dim, + attn_head_dim=attn_head_dim, + num_heads=num_heads, + ffn_dim=ffn_dim, + eps=eps, + fp32_attention=fp32_attention, + ) + for _ in range(num_layers) + ] + ) + self.head = nn.Linear(hidden_dim, action_dim) + self.freqs = precompute_freqs_cis(attn_head_dim, end=1024) + self.fp32_attention = bool(fp32_attention) + + self.use_gradient_checkpointing = use_gradient_checkpointing + + def pre_dit( + self, + action_tokens: torch.Tensor, + timestep: torch.Tensor, + context: torch.Tensor, + context_mask: torch.Tensor | None = None, + ) -> dict[str, Any]: + if action_tokens.ndim != 3: + raise ValueError( + f"`action_tokens` must be 3D [B, T, action_dim], got shape {tuple(action_tokens.shape)}" + ) + if action_tokens.shape[2] != self.action_dim: + raise ValueError( + f"`action_tokens` last dim must be {self.action_dim}, got {action_tokens.shape[2]}" + ) + if timestep.ndim != 1: + raise ValueError(f"`timestep` must be 1D [B] or [1], got shape {tuple(timestep.shape)}") + if context.ndim != 3: + raise ValueError(f"`context` must be 3D [B, L, D], got shape {tuple(context.shape)}") + + batch_size = action_tokens.shape[0] + if context.shape[0] != batch_size: + raise ValueError( + f"Batch mismatch between action tokens and text context: {batch_size} vs {context.shape[0]}" + ) + if timestep.shape[0] not in (1, batch_size): + raise ValueError( + f"`timestep` length must be 1 or batch_size({batch_size}), got {timestep.shape[0]}" + ) + if timestep.shape[0] == 1 and batch_size > 1: + if self.training: + raise ValueError("During training, action timestep length must match batch_size.") + timestep = timestep.expand(batch_size) + + if context_mask is None: + context_mask = torch.ones((batch_size, context.shape[1]), dtype=torch.bool, device=context.device) + else: + if context_mask.ndim != 2: + raise ValueError(f"`context_mask` must be 2D [B, L], got shape {tuple(context_mask.shape)}") + if context_mask.shape[0] != batch_size or context_mask.shape[1] != context.shape[1]: + raise ValueError( + f"`context_mask` shape must match `context` shape [B, L], got {tuple(context_mask.shape)} vs {tuple(context.shape)}" + ) + + seq_len = action_tokens.shape[1] + if seq_len > self.freqs.shape[0]: + raise ValueError(f"Action token length {seq_len} exceeds RoPE cache {self.freqs.shape[0]}.") + + model_dtype = self.action_encoder.weight.dtype + action_tokens = action_tokens.to(dtype=model_dtype) + context = context.to(dtype=model_dtype) + t_emb = sinusoidal_embedding_1d(self.freq_dim, timestep).to(dtype=model_dtype) + t = self.time_embedding(t_emb) + t_mod = self.time_projection(t).unflatten(1, (6, self.hidden_dim)) + + tokens = self.action_encoder(action_tokens) + context_emb = self.text_embedding(context) + context_attn_mask = context_mask.unsqueeze(1).expand(-1, seq_len, -1) + freqs = self.freqs[:seq_len].view(seq_len, 1, -1).to(tokens.device) + + return { + "tokens": tokens, + "freqs": freqs, + "t": t, + "t_mod": t_mod, + "context": context_emb, + "context_mask": context_attn_mask, + "meta": { + "batch_size": batch_size, + "seq_len": seq_len, + }, + } + + def post_dit(self, tokens: torch.Tensor, pre_state: dict[str, Any]) -> torch.Tensor: + return self.head(tokens) + + def forward( + self, + action_tokens: torch.Tensor, + timestep: torch.Tensor, + context: torch.Tensor, + context_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + pre_state = self.pre_dit( + action_tokens=action_tokens, + timestep=timestep, + context=context, + context_mask=context_mask, + ) + x = pre_state["tokens"] + context = pre_state["context"] + t_mod = pre_state["t_mod"] + freqs = pre_state["freqs"] + context_mask = pre_state["context_mask"] + + for block in self.blocks: + if self.use_gradient_checkpointing: + x = gradient_checkpoint_forward( + block, + self.use_gradient_checkpointing, + x, + context, + t_mod, + freqs, + context_mask=context_mask, + ) + else: + x = block(x, context, t_mod, freqs, context_mask=context_mask) + + return self.post_dit(x, pre_state) + + +class MoTLayer(nn.Module): + """A single MoT layer: owns one transformer block per expert and runs the cross-expert + mixed-attention step for that layer. + + This exists as a module — rather than the per-layer work being inlined in ``MoT``'s loop — + so FSDP can wrap each layer as its own unit. FSDP all-gathers a wrapped module's sharded + parameters via a hook on that module's ``forward``/``__call__``. ``MoT`` drives block + submodules directly (the joint mixed attention concatenates Q/K/V across experts, so no + single block's ``forward`` is ever called), so ``MoTLayer.forward`` is the only call + boundary FSDP can hook. All three per-layer paths therefore dispatch through + ``forward(mode=...)`` so each enters via ``__call__``. + """ + + def __init__( + self, + blocks: dict[str, nn.Module], + experts: dict[str, nn.Module], + num_heads: int, + attn_head_dim: int, + fp32_attention: bool, + mot_checkpoint_mixed_attn: bool, + ): + super().__init__() + # Registered owner of this layer's blocks (one per expert) — the FSDP wrap unit. + self.blocks = nn.ModuleDict(blocks) + self.expert_order = list(blocks.keys()) + # Unregistered back-references to the experts: used only to read the live + # `use_gradient_checkpointing` flag, kept out of parameters()/state_dict(). + object.__setattr__(self, "_experts", dict(experts)) + self.num_heads = num_heads + self.attn_head_dim = attn_head_dim + self.fp32_attention = bool(fp32_attention) + self.mot_checkpoint_mixed_attn = bool(mot_checkpoint_mixed_attn) + + @staticmethod + def _split_modulation(block, t_mod: torch.Tensor): + has_seq = len(t_mod.shape) == 4 + chunk_dim = 2 if has_seq else 1 + + base_mod = block.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (base_mod + t_mod).chunk( + 6, dim=chunk_dim + ) + if has_seq: + # means t_mod has separate modulation for each token, otherwise same modulation for all tokens in the block + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + shift_msa.squeeze(2), + scale_msa.squeeze(2), + gate_msa.squeeze(2), + shift_mlp.squeeze(2), + scale_mlp.squeeze(2), + gate_mlp.squeeze(2), + ) + return shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp + + def _mixed_attention( + self, + q_cat: torch.Tensor, + k_cat: torch.Tensor, + v_cat: torch.Tensor, + attention_mask: torch.Tensor, + ) -> torch.Tensor: + attn_mask = attention_mask.to(device=q_cat.device) + + def _forward(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + return fastwam_masked_attention( + q=q, + k=k, + v=v, + num_heads=self.num_heads, + ctx_mask=attn_mask, + fp32_attention=self.fp32_attention, + ) + + if self.mot_checkpoint_mixed_attn and self.training: + return torch.utils.checkpoint.checkpoint( + _forward, + q_cat, + k_cat, + v_cat, + use_reentrant=False, + ) + return _forward(q_cat, k_cat, v_cat) + + @staticmethod + def _apply_expert_post_block( + block, + residual_x: torch.Tensor, + mixed_attn_out: torch.Tensor, + gate_msa: torch.Tensor, + shift_mlp: torch.Tensor, + scale_mlp: torch.Tensor, + gate_mlp: torch.Tensor, + context_payload: dict | None, + ) -> torch.Tensor: + if hasattr(block, "project_self_attention_output"): + projected_attn = block.project_self_attention_output(mixed_attn_out) + else: + projected_attn = block.self_attn.o(mixed_attn_out.to(dtype=block.self_attn.o.weight.dtype)) + x = residual_x + gate_msa * projected_attn + + if context_payload is not None: + context = context_payload.get("context") + if context is not None: + context_mask = context_payload.get("mask") + if context_mask is not None and context_mask.dim() == 3: + context_mask = context_mask.unsqueeze(1) + x = x + block.apply_cross_attention( + _apply_block_norm(block, "norm3", x), + context, + context_mask=context_mask, + ) + + mlp_input = modulate(_apply_block_norm(block, "norm2", x), shift_mlp, scale_mlp) + x = x + gate_mlp * block.ffn(mlp_input) + return x + + def _build_expert_attention_io( + self, + name: str, + x: torch.Tensor, + freqs: torch.Tensor | dict[str, torch.Tensor], + t_mod: torch.Tensor, + ): + """Build this expert's attention tensors and post-block states for the layer. + + Returns (q, k, v, residual_x, gate_msa, shift_mlp, scale_mlp, gate_mlp, use_gc). + """ + block = self.blocks[name] + expert = self._experts[name] + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self._split_modulation(block, t_mod) + attn_input = modulate(_apply_block_norm(block, "norm1", x), shift_msa, scale_msa) + + q, k, v = block.project_self_attention(attn_input, freqs) + + use_gradient_checkpointing = bool(getattr(expert, "use_gradient_checkpointing", False)) + return ( + q, + k, + v, + x, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + use_gradient_checkpointing, + ) + + def _apply_post_with_optional_checkpoint( + self, + block, + residual_x: torch.Tensor, + gate_msa: torch.Tensor, + shift_mlp: torch.Tensor, + scale_mlp: torch.Tensor, + gate_mlp: torch.Tensor, + use_gradient_checkpointing: bool, + mixed_slice: torch.Tensor, + context_payload: dict | None, + ) -> torch.Tensor: + def _post_fn( + _mixed_slice: torch.Tensor, + _x: torch.Tensor, + _gate_msa: torch.Tensor, + _shift_mlp: torch.Tensor, + _scale_mlp: torch.Tensor, + _gate_mlp: torch.Tensor, + _block=block, + _context_payload=context_payload, + ) -> torch.Tensor: + return self._apply_expert_post_block( + block=_block, + residual_x=_x, + mixed_attn_out=_mixed_slice, + gate_msa=_gate_msa, + shift_mlp=_shift_mlp, + scale_mlp=_scale_mlp, + gate_mlp=_gate_mlp, + context_payload=_context_payload, + ) + + if use_gradient_checkpointing and self.training: + return torch.utils.checkpoint.checkpoint( + _post_fn, + mixed_slice, + residual_x, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + use_reentrant=False, + ) + return _post_fn( + mixed_slice, + residual_x, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + ) + + def forward(self, mode: str, **kwargs): + if mode == "joint": + return self._forward_joint(**kwargs) + if mode == "video_prefill": + return self._forward_video_prefill(**kwargs) + if mode == "action_cached": + return self._forward_action_cached(**kwargs) + raise ValueError(f"Unknown MoTLayer forward mode: {mode!r}") + + def _forward_joint( + self, + tokens_all: dict[str, torch.Tensor], + attention_mask: torch.Tensor, + freqs_all: dict[str, torch.Tensor], + context_all: dict[str, dict | None], + t_mod_all: dict[str, torch.Tensor], + ) -> dict[str, torch.Tensor]: + q_chunks = [] + k_chunks = [] + v_chunks = [] + cached = {} + seq_lens = [] + + for name in self.expert_order: + ( + q, + k, + v, + residual_x, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + use_gradient_checkpointing, + ) = self._build_expert_attention_io(name, tokens_all[name], freqs_all[name], t_mod_all[name]) + + q_chunks.append(q) + k_chunks.append(k) + v_chunks.append(v) + seq_lens.append(tokens_all[name].shape[1]) + cached[name] = { + "residual_x": residual_x, + "gate_msa": gate_msa, + "shift_mlp": shift_mlp, + "scale_mlp": scale_mlp, + "gate_mlp": gate_mlp, + "use_gradient_checkpointing": use_gradient_checkpointing, + } + + q_cat = torch.cat(q_chunks, dim=1) + k_cat = torch.cat(k_chunks, dim=1) + v_cat = torch.cat(v_chunks, dim=1) + + total_seq = q_cat.shape[1] + if attention_mask.shape[0] != total_seq: + raise ValueError( + f"Attention mask seq length mismatch: mask={attention_mask.shape[0]} vs tokens={total_seq}" + ) + + mixed = self._mixed_attention(q_cat=q_cat, k_cat=k_cat, v_cat=v_cat, attention_mask=attention_mask) + + out = {} + start = 0 + for name, seq_len in zip(self.expert_order, seq_lens, strict=True): + end = start + seq_len + mixed_slice = mixed[:, start:end, :] + cached_expert = cached[name] + out[name] = self._apply_post_with_optional_checkpoint( + block=self.blocks[name], + residual_x=cached_expert["residual_x"], + gate_msa=cached_expert["gate_msa"], + shift_mlp=cached_expert["shift_mlp"], + scale_mlp=cached_expert["scale_mlp"], + gate_mlp=cached_expert["gate_mlp"], + use_gradient_checkpointing=cached_expert["use_gradient_checkpointing"], + mixed_slice=mixed_slice, + context_payload=context_all.get(name), + ) + start = end + return out + + def _forward_video_prefill( + self, + x: torch.Tensor, + freqs: torch.Tensor, + t_mod: torch.Tensor, + context_payload: dict | None, + video_attention_mask: torch.Tensor, + ): + ( + q, + k, + v, + residual_x, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + use_gradient_checkpointing, + ) = self._build_expert_attention_io("video", x, freqs, t_mod) + # Video prefill uses only video self-attention mask. + mixed = self._mixed_attention(q_cat=q, k_cat=k, v_cat=v, attention_mask=video_attention_mask) + x_out = self._apply_post_with_optional_checkpoint( + block=self.blocks["video"], + residual_x=residual_x, + gate_msa=gate_msa, + shift_mlp=shift_mlp, + scale_mlp=scale_mlp, + gate_mlp=gate_mlp, + use_gradient_checkpointing=use_gradient_checkpointing, + mixed_slice=mixed, + context_payload=context_payload, + ) + return x_out, k, v + + def _forward_action_cached( + self, + x: torch.Tensor, + freqs: torch.Tensor, + t_mod: torch.Tensor, + context_payload: dict | None, + k_video: torch.Tensor, + v_video: torch.Tensor, + action_attention_mask: torch.Tensor, + ) -> torch.Tensor: + ( + q_action, + k_action, + v_action, + residual_x, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + use_gradient_checkpointing, + ) = self._build_expert_attention_io("action", x, freqs, t_mod) + # Mixed attention: action queries attend to cached video K/V plus current action K/V. + k_cat = torch.cat([k_video, k_action], dim=1) + v_cat = torch.cat([v_video, v_action], dim=1) + mixed = self._mixed_attention( + q_cat=q_action, k_cat=k_cat, v_cat=v_cat, attention_mask=action_attention_mask + ) + return self._apply_post_with_optional_checkpoint( + block=self.blocks["action"], + residual_x=residual_x, + gate_msa=gate_msa, + shift_mlp=shift_mlp, + scale_mlp=scale_mlp, + gate_mlp=gate_mlp, + use_gradient_checkpointing=use_gradient_checkpointing, + mixed_slice=mixed, + context_payload=context_payload, + ) + + +class MoT(nn.Module): + def __init__( + self, + mixtures: dict[str, nn.Module], + mot_checkpoint_mixed_attn: bool = True, + ): + super().__init__() + if not mixtures: + raise ValueError("`mixtures` cannot be empty.") + if "video" not in mixtures or "action" not in mixtures: + raise ValueError("`mixtures` must include both 'video' and 'action' experts.") + + self.mixtures = nn.ModuleDict(mixtures) + self.expert_order = list(self.mixtures.keys()) + self.mot_checkpoint_mixed_attn = mot_checkpoint_mixed_attn + if mot_checkpoint_mixed_attn: + logger.info( + "Using gradient checkpointing for mixture attention. This will save memory but use more computation." + ) + + first_expert = self.mixtures[self.expert_order[0]] + self.num_layers = len(first_expert.blocks) + self.num_heads = first_expert.num_heads + self.attn_head_dim = first_expert.attn_head_dim + self.fp32_attention = bool(getattr(first_expert, "fp32_attention", True)) + + for name in self.expert_order[1:]: + expert = self.mixtures[name] + if len(expert.blocks) != self.num_layers: + raise ValueError( + f"All experts must have same number of layers; got {self.num_layers} and {len(expert.blocks)}" + ) + if expert.num_heads != self.num_heads: + raise ValueError( + f"All experts must have same num_heads; got {self.num_heads} and {expert.num_heads}" + ) + if expert.attn_head_dim != self.attn_head_dim: + raise ValueError( + "All experts must have same attn_head_dim; " + f"got {self.attn_head_dim} and {expert.attn_head_dim}" + ) + if bool(getattr(expert, "fp32_attention", True)) != self.fp32_attention: + raise ValueError("All experts must use the same `fp32_attention` setting.") + + logger.info(f"Initialized MoT with experts: {self.expert_order}, num_layers={self.num_layers}") + for name in self.expert_order: + expert = self.mixtures[name] + logger.info( + f" Expert '{name}': num_params={sum(p.numel() for p in expert.parameters()) / 1e9:.2f} B" + ) + + # One MoTLayer per layer, each owning that layer's block from every expert. This is the + # FSDP wrap unit: only MoTLayer.forward is ever called (MoT drives block submodules + # directly for the cross-expert mixed attention), so it is the boundary at which FSDP can + # all-gather a layer's params. The blocks are RE-PARENTED into the layers — removed from + # each expert's module registry — so they have a single owner; leaving them registered + # under both the expert and the layer would make FSDP try to manage the same params twice. + self.layers = nn.ModuleList( + [ + MoTLayer( + blocks={name: self.mixtures[name].blocks[layer_idx] for name in self.expert_order}, + experts={name: self.mixtures[name] for name in self.expert_order}, + num_heads=self.num_heads, + attn_head_dim=self.attn_head_dim, + fp32_attention=self.fp32_attention, + mot_checkpoint_mixed_attn=self.mot_checkpoint_mixed_attn, + ) + for layer_idx in range(self.num_layers) + ] + ) + for name in self.expert_order: + expert = self.mixtures[name] + kept_blocks = list(expert.blocks) + del expert._modules["blocks"] + # Keep an UNREGISTERED reference so the (unused) standalone `expert.forward` and any + # `len(expert.blocks)` still work, without re-adding the params to the expert's + # parameters()/state_dict() (which would double-register them with the MoTLayer owner). + object.__setattr__(expert, "blocks", kept_blocks) + + def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): + # Backward-compat for checkpoints saved before the MoTLayer refactor. Then the per-layer + # blocks were keyed under the experts (`{prefix}mixtures..blocks..`, e.g. + # the released `ZibinDong/fastwam_libero_uncond_2cam224`); now they are owned by the layers + # (`{prefix}layers..blocks..`). Remap legacy keys in place so the recursion + # into `self.layers` finds them and the (now block-less) `self.mixtures` does not flag them. + legacy = re.compile(re.escape(prefix) + r"mixtures\.([^.]+)\.blocks\.(\d+)\.(.+)$") + moved = {} + for key in list(state_dict.keys()): + m = legacy.match(key) + if m is not None: + name, layer_idx, rest = m.group(1), m.group(2), m.group(3) + moved[f"{prefix}layers.{layer_idx}.blocks.{name}.{rest}"] = state_dict.pop(key) + state_dict.update(moved) + super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) + + def prefill_video_cache( + self, + video_tokens: torch.Tensor, + video_freqs: torch.Tensor, + video_t_mod: torch.Tensor, + video_context_payload: dict | None, + video_attention_mask: torch.Tensor, + ) -> list[dict[str, torch.Tensor]]: + """Prefill video branch once and cache per-layer K/V for action denoising. + + Returns a list of length ``num_layers``, each entry ``{"k": ..., "v": ...}``. + """ + if "video" not in self.mixtures: + raise ValueError("MoT requires `video` expert for `prefill_video_cache`.") + if video_attention_mask.ndim != 2: + raise ValueError( + f"`video_attention_mask` must be 2D [S,S], got shape {tuple(video_attention_mask.shape)}" + ) + if video_attention_mask.shape[0] != video_attention_mask.shape[1]: + raise ValueError( + f"`video_attention_mask` must be square, got shape {tuple(video_attention_mask.shape)}" + ) + if video_attention_mask.shape[0] != video_tokens.shape[1]: + raise ValueError( + "`video_attention_mask` seq length mismatch: " + f"mask={video_attention_mask.shape[0]} vs tokens={video_tokens.shape[1]}" + ) + + x = video_tokens + kv_cache: list[dict[str, torch.Tensor]] = [] + for layer in self.layers: + x, k, v = layer( + mode="video_prefill", + x=x, + freqs=video_freqs, + t_mod=video_t_mod, + context_payload=video_context_payload, + video_attention_mask=video_attention_mask, + ) + kv_cache.append({"k": k, "v": v}) + return kv_cache + + def forward_action_with_video_cache( + self, + action_tokens: torch.Tensor, + action_freqs: torch.Tensor, + action_t_mod: torch.Tensor, + action_context_payload: dict | None, + video_kv_cache: list[dict[str, torch.Tensor]], + attention_mask: torch.Tensor, + video_seq_len: int, + ) -> torch.Tensor: + """Run action branch with cached video K/V instead of recomputing video tokens.""" + if "action" not in self.mixtures: + raise ValueError("MoT requires `action` expert for `forward_action_with_video_cache`.") + if len(video_kv_cache) != self.num_layers: + raise ValueError( + f"`video_kv_cache` must contain {self.num_layers} layers, got {len(video_kv_cache)}." + ) + if attention_mask.ndim != 2: + raise ValueError(f"`attention_mask` must be 2D [S,S], got shape {tuple(attention_mask.shape)}") + if attention_mask.shape[0] != attention_mask.shape[1]: + raise ValueError(f"`attention_mask` must be square, got shape {tuple(attention_mask.shape)}") + + action_seq_len = int(action_tokens.shape[1]) + total_seq_len = int(video_seq_len) + action_seq_len + if attention_mask.shape[0] != total_seq_len: + raise ValueError( + "`attention_mask` seq length mismatch: " + f"mask={attention_mask.shape[0]} vs expected_total={total_seq_len}" + ) + # Use the action query rows from the joint [video+action] mask. + action_attention_mask = attention_mask[video_seq_len:total_seq_len, :total_seq_len] + + x = action_tokens + for layer_idx, layer in enumerate(self.layers): + layer_cache = video_kv_cache[layer_idx] + if "k" not in layer_cache or "v" not in layer_cache: + raise ValueError(f"`video_kv_cache[{layer_idx}]` must contain `k` and `v`.") + k_video = layer_cache["k"] + v_video = layer_cache["v"] + if k_video.shape[1] != video_seq_len or v_video.shape[1] != video_seq_len: + raise ValueError(f"`video_kv_cache[{layer_idx}]` seq len mismatch, expected {video_seq_len}.") + x = layer( + mode="action_cached", + x=x, + freqs=action_freqs, + t_mod=action_t_mod, + context_payload=action_context_payload, + k_video=k_video, + v_video=v_video, + action_attention_mask=action_attention_mask, + ) + return x + + def forward( + self, + embeds_all: dict[str, torch.Tensor], + attention_mask: torch.Tensor, + freqs_all: dict[str, torch.Tensor], + context_all: dict[str, dict | None], + t_mod_all: dict[str, torch.Tensor], + ): + missing = [k for k in self.expert_order if k not in embeds_all] + if missing: + raise ValueError(f"Missing expert tokens for {missing}") + missing = [k for k in self.expert_order if k not in freqs_all] + if missing: + raise ValueError(f"Missing expert freqs for {missing}") + missing = [k for k in self.expert_order if k not in t_mod_all] + if missing: + raise ValueError(f"Missing expert t_mod for {missing}") + + if attention_mask.ndim != 2: + raise ValueError(f"`attention_mask` must be 2D [S, S], got shape {tuple(attention_mask.shape)}") + if attention_mask.shape[0] != attention_mask.shape[1]: + raise ValueError(f"`attention_mask` must be square, got shape {tuple(attention_mask.shape)}") + + # Each layer is a MoTLayer module; entering via __call__ lets FSDP all-gather that + # layer's params (the whole point of the per-layer split). + tokens_all = dict(embeds_all) + for layer in self.layers: + tokens_all = layer( + mode="joint", + tokens_all=tokens_all, + attention_mask=attention_mask, + freqs_all=freqs_all, + context_all=context_all, + t_mod_all=t_mod_all, + ) + return tokens_all + + +class FastWAM(torch.nn.Module): + """MoT world model with video/action experts.""" + + def __init__( + self, + video_expert, + action_expert: ActionDiT, + mot: MoT, + vae, + text_encoder=None, + tokenizer=None, + text_dim: int | None = None, + proprio_dim: int | None = None, + device: str = "cpu", + torch_dtype: torch.dtype = torch.float32, + video_train_shift: float = 5.0, + video_infer_shift: float = 5.0, + video_num_train_timesteps: int = 1000, + action_train_shift: float = 5.0, + action_infer_shift: float = 5.0, + action_num_train_timesteps: int = 1000, + loss_lambda_video: float = 1.0, + loss_lambda_action: float = 1.0, + ): + super().__init__() + self.mot = mot + # `video_expert` / `action_expert` are the very same module objects as + # `mot.mixtures["video"]` / `["action"]`, and `dit` is an alias of `mot`. Registering + # them as submodules too would give every expert tensor three names in `state_dict()` + # (`video_expert.*`, `mot.mixtures.video.*`, `dit.mixtures.video.*`) — a 3x-bloated + # gathered FSDP checkpoint and a doubled module tree for FSDP to traverse. Hold them as + # plain (unregistered) attributes instead — bypassing `nn.Module.__setattr__`, like the + # frozen vae/text_encoder below — so `mot` is the single registered owner and each tensor + # has one canonical name (`mot.mixtures.*` / `mot.layers.*`, matching the base checkpoint). + # Forward / freeze / optimizer code still reaches them by attribute, and device/dtype moves + # still apply via `mot`. (optimizer + freeze logic use `model.dit`.) + object.__setattr__(self, "video_expert", video_expert) + object.__setattr__(self, "action_expert", action_expert) + object.__setattr__(self, "dit", self.mot) + + # Frozen Wan2.2 components: bypass `nn.Module.__setattr__` so they are NOT + # registered as submodules. They are therefore excluded from `state_dict()` + # (lean checkpoints), `parameters()`, and DDP gradient sync, and are loaded + # with their real weights from the diffusers/transformers repos at construction. + # Device/dtype moves still reach them via the `_apply` override below. + object.__setattr__(self, "vae", vae) + object.__setattr__(self, "text_encoder", text_encoder) + self.tokenizer = tokenizer + vae.requires_grad_(False) + if text_encoder is not None: + text_encoder.requires_grad_(False) + if text_dim is None: + if self.text_encoder is None: + raise ValueError("`text_dim` is required when `text_encoder` is not loaded.") + text_dim = int(self.text_encoder.dim) + self.text_dim = int(text_dim) + self.proprio_dim = None if proprio_dim is None else int(proprio_dim) + if self.proprio_dim is not None: + self.proprio_encoder = nn.Linear(self.proprio_dim, self.text_dim).to(torch_dtype) + else: + self.proprio_encoder = None + + self.train_video_scheduler = WanContinuousFlowMatchScheduler( + num_train_timesteps=video_num_train_timesteps, + shift=video_train_shift, + ) + self.infer_video_scheduler = WanContinuousFlowMatchScheduler( + num_train_timesteps=video_num_train_timesteps, + shift=video_infer_shift, + ) + self.train_action_scheduler = WanContinuousFlowMatchScheduler( + num_train_timesteps=action_num_train_timesteps, + shift=action_train_shift, + ) + self.infer_action_scheduler = WanContinuousFlowMatchScheduler( + num_train_timesteps=action_num_train_timesteps, + shift=action_infer_shift, + ) + # Optional aliases for consistency with Wan22Core naming. + self.train_scheduler = self.train_video_scheduler + self.infer_scheduler = self.infer_video_scheduler + + self.device = torch.device(device) + self.torch_dtype = torch_dtype + self.loss_lambda_video = float(loss_lambda_video) + self.loss_lambda_action = float(loss_lambda_action) + + self.to(self.device) + + @classmethod + def from_wan22_pretrained( + cls, + device: str = "cuda", + torch_dtype: torch.dtype = torch.bfloat16, + model_id: str = "Wan-AI/Wan2.2-TI2V-5B", + tokenizer_model_id: str = WAN_T5_TOKENIZER, + text_encoder_model_id: str = WAN22_DIFFUSERS_MODEL_ID, + tokenizer_max_len: int = 512, + load_text_encoder: bool = True, + proprio_dim: int | None = None, + video_dit_config: dict[str, Any] | None = None, + action_dit_config: dict[str, Any] | None = None, + mot_checkpoint_mixed_attn: bool = True, + video_train_shift: float = 5.0, + video_infer_shift: float = 5.0, + video_num_train_timesteps: int = 1000, + action_train_shift: float = 5.0, + action_infer_shift: float = 5.0, + action_num_train_timesteps: int = 1000, + loss_lambda_video: float = 1.0, + loss_lambda_action: float = 1.0, + ): + if video_dit_config is None: + raise ValueError("`video_dit_config` is required for FastWAM.from_wan22_pretrained().") + if "text_dim" not in video_dit_config: + raise ValueError("`video_dit_config['text_dim']` is required for FastWAM.") + + # Custom MoT video DiT from the original Wan2.2 repo; frozen VAE / UMT5 from + # the diffusers conversion. This is the offline base-creation path; the + # weights it loads are then bundled into the FastWAM `model.safetensors`. + video_expert = load_wan_video_dit( + resolve_wan_dit_paths(model_id), + dit_config=video_dit_config, + torch_dtype=torch_dtype, + device=device, + ) + action_expert = ActionDiT(**action_dit_config).to(device=device, dtype=torch_dtype) + if int(action_expert.num_heads) != int(video_expert.num_heads): + raise ValueError("ActionDiT `num_heads` must match video expert for MoT mixed attention.") + if int(action_expert.attn_head_dim) != int(video_expert.attn_head_dim): + raise ValueError("ActionDiT `attn_head_dim` must match video expert for MoT mixed attention.") + if int(len(action_expert.blocks)) != int(len(video_expert.blocks)): + raise ValueError("ActionDiT `num_layers` must match video expert.") + + mot = MoT( + mixtures={"video": video_expert, "action": action_expert}, + mot_checkpoint_mixed_attn=mot_checkpoint_mixed_attn, + ) + + vae = load_pretrained_wan_vae(torch_dtype=torch_dtype, device=device) + text_encoder = ( + load_pretrained_wan_text_encoder( + model_id=text_encoder_model_id, torch_dtype=torch_dtype, device=device + ) + if load_text_encoder + else None + ) + tokenizer = build_wan_tokenizer(model_id=tokenizer_model_id, tokenizer_max_len=tokenizer_max_len) + + return cls( + video_expert=video_expert, + action_expert=action_expert, + mot=mot, + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + text_dim=int(video_dit_config["text_dim"]), + proprio_dim=proprio_dim, + device=device, + torch_dtype=torch_dtype, + video_train_shift=video_train_shift, + video_infer_shift=video_infer_shift, + video_num_train_timesteps=video_num_train_timesteps, + action_train_shift=action_train_shift, + action_infer_shift=action_infer_shift, + action_num_train_timesteps=action_num_train_timesteps, + loss_lambda_video=loss_lambda_video, + loss_lambda_action=loss_lambda_action, + ) + + def _apply(self, fn, *args, **kwargs): + # `.to()` / `.cuda()` / `.cpu()` and accelerate/DDP device moves all funnel + # through `_apply`, and the parent policy reaches us via `child._apply(fn)` + # (not `child.to()`). Propagate `fn` to the *unregistered* frozen VAE / text + # encoder here so they follow the rest of the model onto the right device, + # while staying out of `state_dict()` / `parameters()`. + super()._apply(fn, *args, **kwargs) + self.vae._apply(fn) + if self.text_encoder is not None: + self.text_encoder._apply(fn) + return self + + @staticmethod + def _check_resize_height_width(height, width, num_frames): + if height % 16 != 0: + height = (height + 15) // 16 * 16 + if width % 16 != 0: + width = (width + 15) // 16 * 16 + if num_frames % 4 != 1: + num_frames = (num_frames + 3) // 4 * 4 + 1 + return height, width, num_frames + + @torch.no_grad() + def encode_prompt(self, prompt: str | Sequence[str]): + if self.text_encoder is None or self.tokenizer is None: + raise ValueError( + "Prompt encoding requires loaded text encoder/tokenizer. " + "Set `load_text_encoder=true` or provide precomputed `context/context_mask`." + ) + ids, mask = self.tokenizer(prompt, return_mask=True, add_special_tokens=True) + ids = ids.to(self.device) + mask = mask.to(self.device, dtype=torch.bool) + prompt_emb = self.text_encoder(ids, mask) + seq_lens = mask.gt(0).sum(dim=1).long() + for i, v in enumerate(seq_lens): + prompt_emb[i, v:] = 0 + # Match FastWAM/Wan2.2 context semantics: padding embeddings are zeroed, + # while cross-attention still sees a fixed-length context. + mask = torch.ones_like(mask) + return prompt_emb.to(device=self.device), mask + + def _append_proprio_to_context( + self, + context: torch.Tensor, + context_mask: torch.Tensor, + proprio: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.proprio_encoder is None or proprio is None: + return context, context_mask + if proprio.ndim != 2: + raise ValueError(f"`proprio` must be 2D [B, D], got shape {tuple(proprio.shape)}") + if self.proprio_dim is None or proprio.shape[1] != self.proprio_dim: + raise ValueError(f"`proprio` last dim must be {self.proprio_dim}, got {proprio.shape[1]}") + proprio_token = self.proprio_encoder( + proprio.to(device=self.device, dtype=context.dtype).unsqueeze(1) + ).to(dtype=context.dtype) # [B, 1, D] + proprio_mask = torch.ones((context_mask.shape[0], 1), dtype=torch.bool, device=context_mask.device) + return ( + torch.cat([context, proprio_token], dim=1), + torch.cat([context_mask, proprio_mask], dim=1), + ) + + @torch.no_grad() + def _encode_video_latents(self, video_tensor, tiled=False, tile_size=(30, 52), tile_stride=(15, 26)): + # The Wan VAE expects pixels in [-1, 1]; model inputs arrive in [0, 1] (VISUAL is IDENTITY in + # the preprocessor — see configuration_fastwam.normalization_mapping). Map here, at the single + # video-encode boundary, so it is applied exactly once on every path. + video_tensor = video_tensor * 2.0 - 1.0 + z = self.vae.encode( + video_tensor, + device=self.device, + tiled=tiled, + tile_size=tile_size, + tile_stride=tile_stride, + ) + return z + + @torch.no_grad() + def _encode_input_image_latents_tensor( + self, input_image: torch.Tensor, tiled=False, tile_size=(30, 52), tile_stride=(15, 26) + ): + if input_image.ndim == 3: + input_image = input_image.unsqueeze(0) + if input_image.ndim != 4 or input_image.shape[0] != 1 or input_image.shape[1] != 3: + raise ValueError( + f"`input_image` must have shape [1,3,H,W] or [3,H,W], got {tuple(input_image.shape)}" + ) + # [0, 1] -> [-1, 1] for the Wan VAE (mirrors `_encode_video_latents`); single image-encode boundary. + input_image = input_image * 2.0 - 1.0 + image = input_image.to(device=self.device)[0].unsqueeze(1) + z = self.vae.encode( + [image], device=self.device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride + ) + if isinstance(z, list): + z = z[0].unsqueeze(0) + return z + + def _decode_latents(self, latents, tiled=False, tile_size=(30, 52), tile_stride=(15, 26)): + video_tensor = self.vae.decode( + latents, device=self.device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride + ) + video_tensor = video_tensor.squeeze(0).detach().float().clamp(-1, 1) + video_tensor = ((video_tensor + 1.0) * 127.5).to(torch.uint8).cpu() + frames = [] + for t in range(video_tensor.shape[1]): + frame = video_tensor[:, t].permute(1, 2, 0).numpy() + frames.append(Image.fromarray(frame)) + return frames + + def build_inputs(self, sample, tiled: bool = False): + video = sample["video"] + if "context" not in sample or "context_mask" not in sample: + raise ValueError("FastWAM training requires `sample['context']` and `sample['context_mask']`.") + context = sample["context"] + context_mask = sample["context_mask"] + proprio = sample.get("proprio", None) + if video.ndim != 5: + raise ValueError(f"`sample['video']` must be 5D [B, 3, T, H, W], got shape {tuple(video.shape)}") + if video.shape[1] != 3: + raise ValueError(f"`sample['video']` channel dimension must be 3, got shape {tuple(video.shape)}") + + batch_size, _, num_frames, height, width = video.shape + if height % 16 != 0 or width % 16 != 0: + raise ValueError(f"Video spatial dims must be multiples of 16, got H={height}, W={width}") + if num_frames % 4 != 1: + raise ValueError(f"Video T must satisfy T % 4 == 1, got T={num_frames}") + if num_frames <= 1: + raise ValueError(f"Video T must be > 1 for action-conditioned training, got T={num_frames}") + + if "action" not in sample: + raise ValueError("`sample['action']` is required for FastWAM training.") + + action = sample["action"] + if action.ndim != 3: + raise ValueError(f"`sample['action']` must be 3D [B, T, a_dim], got shape {tuple(action.shape)}") + action_horizon = int(action.shape[1]) + if action_horizon % (num_frames - 1) != 0: + raise ValueError( + f"`sample['action']` temporal dimension must be divisible by video transitions ({num_frames - 1}), got {action_horizon}" + ) + + action_is_pad = sample.get("action_is_pad", None) + if action_is_pad is not None: + if action_is_pad.ndim != 2: + raise ValueError( + f"`sample['action_is_pad']` must be 2D [B, T], got shape {tuple(action_is_pad.shape)}" + ) + if action_is_pad.shape[0] != batch_size or action_is_pad.shape[1] != action_horizon: + raise ValueError( + "`sample['action_is_pad']` shape mismatch: " + f"got {tuple(action_is_pad.shape)} vs expected ({batch_size}, {action_horizon})" + ) + + image_is_pad = sample.get("image_is_pad", None) + if image_is_pad is not None: + if image_is_pad.ndim != 2: + raise ValueError( + f"`sample['image_is_pad']` must be 2D [B, T], got shape {tuple(image_is_pad.shape)}" + ) + if image_is_pad.shape[0] != batch_size or image_is_pad.shape[1] != num_frames: + raise ValueError( + "`sample['image_is_pad']` shape mismatch: " + f"got {tuple(image_is_pad.shape)} vs expected ({batch_size}, {num_frames})" + ) + + input_video = video.to(device=self.device, dtype=self.torch_dtype, non_blocking=True) + input_latents = self._encode_video_latents(input_video, tiled=tiled) + + first_frame_latents = None + fuse_flag = False + if getattr(self.video_expert, "fuse_vae_embedding_in_latents", False): + first_frame_latents = input_latents[:, :, 0:1] + fuse_flag = True + + if context.ndim != 3 or context_mask.ndim != 2: + raise ValueError( + f"`context/context_mask` must be [B,L,D]/[B,L], got {tuple(context.shape)} and {tuple(context_mask.shape)}" + ) + context = context.to(device=self.device, dtype=self.torch_dtype, non_blocking=True) + context_mask = context_mask.to(device=self.device, dtype=torch.bool, non_blocking=True) + if self.proprio_encoder is not None: + if proprio is None: + raise ValueError("`sample['proprio']` is required when `proprio_dim` is enabled.") + if proprio.ndim != 3: + raise ValueError( + f"`sample['proprio']` must be 3D [B, T, d], got shape {tuple(proprio.shape)}" + ) + if proprio.shape[2] != self.proprio_dim: + raise ValueError( + f"`sample['proprio']` last dim must be {self.proprio_dim}, got {proprio.shape[2]}" + ) + proprio = proprio[:, 0, :] # [B, D] + context, context_mask = self._append_proprio_to_context( + context=context, + context_mask=context_mask, + proprio=proprio.to(device=self.device, dtype=self.torch_dtype), + ) + action = action.to(device=self.device, dtype=self.torch_dtype, non_blocking=True) + + if action_is_pad is not None: + action_is_pad = action_is_pad.to(device=self.device, dtype=torch.bool, non_blocking=True) + if image_is_pad is not None: + image_is_pad = image_is_pad.to(device=self.device, dtype=torch.bool, non_blocking=True) + + return { + "context": context, + "context_mask": context_mask, + "input_latents": input_latents, + "first_frame_latents": first_frame_latents, + "fuse_vae_embedding_in_latents": fuse_flag, + "action": action, + "action_is_pad": action_is_pad, + "image_is_pad": image_is_pad, + } + + @torch.no_grad() + def _build_mot_attention_mask( + self, + video_seq_len: int, + action_seq_len: int, + video_tokens_per_frame: int, + device: torch.device, + ) -> torch.Tensor: + total_seq_len = video_seq_len + action_seq_len + mask = torch.zeros((total_seq_len, total_seq_len), dtype=torch.bool, device=device) + + # video -> video + mask[:video_seq_len, :video_seq_len] = self.video_expert.build_video_to_video_mask( + video_seq_len=video_seq_len, + video_tokens_per_frame=video_tokens_per_frame, + device=device, + ) + # action -> action + mask[video_seq_len:, video_seq_len:] = True + # action -> first-frame video only + first_frame_tokens = min(video_tokens_per_frame, video_seq_len) + mask[video_seq_len:, :first_frame_tokens] = True + return mask + + def _compute_video_loss_per_sample( + self, + pred_video: torch.Tensor, + target_video: torch.Tensor, + image_is_pad: torch.Tensor | None, + include_initial_video_step: bool, + ) -> torch.Tensor: + video_loss_token = functional.mse_loss( + pred_video.float(), target_video.float(), reduction="none" + ).mean(dim=(1, 3, 4)) + if image_is_pad is None: + return video_loss_token.mean(dim=1) + + temporal_factor = int(self.vae.temporal_downsample_factor) + if temporal_factor <= 0: + raise ValueError(f"`vae.temporal_downsample_factor` must be positive, got {temporal_factor}.") + if image_is_pad.shape[1] < 1: + raise ValueError("`image_is_pad` must contain at least one frame.") + if (image_is_pad.shape[1] - 1) % temporal_factor != 0: + raise ValueError( + "Cannot align `image_is_pad` with video latent steps: " + f"num_frames={image_is_pad.shape[1]}, temporal_downsample_factor={temporal_factor}." + ) + + tail_is_pad = image_is_pad[:, 1:] + latent_tail_is_pad = tail_is_pad.view(image_is_pad.shape[0], -1, temporal_factor).all(dim=2) + if include_initial_video_step: + video_is_pad = torch.cat([image_is_pad[:, :1], latent_tail_is_pad], dim=1) + else: + video_is_pad = latent_tail_is_pad + + if video_is_pad.shape[1] != video_loss_token.shape[1]: + raise ValueError( + "Video-loss mask shape mismatch: " + f"mask steps={video_is_pad.shape[1]}, loss steps={video_loss_token.shape[1]}." + ) + + valid = (~video_is_pad).to(device=video_loss_token.device, dtype=video_loss_token.dtype) + valid_sum = valid.sum(dim=1).clamp(min=1.0) + return (video_loss_token * valid).sum(dim=1) / valid_sum + + def _sample_training_targets(self, inputs: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + input_latents = inputs["input_latents"] + batch_size = input_latents.shape[0] + action = inputs["action"] + noise_video = torch.randn_like(input_latents) + timestep_video = self.train_video_scheduler.sample_training_t( + batch_size=batch_size, + device=self.device, + dtype=input_latents.dtype, + ) + latents = self.train_video_scheduler.add_noise(input_latents, noise_video, timestep_video) + target_video = self.train_video_scheduler.training_target(input_latents, noise_video, timestep_video) + + if inputs["first_frame_latents"] is not None: + latents[:, :, 0:1] = inputs["first_frame_latents"] + noise_action = torch.randn_like(action) + timestep_action = self.train_action_scheduler.sample_training_t( + batch_size=batch_size, + device=self.device, + dtype=action.dtype, + ) + noisy_action = self.train_action_scheduler.add_noise(action, noise_action, timestep_action) + target_action = self.train_action_scheduler.training_target(action, noise_action, timestep_action) + return { + "latents": latents, + "target_video": target_video, + "noisy_action": noisy_action, + "target_action": target_action, + "timestep_video": timestep_video, + "timestep_action": timestep_action, + } + + def _run_training_mot(self, inputs: dict[str, torch.Tensor], targets: dict[str, torch.Tensor]): + video_pre = self.video_expert.pre_dit( + x=targets["latents"], + timestep=targets["timestep_video"], + context=inputs["context"], + context_mask=inputs["context_mask"], + action=inputs["action"], + fuse_vae_embedding_in_latents=inputs["fuse_vae_embedding_in_latents"], + ) + action_pre = self.action_expert.pre_dit( + action_tokens=targets["noisy_action"], + timestep=targets["timestep_action"], + context=inputs["context"], + context_mask=inputs["context_mask"], + ) + video_tokens = video_pre["tokens"] + action_tokens = action_pre["tokens"] + attention_mask = self._build_mot_attention_mask( + video_seq_len=video_tokens.shape[1], + action_seq_len=action_tokens.shape[1], + video_tokens_per_frame=int(video_pre["meta"]["tokens_per_frame"]), + device=video_tokens.device, + ) + tokens_out = self.mot( + embeds_all={ + "video": video_tokens, + "action": action_tokens, + }, + attention_mask=attention_mask, + freqs_all={ + "video": video_pre["freqs"], + "action": action_pre["freqs"], + }, + context_all={ + "video": { + "context": video_pre["context"], + "mask": video_pre["context_mask"], + }, + "action": { + "context": action_pre["context"], + "mask": action_pre["context_mask"], + }, + }, + t_mod_all={ + "video": video_pre["t_mod"], + "action": action_pre["t_mod"], + }, + ) + pred_video = self.video_expert.post_dit(tokens_out["video"], video_pre) + pred_action = self.action_expert.post_dit(tokens_out["action"], action_pre) + return pred_video, pred_action + + def _compute_training_video_loss(self, inputs, pred_video, target_video, timestep_video): + include_initial_video_step = inputs["first_frame_latents"] is None + if inputs["first_frame_latents"] is not None: + pred_video = pred_video[:, :, 1:] + target_video = target_video[:, :, 1:] + loss_video_per_sample = self._compute_video_loss_per_sample( + pred_video=pred_video, + target_video=target_video, + image_is_pad=inputs["image_is_pad"], + include_initial_video_step=include_initial_video_step, + ) + video_weight = self.train_video_scheduler.training_weight(timestep_video).to( + loss_video_per_sample.device, + dtype=loss_video_per_sample.dtype, + ) + return (loss_video_per_sample * video_weight).mean() + + def _compute_training_action_loss(self, inputs, pred_action, target_action, timestep_action): + action_loss_token = functional.mse_loss( + pred_action.float(), target_action.float(), reduction="none" + ).mean(dim=2) + if inputs["action_is_pad"] is not None: + valid = (~inputs["action_is_pad"]).to( + device=action_loss_token.device, + dtype=action_loss_token.dtype, + ) + valid_sum = valid.sum(dim=1).clamp(min=1.0) + action_loss_per_sample = (action_loss_token * valid).sum(dim=1) / valid_sum + else: + action_loss_per_sample = action_loss_token.mean(dim=1) + action_weight = self.train_action_scheduler.training_weight(timestep_action).to( + action_loss_per_sample.device, + dtype=action_loss_per_sample.dtype, + ) + return (action_loss_per_sample * action_weight).mean() + + def training_loss(self, sample, tiled: bool = False): + inputs = self.build_inputs(sample, tiled=tiled) + targets = self._sample_training_targets(inputs) + pred_video, pred_action = self._run_training_mot(inputs=inputs, targets=targets) + loss_video = self._compute_training_video_loss( + inputs=inputs, + pred_video=pred_video, + target_video=targets["target_video"], + timestep_video=targets["timestep_video"], + ) + loss_action = self._compute_training_action_loss( + inputs=inputs, + pred_action=pred_action, + target_action=targets["target_action"], + timestep_action=targets["timestep_action"], + ) + loss_total = self.loss_lambda_video * loss_video + self.loss_lambda_action * loss_action + loss_dict = { + "loss_video": self.loss_lambda_video * float(loss_video.detach().item()), + "loss_action": self.loss_lambda_action * float(loss_action.detach().item()), + } + return loss_total, loss_dict + + @torch.no_grad() + def _predict_joint_noise( + self, + latents_video: torch.Tensor, + latents_action: torch.Tensor, + timestep_video: torch.Tensor, + timestep_action: torch.Tensor, + context: torch.Tensor, + context_mask: torch.Tensor, + fuse_vae_embedding_in_latents: bool, + gt_action: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + video_pre = self.video_expert.pre_dit( + x=latents_video, + timestep=timestep_video, + context=context, + context_mask=context_mask, + action=gt_action, + fuse_vae_embedding_in_latents=fuse_vae_embedding_in_latents, + ) + action_pre = self.action_expert.pre_dit( + action_tokens=latents_action, + timestep=timestep_action, + context=context, + context_mask=context_mask, + ) + + attention_mask = self._build_mot_attention_mask( + video_seq_len=video_pre["tokens"].shape[1], + action_seq_len=action_pre["tokens"].shape[1], + video_tokens_per_frame=int(video_pre["meta"]["tokens_per_frame"]), + device=video_pre["tokens"].device, + ) + + tokens_out = self.mot( + embeds_all={ + "video": video_pre["tokens"], + "action": action_pre["tokens"], + }, + attention_mask=attention_mask, + freqs_all={ + "video": video_pre["freqs"], + "action": action_pre["freqs"], + }, + context_all={ + "video": { + "context": video_pre["context"], + "mask": video_pre["context_mask"], + }, + "action": { + "context": action_pre["context"], + "mask": action_pre["context_mask"], + }, + }, + t_mod_all={ + "video": video_pre["t_mod"], + "action": action_pre["t_mod"], + }, + ) + + pred_video = self.video_expert.post_dit(tokens_out["video"], video_pre) + pred_action = self.action_expert.post_dit(tokens_out["action"], action_pre) + return pred_video, pred_action + + @torch.no_grad() + def _predict_action_noise( + self, + first_frame_latents: torch.Tensor, + latents_action: torch.Tensor, + timestep_action: torch.Tensor, + context: torch.Tensor, + context_mask: torch.Tensor, + fuse_vae_embedding_in_latents: bool, + ) -> torch.Tensor: + timestep_video = torch.zeros_like( + timestep_action, dtype=first_frame_latents.dtype, device=self.device + ) + video_pre = self.video_expert.pre_dit( + x=first_frame_latents, + timestep=timestep_video, + context=context, + context_mask=context_mask, + action=None, + fuse_vae_embedding_in_latents=fuse_vae_embedding_in_latents, + ) + action_pre = self.action_expert.pre_dit( + action_tokens=latents_action, + timestep=timestep_action, + context=context, + context_mask=context_mask, + ) + + attention_mask = self._build_mot_attention_mask( + video_seq_len=video_pre["tokens"].shape[1], + action_seq_len=action_pre["tokens"].shape[1], + video_tokens_per_frame=int(video_pre["meta"]["tokens_per_frame"]), + device=video_pre["tokens"].device, + ) + tokens_out = self.mot( + embeds_all={ + "video": video_pre["tokens"], + "action": action_pre["tokens"], + }, + attention_mask=attention_mask, + freqs_all={ + "video": video_pre["freqs"], + "action": action_pre["freqs"], + }, + context_all={ + "video": { + "context": video_pre["context"], + "mask": video_pre["context_mask"], + }, + "action": { + "context": action_pre["context"], + "mask": action_pre["context_mask"], + }, + }, + t_mod_all={ + "video": video_pre["t_mod"], + "action": action_pre["t_mod"], + }, + ) + pred_action = self.action_expert.post_dit(tokens_out["action"], action_pre) + return pred_action + + @torch.no_grad() + def _predict_action_noise_with_cache( + self, + latents_action: torch.Tensor, + timestep_action: torch.Tensor, + context: torch.Tensor, + context_mask: torch.Tensor, + video_kv_cache: list[dict[str, torch.Tensor]], + attention_mask: torch.Tensor, + video_seq_len: int, + ) -> torch.Tensor: + action_pre = self.action_expert.pre_dit( + action_tokens=latents_action, + timestep=timestep_action, + context=context, + context_mask=context_mask, + ) + action_tokens = self.mot.forward_action_with_video_cache( + action_tokens=action_pre["tokens"], + action_freqs=action_pre["freqs"], + action_t_mod=action_pre["t_mod"], + action_context_payload={ + "context": action_pre["context"], + "mask": action_pre["context_mask"], + }, + video_kv_cache=video_kv_cache, + attention_mask=attention_mask, + video_seq_len=video_seq_len, + ) + return self.action_expert.post_dit(action_tokens, action_pre) + + def _normalize_infer_input_image( + self, + input_image: torch.Tensor, + num_video_frames: int | None = None, + ) -> tuple[torch.Tensor, int, int]: + if input_image.ndim == 3: + input_image = input_image.unsqueeze(0) + if input_image.ndim != 4 or input_image.shape[0] != 1 or input_image.shape[1] != 3: + raise ValueError( + f"`input_image` must have shape [1,3,H,W] or [3,H,W], got {tuple(input_image.shape)}" + ) + _, _, height, width = input_image.shape + if height % 16 != 0 or width % 16 != 0: + raise ValueError( + f"`input_image` must be resized before infer, expected multiples of 16 but got HxW=({height},{width})" + ) + if num_video_frames is not None: + checked_h, checked_w, checked_t = self._check_resize_height_width(height, width, num_video_frames) + if (checked_h, checked_w) != (height, width): + raise ValueError( + f"`input_image` must be resized before infer, expected multiples of 16 but got HxW=({height},{width})" + ) + if checked_t != num_video_frames: + raise ValueError(f"`num_video_frames` must satisfy T % 4 == 1, got {num_video_frames}") + return input_image, height, width + + def _normalize_infer_proprio(self, proprio: torch.Tensor | None) -> torch.Tensor | None: + if proprio is None: + return None + if self.proprio_dim is None: + raise ValueError( + "`proprio` was provided but `proprio_dim=None` so `proprio_encoder` is disabled." + ) + if proprio.ndim == 1: + proprio = proprio.unsqueeze(0) + elif proprio.ndim == 2 and proprio.shape[0] == 1: + pass + else: + raise ValueError(f"`proprio` must be [D] or [1,D], got shape {tuple(proprio.shape)}") + if proprio.shape[1] != self.proprio_dim: + raise ValueError(f"`proprio` last dim must be {self.proprio_dim}, got {proprio.shape[1]}") + return proprio.to(device=self.device, dtype=self.torch_dtype) + + def _prepare_infer_context(self, prompt, context, context_mask, proprio): + use_prompt = prompt is not None + use_context = context is not None or context_mask is not None + if use_prompt and use_context: + raise ValueError("`prompt` and `context/context_mask` are mutually exclusive.") + if not use_prompt and not use_context: + raise ValueError("Either `prompt` or both `context/context_mask` must be provided.") + if use_prompt: + context, context_mask = self.encode_prompt(prompt) + else: + context, context_mask = self._normalize_context_tensors(context, context_mask) + if proprio is not None: + context, context_mask = self._append_proprio_to_context( + context=context, + context_mask=context_mask, + proprio=proprio, + ) + return context, context_mask + + def _normalize_context_tensors(self, context, context_mask): + if context is None or context_mask is None: + raise ValueError("`context` and `context_mask` must be both provided together.") + if context.ndim == 2: + context = context.unsqueeze(0) + if context_mask.ndim == 1: + context_mask = context_mask.unsqueeze(0) + if context.ndim != 3 or context_mask.ndim != 2: + raise ValueError( + f"`context/context_mask` must be [B,L,D]/[B,L], got {tuple(context.shape)} and {tuple(context_mask.shape)}" + ) + context = context.to(device=self.device, dtype=self.torch_dtype, non_blocking=True) + context_mask = context_mask.to(device=self.device, dtype=torch.bool, non_blocking=True) + return context, context_mask + + def _make_action_latents(self, action_horizon: int, seed: int | None, rand_device: str): + generator = None if seed is None else torch.Generator(device=rand_device).manual_seed(seed) + return torch.randn( + (1, action_horizon, self.action_expert.action_dim), + generator=generator, + device=rand_device, + dtype=torch.float32, + ).to(device=self.device, dtype=self.torch_dtype) + + def _make_video_latents(self, num_video_frames: int, height: int, width: int, seed, rand_device): + latent_t = (num_video_frames - 1) // self.vae.temporal_downsample_factor + 1 + latent_h = height // self.vae.upsampling_factor + latent_w = width // self.vae.upsampling_factor + generator = None if seed is None else torch.Generator(device=rand_device).manual_seed(seed) + return torch.randn( + (1, self.vae.z_dim, latent_t, latent_h, latent_w), + generator=generator, + device=rand_device, + dtype=torch.float32, + ).to(device=self.device, dtype=self.torch_dtype) + + @torch.no_grad() + def infer_joint( + self, + prompt: str | None, + input_image: torch.Tensor, + num_video_frames: int, + action_horizon: int, + action: torch.Tensor + | None = None, # NOTE: this is gt action for conditioning videos, not for action expert + proprio: torch.Tensor | None = None, + context: torch.Tensor | None = None, + context_mask: torch.Tensor | None = None, + negative_prompt: str | None = None, + text_cfg_scale: float = 1.0, + num_inference_steps: int = 20, + sigma_shift: float | None = None, + seed: int | None = None, + rand_device: str = "cpu", + tiled: bool = False, + test_action_with_infer_action: bool = True, + ) -> dict[str, Any]: + self.eval() + if test_action_with_infer_action: + if seed is None: + raise ValueError("`test_action_with_infer_action=True` requires non-null `seed`.") + action_only_out = self.infer_action( + prompt=prompt, + input_image=input_image.clone(), + action_horizon=action_horizon, + context=context.clone() if context is not None else None, + context_mask=context_mask.clone() if context_mask is not None else None, + num_inference_steps=num_inference_steps, + sigma_shift=sigma_shift, + seed=seed, + rand_device=rand_device, + tiled=tiled, + proprio=proprio.clone() if proprio is not None else None, + )["action"] + + input_image, height, width = self._normalize_infer_input_image(input_image, num_video_frames) + if action is not None: + if action.ndim == 2: + action = action.unsqueeze(0) + if action.ndim != 3 or action.shape[0] != 1 or action.shape[1] != action_horizon: + # NOTE: This enforces action condition to have the same shape as action horizon to predict, which may be unnecessary + raise ValueError( + f"`action` must have shape [1, T, a_dim] or [T, a_dim], got {tuple(action.shape)} with action_horizon={action_horizon}" + ) + action = action.to(device=self.device, dtype=self.torch_dtype) + proprio = self._normalize_infer_proprio(proprio) + latents_video = self._make_video_latents(num_video_frames, height, width, seed, rand_device) + latents_action = self._make_action_latents(action_horizon, seed, rand_device) + + input_image = input_image.to(device=self.device, dtype=self.torch_dtype) + first_frame_latents = self._encode_input_image_latents_tensor(input_image=input_image, tiled=tiled) + latents_video[:, :, 0:1] = first_frame_latents.clone() + fuse_flag = bool(getattr(self.video_expert, "fuse_vae_embedding_in_latents", False)) + context, context_mask = self._prepare_infer_context(prompt, context, context_mask, proprio) + + infer_timesteps_video, infer_deltas_video = self.infer_video_scheduler.build_inference_schedule( + num_inference_steps=num_inference_steps, + device=self.device, + dtype=latents_video.dtype, + shift_override=sigma_shift, + ) + infer_timesteps_action, infer_deltas_action = self.infer_action_scheduler.build_inference_schedule( + num_inference_steps=num_inference_steps, + device=self.device, + dtype=latents_action.dtype, + shift_override=sigma_shift, + ) + for step_t_video, step_delta_video, step_t_action, step_delta_action in zip( + infer_timesteps_video, + infer_deltas_video, + infer_timesteps_action, + infer_deltas_action, + strict=True, + ): + timestep_video = step_t_video.unsqueeze(0).to(dtype=latents_video.dtype, device=self.device) + timestep_action = step_t_action.unsqueeze(0).to(dtype=latents_action.dtype, device=self.device) + + pred_video, pred_action = self._predict_joint_noise( + latents_video=latents_video, + latents_action=latents_action, + timestep_video=timestep_video, + timestep_action=timestep_action, + context=context, + context_mask=context_mask, + fuse_vae_embedding_in_latents=fuse_flag, + gt_action=action, + ) + + latents_video = self.infer_video_scheduler.step(pred_video, step_delta_video, latents_video) + latents_action = self.infer_action_scheduler.step(pred_action, step_delta_action, latents_action) + latents_video[:, :, 0:1] = first_frame_latents.clone() + + action_out = latents_action[0].detach().to(device="cpu", dtype=torch.float32) + if test_action_with_infer_action and not torch.allclose( + action_out, action_only_out, atol=1e-2, rtol=1e-2 + ): + max_abs_diff = (action_out - action_only_out).abs().max().item() + logger.warning( + f"Action from infer_joint and infer_action differ with max abs diff {max_abs_diff:.6f}. " + ) + + return { + "video": self._decode_latents(latents_video, tiled=tiled), + "action": action_out, + } + + @torch.no_grad() + def infer_action( + self, + prompt: str | None, + input_image: torch.Tensor, + action_horizon: int, + proprio: torch.Tensor | None = None, + context: torch.Tensor | None = None, + context_mask: torch.Tensor | None = None, + negative_prompt: str | None = None, + text_cfg_scale: float = 1.0, + num_inference_steps: int = 20, + sigma_shift: float | None = None, + seed: int | None = None, + rand_device: str = "cpu", + tiled: bool = False, + ) -> dict[str, Any]: + self.eval() + if str(getattr(self.video_expert, "video_attention_mask_mode", "")) != "first_frame_causal": + raise ValueError("`infer_action` requires `video_attention_mask_mode='first_frame_causal'`.") + + input_image, _, _ = self._normalize_infer_input_image(input_image) + proprio = self._normalize_infer_proprio(proprio) + latents_action = self._make_action_latents(action_horizon, seed, rand_device) + + input_image = input_image.to(device=self.device, dtype=self.torch_dtype) + first_frame_latents = self._encode_input_image_latents_tensor(input_image=input_image, tiled=tiled) + fuse_flag = bool(getattr(self.video_expert, "fuse_vae_embedding_in_latents", False)) + + context, context_mask = self._prepare_infer_context(prompt, context, context_mask, proprio) + + timestep_video = torch.zeros( + (first_frame_latents.shape[0],), + dtype=first_frame_latents.dtype, + device=self.device, + ) + video_pre = self.video_expert.pre_dit( + x=first_frame_latents, + timestep=timestep_video, + context=context, + context_mask=context_mask, + action=None, + fuse_vae_embedding_in_latents=fuse_flag, + ) + video_seq_len = int(video_pre["tokens"].shape[1]) + attention_mask = self._build_mot_attention_mask( + video_seq_len=video_seq_len, + action_seq_len=latents_action.shape[1], + video_tokens_per_frame=int(video_pre["meta"]["tokens_per_frame"]), + device=video_pre["tokens"].device, + ) + video_kv_cache = self.mot.prefill_video_cache( + video_tokens=video_pre["tokens"], + video_freqs=video_pre["freqs"], + video_t_mod=video_pre["t_mod"], + video_context_payload={ + "context": video_pre["context"], + "mask": video_pre["context_mask"], + }, + video_attention_mask=attention_mask[:video_seq_len, :video_seq_len], + ) + + infer_timesteps_action, infer_deltas_action = self.infer_action_scheduler.build_inference_schedule( + num_inference_steps=num_inference_steps, + device=self.device, + dtype=latents_action.dtype, + shift_override=sigma_shift, + ) + for step_t_action, step_delta_action in zip(infer_timesteps_action, infer_deltas_action, strict=True): + timestep_action = step_t_action.unsqueeze(0).to(dtype=latents_action.dtype, device=self.device) + + pred_action = self._predict_action_noise_with_cache( + latents_action=latents_action, + timestep_action=timestep_action, + context=context, + context_mask=context_mask, + video_kv_cache=video_kv_cache, + attention_mask=attention_mask, + video_seq_len=video_seq_len, + ) + + latents_action = self.infer_action_scheduler.step(pred_action, step_delta_action, latents_action) + + return { + "action": latents_action[0].detach().to(device="cpu", dtype=torch.float32), + } + + @torch.no_grad() + def infer( + self, + prompt: str | None, + input_image: torch.Tensor, + num_frames: int, + action: torch.Tensor | None = None, + action_horizon: int | None = None, + proprio: torch.Tensor | None = None, + context: torch.Tensor | None = None, + context_mask: torch.Tensor | None = None, + negative_prompt: str | None = None, + text_cfg_scale: float = 5.0, + action_cfg_scale: float = 1.0, + num_inference_steps: int = 20, + sigma_shift: float | None = None, + seed: int | None = None, + rand_device: str = "cpu", + tiled: bool = False, + ): + return self.infer_joint( + prompt=prompt, + input_image=input_image, + num_video_frames=num_frames, + action_horizon=action_horizon, + action=action, + proprio=proprio, + context=context, + context_mask=context_mask, + negative_prompt=negative_prompt, + text_cfg_scale=text_cfg_scale, + num_inference_steps=num_inference_steps, + sigma_shift=sigma_shift, + seed=seed, + rand_device=rand_device, + tiled=tiled, + ) + + def forward(self, *args, **kwargs): + return self.training_loss(*args, **kwargs) diff --git a/src/lerobot/policies/fastwam/wan/video_dit.py b/src/lerobot/policies/fastwam/wan/video_dit.py new file mode 100644 index 000000000..a98f06cde --- /dev/null +++ b/src/lerobot/policies/fastwam/wan/video_dit.py @@ -0,0 +1,800 @@ +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +from typing import Any + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as functional +from einops import rearrange + +from .model import ( + WanAttentionBlock, + WanLayerNorm, + WanModel, + WanRMSNorm, + rope_apply, + rope_params, + sinusoidal_embedding_1d, +) + +logger = logging.getLogger(__name__) + + +def get_sampling_sigmas(sampling_steps, shift): + # Vendored from Wan2.2 (formerly wan/utils/fm_solvers.py); computes the + # noise-level (sigma) schedule for Wan-compatible flow-matching inference. + sigma = np.linspace(1, 0, sampling_steps + 1)[:sampling_steps] + sigma = shift * sigma / (1 + (shift - 1) * sigma) + return sigma + + +def create_custom_forward(module): + def custom_forward(*inputs, **kwargs): + return module(*inputs, **kwargs) + + return custom_forward + + +def gradient_checkpoint_forward( + model, + use_gradient_checkpointing, + *args, + **kwargs, +): + if use_gradient_checkpointing: + model_output = torch.utils.checkpoint.checkpoint( + create_custom_forward(model), + *args, + **kwargs, + use_reentrant=False, + ) + else: + model_output = model(*args, **kwargs) + return model_output + + +def fastwam_masked_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + num_heads: int, + ctx_mask: torch.Tensor | None = None, + fp32_attention: bool = True, +) -> torch.Tensor: + """FastWAM masked attention wrapper for MoT masks and CPU test coverage. + + The official Wan attention implementation is still used as the source of + the projection/norm modules. This wrapper only replaces the final attention + kernel because FastWAM needs explicit boolean masks for video/action MoT + routing, while the upstream FlashAttention path accepts sequence lengths + but not arbitrary [query, key] masks. + """ + + q = rearrange(q, "b s (n d) -> b n s d", n=num_heads) + k = rearrange(k, "b s (n d) -> b n s d", n=num_heads) + v = rearrange(v, "b s (n d) -> b n s d", n=num_heads) + if fp32_attention: + q = q.float() + k = k.float() + v = v.float() + else: + q = q.to(dtype=v.dtype) + k = k.to(dtype=v.dtype) + x = functional.scaled_dot_product_attention(q, k, v, attn_mask=ctx_mask) + return rearrange(x, "b n s d -> b s (n d)", n=num_heads) + + +def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor): + return x * (1 + scale) + shift + + +class WanContinuousFlowMatchScheduler: + """Continuous-time Flow-Matching scheduler with shift-based Wan sampling.""" + + def __init__(self, num_train_timesteps: int = 1000, shift: float = 5.0, eps: float = 1e-10): + if num_train_timesteps <= 0: + raise ValueError(f"`num_train_timesteps` must be positive, got {num_train_timesteps}") + if shift <= 0: + raise ValueError(f"`shift` must be positive, got {shift}") + self.num_train_timesteps = int(num_train_timesteps) + self.shift = float(shift) + self.eps = float(eps) + self._y_min, self._weight_norm_const = self._precompute_training_weight_stats() + + @staticmethod + def _phi(u: torch.Tensor, shift: float) -> torch.Tensor: + return shift * u / (1.0 + (shift - 1.0) * u) + + def _precompute_training_weight_stats(self) -> tuple[float, float]: + steps = self.num_train_timesteps + u_grid = torch.linspace(1.0, 0.0, steps + 1, dtype=torch.float64)[:-1] + t_grid = self._phi(u_grid, self.shift) * float(steps) + y_grid = torch.exp(-2.0 * ((t_grid - (steps / 2.0)) / steps) ** 2) + y_min = float(y_grid.min().item()) + y_shifted_grid = y_grid - y_min + norm_const = float(y_shifted_grid.mean().item()) + return y_min, norm_const + + def sample_training_t(self, batch_size: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + if batch_size <= 0: + raise ValueError(f"`batch_size` must be positive, got {batch_size}") + u = torch.rand((batch_size,), device=device, dtype=torch.float32) + sigma = self._phi(u, self.shift) + timestep = sigma * float(self.num_train_timesteps) + return timestep.to(dtype=dtype) + + def training_weight(self, timestep: torch.Tensor) -> torch.Tensor: + t = timestep.to(dtype=torch.float32) + steps = float(self.num_train_timesteps) + y = torch.exp(-2.0 * ((t - (steps / 2.0)) / steps) ** 2) + y_shifted = y - self._y_min + weight = y_shifted / (self._weight_norm_const + self.eps) + if weight.numel() == 1: + return weight.reshape(()) + return weight + + def add_noise( + self, original_samples: torch.Tensor, noise: torch.Tensor, timestep: torch.Tensor + ) -> torch.Tensor: + sigma = (timestep / float(self.num_train_timesteps)).to( + original_samples.device, dtype=original_samples.dtype + ) + if sigma.ndim == 0: + return (1 - sigma) * original_samples + sigma * noise + sigma = sigma.view(-1, *([1] * (original_samples.ndim - 1))) + return (1 - sigma) * original_samples + sigma * noise + + @staticmethod + def training_target(sample: torch.Tensor, noise: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor: + del timestep + return noise - sample + + def build_inference_schedule( + self, + num_inference_steps: int, + device: torch.device, + dtype: torch.dtype, + shift_override: float | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if num_inference_steps <= 0: + raise ValueError(f"`num_inference_steps` must be positive, got {num_inference_steps}") + shift = self.shift if shift_override is None else float(shift_override) + if shift <= 0: + raise ValueError(f"`shift` must be positive, got {shift}") + + sigma_steps = torch.as_tensor( + get_sampling_sigmas(num_inference_steps, shift), + device=device, + dtype=torch.float32, + ) + timesteps = sigma_steps * float(self.num_train_timesteps) + sigma_next = torch.cat([sigma_steps[1:], sigma_steps.new_zeros(1)]) + deltas = sigma_next - sigma_steps + return timesteps.to(dtype=dtype), deltas.to(dtype=dtype) + + @staticmethod + def step(model_output: torch.Tensor, delta: torch.Tensor, sample: torch.Tensor) -> torch.Tensor: + delta = delta.to(sample.device, dtype=sample.dtype) + if delta.ndim == 0: + return sample + model_output * delta + delta = delta.view(-1, *([1] * (sample.ndim - 1))) + return sample + model_output * delta + + +def precompute_freqs_cis(dim: int, end: int = 1024, theta: float = 10000.0): + return rope_params(end, dim, theta) + + +def apply_dense_rope(x: torch.Tensor, freqs: torch.Tensor, num_heads: int) -> torch.Tensor: + x = rearrange(x, "b s (n d) -> b s n d", n=num_heads) + x_out = torch.view_as_complex(x.to(torch.float32).reshape(x.shape[0], x.shape[1], x.shape[2], -1, 2)) + freqs = freqs.to(torch.complex64) if freqs.device.type == "npu" else freqs + x_out = torch.view_as_real(x_out * freqs).flatten(2) + return x_out.to(x.dtype) + + +def _linear_input(linear: nn.Linear, x: torch.Tensor) -> torch.Tensor: + return x.to(dtype=linear.weight.dtype) + + +def _wan_layer_norm(norm: nn.Module, x: torch.Tensor) -> torch.Tensor: + if isinstance(norm, WanLayerNorm) and norm.weight is not None: + weight = norm.weight.float() + bias = norm.bias.float() if norm.bias is not None else None + return functional.layer_norm(x.float(), norm.normalized_shape, weight, bias, norm.eps).to( + dtype=x.dtype + ) + return norm(x) + + +def create_group_causal_attn_mask( + num_temporal_groups: int, num_query_per_group: int, num_key_per_group: int, mode: str = "causal" +) -> torch.Tensor: + if mode not in ["causal", "group_diagonal"]: + raise ValueError(f"`mode` must be 'causal' or 'group_diagonal', got {mode}.") + if num_temporal_groups <= 0: + raise ValueError(f"`num_temporal_groups` must be positive, got {num_temporal_groups}.") + if num_query_per_group <= 0: + raise ValueError(f"`num_query_per_group` must be positive, got {num_query_per_group}.") + if num_key_per_group <= 0: + raise ValueError(f"`num_key_per_group` must be positive, got {num_key_per_group}.") + + total_num_query_tokens = num_temporal_groups * num_query_per_group + total_num_key_tokens = num_temporal_groups * num_key_per_group + query_time_indices = torch.arange(num_temporal_groups).repeat_interleave(num_query_per_group).unsqueeze(1) + key_time_indices = torch.arange(num_temporal_groups).repeat_interleave(num_key_per_group).unsqueeze(0) + + if mode == "causal": + attn_mask = query_time_indices >= key_time_indices + else: + attn_mask = query_time_indices == key_time_indices + + if attn_mask.shape != (total_num_query_tokens, total_num_key_tokens): + raise RuntimeError("Attention mask shape mismatch.") + return attn_mask + + +class FastWAMAttentionBlock(WanAttentionBlock): + """Wan attention block with FastWAM's arbitrary boolean mask support.""" + + def __init__( + self, + hidden_dim: int, + attn_head_dim: int, + num_heads: int, + ffn_dim: int, + eps: float = 1e-6, + fp32_attention: bool = True, + ): + attention_dim = attn_head_dim * num_heads + if hidden_dim == attention_dim: + super().__init__( + dim=hidden_dim, + ffn_dim=ffn_dim, + num_heads=num_heads, + qk_norm=True, + cross_attn_norm=True, + eps=eps, + ) + else: + nn.Module.__init__(self) + self.dim = hidden_dim + self.ffn_dim = ffn_dim + self.num_heads = num_heads + self.qk_norm = True + self.cross_attn_norm = True + self.eps = eps + self.norm1 = WanLayerNorm(hidden_dim, eps) + self.self_attn = _FastWAMProjectedAttention(hidden_dim, attention_dim, num_heads, eps) + self.norm3 = WanLayerNorm(hidden_dim, eps, elementwise_affine=True) + self.cross_attn = _FastWAMProjectedAttention(hidden_dim, attention_dim, num_heads, eps) + self.norm2 = WanLayerNorm(hidden_dim, eps) + self.ffn = nn.Sequential( + nn.Linear(hidden_dim, ffn_dim), + nn.GELU(approximate="tanh"), + nn.Linear(ffn_dim, hidden_dim), + ) + self.modulation = nn.Parameter(torch.randn(1, 6, hidden_dim) / hidden_dim**0.5) + self.attn_head_dim = attn_head_dim + self.fp32_attention = bool(fp32_attention) + + @staticmethod + def split_modulation(block, t_mod: torch.Tensor): + has_seq = len(t_mod.shape) == 4 + chunk_dim = 2 if has_seq else 1 + + base_mod = block.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (base_mod + t_mod).chunk( + 6, dim=chunk_dim + ) + if has_seq: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + shift_msa.squeeze(2), + scale_msa.squeeze(2), + gate_msa.squeeze(2), + shift_mlp.squeeze(2), + scale_mlp.squeeze(2), + gate_mlp.squeeze(2), + ) + return shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp + + def project_self_attention( + self, x: torch.Tensor, freqs: torch.Tensor | dict[str, torch.Tensor] + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + q = self.self_attn.norm_q(self.self_attn.q(x)) + k = self.self_attn.norm_k(self.self_attn.k(x)) + v = self.self_attn.v(x) + if isinstance(freqs, dict): + b, s = x.shape[:2] + q = rope_apply( + q.view(b, s, self.num_heads, self.attn_head_dim), + freqs["grid_sizes"], + freqs["freqs"], + ).flatten(2) + k = rope_apply( + k.view(b, s, self.num_heads, self.attn_head_dim), + freqs["grid_sizes"], + freqs["freqs"], + ).flatten(2) + else: + q = apply_dense_rope(q, freqs, self.num_heads) + k = apply_dense_rope(k, freqs, self.num_heads) + return q, k, v + + def apply_cross_attention( + self, x: torch.Tensor, context: torch.Tensor, context_mask: torch.Tensor | None = None + ) -> torch.Tensor: + if context_mask is not None and context_mask.dim() == 3: + context_mask = context_mask.unsqueeze(1) + attn = self.cross_attn + b, n, d = x.size(0), attn.num_heads, attn.head_dim + q = attn.norm_q(attn.q(x)).view(b, -1, n * d) + k = attn.norm_k(attn.k(context)).view(b, -1, n * d) + v = attn.v(context).view(b, -1, n * d) + x = fastwam_masked_attention( + q=q, + k=k, + v=v, + num_heads=n, + ctx_mask=context_mask, + fp32_attention=self.fp32_attention, + ) + return attn.o(_linear_input(attn.o, x)) + + def project_self_attention_output(self, x: torch.Tensor) -> torch.Tensor: + return self.self_attn.o(_linear_input(self.self_attn.o, x)) + + def apply_norm1(self, x: torch.Tensor) -> torch.Tensor: + return _wan_layer_norm(self.norm1, x) + + def apply_norm2(self, x: torch.Tensor) -> torch.Tensor: + return _wan_layer_norm(self.norm2, x) + + def apply_norm3(self, x: torch.Tensor) -> torch.Tensor: + return _wan_layer_norm(self.norm3, x) + + def forward( + self, + x: torch.Tensor, + context: torch.Tensor, + t_mod: torch.Tensor, + freqs: torch.Tensor, + context_mask: torch.Tensor | None = None, + self_attn_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.split_modulation(self, t_mod) + residual_x = x + attn_input = modulate(self.apply_norm1(x), shift_msa, scale_msa) + q, k, v = self.project_self_attention(attn_input, freqs) + y = fastwam_masked_attention( + q=q, + k=k, + v=v, + num_heads=self.num_heads, + ctx_mask=self_attn_mask, + fp32_attention=self.fp32_attention, + ) + x = residual_x + gate_msa * self.project_self_attention_output(y) + x = x + self.apply_cross_attention(self.apply_norm3(x), context, context_mask=context_mask) + mlp_input = modulate(self.apply_norm2(x), shift_mlp, scale_mlp) + return x + gate_mlp * self.ffn(mlp_input) + + +class _FastWAMProjectedAttention(nn.Module): + def __init__(self, hidden_dim: int, attention_dim: int, num_heads: int, eps: float): + super().__init__() + self.dim = hidden_dim + self.num_heads = num_heads + self.head_dim = attention_dim // num_heads + self.q = nn.Linear(hidden_dim, attention_dim) + self.k = nn.Linear(hidden_dim, attention_dim) + self.v = nn.Linear(hidden_dim, attention_dim) + self.o = nn.Linear(attention_dim, hidden_dim) + self.norm_q = WanRMSNorm(attention_dim, eps=eps) + self.norm_k = WanRMSNorm(attention_dim, eps=eps) + + +class WanVideoDiT(WanModel): + def __init__( + self, + hidden_dim: int, + in_dim: int, + ffn_dim: int, + out_dim: int, + text_dim: int, + freq_dim: int, + eps: float, + patch_size: tuple[int, int, int], + num_heads: int, + attn_head_dim: int, + num_layers: int, + has_image_input: bool = False, + has_image_pos_emb: bool = False, + has_ref_conv: bool = False, + add_control_adapter: bool = False, + in_dim_control_adapter: int = 24, + seperated_timestep: bool = False, + require_vae_embedding: bool = False, + require_clip_embedding: bool = False, + fuse_vae_embedding_in_latents: bool = True, + action_conditioned: bool = False, + action_dim: int = 7, + action_group_causal_mask_mode="causal", + video_attention_mask_mode: str = "bidirectional", + use_gradient_checkpointing: bool = False, + fp32_attention: bool = True, + ): + del in_dim_control_adapter + if has_image_input: + raise ValueError("FastWAM currently expects Wan2.2 TI2V latents with fused image conditioning.") + if has_image_pos_emb: + raise ValueError("FastWAM does not support extra image positional embeddings in WanVideoDiT.") + if has_ref_conv: + raise ValueError("FastWAM does not support reference convolutions in WanVideoDiT.") + if add_control_adapter: + raise ValueError("FastWAM does not support control adapters in WanVideoDiT.") + if require_clip_embedding: + raise ValueError("FastWAM does not support CLIP embedding conditioning in WanVideoDiT.") + if require_vae_embedding or not fuse_vae_embedding_in_latents: + raise ValueError("FastWAM expects VAE conditioning to be fused in latents.") + if attn_head_dim != hidden_dim // num_heads: + raise ValueError( + "`attn_head_dim` must match the upstream Wan head dimension `hidden_dim // num_heads`; " + f"got {attn_head_dim} vs {hidden_dim // num_heads}." + ) + + super().__init__( + model_type="ti2v", + patch_size=patch_size, + text_len=512, + in_dim=in_dim, + dim=hidden_dim, + ffn_dim=ffn_dim, + freq_dim=freq_dim, + text_dim=text_dim, + out_dim=out_dim, + num_heads=num_heads, + num_layers=num_layers, + qk_norm=True, + cross_attn_norm=True, + eps=eps, + ) + self.blocks = torch.nn.ModuleList( + [ + FastWAMAttentionBlock( + hidden_dim=hidden_dim, + attn_head_dim=attn_head_dim, + num_heads=num_heads, + ffn_dim=ffn_dim, + eps=eps, + fp32_attention=fp32_attention, + ) + for _ in range(num_layers) + ] + ) + self.init_weights() + + self.hidden_dim = hidden_dim + self.attn_head_dim = attn_head_dim + self.seperated_timestep = seperated_timestep + self.fuse_vae_embedding_in_latents = fuse_vae_embedding_in_latents + self.video_attention_mask_mode = str(video_attention_mask_mode) + self.action_conditioned = action_conditioned + self.action_dim = action_dim + self.fp32_attention = bool(fp32_attention) + + if self.action_conditioned: + self.action_embedding = torch.nn.Linear(action_dim, hidden_dim) + self.action_group_causal_mask_mode = action_group_causal_mask_mode + + self.use_gradient_checkpointing = use_gradient_checkpointing + if self.use_gradient_checkpointing: + logger.info( + "Using gradient checkpointing for DiT blocks. This will save memory but use more computation." + ) + + def patchify(self, x: torch.Tensor): + return self.patch_embedding(x) + + def _validate_forward_inputs( + self, + x: torch.Tensor, + timestep: torch.Tensor, + context: torch.Tensor, + context_mask: torch.Tensor | None, + action: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if x.ndim != 5: + raise ValueError(f"`latents` must be 5D [B, C, T, H, W], got shape {tuple(x.shape)}") + num_latent_frames = x.shape[2] + if context.ndim != 3: + raise ValueError(f"`context` must be 3D [B, L, D], got shape {tuple(context.shape)}") + if timestep.ndim != 1: + raise ValueError(f"`timestep` must be 1D [B] or [1], got shape {tuple(timestep.shape)}") + if self.action_conditioned: + allow_text_only_single_frame = num_latent_frames == 1 and action is None + if not allow_text_only_single_frame: + if action is None: + raise ValueError("Action input is required for action-conditioned model.") + if action.ndim != 3: + raise ValueError( + f"`action` must be 3D [B, action_horizon, action_dim], got shape {tuple(action.shape)}" + ) + if action.shape[2] != self.action_dim: + raise ValueError( + f"`action` last dimension must be {self.action_dim}, got {action.shape[2]}" + ) + if num_latent_frames <= 1: + raise ValueError( + f"video length must be > 1 for action-conditioned model, got {num_latent_frames}" + ) + if action.shape[1] % (num_latent_frames - 1) != 0: + raise ValueError( + "action horizon must be divisible by (num_latent_frames - 1), " + f"got action_horizon={action.shape[1]}" + ) + if context_mask is None: + context_mask = torch.ones( + (context.shape[0], context.shape[1]), dtype=torch.bool, device=context.device + ) + else: + if context_mask.ndim != 2: + raise ValueError(f"`context_mask` must be 2D [B, L], got shape {tuple(context_mask.shape)}") + if context_mask.shape[0] != context.shape[0] or context_mask.shape[1] != context.shape[1]: + raise ValueError( + "`context_mask` shape must match `context` shape [B, L], " + f"got {tuple(context_mask.shape)} vs {tuple(context.shape)}" + ) + + batch_size = x.shape[0] + if batch_size != context.shape[0]: + if not self.training and batch_size == 1: + x = x.expand(context.shape[0], -1, -1, -1, -1) + batch_size = context.shape[0] + else: + raise ValueError( + f"Batch mismatch between latents and context: {batch_size} vs {context.shape[0]}." + ) + + if timestep.shape[0] not in (1, batch_size): + raise ValueError( + f"`timestep` length must be 1 or batch_size({batch_size}), got {timestep.shape[0]}" + ) + if timestep.shape[0] == 1 and batch_size > 1: + if self.training: + raise ValueError("During training, timestep length must match batch_size.") + timestep = timestep.expand(batch_size) + return x, timestep, context_mask + + def build_video_to_video_mask( + self, + video_seq_len: int, + video_tokens_per_frame: int, + device: torch.device, + ) -> torch.Tensor: + if video_seq_len <= 0: + raise ValueError(f"`video_seq_len` must be positive, got {video_seq_len}") + if video_tokens_per_frame <= 0: + raise ValueError(f"`video_tokens_per_frame` must be positive, got {video_tokens_per_frame}") + + if self.video_attention_mask_mode == "bidirectional": + return torch.ones((video_seq_len, video_seq_len), dtype=torch.bool, device=device) + + if self.video_attention_mask_mode == "per_frame_causal": + if video_seq_len % video_tokens_per_frame != 0: + raise ValueError( + "`video_seq_len` must be divisible by `video_tokens_per_frame` in `per_frame_causal` mode, " + f"got {video_seq_len} and {video_tokens_per_frame}" + ) + num_video_frames = video_seq_len // video_tokens_per_frame + frame_causal = torch.tril( + torch.ones((num_video_frames, num_video_frames), dtype=torch.bool, device=device) + ) + return frame_causal.repeat_interleave(video_tokens_per_frame, dim=0).repeat_interleave( + video_tokens_per_frame, dim=1 + ) + + if self.video_attention_mask_mode == "first_frame_causal": + video_mask = torch.ones((video_seq_len, video_seq_len), dtype=torch.bool, device=device) + first_frame_tokens = min(video_tokens_per_frame, video_seq_len) + video_mask[:first_frame_tokens, first_frame_tokens:] = False + return video_mask + + raise ValueError(f"Unsupported video attention mask mode: {self.video_attention_mask_mode}") + + def pre_dit( + self, + x: torch.Tensor, + timestep: torch.Tensor, + context: torch.Tensor, + context_mask: torch.Tensor | None = None, + action: torch.Tensor | None = None, + fuse_vae_embedding_in_latents: bool = False, + ) -> dict[str, Any]: + x, timestep, context_mask = self._validate_forward_inputs( + x=x, + timestep=timestep, + context=context, + context_mask=context_mask, + action=action, + ) + model_dtype = self.patch_embedding.weight.dtype + x = x.to(dtype=model_dtype) + context = context.to(dtype=model_dtype) + if action is not None: + action = action.to(dtype=model_dtype) + + batch_size = x.shape[0] + patch_h = int(self.patch_size[1]) + patch_w = int(self.patch_size[2]) + if x.shape[3] % patch_h != 0 or x.shape[4] % patch_w != 0: + raise ValueError( + "Latent spatial shape must be divisible by DiT patch size, " + f"got HxW=({x.shape[3]}, {x.shape[4]}), patch=({patch_h}, {patch_w})" + ) + tokens_per_frame = (x.shape[3] // patch_h) * (x.shape[4] // patch_w) + + if not (self.seperated_timestep and fuse_vae_embedding_in_latents): + raise NotImplementedError( + "FastWAM currently requires separated timesteps with fused VAE latents." + ) + + token_timesteps = torch.ones( + (batch_size, x.shape[2], tokens_per_frame), + dtype=model_dtype, + device=timestep.device, + ) * timestep.to(dtype=model_dtype).view(batch_size, 1, 1) + token_timesteps[:, 0, :] = 0 + token_timesteps = token_timesteps.reshape(batch_size, -1) + # Wan keeps the time embedding in fp32: the AdaLN modulation in the vendored + # Head/Block asserts e.dtype == float32 (numerical stability of the scale/shift). + # Upstream guarantees this via an fp32 autocast region, so it holds even when the + # model runs in bf16. Mirror that here, then cast the per-block modulation back to + # model_dtype so the bf16 attention blocks are not upcast to fp32. + with torch.amp.autocast("cuda", dtype=torch.float32): + token_t_emb = sinusoidal_embedding_1d(self.freq_dim, token_timesteps.reshape(-1)).float() + t = self.time_embedding(token_t_emb).reshape(batch_size, -1, self.hidden_dim) + t_mod = self.time_projection(t).unflatten(2, (6, self.hidden_dim)) + t_mod = t_mod.to(dtype=model_dtype) + + x = self.patchify(x) + f, h, w = x.shape[2:] + + context = self.text_embedding(context) + context_len = context.shape[1] + if self.action_conditioned and action is not None: + action_len = action.shape[1] + action_emb = self.action_embedding(action) + action_pos_embed = sinusoidal_embedding_1d( + self.hidden_dim, torch.arange(action_len, device=action_emb.device) + ).to(dtype=action_emb.dtype) + action_emb = action_emb + action_pos_embed.unsqueeze(0) + context = torch.cat([context, action_emb], dim=1) + + num_temporal_groups = f - 1 + if num_temporal_groups <= 0: + raise ValueError( + "Action-conditioned context mask requires at least 2 latent frames when `action` is provided." + ) + if action_emb.shape[1] % num_temporal_groups != 0: + raise ValueError( + f"Action embedding length {action_emb.shape[1]} must be divisible by " + f"number of temporal groups {num_temporal_groups}" + ) + action_group_mask = create_group_causal_attn_mask( + num_temporal_groups=num_temporal_groups, + num_query_per_group=tokens_per_frame, + num_key_per_group=action_len // num_temporal_groups, + mode=self.action_group_causal_mask_mode, + ).to(context.device) + + seq_len = f * h * w + final_context_mask = torch.zeros( + (batch_size, seq_len, context.shape[1]), dtype=torch.bool, device=context.device + ) + final_context_mask[:, :, :context_len] = context_mask.unsqueeze(1).expand(-1, seq_len, -1) + final_context_mask[:, tokens_per_frame:, context_len:] = action_group_mask.unsqueeze(0).expand( + batch_size, -1, -1 + ) + context_mask = final_context_mask + elif self.action_conditioned and action is None: + if f != 1: + raise ValueError( + "Action-conditioned model requires `action` unless running single-frame text-only mode " + "with num_latent_frames=1." + ) + context_mask = context_mask.unsqueeze(1).expand(-1, f * h * w, -1) + else: + context_mask = context_mask.unsqueeze(1).expand(-1, f * h * w, -1) + + x_tokens = rearrange(x, "b c f h w -> b (f h w) c").contiguous() + grid_sizes = torch.tensor([[f, h, w]] * batch_size, dtype=torch.long, device=x_tokens.device) + freqs = {"grid_sizes": grid_sizes, "freqs": self.freqs.to(x_tokens.device)} + + return { + "tokens": x_tokens, + "freqs": freqs, + "t": t, + "t_mod": t_mod, + "context": context, + "context_mask": context_mask, + "meta": { + "grid_sizes": grid_sizes, + "tokens_per_frame": tokens_per_frame, + "batch_size": batch_size, + }, + } + + def post_dit(self, x_tokens: torch.Tensor, pre_state: dict[str, Any]) -> torch.Tensor: + x = self.head(x_tokens, pre_state["t"]) + return torch.stack(super().unpatchify(x, pre_state["meta"]["grid_sizes"])) + + def forward( + self, + x: torch.Tensor, + timestep: torch.Tensor, + context: torch.Tensor, + context_mask: torch.Tensor | None = None, + action: torch.Tensor | None = None, + fuse_vae_embedding_in_latents: bool = False, + ): + pre_state = self.pre_dit( + x=x, + timestep=timestep, + context=context, + context_mask=context_mask, + action=action, + fuse_vae_embedding_in_latents=fuse_vae_embedding_in_latents, + ) + x_tokens = pre_state["tokens"] + context_emb = pre_state["context"] + t_mod = pre_state["t_mod"] + freqs = pre_state["freqs"] + context_attn_mask = pre_state["context_mask"] + self_attn_mask = ( + self.build_video_to_video_mask( + video_seq_len=x_tokens.shape[1], + video_tokens_per_frame=int(pre_state["meta"]["tokens_per_frame"]), + device=x_tokens.device, + ) + if self.video_attention_mask_mode != "bidirectional" + else None + ) + + for block in self.blocks: + if self.use_gradient_checkpointing: + x_tokens = gradient_checkpoint_forward( + block, + self.use_gradient_checkpointing, + x_tokens, + context_emb, + t_mod, + freqs, + context_mask=context_attn_mask, + self_attn_mask=self_attn_mask, + ) + else: + x_tokens = block( + x_tokens, + context_emb, + t_mod, + freqs, + context_mask=context_attn_mask, + self_attn_mask=self_attn_mask, + ) + + return self.post_dit(x_tokens, pre_state) diff --git a/tests/policies/fastwam/test_fastwam_policy.py b/tests/policies/fastwam/test_fastwam_policy.py new file mode 100644 index 000000000..05e86b7f4 --- /dev/null +++ b/tests/policies/fastwam/test_fastwam_policy.py @@ -0,0 +1,391 @@ +#!/usr/bin/env python + +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json + +import pytest +import torch +from safetensors import safe_open +from torch import nn + +pytest.importorskip("transformers", reason="fastwam requires the `fastwam` extra (transformers)") +pytest.importorskip("diffusers", reason="fastwam requires the `fastwam` extra (diffusers)") + +from lerobot.configs import FeatureType, PolicyFeature, PreTrainedConfig +from lerobot.policies import FastWAMConfig, get_policy_class, make_policy_config, make_pre_post_processors +from lerobot.policies.fastwam.modeling_fastwam import FastWAMPolicy +from lerobot.policies.fastwam.processor_fastwam import FastWAMActionToggleProcessorStep +from lerobot.utils.constants import ACTION, OBS_STATE + + +class FakeFastWAMCore(nn.Module): + def __init__(self): + super().__init__() + self.dit = nn.Linear(2, 2) + + def training_loss(self, sample): + assert sample["video"].ndim == 5 + assert sample["context"].ndim == 3 + return sample[ACTION].sum() * 0.0 + torch.tensor(1.0), {"loss_action": 1.0} + + def infer_action(self, **kwargs): + return {"action": torch.ones(1, kwargs["action_horizon"], 3)} + + +def test_fastwam_is_registered_and_publicly_exported(): + cfg = make_policy_config( + "fastwam", + action_dim=3, + proprio_dim=2, + action_horizon=4, + n_action_steps=2, + num_video_frames=5, + action_video_freq_ratio=1, + base_model_id=None, + ) + + assert isinstance(cfg, FastWAMConfig) + assert cfg.type == "fastwam" + assert get_policy_class("fastwam") is FastWAMPolicy + + +def test_config_validates_features_model_ids_and_saved_auto_route(tmp_path): + cfg = FastWAMConfig() + cfg.save_pretrained(tmp_path) + saved = json.loads((tmp_path / "config.json").read_text()) + + assert saved["pretrained_path"] is None + assert cfg.image_features["observation.images.image"].type == FeatureType.VISUAL + assert cfg.action_feature.shape == (7,) + assert cfg.robot_state_feature.shape == (8,) + with pytest.raises(ValueError, match="image feature"): + FastWAMConfig(input_features={OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(8,))}) + assert FastWAMConfig(tokenizer_model_id="somebody/other-tokenizer").tokenizer_model_id == ( + "somebody/other-tokenizer" + ) + + +def test_preprocessor_passes_images_through_and_postprocessor_toggles_actions(tmp_path): + cfg = FastWAMConfig( + action_dim=3, + proprio_dim=2, + action_horizon=4, + n_action_steps=2, + num_video_frames=5, + action_video_freq_ratio=1, + image_size=(2, 2), + device="cpu", + toggle_action_dimensions=[-1], + input_features={ + "observation.images.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 2, 2)), + OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(2,)), + }, + output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(3,))}, + base_model_id=None, + ) + dataset_stats = { + "observation.images.image": { + "mean": torch.full((3, 1, 1), 0.2), + "std": torch.full((3, 1, 1), 0.1), + }, + OBS_STATE: { + "mean": torch.tensor([1.0, 3.0]), + "std": torch.tensor([2.0, 4.0]), + }, + ACTION: { + "mean": torch.zeros(3), + "std": torch.ones(3), + }, + } + + preprocessor, postprocessor = make_pre_post_processors(cfg, dataset_stats=dataset_stats) + processed = preprocessor( + { + "observation.images.image": torch.tensor( + [ + [[0.0, 0.5], [1.0, 0.5]], + [[0.0, 0.5], [1.0, 0.5]], + [[0.0, 0.5], [1.0, 0.5]], + ] + ), + OBS_STATE: torch.tensor([3.0, 7.0]), + } + ) + preprocessor.save_pretrained(tmp_path, config_filename="policy_preprocessor.json") + postprocessor.save_pretrained(tmp_path, config_filename="policy_postprocessor.json") + _, loaded_postprocessor = make_pre_post_processors(cfg, pretrained_path=str(tmp_path)) + + # VISUAL normalization is IDENTITY + expected_image = torch.tensor( + [[[[0.0, 0.5], [1.0, 0.5]], [[0.0, 0.5], [1.0, 0.5]], [[0.0, 0.5], [1.0, 0.5]]]] + ) + assert preprocessor.name == "policy_preprocessor" + assert postprocessor.name == "policy_postprocessor" + assert torch.allclose(processed["observation.images.image"], expected_image) + assert torch.allclose(processed[OBS_STATE], torch.tensor([[1.0, 1.0]])) + assert torch.equal(dataset_stats["observation.images.image"]["mean"], torch.full((3, 1, 1), 0.2)) + assert any(isinstance(step, FastWAMActionToggleProcessorStep) for step in loaded_postprocessor.steps) + assert torch.equal( + loaded_postprocessor(torch.tensor([[0.25, 0.5, 1.0]])), torch.tensor([[0.25, 0.5, -1.0]]) + ) + + +def test_policy_forward_and_predict_action_adapt_lerobot_batches(monkeypatch): + captured = [] + + class CapturingCore(FakeFastWAMCore): + def infer_action(self, **kwargs): + captured.append( + { + "image_shape": tuple(kwargs["input_image"].shape), + "proprio_shape": tuple(kwargs["proprio"].shape), + "prompt": kwargs["prompt"], + } + ) + return {"action": torch.full((1, kwargs["action_horizon"], 3), float(len(captured)))} + + monkeypatch.setattr(FastWAMPolicy, "_build_core_model", lambda self, config: CapturingCore()) + cfg = FastWAMConfig( + action_dim=3, + proprio_dim=2, + action_horizon=4, + n_action_steps=2, + num_video_frames=5, + action_video_freq_ratio=1, + image_size=(16, 16), + input_features={ + "observation.images.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 16, 16)), + OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(2,)), + }, + output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(3,))}, + base_model_id=None, + ) + policy = FastWAMPolicy(cfg) + + loss, metrics = policy.forward( + { + "observation.images.image": torch.zeros(1, 3, 16, 16), + OBS_STATE: torch.zeros(1, 2), + ACTION: torch.zeros(1, 4, 3), + "context": torch.zeros(1, 5, 4096), + "context_mask": torch.ones(1, 5, dtype=torch.bool), + } + ) + action = policy.predict_action_chunk( + { + "observation.images.image": torch.stack( + [ + torch.zeros(3, 16, 16), + torch.ones(3, 16, 16), + ] + ), + OBS_STATE: torch.tensor([[0.0, 1.0], [2.0, 3.0]]), + "task": ["task 0", "task 1"], + } + ) + + assert loss.item() == 1.0 + assert metrics["loss_action"] == 1.0 + assert action.shape == (2, 4, 3) + assert action[:, 0, 0].tolist() == [1.0, 2.0] + assert [item["image_shape"] for item in captured] == [(1, 3, 16, 16), (1, 3, 16, 16)] + assert [item["proprio_shape"] for item in captured] == [(1, 2), (1, 2)] + assert [item["prompt"] for item in captured] == [ + cfg.prompt_template.format(task="task 0"), + cfg.prompt_template.format(task="task 1"), + ] + + +class CoreWithFrozenComponents(FakeFastWAMCore): + """Fake core mirroring the real one: frozen VAE / text encoder held as + *unregistered* attributes (via `object.__setattr__`) so they are excluded from + `state_dict()` and the saved checkpoint, but still moved by the `_apply` override.""" + + def __init__(self): + super().__init__() + object.__setattr__(self, "vae", nn.Linear(2, 2)) + object.__setattr__(self, "text_encoder", nn.Linear(2, 2)) + self.vae.requires_grad_(False) + self.text_encoder.requires_grad_(False) + + def _apply(self, fn, *args, **kwargs): + super()._apply(fn, *args, **kwargs) + self.vae._apply(fn) + self.text_encoder._apply(fn) + return self + + +def test_from_pretrained_uses_base_loader_and_skips_wan_backbone(monkeypatch, tmp_path): + cfg = FastWAMConfig( + action_dim=3, + proprio_dim=2, + action_horizon=4, + n_action_steps=2, + num_video_frames=5, + action_video_freq_ratio=1, + base_model_id=None, + ) + + def build_core(self, config): + core = CoreWithFrozenComponents() + with torch.no_grad(): + core.dit.weight.fill_(0.5) + return core + + monkeypatch.setattr(FastWAMPolicy, "_build_core_model", build_core) + + reference = FastWAMPolicy(cfg) + with torch.no_grad(): + reference.model.dit.weight.fill_(1.25) # a distinctive, trained-looking weight + reference.save_pretrained(tmp_path) + + # Building from Wan2.2 must never happen on a checkpoint load. + def fail_if_wan_pretrained_is_loaded(*args, **kwargs): + raise AssertionError("from_pretrained must not initialize or download the Wan2.2 backbone") + + monkeypatch.setattr( + "lerobot.policies.fastwam.wan.modular.FastWAM.from_wan22_pretrained", + fail_if_wan_pretrained_is_loaded, + ) + + policy = FastWAMPolicy.from_pretrained(tmp_path) + + assert isinstance(policy.model, CoreWithFrozenComponents) + # The bundled checkpoint weights overwrote the freshly built (0.5) DiT weights. + assert torch.allclose(policy.model.dit.weight, torch.full_like(policy.model.dit.weight, 1.25)) + + +def test_save_pretrained_excludes_frozen_components(monkeypatch, tmp_path): + cfg = FastWAMConfig( + action_dim=3, + proprio_dim=2, + action_horizon=4, + n_action_steps=2, + num_video_frames=5, + action_video_freq_ratio=1, + base_model_id=None, + ) + monkeypatch.setattr(FastWAMPolicy, "_build_core_model", lambda self, config: CoreWithFrozenComponents()) + policy = FastWAMPolicy(cfg) + + save_dir = tmp_path / "saved" + policy.save_pretrained(save_dir) + + assert (save_dir / "model.safetensors").is_file() + # No Wan sidecar files either: the frozen backbone comes from the diffusers repo. + assert not (save_dir / "Wan2.2_VAE.safetensors").exists() + assert not (save_dir / "google").exists() + + with safe_open(save_dir / "model.safetensors", framework="pt") as f: + keys = set(f.keys()) + # Lean checkpoint: only the trainable DiT is saved; the frozen VAE / UMT5 text + # encoder are excluded (loaded from the diffusers/transformers repos at init). + assert any(key.startswith("model.dit.") for key in keys) + assert not any(key.startswith("model.vae.") for key in keys) + assert not any(key.startswith("model.text_encoder.") for key in keys) + + +def test_frozen_components_excluded_from_params_but_follow_device_moves(monkeypatch): + cfg = FastWAMConfig( + action_dim=3, + proprio_dim=2, + action_horizon=4, + n_action_steps=2, + num_video_frames=5, + action_video_freq_ratio=1, + base_model_id=None, + ) + monkeypatch.setattr(FastWAMPolicy, "_build_core_model", lambda self, config: CoreWithFrozenComponents()) + policy = FastWAMPolicy(cfg) + + # Unregistered: excluded from state_dict and from the optimizer's parameter set. + sd = policy.state_dict() + assert not any(k.startswith("model.vae.") or k.startswith("model.text_encoder.") for k in sd) + param_names = [n for n, _ in policy.named_parameters()] + assert not any("vae" in n or "text_encoder" in n for n in param_names) + + # ...but the `_apply` override still carries them through `.to()` (dtype stands in + # for device on a CPU box), so they never strand off the rest of the model. + policy.to(torch.float64) + assert policy.model.dit.weight.dtype == torch.float64 # registered + assert policy.model.vae.weight.dtype == torch.float64 # unregistered, moved via _apply + assert policy.model.text_encoder.weight.dtype == torch.float64 + + +def test_pretrained_config_round_trips_fastwam_features(tmp_path): + cfg = FastWAMConfig(action_dim=7, proprio_dim=8, image_size=(224, 448), base_model_id=None) + cfg.save_pretrained(tmp_path) + + loaded = PreTrainedConfig.from_pretrained(tmp_path) + + assert loaded.type == "fastwam" + assert loaded.image_features["observation.images.image"].type == FeatureType.VISUAL + assert loaded.action_feature.shape == (7,) + assert loaded.robot_state_feature.shape == (8,) + + +def test_vae_adapter_empty_build_encode_decode_shapes(): + """Offline glue check of the diffusers-backed VAE adapter (random weights). + + Validates the encode/decode contract — 48 latent channels, 16x spatial / 4x + temporal compression, list-or-batch input, scaling round-trip — without any + weight download. (Numerical fidelity vs the original Wan VAE is a separate, + GPU + real-weights verification step.) + """ + pytest.importorskip("diffusers") + from diffusers import AutoencoderKLWan + + from lerobot.policies.fastwam.wan import WanVideoVAE38 + + # Production always loads a real pretrained VAE from the diffusers repo; here we + # build the same architecture with random weights and dummy standardization stats + # to exercise the adapter's shape/scaling contract offline (fidelity is checked + # separately, with real weights, on GPU). + arch = { + "base_dim": 160, + "decoder_base_dim": 256, + "z_dim": 48, + "dim_mult": [1, 2, 4, 4], + "num_res_blocks": 2, + "attn_scales": [], + "temporal_downsample": [False, True, True], + "dropout": 0.0, + "is_residual": True, + "in_channels": 12, + "out_channels": 12, + "patch_size": 2, + "scale_factor_spatial": 16, + "scale_factor_temporal": 4, + "clip_output": False, + "latents_mean": [0.0] * 48, + "latents_std": [1.0] * 48, + } + raw = AutoencoderKLWan.from_config(arch) + vae = WanVideoVAE38(dtype=torch.float32, device="cpu", pretrained=raw) + assert vae.z_dim == 48 + assert vae.upsampling_factor == 16 + assert vae.temporal_downsample_factor == 4 + + video = torch.rand(1, 3, 5, 32, 32) * 2 - 1 # [B,C,T,H,W] in [-1,1] + latents = vae.encode(video) + assert latents.shape == (1, 48, 2, 2, 2) # T'=(5-1)//4+1, H'=W'=32//16 + + decoded = vae.decode(latents) + assert decoded.shape[0] == 1 and decoded.shape[1] == 3 and decoded.shape[-2:] == (32, 32) + assert decoded.min() >= -1.0 and decoded.max() <= 1.0 + + # list input is accepted and equals the batched path + assert torch.equal(vae.encode([video[0]]), latents) diff --git a/uv.lock b/uv.lock index 5a76fcbf8..076d021f2 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 2 +revision = 3 requires-python = ">=3.12" resolution-markers = [ "(python_full_version >= '3.15' and platform_machine == 'AMD64' and sys_platform == 'linux') or (python_full_version >= '3.15' and platform_machine == 'x86_64' and sys_platform == 'linux')", @@ -2955,6 +2955,10 @@ eo1 = [ evaluation = [ { name = "av" }, ] +fastwam = [ + { name = "diffusers" }, + { name = "transformers" }, +] feetech = [ { name = "deepdiff" }, { name = "feetech-servo-sdk" }, @@ -3261,11 +3265,13 @@ requires-dist = [ { name = "lerobot", extras = ["deepdiff-dep"], marker = "extra == 'hardware'" }, { name = "lerobot", extras = ["dev"], marker = "extra == 'all'" }, { name = "lerobot", extras = ["diffusers-dep"], marker = "extra == 'diffusion'" }, + { name = "lerobot", extras = ["diffusers-dep"], marker = "extra == 'fastwam'" }, { name = "lerobot", extras = ["diffusers-dep"], marker = "extra == 'groot'" }, { name = "lerobot", extras = ["diffusers-dep"], marker = "extra == 'multi-task-dit'" }, { name = "lerobot", extras = ["diffusers-dep"], marker = "extra == 'vla-jepa'" }, { name = "lerobot", extras = ["diffusion"], marker = "extra == 'all'" }, { name = "lerobot", extras = ["dynamixel"], marker = "extra == 'all'" }, + { name = "lerobot", extras = ["fastwam"], marker = "extra == 'all'" }, { name = "lerobot", extras = ["feetech"], marker = "extra == 'all'" }, { name = "lerobot", extras = ["feetech"], marker = "extra == 'hopejr'" }, { name = "lerobot", extras = ["feetech"], marker = "extra == 'lekiwi'" }, @@ -3335,6 +3341,7 @@ requires-dist = [ { name = "lerobot", extras = ["training"], marker = "extra == 'all'" }, { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'annotations'" }, { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'eo1'" }, + { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'fastwam'" }, { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'groot'" }, { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'hilserl'" }, { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'libero'" }, @@ -3417,7 +3424,7 @@ requires-dist = [ { name = "transformers", marker = "extra == 'transformers-dep'", specifier = ">=5.4.0,<5.6.0" }, { name = "wandb", marker = "extra == 'training'", specifier = ">=0.24.0,<0.28.0" }, ] -provides-extras = ["dataset", "training", "hardware", "viz", "core-scripts", "evaluation", "dataset-viz", "av-dep", "pygame-dep", "placo-dep", "transformers-dep", "grpcio-dep", "accelerate-dep", "can-dep", "peft-dep", "scipy-dep", "diffusers-dep", "qwen-vl-utils-dep", "matplotlib-dep", "pyserial-dep", "deepdiff-dep", "pynput-dep", "pyzmq-dep", "motorbridge-dep", "motorbridge-smart-servo-dep", "feetech", "dynamixel", "damiao", "robstride", "openarms", "gamepad", "hopejr", "lekiwi", "unitree-g1", "reachy2", "rebot", "kinematics", "intelrealsense", "phone", "diffusion", "wallx", "pi", "molmoact2", "smolvla", "multi-task-dit", "groot", "sarm", "robometer", "topreward", "xvla", "eo1", "hilserl", "vla-jepa", "async", "peft", "annotations", "dev", "notebook", "test", "video-benchmark", "aloha", "pusht", "libero", "metaworld", "all"] +provides-extras = ["dataset", "training", "hardware", "viz", "core-scripts", "evaluation", "dataset-viz", "av-dep", "pygame-dep", "placo-dep", "transformers-dep", "grpcio-dep", "accelerate-dep", "can-dep", "peft-dep", "scipy-dep", "diffusers-dep", "qwen-vl-utils-dep", "matplotlib-dep", "pyserial-dep", "deepdiff-dep", "pynput-dep", "pyzmq-dep", "motorbridge-dep", "motorbridge-smart-servo-dep", "feetech", "dynamixel", "damiao", "robstride", "openarms", "gamepad", "hopejr", "lekiwi", "unitree-g1", "reachy2", "rebot", "kinematics", "intelrealsense", "phone", "diffusion", "wallx", "pi", "molmoact2", "smolvla", "multi-task-dit", "groot", "sarm", "robometer", "topreward", "xvla", "eo1", "fastwam", "hilserl", "vla-jepa", "async", "peft", "annotations", "dev", "notebook", "test", "video-benchmark", "aloha", "pusht", "libero", "metaworld", "all"] [[package]] name = "librt"