Add FastWAM policy

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

docs/source/fastwam.mdx

@@ -0,0 +1,179 @@
+# 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=256 \
+  --env.observation_width=256 \
+  --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`.
+
+### 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]'
+```
+
+`policy.invert_dimensions` remains available for older checkpoints or robot
+setups that only need a sign inversion.
+
+## 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).
+
+## Reproducibility Checklist
+
+For a PR adding or updating FastWAM results, include:
+
+- the training dataset repo id
+- the LeRobot-format checkpoint repo id
+- the exact `lerobot-train` command
+- the exact `lerobot-eval` or `lerobot-rollout` command
+- the number of evaluation episodes
+- the GPU type and count
+
+The upstream Fast-WAM release provides reference checkpoints and benchmark assets at `yuanty/fastwam`; LeRobot eval numbers should be reported from a converted LeRobot-format checkpoint so reviewers can reproduce them with the commands above.
+
+## 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}
+}
+```

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
+```

pyproject.toml

@@ -222,6 +222,12 @@ 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]",
+    "ftfy>=6.1.1,<7.0.0",
+    "regex>=2024.0.0,<2027.0.0",
+]
 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]"]
 
@@ -301,6 +307,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]",

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",

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)
@@ -323,6 +330,10 @@ def make_pre_post_processors(
             revision=pretrained_revision,
         )
         _reconnect_relative_absolute_steps(preprocessor, postprocessor)
+        if isinstance(policy_cfg, FastWAMConfig):
+            from .fastwam.processor_fastwam import migrate_fastwam_postprocessor
+
+            postprocessor = migrate_fastwam_postprocessor(postprocessor, policy_cfg)
         return preprocessor, postprocessor
 
     # Create a new processor based on policy type
@@ -451,6 +462,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(

src/lerobot/policies/fastwam/README.md

@@ -0,0 +1 @@
+../../../../docs/source/policy_fastwam_README.md

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",
+]

src/lerobot/policies/fastwam/configuration_fastwam.py

@@ -0,0 +1,284 @@
+# 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 json
+from dataclasses import dataclass, field, fields
+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"
+
+
+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",
+    }
+
+
+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,
+    }
+
+
+def _coerce_enum(enum_cls: type, value: Any) -> Any:
+    if isinstance(value, enum_cls):
+        return value
+    try:
+        return enum_cls(value)
+    except (TypeError, ValueError):
+        return getattr(enum_cls, str(value), value)
+
+
+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 _coerce_normalization_mapping(mapping: dict[str, Any]) -> dict[str, Any]:
+    return {key: _coerce_enum(NormalizationMode, value) for key, value in mapping.items()}
+
+
+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 _coerce_pretrained_tokenizer_model_id(payload: dict[str, Any]) -> None:
+    tokenizer_model_id = payload.get("tokenizer_model_id")
+    if tokenizer_model_id is None:
+        return
+    if tokenizer_model_id == WAN22_MODEL_ID or _is_local_model_id(tokenizer_model_id):
+        return
+    payload["tokenizer_model_id"] = WAN22_MODEL_ID
+
+
+@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): Number of video frames used by FastWAM rollout.
+        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.
+        invert_dimensions (list[int]): Action dimensions to multiply by -1
+            during postprocessing. Supports negative indices.
+    """
+
+    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
+    image_size: tuple[int, int] = (224, 448)
+    context_len: int = 128
+    model_id: str = WAN22_MODEL_ID
+    tokenizer_model_id: str = WAN22_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
+    toggle_action_dimensions: list[int] = field(default_factory=lambda: [-1])
+    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
+    invert_dimensions: list[int] = field(default_factory=list)
+    normalization_mapping: dict[str, Any] = field(
+        default_factory=lambda: {
+            "VISUAL": NormalizationMode.MEAN_STD,
+            "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:
+        parent_post_init = getattr(super(), "__post_init__", None)
+        if parent_post_init is not None:
+            parent_post_init()
+        self.image_size = tuple(self.image_size)
+        self.model_id = _validate_wan_model_id(self.model_id, "model_id")
+        self.tokenizer_model_id = _validate_wan_model_id(self.tokenizer_model_id, "tokenizer_model_id")
+        self.input_features = _coerce_policy_features(self.input_features)
+        self.output_features = _coerce_policy_features(self.output_features)
+        self.invert_dimensions = [int(dim) for dim in self.invert_dimensions]
+        self.toggle_action_dimensions = [int(dim) for dim in self.toggle_action_dimensions]
+        self.normalization_mapping = _coerce_normalization_mapping(self.normalization_mapping)
+        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.input_features = self.input_features or self._default_input_features()
+        self.output_features = self.output_features or self._default_output_features()
+        self.validate_features()
+
+    @classmethod
+    def from_pretrained(cls, pretrained_name_or_path: str | Path, **_: Any) -> FastWAMConfig:
+        config_path = Path(pretrained_name_or_path) / "config.json"
+        with open(config_path, encoding="utf-8") as f:
+            payload = json.load(f)
+        payload.pop("type", None)
+        known_fields = {field.name for field in fields(cls)}
+        payload = {key: value for key, value in payload.items() if key in known_fields}
+        _coerce_pretrained_tokenizer_model_id(payload)
+        return cls(**payload)
+
+    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 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.num_video_frames % 4 != 1:
+            raise ValueError(f"`num_video_frames` must satisfy T % 4 == 1, got {self.num_video_frames}.")
+        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}."
+                )
+        self._validate_image_feature_shapes()
+
+    def _validate_image_feature_shapes(self) -> None:
+        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 observation_delta_indices(self) -> None:
+        return None
+
+    @property
+    def action_delta_indices(self) -> list[int]:
+        return list(range(self.action_horizon))
+
+    @property
+    def reward_delta_indices(self) -> None:
+        return None
+
+    def _default_input_features(self) -> dict[str, PolicyFeature]:
+        height, width = self.image_size
+        features = {
+            "observation.images.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, height, width))
+        }
+        if self.proprio_dim is not None:
+            features[OBS_STATE] = PolicyFeature(type=FeatureType.STATE, shape=(self.proprio_dim,))
+        return features
+
+    def _default_output_features(self) -> dict[str, PolicyFeature]:
+        return {ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(self.action_dim,))}

src/lerobot/policies/fastwam/modeling_fastwam.py

@@ -0,0 +1,583 @@
+# 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 shutil
+from collections import deque
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import torch
+import torch.nn.functional as functional
+from torch import Tensor
+
+from lerobot.policies.pretrained import PreTrainedPolicy
+from lerobot.utils.constants import ACTION, OBS_STATE
+
+from .configuration_fastwam import FastWAMConfig
+
+if TYPE_CHECKING:
+    from .wan_components import WanCheckpointPaths
+
+
+class FastWAMPolicy(PreTrainedPolicy):
+    """LeRobot policy wrapper for FastWAM.
+
+    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,
+    ):
+        skip_wan_init = bool(kwargs.pop("_skip_wan_init", False))
+        super().__init__(config, dataset_stats)
+        config.validate_features()
+        self.config = config
+        self.dataset_stats = dataset_stats
+        if skip_wan_init:
+            self.model = _build_core_model_from_architecture(config)
+        else:
+            self.model = self._build_core_model(config)
+        self.reset()
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_name_or_path: str | Path,
+        *,
+        config: FastWAMConfig | None = None,
+        force_download: bool = False,
+        resume_download: bool | None = None,
+        proxies: dict | None = None,
+        token: str | bool | None = None,
+        cache_dir: str | Path | None = None,
+        local_files_only: bool = False,
+        revision: str | None = None,
+        strict: bool = False,
+        **kwargs: Any,
+    ) -> FastWAMPolicy:
+        """Load FastWAM weights and local Wan components from one HF directory.
+
+        Args:
+            pretrained_name_or_path (str | Path): HF-format policy directory
+                containing `config.json`, `model.safetensors`, local Wan VAE,
+                local UMT5 text encoder, and tokenizer files.
+            config (FastWAMConfig | None): Optional config override. When
+                omitted, `config.json` is read from `pretrained_name_or_path`.
+            force_download (bool): Forwarded to LeRobot's pretrained loader.
+            resume_download (bool | None): Forwarded to LeRobot's pretrained loader.
+            proxies (dict | None): Forwarded to LeRobot's pretrained loader.
+            token (str | bool | None): Forwarded to LeRobot's pretrained loader.
+            cache_dir (str | Path | None): Forwarded to LeRobot's pretrained loader.
+            local_files_only (bool): Forwarded to LeRobot's pretrained loader.
+            revision (str | None): Forwarded to LeRobot's pretrained loader.
+            strict (bool): Whether safetensors loading should require an exact
+                match between checkpoint keys and policy module keys.
+            **kwargs (Any): Extra constructor arguments forwarded to
+                `FastWAMPolicy`.
+        """
+
+        pretrained_path = _resolve_pretrained_directory(
+            pretrained_name_or_path=pretrained_name_or_path,
+            force_download=force_download,
+            token=token,
+            cache_dir=cache_dir,
+            local_files_only=local_files_only,
+            revision=revision,
+        )
+        if config is None:
+            config = cls.config_class.from_pretrained(pretrained_path)
+        kwargs["_skip_wan_init"] = True
+        policy = super().from_pretrained(
+            pretrained_path,
+            config=config,
+            force_download=force_download,
+            resume_download=resume_download,
+            proxies=proxies,
+            token=token,
+            cache_dir=cache_dir,
+            local_files_only=local_files_only,
+            revision=revision,
+            strict=strict,
+            **kwargs,
+        )
+        policy.load_wan_components_from_pretrained(pretrained_path)
+        policy.eval()
+        return policy
+
+    def _save_pretrained(self, save_directory: Path) -> None:
+        super()._save_pretrained(save_directory)
+        _copy_wan_components_from_policy(policy=self, save_directory=save_directory)
+
+    def get_optim_params(self) -> dict[str, Any]:
+        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 {"params": [p for p in params if p.requires_grad]}
+
+    def load_wan_components_from_pretrained(self, pretrained_name_or_path: str | Path) -> None:
+        """Attach local Wan VAE, text encoder, and tokenizer from a HF directory.
+
+        Args:
+            pretrained_name_or_path (str | Path): Directory containing
+                `Wan2.2_VAE.pth`, `models_t5_umt5-xxl-enc-bf16.pth`,
+                and `google/umt5-xxl/` tokenizer files.
+        """
+
+        paths = resolve_wan_component_paths(pretrained_name_or_path)
+        _load_wan_components_into_policy(policy=self, paths=paths)
+
+    def reset(self) -> None:
+        self._action_queue: deque[Tensor] = deque([], maxlen=self.config.n_action_steps)
+
+    def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
+        """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:
+            dict[str, Tensor]: Output dictionary containing the scalar `loss`
+            key required by LeRobot and optional tensor metrics.
+        """
+
+        sample = _batch_to_training_sample(batch=batch, config=self.config)
+        loss, metrics = self.model.training_loss(sample)
+        output = {"loss": loss}
+        output.update(_metrics_to_tensors(metrics=metrics, device=loss.device))
+        return output
+
+    @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) -> torch.nn.Module:
+        return _build_core_model_from_wan22(config)
+
+
+def _resolve_pretrained_directory(
+    pretrained_name_or_path: str | Path,
+    *,
+    force_download: bool,
+    token: str | bool | None,
+    cache_dir: str | Path | None,
+    local_files_only: bool,
+    revision: str | None,
+) -> Path:
+    path = Path(pretrained_name_or_path)
+    if path.is_dir():
+        return path
+
+    from huggingface_hub import snapshot_download
+
+    snapshot_path = snapshot_download(
+        repo_id=str(pretrained_name_or_path),
+        revision=revision,
+        cache_dir=cache_dir,
+        force_download=force_download,
+        token=token,
+        local_files_only=local_files_only,
+        allow_patterns=[
+            "config.json",
+            "model.safetensors",
+            "Wan2.2_VAE.pth",
+            "models_t5_umt5-xxl-enc-bf16.pth",
+            "google/umt5-xxl/**",
+        ],
+    )
+    return Path(snapshot_path)
+
+
+def resolve_wan_component_paths(pretrained_name_or_path: str | Path) -> WanCheckpointPaths:
+    """Resolve local Wan component paths stored beside FastWAM HF weights.
+
+    Args:
+        pretrained_name_or_path (str | Path): HF-format FastWAM directory.
+
+    Returns:
+        WanCheckpointPaths: Existing VAE, text encoder, and tokenizer paths.
+        DiT shards are intentionally optional here because FastWAM HF
+        checkpoints store trainable DiT weights in `model.safetensors`.
+    """
+
+    from .wan_components import resolve_wan_checkpoint_paths
+
+    return resolve_wan_checkpoint_paths(
+        pretrained_name_or_path,
+        load_dit=False,
+        load_text_encoder=True,
+    )
+
+
+def _load_wan_components_into_policy(policy: FastWAMPolicy, paths: WanCheckpointPaths) -> None:
+    from .wan_components import load_wan_text_encoder, load_wan_tokenizer, load_wan_vae
+
+    if paths.text_encoder is None or paths.tokenizer is None:
+        raise FileNotFoundError("FastWAM HF checkpoint requires Wan text encoder and tokenizer sidecars.")
+    dtype = _dtype_from_name(policy.config.torch_dtype)
+    device = str(policy.config.device)
+    policy.model.vae = load_wan_vae(paths.vae, torch_dtype=dtype, device=device)
+    policy.model.text_encoder = load_wan_text_encoder(paths.text_encoder, torch_dtype=dtype, device=device)
+    policy.model.tokenizer = load_wan_tokenizer(
+        paths.tokenizer,
+        tokenizer_max_len=int(policy.config.tokenizer_max_len),
+    )
+    _record_wan_component_paths(policy=policy, paths=paths)
+
+
+def _record_wan_component_paths(policy: FastWAMPolicy, paths: WanCheckpointPaths) -> None:
+    model_paths = dict(getattr(policy.model, "model_paths", {}) or {})
+    model_paths.update(
+        {
+            "vae": str(paths.vae),
+            "text_encoder": str(paths.text_encoder),
+            "tokenizer": str(paths.tokenizer),
+        }
+    )
+    policy.model.model_paths = model_paths
+
+
+def _copy_wan_components_from_policy(policy: FastWAMPolicy, save_directory: Path) -> None:
+    model_paths = getattr(policy.model, "model_paths", {}) or {}
+    paths = {
+        "vae": model_paths.get("vae"),
+        "text_encoder": model_paths.get("text_encoder"),
+        "tokenizer": model_paths.get("tokenizer"),
+    }
+    missing = [name for name, path in paths.items() if path is None]
+    if missing:
+        raise RuntimeError(
+            "FastWAM save_pretrained requires local Wan component paths for "
+            f"{missing}. Load or initialize the policy with local Wan VAE, text encoder, and tokenizer files."
+        )
+    _copy_component_path(Path(paths["vae"]), save_directory / Path(paths["vae"]).name)
+    _copy_component_path(Path(paths["text_encoder"]), save_directory / Path(paths["text_encoder"]).name)
+    tokenizer_source = Path(paths["tokenizer"])
+    _copy_component_path(tokenizer_source, save_directory / "google" / "umt5-xxl")
+
+
+def _copy_component_path(source: Path, destination: Path) -> None:
+    source = source.expanduser()
+    if not source.exists():
+        raise FileNotFoundError(f"FastWAM component path does not exist: {source}")
+    if source.resolve() == destination.resolve():
+        return
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    if source.is_dir():
+        shutil.copytree(source, destination, dirs_exist_ok=True)
+    else:
+        shutil.copy2(source, destination)
+
+
+def _build_core_model_from_wan22(config: FastWAMConfig) -> torch.nn.Module:
+    from .modular_fastwam import FastWAM
+
+    dtype = _dtype_from_name(config.torch_dtype)
+    return FastWAM.from_wan22_pretrained(
+        device=config.device,
+        torch_dtype=dtype,
+        model_id=config.model_id,
+        tokenizer_model_id=config.tokenizer_model_id,
+        tokenizer_max_len=config.tokenizer_max_len,
+        load_text_encoder=config.load_text_encoder,
+        proprio_dim=config.proprio_dim,
+        video_dit_config=config.video_dit_config,
+        action_dit_config=config.action_dit_config,
+        mot_checkpoint_mixed_attn=config.mot_checkpoint_mixed_attn,
+        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 _build_core_model_from_architecture(config: FastWAMConfig) -> torch.nn.Module:
+    from .modular_fastwam import ActionDiT, FastWAM, MoT
+    from .wan_video_dit import WanVideoDiT
+
+    dtype = _dtype_from_name(config.torch_dtype)
+    video_expert = WanVideoDiT(**config.video_dit_config).to(device=config.device, dtype=dtype)
+    action_expert = ActionDiT(**config.action_dit_config).to(device=config.device, dtype=dtype)
+    mot = MoT(
+        mixtures={"video": video_expert, "action": action_expert},
+        mot_checkpoint_mixed_attn=config.mot_checkpoint_mixed_attn,
+    )
+    return FastWAM(
+        video_expert=video_expert,
+        action_expert=action_expert,
+        mot=mot,
+        vae=_FastWAMVAEPlaceholder(),
+        text_encoder=None,
+        tokenizer=None,
+        text_dim=int(config.video_dit_config["text_dim"]),
+        proprio_dim=config.proprio_dim,
+        device=config.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"]),
+    )
+
+
+class _FastWAMVAEPlaceholder(torch.nn.Module):
+    """Minimal VAE placeholder for checkpoint loading without Wan2.2 VAE.
+
+    Args:
+        temporal_downsample_factor (int): Temporal compression factor expected
+            by FastWAM latent shape logic.
+        upsampling_factor (int): Spatial compression factor expected by FastWAM.
+        z_dim (int): Latent channel count used by Wan2.2 TI2V VAE.
+    """
+
+    temporal_downsample_factor: int = 4
+    upsampling_factor: int = 8
+
+    def __init__(self, z_dim: int = 48):
+        super().__init__()
+        self.model = type("VAEModelShape", (), {"z_dim": int(z_dim)})()
+
+    def encode(self, *args, **kwargs):
+        raise RuntimeError(
+            "FastWAM VAE placeholder cannot encode images; load Wan2.2 VAE for image inference."
+        )
+
+    def decode(self, *args, **kwargs):
+        raise RuntimeError(
+            "FastWAM VAE placeholder cannot decode latents; load Wan2.2 VAE for video inference."
+        )
+
+
+def _batch_to_training_sample(batch: dict[str, Tensor], config: FastWAMConfig) -> dict[str, Tensor]:
+    sample = dict(batch)
+    if "video" not in sample:
+        sample["video"] = _stack_video_from_images(batch, config)
+    if "proprio" not in sample and OBS_STATE in batch:
+        sample["proprio"] = batch[OBS_STATE]
+    required = {"video", ACTION, "context", "context_mask"}
+    missing = sorted(required - set(sample))
+    if missing:
+        raise KeyError(f"FastWAM training batch is missing keys: {missing}.")
+    return sample
+
+
+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(batch.get("text_cfg_scale", config.text_cfg_scale)),
+        "num_inference_steps": int(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 _metrics_to_tensors(metrics: dict[str, Any] | None, device: torch.device) -> dict[str, Tensor]:
+    if metrics is None:
+        return {}
+    tensor_metrics = {}
+    for key, value in metrics.items():
+        if isinstance(value, Tensor):
+            tensor_metrics[key] = value.to(device=device)
+        else:
+            tensor_metrics[key] = torch.as_tensor(value, device=device)
+    return tensor_metrics
+
+
+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 _stack_video_from_images(batch: dict[str, Tensor], config: FastWAMConfig) -> Tensor:
+    image_keys = sorted(k for k in batch if k.startswith("observation.images."))
+    if not image_keys:
+        raise KeyError("FastWAM batch must contain `video` or `observation.images.*` keys.")
+    images = [batch[key] for key in image_keys]
+    image = torch.cat(images, dim=-1) if len(images) > 1 else images[0]
+    if image.ndim == 4:
+        image = image.unsqueeze(2).repeat(1, 1, config.num_video_frames, 1, 1)
+    if image.ndim != 5:
+        raise ValueError(f"Expected image batch [B,C,H,W] or video [B,C,T,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)}.")
+
+    if image.dtype == torch.uint8:
+        image = image.to(dtype=torch.float32).div(255.0).mul(2.0).sub(1.0)
+    else:
+        image = image.to(dtype=torch.float32)
+        image_min = float(image.detach().amin().cpu())
+        image_max = float(image.detach().amax().cpu())
+        if image_min >= 0.0 and image_max <= 1.0:
+            image = image.mul(2.0).sub(1.0)
+        elif image_max > 2.0:
+            image = image.div(255.0).mul(2.0).sub(1.0)
+
+    target_h, target_w = config.image_size
+    if tuple(image.shape[-2:]) != (target_h, target_w):
+        image = _center_crop_resize(image, target_h=target_h, target_w=target_w)
+    return image
+
+
+def _center_crop_resize(image: Tensor, *, target_h: int, target_w: int) -> Tensor:
+    _, _, height, width = image.shape
+    scale = max(target_h / height, target_w / width)
+    resized_h = round(height * scale)
+    resized_w = round(width * scale)
+    image = functional.interpolate(image, size=(resized_h, resized_w), mode="bilinear", align_corners=False)
+    top = max((resized_h - target_h) // 2, 0)
+    left = max((resized_w - target_w) // 2, 0)
+    return image[:, :, top : top + target_h, left : left + target_w].contiguous()

src/lerobot/policies/fastwam/processor_fastwam.py

@@ -0,0 +1,186 @@
+# 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_inversion_processor")
+class FastWAMActionInversionProcessorStep(ActionProcessorStep):
+    """Invert configured FastWAM action dimensions during postprocessing."""
+
+    invert_dimensions: list[int]
+
+    def action(self, action: PolicyAction) -> PolicyAction:
+        if not self.invert_dimensions:
+            return action
+        processed_action = action.clone()
+        action_dim = int(processed_action.shape[-1])
+        for dim in self.invert_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 inversion dimension {dim} is out of bounds for action dim {action_dim}."
+                )
+            processed_action[..., resolved_dim] = -processed_action[..., resolved_dim]
+        return processed_action
+
+    def get_config(self) -> dict[str, Any]:
+        return {"invert_dimensions": self.invert_dimensions}
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        return features
+
+
+@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.
+    """
+
+    input_steps = [
+        RenameObservationsProcessorStep(rename_map={}),
+        AddBatchDimensionProcessorStep(),
+        DeviceProcessorStep(device=config.device),
+        NormalizerProcessorStep(
+            features={**config.input_features, **config.output_features},
+            norm_map=config.normalization_mapping,
+            stats=dataset_stats,
+            device=config.device,
+        ),
+    ]
+    output_steps = [
+        UnnormalizerProcessorStep(
+            features=config.output_features,
+            norm_map=config.normalization_mapping,
+            stats=dataset_stats,
+        ),
+    ]
+    if config.toggle_action_dimensions:
+        output_steps.append(
+            FastWAMActionToggleProcessorStep(toggle_dimensions=config.toggle_action_dimensions)
+        )
+    elif config.invert_dimensions:
+        output_steps.append(FastWAMActionInversionProcessorStep(invert_dimensions=config.invert_dimensions))
+    output_steps.append(DeviceProcessorStep(device="cpu"))
+    return _build_lerobot_pipelines(input_steps=input_steps, output_steps=output_steps)
+
+
+def migrate_fastwam_postprocessor(
+    postprocessor: PolicyProcessorPipeline[PolicyAction, PolicyAction],
+    config: FastWAMConfig,
+) -> PolicyProcessorPipeline[PolicyAction, PolicyAction]:
+    """Upgrade old FastWAM postprocessor pipelines to the LIBERO toggle step."""
+
+    if not config.toggle_action_dimensions:
+        return postprocessor
+
+    toggle_step = FastWAMActionToggleProcessorStep(toggle_dimensions=config.toggle_action_dimensions)
+    steps = [
+        step
+        for step in postprocessor.steps
+        if not isinstance(step, (FastWAMActionInversionProcessorStep, FastWAMActionToggleProcessorStep))
+    ]
+    insert_at = next(
+        (idx for idx, step in enumerate(steps) if isinstance(step, DeviceProcessorStep)), len(steps)
+    )
+    steps.insert(insert_at, toggle_step)
+    postprocessor.steps = steps
+    return postprocessor
+
+
+def _build_lerobot_pipelines(input_steps: list[Any], output_steps: list[Any]) -> tuple[Any, Any]:
+    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,
+        ),
+    )

src/lerobot/policies/fastwam/wan/README.md

@@ -0,0 +1,28 @@
+# Wan2.2 Upstream Subset
+
+This directory contains an unmodified subset of the official Wan2.2 source tree.
+
+- 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:
+
+- `wan/modules/attention.py`
+- `wan/modules/model.py`
+- `wan/modules/t5.py`
+- `wan/modules/tokenizers.py`
+- `wan/modules/vae2_1.py`
+- `wan/modules/vae2_2.py`
+- `wan/modules/__init__.py`
+- `wan/utils/fm_solvers.py`
+- `wan/utils/fm_solvers_unipc.py`
+- `wan/utils/__init__.py`
+
+FastWAM-specific model glue and any code adapted from these modules live outside this directory. This keeps the upstream Wan2.2 code reviewable as a vendored reference subset and makes it straightforward to replace this directory with an external Wan2.2 dependency by changing import paths.
+
+Current FastWAM adapters that directly reuse this vendored subset:
+
+- `../wan_components.py` instantiates the upstream `wan.modules.t5.umt5_xxl` encoder factory and uses `wan.modules.tokenizers.HuggingfaceTokenizer`.
+- `../wan_adapters.py` wraps `wan.modules.vae2_2.Wan2_2_VAE` with the FastWAM tensor-batch encode/decode API.
+- `../modular_fastwam.py` reuses `wan.utils.fm_solvers.get_sampling_sigmas` for Wan-compatible inference timesteps.

src/lerobot/policies/fastwam/wan/modules/__init__.py

@@ -0,0 +1,19 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+from .attention import flash_attention
+from .model import WanModel
+from .t5 import T5Decoder, T5Encoder, T5EncoderModel, T5Model
+from .tokenizers import HuggingfaceTokenizer
+from .vae2_1 import Wan2_1_VAE
+from .vae2_2 import Wan2_2_VAE
+
+__all__ = [
+    "Wan2_1_VAE",
+    "Wan2_2_VAE",
+    "WanModel",
+    "T5Model",
+    "T5Encoder",
+    "T5Decoder",
+    "T5EncoderModel",
+    "HuggingfaceTokenizer",
+    "flash_attention",
+]

src/lerobot/policies/fastwam/wan/modules/attention.py

@@ -0,0 +1,182 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import torch
+
+try:
+    import flash_attn_interface
+
+    FLASH_ATTN_3_AVAILABLE = True
+except ModuleNotFoundError:
+    FLASH_ATTN_3_AVAILABLE = False
+
+try:
+    import flash_attn
+
+    FLASH_ATTN_2_AVAILABLE = True
+except ModuleNotFoundError:
+    FLASH_ATTN_2_AVAILABLE = False
+
+import warnings
+
+__all__ = [
+    "flash_attention",
+    "attention",
+]
+
+
+def flash_attention(
+    q,
+    k,
+    v,
+    q_lens=None,
+    k_lens=None,
+    dropout_p=0.0,
+    softmax_scale=None,
+    q_scale=None,
+    causal=False,
+    window_size=(-1, -1),
+    deterministic=False,
+    dtype=torch.bfloat16,
+    version=None,
+):
+    """
+    q:              [B, Lq, Nq, C1].
+    k:              [B, Lk, Nk, C1].
+    v:              [B, Lk, Nk, C2]. Nq must be divisible by Nk.
+    q_lens:         [B].
+    k_lens:         [B].
+    dropout_p:      float. Dropout probability.
+    softmax_scale:  float. The scaling of QK^T before applying softmax.
+    causal:         bool. Whether to apply causal attention mask.
+    window_size:    (left right). If not (-1, -1), apply sliding window local attention.
+    deterministic:  bool. If True, slightly slower and uses more memory.
+    dtype:          torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16.
+    """
+    half_dtypes = (torch.float16, torch.bfloat16)
+    assert dtype in half_dtypes
+    assert q.device.type == "cuda" and q.size(-1) <= 256
+
+    # params
+    b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype
+
+    def half(x):
+        return x if x.dtype in half_dtypes else x.to(dtype)
+
+    # preprocess query
+    if q_lens is None:
+        q = half(q.flatten(0, 1))
+        q_lens = torch.tensor([lq] * b, dtype=torch.int32).to(device=q.device, non_blocking=True)
+    else:
+        q = half(torch.cat([u[:v] for u, v in zip(q, q_lens)]))
+
+    # preprocess key, value
+    if k_lens is None:
+        k = half(k.flatten(0, 1))
+        v = half(v.flatten(0, 1))
+        k_lens = torch.tensor([lk] * b, dtype=torch.int32).to(device=k.device, non_blocking=True)
+    else:
+        k = half(torch.cat([u[:v] for u, v in zip(k, k_lens)]))
+        v = half(torch.cat([u[:v] for u, v in zip(v, k_lens)]))
+
+    q = q.to(v.dtype)
+    k = k.to(v.dtype)
+
+    if q_scale is not None:
+        q = q * q_scale
+
+    if version is not None and version == 3 and not FLASH_ATTN_3_AVAILABLE:
+        warnings.warn("Flash attention 3 is not available, use flash attention 2 instead.")
+
+    # apply attention
+    if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE:
+        # Note: dropout_p, window_size are not supported in FA3 now.
+        x = flash_attn_interface.flash_attn_varlen_func(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens])
+            .cumsum(0, dtype=torch.int32)
+            .to(q.device, non_blocking=True),
+            cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens])
+            .cumsum(0, dtype=torch.int32)
+            .to(q.device, non_blocking=True),
+            seqused_q=None,
+            seqused_k=None,
+            max_seqlen_q=lq,
+            max_seqlen_k=lk,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            deterministic=deterministic,
+        )[0].unflatten(0, (b, lq))
+    else:
+        assert FLASH_ATTN_2_AVAILABLE
+        x = flash_attn.flash_attn_varlen_func(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens])
+            .cumsum(0, dtype=torch.int32)
+            .to(q.device, non_blocking=True),
+            cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens])
+            .cumsum(0, dtype=torch.int32)
+            .to(q.device, non_blocking=True),
+            max_seqlen_q=lq,
+            max_seqlen_k=lk,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            window_size=window_size,
+            deterministic=deterministic,
+        ).unflatten(0, (b, lq))
+
+    # output
+    return x.type(out_dtype)
+
+
+def attention(
+    q,
+    k,
+    v,
+    q_lens=None,
+    k_lens=None,
+    dropout_p=0.0,
+    softmax_scale=None,
+    q_scale=None,
+    causal=False,
+    window_size=(-1, -1),
+    deterministic=False,
+    dtype=torch.bfloat16,
+    fa_version=None,
+):
+    if FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE:
+        return flash_attention(
+            q=q,
+            k=k,
+            v=v,
+            q_lens=q_lens,
+            k_lens=k_lens,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            q_scale=q_scale,
+            causal=causal,
+            window_size=window_size,
+            deterministic=deterministic,
+            dtype=dtype,
+            version=fa_version,
+        )
+    else:
+        if q_lens is not None or k_lens is not None:
+            warnings.warn(
+                "Padding mask is disabled when using scaled_dot_product_attention. It can have a significant impact on performance."
+            )
+        attn_mask = None
+
+        q = q.transpose(1, 2).to(dtype)
+        k = k.transpose(1, 2).to(dtype)
+        v = v.transpose(1, 2).to(dtype)
+
+        out = torch.nn.functional.scaled_dot_product_attention(
+            q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p
+        )
+
+        out = out.transpose(1, 2).contiguous()
+        return out

src/lerobot/policies/fastwam/wan/modules/model.py

@@ -0,0 +1,519 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import math
+
+import torch
+import torch.nn as nn
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+
+from .attention import flash_attention
+
+__all__ = ["WanModel"]
+
+
+def sinusoidal_embedding_1d(dim, position):
+    # preprocess
+    assert dim % 2 == 0
+    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):
+    assert dim % 2 == 0
+    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, window_size=(-1, -1), qk_norm=True, eps=1e-6):
+        assert dim % num_heads == 0
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.eps = eps
+
+        # 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()
+
+    def forward(self, x, seq_lens, grid_sizes, freqs):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, num_heads, C / num_heads]
+            seq_lens(Tensor): Shape [B]
+            grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+        """
+        b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+
+        # query, key, value function
+        def qkv_fn(x):
+            q = self.norm_q(self.q(x)).view(b, s, n, d)
+            k = self.norm_k(self.k(x)).view(b, s, n, d)
+            v = self.v(x).view(b, s, n, d)
+            return q, k, v
+
+        q, k, v = qkv_fn(x)
+
+        x = flash_attention(
+            q=rope_apply(q, grid_sizes, freqs),
+            k=rope_apply(k, grid_sizes, freqs),
+            v=v,
+            k_lens=seq_lens,
+            window_size=self.window_size,
+        )
+
+        # output
+        x = x.flatten(2)
+        x = self.o(x)
+        return x
+
+
+class WanCrossAttention(WanSelfAttention):
+    def forward(self, x, context, context_lens):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L1, C]
+            context(Tensor): Shape [B, L2, C]
+            context_lens(Tensor): Shape [B]
+        """
+        b, n, d = x.size(0), self.num_heads, self.head_dim
+
+        # compute query, key, value
+        q = self.norm_q(self.q(x)).view(b, -1, n, d)
+        k = self.norm_k(self.k(context)).view(b, -1, n, d)
+        v = self.v(context).view(b, -1, n, d)
+
+        # compute attention
+        x = flash_attention(q, k, v, k_lens=context_lens)
+
+        # output
+        x = x.flatten(2)
+        x = self.o(x)
+        return x
+
+
+class WanAttentionBlock(nn.Module):
+    def __init__(
+        self, dim, ffn_dim, num_heads, window_size=(-1, -1), 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.window_size = window_size
+        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, window_size, qk_norm, eps)
+        self.norm3 = WanLayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()
+        self.cross_attn = WanCrossAttention(dim, num_heads, (-1, -1), 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)
+
+    def forward(
+        self,
+        x,
+        e,
+        seq_lens,
+        grid_sizes,
+        freqs,
+        context,
+        context_lens,
+    ):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, C]
+            e(Tensor): Shape [B, L1, 6, C]
+            seq_lens(Tensor): Shape [B], length of each sequence in batch
+            grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+        """
+        assert e.dtype == torch.float32
+        with torch.amp.autocast("cuda", dtype=torch.float32):
+            e = (self.modulation.unsqueeze(0) + e).chunk(6, dim=2)
+        assert e[0].dtype == torch.float32
+
+        # self-attention
+        y = self.self_attn(
+            self.norm1(x).float() * (1 + e[1].squeeze(2)) + e[0].squeeze(2), seq_lens, grid_sizes, freqs
+        )
+        with torch.amp.autocast("cuda", dtype=torch.float32):
+            x = x + y * e[2].squeeze(2)
+
+        # cross-attention & ffn function
+        def cross_attn_ffn(x, context, context_lens, e):
+            x = x + self.cross_attn(self.norm3(x), context, context_lens)
+            y = self.ffn(self.norm2(x).float() * (1 + e[4].squeeze(2)) + e[3].squeeze(2))
+            with torch.amp.autocast("cuda", dtype=torch.float32):
+                x = x + y * e[5].squeeze(2)
+            return x
+
+        x = cross_attn_ffn(x, context, context_lens, e)
+        return x
+
+
+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]
+        """
+        assert e.dtype == torch.float32
+        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(ModelMixin, ConfigMixin):
+    r"""
+    Wan diffusion backbone supporting both text-to-video and image-to-video.
+    """
+
+    ignore_for_config = ["patch_size", "cross_attn_norm", "qk_norm", "text_dim", "window_size"]
+    _no_split_modules = ["WanAttentionBlock"]
+
+    @register_to_config
+    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,
+        window_size=(-1, -1),
+        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
+            window_size (`tuple`, *optional*, defaults to (-1, -1)):
+                Window size for local attention (-1 indicates global attention)
+            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__()
+
+        assert model_type in ["t2v", "i2v", "ti2v", "s2v"]
+        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.window_size = window_size
+        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, window_size, 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())
+        assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0
+        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()
+
+    def forward(
+        self,
+        x,
+        t,
+        context,
+        seq_len,
+        y=None,
+    ):
+        r"""
+        Forward pass through the diffusion model
+
+        Args:
+            x (List[Tensor]):
+                List of input video tensors, each with shape [C_in, F, H, W]
+            t (Tensor):
+                Diffusion timesteps tensor of shape [B]
+            context (List[Tensor]):
+                List of text embeddings each with shape [L, C]
+            seq_len (`int`):
+                Maximum sequence length for positional encoding
+            y (List[Tensor], *optional*):
+                Conditional video inputs for image-to-video mode, same shape as x
+
+        Returns:
+            List[Tensor]:
+                List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8]
+        """
+        if self.model_type == "i2v":
+            assert y is not None
+        # params
+        device = self.patch_embedding.weight.device
+        if self.freqs.device != device:
+            self.freqs = self.freqs.to(device)
+
+        if y is not None:
+            x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]
+
+        # embeddings
+        x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
+        grid_sizes = torch.stack([torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
+        x = [u.flatten(2).transpose(1, 2) for u in x]
+        seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
+        assert seq_lens.max() <= seq_len
+        x = torch.cat([torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1) for u in x])
+
+        # time embeddings
+        if t.dim() == 1:
+            t = t.expand(t.size(0), seq_len)
+        with torch.amp.autocast("cuda", dtype=torch.float32):
+            bt = t.size(0)
+            t = t.flatten()
+            e = self.time_embedding(
+                sinusoidal_embedding_1d(self.freq_dim, t).unflatten(0, (bt, seq_len)).float()
+            )
+            e0 = self.time_projection(e).unflatten(2, (6, self.dim))
+            assert e.dtype == torch.float32 and e0.dtype == torch.float32
+
+        # context
+        context_lens = None
+        context = self.text_embedding(
+            torch.stack([torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) for u in context])
+        )
+
+        # arguments
+        kwargs = dict(
+            e=e0,
+            seq_lens=seq_lens,
+            grid_sizes=grid_sizes,
+            freqs=self.freqs,
+            context=context,
+            context_lens=context_lens,
+        )
+
+        for block in self.blocks:
+            x = block(x, **kwargs)
+
+        # head
+        x = self.head(x, e)
+
+        # unpatchify
+        x = self.unpatchify(x, grid_sizes)
+        return [u.float() for u in x]
+
+    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()):
+            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)])
+            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)

src/lerobot/policies/fastwam/wan/modules/t5.py

@@ -0,0 +1,488 @@
+# Modified from transformers.models.t5.modeling_t5
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import logging
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .tokenizers import HuggingfaceTokenizer
+
+__all__ = [
+    "T5Model",
+    "T5Encoder",
+    "T5Decoder",
+    "T5EncoderModel",
+]
+
+
+def fp16_clamp(x):
+    if x.dtype == torch.float16 and torch.isinf(x).any():
+        clamp = torch.finfo(x.dtype).max - 1000
+        x = torch.clamp(x, min=-clamp, max=clamp)
+    return x
+
+
+def init_weights(m):
+    if isinstance(m, T5LayerNorm):
+        nn.init.ones_(m.weight)
+    elif isinstance(m, T5Model):
+        nn.init.normal_(m.token_embedding.weight, std=1.0)
+    elif isinstance(m, T5FeedForward):
+        nn.init.normal_(m.gate[0].weight, std=m.dim**-0.5)
+        nn.init.normal_(m.fc1.weight, std=m.dim**-0.5)
+        nn.init.normal_(m.fc2.weight, std=m.dim_ffn**-0.5)
+    elif isinstance(m, T5Attention):
+        nn.init.normal_(m.q.weight, std=(m.dim * m.dim_attn) ** -0.5)
+        nn.init.normal_(m.k.weight, std=m.dim**-0.5)
+        nn.init.normal_(m.v.weight, std=m.dim**-0.5)
+        nn.init.normal_(m.o.weight, std=(m.num_heads * m.dim_attn) ** -0.5)
+    elif isinstance(m, T5RelativeEmbedding):
+        nn.init.normal_(m.embedding.weight, std=(2 * m.num_buckets * m.num_heads) ** -0.5)
+
+
+class GELU(nn.Module):
+    def forward(self, x):
+        return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))
+
+
+class T5LayerNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super(T5LayerNorm, self).__init__()
+        self.dim = dim
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        x = x * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + self.eps)
+        if self.weight.dtype in [torch.float16, torch.bfloat16]:
+            x = x.type_as(self.weight)
+        return self.weight * x
+
+
+class T5Attention(nn.Module):
+    def __init__(self, dim, dim_attn, num_heads, dropout=0.1):
+        assert dim_attn % num_heads == 0
+        super(T5Attention, self).__init__()
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.num_heads = num_heads
+        self.head_dim = dim_attn // num_heads
+
+        # layers
+        self.q = nn.Linear(dim, dim_attn, bias=False)
+        self.k = nn.Linear(dim, dim_attn, bias=False)
+        self.v = nn.Linear(dim, dim_attn, bias=False)
+        self.o = nn.Linear(dim_attn, dim, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, context=None, mask=None, pos_bias=None):
+        """
+        x:          [B, L1, C].
+        context:    [B, L2, C] or None.
+        mask:       [B, L2] or [B, L1, L2] or None.
+        """
+        # check inputs
+        context = x if context is None else context
+        b, n, c = x.size(0), self.num_heads, self.head_dim
+
+        # compute query, key, value
+        q = self.q(x).view(b, -1, n, c)
+        k = self.k(context).view(b, -1, n, c)
+        v = self.v(context).view(b, -1, n, c)
+
+        # attention bias
+        attn_bias = x.new_zeros(b, n, q.size(1), k.size(1))
+        if pos_bias is not None:
+            attn_bias += pos_bias
+        if mask is not None:
+            assert mask.ndim in [2, 3]
+            mask = mask.view(b, 1, 1, -1) if mask.ndim == 2 else mask.unsqueeze(1)
+            attn_bias.masked_fill_(mask == 0, torch.finfo(x.dtype).min)
+
+        # compute attention (T5 does not use scaling)
+        attn = torch.einsum("binc,bjnc->bnij", q, k) + attn_bias
+        attn = F.softmax(attn.float(), dim=-1).type_as(attn)
+        x = torch.einsum("bnij,bjnc->binc", attn, v)
+
+        # output
+        x = x.reshape(b, -1, n * c)
+        x = self.o(x)
+        x = self.dropout(x)
+        return x
+
+
+class T5FeedForward(nn.Module):
+    def __init__(self, dim, dim_ffn, dropout=0.1):
+        super(T5FeedForward, self).__init__()
+        self.dim = dim
+        self.dim_ffn = dim_ffn
+
+        # layers
+        self.gate = nn.Sequential(nn.Linear(dim, dim_ffn, bias=False), GELU())
+        self.fc1 = nn.Linear(dim, dim_ffn, bias=False)
+        self.fc2 = nn.Linear(dim_ffn, dim, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        x = self.fc1(x) * self.gate(x)
+        x = self.dropout(x)
+        x = self.fc2(x)
+        x = self.dropout(x)
+        return x
+
+
+class T5SelfAttention(nn.Module):
+    def __init__(self, dim, dim_attn, dim_ffn, num_heads, num_buckets, shared_pos=True, dropout=0.1):
+        super(T5SelfAttention, self).__init__()
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.dim_ffn = dim_ffn
+        self.num_heads = num_heads
+        self.num_buckets = num_buckets
+        self.shared_pos = shared_pos
+
+        # layers
+        self.norm1 = T5LayerNorm(dim)
+        self.attn = T5Attention(dim, dim_attn, num_heads, dropout)
+        self.norm2 = T5LayerNorm(dim)
+        self.ffn = T5FeedForward(dim, dim_ffn, dropout)
+        self.pos_embedding = (
+            None if shared_pos else T5RelativeEmbedding(num_buckets, num_heads, bidirectional=True)
+        )
+
+    def forward(self, x, mask=None, pos_bias=None):
+        e = pos_bias if self.shared_pos else self.pos_embedding(x.size(1), x.size(1))
+        x = fp16_clamp(x + self.attn(self.norm1(x), mask=mask, pos_bias=e))
+        x = fp16_clamp(x + self.ffn(self.norm2(x)))
+        return x
+
+
+class T5CrossAttention(nn.Module):
+    def __init__(self, dim, dim_attn, dim_ffn, num_heads, num_buckets, shared_pos=True, dropout=0.1):
+        super(T5CrossAttention, self).__init__()
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.dim_ffn = dim_ffn
+        self.num_heads = num_heads
+        self.num_buckets = num_buckets
+        self.shared_pos = shared_pos
+
+        # layers
+        self.norm1 = T5LayerNorm(dim)
+        self.self_attn = T5Attention(dim, dim_attn, num_heads, dropout)
+        self.norm2 = T5LayerNorm(dim)
+        self.cross_attn = T5Attention(dim, dim_attn, num_heads, dropout)
+        self.norm3 = T5LayerNorm(dim)
+        self.ffn = T5FeedForward(dim, dim_ffn, dropout)
+        self.pos_embedding = (
+            None if shared_pos else T5RelativeEmbedding(num_buckets, num_heads, bidirectional=False)
+        )
+
+    def forward(self, x, mask=None, encoder_states=None, encoder_mask=None, pos_bias=None):
+        e = pos_bias if self.shared_pos else self.pos_embedding(x.size(1), x.size(1))
+        x = fp16_clamp(x + self.self_attn(self.norm1(x), mask=mask, pos_bias=e))
+        x = fp16_clamp(x + self.cross_attn(self.norm2(x), context=encoder_states, mask=encoder_mask))
+        x = fp16_clamp(x + self.ffn(self.norm3(x)))
+        return x
+
+
+class T5RelativeEmbedding(nn.Module):
+    def __init__(self, num_buckets, num_heads, bidirectional, max_dist=128):
+        super(T5RelativeEmbedding, self).__init__()
+        self.num_buckets = num_buckets
+        self.num_heads = num_heads
+        self.bidirectional = bidirectional
+        self.max_dist = max_dist
+
+        # layers
+        self.embedding = nn.Embedding(num_buckets, num_heads)
+
+    def forward(self, lq, lk):
+        device = self.embedding.weight.device
+        # rel_pos = torch.arange(lk).unsqueeze(0).to(device) - \
+        #     torch.arange(lq).unsqueeze(1).to(device)
+        rel_pos = torch.arange(lk, device=device).unsqueeze(0) - torch.arange(lq, device=device).unsqueeze(1)
+        rel_pos = self._relative_position_bucket(rel_pos)
+        rel_pos_embeds = self.embedding(rel_pos)
+        rel_pos_embeds = rel_pos_embeds.permute(2, 0, 1).unsqueeze(0)  # [1, N, Lq, Lk]
+        return rel_pos_embeds.contiguous()
+
+    def _relative_position_bucket(self, rel_pos):
+        # preprocess
+        if self.bidirectional:
+            num_buckets = self.num_buckets // 2
+            rel_buckets = (rel_pos > 0).long() * num_buckets
+            rel_pos = torch.abs(rel_pos)
+        else:
+            num_buckets = self.num_buckets
+            rel_buckets = 0
+            rel_pos = -torch.min(rel_pos, torch.zeros_like(rel_pos))
+
+        # embeddings for small and large positions
+        max_exact = num_buckets // 2
+        rel_pos_large = (
+            max_exact
+            + (
+                torch.log(rel_pos.float() / max_exact)
+                / math.log(self.max_dist / max_exact)
+                * (num_buckets - max_exact)
+            ).long()
+        )
+        rel_pos_large = torch.min(rel_pos_large, torch.full_like(rel_pos_large, num_buckets - 1))
+        rel_buckets += torch.where(rel_pos < max_exact, rel_pos, rel_pos_large)
+        return rel_buckets
+
+
+class T5Encoder(nn.Module):
+    def __init__(
+        self, vocab, dim, dim_attn, dim_ffn, num_heads, num_layers, num_buckets, shared_pos=True, dropout=0.1
+    ):
+        super(T5Encoder, self).__init__()
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.dim_ffn = dim_ffn
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.num_buckets = num_buckets
+        self.shared_pos = shared_pos
+
+        # layers
+        self.token_embedding = vocab if isinstance(vocab, nn.Embedding) else nn.Embedding(vocab, dim)
+        self.pos_embedding = (
+            T5RelativeEmbedding(num_buckets, num_heads, bidirectional=True) if shared_pos else None
+        )
+        self.dropout = nn.Dropout(dropout)
+        self.blocks = nn.ModuleList(
+            [
+                T5SelfAttention(dim, dim_attn, dim_ffn, num_heads, num_buckets, shared_pos, dropout)
+                for _ in range(num_layers)
+            ]
+        )
+        self.norm = T5LayerNorm(dim)
+
+        # initialize weights
+        self.apply(init_weights)
+
+    def forward(self, ids, mask=None):
+        x = self.token_embedding(ids)
+        x = self.dropout(x)
+        e = self.pos_embedding(x.size(1), x.size(1)) if self.shared_pos else None
+        for block in self.blocks:
+            x = block(x, mask, pos_bias=e)
+        x = self.norm(x)
+        x = self.dropout(x)
+        return x
+
+
+class T5Decoder(nn.Module):
+    def __init__(
+        self, vocab, dim, dim_attn, dim_ffn, num_heads, num_layers, num_buckets, shared_pos=True, dropout=0.1
+    ):
+        super(T5Decoder, self).__init__()
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.dim_ffn = dim_ffn
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.num_buckets = num_buckets
+        self.shared_pos = shared_pos
+
+        # layers
+        self.token_embedding = vocab if isinstance(vocab, nn.Embedding) else nn.Embedding(vocab, dim)
+        self.pos_embedding = (
+            T5RelativeEmbedding(num_buckets, num_heads, bidirectional=False) if shared_pos else None
+        )
+        self.dropout = nn.Dropout(dropout)
+        self.blocks = nn.ModuleList(
+            [
+                T5CrossAttention(dim, dim_attn, dim_ffn, num_heads, num_buckets, shared_pos, dropout)
+                for _ in range(num_layers)
+            ]
+        )
+        self.norm = T5LayerNorm(dim)
+
+        # initialize weights
+        self.apply(init_weights)
+
+    def forward(self, ids, mask=None, encoder_states=None, encoder_mask=None):
+        b, s = ids.size()
+
+        # causal mask
+        if mask is None:
+            mask = torch.tril(torch.ones(1, s, s).to(ids.device))
+        elif mask.ndim == 2:
+            mask = torch.tril(mask.unsqueeze(1).expand(-1, s, -1))
+
+        # layers
+        x = self.token_embedding(ids)
+        x = self.dropout(x)
+        e = self.pos_embedding(x.size(1), x.size(1)) if self.shared_pos else None
+        for block in self.blocks:
+            x = block(x, mask, encoder_states, encoder_mask, pos_bias=e)
+        x = self.norm(x)
+        x = self.dropout(x)
+        return x
+
+
+class T5Model(nn.Module):
+    def __init__(
+        self,
+        vocab_size,
+        dim,
+        dim_attn,
+        dim_ffn,
+        num_heads,
+        encoder_layers,
+        decoder_layers,
+        num_buckets,
+        shared_pos=True,
+        dropout=0.1,
+    ):
+        super(T5Model, self).__init__()
+        self.vocab_size = vocab_size
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.dim_ffn = dim_ffn
+        self.num_heads = num_heads
+        self.encoder_layers = encoder_layers
+        self.decoder_layers = decoder_layers
+        self.num_buckets = num_buckets
+
+        # layers
+        self.token_embedding = nn.Embedding(vocab_size, dim)
+        self.encoder = T5Encoder(
+            self.token_embedding,
+            dim,
+            dim_attn,
+            dim_ffn,
+            num_heads,
+            encoder_layers,
+            num_buckets,
+            shared_pos,
+            dropout,
+        )
+        self.decoder = T5Decoder(
+            self.token_embedding,
+            dim,
+            dim_attn,
+            dim_ffn,
+            num_heads,
+            decoder_layers,
+            num_buckets,
+            shared_pos,
+            dropout,
+        )
+        self.head = nn.Linear(dim, vocab_size, bias=False)
+
+        # initialize weights
+        self.apply(init_weights)
+
+    def forward(self, encoder_ids, encoder_mask, decoder_ids, decoder_mask):
+        x = self.encoder(encoder_ids, encoder_mask)
+        x = self.decoder(decoder_ids, decoder_mask, x, encoder_mask)
+        x = self.head(x)
+        return x
+
+
+def _t5(
+    name,
+    encoder_only=False,
+    decoder_only=False,
+    return_tokenizer=False,
+    tokenizer_kwargs={},
+    dtype=torch.float32,
+    device="cpu",
+    **kwargs,
+):
+    # sanity check
+    assert not (encoder_only and decoder_only)
+
+    # params
+    if encoder_only:
+        model_cls = T5Encoder
+        kwargs["vocab"] = kwargs.pop("vocab_size")
+        kwargs["num_layers"] = kwargs.pop("encoder_layers")
+        _ = kwargs.pop("decoder_layers")
+    elif decoder_only:
+        model_cls = T5Decoder
+        kwargs["vocab"] = kwargs.pop("vocab_size")
+        kwargs["num_layers"] = kwargs.pop("decoder_layers")
+        _ = kwargs.pop("encoder_layers")
+    else:
+        model_cls = T5Model
+
+    # init model
+    with torch.device(device):
+        model = model_cls(**kwargs)
+
+    # set device
+    model = model.to(dtype=dtype, device=device)
+
+    # init tokenizer
+    if return_tokenizer:
+        from .tokenizers import HuggingfaceTokenizer
+
+        tokenizer = HuggingfaceTokenizer(f"google/{name}", **tokenizer_kwargs)
+        return model, tokenizer
+    else:
+        return model
+
+
+def umt5_xxl(**kwargs):
+    cfg = dict(
+        vocab_size=256384,
+        dim=4096,
+        dim_attn=4096,
+        dim_ffn=10240,
+        num_heads=64,
+        encoder_layers=24,
+        decoder_layers=24,
+        num_buckets=32,
+        shared_pos=False,
+        dropout=0.1,
+    )
+    cfg.update(**kwargs)
+    return _t5("umt5-xxl", **cfg)
+
+
+class T5EncoderModel:
+    def __init__(
+        self,
+        text_len,
+        dtype=torch.bfloat16,
+        device=torch.cuda.current_device(),
+        checkpoint_path=None,
+        tokenizer_path=None,
+        shard_fn=None,
+    ):
+        self.text_len = text_len
+        self.dtype = dtype
+        self.device = device
+        self.checkpoint_path = checkpoint_path
+        self.tokenizer_path = tokenizer_path
+
+        # init model
+        model = (
+            umt5_xxl(encoder_only=True, return_tokenizer=False, dtype=dtype, device=device)
+            .eval()
+            .requires_grad_(False)
+        )
+        logging.info(f"loading {checkpoint_path}")
+        model.load_state_dict(torch.load(checkpoint_path, map_location="cpu"))
+        self.model = model
+        if shard_fn is not None:
+            self.model = shard_fn(self.model, sync_module_states=False)
+        else:
+            self.model.to(self.device)
+        # init tokenizer
+        self.tokenizer = HuggingfaceTokenizer(name=tokenizer_path, seq_len=text_len, clean="whitespace")
+
+    def __call__(self, texts, device):
+        ids, mask = self.tokenizer(texts, return_mask=True, add_special_tokens=True)
+        ids = ids.to(device)
+        mask = mask.to(device)
+        seq_lens = mask.gt(0).sum(dim=1).long()
+        context = self.model(ids, mask)
+        return [u[:v] for u, v in zip(context, seq_lens)]

src/lerobot/policies/fastwam/wan/modules/tokenizers.py

@@ -0,0 +1,78 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import html
+import string
+
+import ftfy
+import regex as re
+from transformers import AutoTokenizer
+
+__all__ = ["HuggingfaceTokenizer"]
+
+
+def basic_clean(text):
+    text = ftfy.fix_text(text)
+    text = html.unescape(html.unescape(text))
+    return text.strip()
+
+
+def whitespace_clean(text):
+    text = re.sub(r"\s+", " ", text)
+    text = text.strip()
+    return text
+
+
+def canonicalize(text, keep_punctuation_exact_string=None):
+    text = text.replace("_", " ")
+    if keep_punctuation_exact_string:
+        text = keep_punctuation_exact_string.join(
+            part.translate(str.maketrans("", "", string.punctuation))
+            for part in text.split(keep_punctuation_exact_string)
+        )
+    else:
+        text = text.translate(str.maketrans("", "", string.punctuation))
+    text = text.lower()
+    text = re.sub(r"\s+", " ", text)
+    return text.strip()
+
+
+class HuggingfaceTokenizer:
+    def __init__(self, name, seq_len=None, clean=None, **kwargs):
+        assert clean in (None, "whitespace", "lower", "canonicalize")
+        self.name = name
+        self.seq_len = seq_len
+        self.clean = clean
+
+        # init tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(name, **kwargs)
+        self.vocab_size = self.tokenizer.vocab_size
+
+    def __call__(self, sequence, **kwargs):
+        return_mask = kwargs.pop("return_mask", False)
+
+        # arguments
+        _kwargs = {"return_tensors": "pt"}
+        if self.seq_len is not None:
+            _kwargs.update({"padding": "max_length", "truncation": True, "max_length": self.seq_len})
+        _kwargs.update(**kwargs)
+
+        # tokenization
+        if isinstance(sequence, str):
+            sequence = [sequence]
+        if self.clean:
+            sequence = [self._clean(u) for u in sequence]
+        ids = self.tokenizer(sequence, **_kwargs)
+
+        # output
+        if return_mask:
+            return ids.input_ids, ids.attention_mask
+        else:
+            return ids.input_ids
+
+    def _clean(self, text):
+        if self.clean == "whitespace":
+            text = whitespace_clean(basic_clean(text))
+        elif self.clean == "lower":
+            text = whitespace_clean(basic_clean(text)).lower()
+        elif self.clean == "canonicalize":
+            text = canonicalize(basic_clean(text))
+        return text

src/lerobot/policies/fastwam/wan/modules/vae2_1.py

@@ -0,0 +1,665 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import logging
+
+import torch
+import torch.cuda.amp as amp
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+__all__ = [
+    "Wan2_1_VAE",
+]
+
+CACHE_T = 2
+
+
+class CausalConv3d(nn.Conv3d):
+    """
+    Causal 3d convolusion.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._padding = (
+            self.padding[2],
+            self.padding[2],
+            self.padding[1],
+            self.padding[1],
+            2 * self.padding[0],
+            0,
+        )
+        self.padding = (0, 0, 0)
+
+    def forward(self, x, cache_x=None):
+        padding = list(self._padding)
+        if cache_x is not None and self._padding[4] > 0:
+            cache_x = cache_x.to(x.device)
+            x = torch.cat([cache_x, x], dim=2)
+            padding[4] -= cache_x.shape[2]
+        x = F.pad(x, padding)
+
+        return super().forward(x)
+
+
+class RMS_norm(nn.Module):
+    def __init__(self, dim, channel_first=True, images=True, bias=False):
+        super().__init__()
+        broadcastable_dims = (1, 1, 1) if not images else (1, 1)
+        shape = (dim, *broadcastable_dims) if channel_first else (dim,)
+
+        self.channel_first = channel_first
+        self.scale = dim**0.5
+        self.gamma = nn.Parameter(torch.ones(shape))
+        self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0
+
+    def forward(self, x):
+        return F.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias
+
+
+class Upsample(nn.Upsample):
+    def forward(self, x):
+        """
+        Fix bfloat16 support for nearest neighbor interpolation.
+        """
+        return super().forward(x.float()).type_as(x)
+
+
+class Resample(nn.Module):
+    def __init__(self, dim, mode):
+        assert mode in ("none", "upsample2d", "upsample3d", "downsample2d", "downsample3d")
+        super().__init__()
+        self.dim = dim
+        self.mode = mode
+
+        # layers
+        if mode == "upsample2d":
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
+                nn.Conv2d(dim, dim // 2, 3, padding=1),
+            )
+        elif mode == "upsample3d":
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
+                nn.Conv2d(dim, dim // 2, 3, padding=1),
+            )
+            self.time_conv = CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0))
+
+        elif mode == "downsample2d":
+            self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+        elif mode == "downsample3d":
+            self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+            self.time_conv = CausalConv3d(dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))
+
+        else:
+            self.resample = nn.Identity()
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        b, c, t, h, w = x.size()
+        if self.mode == "upsample3d":
+            if feat_cache is not None:
+                idx = feat_idx[0]
+                if feat_cache[idx] is None:
+                    feat_cache[idx] = "Rep"
+                    feat_idx[0] += 1
+                else:
+                    cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                    if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != "Rep":
+                        # cache last frame of last two chunk
+                        cache_x = torch.cat(
+                            [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2
+                        )
+                    if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] == "Rep":
+                        cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2)
+                    if feat_cache[idx] == "Rep":
+                        x = self.time_conv(x)
+                    else:
+                        x = self.time_conv(x, feat_cache[idx])
+                    feat_cache[idx] = cache_x
+                    feat_idx[0] += 1
+
+                    x = x.reshape(b, 2, c, t, h, w)
+                    x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3)
+                    x = x.reshape(b, c, t * 2, h, w)
+        t = x.shape[2]
+        x = rearrange(x, "b c t h w -> (b t) c h w")
+        x = self.resample(x)
+        x = rearrange(x, "(b t) c h w -> b c t h w", t=t)
+
+        if self.mode == "downsample3d":
+            if feat_cache is not None:
+                idx = feat_idx[0]
+                if feat_cache[idx] is None:
+                    feat_cache[idx] = x.clone()
+                    feat_idx[0] += 1
+                else:
+                    cache_x = x[:, :, -1:, :, :].clone()
+                    # if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx]!='Rep':
+                    #     # cache last frame of last two chunk
+                    #     cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
+
+                    x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))
+                    feat_cache[idx] = cache_x
+                    feat_idx[0] += 1
+        return x
+
+    def init_weight(self, conv):
+        conv_weight = conv.weight
+        nn.init.zeros_(conv_weight)
+        c1, c2, t, h, w = conv_weight.size()
+        one_matrix = torch.eye(c1, c2)
+        init_matrix = one_matrix
+        nn.init.zeros_(conv_weight)
+        # conv_weight.data[:,:,-1,1,1] = init_matrix * 0.5
+        conv_weight.data[:, :, 1, 0, 0] = init_matrix  # * 0.5
+        conv.weight.data.copy_(conv_weight)
+        nn.init.zeros_(conv.bias.data)
+
+    def init_weight2(self, conv):
+        conv_weight = conv.weight.data
+        nn.init.zeros_(conv_weight)
+        c1, c2, t, h, w = conv_weight.size()
+        init_matrix = torch.eye(c1 // 2, c2)
+        # init_matrix = repeat(init_matrix, 'o ... -> (o 2) ...').permute(1,0,2).contiguous().reshape(c1,c2)
+        conv_weight[: c1 // 2, :, -1, 0, 0] = init_matrix
+        conv_weight[c1 // 2 :, :, -1, 0, 0] = init_matrix
+        conv.weight.data.copy_(conv_weight)
+        nn.init.zeros_(conv.bias.data)
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_dim, out_dim, dropout=0.0):
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+
+        # layers
+        self.residual = nn.Sequential(
+            RMS_norm(in_dim, images=False),
+            nn.SiLU(),
+            CausalConv3d(in_dim, out_dim, 3, padding=1),
+            RMS_norm(out_dim, images=False),
+            nn.SiLU(),
+            nn.Dropout(dropout),
+            CausalConv3d(out_dim, out_dim, 3, padding=1),
+        )
+        self.shortcut = CausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity()
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        h = self.shortcut(x)
+        for layer in self.residual:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    # cache last frame of last two chunk
+                    cache_x = torch.cat(
+                        [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2
+                    )
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x + h
+
+
+class AttentionBlock(nn.Module):
+    """
+    Causal self-attention with a single head.
+    """
+
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+        # layers
+        self.norm = RMS_norm(dim)
+        self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
+        self.proj = nn.Conv2d(dim, dim, 1)
+
+        # zero out the last layer params
+        nn.init.zeros_(self.proj.weight)
+
+    def forward(self, x):
+        identity = x
+        b, c, t, h, w = x.size()
+        x = rearrange(x, "b c t h w -> (b t) c h w")
+        x = self.norm(x)
+        # compute query, key, value
+        q, k, v = (
+            self.to_qkv(x).reshape(b * t, 1, c * 3, -1).permute(0, 1, 3, 2).contiguous().chunk(3, dim=-1)
+        )
+
+        # apply attention
+        x = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+        )
+        x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w)
+
+        # output
+        x = self.proj(x)
+        x = rearrange(x, "(b t) c h w-> b c t h w", t=t)
+        return x + identity
+
+
+class Encoder3d(nn.Module):
+    def __init__(
+        self,
+        dim=128,
+        z_dim=4,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_downsample=[True, True, False],
+        dropout=0.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_downsample = temperal_downsample
+
+        # dimensions
+        dims = [dim * u for u in [1] + dim_mult]
+        scale = 1.0
+
+        # init block
+        self.conv1 = CausalConv3d(3, dims[0], 3, padding=1)
+
+        # downsample blocks
+        downsamples = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            # residual (+attention) blocks
+            for _ in range(num_res_blocks):
+                downsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+                if scale in attn_scales:
+                    downsamples.append(AttentionBlock(out_dim))
+                in_dim = out_dim
+
+            # downsample block
+            if i != len(dim_mult) - 1:
+                mode = "downsample3d" if temperal_downsample[i] else "downsample2d"
+                downsamples.append(Resample(out_dim, mode=mode))
+                scale /= 2.0
+        self.downsamples = nn.Sequential(*downsamples)
+
+        # middle blocks
+        self.middle = nn.Sequential(
+            ResidualBlock(out_dim, out_dim, dropout),
+            AttentionBlock(out_dim),
+            ResidualBlock(out_dim, out_dim, dropout),
+        )
+
+        # output blocks
+        self.head = nn.Sequential(
+            RMS_norm(out_dim, images=False), nn.SiLU(), CausalConv3d(out_dim, z_dim, 3, padding=1)
+        )
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                # cache last frame of last two chunk
+                cache_x = torch.cat(
+                    [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2
+                )
+            x = self.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(x)
+
+        ## downsamples
+        for layer in self.downsamples:
+            if feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        ## middle
+        for layer in self.middle:
+            if isinstance(layer, ResidualBlock) and feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        ## head
+        for layer in self.head:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    # cache last frame of last two chunk
+                    cache_x = torch.cat(
+                        [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2
+                    )
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x
+
+
+class Decoder3d(nn.Module):
+    def __init__(
+        self,
+        dim=128,
+        z_dim=4,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_upsample=[False, True, True],
+        dropout=0.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_upsample = temperal_upsample
+
+        # dimensions
+        dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]
+        scale = 1.0 / 2 ** (len(dim_mult) - 2)
+
+        # init block
+        self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1)
+
+        # middle blocks
+        self.middle = nn.Sequential(
+            ResidualBlock(dims[0], dims[0], dropout),
+            AttentionBlock(dims[0]),
+            ResidualBlock(dims[0], dims[0], dropout),
+        )
+
+        # upsample blocks
+        upsamples = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            # residual (+attention) blocks
+            if i == 1 or i == 2 or i == 3:
+                in_dim = in_dim // 2
+            for _ in range(num_res_blocks + 1):
+                upsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+                if scale in attn_scales:
+                    upsamples.append(AttentionBlock(out_dim))
+                in_dim = out_dim
+
+            # upsample block
+            if i != len(dim_mult) - 1:
+                mode = "upsample3d" if temperal_upsample[i] else "upsample2d"
+                upsamples.append(Resample(out_dim, mode=mode))
+                scale *= 2.0
+        self.upsamples = nn.Sequential(*upsamples)
+
+        # output blocks
+        self.head = nn.Sequential(
+            RMS_norm(out_dim, images=False), nn.SiLU(), CausalConv3d(out_dim, 3, 3, padding=1)
+        )
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        ## conv1
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                # cache last frame of last two chunk
+                cache_x = torch.cat(
+                    [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2
+                )
+            x = self.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(x)
+
+        ## middle
+        for layer in self.middle:
+            if isinstance(layer, ResidualBlock) and feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        ## upsamples
+        for layer in self.upsamples:
+            if feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        ## head
+        for layer in self.head:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    # cache last frame of last two chunk
+                    cache_x = torch.cat(
+                        [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2
+                    )
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x
+
+
+def count_conv3d(model):
+    count = 0
+    for m in model.modules():
+        if isinstance(m, CausalConv3d):
+            count += 1
+    return count
+
+
+class WanVAE_(nn.Module):
+    def __init__(
+        self,
+        dim=128,
+        z_dim=4,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_downsample=[True, True, False],
+        dropout=0.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_downsample = temperal_downsample
+        self.temperal_upsample = temperal_downsample[::-1]
+
+        # modules
+        self.encoder = Encoder3d(
+            dim, z_dim * 2, dim_mult, num_res_blocks, attn_scales, self.temperal_downsample, dropout
+        )
+        self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1)
+        self.conv2 = CausalConv3d(z_dim, z_dim, 1)
+        self.decoder = Decoder3d(
+            dim, z_dim, dim_mult, num_res_blocks, attn_scales, self.temperal_upsample, dropout
+        )
+
+    def forward(self, x):
+        mu, log_var = self.encode(x)
+        z = self.reparameterize(mu, log_var)
+        x_recon = self.decode(z)
+        return x_recon, mu, log_var
+
+    def encode(self, x, scale):
+        self.clear_cache()
+        ## cache
+        t = x.shape[2]
+        iter_ = 1 + (t - 1) // 4
+        ## 对encode输入的x，按时间拆分为1、4、4、4....
+        for i in range(iter_):
+            self._enc_conv_idx = [0]
+            if i == 0:
+                out = self.encoder(
+                    x[:, :, :1, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx
+                )
+            else:
+                out_ = self.encoder(
+                    x[:, :, 1 + 4 * (i - 1) : 1 + 4 * i, :, :],
+                    feat_cache=self._enc_feat_map,
+                    feat_idx=self._enc_conv_idx,
+                )
+                out = torch.cat([out, out_], 2)
+        mu, log_var = self.conv1(out).chunk(2, dim=1)
+        if isinstance(scale[0], torch.Tensor):
+            mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(1, self.z_dim, 1, 1, 1)
+        else:
+            mu = (mu - scale[0]) * scale[1]
+        self.clear_cache()
+        return mu
+
+    def decode(self, z, scale):
+        self.clear_cache()
+        # z: [b,c,t,h,w]
+        if isinstance(scale[0], torch.Tensor):
+            z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(1, self.z_dim, 1, 1, 1)
+        else:
+            z = z / scale[1] + scale[0]
+        iter_ = z.shape[2]
+        x = self.conv2(z)
+        for i in range(iter_):
+            self._conv_idx = [0]
+            if i == 0:
+                out = self.decoder(
+                    x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx
+                )
+            else:
+                out_ = self.decoder(
+                    x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx
+                )
+                out = torch.cat([out, out_], 2)
+        self.clear_cache()
+        return out
+
+    def reparameterize(self, mu, log_var):
+        std = torch.exp(0.5 * log_var)
+        eps = torch.randn_like(std)
+        return eps * std + mu
+
+    def sample(self, imgs, deterministic=False):
+        mu, log_var = self.encode(imgs)
+        if deterministic:
+            return mu
+        std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0))
+        return mu + std * torch.randn_like(std)
+
+    def clear_cache(self):
+        self._conv_num = count_conv3d(self.decoder)
+        self._conv_idx = [0]
+        self._feat_map = [None] * self._conv_num
+        # cache encode
+        self._enc_conv_num = count_conv3d(self.encoder)
+        self._enc_conv_idx = [0]
+        self._enc_feat_map = [None] * self._enc_conv_num
+
+
+def _video_vae(pretrained_path=None, z_dim=None, device="cpu", **kwargs):
+    """
+    Autoencoder3d adapted from Stable Diffusion 1.x, 2.x and XL.
+    """
+    # params
+    cfg = dict(
+        dim=96,
+        z_dim=z_dim,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_downsample=[False, True, True],
+        dropout=0.0,
+    )
+    cfg.update(**kwargs)
+
+    # init model
+    with torch.device("meta"):
+        model = WanVAE_(**cfg)
+
+    # load checkpoint
+    logging.info(f"loading {pretrained_path}")
+    model.load_state_dict(torch.load(pretrained_path, map_location=device), assign=True)
+
+    return model
+
+
+class Wan2_1_VAE:
+    def __init__(self, z_dim=16, vae_pth="cache/vae_step_411000.pth", dtype=torch.float, device="cuda"):
+        self.dtype = dtype
+        self.device = device
+
+        mean = [
+            -0.7571,
+            -0.7089,
+            -0.9113,
+            0.1075,
+            -0.1745,
+            0.9653,
+            -0.1517,
+            1.5508,
+            0.4134,
+            -0.0715,
+            0.5517,
+            -0.3632,
+            -0.1922,
+            -0.9497,
+            0.2503,
+            -0.2921,
+        ]
+        std = [
+            2.8184,
+            1.4541,
+            2.3275,
+            2.6558,
+            1.2196,
+            1.7708,
+            2.6052,
+            2.0743,
+            3.2687,
+            2.1526,
+            2.8652,
+            1.5579,
+            1.6382,
+            1.1253,
+            2.8251,
+            1.9160,
+        ]
+        self.mean = torch.tensor(mean, dtype=dtype, device=device)
+        self.std = torch.tensor(std, dtype=dtype, device=device)
+        self.scale = [self.mean, 1.0 / self.std]
+
+        # init model
+        self.model = (
+            _video_vae(
+                pretrained_path=vae_pth,
+                z_dim=z_dim,
+            )
+            .eval()
+            .requires_grad_(False)
+            .to(device)
+        )
+
+    def encode(self, videos):
+        """
+        videos: A list of videos each with shape [C, T, H, W].
+        """
+        with amp.autocast(dtype=self.dtype):
+            return [self.model.encode(u.unsqueeze(0), self.scale).float().squeeze(0) for u in videos]
+
+    def decode(self, zs):
+        with amp.autocast(dtype=self.dtype):
+            return [
+                self.model.decode(u.unsqueeze(0), self.scale).float().clamp_(-1, 1).squeeze(0) for u in zs
+            ]

src/lerobot/policies/fastwam/wan/modules/vae2_2.py

@@ -0,0 +1,995 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import logging
+
+import torch
+import torch.cuda.amp as amp
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+__all__ = [
+    "Wan2_2_VAE",
+]
+
+CACHE_T = 2
+
+
+class CausalConv3d(nn.Conv3d):
+    """
+    Causal 3d convolusion.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._padding = (
+            self.padding[2],
+            self.padding[2],
+            self.padding[1],
+            self.padding[1],
+            2 * self.padding[0],
+            0,
+        )
+        self.padding = (0, 0, 0)
+
+    def forward(self, x, cache_x=None):
+        padding = list(self._padding)
+        if cache_x is not None and self._padding[4] > 0:
+            cache_x = cache_x.to(x.device)
+            x = torch.cat([cache_x, x], dim=2)
+            padding[4] -= cache_x.shape[2]
+        x = F.pad(x, padding)
+
+        return super().forward(x)
+
+
+class RMS_norm(nn.Module):
+    def __init__(self, dim, channel_first=True, images=True, bias=False):
+        super().__init__()
+        broadcastable_dims = (1, 1, 1) if not images else (1, 1)
+        shape = (dim, *broadcastable_dims) if channel_first else (dim,)
+
+        self.channel_first = channel_first
+        self.scale = dim**0.5
+        self.gamma = nn.Parameter(torch.ones(shape))
+        self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0
+
+    def forward(self, x):
+        return F.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias
+
+
+class Upsample(nn.Upsample):
+    def forward(self, x):
+        """
+        Fix bfloat16 support for nearest neighbor interpolation.
+        """
+        return super().forward(x.float()).type_as(x)
+
+
+class Resample(nn.Module):
+    def __init__(self, dim, mode):
+        assert mode in (
+            "none",
+            "upsample2d",
+            "upsample3d",
+            "downsample2d",
+            "downsample3d",
+        )
+        super().__init__()
+        self.dim = dim
+        self.mode = mode
+
+        # layers
+        if mode == "upsample2d":
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
+                nn.Conv2d(dim, dim, 3, padding=1),
+            )
+        elif mode == "upsample3d":
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
+                nn.Conv2d(dim, dim, 3, padding=1),
+                # nn.Conv2d(dim, dim//2, 3, padding=1)
+            )
+            self.time_conv = CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0))
+        elif mode == "downsample2d":
+            self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+        elif mode == "downsample3d":
+            self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+            self.time_conv = CausalConv3d(dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))
+        else:
+            self.resample = nn.Identity()
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        b, c, t, h, w = x.size()
+        if self.mode == "upsample3d":
+            if feat_cache is not None:
+                idx = feat_idx[0]
+                if feat_cache[idx] is None:
+                    feat_cache[idx] = "Rep"
+                    feat_idx[0] += 1
+                else:
+                    cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                    if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != "Rep":
+                        # cache last frame of last two chunk
+                        cache_x = torch.cat(
+                            [
+                                feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device),
+                                cache_x,
+                            ],
+                            dim=2,
+                        )
+                    if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] == "Rep":
+                        cache_x = torch.cat(
+                            [torch.zeros_like(cache_x).to(cache_x.device), cache_x],
+                            dim=2,
+                        )
+                    if feat_cache[idx] == "Rep":
+                        x = self.time_conv(x)
+                    else:
+                        x = self.time_conv(x, feat_cache[idx])
+                    feat_cache[idx] = cache_x
+                    feat_idx[0] += 1
+                    x = x.reshape(b, 2, c, t, h, w)
+                    x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3)
+                    x = x.reshape(b, c, t * 2, h, w)
+        t = x.shape[2]
+        x = rearrange(x, "b c t h w -> (b t) c h w")
+        x = self.resample(x)
+        x = rearrange(x, "(b t) c h w -> b c t h w", t=t)
+
+        if self.mode == "downsample3d":
+            if feat_cache is not None:
+                idx = feat_idx[0]
+                if feat_cache[idx] is None:
+                    feat_cache[idx] = x.clone()
+                    feat_idx[0] += 1
+                else:
+                    cache_x = x[:, :, -1:, :, :].clone()
+                    x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))
+                    feat_cache[idx] = cache_x
+                    feat_idx[0] += 1
+        return x
+
+    def init_weight(self, conv):
+        conv_weight = conv.weight.detach().clone()
+        nn.init.zeros_(conv_weight)
+        c1, c2, t, h, w = conv_weight.size()
+        one_matrix = torch.eye(c1, c2)
+        init_matrix = one_matrix
+        nn.init.zeros_(conv_weight)
+        conv_weight.data[:, :, 1, 0, 0] = init_matrix  # * 0.5
+        conv.weight = nn.Parameter(conv_weight)
+        nn.init.zeros_(conv.bias.data)
+
+    def init_weight2(self, conv):
+        conv_weight = conv.weight.data.detach().clone()
+        nn.init.zeros_(conv_weight)
+        c1, c2, t, h, w = conv_weight.size()
+        init_matrix = torch.eye(c1 // 2, c2)
+        conv_weight[: c1 // 2, :, -1, 0, 0] = init_matrix
+        conv_weight[c1 // 2 :, :, -1, 0, 0] = init_matrix
+        conv.weight = nn.Parameter(conv_weight)
+        nn.init.zeros_(conv.bias.data)
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_dim, out_dim, dropout=0.0):
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+
+        # layers
+        self.residual = nn.Sequential(
+            RMS_norm(in_dim, images=False),
+            nn.SiLU(),
+            CausalConv3d(in_dim, out_dim, 3, padding=1),
+            RMS_norm(out_dim, images=False),
+            nn.SiLU(),
+            nn.Dropout(dropout),
+            CausalConv3d(out_dim, out_dim, 3, padding=1),
+        )
+        self.shortcut = CausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity()
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        h = self.shortcut(x)
+        for layer in self.residual:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    # cache last frame of last two chunk
+                    cache_x = torch.cat(
+                        [
+                            feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device),
+                            cache_x,
+                        ],
+                        dim=2,
+                    )
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x + h
+
+
+class AttentionBlock(nn.Module):
+    """
+    Causal self-attention with a single head.
+    """
+
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+        # layers
+        self.norm = RMS_norm(dim)
+        self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
+        self.proj = nn.Conv2d(dim, dim, 1)
+
+        # zero out the last layer params
+        nn.init.zeros_(self.proj.weight)
+
+    def forward(self, x):
+        identity = x
+        b, c, t, h, w = x.size()
+        x = rearrange(x, "b c t h w -> (b t) c h w")
+        x = self.norm(x)
+        # compute query, key, value
+        q, k, v = (
+            self.to_qkv(x).reshape(b * t, 1, c * 3, -1).permute(0, 1, 3, 2).contiguous().chunk(3, dim=-1)
+        )
+
+        # apply attention
+        x = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+        )
+        x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w)
+
+        # output
+        x = self.proj(x)
+        x = rearrange(x, "(b t) c h w-> b c t h w", t=t)
+        return x + identity
+
+
+def patchify(x, patch_size):
+    if patch_size == 1:
+        return x
+    if x.dim() == 4:
+        x = rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size)
+    elif x.dim() == 5:
+        x = rearrange(
+            x,
+            "b c f (h q) (w r) -> b (c r q) f h w",
+            q=patch_size,
+            r=patch_size,
+        )
+    else:
+        raise ValueError(f"Invalid input shape: {x.shape}")
+
+    return x
+
+
+def unpatchify(x, patch_size):
+    if patch_size == 1:
+        return x
+
+    if x.dim() == 4:
+        x = rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size)
+    elif x.dim() == 5:
+        x = rearrange(
+            x,
+            "b (c r q) f h w -> b c f (h q) (w r)",
+            q=patch_size,
+            r=patch_size,
+        )
+    return x
+
+
+class AvgDown3D(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        factor_t,
+        factor_s=1,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.factor_t = factor_t
+        self.factor_s = factor_s
+        self.factor = self.factor_t * self.factor_s * self.factor_s
+
+        assert in_channels * self.factor % out_channels == 0
+        self.group_size = in_channels * self.factor // out_channels
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t
+        pad = (0, 0, 0, 0, pad_t, 0)
+        x = F.pad(x, pad)
+        B, C, T, H, W = x.shape
+        x = x.view(
+            B,
+            C,
+            T // self.factor_t,
+            self.factor_t,
+            H // self.factor_s,
+            self.factor_s,
+            W // self.factor_s,
+            self.factor_s,
+        )
+        x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous()
+        x = x.view(
+            B,
+            C * self.factor,
+            T // self.factor_t,
+            H // self.factor_s,
+            W // self.factor_s,
+        )
+        x = x.view(
+            B,
+            self.out_channels,
+            self.group_size,
+            T // self.factor_t,
+            H // self.factor_s,
+            W // self.factor_s,
+        )
+        x = x.mean(dim=2)
+        return x
+
+
+class DupUp3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        factor_t,
+        factor_s=1,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        self.factor_t = factor_t
+        self.factor_s = factor_s
+        self.factor = self.factor_t * self.factor_s * self.factor_s
+
+        assert out_channels * self.factor % in_channels == 0
+        self.repeats = out_channels * self.factor // in_channels
+
+    def forward(self, x: torch.Tensor, first_chunk=False) -> torch.Tensor:
+        x = x.repeat_interleave(self.repeats, dim=1)
+        x = x.view(
+            x.size(0),
+            self.out_channels,
+            self.factor_t,
+            self.factor_s,
+            self.factor_s,
+            x.size(2),
+            x.size(3),
+            x.size(4),
+        )
+        x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
+        x = x.view(
+            x.size(0),
+            self.out_channels,
+            x.size(2) * self.factor_t,
+            x.size(4) * self.factor_s,
+            x.size(6) * self.factor_s,
+        )
+        if first_chunk:
+            x = x[:, :, self.factor_t - 1 :, :, :]
+        return x
+
+
+class Down_ResidualBlock(nn.Module):
+    def __init__(self, in_dim, out_dim, dropout, mult, temperal_downsample=False, down_flag=False):
+        super().__init__()
+
+        # Shortcut path with downsample
+        self.avg_shortcut = AvgDown3D(
+            in_dim,
+            out_dim,
+            factor_t=2 if temperal_downsample else 1,
+            factor_s=2 if down_flag else 1,
+        )
+
+        # Main path with residual blocks and downsample
+        downsamples = []
+        for _ in range(mult):
+            downsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+            in_dim = out_dim
+
+        # Add the final downsample block
+        if down_flag:
+            mode = "downsample3d" if temperal_downsample else "downsample2d"
+            downsamples.append(Resample(out_dim, mode=mode))
+
+        self.downsamples = nn.Sequential(*downsamples)
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        x_copy = x.clone()
+        for module in self.downsamples:
+            x = module(x, feat_cache, feat_idx)
+
+        return x + self.avg_shortcut(x_copy)
+
+
+class Up_ResidualBlock(nn.Module):
+    def __init__(self, in_dim, out_dim, dropout, mult, temperal_upsample=False, up_flag=False):
+        super().__init__()
+        # Shortcut path with upsample
+        if up_flag:
+            self.avg_shortcut = DupUp3D(
+                in_dim,
+                out_dim,
+                factor_t=2 if temperal_upsample else 1,
+                factor_s=2 if up_flag else 1,
+            )
+        else:
+            self.avg_shortcut = None
+
+        # Main path with residual blocks and upsample
+        upsamples = []
+        for _ in range(mult):
+            upsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+            in_dim = out_dim
+
+        # Add the final upsample block
+        if up_flag:
+            mode = "upsample3d" if temperal_upsample else "upsample2d"
+            upsamples.append(Resample(out_dim, mode=mode))
+
+        self.upsamples = nn.Sequential(*upsamples)
+
+    def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False):
+        x_main = x.clone()
+        for module in self.upsamples:
+            x_main = module(x_main, feat_cache, feat_idx)
+        if self.avg_shortcut is not None:
+            x_shortcut = self.avg_shortcut(x, first_chunk)
+            return x_main + x_shortcut
+        else:
+            return x_main
+
+
+class Encoder3d(nn.Module):
+    def __init__(
+        self,
+        dim=128,
+        z_dim=4,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_downsample=[True, True, False],
+        dropout=0.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_downsample = temperal_downsample
+
+        # dimensions
+        dims = [dim * u for u in [1] + dim_mult]
+        scale = 1.0
+
+        # init block
+        self.conv1 = CausalConv3d(12, dims[0], 3, padding=1)
+
+        # downsample blocks
+        downsamples = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            t_down_flag = temperal_downsample[i] if i < len(temperal_downsample) else False
+            downsamples.append(
+                Down_ResidualBlock(
+                    in_dim=in_dim,
+                    out_dim=out_dim,
+                    dropout=dropout,
+                    mult=num_res_blocks,
+                    temperal_downsample=t_down_flag,
+                    down_flag=i != len(dim_mult) - 1,
+                )
+            )
+            scale /= 2.0
+        self.downsamples = nn.Sequential(*downsamples)
+
+        # middle blocks
+        self.middle = nn.Sequential(
+            ResidualBlock(out_dim, out_dim, dropout),
+            AttentionBlock(out_dim),
+            ResidualBlock(out_dim, out_dim, dropout),
+        )
+
+        # # output blocks
+        self.head = nn.Sequential(
+            RMS_norm(out_dim, images=False),
+            nn.SiLU(),
+            CausalConv3d(out_dim, z_dim, 3, padding=1),
+        )
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                cache_x = torch.cat(
+                    [
+                        feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device),
+                        cache_x,
+                    ],
+                    dim=2,
+                )
+            x = self.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(x)
+
+        ## downsamples
+        for layer in self.downsamples:
+            if feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        ## middle
+        for layer in self.middle:
+            if isinstance(layer, ResidualBlock) and feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        ## head
+        for layer in self.head:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    cache_x = torch.cat(
+                        [
+                            feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device),
+                            cache_x,
+                        ],
+                        dim=2,
+                    )
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+
+        return x
+
+
+class Decoder3d(nn.Module):
+    def __init__(
+        self,
+        dim=128,
+        z_dim=4,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_upsample=[False, True, True],
+        dropout=0.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_upsample = temperal_upsample
+
+        # dimensions
+        dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]
+        scale = 1.0 / 2 ** (len(dim_mult) - 2)
+        # init block
+        self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1)
+
+        # middle blocks
+        self.middle = nn.Sequential(
+            ResidualBlock(dims[0], dims[0], dropout),
+            AttentionBlock(dims[0]),
+            ResidualBlock(dims[0], dims[0], dropout),
+        )
+
+        # upsample blocks
+        upsamples = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            t_up_flag = temperal_upsample[i] if i < len(temperal_upsample) else False
+            upsamples.append(
+                Up_ResidualBlock(
+                    in_dim=in_dim,
+                    out_dim=out_dim,
+                    dropout=dropout,
+                    mult=num_res_blocks + 1,
+                    temperal_upsample=t_up_flag,
+                    up_flag=i != len(dim_mult) - 1,
+                )
+            )
+        self.upsamples = nn.Sequential(*upsamples)
+
+        # output blocks
+        self.head = nn.Sequential(
+            RMS_norm(out_dim, images=False),
+            nn.SiLU(),
+            CausalConv3d(out_dim, 12, 3, padding=1),
+        )
+
+    def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False):
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                cache_x = torch.cat(
+                    [
+                        feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device),
+                        cache_x,
+                    ],
+                    dim=2,
+                )
+            x = self.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(x)
+
+        for layer in self.middle:
+            if isinstance(layer, ResidualBlock) and feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        ## upsamples
+        for layer in self.upsamples:
+            if feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx, first_chunk)
+            else:
+                x = layer(x)
+
+        ## head
+        for layer in self.head:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    cache_x = torch.cat(
+                        [
+                            feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device),
+                            cache_x,
+                        ],
+                        dim=2,
+                    )
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x
+
+
+def count_conv3d(model):
+    count = 0
+    for m in model.modules():
+        if isinstance(m, CausalConv3d):
+            count += 1
+    return count
+
+
+class WanVAE_(nn.Module):
+    def __init__(
+        self,
+        dim=160,
+        dec_dim=256,
+        z_dim=16,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_downsample=[True, True, False],
+        dropout=0.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_downsample = temperal_downsample
+        self.temperal_upsample = temperal_downsample[::-1]
+
+        # modules
+        self.encoder = Encoder3d(
+            dim,
+            z_dim * 2,
+            dim_mult,
+            num_res_blocks,
+            attn_scales,
+            self.temperal_downsample,
+            dropout,
+        )
+        self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1)
+        self.conv2 = CausalConv3d(z_dim, z_dim, 1)
+        self.decoder = Decoder3d(
+            dec_dim,
+            z_dim,
+            dim_mult,
+            num_res_blocks,
+            attn_scales,
+            self.temperal_upsample,
+            dropout,
+        )
+
+    def forward(self, x, scale=[0, 1]):
+        mu = self.encode(x, scale)
+        x_recon = self.decode(mu, scale)
+        return x_recon, mu
+
+    def encode(self, x, scale):
+        self.clear_cache()
+        x = patchify(x, patch_size=2)
+        t = x.shape[2]
+        iter_ = 1 + (t - 1) // 4
+        for i in range(iter_):
+            self._enc_conv_idx = [0]
+            if i == 0:
+                out = self.encoder(
+                    x[:, :, :1, :, :],
+                    feat_cache=self._enc_feat_map,
+                    feat_idx=self._enc_conv_idx,
+                )
+            else:
+                out_ = self.encoder(
+                    x[:, :, 1 + 4 * (i - 1) : 1 + 4 * i, :, :],
+                    feat_cache=self._enc_feat_map,
+                    feat_idx=self._enc_conv_idx,
+                )
+                out = torch.cat([out, out_], 2)
+        mu, log_var = self.conv1(out).chunk(2, dim=1)
+        if isinstance(scale[0], torch.Tensor):
+            mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(1, self.z_dim, 1, 1, 1)
+        else:
+            mu = (mu - scale[0]) * scale[1]
+        self.clear_cache()
+        return mu
+
+    def decode(self, z, scale):
+        self.clear_cache()
+        if isinstance(scale[0], torch.Tensor):
+            z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(1, self.z_dim, 1, 1, 1)
+        else:
+            z = z / scale[1] + scale[0]
+        iter_ = z.shape[2]
+        x = self.conv2(z)
+        for i in range(iter_):
+            self._conv_idx = [0]
+            if i == 0:
+                out = self.decoder(
+                    x[:, :, i : i + 1, :, :],
+                    feat_cache=self._feat_map,
+                    feat_idx=self._conv_idx,
+                    first_chunk=True,
+                )
+            else:
+                out_ = self.decoder(
+                    x[:, :, i : i + 1, :, :],
+                    feat_cache=self._feat_map,
+                    feat_idx=self._conv_idx,
+                )
+                out = torch.cat([out, out_], 2)
+        out = unpatchify(out, patch_size=2)
+        self.clear_cache()
+        return out
+
+    def reparameterize(self, mu, log_var):
+        std = torch.exp(0.5 * log_var)
+        eps = torch.randn_like(std)
+        return eps * std + mu
+
+    def sample(self, imgs, deterministic=False):
+        mu, log_var = self.encode(imgs)
+        if deterministic:
+            return mu
+        std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0))
+        return mu + std * torch.randn_like(std)
+
+    def clear_cache(self):
+        self._conv_num = count_conv3d(self.decoder)
+        self._conv_idx = [0]
+        self._feat_map = [None] * self._conv_num
+        # cache encode
+        self._enc_conv_num = count_conv3d(self.encoder)
+        self._enc_conv_idx = [0]
+        self._enc_feat_map = [None] * self._enc_conv_num
+
+
+def _video_vae(pretrained_path=None, z_dim=16, dim=160, device="cpu", **kwargs):
+    # params
+    cfg = dict(
+        dim=dim,
+        z_dim=z_dim,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_downsample=[True, True, True],
+        dropout=0.0,
+    )
+    cfg.update(**kwargs)
+
+    # init model
+    with torch.device("meta"):
+        model = WanVAE_(**cfg)
+
+    # load checkpoint
+    logging.info(f"loading {pretrained_path}")
+    model.load_state_dict(torch.load(pretrained_path, map_location=device), assign=True)
+
+    return model
+
+
+class Wan2_2_VAE:
+    def __init__(
+        self,
+        z_dim=48,
+        c_dim=160,
+        vae_pth=None,
+        dim_mult=[1, 2, 4, 4],
+        temperal_downsample=[False, True, True],
+        dtype=torch.float,
+        device="cuda",
+    ):
+
+        self.dtype = dtype
+        self.device = device
+
+        mean = torch.tensor(
+            [
+                -0.2289,
+                -0.0052,
+                -0.1323,
+                -0.2339,
+                -0.2799,
+                0.0174,
+                0.1838,
+                0.1557,
+                -0.1382,
+                0.0542,
+                0.2813,
+                0.0891,
+                0.1570,
+                -0.0098,
+                0.0375,
+                -0.1825,
+                -0.2246,
+                -0.1207,
+                -0.0698,
+                0.5109,
+                0.2665,
+                -0.2108,
+                -0.2158,
+                0.2502,
+                -0.2055,
+                -0.0322,
+                0.1109,
+                0.1567,
+                -0.0729,
+                0.0899,
+                -0.2799,
+                -0.1230,
+                -0.0313,
+                -0.1649,
+                0.0117,
+                0.0723,
+                -0.2839,
+                -0.2083,
+                -0.0520,
+                0.3748,
+                0.0152,
+                0.1957,
+                0.1433,
+                -0.2944,
+                0.3573,
+                -0.0548,
+                -0.1681,
+                -0.0667,
+            ],
+            dtype=dtype,
+            device=device,
+        )
+        std = torch.tensor(
+            [
+                0.4765,
+                1.0364,
+                0.4514,
+                1.1677,
+                0.5313,
+                0.4990,
+                0.4818,
+                0.5013,
+                0.8158,
+                1.0344,
+                0.5894,
+                1.0901,
+                0.6885,
+                0.6165,
+                0.8454,
+                0.4978,
+                0.5759,
+                0.3523,
+                0.7135,
+                0.6804,
+                0.5833,
+                1.4146,
+                0.8986,
+                0.5659,
+                0.7069,
+                0.5338,
+                0.4889,
+                0.4917,
+                0.4069,
+                0.4999,
+                0.6866,
+                0.4093,
+                0.5709,
+                0.6065,
+                0.6415,
+                0.4944,
+                0.5726,
+                1.2042,
+                0.5458,
+                1.6887,
+                0.3971,
+                1.0600,
+                0.3943,
+                0.5537,
+                0.5444,
+                0.4089,
+                0.7468,
+                0.7744,
+            ],
+            dtype=dtype,
+            device=device,
+        )
+        self.scale = [mean, 1.0 / std]
+
+        # init model
+        self.model = (
+            _video_vae(
+                pretrained_path=vae_pth,
+                z_dim=z_dim,
+                dim=c_dim,
+                dim_mult=dim_mult,
+                temperal_downsample=temperal_downsample,
+            )
+            .eval()
+            .requires_grad_(False)
+            .to(device)
+        )
+
+    def encode(self, videos):
+        try:
+            if not isinstance(videos, list):
+                raise TypeError("videos should be a list")
+            with amp.autocast(dtype=self.dtype):
+                return [self.model.encode(u.unsqueeze(0), self.scale).float().squeeze(0) for u in videos]
+        except TypeError as e:
+            logging.info(e)
+            return None
+
+    def decode(self, zs):
+        try:
+            if not isinstance(zs, list):
+                raise TypeError("zs should be a list")
+            with amp.autocast(dtype=self.dtype):
+                return [
+                    self.model.decode(u.unsqueeze(0), self.scale).float().clamp_(-1, 1).squeeze(0) for u in zs
+                ]
+        except TypeError as e:
+            logging.info(e)
+            return None

src/lerobot/policies/fastwam/wan/utils/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+from .fm_solvers import (
+    FlowDPMSolverMultistepScheduler,
+    get_sampling_sigmas,
+    retrieve_timesteps,
+)
+from .fm_solvers_unipc import FlowUniPCMultistepScheduler
+
+__all__ = [
+    "HuggingfaceTokenizer",
+    "get_sampling_sigmas",
+    "retrieve_timesteps",
+    "FlowDPMSolverMultistepScheduler",
+    "FlowUniPCMultistepScheduler",
+]

src/lerobot/policies/fastwam/wan/utils/fm_solvers.py

@@ -0,0 +1,837 @@
+# Copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py
+# Convert dpm solver for flow matching
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+
+import inspect
+import math
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import (
+    KarrasDiffusionSchedulers,
+    SchedulerMixin,
+    SchedulerOutput,
+)
+from diffusers.utils import deprecate, is_scipy_available
+from diffusers.utils.torch_utils import randn_tensor
+
+if is_scipy_available():
+    pass
+
+
+def get_sampling_sigmas(sampling_steps, shift):
+    sigma = np.linspace(1, 0, sampling_steps + 1)[:sampling_steps]
+    sigma = shift * sigma / (1 + (shift - 1) * sigma)
+
+    return sigma
+
+
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps=None,
+    device=None,
+    timesteps=None,
+    sigmas=None,
+    **kwargs,
+):
+    if timesteps is not None and sigmas is not None:
+        raise ValueError(
+            "Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values"
+        )
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+class FlowDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
+    """
+    `FlowDPMSolverMultistepScheduler` is a fast dedicated high-order solver for diffusion ODEs.
+    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
+    methods the library implements for all schedulers such as loading and saving.
+    Args:
+        num_train_timesteps (`int`, defaults to 1000):
+            The number of diffusion steps to train the model. This determines the resolution of the diffusion process.
+        solver_order (`int`, defaults to 2):
+            The DPMSolver order which can be `1`, `2`, or `3`. It is recommended to use `solver_order=2` for guided
+            sampling, and `solver_order=3` for unconditional sampling. This affects the number of model outputs stored
+            and used in multistep updates.
+        prediction_type (`str`, defaults to "flow_prediction"):
+            Prediction type of the scheduler function; must be `flow_prediction` for this scheduler, which predicts
+            the flow of the diffusion process.
+        shift (`float`, *optional*, defaults to 1.0):
+            A factor used to adjust the sigmas in the noise schedule. It modifies the step sizes during the sampling
+            process.
+        use_dynamic_shifting (`bool`, defaults to `False`):
+            Whether to apply dynamic shifting to the timesteps based on image resolution. If `True`, the shifting is
+            applied on the fly.
+        thresholding (`bool`, defaults to `False`):
+            Whether to use the "dynamic thresholding" method. This method adjusts the predicted sample to prevent
+            saturation and improve photorealism.
+        dynamic_thresholding_ratio (`float`, defaults to 0.995):
+            The ratio for the dynamic thresholding method. Valid only when `thresholding=True`.
+        sample_max_value (`float`, defaults to 1.0):
+            The threshold value for dynamic thresholding. Valid only when `thresholding=True` and
+            `algorithm_type="dpmsolver++"`.
+        algorithm_type (`str`, defaults to `dpmsolver++`):
+            Algorithm type for the solver; can be `dpmsolver`, `dpmsolver++`, `sde-dpmsolver` or `sde-dpmsolver++`. The
+            `dpmsolver` type implements the algorithms in the [DPMSolver](https://huggingface.co/papers/2206.00927)
+            paper, and the `dpmsolver++` type implements the algorithms in the
+            [DPMSolver++](https://huggingface.co/papers/2211.01095) paper. It is recommended to use `dpmsolver++` or
+            `sde-dpmsolver++` with `solver_order=2` for guided sampling like in Stable Diffusion.
+        solver_type (`str`, defaults to `midpoint`):
+            Solver type for the second-order solver; can be `midpoint` or `heun`. The solver type slightly affects the
+            sample quality, especially for a small number of steps. It is recommended to use `midpoint` solvers.
+        lower_order_final (`bool`, defaults to `True`):
+            Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. This can
+            stabilize the sampling of DPMSolver for steps < 15, especially for steps <= 10.
+        euler_at_final (`bool`, defaults to `False`):
+            Whether to use Euler's method in the final step. It is a trade-off between numerical stability and detail
+            richness. This can stabilize the sampling of the SDE variant of DPMSolver for small number of inference
+            steps, but sometimes may result in blurring.
+        final_sigmas_type (`str`, *optional*, defaults to "zero"):
+            The final `sigma` value for the noise schedule during the sampling process. If `"sigma_min"`, the final
+            sigma is the same as the last sigma in the training schedule. If `zero`, the final sigma is set to 0.
+        lambda_min_clipped (`float`, defaults to `-inf`):
+            Clipping threshold for the minimum value of `lambda(t)` for numerical stability. This is critical for the
+            cosine (`squaredcos_cap_v2`) noise schedule.
+        variance_type (`str`, *optional*):
+            Set to "learned" or "learned_range" for diffusion models that predict variance. If set, the model's output
+            contains the predicted Gaussian variance.
+    """
+
+    _compatibles = [e.name for e in KarrasDiffusionSchedulers]
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        solver_order: int = 2,
+        prediction_type: str = "flow_prediction",
+        shift: Optional[float] = 1.0,
+        use_dynamic_shifting=False,
+        thresholding: bool = False,
+        dynamic_thresholding_ratio: float = 0.995,
+        sample_max_value: float = 1.0,
+        algorithm_type: str = "dpmsolver++",
+        solver_type: str = "midpoint",
+        lower_order_final: bool = True,
+        euler_at_final: bool = False,
+        final_sigmas_type: Optional[str] = "zero",  # "zero", "sigma_min"
+        lambda_min_clipped: float = -float("inf"),
+        variance_type: Optional[str] = None,
+        invert_sigmas: bool = False,
+    ):
+        if algorithm_type in ["dpmsolver", "sde-dpmsolver"]:
+            deprecation_message = f"algorithm_type {algorithm_type} is deprecated and will be removed in a future version. Choose from `dpmsolver++` or `sde-dpmsolver++` instead"
+            deprecate("algorithm_types dpmsolver and sde-dpmsolver", "1.0.0", deprecation_message)
+
+        # settings for DPM-Solver
+        if algorithm_type not in ["dpmsolver", "dpmsolver++", "sde-dpmsolver", "sde-dpmsolver++"]:
+            if algorithm_type == "deis":
+                self.register_to_config(algorithm_type="dpmsolver++")
+            else:
+                raise NotImplementedError(f"{algorithm_type} is not implemented for {self.__class__}")
+
+        if solver_type not in ["midpoint", "heun"]:
+            if solver_type in ["logrho", "bh1", "bh2"]:
+                self.register_to_config(solver_type="midpoint")
+            else:
+                raise NotImplementedError(f"{solver_type} is not implemented for {self.__class__}")
+
+        if algorithm_type not in ["dpmsolver++", "sde-dpmsolver++"] and final_sigmas_type == "zero":
+            raise ValueError(
+                f"`final_sigmas_type` {final_sigmas_type} is not supported for `algorithm_type` {algorithm_type}. Please choose `sigma_min` instead."
+            )
+
+        # setable values
+        self.num_inference_steps = None
+        alphas = np.linspace(1, 1 / num_train_timesteps, num_train_timesteps)[::-1].copy()
+        sigmas = 1.0 - alphas
+        sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32)
+
+        if not use_dynamic_shifting:
+            # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution
+            sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)  # pyright: ignore
+
+        self.sigmas = sigmas
+        self.timesteps = sigmas * num_train_timesteps
+
+        self.model_outputs = [None] * solver_order
+        self.lower_order_nums = 0
+        self._step_index = None
+        self._begin_index = None
+
+        # self.sigmas = self.sigmas.to(
+        #     "cpu")  # to avoid too much CPU/GPU communication
+        self.sigma_min = self.sigmas[-1].item()
+        self.sigma_max = self.sigmas[0].item()
+
+    @property
+    def step_index(self):
+        """
+        The index counter for current timestep. It will increase 1 after each scheduler step.
+        """
+        return self._step_index
+
+    @property
+    def begin_index(self):
+        """
+        The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
+        """
+        return self._begin_index
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
+    def set_begin_index(self, begin_index: int = 0):
+        """
+        Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
+        Args:
+            begin_index (`int`):
+                The begin index for the scheduler.
+        """
+        self._begin_index = begin_index
+
+    # Modified from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.set_timesteps
+    def set_timesteps(
+        self,
+        num_inference_steps: Union[int, None] = None,
+        device: Union[str, torch.device] = None,
+        sigmas: Optional[List[float]] = None,
+        mu: Optional[Union[float, None]] = None,
+        shift: Optional[Union[float, None]] = None,
+    ):
+        """
+        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
+        Args:
+            num_inference_steps (`int`):
+                Total number of the spacing of the time steps.
+            device (`str` or `torch.device`, *optional*):
+                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        """
+
+        if self.config.use_dynamic_shifting and mu is None:
+            raise ValueError(
+                " you have to pass a value for `mu` when `use_dynamic_shifting` is set to be `True`"
+            )
+
+        if sigmas is None:
+            sigmas = np.linspace(self.sigma_max, self.sigma_min, num_inference_steps + 1).copy()[:-1]  # pyright: ignore
+
+        if self.config.use_dynamic_shifting:
+            sigmas = self.time_shift(mu, 1.0, sigmas)  # pyright: ignore
+        else:
+            if shift is None:
+                shift = self.config.shift
+            sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)  # pyright: ignore
+
+        if self.config.final_sigmas_type == "sigma_min":
+            sigma_last = ((1 - self.alphas_cumprod[0]) / self.alphas_cumprod[0]) ** 0.5
+        elif self.config.final_sigmas_type == "zero":
+            sigma_last = 0
+        else:
+            raise ValueError(
+                f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}"
+            )
+
+        timesteps = sigmas * self.config.num_train_timesteps
+        sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32)  # pyright: ignore
+
+        self.sigmas = torch.from_numpy(sigmas)
+        self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.int64)
+
+        self.num_inference_steps = len(timesteps)
+
+        self.model_outputs = [
+            None,
+        ] * self.config.solver_order
+        self.lower_order_nums = 0
+
+        self._step_index = None
+        self._begin_index = None
+        # self.sigmas = self.sigmas.to(
+        #     "cpu")  # to avoid too much CPU/GPU communication
+
+    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample
+    def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
+        """
+        "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the
+        prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by
+        s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing
+        pixels from saturation at each step. We find that dynamic thresholding results in significantly better
+        photorealism as well as better image-text alignment, especially when using very large guidance weights."
+        https://arxiv.org/abs/2205.11487
+        """
+        dtype = sample.dtype
+        batch_size, channels, *remaining_dims = sample.shape
+
+        if dtype not in (torch.float32, torch.float64):
+            sample = sample.float()  # upcast for quantile calculation, and clamp not implemented for cpu half
+
+        # Flatten sample for doing quantile calculation along each image
+        sample = sample.reshape(batch_size, channels * np.prod(remaining_dims))
+
+        abs_sample = sample.abs()  # "a certain percentile absolute pixel value"
+
+        s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1)
+        s = torch.clamp(
+            s, min=1, max=self.config.sample_max_value
+        )  # When clamped to min=1, equivalent to standard clipping to [-1, 1]
+        s = s.unsqueeze(1)  # (batch_size, 1) because clamp will broadcast along dim=0
+        sample = (
+            torch.clamp(sample, -s, s) / s
+        )  # "we threshold xt0 to the range [-s, s] and then divide by s"
+
+        sample = sample.reshape(batch_size, channels, *remaining_dims)
+        sample = sample.to(dtype)
+
+        return sample
+
+    # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler._sigma_to_t
+    def _sigma_to_t(self, sigma):
+        return sigma * self.config.num_train_timesteps
+
+    def _sigma_to_alpha_sigma_t(self, sigma):
+        return 1 - sigma, sigma
+
+    # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.set_timesteps
+    def time_shift(self, mu: float, sigma: float, t: torch.Tensor):
+        return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.convert_model_output
+    def convert_model_output(
+        self,
+        model_output: torch.Tensor,
+        *args,
+        sample: torch.Tensor = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        Convert the model output to the corresponding type the DPMSolver/DPMSolver++ algorithm needs. DPM-Solver is
+        designed to discretize an integral of the noise prediction model, and DPM-Solver++ is designed to discretize an
+        integral of the data prediction model.
+        <Tip>
+        The algorithm and model type are decoupled. You can use either DPMSolver or DPMSolver++ for both noise
+        prediction and data prediction models.
+        </Tip>
+        Args:
+            model_output (`torch.Tensor`):
+                The direct output from the learned diffusion model.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+        Returns:
+            `torch.Tensor`:
+                The converted model output.
+        """
+        timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None)
+        if sample is None:
+            if len(args) > 1:
+                sample = args[1]
+            else:
+                raise ValueError("missing `sample` as a required keyward argument")
+        if timestep is not None:
+            deprecate(
+                "timesteps",
+                "1.0.0",
+                "Passing `timesteps` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        # DPM-Solver++ needs to solve an integral of the data prediction model.
+        if self.config.algorithm_type in ["dpmsolver++", "sde-dpmsolver++"]:
+            if self.config.prediction_type == "flow_prediction":
+                sigma_t = self.sigmas[self.step_index]
+                x0_pred = sample - sigma_t * model_output
+            else:
+                raise ValueError(
+                    f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`,"
+                    " `v_prediction`, or `flow_prediction` for the FlowDPMSolverMultistepScheduler."
+                )
+
+            if self.config.thresholding:
+                x0_pred = self._threshold_sample(x0_pred)
+
+            return x0_pred
+
+        # DPM-Solver needs to solve an integral of the noise prediction model.
+        elif self.config.algorithm_type in ["dpmsolver", "sde-dpmsolver"]:
+            if self.config.prediction_type == "flow_prediction":
+                sigma_t = self.sigmas[self.step_index]
+                epsilon = sample - (1 - sigma_t) * model_output
+            else:
+                raise ValueError(
+                    f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`,"
+                    " `v_prediction` or `flow_prediction` for the FlowDPMSolverMultistepScheduler."
+                )
+
+            if self.config.thresholding:
+                sigma_t = self.sigmas[self.step_index]
+                x0_pred = sample - sigma_t * model_output
+                x0_pred = self._threshold_sample(x0_pred)
+                epsilon = model_output + x0_pred
+
+            return epsilon
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.dpm_solver_first_order_update
+    def dpm_solver_first_order_update(
+        self,
+        model_output: torch.Tensor,
+        *args,
+        sample: torch.Tensor = None,
+        noise: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        One step for the first-order DPMSolver (equivalent to DDIM).
+        Args:
+            model_output (`torch.Tensor`):
+                The direct output from the learned diffusion model.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+        Returns:
+            `torch.Tensor`:
+                The sample tensor at the previous timestep.
+        """
+        timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None)
+        prev_timestep = args[1] if len(args) > 1 else kwargs.pop("prev_timestep", None)
+        if sample is None:
+            if len(args) > 2:
+                sample = args[2]
+            else:
+                raise ValueError(" missing `sample` as a required keyward argument")
+        if timestep is not None:
+            deprecate(
+                "timesteps",
+                "1.0.0",
+                "Passing `timesteps` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        if prev_timestep is not None:
+            deprecate(
+                "prev_timestep",
+                "1.0.0",
+                "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        sigma_t, sigma_s = self.sigmas[self.step_index + 1], self.sigmas[self.step_index]  # pyright: ignore
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+        alpha_s, sigma_s = self._sigma_to_alpha_sigma_t(sigma_s)
+        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+        lambda_s = torch.log(alpha_s) - torch.log(sigma_s)
+
+        h = lambda_t - lambda_s
+        if self.config.algorithm_type == "dpmsolver++":
+            x_t = (sigma_t / sigma_s) * sample - (alpha_t * (torch.exp(-h) - 1.0)) * model_output
+        elif self.config.algorithm_type == "dpmsolver":
+            x_t = (alpha_t / alpha_s) * sample - (sigma_t * (torch.exp(h) - 1.0)) * model_output
+        elif self.config.algorithm_type == "sde-dpmsolver++":
+            assert noise is not None
+            x_t = (
+                (sigma_t / sigma_s * torch.exp(-h)) * sample
+                + (alpha_t * (1 - torch.exp(-2.0 * h))) * model_output
+                + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise
+            )
+        elif self.config.algorithm_type == "sde-dpmsolver":
+            assert noise is not None
+            x_t = (
+                (alpha_t / alpha_s) * sample
+                - 2.0 * (sigma_t * (torch.exp(h) - 1.0)) * model_output
+                + sigma_t * torch.sqrt(torch.exp(2 * h) - 1.0) * noise
+            )
+        return x_t  # pyright: ignore
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.multistep_dpm_solver_second_order_update
+    def multistep_dpm_solver_second_order_update(
+        self,
+        model_output_list: List[torch.Tensor],
+        *args,
+        sample: torch.Tensor = None,
+        noise: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        One step for the second-order multistep DPMSolver.
+        Args:
+            model_output_list (`List[torch.Tensor]`):
+                The direct outputs from learned diffusion model at current and latter timesteps.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+        Returns:
+            `torch.Tensor`:
+                The sample tensor at the previous timestep.
+        """
+        timestep_list = args[0] if len(args) > 0 else kwargs.pop("timestep_list", None)
+        prev_timestep = args[1] if len(args) > 1 else kwargs.pop("prev_timestep", None)
+        if sample is None:
+            if len(args) > 2:
+                sample = args[2]
+            else:
+                raise ValueError(" missing `sample` as a required keyward argument")
+        if timestep_list is not None:
+            deprecate(
+                "timestep_list",
+                "1.0.0",
+                "Passing `timestep_list` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        if prev_timestep is not None:
+            deprecate(
+                "prev_timestep",
+                "1.0.0",
+                "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        sigma_t, sigma_s0, sigma_s1 = (
+            self.sigmas[self.step_index + 1],  # pyright: ignore
+            self.sigmas[self.step_index],
+            self.sigmas[self.step_index - 1],  # pyright: ignore
+        )
+
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+        alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
+        alpha_s1, sigma_s1 = self._sigma_to_alpha_sigma_t(sigma_s1)
+
+        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+        lambda_s1 = torch.log(alpha_s1) - torch.log(sigma_s1)
+
+        m0, m1 = model_output_list[-1], model_output_list[-2]
+
+        h, h_0 = lambda_t - lambda_s0, lambda_s0 - lambda_s1
+        r0 = h_0 / h
+        D0, D1 = m0, (1.0 / r0) * (m0 - m1)
+        if self.config.algorithm_type == "dpmsolver++":
+            # See https://arxiv.org/abs/2211.01095 for detailed derivations
+            if self.config.solver_type == "midpoint":
+                x_t = (
+                    (sigma_t / sigma_s0) * sample
+                    - (alpha_t * (torch.exp(-h) - 1.0)) * D0
+                    - 0.5 * (alpha_t * (torch.exp(-h) - 1.0)) * D1
+                )
+            elif self.config.solver_type == "heun":
+                x_t = (
+                    (sigma_t / sigma_s0) * sample
+                    - (alpha_t * (torch.exp(-h) - 1.0)) * D0
+                    + (alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0)) * D1
+                )
+        elif self.config.algorithm_type == "dpmsolver":
+            # See https://arxiv.org/abs/2206.00927 for detailed derivations
+            if self.config.solver_type == "midpoint":
+                x_t = (
+                    (alpha_t / alpha_s0) * sample
+                    - (sigma_t * (torch.exp(h) - 1.0)) * D0
+                    - 0.5 * (sigma_t * (torch.exp(h) - 1.0)) * D1
+                )
+            elif self.config.solver_type == "heun":
+                x_t = (
+                    (alpha_t / alpha_s0) * sample
+                    - (sigma_t * (torch.exp(h) - 1.0)) * D0
+                    - (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1
+                )
+        elif self.config.algorithm_type == "sde-dpmsolver++":
+            assert noise is not None
+            if self.config.solver_type == "midpoint":
+                x_t = (
+                    (sigma_t / sigma_s0 * torch.exp(-h)) * sample
+                    + (alpha_t * (1 - torch.exp(-2.0 * h))) * D0
+                    + 0.5 * (alpha_t * (1 - torch.exp(-2.0 * h))) * D1
+                    + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise
+                )
+            elif self.config.solver_type == "heun":
+                x_t = (
+                    (sigma_t / sigma_s0 * torch.exp(-h)) * sample
+                    + (alpha_t * (1 - torch.exp(-2.0 * h))) * D0
+                    + (alpha_t * ((1.0 - torch.exp(-2.0 * h)) / (-2.0 * h) + 1.0)) * D1
+                    + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise
+                )
+        elif self.config.algorithm_type == "sde-dpmsolver":
+            assert noise is not None
+            if self.config.solver_type == "midpoint":
+                x_t = (
+                    (alpha_t / alpha_s0) * sample
+                    - 2.0 * (sigma_t * (torch.exp(h) - 1.0)) * D0
+                    - (sigma_t * (torch.exp(h) - 1.0)) * D1
+                    + sigma_t * torch.sqrt(torch.exp(2 * h) - 1.0) * noise
+                )
+            elif self.config.solver_type == "heun":
+                x_t = (
+                    (alpha_t / alpha_s0) * sample
+                    - 2.0 * (sigma_t * (torch.exp(h) - 1.0)) * D0
+                    - 2.0 * (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1
+                    + sigma_t * torch.sqrt(torch.exp(2 * h) - 1.0) * noise
+                )
+        return x_t  # pyright: ignore
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.multistep_dpm_solver_third_order_update
+    def multistep_dpm_solver_third_order_update(
+        self,
+        model_output_list: List[torch.Tensor],
+        *args,
+        sample: torch.Tensor = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        One step for the third-order multistep DPMSolver.
+        Args:
+            model_output_list (`List[torch.Tensor]`):
+                The direct outputs from learned diffusion model at current and latter timesteps.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by diffusion process.
+        Returns:
+            `torch.Tensor`:
+                The sample tensor at the previous timestep.
+        """
+
+        timestep_list = args[0] if len(args) > 0 else kwargs.pop("timestep_list", None)
+        prev_timestep = args[1] if len(args) > 1 else kwargs.pop("prev_timestep", None)
+        if sample is None:
+            if len(args) > 2:
+                sample = args[2]
+            else:
+                raise ValueError(" missing`sample` as a required keyward argument")
+        if timestep_list is not None:
+            deprecate(
+                "timestep_list",
+                "1.0.0",
+                "Passing `timestep_list` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        if prev_timestep is not None:
+            deprecate(
+                "prev_timestep",
+                "1.0.0",
+                "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        sigma_t, sigma_s0, sigma_s1, sigma_s2 = (
+            self.sigmas[self.step_index + 1],  # pyright: ignore
+            self.sigmas[self.step_index],
+            self.sigmas[self.step_index - 1],  # pyright: ignore
+            self.sigmas[self.step_index - 2],  # pyright: ignore
+        )
+
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+        alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
+        alpha_s1, sigma_s1 = self._sigma_to_alpha_sigma_t(sigma_s1)
+        alpha_s2, sigma_s2 = self._sigma_to_alpha_sigma_t(sigma_s2)
+
+        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+        lambda_s1 = torch.log(alpha_s1) - torch.log(sigma_s1)
+        lambda_s2 = torch.log(alpha_s2) - torch.log(sigma_s2)
+
+        m0, m1, m2 = model_output_list[-1], model_output_list[-2], model_output_list[-3]
+
+        h, h_0, h_1 = lambda_t - lambda_s0, lambda_s0 - lambda_s1, lambda_s1 - lambda_s2
+        r0, r1 = h_0 / h, h_1 / h
+        D0 = m0
+        D1_0, D1_1 = (1.0 / r0) * (m0 - m1), (1.0 / r1) * (m1 - m2)
+        D1 = D1_0 + (r0 / (r0 + r1)) * (D1_0 - D1_1)
+        D2 = (1.0 / (r0 + r1)) * (D1_0 - D1_1)
+        if self.config.algorithm_type == "dpmsolver++":
+            # See https://arxiv.org/abs/2206.00927 for detailed derivations
+            x_t = (
+                (sigma_t / sigma_s0) * sample
+                - (alpha_t * (torch.exp(-h) - 1.0)) * D0
+                + (alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0)) * D1
+                - (alpha_t * ((torch.exp(-h) - 1.0 + h) / h**2 - 0.5)) * D2
+            )
+        elif self.config.algorithm_type == "dpmsolver":
+            # See https://arxiv.org/abs/2206.00927 for detailed derivations
+            x_t = (
+                (alpha_t / alpha_s0) * sample
+                - (sigma_t * (torch.exp(h) - 1.0)) * D0
+                - (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1
+                - (sigma_t * ((torch.exp(h) - 1.0 - h) / h**2 - 0.5)) * D2
+            )
+        return x_t  # pyright: ignore
+
+    def index_for_timestep(self, timestep, schedule_timesteps=None):
+        if schedule_timesteps is None:
+            schedule_timesteps = self.timesteps
+
+        indices = (schedule_timesteps == timestep).nonzero()
+
+        # The sigma index that is taken for the **very** first `step`
+        # is always the second index (or the last index if there is only 1)
+        # This way we can ensure we don't accidentally skip a sigma in
+        # case we start in the middle of the denoising schedule (e.g. for image-to-image)
+        pos = 1 if len(indices) > 1 else 0
+
+        return indices[pos].item()
+
+    def _init_step_index(self, timestep):
+        """
+        Initialize the step_index counter for the scheduler.
+        """
+
+        if self.begin_index is None:
+            if isinstance(timestep, torch.Tensor):
+                timestep = timestep.to(self.timesteps.device)
+            self._step_index = self.index_for_timestep(timestep)
+        else:
+            self._step_index = self._begin_index
+
+    # Modified from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.step
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: Union[int, torch.Tensor],
+        sample: torch.Tensor,
+        generator=None,
+        variance_noise: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+    ) -> Union[SchedulerOutput, Tuple]:
+        """
+        Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
+        the multistep DPMSolver.
+        Args:
+            model_output (`torch.Tensor`):
+                The direct output from learned diffusion model.
+            timestep (`int`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+            generator (`torch.Generator`, *optional*):
+                A random number generator.
+            variance_noise (`torch.Tensor`):
+                Alternative to generating noise with `generator` by directly providing the noise for the variance
+                itself. Useful for methods such as [`LEdits++`].
+            return_dict (`bool`):
+                Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`.
+        Returns:
+            [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`:
+                If return_dict is `True`, [`~schedulers.scheduling_utils.SchedulerOutput`] is returned, otherwise a
+                tuple is returned where the first element is the sample tensor.
+        """
+        if self.num_inference_steps is None:
+            raise ValueError(
+                "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
+            )
+
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        # Improve numerical stability for small number of steps
+        lower_order_final = (self.step_index == len(self.timesteps) - 1) and (
+            self.config.euler_at_final
+            or (self.config.lower_order_final and len(self.timesteps) < 15)
+            or self.config.final_sigmas_type == "zero"
+        )
+        lower_order_second = (
+            (self.step_index == len(self.timesteps) - 2)
+            and self.config.lower_order_final
+            and len(self.timesteps) < 15
+        )
+
+        model_output = self.convert_model_output(model_output, sample=sample)
+        for i in range(self.config.solver_order - 1):
+            self.model_outputs[i] = self.model_outputs[i + 1]
+        self.model_outputs[-1] = model_output
+
+        # Upcast to avoid precision issues when computing prev_sample
+        sample = sample.to(torch.float32)
+        if self.config.algorithm_type in ["sde-dpmsolver", "sde-dpmsolver++"] and variance_noise is None:
+            noise = randn_tensor(
+                model_output.shape, generator=generator, device=model_output.device, dtype=torch.float32
+            )
+        elif self.config.algorithm_type in ["sde-dpmsolver", "sde-dpmsolver++"]:
+            noise = variance_noise.to(device=model_output.device, dtype=torch.float32)  # pyright: ignore
+        else:
+            noise = None
+
+        if self.config.solver_order == 1 or self.lower_order_nums < 1 or lower_order_final:
+            prev_sample = self.dpm_solver_first_order_update(model_output, sample=sample, noise=noise)
+        elif self.config.solver_order == 2 or self.lower_order_nums < 2 or lower_order_second:
+            prev_sample = self.multistep_dpm_solver_second_order_update(
+                self.model_outputs, sample=sample, noise=noise
+            )
+        else:
+            prev_sample = self.multistep_dpm_solver_third_order_update(self.model_outputs, sample=sample)
+
+        if self.lower_order_nums < self.config.solver_order:
+            self.lower_order_nums += 1
+
+        # Cast sample back to expected dtype
+        prev_sample = prev_sample.to(model_output.dtype)
+
+        # upon completion increase step index by one
+        self._step_index += 1  # pyright: ignore
+
+        if not return_dict:
+            return (prev_sample,)
+
+        return SchedulerOutput(prev_sample=prev_sample)
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.scale_model_input
+    def scale_model_input(self, sample: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+        """
+        Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
+        current timestep.
+        Args:
+            sample (`torch.Tensor`):
+                The input sample.
+        Returns:
+            `torch.Tensor`:
+                A scaled input sample.
+        """
+        return sample
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.scale_model_input
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.IntTensor,
+    ) -> torch.Tensor:
+        # Make sure sigmas and timesteps have the same device and dtype as original_samples
+        sigmas = self.sigmas.to(device=original_samples.device, dtype=original_samples.dtype)
+        if original_samples.device.type == "mps" and torch.is_floating_point(timesteps):
+            # mps does not support float64
+            schedule_timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32)
+            timesteps = timesteps.to(original_samples.device, dtype=torch.float32)
+        else:
+            schedule_timesteps = self.timesteps.to(original_samples.device)
+            timesteps = timesteps.to(original_samples.device)
+
+        # begin_index is None when the scheduler is used for training or pipeline does not implement set_begin_index
+        if self.begin_index is None:
+            step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timesteps]
+        elif self.step_index is not None:
+            # add_noise is called after first denoising step (for inpainting)
+            step_indices = [self.step_index] * timesteps.shape[0]
+        else:
+            # add noise is called before first denoising step to create initial latent(img2img)
+            step_indices = [self.begin_index] * timesteps.shape[0]
+
+        sigma = sigmas[step_indices].flatten()
+        while len(sigma.shape) < len(original_samples.shape):
+            sigma = sigma.unsqueeze(-1)
+
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+        noisy_samples = alpha_t * original_samples + sigma_t * noise
+        return noisy_samples
+
+    def __len__(self):
+        return self.config.num_train_timesteps

src/lerobot/policies/fastwam/wan/utils/fm_solvers_unipc.py

@@ -0,0 +1,765 @@
+# Copied from https://github.com/huggingface/diffusers/blob/v0.31.0/src/diffusers/schedulers/scheduling_unipc_multistep.py
+# Convert unipc for flow matching
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+
+import math
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import (
+    KarrasDiffusionSchedulers,
+    SchedulerMixin,
+    SchedulerOutput,
+)
+from diffusers.utils import deprecate, is_scipy_available
+
+if is_scipy_available():
+    import scipy.stats
+
+
+class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
+    """
+    `UniPCMultistepScheduler` is a training-free framework designed for the fast sampling of diffusion models.
+
+    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
+    methods the library implements for all schedulers such as loading and saving.
+
+    Args:
+        num_train_timesteps (`int`, defaults to 1000):
+            The number of diffusion steps to train the model.
+        solver_order (`int`, default `2`):
+            The UniPC order which can be any positive integer. The effective order of accuracy is `solver_order + 1`
+            due to the UniC. It is recommended to use `solver_order=2` for guided sampling, and `solver_order=3` for
+            unconditional sampling.
+        prediction_type (`str`, defaults to "flow_prediction"):
+            Prediction type of the scheduler function; must be `flow_prediction` for this scheduler, which predicts
+            the flow of the diffusion process.
+        thresholding (`bool`, defaults to `False`):
+            Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such
+            as Stable Diffusion.
+        dynamic_thresholding_ratio (`float`, defaults to 0.995):
+            The ratio for the dynamic thresholding method. Valid only when `thresholding=True`.
+        sample_max_value (`float`, defaults to 1.0):
+            The threshold value for dynamic thresholding. Valid only when `thresholding=True` and `predict_x0=True`.
+        predict_x0 (`bool`, defaults to `True`):
+            Whether to use the updating algorithm on the predicted x0.
+        solver_type (`str`, default `bh2`):
+            Solver type for UniPC. It is recommended to use `bh1` for unconditional sampling when steps < 10, and `bh2`
+            otherwise.
+        lower_order_final (`bool`, default `True`):
+            Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. This can
+            stabilize the sampling of DPMSolver for steps < 15, especially for steps <= 10.
+        disable_corrector (`list`, default `[]`):
+            Decides which step to disable the corrector to mitigate the misalignment between `epsilon_theta(x_t, c)`
+            and `epsilon_theta(x_t^c, c)` which can influence convergence for a large guidance scale. Corrector is
+            usually disabled during the first few steps.
+        solver_p (`SchedulerMixin`, default `None`):
+            Any other scheduler that if specified, the algorithm becomes `solver_p + UniC`.
+        use_karras_sigmas (`bool`, *optional*, defaults to `False`):
+            Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`,
+            the sigmas are determined according to a sequence of noise levels {σi}.
+        use_exponential_sigmas (`bool`, *optional*, defaults to `False`):
+            Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process.
+        timestep_spacing (`str`, defaults to `"linspace"`):
+            The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
+            Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
+        steps_offset (`int`, defaults to 0):
+            An offset added to the inference steps, as required by some model families.
+        final_sigmas_type (`str`, defaults to `"zero"`):
+            The final `sigma` value for the noise schedule during the sampling process. If `"sigma_min"`, the final
+            sigma is the same as the last sigma in the training schedule. If `zero`, the final sigma is set to 0.
+    """
+
+    _compatibles = [e.name for e in KarrasDiffusionSchedulers]
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        solver_order: int = 2,
+        prediction_type: str = "flow_prediction",
+        shift: Optional[float] = 1.0,
+        use_dynamic_shifting=False,
+        thresholding: bool = False,
+        dynamic_thresholding_ratio: float = 0.995,
+        sample_max_value: float = 1.0,
+        predict_x0: bool = True,
+        solver_type: str = "bh2",
+        lower_order_final: bool = True,
+        disable_corrector: List[int] = [],
+        solver_p: SchedulerMixin = None,
+        timestep_spacing: str = "linspace",
+        steps_offset: int = 0,
+        final_sigmas_type: Optional[str] = "zero",  # "zero", "sigma_min"
+    ):
+
+        if solver_type not in ["bh1", "bh2"]:
+            if solver_type in ["midpoint", "heun", "logrho"]:
+                self.register_to_config(solver_type="bh2")
+            else:
+                raise NotImplementedError(f"{solver_type} is not implemented for {self.__class__}")
+
+        self.predict_x0 = predict_x0
+        # setable values
+        self.num_inference_steps = None
+        alphas = np.linspace(1, 1 / num_train_timesteps, num_train_timesteps)[::-1].copy()
+        sigmas = 1.0 - alphas
+        sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32)
+
+        if not use_dynamic_shifting:
+            # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution
+            sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)  # pyright: ignore
+
+        self.sigmas = sigmas
+        self.timesteps = sigmas * num_train_timesteps
+
+        self.model_outputs = [None] * solver_order
+        self.timestep_list = [None] * solver_order
+        self.lower_order_nums = 0
+        self.disable_corrector = disable_corrector
+        self.solver_p = solver_p
+        self.last_sample = None
+        self._step_index = None
+        self._begin_index = None
+
+        self.sigmas = self.sigmas.to("cpu")  # to avoid too much CPU/GPU communication
+        self.sigma_min = self.sigmas[-1].item()
+        self.sigma_max = self.sigmas[0].item()
+
+    @property
+    def step_index(self):
+        """
+        The index counter for current timestep. It will increase 1 after each scheduler step.
+        """
+        return self._step_index
+
+    @property
+    def begin_index(self):
+        """
+        The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
+        """
+        return self._begin_index
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
+    def set_begin_index(self, begin_index: int = 0):
+        """
+        Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
+
+        Args:
+            begin_index (`int`):
+                The begin index for the scheduler.
+        """
+        self._begin_index = begin_index
+
+    # Modified from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.set_timesteps
+    def set_timesteps(
+        self,
+        num_inference_steps: Union[int, None] = None,
+        device: Union[str, torch.device] = None,
+        sigmas: Optional[List[float]] = None,
+        mu: Optional[Union[float, None]] = None,
+        shift: Optional[Union[float, None]] = None,
+    ):
+        """
+        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
+        Args:
+            num_inference_steps (`int`):
+                Total number of the spacing of the time steps.
+            device (`str` or `torch.device`, *optional*):
+                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        """
+
+        if self.config.use_dynamic_shifting and mu is None:
+            raise ValueError(
+                " you have to pass a value for `mu` when `use_dynamic_shifting` is set to be `True`"
+            )
+
+        if sigmas is None:
+            sigmas = np.linspace(self.sigma_max, self.sigma_min, num_inference_steps + 1).copy()[:-1]  # pyright: ignore
+
+        if self.config.use_dynamic_shifting:
+            sigmas = self.time_shift(mu, 1.0, sigmas)  # pyright: ignore
+        else:
+            if shift is None:
+                shift = self.config.shift
+            sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)  # pyright: ignore
+
+        if self.config.final_sigmas_type == "sigma_min":
+            sigma_last = ((1 - self.alphas_cumprod[0]) / self.alphas_cumprod[0]) ** 0.5
+        elif self.config.final_sigmas_type == "zero":
+            sigma_last = 0
+        else:
+            raise ValueError(
+                f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}"
+            )
+
+        timesteps = sigmas * self.config.num_train_timesteps
+        sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32)  # pyright: ignore
+
+        self.sigmas = torch.from_numpy(sigmas)
+        self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.int64)
+
+        self.num_inference_steps = len(timesteps)
+
+        self.model_outputs = [
+            None,
+        ] * self.config.solver_order
+        self.lower_order_nums = 0
+        self.last_sample = None
+        if self.solver_p:
+            self.solver_p.set_timesteps(self.num_inference_steps, device=device)
+
+        # add an index counter for schedulers that allow duplicated timesteps
+        self._step_index = None
+        self._begin_index = None
+        self.sigmas = self.sigmas.to("cpu")  # to avoid too much CPU/GPU communication
+
+    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample
+    def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
+        """
+        "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the
+        prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by
+        s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing
+        pixels from saturation at each step. We find that dynamic thresholding results in significantly better
+        photorealism as well as better image-text alignment, especially when using very large guidance weights."
+
+        https://arxiv.org/abs/2205.11487
+        """
+        dtype = sample.dtype
+        batch_size, channels, *remaining_dims = sample.shape
+
+        if dtype not in (torch.float32, torch.float64):
+            sample = sample.float()  # upcast for quantile calculation, and clamp not implemented for cpu half
+
+        # Flatten sample for doing quantile calculation along each image
+        sample = sample.reshape(batch_size, channels * np.prod(remaining_dims))
+
+        abs_sample = sample.abs()  # "a certain percentile absolute pixel value"
+
+        s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1)
+        s = torch.clamp(
+            s, min=1, max=self.config.sample_max_value
+        )  # When clamped to min=1, equivalent to standard clipping to [-1, 1]
+        s = s.unsqueeze(1)  # (batch_size, 1) because clamp will broadcast along dim=0
+        sample = (
+            torch.clamp(sample, -s, s) / s
+        )  # "we threshold xt0 to the range [-s, s] and then divide by s"
+
+        sample = sample.reshape(batch_size, channels, *remaining_dims)
+        sample = sample.to(dtype)
+
+        return sample
+
+    # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler._sigma_to_t
+    def _sigma_to_t(self, sigma):
+        return sigma * self.config.num_train_timesteps
+
+    def _sigma_to_alpha_sigma_t(self, sigma):
+        return 1 - sigma, sigma
+
+    # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.set_timesteps
+    def time_shift(self, mu: float, sigma: float, t: torch.Tensor):
+        return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
+
+    def convert_model_output(
+        self,
+        model_output: torch.Tensor,
+        *args,
+        sample: torch.Tensor = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        r"""
+        Convert the model output to the corresponding type the UniPC algorithm needs.
+
+        Args:
+            model_output (`torch.Tensor`):
+                The direct output from the learned diffusion model.
+            timestep (`int`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+
+        Returns:
+            `torch.Tensor`:
+                The converted model output.
+        """
+        timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None)
+        if sample is None:
+            if len(args) > 1:
+                sample = args[1]
+            else:
+                raise ValueError("missing `sample` as a required keyward argument")
+        if timestep is not None:
+            deprecate(
+                "timesteps",
+                "1.0.0",
+                "Passing `timesteps` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        sigma = self.sigmas[self.step_index]
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+
+        if self.predict_x0:
+            if self.config.prediction_type == "flow_prediction":
+                sigma_t = self.sigmas[self.step_index]
+                x0_pred = sample - sigma_t * model_output
+            else:
+                raise ValueError(
+                    f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`,"
+                    " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler."
+                )
+
+            if self.config.thresholding:
+                x0_pred = self._threshold_sample(x0_pred)
+
+            return x0_pred
+        else:
+            if self.config.prediction_type == "flow_prediction":
+                sigma_t = self.sigmas[self.step_index]
+                epsilon = sample - (1 - sigma_t) * model_output
+            else:
+                raise ValueError(
+                    f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`,"
+                    " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler."
+                )
+
+            if self.config.thresholding:
+                sigma_t = self.sigmas[self.step_index]
+                x0_pred = sample - sigma_t * model_output
+                x0_pred = self._threshold_sample(x0_pred)
+                epsilon = model_output + x0_pred
+
+            return epsilon
+
+    def multistep_uni_p_bh_update(
+        self,
+        model_output: torch.Tensor,
+        *args,
+        sample: torch.Tensor = None,
+        order: int = None,  # pyright: ignore
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        One step for the UniP (B(h) version). Alternatively, `self.solver_p` is used if is specified.
+
+        Args:
+            model_output (`torch.Tensor`):
+                The direct output from the learned diffusion model at the current timestep.
+            prev_timestep (`int`):
+                The previous discrete timestep in the diffusion chain.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+            order (`int`):
+                The order of UniP at this timestep (corresponds to the *p* in UniPC-p).
+
+        Returns:
+            `torch.Tensor`:
+                The sample tensor at the previous timestep.
+        """
+        prev_timestep = args[0] if len(args) > 0 else kwargs.pop("prev_timestep", None)
+        if sample is None:
+            if len(args) > 1:
+                sample = args[1]
+            else:
+                raise ValueError(" missing `sample` as a required keyward argument")
+        if order is None:
+            if len(args) > 2:
+                order = args[2]
+            else:
+                raise ValueError(" missing `order` as a required keyward argument")
+        if prev_timestep is not None:
+            deprecate(
+                "prev_timestep",
+                "1.0.0",
+                "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+        model_output_list = self.model_outputs
+
+        s0 = self.timestep_list[-1]
+        m0 = model_output_list[-1]
+        x = sample
+
+        if self.solver_p:
+            x_t = self.solver_p.step(model_output, s0, x).prev_sample
+            return x_t
+
+        sigma_t, sigma_s0 = self.sigmas[self.step_index + 1], self.sigmas[self.step_index]  # pyright: ignore
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+        alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
+
+        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+
+        h = lambda_t - lambda_s0
+        device = sample.device
+
+        rks = []
+        D1s = []
+        for i in range(1, order):
+            si = self.step_index - i  # pyright: ignore
+            mi = model_output_list[-(i + 1)]
+            alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
+            lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
+            rk = (lambda_si - lambda_s0) / h
+            rks.append(rk)
+            D1s.append((mi - m0) / rk)  # pyright: ignore
+
+        rks.append(1.0)
+        rks = torch.tensor(rks, device=device)
+
+        R = []
+        b = []
+
+        hh = -h if self.predict_x0 else h
+        h_phi_1 = torch.expm1(hh)  # h\phi_1(h) = e^h - 1
+        h_phi_k = h_phi_1 / hh - 1
+
+        factorial_i = 1
+
+        if self.config.solver_type == "bh1":
+            B_h = hh
+        elif self.config.solver_type == "bh2":
+            B_h = torch.expm1(hh)
+        else:
+            raise NotImplementedError()
+
+        for i in range(1, order + 1):
+            R.append(torch.pow(rks, i - 1))
+            b.append(h_phi_k * factorial_i / B_h)
+            factorial_i *= i + 1
+            h_phi_k = h_phi_k / hh - 1 / factorial_i
+
+        R = torch.stack(R)
+        b = torch.tensor(b, device=device)
+
+        if len(D1s) > 0:
+            D1s = torch.stack(D1s, dim=1)  # (B, K)
+            # for order 2, we use a simplified version
+            if order == 2:
+                rhos_p = torch.tensor([0.5], dtype=x.dtype, device=device)
+            else:
+                rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1]).to(device).to(x.dtype)
+        else:
+            D1s = None
+
+        if self.predict_x0:
+            x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
+            if D1s is not None:
+                pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s)  # pyright: ignore
+            else:
+                pred_res = 0
+            x_t = x_t_ - alpha_t * B_h * pred_res
+        else:
+            x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
+            if D1s is not None:
+                pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s)  # pyright: ignore
+            else:
+                pred_res = 0
+            x_t = x_t_ - sigma_t * B_h * pred_res
+
+        x_t = x_t.to(x.dtype)
+        return x_t
+
+    def multistep_uni_c_bh_update(
+        self,
+        this_model_output: torch.Tensor,
+        *args,
+        last_sample: torch.Tensor = None,
+        this_sample: torch.Tensor = None,
+        order: int = None,  # pyright: ignore
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        One step for the UniC (B(h) version).
+
+        Args:
+            this_model_output (`torch.Tensor`):
+                The model outputs at `x_t`.
+            this_timestep (`int`):
+                The current timestep `t`.
+            last_sample (`torch.Tensor`):
+                The generated sample before the last predictor `x_{t-1}`.
+            this_sample (`torch.Tensor`):
+                The generated sample after the last predictor `x_{t}`.
+            order (`int`):
+                The `p` of UniC-p at this step. The effective order of accuracy should be `order + 1`.
+
+        Returns:
+            `torch.Tensor`:
+                The corrected sample tensor at the current timestep.
+        """
+        this_timestep = args[0] if len(args) > 0 else kwargs.pop("this_timestep", None)
+        if last_sample is None:
+            if len(args) > 1:
+                last_sample = args[1]
+            else:
+                raise ValueError(" missing`last_sample` as a required keyward argument")
+        if this_sample is None:
+            if len(args) > 2:
+                this_sample = args[2]
+            else:
+                raise ValueError(" missing`this_sample` as a required keyward argument")
+        if order is None:
+            if len(args) > 3:
+                order = args[3]
+            else:
+                raise ValueError(" missing`order` as a required keyward argument")
+        if this_timestep is not None:
+            deprecate(
+                "this_timestep",
+                "1.0.0",
+                "Passing `this_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        model_output_list = self.model_outputs
+
+        m0 = model_output_list[-1]
+        x = last_sample
+        x_t = this_sample
+        model_t = this_model_output
+
+        sigma_t, sigma_s0 = self.sigmas[self.step_index], self.sigmas[self.step_index - 1]  # pyright: ignore
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+        alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
+
+        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+
+        h = lambda_t - lambda_s0
+        device = this_sample.device
+
+        rks = []
+        D1s = []
+        for i in range(1, order):
+            si = self.step_index - (i + 1)  # pyright: ignore
+            mi = model_output_list[-(i + 1)]
+            alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
+            lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
+            rk = (lambda_si - lambda_s0) / h
+            rks.append(rk)
+            D1s.append((mi - m0) / rk)  # pyright: ignore
+
+        rks.append(1.0)
+        rks = torch.tensor(rks, device=device)
+
+        R = []
+        b = []
+
+        hh = -h if self.predict_x0 else h
+        h_phi_1 = torch.expm1(hh)  # h\phi_1(h) = e^h - 1
+        h_phi_k = h_phi_1 / hh - 1
+
+        factorial_i = 1
+
+        if self.config.solver_type == "bh1":
+            B_h = hh
+        elif self.config.solver_type == "bh2":
+            B_h = torch.expm1(hh)
+        else:
+            raise NotImplementedError()
+
+        for i in range(1, order + 1):
+            R.append(torch.pow(rks, i - 1))
+            b.append(h_phi_k * factorial_i / B_h)
+            factorial_i *= i + 1
+            h_phi_k = h_phi_k / hh - 1 / factorial_i
+
+        R = torch.stack(R)
+        b = torch.tensor(b, device=device)
+
+        if len(D1s) > 0:
+            D1s = torch.stack(D1s, dim=1)
+        else:
+            D1s = None
+
+        # for order 1, we use a simplified version
+        if order == 1:
+            rhos_c = torch.tensor([0.5], dtype=x.dtype, device=device)
+        else:
+            rhos_c = torch.linalg.solve(R, b).to(device).to(x.dtype)
+
+        if self.predict_x0:
+            x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
+            if D1s is not None:
+                corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s)
+            else:
+                corr_res = 0
+            D1_t = model_t - m0
+            x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t)
+        else:
+            x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
+            if D1s is not None:
+                corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s)
+            else:
+                corr_res = 0
+            D1_t = model_t - m0
+            x_t = x_t_ - sigma_t * B_h * (corr_res + rhos_c[-1] * D1_t)
+        x_t = x_t.to(x.dtype)
+        return x_t
+
+    def index_for_timestep(self, timestep, schedule_timesteps=None):
+        if schedule_timesteps is None:
+            schedule_timesteps = self.timesteps
+
+        indices = (schedule_timesteps == timestep).nonzero()
+
+        # The sigma index that is taken for the **very** first `step`
+        # is always the second index (or the last index if there is only 1)
+        # This way we can ensure we don't accidentally skip a sigma in
+        # case we start in the middle of the denoising schedule (e.g. for image-to-image)
+        pos = 1 if len(indices) > 1 else 0
+
+        return indices[pos].item()
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index
+    def _init_step_index(self, timestep):
+        """
+        Initialize the step_index counter for the scheduler.
+        """
+
+        if self.begin_index is None:
+            if isinstance(timestep, torch.Tensor):
+                timestep = timestep.to(self.timesteps.device)
+            self._step_index = self.index_for_timestep(timestep)
+        else:
+            self._step_index = self._begin_index
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: Union[int, torch.Tensor],
+        sample: torch.Tensor,
+        return_dict: bool = True,
+        generator=None,
+    ) -> Union[SchedulerOutput, Tuple]:
+        """
+        Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
+        the multistep UniPC.
+
+        Args:
+            model_output (`torch.Tensor`):
+                The direct output from learned diffusion model.
+            timestep (`int`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+            return_dict (`bool`):
+                Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`.
+
+        Returns:
+            [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`:
+                If return_dict is `True`, [`~schedulers.scheduling_utils.SchedulerOutput`] is returned, otherwise a
+                tuple is returned where the first element is the sample tensor.
+
+        """
+        if self.num_inference_steps is None:
+            raise ValueError(
+                "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
+            )
+
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        use_corrector = (
+            self.step_index > 0
+            and self.step_index - 1 not in self.disable_corrector
+            and self.last_sample is not None  # pyright: ignore
+        )
+
+        model_output_convert = self.convert_model_output(model_output, sample=sample)
+        if use_corrector:
+            sample = self.multistep_uni_c_bh_update(
+                this_model_output=model_output_convert,
+                last_sample=self.last_sample,
+                this_sample=sample,
+                order=self.this_order,
+            )
+
+        for i in range(self.config.solver_order - 1):
+            self.model_outputs[i] = self.model_outputs[i + 1]
+            self.timestep_list[i] = self.timestep_list[i + 1]
+
+        self.model_outputs[-1] = model_output_convert
+        self.timestep_list[-1] = timestep  # pyright: ignore
+
+        if self.config.lower_order_final:
+            this_order = min(self.config.solver_order, len(self.timesteps) - self.step_index)  # pyright: ignore
+        else:
+            this_order = self.config.solver_order
+
+        self.this_order = min(this_order, self.lower_order_nums + 1)  # warmup for multistep
+        assert self.this_order > 0
+
+        self.last_sample = sample
+        prev_sample = self.multistep_uni_p_bh_update(
+            model_output=model_output,  # pass the original non-converted model output, in case solver-p is used
+            sample=sample,
+            order=self.this_order,
+        )
+
+        if self.lower_order_nums < self.config.solver_order:
+            self.lower_order_nums += 1
+
+        # upon completion increase step index by one
+        self._step_index += 1  # pyright: ignore
+
+        if not return_dict:
+            return (prev_sample,)
+
+        return SchedulerOutput(prev_sample=prev_sample)
+
+    def scale_model_input(self, sample: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+        """
+        Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
+        current timestep.
+
+        Args:
+            sample (`torch.Tensor`):
+                The input sample.
+
+        Returns:
+            `torch.Tensor`:
+                A scaled input sample.
+        """
+        return sample
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.add_noise
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.IntTensor,
+    ) -> torch.Tensor:
+        # Make sure sigmas and timesteps have the same device and dtype as original_samples
+        sigmas = self.sigmas.to(device=original_samples.device, dtype=original_samples.dtype)
+        if original_samples.device.type == "mps" and torch.is_floating_point(timesteps):
+            # mps does not support float64
+            schedule_timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32)
+            timesteps = timesteps.to(original_samples.device, dtype=torch.float32)
+        else:
+            schedule_timesteps = self.timesteps.to(original_samples.device)
+            timesteps = timesteps.to(original_samples.device)
+
+        # begin_index is None when the scheduler is used for training or pipeline does not implement set_begin_index
+        if self.begin_index is None:
+            step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timesteps]
+        elif self.step_index is not None:
+            # add_noise is called after first denoising step (for inpainting)
+            step_indices = [self.step_index] * timesteps.shape[0]
+        else:
+            # add noise is called before first denoising step to create initial latent(img2img)
+            step_indices = [self.begin_index] * timesteps.shape[0]
+
+        sigma = sigmas[step_indices].flatten()
+        while len(sigma.shape) < len(original_samples.shape):
+            sigma = sigma.unsqueeze(-1)
+
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+        noisy_samples = alpha_t * original_samples + sigma_t * noise
+        return noisy_samples
+
+    def __len__(self):
+        return self.config.num_train_timesteps

src/lerobot/policies/fastwam/wan_adapters.py

@@ -0,0 +1,99 @@
+# 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 pathlib import Path
+from typing import Any
+
+import torch
+
+from .wan.modules.vae2_2 import Wan2_2_VAE
+
+
+class WanVideoVAE38(torch.nn.Module):
+    """Tensor-batch adapter around the official Wan2.2 VAE wrapper."""
+
+    upsampling_factor = 16
+    temporal_downsample_factor = 4
+    z_dim = 48
+
+    def __init__(
+        self,
+        vae_pth: str | Path,
+        dtype: torch.dtype = torch.float32,
+        device: str | torch.device = "cuda",
+    ) -> None:
+        super().__init__()
+        self.wan_vae = Wan2_2_VAE(vae_pth=str(vae_pth), dtype=dtype, device=str(device))
+        self.model = self.wan_vae.model
+        self.dtype = dtype
+        self.device = torch.device(device)
+
+    def to(self, *args: Any, **kwargs: Any):
+        super().to(*args, **kwargs)
+        self.model.to(*args, **kwargs)
+        param = next(self.model.parameters())
+        self.device = param.device
+        self.dtype = param.dtype
+        self.wan_vae.device = self.device
+        self.wan_vae.dtype = self.dtype
+        self.wan_vae.scale = [scale.to(device=self.device, dtype=self.dtype) for scale in self.wan_vae.scale]
+        self.wan_vae.model = self.model
+        return self
+
+    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 tile_size, tile_stride
+        if tiled:
+            raise NotImplementedError("Tiled Wan2.2 VAE encoding is not supported by the FastWAM adapter.")
+        target_device = self.device if device is None else torch.device(device)
+        if target_device != self.device:
+            self.to(device=target_device)
+        if isinstance(videos, torch.Tensor):
+            videos = list(videos)
+        hidden_states = self.wan_vae.encode([video.to(self.device) for video in videos])
+        if hidden_states is None:
+            raise RuntimeError("Wan2.2 VAE encode failed; expected a list of video tensors.")
+        return torch.stack(hidden_states)
+
+    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 tile_size, tile_stride
+        if tiled:
+            raise NotImplementedError("Tiled Wan2.2 VAE decoding is not supported by the FastWAM adapter.")
+        target_device = self.device if device is None else torch.device(device)
+        if target_device != self.device:
+            self.to(device=target_device)
+        if isinstance(hidden_states, torch.Tensor):
+            hidden_states = list(hidden_states)
+        videos = self.wan_vae.decode([hidden_state.to(self.device) for hidden_state in hidden_states])
+        if videos is None:
+            raise RuntimeError("Wan2.2 VAE decode failed; expected a list of latent tensors.")
+        return torch.stack(videos)
+
+
+__all__ = ["WanVideoVAE38"]

src/lerobot/policies/fastwam/wan_components.py

@@ -0,0 +1,266 @@
+# 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 time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import torch
+from safetensors.torch import load_file
+
+if TYPE_CHECKING:
+    from .wan.modules.tokenizers import HuggingfaceTokenizer
+    from .wan_adapters import WanVideoVAE38
+    from .wan_video_dit import WanVideoDiT
+
+logger = logging.getLogger(__name__)
+
+WAN_DIT_PATTERN = "diffusion_pytorch_model*.safetensors"
+WAN_T5_CHECKPOINT = "models_t5_umt5-xxl-enc-bf16.pth"
+WAN_T5_TOKENIZER = "google/umt5-xxl"
+WAN_VAE_CHECKPOINT = "Wan2.2_VAE.pth"
+
+
+@dataclass(frozen=True)
+class WanCheckpointPaths:
+    root: Path
+    dit: list[Path]
+    vae: Path
+    text_encoder: Path | None
+    tokenizer: Path | None
+
+
+@dataclass
+class Wan22LoadedComponents:
+    dit: WanVideoDiT
+    vae: WanVideoVAE38
+    text_encoder: torch.nn.Module | None
+    tokenizer: HuggingfaceTokenizer | None
+    dit_path: list[str]
+    vae_path: str
+    text_encoder_path: str | None
+    tokenizer_path: str | None
+
+
+def resolve_wan_checkpoint_dir(
+    model_id_or_path: str | Path,
+    *,
+    cache_dir: str | Path | None = None,
+    local_files_only: bool = False,
+    revision: str | None = None,
+) -> Path:
+    """Return a local Wan2.2 checkpoint directory.
+
+    Local paths are used directly. Hub repos are downloaded with the same fixed
+    component names used by the upstream Wan2.2 inference code.
+    """
+
+    path = Path(model_id_or_path).expanduser()
+    if path.is_dir():
+        return path
+
+    from huggingface_hub import snapshot_download
+
+    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,
+            WAN_T5_CHECKPOINT,
+            WAN_VAE_CHECKPOINT,
+            f"{WAN_T5_TOKENIZER}/**",
+        ],
+    )
+    return Path(snapshot_path)
+
+
+def resolve_wan_checkpoint_paths(
+    checkpoint_dir: str | Path,
+    *,
+    tokenizer_dir: str | Path | None = None,
+    load_dit: bool = True,
+    load_text_encoder: bool = True,
+) -> WanCheckpointPaths:
+    root = Path(checkpoint_dir).expanduser()
+    tokenizer_root = Path(tokenizer_dir).expanduser() if tokenizer_dir is not None else root
+    dit = sorted(root.glob(WAN_DIT_PATTERN)) if load_dit else []
+    vae = root / WAN_VAE_CHECKPOINT
+    text_encoder = root / WAN_T5_CHECKPOINT if load_text_encoder else None
+    tokenizer = tokenizer_root / WAN_T5_TOKENIZER if load_text_encoder else None
+
+    missing = []
+    if load_dit and len(dit) == 0:
+        missing.append(f"DiT ({WAN_DIT_PATTERN})")
+    if not vae.exists():
+        missing.append(f"VAE ({WAN_VAE_CHECKPOINT})")
+    if load_text_encoder:
+        if text_encoder is None or not text_encoder.exists():
+            missing.append(f"text encoder ({WAN_T5_CHECKPOINT})")
+        if tokenizer is None or not tokenizer.exists():
+            missing.append(f"tokenizer ({WAN_T5_TOKENIZER})")
+    if missing:
+        raise FileNotFoundError(
+            f"Incomplete Wan2.2 checkpoint directory {root}: missing {', '.join(missing)}."
+        )
+
+    return WanCheckpointPaths(
+        root=root,
+        dit=dit,
+        vae=vae,
+        text_encoder=text_encoder,
+        tokenizer=tokenizer,
+    )
+
+
+def load_wan_video_dit(
+    paths: list[str | Path],
+    *,
+    dit_config: dict[str, Any],
+    torch_dtype: torch.dtype,
+    device: str,
+) -> WanVideoDiT:
+    from .wan_video_dit import 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 load_wan_text_encoder(
+    checkpoint_path: str | Path,
+    *,
+    torch_dtype: torch.dtype,
+    device: str,
+) -> torch.nn.Module:
+    from .wan.modules.t5 import umt5_xxl
+
+    model = umt5_xxl(
+        encoder_only=True,
+        return_tokenizer=False,
+        dtype=torch_dtype,
+        device=device,
+    )
+    state_dict = torch.load(checkpoint_path, map_location="cpu")
+    model.load_state_dict(state_dict)
+    return model.to(device=device, dtype=torch_dtype)
+
+
+def load_wan_tokenizer(tokenizer_path: str | Path, *, tokenizer_max_len: int) -> HuggingfaceTokenizer:
+    from .wan.modules.tokenizers import HuggingfaceTokenizer
+
+    return HuggingfaceTokenizer(
+        name=str(tokenizer_path),
+        seq_len=int(tokenizer_max_len),
+        clean="whitespace",
+    )
+
+
+def load_wan_vae(checkpoint_path: str | Path, *, torch_dtype: torch.dtype, device: str) -> WanVideoVAE38:
+    from .wan_adapters import WanVideoVAE38
+
+    return WanVideoVAE38(vae_pth=str(checkpoint_path), dtype=torch_dtype, device=device)
+
+
+def load_wan22_ti2v_5b_components(
+    device: str = "cuda",
+    torch_dtype: torch.dtype = torch.bfloat16,
+    model_id: str = "Wan-AI/Wan2.2-TI2V-5B",
+    tokenizer_model_id: str = "Wan-AI/Wan2.2-TI2V-5B",
+    tokenizer_max_len: int = 512,
+    dit_config: dict[str, Any] | None = None,
+    load_text_encoder: bool = True,
+):
+    logger.info("Loading Wan2.2-TI2V-5B components...")
+    start = time.time()
+
+    if dit_config is None:
+        raise ValueError("`dit_config` is required for Wan2.2-TI2V-5B loading.")
+
+    checkpoint_dir = resolve_wan_checkpoint_dir(model_id)
+    tokenizer_dir = (
+        checkpoint_dir if tokenizer_model_id == model_id else resolve_wan_checkpoint_dir(tokenizer_model_id)
+    )
+    paths = resolve_wan_checkpoint_paths(
+        checkpoint_dir,
+        tokenizer_dir=tokenizer_dir,
+        load_text_encoder=load_text_encoder,
+    )
+
+    dit = load_wan_video_dit(
+        paths.dit,
+        dit_config=dit_config,
+        torch_dtype=torch_dtype,
+        device=device,
+    )
+    vae = load_wan_vae(paths.vae, torch_dtype=torch_dtype, device=device)
+
+    text_encoder: torch.nn.Module | None = None
+    tokenizer: HuggingfaceTokenizer | None = None
+    if load_text_encoder:
+        if paths.text_encoder is None or paths.tokenizer is None:
+            raise FileNotFoundError("Wan2.2 text encoder/tokenizer paths were not resolved.")
+        text_encoder = load_wan_text_encoder(
+            paths.text_encoder,
+            torch_dtype=torch_dtype,
+            device=device,
+        )
+        tokenizer = load_wan_tokenizer(paths.tokenizer, tokenizer_max_len=tokenizer_max_len)
+    else:
+        logger.info(
+            "Skipping pretrained text encoder/tokenizer load (`load_text_encoder=False`); "
+            "training must provide cached `context/context_mask`."
+        )
+
+    logger.info("Finished loading Wan2.2-TI2V-5B components in %.2f seconds.", time.time() - start)
+    return Wan22LoadedComponents(
+        dit=dit,
+        vae=vae,
+        text_encoder=text_encoder,
+        tokenizer=tokenizer,
+        dit_path=[str(path) for path in paths.dit],
+        vae_path=str(paths.vae),
+        text_encoder_path=str(paths.text_encoder) if paths.text_encoder is not None else None,
+        tokenizer_path=str(paths.tokenizer) if paths.tokenizer is not None else None,
+    )
+
+
+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
+
+
+__all__ = [
+    "WAN_DIT_PATTERN",
+    "WAN_T5_CHECKPOINT",
+    "WAN_T5_TOKENIZER",
+    "WAN_VAE_CHECKPOINT",
+    "Wan22LoadedComponents",
+    "WanCheckpointPaths",
+    "load_wan22_ti2v_5b_components",
+    "load_wan_text_encoder",
+    "load_wan_tokenizer",
+    "load_wan_vae",
+    "load_wan_video_dit",
+    "resolve_wan_checkpoint_dir",
+    "resolve_wan_checkpoint_paths",
+]

src/lerobot/policies/fastwam/wan_video_dit.py

@@ -0,0 +1,678 @@
+# 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 torch
+import torch.nn as nn
+import torch.nn.functional as functional
+from einops import rearrange
+
+from .wan.modules.model import (
+    WanAttentionBlock,
+    WanLayerNorm,
+    WanModel,
+    WanRMSNorm,
+    rope_apply,
+    rope_params,
+    sinusoidal_embedding_1d,
+)
+
+logger = logging.getLogger(__name__)
+
+
+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,
+) -> 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)
+    q = q.float()
+    k = k.float()
+    v = v.float()
+    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
+
+
+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.float64).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):
+        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,
+                window_size=(-1, -1),
+                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.window_size = (-1, -1)
+            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
+
+    @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)
+        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)
+        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,
+    ):
+        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,
+            window_size=(-1, -1),
+            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,
+                )
+                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
+
+        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)
+        token_t_emb = sinusoidal_embedding_1d(self.freq_dim, token_timesteps.reshape(-1)).to(
+            dtype=model_dtype
+        )
+        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))
+
+        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)
+
+
+__all__ = [
+    "FastWAMAttentionBlock",
+    "WanVideoDiT",
+    "apply_dense_rope",
+    "create_group_causal_attn_mask",
+    "fastwam_masked_attention",
+    "gradient_checkpoint_forward",
+    "modulate",
+    "precompute_freqs_cis",
+    "sinusoidal_embedding_1d",
+]

tests/policies/fastwam/test_fastwam_compliance.py

@@ -0,0 +1,254 @@
+#!/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 inspect
+import typing
+from pathlib import Path
+
+import pytest
+import torch
+from torch import nn
+
+from lerobot.configs import FeatureType, PolicyFeature
+from lerobot.policies.fastwam.configuration_fastwam import FastWAMConfig
+from lerobot.policies.fastwam.modeling_fastwam import FastWAMPolicy
+from lerobot.policies.fastwam.processor_fastwam import make_fastwam_pre_post_processors
+from lerobot.utils.constants import OBS_STATE
+
+ROOT = Path(__file__).resolve().parents[3]
+
+
+def test_package_init_exports_required_symbols():
+    init_source = (ROOT / "src" / "lerobot" / "policies" / "fastwam" / "__init__.py").read_text()
+
+    assert "FastWAMConfig" in init_source
+    assert "make_fastwam_pre_post_processors" in init_source
+
+
+def test_policy_config_is_exported_from_public_policies_package():
+    import lerobot.policies as policies
+
+    assert policies.FastWAMConfig is FastWAMConfig
+    assert "FastWAMConfig" in policies.__all__
+
+
+def test_fastwam_policy_docs_are_registered():
+    readme_path = ROOT / "src" / "lerobot" / "policies" / "fastwam" / "README.md"
+    wan_readme_path = ROOT / "src" / "lerobot" / "policies" / "fastwam" / "wan" / "README.md"
+    policy_readme_path = ROOT / "docs" / "source" / "policy_fastwam_README.md"
+    guide_path = ROOT / "docs" / "source" / "fastwam.mdx"
+    toctree_path = ROOT / "docs" / "source" / "_toctree.yml"
+
+    assert readme_path.is_symlink()
+    assert readme_path.resolve() == policy_readme_path.resolve()
+    assert wan_readme_path.exists()
+    wan_readme = wan_readme_path.read_text()
+    assert "Wan-Video/Wan2.2" in wan_readme
+    assert "42bf4cfaa384bc21833865abc2f9e6c0e67233dc" in wan_readme
+    assert policy_readme_path.exists()
+    assert guide_path.exists()
+    assert "local: fastwam" in toctree_path.read_text()
+
+
+def test_wan_backbone_code_is_isolated_from_lerobot_adapter():
+    wan_dir = ROOT / "src" / "lerobot" / "policies" / "fastwam" / "wan"
+
+    assert (wan_dir / "modules" / "attention.py").exists()
+    assert (wan_dir / "modules" / "model.py").exists()
+    assert (wan_dir / "modules" / "t5.py").exists()
+    assert (wan_dir / "modules" / "tokenizers.py").exists()
+    assert (wan_dir / "modules" / "vae2_1.py").exists()
+    assert (wan_dir / "modules" / "vae2_2.py").exists()
+    assert (wan_dir / "utils" / "fm_solvers.py").exists()
+    assert (wan_dir / "utils" / "fm_solvers_unipc.py").exists()
+
+    assert (wan_dir.parent / "wan_video_dit.py").exists()
+    assert (wan_dir.parent / "wan_adapters.py").exists()
+    assert (wan_dir.parent / "wan_components.py").exists()
+    assert not (wan_dir / "wan_video_dit.py").exists()
+    assert not (wan_dir / "wan_adapters.py").exists()
+    assert not (wan_dir / "wan_components.py").exists()
+
+
+def test_fastwam_text_encoder_uses_upstream_wan_modules_directly():
+    fastwam_dir = ROOT / "src" / "lerobot" / "policies" / "fastwam"
+    modular_source = (fastwam_dir / "modular_fastwam.py").read_text()
+    components_source = (fastwam_dir / "wan_components.py").read_text()
+
+    assert not (fastwam_dir / "wan_video_text_encoder.py").exists()
+    assert "from .wan.modules.t5 import umt5_xxl" in components_source
+    assert "from .wan.modules.tokenizers import HuggingfaceTokenizer" in components_source
+    assert "WAN_T5_ENCODER_KWARGS" not in components_source
+    assert "wan_video_text_encoder" not in modular_source
+
+
+def test_fastwam_vae_reuses_upstream_wan_modules():
+    fastwam_dir = ROOT / "src" / "lerobot" / "policies" / "fastwam"
+    vae_source = (fastwam_dir / "wan_adapters.py").read_text()
+
+    assert not (fastwam_dir / "wan_video_vae.py").exists()
+    assert "from .wan.modules.vae2_2 import Wan2_2_VAE" in vae_source
+    assert "mean = [" not in vae_source
+    assert "std = [" not in vae_source
+    assert "class Encoder3d_38" not in vae_source
+    assert "class Decoder3d_38" not in vae_source
+    assert "class VideoVAE38_" not in vae_source
+
+
+def test_fastwam_component_loading_uses_fixed_wan_checkpoint_layout():
+    modular_source = (ROOT / "src" / "lerobot" / "policies" / "fastwam" / "modular_fastwam.py").read_text()
+    modeling_source = (ROOT / "src" / "lerobot" / "policies" / "fastwam" / "modeling_fastwam.py").read_text()
+    components_source = (ROOT / "src" / "lerobot" / "policies" / "fastwam" / "wan_components.py").read_text()
+
+    assert "class ModelConfig" not in modular_source
+    assert "def load_state_dict" not in modular_source
+    assert "WAN22_MODEL_REGISTRY" not in modular_source
+    assert "class ModelConfig" not in components_source
+    assert "class WanComponentSource" not in components_source
+    assert "def load_state_dict" not in components_source
+    assert "WAN22_MODEL_REGISTRY" not in components_source
+    assert "hash_model_file" not in components_source
+    assert "_resolve_component_sources" not in components_source
+    assert "origin_file_pattern" not in components_source
+    assert "inspect.signature" not in components_source
+    assert "class FastWAMWanComponentPaths" not in modeling_source
+    assert "def _first_existing" not in modeling_source
+    assert "def _missing_wan_component_names" not in modeling_source
+    assert "WAN_T5_CHECKPOINT" in components_source
+    assert "WAN_VAE_CHECKPOINT" in components_source
+    assert "WAN_DIT_PATTERN" in components_source
+
+
+def test_fastwam_dit_reuses_upstream_wan_primitives():
+    dit_source = (ROOT / "src" / "lerobot" / "policies" / "fastwam" / "wan_video_dit.py").read_text()
+
+    assert "from .wan.modules.model import" in dit_source
+    assert "WanModel" in dit_source
+    for duplicated_symbol in [
+        "def flash_attention(",
+        "def sinusoidal_embedding_1d(",
+        "def rope_apply(",
+        "def unpatchify(",
+        "def _dense_video_freqs(",
+        "class RMSNorm(",
+        "class SelfAttention(",
+        "class CrossAttention(",
+        "class Head(",
+    ]:
+        assert duplicated_symbol not in dit_source
+
+
+def test_fastwam_inference_schedule_reuses_upstream_wan_sigmas():
+    modular_source = (ROOT / "src" / "lerobot" / "policies" / "fastwam" / "modular_fastwam.py").read_text()
+
+    assert "def _get_wan_sampling_sigmas" in modular_source
+    assert "from .wan.utils.fm_solvers import get_sampling_sigmas" in modular_source
+    assert "_get_wan_sampling_sigmas(num_inference_steps, shift)" in modular_source
+
+
+def test_policy_config_rejects_missing_required_image_and_action_features():
+    with pytest.raises(ValueError, match="image feature"):
+        FastWAMConfig(
+            input_features={OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(8,))},
+        )
+
+    with pytest.raises(ValueError, match="action"):
+        FastWAMConfig(
+            output_features={"not_action": PolicyFeature(type=FeatureType.ACTION, shape=(7,))},
+        )
+
+
+def test_policy_init_calls_validate_features_even_for_prebuilt_configs(monkeypatch):
+    cfg = FastWAMConfig(action_dim=3, proprio_dim=2, action_horizon=4, n_action_steps=2)
+    calls = []
+
+    def record_validate_features():
+        calls.append("called")
+
+    monkeypatch.setattr(cfg, "validate_features", record_validate_features)
+    monkeypatch.setattr(
+        FastWAMPolicy,
+        "_build_core_model",
+        lambda self, config: nn.Linear(1, 1),
+    )
+    FastWAMPolicy(cfg)
+
+    assert calls == ["called"]
+
+
+def test_required_policy_entrypoints_exist_with_discoverable_names():
+    assert FastWAMPolicy.config_class is FastWAMConfig
+    assert FastWAMPolicy.name == "fastwam"
+    assert callable(FastWAMPolicy.reset)
+    assert callable(FastWAMPolicy.get_optim_params)
+    assert callable(FastWAMPolicy.predict_action_chunk)
+    assert callable(FastWAMPolicy.select_action)
+    assert callable(FastWAMPolicy.forward)
+    assert callable(make_fastwam_pre_post_processors)
+    assert make_fastwam_pre_post_processors.__name__ == "make_fastwam_pre_post_processors"
+
+
+def test_policy_constructor_and_forward_match_byo_template_contract():
+    init_signature = inspect.signature(FastWAMPolicy.__init__)
+
+    assert "dataset_stats" in init_signature.parameters
+    assert "core_model" not in init_signature.parameters
+    assert typing.get_type_hints(FastWAMPolicy.forward)["return"] == dict[str, torch.Tensor]
+
+
+def test_saved_config_round_trips_policy_features(tmp_path):
+    cfg = FastWAMConfig(action_dim=7, proprio_dim=8, image_size=(224, 448))
+    cfg.save_pretrained(tmp_path)
+
+    loaded = FastWAMConfig.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_config_from_pretrained_ignores_unknown_fields(tmp_path):
+    cfg = FastWAMConfig()
+    cfg.save_pretrained(tmp_path)
+    config_path = tmp_path / "config.json"
+    payload = config_path.read_text()
+    payload = payload.replace(
+        '"torch_dtype": "bfloat16"',
+        '"torch_dtype": "bfloat16",\n  "unknown_fastwam_field": true',
+    )
+    config_path.write_text(payload)
+
+    loaded = FastWAMConfig.from_pretrained(tmp_path)
+
+    assert loaded.type == "fastwam"
+    assert not hasattr(loaded, "unknown_fastwam_field")
+
+
+def test_config_from_pretrained_does_not_use_non_wan22_tokenizer_repo_id(tmp_path):
+    cfg = FastWAMConfig()
+    cfg.save_pretrained(tmp_path)
+    config_path = tmp_path / "config.json"
+    payload = config_path.read_text()
+    payload = payload.replace(
+        '"tokenizer_model_id": "Wan-AI/Wan2.2-TI2V-5B"',
+        '"tokenizer_model_id": "somebody/old-tokenizer"',
+    )
+    config_path.write_text(payload)
+
+    loaded = FastWAMConfig.from_pretrained(tmp_path)
+
+    assert loaded.tokenizer_model_id == "Wan-AI/Wan2.2-TI2V-5B"

tests/policies/fastwam/test_fastwam_factory.py

@@ -0,0 +1,89 @@
+#!/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 pytest
+import torch
+
+from lerobot.policies.factory import get_policy_class, make_policy_config, make_pre_post_processors
+
+
+def test_fastwam_is_registered_in_policy_factory():
+    from lerobot.policies.fastwam.configuration_fastwam import FastWAMConfig
+    from lerobot.policies.fastwam.modeling_fastwam import FastWAMPolicy
+
+    cfg = make_policy_config("fastwam", action_dim=3, proprio_dim=2, action_horizon=4, n_action_steps=2)
+
+    assert isinstance(cfg, FastWAMConfig)
+    assert cfg.type == "fastwam"
+    assert get_policy_class("fastwam") is FastWAMPolicy
+
+
+def test_fastwam_pre_post_processors_are_available():
+    cfg = make_policy_config("fastwam", action_dim=3, proprio_dim=2, action_horizon=4, n_action_steps=2)
+
+    preprocessor, postprocessor = make_pre_post_processors(cfg)
+
+    assert preprocessor.name == "policy_preprocessor"
+    assert postprocessor.name == "policy_postprocessor"
+
+
+def test_fastwam_postprocessor_only_adds_action_inversion_when_configured():
+    from lerobot.policies.fastwam.processor_fastwam import (
+        FastWAMActionInversionProcessorStep,
+        FastWAMActionToggleProcessorStep,
+    )
+
+    default_cfg = make_policy_config(
+        "fastwam", action_dim=3, proprio_dim=2, action_horizon=4, n_action_steps=2
+    )
+    _, default_postprocessor = make_pre_post_processors(default_cfg)
+
+    assert any(isinstance(step, FastWAMActionToggleProcessorStep) for step in default_postprocessor.steps)
+    assert not any(
+        isinstance(step, FastWAMActionInversionProcessorStep) for step in default_postprocessor.steps
+    )
+
+    inverted_cfg = make_policy_config(
+        "fastwam",
+        action_dim=3,
+        proprio_dim=2,
+        action_horizon=4,
+        n_action_steps=2,
+        toggle_action_dimensions=[],
+        invert_dimensions=[-1],
+    )
+    _, inverted_postprocessor = make_pre_post_processors(inverted_cfg)
+
+    assert any(isinstance(step, FastWAMActionInversionProcessorStep) for step in inverted_postprocessor.steps)
+
+
+def test_fastwam_action_inversion_processor_flips_configured_dimensions():
+    from lerobot.policies.fastwam.processor_fastwam import FastWAMActionInversionProcessorStep
+
+    processor = FastWAMActionInversionProcessorStep(invert_dimensions=[0, -1])
+    action = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+
+    processed = processor.action(action)
+
+    assert torch.equal(processed, torch.tensor([[-1.0, 2.0, -3.0], [-4.0, 5.0, -6.0]]))
+    assert torch.equal(action, torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]))
+
+
+def test_fastwam_rejects_non_wan22_hub_model_ids():
+    from lerobot.policies.fastwam.configuration_fastwam import FastWAMConfig
+
+    with pytest.raises(ValueError, match="model_id"):
+        FastWAMConfig(model_id="somebody/other-model")

tests/policies/fastwam/test_fastwam_policy.py

@@ -0,0 +1,467 @@
+#!/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 torch
+from safetensors.torch import save_model
+from torch import nn
+
+from lerobot.policies.fastwam import modeling_fastwam
+from lerobot.policies.fastwam.configuration_fastwam import FastWAMConfig
+from lerobot.policies.fastwam.modeling_fastwam import FastWAMPolicy
+from lerobot.policies.fastwam.modular_fastwam import ActionDiT, MoT
+from lerobot.policies.fastwam.wan_video_dit import (
+    FastWAMAttentionBlock,
+    WanVideoDiT,
+    fastwam_masked_attention,
+    precompute_freqs_cis,
+)
+from lerobot.policies.pretrained import PreTrainedPolicy
+
+
+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):
+        horizon = kwargs["action_horizon"]
+        return {"action": torch.ones(horizon, 3)}
+
+
+def _patch_core_builder(monkeypatch):
+    monkeypatch.setattr(
+        FastWAMPolicy,
+        "_build_core_model",
+        lambda self, config: FakeFastWAMCore(),
+    )
+
+
+def test_action_attention_block_supports_mot_attention_dim_larger_than_hidden_dim():
+    block = FastWAMAttentionBlock(hidden_dim=16, attn_head_dim=8, num_heads=4, ffn_dim=32)
+    x = torch.zeros(1, 2, 16)
+    context = torch.zeros(1, 3, 16)
+    t_mod = torch.zeros(1, 6, 16)
+    freqs = precompute_freqs_cis(8, end=2).view(2, 1, -1)
+
+    output = block(x, context, t_mod, freqs)
+
+    assert output.shape == x.shape
+    assert block.self_attn.q.out_features == 32
+    assert block.self_attn.o.out_features == 16
+
+
+def test_fastwam_masked_attention_accepts_rope_float32_qk_with_bfloat16_values():
+    q = torch.zeros(1, 2, 32, dtype=torch.float32)
+    k = torch.zeros(1, 2, 32, dtype=torch.float32)
+    v = torch.zeros(1, 2, 32, dtype=torch.bfloat16)
+
+    out = fastwam_masked_attention(q=q, k=k, v=v, num_heads=4)
+
+    assert out.dtype == torch.float32
+    assert out.shape == v.shape
+
+
+def test_fastwam_masked_attention_runs_fp32_when_cache_promotes_keys():
+    q = torch.zeros(1, 2, 32, dtype=torch.bfloat16)
+    k = torch.zeros(1, 4, 32, dtype=torch.float32)
+    v = torch.zeros(1, 4, 32, dtype=torch.bfloat16)
+    mask = torch.ones(2, 4, dtype=torch.bool)
+
+    out = fastwam_masked_attention(q=q, k=k, v=v, num_heads=4, ctx_mask=mask)
+
+    assert out.dtype == torch.float32
+    assert out.shape == q.shape
+
+
+def test_attention_post_projection_casts_fp32_attention_to_block_dtype():
+    block = FastWAMAttentionBlock(hidden_dim=16, attn_head_dim=8, num_heads=4, ffn_dim=32).to(
+        dtype=torch.bfloat16
+    )
+    residual = torch.zeros(1, 2, 16, dtype=torch.bfloat16)
+    mixed_attn = torch.zeros(1, 2, 32, dtype=torch.float32)
+    gate_msa = torch.ones(1, 16, dtype=torch.bfloat16)
+    shift_mlp = torch.zeros(1, 16, dtype=torch.bfloat16)
+    scale_mlp = torch.zeros(1, 16, dtype=torch.bfloat16)
+    gate_mlp = torch.zeros(1, 16, dtype=torch.bfloat16)
+
+    out = MoT._apply_expert_post_block(
+        block=block,
+        residual_x=residual,
+        mixed_attn_out=mixed_attn,
+        gate_msa=gate_msa,
+        shift_mlp=shift_mlp,
+        scale_mlp=scale_mlp,
+        gate_mlp=gate_mlp,
+        context_payload=None,
+    )
+
+    assert out.dtype == torch.bfloat16
+    assert out.shape == residual.shape
+
+
+def test_attention_cross_projection_casts_fp32_attention_to_block_dtype():
+    block = FastWAMAttentionBlock(hidden_dim=16, attn_head_dim=8, num_heads=4, ffn_dim=32).to(
+        dtype=torch.bfloat16
+    )
+    x = torch.zeros(1, 2, 16, dtype=torch.bfloat16)
+    context = torch.zeros(1, 3, 16, dtype=torch.bfloat16)
+
+    out = block.apply_cross_attention(x, context)
+
+    assert out.dtype == torch.bfloat16
+    assert out.shape == x.shape
+
+
+def test_attention_norm3_handles_bfloat16_affine_weights():
+    block = FastWAMAttentionBlock(hidden_dim=16, attn_head_dim=8, num_heads=4, ffn_dim=32).to(
+        dtype=torch.bfloat16
+    )
+    x = torch.zeros(1, 2, 16, dtype=torch.bfloat16)
+
+    out = block.apply_norm3(x)
+
+    assert out.dtype == torch.bfloat16
+    assert out.shape == x.shape
+
+
+def test_attention_post_block_handles_bfloat16_cross_attention_norm():
+    block = FastWAMAttentionBlock(hidden_dim=16, attn_head_dim=8, num_heads=4, ffn_dim=32).to(
+        dtype=torch.bfloat16
+    )
+    residual = torch.zeros(1, 2, 16, dtype=torch.bfloat16)
+    mixed_attn = torch.zeros(1, 2, 32, dtype=torch.float32)
+    gate_msa = torch.ones(1, 16, dtype=torch.bfloat16)
+    shift_mlp = torch.zeros(1, 16, dtype=torch.bfloat16)
+    scale_mlp = torch.zeros(1, 16, dtype=torch.bfloat16)
+    gate_mlp = torch.zeros(1, 16, dtype=torch.bfloat16)
+    context_payload = {"context": torch.zeros(1, 3, 16, dtype=torch.bfloat16), "mask": None}
+
+    out = MoT._apply_expert_post_block(
+        block=block,
+        residual_x=residual,
+        mixed_attn_out=mixed_attn,
+        gate_msa=gate_msa,
+        shift_mlp=shift_mlp,
+        scale_mlp=scale_mlp,
+        gate_mlp=gate_mlp,
+        context_payload=context_payload,
+    )
+
+    assert out.dtype == torch.bfloat16
+    assert out.shape == residual.shape
+
+
+def test_video_dit_pre_dit_casts_double_latents_to_model_dtype():
+    model = WanVideoDiT(
+        hidden_dim=4,
+        in_dim=48,
+        ffn_dim=8,
+        out_dim=48,
+        text_dim=6,
+        freq_dim=4,
+        eps=1e-6,
+        patch_size=(1, 2, 2),
+        num_heads=1,
+        attn_head_dim=4,
+        num_layers=0,
+        seperated_timestep=True,
+        fuse_vae_embedding_in_latents=True,
+        video_attention_mask_mode="first_frame_causal",
+    ).to(dtype=torch.bfloat16)
+
+    state = model.pre_dit(
+        x=torch.zeros(1, 48, 1, 2, 2, dtype=torch.float64),
+        timestep=torch.zeros(1, dtype=torch.float64),
+        context=torch.zeros(1, 2, 6, dtype=torch.float64),
+        fuse_vae_embedding_in_latents=True,
+    )
+
+    assert state["tokens"].dtype == torch.bfloat16
+    assert state["context"].dtype == torch.bfloat16
+    assert state["t_mod"].dtype == torch.bfloat16
+
+
+def test_action_dit_pre_dit_casts_double_inputs_to_model_dtype():
+    model = ActionDiT(
+        hidden_dim=16,
+        action_dim=3,
+        ffn_dim=32,
+        text_dim=6,
+        freq_dim=4,
+        eps=1e-6,
+        num_heads=4,
+        attn_head_dim=8,
+        num_layers=0,
+    ).to(dtype=torch.bfloat16)
+
+    state = model.pre_dit(
+        action_tokens=torch.zeros(1, 2, 3, dtype=torch.float64),
+        timestep=torch.zeros(1, dtype=torch.float64),
+        context=torch.zeros(1, 2, 6, dtype=torch.float64),
+    )
+
+    assert state["tokens"].dtype == torch.bfloat16
+    assert state["context"].dtype == torch.bfloat16
+    assert state["t_mod"].dtype == torch.bfloat16
+
+
+def test_forward_adapts_lerobot_batch_to_fastwam_sample(monkeypatch):
+    _patch_core_builder(monkeypatch)
+    cfg = FastWAMConfig(action_dim=3, proprio_dim=2, action_horizon=4, n_action_steps=2)
+    policy = FastWAMPolicy(cfg)
+    batch = {
+        "observation.images.image": torch.zeros(1, 3, 16, 16),
+        "observation.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),
+    }
+
+    output = policy.forward(batch)
+
+    assert set(output) == {"loss", "loss_action"}
+    assert output["loss"].item() == 1.0
+    assert output["loss_action"].item() == 1.0
+
+
+def test_get_optim_params_returns_lerobot_optimizer_dict(monkeypatch):
+    _patch_core_builder(monkeypatch)
+    cfg = FastWAMConfig(action_dim=3, proprio_dim=2, action_horizon=4, n_action_steps=2)
+    policy = FastWAMPolicy(cfg)
+
+    optim_params = policy.get_optim_params()
+
+    assert isinstance(optim_params, dict)
+    assert set(optim_params) == {"params"}
+    assert list(optim_params["params"])
+
+
+def test_select_action_uses_action_queue(monkeypatch):
+    _patch_core_builder(monkeypatch)
+    cfg = FastWAMConfig(action_dim=3, proprio_dim=2, action_horizon=4, n_action_steps=2)
+    policy = FastWAMPolicy(cfg)
+    batch = {
+        "input_image": torch.zeros(1, 3, 16, 16),
+        "observation.state": torch.zeros(1, 2),
+        "context": torch.zeros(1, 5, 4096),
+        "context_mask": torch.ones(1, 5, dtype=torch.bool),
+    }
+
+    first = policy.select_action(batch)
+    second = policy.select_action(batch)
+
+    assert first.shape == (1, 3)
+    assert second.shape == (1, 3)
+
+
+def test_predict_action_prepares_lerobot_libero_observation(monkeypatch):
+    captured = {}
+
+    class CapturingCore(FakeFastWAMCore):
+        def infer_action(self, **kwargs):
+            captured.update(kwargs)
+            return {"action": torch.ones(1, 4, 3)}
+
+    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,
+        image_size=(16, 32),
+        input_features={
+            "observation.images.image": {"type": "VISUAL", "shape": (3, 16, 32)},
+            "observation.state": {"type": "STATE", "shape": (2,)},
+        },
+    )
+    policy = FastWAMPolicy(cfg)
+    batch = {
+        "observation.images.image": torch.ones(1, 3, 20, 20),
+        "observation.images.image2": torch.zeros(1, 3, 20, 20),
+        "observation.state": torch.zeros(1, 2),
+        "task": ["pick up the bowl"],
+    }
+
+    action = policy.predict_action_chunk(batch)
+
+    assert action.shape == (1, 4, 3)
+    assert captured["prompt"] == [cfg.prompt_template.format(task="pick up the bowl")]
+    assert tuple(captured["input_image"].shape) == (1, 3, 16, 32)
+    assert captured["input_image"].amin().item() == -1.0
+    assert captured["input_image"].amax().item() == 1.0
+    assert "num_video_frames" not in captured
+
+
+def test_predict_action_splits_parallel_eval_batch_into_single_infer_calls(monkeypatch):
+    captured = []
+
+    class CapturingCore(FakeFastWAMCore):
+        def infer_action(self, **kwargs):
+            captured.append(
+                {
+                    "input_image_shape": tuple(kwargs["input_image"].shape),
+                    "input_image_sum": float(kwargs["input_image"].sum()),
+                    "proprio_shape": tuple(kwargs["proprio"].shape),
+                    "proprio_sum": float(kwargs["proprio"].sum()),
+                    "prompt": kwargs["prompt"],
+                }
+            )
+            action = torch.full((1, kwargs["action_horizon"], 3), float(len(captured)))
+            return {"action": action}
+
+    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,
+        image_size=(16, 16),
+        input_features={
+            "observation.images.image": {"type": "VISUAL", "shape": (3, 16, 16)},
+            "observation.state": {"type": "STATE", "shape": (2,)},
+        },
+    )
+    policy = FastWAMPolicy(cfg)
+    batch = {
+        "observation.images.image": torch.stack(
+            [
+                torch.zeros(3, 16, 16),
+                torch.ones(3, 16, 16),
+                torch.full((3, 16, 16), 2.0),
+            ]
+        ),
+        "observation.state": torch.tensor([[0.0, 1.0], [2.0, 3.0], [4.0, 5.0]]),
+        "task": ["task 0", "task 1", "task 2"],
+    }
+
+    action = policy.predict_action_chunk(batch)
+
+    assert action.shape == (3, 4, 3)
+    assert action[:, 0, 0].tolist() == [1.0, 2.0, 3.0]
+    assert len(captured) == 3
+    assert [item["input_image_shape"] for item in captured] == [(1, 3, 16, 16)] * 3
+    assert [item["proprio_shape"] for item in captured] == [(1, 2)] * 3
+    assert [item["prompt"] for item in captured] == [
+        cfg.prompt_template.format(task="task 0"),
+        cfg.prompt_template.format(task="task 1"),
+        cfg.prompt_template.format(task="task 2"),
+    ]
+
+
+def test_from_pretrained_does_not_initialize_wan_backbone(monkeypatch, tmp_path):
+    cfg = FastWAMConfig(action_dim=3, proprio_dim=2, action_horizon=4, n_action_steps=2)
+    cfg.save_pretrained(tmp_path)
+    _patch_core_builder(monkeypatch)
+    reference_policy = FastWAMPolicy(cfg)
+    save_model(reference_policy, str(tmp_path / "model.safetensors"))
+
+    def fail_if_wan_pretrained_is_loaded(*args, **kwargs):
+        raise AssertionError("from_pretrained must not initialize or download Wan2.2 backbone components")
+
+    monkeypatch.setattr(
+        "lerobot.policies.fastwam.modular_fastwam.FastWAM.from_wan22_pretrained",
+        fail_if_wan_pretrained_is_loaded,
+    )
+    monkeypatch.setattr(
+        modeling_fastwam,
+        "_build_core_model_from_architecture",
+        lambda config: FakeFastWAMCore(),
+        raising=False,
+    )
+    loaded_components_from = []
+    monkeypatch.setattr(
+        FastWAMPolicy,
+        "load_wan_components_from_pretrained",
+        lambda self, path: loaded_components_from.append(path),
+    )
+
+    policy = FastWAMPolicy.from_pretrained(tmp_path, strict=False)
+
+    assert isinstance(policy.model, FakeFastWAMCore)
+    assert loaded_components_from == [tmp_path]
+
+
+def test_from_pretrained_resolves_hub_repo_to_snapshot_before_loading_sidecars(monkeypatch, tmp_path):
+    cfg = FastWAMConfig(action_dim=3, proprio_dim=2, action_horizon=4, n_action_steps=2)
+    cfg.save_pretrained(tmp_path)
+    snapshot_calls = []
+
+    def fake_snapshot_download(**kwargs):
+        snapshot_calls.append(kwargs)
+        return str(tmp_path)
+
+    @classmethod
+    def fake_base_from_pretrained(cls, pretrained_name_or_path, *, config=None, **kwargs):
+        assert pretrained_name_or_path == tmp_path
+        assert kwargs.pop("_skip_wan_init") is True
+        assert kwargs["strict"] is False
+        return cls(config, _skip_wan_init=True)
+
+    monkeypatch.setattr("huggingface_hub.snapshot_download", fake_snapshot_download)
+    monkeypatch.setattr(PreTrainedPolicy, "from_pretrained", fake_base_from_pretrained)
+    monkeypatch.setattr(
+        modeling_fastwam,
+        "_build_core_model_from_architecture",
+        lambda config: FakeFastWAMCore(),
+        raising=False,
+    )
+    loaded_components_from = []
+    monkeypatch.setattr(
+        FastWAMPolicy,
+        "load_wan_components_from_pretrained",
+        lambda self, path: loaded_components_from.append(path),
+    )
+
+    FastWAMPolicy.from_pretrained("org/fastwam", strict=False, local_files_only=True, revision="main")
+
+    assert snapshot_calls[0]["repo_id"] == "org/fastwam"
+    assert snapshot_calls[0]["local_files_only"] is True
+    assert snapshot_calls[0]["revision"] == "main"
+    assert loaded_components_from == [tmp_path]
+
+
+def test_save_pretrained_copies_wan_components(monkeypatch, tmp_path):
+    cfg = FastWAMConfig(action_dim=3, proprio_dim=2, action_horizon=4, n_action_steps=2)
+    source = tmp_path / "source"
+    tokenizer = source / "google" / "umt5-xxl"
+    tokenizer.mkdir(parents=True)
+    vae = source / "Wan2.2_VAE.pth"
+    text_encoder = source / "models_t5_umt5-xxl-enc-bf16.pth"
+    tokenizer_file = tokenizer / "tokenizer.json"
+    vae.write_bytes(b"vae")
+    text_encoder.write_bytes(b"text")
+    tokenizer_file.write_text("{}")
+    core = FakeFastWAMCore()
+    core.model_paths = {
+        "vae": str(vae),
+        "text_encoder": str(text_encoder),
+        "tokenizer": str(tokenizer),
+    }
+    monkeypatch.setattr(FastWAMPolicy, "_build_core_model", lambda self, config: core)
+    policy = FastWAMPolicy(cfg)
+
+    save_dir = tmp_path / "saved"
+    policy.save_pretrained(save_dir)
+
+    assert (save_dir / "model.safetensors").is_file()
+    assert (save_dir / "Wan2.2_VAE.pth").read_bytes() == b"vae"
+    assert (save_dir / "models_t5_umt5-xxl-enc-bf16.pth").read_bytes() == b"text"
+    assert (save_dir / "google" / "umt5-xxl" / "tokenizer.json").read_text() == "{}"

tests/policies/fastwam/test_fastwam_wan_components.py

@@ -0,0 +1,162 @@
+#!/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.
+
+from pathlib import Path
+
+import pytest
+from torch import nn
+
+from lerobot.policies.fastwam import modeling_fastwam
+from lerobot.policies.fastwam.configuration_fastwam import FastWAMConfig
+from lerobot.policies.fastwam.modeling_fastwam import (
+    FastWAMPolicy,
+    resolve_wan_component_paths,
+)
+from lerobot.policies.fastwam.wan_components import (
+    WAN_DIT_PATTERN,
+    WAN_T5_CHECKPOINT,
+    WAN_T5_TOKENIZER,
+    WAN_VAE_CHECKPOINT,
+    resolve_wan_checkpoint_paths,
+)
+
+
+def _make_wan_component_tree(root: Path) -> None:
+    tokenizer = root / WAN_T5_TOKENIZER
+    tokenizer.mkdir(parents=True)
+    (root / WAN_VAE_CHECKPOINT).touch()
+    (root / WAN_T5_CHECKPOINT).touch()
+    (root / "diffusion_pytorch_model-00001-of-00001.safetensors").touch()
+    (tokenizer / "tokenizer.json").touch()
+
+
+def test_resolve_wan_component_paths_finds_complete_local_directory(tmp_path):
+    _make_wan_component_tree(tmp_path)
+
+    paths = resolve_wan_component_paths(tmp_path)
+
+    assert paths.vae == tmp_path / WAN_VAE_CHECKPOINT
+    assert paths.text_encoder == tmp_path / WAN_T5_CHECKPOINT
+    assert paths.tokenizer == tmp_path / WAN_T5_TOKENIZER
+
+
+def test_resolve_wan_component_paths_does_not_require_original_dit_shards(tmp_path):
+    _make_wan_component_tree(tmp_path)
+    for shard in tmp_path.glob(WAN_DIT_PATTERN):
+        shard.unlink()
+
+    paths = resolve_wan_component_paths(tmp_path)
+
+    assert paths.dit == []
+    assert paths.vae == tmp_path / WAN_VAE_CHECKPOINT
+    assert paths.text_encoder == tmp_path / WAN_T5_CHECKPOINT
+    assert paths.tokenizer == tmp_path / WAN_T5_TOKENIZER
+
+
+def test_resolve_wan_checkpoint_paths_uses_official_wan_layout(tmp_path):
+    _make_wan_component_tree(tmp_path)
+
+    paths = resolve_wan_checkpoint_paths(tmp_path)
+
+    assert paths.root == tmp_path
+    assert paths.dit == [tmp_path / "diffusion_pytorch_model-00001-of-00001.safetensors"]
+    assert paths.vae == tmp_path / WAN_VAE_CHECKPOINT
+    assert paths.text_encoder == tmp_path / WAN_T5_CHECKPOINT
+    assert paths.tokenizer == tmp_path / WAN_T5_TOKENIZER
+    assert WAN_DIT_PATTERN == "diffusion_pytorch_model*.safetensors"
+
+
+def test_resolve_wan_component_paths_rejects_partial_local_directory(tmp_path):
+    _make_wan_component_tree(tmp_path)
+    (tmp_path / WAN_T5_CHECKPOINT).unlink()
+
+    with pytest.raises(FileNotFoundError, match="text encoder"):
+        resolve_wan_component_paths(tmp_path)
+
+
+def test_policy_config_construction_loads_wan22_backbone_from_config(monkeypatch):
+    class TinyCore(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.text_encoder = None
+
+    calls = []
+
+    def fake_from_wan22_pretrained(**kwargs):
+        calls.append(kwargs)
+        return TinyCore()
+
+    monkeypatch.setattr(
+        "lerobot.policies.fastwam.modular_fastwam.FastWAM.from_wan22_pretrained",
+        fake_from_wan22_pretrained,
+    )
+
+    cfg = FastWAMConfig()
+    policy = FastWAMPolicy(cfg)
+
+    assert policy.model.text_encoder is None
+    assert calls == [
+        {
+            "device": cfg.device,
+            "torch_dtype": modeling_fastwam._dtype_from_name(cfg.torch_dtype),
+            "model_id": "Wan-AI/Wan2.2-TI2V-5B",
+            "tokenizer_model_id": "Wan-AI/Wan2.2-TI2V-5B",
+            "tokenizer_max_len": cfg.tokenizer_max_len,
+            "load_text_encoder": cfg.load_text_encoder,
+            "proprio_dim": cfg.proprio_dim,
+            "video_dit_config": cfg.video_dit_config,
+            "action_dit_config": cfg.action_dit_config,
+            "mot_checkpoint_mixed_attn": cfg.mot_checkpoint_mixed_attn,
+            "video_train_shift": float(cfg.video_scheduler["train_shift"]),
+            "video_infer_shift": float(cfg.video_scheduler["infer_shift"]),
+            "video_num_train_timesteps": int(cfg.video_scheduler["num_train_timesteps"]),
+            "action_train_shift": float(cfg.action_scheduler["train_shift"]),
+            "action_infer_shift": float(cfg.action_scheduler["infer_shift"]),
+            "action_num_train_timesteps": int(cfg.action_scheduler["num_train_timesteps"]),
+            "loss_lambda_video": float(cfg.loss["lambda_video"]),
+            "loss_lambda_action": float(cfg.loss["lambda_action"]),
+        }
+    ]
+
+
+def test_explicit_local_wan_path_is_preserved(tmp_path):
+    cfg = FastWAMConfig(model_id=str(tmp_path), tokenizer_model_id=str(tmp_path))
+
+    assert cfg.model_id == str(tmp_path)
+    assert cfg.tokenizer_model_id == str(tmp_path)
+
+
+def test_other_hub_model_ids_are_rejected():
+    with pytest.raises(ValueError, match="model_id"):
+        FastWAMConfig(model_id="somebody/other-model")
+
+    with pytest.raises(ValueError, match="tokenizer_model_id"):
+        FastWAMConfig(tokenizer_model_id="somebody/other-tokenizer")
+
+
+def test_resolve_wan_checkpoint_paths_can_skip_text_encoder(tmp_path):
+    _make_wan_component_tree(tmp_path)
+    (tmp_path / WAN_T5_CHECKPOINT).unlink()
+    shutil_tokenizer = tmp_path / WAN_T5_TOKENIZER
+    for child in shutil_tokenizer.iterdir():
+        child.unlink()
+    shutil_tokenizer.rmdir()
+    shutil_tokenizer.parent.rmdir()
+
+    paths = resolve_wan_checkpoint_paths(tmp_path, load_text_encoder=False)
+
+    assert paths.text_encoder is None
+    assert paths.tokenizer is None