Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/franka-coffee.webp filter=lfs diff=lfs merge=lfs -text
+examples/franka-tennis.webp filter=lfs diff=lfs merge=lfs -text
+examples/teaser.webp filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,4 @@
+__pycache__/
+*.pyc
+.DS_Store
+.venv/

README.md

@@ -1,13 +1,31 @@
 ---
-title: Dynamic Vla DOM
-emoji: 🏢
-colorFrom: indigo
-colorTo: blue
+title: DynamicVLA DOM Demo
+emoji: 🦾
+colorFrom: green
+colorTo: indigo
 sdk: gradio
 sdk_version: 6.24.0
-python_version: '3.12'
 app_file: app.py
-pinned: false
+short_description: Predict DynamicVLA action chunks on DOM
+python_version: "3.12"
+startup_duration_timeout: 1h
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# DynamicVLA · DOM Demo
+
+Unofficial Hugging Face Space for [`hzxie/dynamic-vla-DOM`](https://huggingface.co/hzxie/dynamic-vla-DOM) —
+a 0.4B vision-language-action policy for **dynamic object manipulation**.
+
+Upload wrist + opposite-camera frames and a language instruction. The Space
+loads the official checkpoint and returns a **20-step end-effector action chunk**
+(xyz + Euler + gripper) with a 3D path plot.
+
+This is a single-chunk policy demo. Closed-loop evaluation still needs the
+official [Isaac Lab stack](https://github.com/hzxie/DynamicVLA).
+
+- Paper: [arXiv:2601.22153](https://arxiv.org/abs/2601.22153)
+- Dataset: [`hzxie/DOM`](https://huggingface.co/datasets/hzxie/DOM)
+- Project: [infinitescript.com/project/dynamic-vla](https://www.infinitescript.com/project/dynamic-vla/)
+
+Weights are released under NTU S-Lab License 1.0. This Space vendors the
+policy code from the official repo for inference only.

app.py

@@ -0,0 +1,404 @@
+"""DynamicVLA (hzxie/dynamic-vla-DOM) action-chunk demo."""
+
+from __future__ import annotations
+
+import importlib.util
+import json
+import os
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+
+import spaces  # noqa: E402  — must precede any CUDA-touching import
+
+if importlib.util.find_spec("lerobot") is None:
+    subprocess.check_call(
+        [sys.executable, "-m", "pip", "install", "--no-deps", "lerobot==0.3.3"]
+    )
+
+import numpy as np  # noqa: E402
+import pandas as pd  # noqa: E402
+import plotly.graph_objects as go  # noqa: E402
+import torch  # noqa: E402
+from huggingface_hub import snapshot_download  # noqa: E402
+from lerobot.configs.types import FeatureType, PolicyFeature  # noqa: E402
+from PIL import Image  # noqa: E402
+
+from policies.dynamicvla.configuration_dynamicvla import DynamicVLAConfig  # noqa: E402
+from policies.dynamicvla.modeling_dynamicvla import (  # noqa: E402
+    DynamicVLAPolicy,
+    load_dynamicvla,
+)
+
+import gradio as gr  # noqa: E402
+
+MODEL_ID = "hzxie/dynamic-vla-DOM"
+IMG_H, IMG_W = 360, 480
+N_OBS = 2
+ACTION_COLS = ["x", "y", "z", "roll", "pitch", "yaw", "gripper"]
+ROOT = Path(__file__).resolve().parent
+EXAMPLES = ROOT / "examples"
+
+FEATURE_TYPES = {
+    "STATE": FeatureType.STATE,
+    "VISUAL": FeatureType.VISUAL,
+    "ACTION": FeatureType.ACTION,
+}
+
+
+def _features(spec: dict) -> dict:
+    return {
+        key: PolicyFeature(type=FEATURE_TYPES[ft["type"]], shape=tuple(ft["shape"]))
+        for key, ft in spec.items()
+    }
+
+
+def _identity_stats(cfg: DynamicVLAConfig) -> dict[str, dict[str, torch.Tensor]]:
+    """MEAN_STD with mean=0 / std=1 is a passthrough. Checkpoint strips norm buffers."""
+    stats: dict[str, dict[str, torch.Tensor]] = {}
+    for name, feat in {**cfg.input_features, **cfg.output_features}.items():
+        if feat.type in (FeatureType.STATE, FeatureType.ACTION):
+            shape = tuple(feat.shape)
+            stats[name] = {
+                "mean": torch.zeros(shape, dtype=torch.float32),
+                "std": torch.ones(shape, dtype=torch.float32),
+                "min": torch.full(shape, -1.0, dtype=torch.float32),
+                "max": torch.ones(shape, dtype=torch.float32),
+            }
+    return stats
+
+
+def _build_config(ckpt_dir: str) -> DynamicVLAConfig:
+    with open(os.path.join(ckpt_dir, "config.json"), encoding="utf-8") as fh:
+        raw = json.load(fh)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    cfg = DynamicVLAConfig(
+        input_features=_features(raw["input_features"]),
+        output_features=_features(raw["output_features"]),
+        device=device,
+    )
+    skip = {"type", "device", "input_features", "output_features"}
+    for key, value in raw.items():
+        attr = key.lower()
+        if attr in skip or value is None or not hasattr(cfg, attr):
+            continue
+        setattr(cfg, attr, value)
+    cfg.enable_streaming = False
+    return cfg
+
+
+def load_policy() -> DynamicVLAPolicy:
+    ckpt_dir = snapshot_download(MODEL_ID)
+    cfg = _build_config(ckpt_dir)
+    policy = DynamicVLAPolicy(cfg, dataset_stats=_identity_stats(cfg))
+    load_dynamicvla(
+        policy,
+        os.path.join(ckpt_dir, "model.safetensors"),
+        device="cpu",
+        checkpoint_keys_mapping="model._orig_mod.//model.",
+    )
+    policy.eval()
+    return policy.to("cuda")
+
+
+POLICY = load_policy()
+
+
+def _as_image(img) -> Image.Image | None:
+    if img is None:
+        return None
+    if isinstance(img, Image.Image):
+        return img.convert("RGB")
+    if isinstance(img, np.ndarray):
+        if img.ndim == 3 and img.shape[-1] == 4:
+            img = img[..., :3]
+        if img.dtype != np.uint8:
+            img = np.clip(img, 0, 255).astype(np.uint8) if img.max() > 1.5 else (
+                np.clip(img * 255.0, 0, 255).astype(np.uint8)
+            )
+        return Image.fromarray(img).convert("RGB")
+    return Image.open(img).convert("RGB")
+
+
+def _to_nchw(img: Image.Image) -> torch.Tensor:
+    resized = img.resize((IMG_W, IMG_H), Image.BILINEAR)
+    arr = np.asarray(resized, dtype=np.float32) / 255.0
+    return torch.from_numpy(arr).permute(2, 0, 1)
+
+
+def _stack_obs(current: Image.Image, previous: Image.Image | None) -> torch.Tensor:
+    cur = _to_nchw(current)
+    prev = _to_nchw(previous) if previous is not None else cur
+    return torch.stack([prev, cur], dim=0)  # (n_obs, C, H, W)
+
+
+def _plot_path(actions: np.ndarray) -> go.Figure:
+    xs, ys, zs = actions[:, 0], actions[:, 1], actions[:, 2]
+    fig = go.Figure(
+        data=[
+            go.Scatter3d(
+                x=xs,
+                y=ys,
+                z=zs,
+                mode="lines+markers",
+                marker={"size": 4, "color": np.arange(len(xs)), "colorscale": "Viridis"},
+                line={"width": 5, "color": "#2ecc71"},
+                name="EE path",
+            ),
+            go.Scatter3d(
+                x=[xs[0]],
+                y=[ys[0]],
+                z=[zs[0]],
+                mode="markers",
+                marker={"size": 8, "color": "#27ae60"},
+                name="start",
+            ),
+            go.Scatter3d(
+                x=[xs[-1]],
+                y=[ys[-1]],
+                z=[zs[-1]],
+                mode="markers",
+                marker={"size": 8, "color": "#e74c3c"},
+                name="end",
+            ),
+        ]
+    )
+    fig.update_layout(
+        template="plotly_dark",
+        height=420,
+        margin={"l": 0, "r": 0, "t": 30, "b": 0},
+        scene={
+            "xaxis_title": "x (m)",
+            "yaxis_title": "y (m)",
+            "zaxis_title": "z (m)",
+            "aspectmode": "data",
+        },
+        title="Predicted 20-step end-effector chunk",
+        paper_bgcolor="rgba(0,0,0,0)",
+        plot_bgcolor="rgba(0,0,0,0)",
+        legend={"orientation": "h"},
+    )
+    return fig
+
+
+def _summarize(actions: np.ndarray, elapsed: float, instruction: str) -> str:
+    delta = actions[-1, :3] - actions[0, :3]
+    grip = actions[:, -1]
+    return (
+        f"**Instruction:** {instruction.strip()}\n\n"
+        f"**Chunk:** {len(actions)} steps · **{elapsed:.2f}s** GPU\n\n"
+        f"- Start xyz: `{actions[0, :3].round(4).tolist()}`\n"
+        f"- End xyz: `{actions[-1, :3].round(4).tolist()}`\n"
+        f"- Net Δxyz: `{delta.round(4).tolist()}`\n"
+        f"- Gripper: min `{grip.min():.3f}` → max `{grip.max():.3f}` "
+        f"(last `{grip[-1]:.3f}`)\n"
+        f"- Rotation (last rpy): `{actions[-1, 3:6].round(4).tolist()}`"
+    )
+
+
+@spaces.GPU(duration=60)
+def predict_action_chunk(
+    wrist: np.ndarray | Image.Image | None,
+    opposite: np.ndarray | Image.Image | None,
+    wrist_prev: np.ndarray | Image.Image | None,
+    opposite_prev: np.ndarray | Image.Image | None,
+    instruction: str,
+    x: float,
+    y: float,
+    z: float,
+    roll: float,
+    pitch: float,
+    yaw: float,
+    apply_delta: bool,
+) -> tuple[pd.DataFrame, go.Figure, str]:
+    """Predict a 20-step DynamicVLA action chunk from dual-camera frames."""
+    if not instruction or not instruction.strip():
+        raise gr.Error("Provide a language instruction.")
+    wrist_img = _as_image(wrist)
+    opp_img = _as_image(opposite)
+    if wrist_img is None or opp_img is None:
+        raise gr.Error("Upload both wrist and opposite camera frames.")
+
+    device = "cuda"
+    wrist_t = _stack_obs(wrist_img, _as_image(wrist_prev)).unsqueeze(0).to(device)
+    opp_t = _stack_obs(opp_img, _as_image(opposite_prev)).unsqueeze(0).to(device)
+    state = torch.tensor(
+        [[[x, y, z, roll, pitch, yaw]] * N_OBS], dtype=torch.float32, device=device
+    )
+
+    batch = {
+        "observation.images.wrist_cam": wrist_t,
+        "observation.images.opst_cam": opp_t,
+        "observation.state": state,
+        "task": [instruction.strip()],
+    }
+
+    POLICY.reset()
+    tick = time.perf_counter()
+    with torch.inference_mode():
+        actions = POLICY.predict_action_chunk(batch)
+        if apply_delta and getattr(POLICY.config, "use_delta_action", True):
+            actions = actions.clone()
+            actions[..., :6] = actions[..., :6] + state[:, -1:, :6]
+    elapsed = time.perf_counter() - tick
+    acts = actions[0].detach().float().cpu().numpy()
+
+    table = pd.DataFrame(acts, columns=ACTION_COLS)
+    table.insert(0, "step", np.arange(len(table)))
+    return table, _plot_path(acts), _summarize(acts, elapsed, instruction)
+
+
+def _example_row(stem: str, instruction: str) -> list:
+    img = str(EXAMPLES / f"{stem}.png")
+    return [img, img, None, None, instruction, 0.40, 0.00, 0.30, 0.0, 0.0, 0.0, True]
+
+
+GALLERY_MD = """
+## DynamicVLA on DOM
+
+**DynamicVLA** (0.4B, SmolLM2-360M + FastViT) is a VLA for *moving* objects.
+It adds **Continuous Inference** and **Latent-aware Action Streaming** so the
+policy does not freeze between action chunks.
+
+This Space runs the official [`hzxie/dynamic-vla-DOM`](https://huggingface.co/hzxie/dynamic-vla-DOM)
+checkpoint and predicts a **20-step** 7-DoF end-effector chunk
+(`[x, y, z, roll, pitch, yaw, gripper]`).
+
+Closed-loop Isaac Lab eval is *not* hosted here — use
+[hzxie/DynamicVLA](https://github.com/hzxie/DynamicVLA) for that.
+
+| | |
+|---|---|
+| Paper | [arXiv:2601.22153](https://arxiv.org/abs/2601.22153) |
+| Dataset | [`hzxie/DOM`](https://huggingface.co/datasets/hzxie/DOM) — 200K episodes, 2.8K scenes, 206 objects |
+| Weights | [`hzxie/dynamic-vla-DOM`](https://huggingface.co/hzxie/dynamic-vla-DOM) |
+| Project | [infinitescript.com/project/dynamic-vla](https://www.infinitescript.com/project/dynamic-vla/) |
+| Spotlight | [YouTube](https://youtu.be/NmJnHcI04_Q) |
+
+<iframe width="100%" height="360" src="https://www.youtube.com/embed/NmJnHcI04_Q"
+title="DynamicVLA spotlight" frameborder="0"
+allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture"
+allowfullscreen></iframe>
+"""
+
+
+def build_ui() -> gr.Blocks:
+    theme = gr.themes.Soft(primary_hue="green", neutral_hue="zinc").set(
+        body_background_fill="#0b0f0c",
+        block_background_fill="#121814",
+    )
+    with gr.Blocks(theme=theme, title="DynamicVLA · DOM") as demo:
+        gr.Markdown(
+            "# DynamicVLA — DOM action-chunk demo\n"
+            "0.4B VLA for dynamic object manipulation. "
+            "Upload wrist + scene cameras, write an instruction, get a 20-step EE chunk."
+        )
+        with gr.Tabs():
+            with gr.Tab("Predict"):
+                with gr.Row():
+                    with gr.Column():
+                        wrist = gr.Image(label="Wrist camera (current)", type="numpy")
+                        opposite = gr.Image(
+                            label="Opposite / scene camera (current)", type="numpy"
+                        )
+                        with gr.Accordion("Previous frames (optional, n_obs=2)", open=False):
+                            wrist_prev = gr.Image(
+                                label="Wrist camera (t−1)", type="numpy"
+                            )
+                            opposite_prev = gr.Image(
+                                label="Opposite camera (t−1)", type="numpy"
+                            )
+                        instruction = gr.Textbox(
+                            label="Language instruction",
+                            placeholder="Pick up the rolling cylinder and place it onto the wooden block.",
+                            lines=2,
+                        )
+                        with gr.Accordion("Current EE state (meters / rad)", open=False):
+                            with gr.Row():
+                                x = gr.Number(value=0.40, label="x")
+                                y = gr.Number(value=0.00, label="y")
+                                z = gr.Number(value=0.30, label="z")
+                            with gr.Row():
+                                roll = gr.Number(value=0.0, label="roll")
+                                pitch = gr.Number(value=0.0, label="pitch")
+                                yaw = gr.Number(value=0.0, label="yaw")
+                            apply_delta = gr.Checkbox(
+                                value=True,
+                                label="Add delta actions to current EE state",
+                            )
+                        run = gr.Button("Predict action chunk", variant="primary")
+                    with gr.Column():
+                        summary = gr.Markdown("Upload both views and run.")
+                        path = gr.Plot(label="EE trajectory")
+                        table = gr.Dataframe(label="Action chunk (20 × 7)")
+
+                inputs = [
+                    wrist,
+                    opposite,
+                    wrist_prev,
+                    opposite_prev,
+                    instruction,
+                    x,
+                    y,
+                    z,
+                    roll,
+                    pitch,
+                    yaw,
+                    apply_delta,
+                ]
+                run.click(
+                    fn=predict_action_chunk,
+                    inputs=inputs,
+                    outputs=[table, path, summary],
+                )
+                gr.Examples(
+                    examples=[
+                        _example_row(
+                            "franka-coffee",
+                            "Pick up the rolling cylinder and place it onto the wooden block.",
+                        ),
+                        _example_row(
+                            "piper-sesame",
+                            "Grasp the rolling roasted sesame container and place it onto the blue frisbee.",
+                        ),
+                        _example_row(
+                            "franka-tennis",
+                            "Get hold of the moving tennis ball and position it into the paper bowl.",
+                        ),
+                    ],
+                    inputs=inputs,
+                    label="DOM-style prompts (official comparison stills as both views)",
+                    cache_examples=True,
+                    cache_mode="lazy",
+                )
+
+            with gr.Tab("About"):
+                gr.Markdown(GALLERY_MD)
+                if (EXAMPLES / "teaser.webp").exists():
+                    gr.Image(
+                        value=str(EXAMPLES / "teaser.webp"),
+                        label="Official teaser",
+                        interactive=False,
+                    )
+                gr.Markdown(
+                    "```bibtex\n"
+                    "@article{xie2026dynamicvla,\n"
+                    "  title   = {DynamicVLA: A Vision-Language-Action Model for Dynamic Object Manipulation},\n"
+                    "  author  = {Xie, Haozhe and Wen, Beichen and Zheng, Jiarui and Chen, Zhaoxi\n"
+                    "             and Hong, Fangzhou and Diao, Haiwen and Liu, Ziwei},\n"
+                    "  journal = {arXiv preprint arXiv:2601.22153},\n"
+                    "  year    = {2026}\n"
+                    "}\n"
+                    "```"
+                )
+    return demo
+
+
+demo = build_ui()
+
+if __name__ == "__main__":
+    demo.launch(mcp_server=True)

policies/dynamicvla/__init__.py

@@ -0,0 +1 @@
+"""Vendored DynamicVLA policy package from https://github.com/hzxie/DynamicVLA."""

policies/dynamicvla/configuration_dynamicvla.py

@@ -0,0 +1,152 @@
+# -*- coding: utf-8 -*-
+#
+# @File:   configuration_dynamicvla.py
+# @Author: Haozhe Xie
+# @Date:   2025-08-21 15:22:31
+# @Last Modified by: Haozhe Xie
+# @Last Modified at: 2026-02-26 10:04:17
+# @Email:  root@haozhexie.com
+
+from dataclasses import dataclass, field
+
+from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.optim.optimizers import AdamWConfig
+from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
+
+
+@PreTrainedConfig.register_subclass("dynamicvla")
+@dataclass
+class DynamicVLAConfig(PreTrainedConfig):
+    # Input / output structure.
+    n_obs_steps: int = 1
+    chunk_size: int = 50
+    n_action_steps: int = 50
+
+    normalization_mapping: dict[str, NormalizationMode] = field(
+        default_factory=lambda: {
+            "VISUAL": NormalizationMode.IDENTITY,
+            "STATE": NormalizationMode.MEAN_STD,
+            "ACTION": NormalizationMode.MEAN_STD,
+        }
+    )
+
+    # Shorter state and action vectors will be padded
+    max_state_dim: int = 32
+    max_action_dim: int = 32
+    # Image preprocessing
+    resize_imgs_with_padding: tuple[int, int] = (512, 512)
+    # Add empty images. Used by dynamicvla_aloha_sim which adds the empty
+    # left and right wrist cameras in addition to the top camera.
+    empty_cameras: int = 0
+    # Converts the joint and gripper values from the standard Aloha space to the space
+    # used by the pi internal runtime which was used to train the base model.
+    adapt_to_pi_aloha: bool = False
+    # Converts joint dimensions to deltas with respect to the current state before
+    # passing to the model. Gripper dimensions will remain in absolute values.
+    use_delta_joint_actions_aloha: bool = False
+    # Use delta action prediction (relative to the current robot state)
+    use_delta_action: bool = True
+    # Streaming inference
+    enable_streaming: bool = False
+    # Multi-timestep fusion ("conv" or "attn")
+    temporal_fusion: str = "conv"
+    # Tokenizer
+    tokenizer_max_length: int = 48
+    # Decoding
+    num_steps: int = 10
+    # Attention utils
+    use_cache: bool = True
+    # Finetuning settings
+    freeze_vision_model: bool = True
+    freeze_connector: bool = True
+    freeze_text_model: bool = True
+    train_state_proj: bool = True
+    # Training presets
+    optimizer_lr: float = 1e-4
+    optimizer_betas: tuple[float, float] = (0.9, 0.95)
+    optimizer_eps: float = 1e-8
+    optimizer_weight_decay: float = 1e-10
+    optimizer_grad_clip_norm: float = 10
+    scheduler_warmup_steps: int = 1_000
+    scheduler_decay_steps: int = 30_000
+    scheduler_decay_lr: float = 2.5e-6
+    # Select the VLM backbone.
+    attention_mode: str = "cross_attn"
+    prefix_length: int = -1
+    pad_language_to: str = "longest"  # "max_length"
+    # Less or equal to 0 is the default where the action expert has the same number of
+    # layers of VLM. Otherwise the expert have less layers.
+    num_expert_layers: int = -1
+    num_expert_skip_layers: int = 0
+    # VLM settings
+    vlm_model_name: str = "HuggingFaceTB/SmolLM2-360M-Instruct"
+    num_vlm_layers: int = 16
+    # SmolVLM Settings
+    smolvlm_patch_size: int = 16
+    smolvlm_attention_heads: int = 12
+    smolvlm_hidden_size: int = 768
+    smolvlm_intermediate_size: int = 3072
+    # FastVLM Settings
+    fastvlm_inference_mode: bool = True
+    # Interleave SA layers each self_attn_every_n_layers
+    self_attn_every_n_layers: int = 2
+    # The action expert hidden size (wrt to the VLM)
+    expert_width_multiplier: float = 0.75
+    # sensitivity range for the timestep used in sine-cosine positional encoding
+    min_period: float = 4e-3
+    max_period: float = 4.0
+
+    def __post_init__(self):
+        super().__post_init__()
+
+        """Input validation (not exhaustive)."""
+        if self.n_action_steps > self.chunk_size:
+            raise ValueError(
+                "The chunk size is the upper bound for the number of action steps per"
+                f" model invocation. Got {self.n_action_steps} for `n_action_steps` and"
+                f" {self.chunk_size} for `chunk_size`."
+            )
+        if self.use_delta_joint_actions_aloha:
+            raise NotImplementedError(
+                "`use_delta_joint_actions_aloha` is used by dynamicvla for aloha real"
+                " models. It is not ported yet in LeRobot."
+            )
+
+    def validate_features(self) -> None:
+        for i in range(self.empty_cameras):
+            key = f"observation.images.empty_camera_{i}"
+            empty_camera = PolicyFeature(
+                type=FeatureType.VISUAL,
+                shape=(3, 480, 640),
+            )
+            self.input_features[key] = empty_camera
+
+    def get_optimizer_preset(self) -> AdamWConfig:
+        return AdamWConfig(
+            lr=self.optimizer_lr,
+            betas=self.optimizer_betas,
+            eps=self.optimizer_eps,
+            weight_decay=self.optimizer_weight_decay,
+            grad_clip_norm=self.optimizer_grad_clip_norm,
+        )
+
+    def get_scheduler_preset(self):
+        return CosineDecayWithWarmupSchedulerConfig(
+            peak_lr=self.optimizer_lr,
+            decay_lr=self.scheduler_decay_lr,
+            num_warmup_steps=self.scheduler_warmup_steps,
+            num_decay_steps=self.scheduler_decay_steps,
+        )
+
+    @property
+    def observation_delta_indices(self) -> list:
+        return [0]
+
+    @property
+    def action_delta_indices(self) -> list:
+        return list(range(self.chunk_size))
+
+    @property
+    def reward_delta_indices(self) -> None:
+        return None

policies/dynamicvla/modeling_dynamicvla.py

@@ -0,0 +1,1189 @@
+# -*- coding: utf-8 -*-
+#
+# @File:   modeling_dynamicvla.py
+# @Author: Haozhe Xie
+# @Date:   2025-08-21 15:23:45
+# @Last Modified by: Haozhe Xie
+# @Last Modified at: 2026-02-26 10:10:57
+# @Email:  root@haozhexie.com
+
+import logging
+import math
+import os
+import re
+import time
+from collections import deque
+
+import safetensors
+import torch
+import torch.multiprocessing as mp
+import torch.nn.functional as F  # noqa: N812
+from lerobot.constants import ACTION, OBS_STATE
+from lerobot.policies.normalize import Normalize, Unnormalize
+from lerobot.policies.pretrained import PreTrainedPolicy
+from lerobot.policies.utils import populate_queues
+from lerobot.utils.utils import get_safe_dtype
+from transformers import AutoConfig, SmolVLMForConditionalGeneration
+
+from policies.dynamicvla.configuration_dynamicvla import DynamicVLAConfig
+from policies.dynamicvla.modeling_fastvlm import (
+    FastViTConfig,
+    FastVLMConfig,
+    FastVLMForConditionalGeneration,
+)
+from policies.dynamicvla.modeling_vlm_with_expert import VLMWithExpertModel
+
+# Matches ".soNNN", optionally followed by "-something", up to the "_buffer_" marker
+_VARIANT_RE = re.compile(r"\.so\d+(?:-[\w]+)?_buffer_")
+
+
+def canonicalise(k: str) -> str:
+    """
+    Remove dataset-variant markers like '.so100-blue_' or '.so100_' from a
+    normalisation-buffer key.
+    """
+    return _VARIANT_RE.sub(".buffer_", k)
+
+
+def standardise_state_dict(
+    checkpoint: dict[str, torch.Tensor], ref_keys: set[str], *, verbose: bool = True
+) -> tuple[dict[str, torch.Tensor], list[str]]:
+    """
+    • Re-keys `checkpoint ` so that every entry matches the *reference* key set.
+    • If several variant keys collapse to the same canonical name we keep the
+      first one and log the collision.
+    • Returns the new dict + a list of entries that could not be matched.
+    """
+    out, collisions, unmatched = {}, {}, []
+
+    for k, v in checkpoint.items():
+        canon = canonicalise(k)
+        if canon in ref_keys:
+            if canon in out:  # duplicate after collapsing
+                collisions.setdefault(canon, []).append(k)
+            else:
+                out[canon] = v
+        else:
+            unmatched.append(k)
+
+    if verbose:
+        for canon, variants in collisions.items():
+            logging.info(f"[standardise_state_dict] '{canon}'  ←  {variants}")
+        if unmatched:
+            logging.info(
+                f"[standardise_state_dict] kept {len(unmatched)} unmatched keys"
+            )
+
+    out.update({k: checkpoint[k] for k in unmatched})
+    return out, unmatched
+
+
+def rename_checkpoint_keys(checkpoint: dict, rename_str: str):
+    """
+    Renames keys in a checkpoint dictionary based on the given rename string.
+
+    Args:
+        checkpoint (dict): The checkpoint dictionary.
+        rename_str (str): A string specifying key mappings in the format "old1//new1,old2//new2".
+
+    Returns:
+        dict: The modified checkpoint with renamed keys.
+    """
+
+    rename_dict = dict(pair.split("//") for pair in rename_str.split(","))
+
+    new_checkpoint = {}
+    for k, v in checkpoint.items():
+        for old_key, new_key in rename_dict.items():
+            if old_key in k:
+                k = k.replace(old_key, new_key)
+        new_checkpoint[k] = v
+    return new_checkpoint
+
+
+def load_dynamicvla(
+    model: torch.nn.Module,
+    filename: str | os.PathLike,
+    *,
+    device: str = "cpu",
+    checkpoint_keys_mapping: str = "",
+) -> torch.nn.Module:
+    state_dict = safetensors.torch.load_file(filename, device=device)
+
+    # Optional user-supplied renames (e.g. "model._orig_mod.//model.")
+    if checkpoint_keys_mapping and "//" in checkpoint_keys_mapping:
+        state_dict = rename_checkpoint_keys(state_dict, checkpoint_keys_mapping)
+
+    state_dict, _ = standardise_state_dict(state_dict, set(model.state_dict().keys()))
+
+    # HACK(aliberts): to not overwrite normalization parameters as they should come from the dataset
+    norm_keys = ("normalize_inputs", "normalize_targets", "unnormalize_outputs")
+    state_dict = {k: v for k, v in state_dict.items() if not k.startswith(norm_keys)}
+
+    missing, unexpected = model.load_state_dict(state_dict, strict=False)
+    if not all(key.startswith(norm_keys) for key in missing) or unexpected:
+        raise RuntimeError(
+            "DynamicVLA %d missing / %d unexpected keys"
+            % (len(missing), len(unexpected))
+        )
+
+    return model
+
+
+def create_sinusoidal_pos_embedding(
+    time: torch.tensor,
+    dimension: int,
+    min_period: float,
+    max_period: float,
+    device="cpu",
+) -> torch.Tensor:
+    """Computes sine-cosine positional embedding vectors for scalar positions."""
+    if dimension % 2 != 0:
+        raise ValueError(f"dimension ({dimension}) must be divisible by 2")
+
+    if time.ndim != 1:
+        raise ValueError(
+            "The time torch.Tensor is expected to be of shape `(batch_size, )`."
+        )
+
+    dtype = get_safe_dtype(torch.float64, device.type)
+    fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device)
+    period = min_period * (max_period / min_period) ** fraction
+
+    # Compute the outer product
+    scaling_factor = 1.0 / period * 2 * math.pi
+    sin_input = scaling_factor[None, :] * time[:, None]
+    pos_emb = torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1)
+    return pos_emb
+
+
+def make_att_2d_masks(pad_masks, att_masks):
+    """Copied from big_vision.
+
+    Tokens can attend to valid inputs tokens which have a cumulative mask_ar
+    smaller or equal to theirs. This way `mask_ar` int[B, N] can be used to
+    setup several types of attention, for example:
+
+      [[1 1 1 1 1 1]]: pure causal attention.
+
+      [[0 0 0 1 1 1]]: prefix-lm attention. The first 3 tokens can attend between
+          themselves and the last 3 tokens have a causal attention. The first
+          entry could also be a 1 without changing behaviour.
+
+      [[1 0 1 0 1 0 0 1 0 0]]: causal attention between 4 blocks. Tokens of a
+          block can attend all previous blocks and all tokens on the same block.
+
+    Args:
+      input_mask: bool[B, N] true if its part of the input, false if padding.
+      mask_ar: int32[B, N] mask that's 1 where previous tokens cannot depend on
+        it and 0 where it shares the same attention mask as the previous token.
+    """
+    if att_masks.ndim != 2:
+        raise ValueError(att_masks.ndim)
+    if pad_masks.ndim != 2:
+        raise ValueError(pad_masks.ndim)
+
+    cumsum = torch.cumsum(att_masks, dim=1)
+    att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None]
+    pad_2d_masks = pad_masks[:, None, :] * pad_masks[:, :, None]
+    att_2d_masks = att_2d_masks & pad_2d_masks
+    return att_2d_masks
+
+
+def resize_with_pad(img, width, height, pad_value=-1):
+    # assume no-op when width height fits already
+    if img.ndim != 4:
+        raise ValueError(f"(b,c,h,w) expected, but {img.shape}")
+
+    cur_height, cur_width = img.shape[2:]
+
+    ratio = max(cur_width / width, cur_height / height)
+    resized_height = int(cur_height / ratio)
+    resized_width = int(cur_width / ratio)
+    resized_img = F.interpolate(
+        img, size=(resized_height, resized_width), mode="bilinear", align_corners=False
+    )
+
+    pad_height = max(0, int(height - resized_height))
+    pad_width = max(0, int(width - resized_width))
+
+    # pad on left and top of image
+    padded_img = F.pad(resized_img, (pad_width, 0, pad_height, 0), value=pad_value)
+    return padded_img
+
+
+def pad_vector(vector, new_dim):
+    """Can be (batch_size x sequence_length x features_dimension)
+    or (batch_size x features_dimension)
+    """
+    if vector.shape[-1] == new_dim:
+        return vector
+    shape = list(vector.shape)
+    current_dim = shape[-1]
+    shape[-1] = new_dim
+    new_vector = torch.zeros(*shape, dtype=vector.dtype, device=vector.device)
+    new_vector[..., :current_dim] = vector
+    return new_vector
+
+
+def normalize(x, min_val, max_val):
+    return (x - min_val) / (max_val - min_val)
+
+
+def unnormalize(x, min_val, max_val):
+    return x * (max_val - min_val) + min_val
+
+
+def safe_arcsin(value):
+    # This ensures that the input stays within
+    # [−1,1] to avoid invalid values for arcsin
+    return torch.arcsin(torch.clamp(value, -1.0, 1.0))
+
+
+def aloha_gripper_to_angular(value):
+    # Aloha transforms the gripper positions into a linear space. The following code
+    # reverses this transformation to be consistent with dynamicvla which is pretrained in
+    # angular space.
+    #
+    # These values are coming from the Aloha code:
+    # PUPPET_GRIPPER_POSITION_OPEN, PUPPET_GRIPPER_POSITION_CLOSED
+    value = unnormalize(value, min_val=0.01844, max_val=0.05800)
+
+    # This is the inverse of the angular to linear transformation inside the Interbotix code.
+    def linear_to_radian(linear_position, arm_length, horn_radius):
+        value = (horn_radius**2 + linear_position**2 - arm_length**2) / (
+            2 * horn_radius * linear_position
+        )
+        return safe_arcsin(value)
+
+    # The constants are taken from the Interbotix code.
+    value = linear_to_radian(value, arm_length=0.036, horn_radius=0.022)
+
+    # Normalize to [0, 1].
+    # The values 0.4 and 1.5 were measured on an actual Trossen robot.
+    return normalize(value, min_val=0.4, max_val=1.5)
+
+
+def aloha_gripper_from_angular(value):
+    # Convert from the gripper position used by dynamicvla to the gripper position that is used by Aloha.
+    # Note that the units are still angular but the range is different.
+
+    # The values 0.4 and 1.5 were measured on an actual Trossen robot.
+    value = unnormalize(value, min_val=0.4, max_val=1.5)
+
+    # These values are coming from the Aloha code:
+    # PUPPET_GRIPPER_JOINT_OPEN, PUPPET_GRIPPER_JOINT_CLOSE
+    return normalize(value, min_val=-0.6213, max_val=1.4910)
+
+
+def aloha_gripper_from_angular_inv(value):
+    # Directly inverts the gripper_from_angular function.
+    value = unnormalize(value, min_val=-0.6213, max_val=1.4910)
+    return normalize(value, min_val=0.4, max_val=1.5)
+
+
+class DynamicVLAPolicy(PreTrainedPolicy):
+    """Wrapper class around VLAFlowMatching model to train and run inference within LeRobot."""
+
+    config_class = DynamicVLAConfig
+    name = "dynamicvla"
+
+    def __init__(
+        self,
+        config: DynamicVLAConfig,
+        dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
+        ckpt_filename: str | None = None,
+    ):
+        """
+        Args:
+            config: Policy configuration class instance or None, in which case the default instantiation of
+                    the configuration class is used.
+            dataset_stats: Dataset statistics to be used for normalization. If not passed here, it is expected
+                that they will be passed with a call to `load_state_dict` before the policy is used.
+            ckpt_filename: The file path of the pretrained model checkpoint. This is only used when
+                `config.enable_streaming` is set to True, in which case a separate process
+                is spawned to run the VLA model for streaming inference.
+        """
+
+        super().__init__(config)
+        config.validate_features()
+        self.config = config
+        self.normalize_inputs = Normalize(
+            config.input_features, config.normalization_mapping, dataset_stats
+        )
+        self.normalize_targets = Normalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+        self.unnormalize_outputs = Unnormalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+        self.model = VLAFlowMatching(config)
+        self.language_tokenizer = self.model.vlm_with_expert.tokenizer
+        self.reset()
+
+        # ckpt_filename is used to initlialize the streaming process.
+        # The variable is only set once in the get_streaming_model function.
+        if config.enable_streaming and ckpt_filename is not None:
+            # The initialization of the wrapper process of _inference_loop
+            ctx = mp.get_context("spawn")
+            self.q_in = ctx.Manager().dict()
+            self.q_out = ctx.Queue(maxsize=1)
+            self.worker = ctx.Process(
+                target=self._inference_loop,
+                args=(ckpt_filename, config, self.q_in, self.q_out),
+            )
+            self.worker.daemon = True
+            self.worker.start()
+            # Wait for the VLA model to be initialized
+            _ = self.q_out.get()
+            assert "initialized" in _
+
+    def reset(self):
+        """This should be called whenever the environment is reset."""
+        self._queues = {ACTION: deque(maxlen=self.config.n_action_steps)}
+        self._obs_index = 0
+        self._action_index = 0
+        if hasattr(self, "q_in"):
+            self.q_in.clear()  # Clear the input queue
+        if hasattr(self, "q_out") and not self.q_out.empty():
+            self.q_out.get_nowait()  # Clear the output queue
+
+    @staticmethod
+    def get_streaming_model(pretrained_model: str, vla_cfg: DynamicVLAConfig):
+        wrapper = DynamicVLAPolicy.from_pretrained(
+            pretrained_model, config=vla_cfg, ckpt_filename=pretrained_model
+        )
+        # Remove unnecessary components to reduce VRAM
+        del wrapper.language_tokenizer, wrapper.model
+        torch.cuda.empty_cache()
+
+        return wrapper
+
+    @staticmethod
+    @torch.no_grad()
+    def _inference_loop(
+        pretrained_model: str, vla_cfg: DynamicVLAConfig, q_in: dict, q_out: mp.Queue
+    ):
+        logging.basicConfig(
+            level=logging.INFO, format="[%(levelname)s] %(asctime)s %(message)s"
+        )
+        # Initialization
+        vla_model = DynamicVLAPolicy.from_pretrained(pretrained_model, config=vla_cfg)
+        vla_model.eval()
+        if torch.cuda.is_available():
+            vla_model = vla_model.cuda()
+        # Warm-up (to accelerate the first inference)
+        dummy_batch = {
+            k: torch.zeros(
+                1,
+                vla_cfg.n_obs_steps,
+                *vla_cfg.input_features[k].shape,
+                dtype=torch.float32,
+                device="cuda" if torch.cuda.is_available() else "cpu",
+            )
+            for k in vla_cfg.input_features
+        }
+        dummy_batch["task"] = ["dummy text input"]
+        vla_model._get_action_chunk(dummy_batch)
+        q_out.put({"initialized": True})
+
+        # The streaming inference loop
+        while True:
+            try:
+                latest_obs = q_in.get("obs")
+            except:
+                latest_obs = None
+
+            if latest_obs is None:
+                continue
+
+            q_in.clear()
+            noise = latest_obs[1].cuda() if latest_obs[1] is not None else None
+            batch = {}
+            for k, v in latest_obs[0].items():
+                if isinstance(v, torch.Tensor):
+                    batch[k] = v.cuda() if torch.cuda.is_available() else v
+                else:
+                    batch[k] = v
+
+            index = batch["index"]
+            latest_state = batch[OBS_STATE][:, -1:, :]
+            batch = vla_model._prepare_batch(batch)
+
+            actions = vla_model._get_action_chunk(batch, noise)
+            if vla_model.config.use_delta_action:
+                action_dim = actions.shape[-1] - 1
+                actions[..., :action_dim] += latest_state[..., :action_dim]
+
+            if q_out.full():
+                logging.warning("The output queue is full. Skipping an action.")
+                continue
+
+            # NOTE: All torch.Tensors are on CPU if streaming is enabled. Because IPC with
+            #       CUDA torch.Tensors is not supported.
+            q_out.put_nowait({"actions": actions.transpose(0, 1).cpu(), "index": index})
+
+    @classmethod
+    def _load_as_safetensor(
+        cls,
+        model: "DynamicVLAPolicy",
+        model_file: str,
+        map_location: str,
+        strict: bool,
+    ):
+        safetensors.torch.load_model(
+            model, model_file, strict=strict, device=map_location
+        )
+        return load_dynamicvla(
+            model,
+            model_file,
+            device=map_location,
+            checkpoint_keys_mapping="model._orig_mod.//model.",
+        )
+
+    def get_optim_params(self) -> dict:
+        return self.parameters()
+
+    def _get_action_chunk(
+        self, batch: dict[str, torch.Tensor], noise: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        tick = time.perf_counter()
+        for k in batch:
+            if k in self._queues and k != ACTION:
+                batch[k] = torch.stack(list(self._queues[k]), dim=1)
+
+        images, img_masks = self.prepare_images(batch)
+        state = self.prepare_state(batch)
+        lang_tokens, lang_masks = self.prepare_language(batch)
+
+        actions = self.model.sample_actions(
+            images, img_masks, lang_tokens, lang_masks, state, noise=noise
+        )
+
+        # Unpad actions
+        original_action_dim = self.config.action_feature.shape[0]
+        actions = actions[:, :, :original_action_dim]
+        actions = self.unnormalize_outputs({ACTION: actions})[ACTION]
+
+        if self.config.adapt_to_pi_aloha:
+            actions = self._pi_aloha_encode_actions(actions)
+
+        inference_time = time.perf_counter() - tick
+        if "dt_scale" in batch and batch["dt_scale"] > 1.0:
+            # IMPORTANT: To align the inference time with the simulation time
+            sleep_time = inference_time * (batch["dt_scale"] - 1)
+            logging.info(
+                "[Step%03d] Inference Time: %.4fs; Sleep Time: %.4fs"
+                % (batch["index"], inference_time, sleep_time)
+            )
+            time.sleep(sleep_time)
+
+        return actions
+
+    def _prepare_batch(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+        if self.config.adapt_to_pi_aloha:
+            batch[OBS_STATE] = self._pi_aloha_decode_state(batch[OBS_STATE])
+
+        return self.normalize_inputs(batch)
+
+    @torch.no_grad()
+    def predict_action_chunk(
+        self, batch: dict[str, torch.Tensor], noise: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        self.eval()
+        batch = self._prepare_batch(batch)
+        self._queues = populate_queues(self._queues, batch, exclude_keys=[ACTION])
+
+        actions = self._get_action_chunk(batch, noise)
+        return actions
+
+    @torch.no_grad()
+    def select_action(
+        self, batch: dict[str, torch.Tensor], noise: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        if self.config.enable_streaming:
+            return self._get_streaming_action(batch, noise)
+        else:
+            return self._get_non_streaming_action(batch, noise)
+
+    @torch.no_grad()
+    def _get_streaming_action(
+        self, batch: dict[str, torch.Tensor], noise: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        # NOTE: This function does not do any GPU computation.
+        assert "index" in batch
+        # Put the latest observation into the input dict
+        self.q_in.update({"obs": (batch, noise)})
+
+        actions = None
+        if not self.q_out.empty():
+            actions = self.q_out.get_nowait()
+
+        # Merge actions into the queue
+        if actions is not None:
+            assert actions["actions"].size(0) == self.config.n_action_steps
+            skip_n_actions = batch["index"] - actions["index"]
+            logging.debug(
+                "Curr. Step: %03d; Act. Step: %03d; Skip Steps: %03d"
+                % (batch["index"], actions["index"], skip_n_actions)
+            )
+
+            actions["actions"] = actions["actions"][skip_n_actions:]
+            actions["index"] += skip_n_actions
+
+            prev_action_chunk = list(self._queues[ACTION])
+            curr_action_chunk = [
+                {"index": actions["index"] + i, "action": a}
+                for i, a in enumerate(actions["actions"])
+            ]
+            if not prev_action_chunk:
+                # The action queue is empty
+                self._queues[ACTION].extend(
+                    [a for a in curr_action_chunk if a["index"] > self._action_index]
+                )
+            else:
+                self._queues[ACTION].clear()
+                prev_index_start = prev_action_chunk[0]["index"]
+                prev_index_end = prev_action_chunk[-1]["index"]
+                curr_index_start = curr_action_chunk[0]["index"]
+                if curr_index_start > prev_index_end:
+                    self._queues[ACTION].extend(curr_action_chunk)
+                elif curr_index_start > prev_index_start:
+                    keeplen = curr_index_start - prev_index_end
+                    self._queues[ACTION].extend(
+                        prev_action_chunk[:keeplen] + curr_action_chunk
+                    )
+                elif curr_index_start <= prev_index_start:
+                    droplen = prev_index_start - curr_index_start
+                    self._queues[ACTION].extend(curr_action_chunk[droplen:])
+
+        if len(self._queues[ACTION]) == 0:
+            return None
+        else:
+            action = self._queues[ACTION].popleft()
+            self._action_index = action["index"]
+            return action["action"]
+
+    @torch.no_grad()
+    def _get_non_streaming_action(
+        self, batch: dict[str, torch.Tensor], noise: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        self.eval()
+        # Save the state before normalization
+        latest_state = batch[OBS_STATE][:, -1:, :]
+
+        batch = self._prepare_batch(batch)
+        self._queues = populate_queues(self._queues, batch, exclude_keys=[ACTION])
+        # Action queue logic for n_action_steps > 1. When the action_queue is depleted,
+        # populate it by querying the policy.
+        if len(self._queues[ACTION]) == 0:
+            actions = self._get_action_chunk(batch, noise)
+            # `self.predict_action_chunk` returns a (batch_size, n_action_steps, action_dim)
+            # torch.Tensor, but the queue effectively has shape (n_action_steps, batch_size, *),
+            # hence the transpose.
+            if self.config.use_delta_action:
+                action_dim = actions.shape[-1] - 1
+                actions[..., :action_dim] += latest_state[..., :action_dim]
+
+            self._queues[ACTION].extend(
+                actions.transpose(0, 1)[: self.config.n_action_steps]
+            )
+
+        return self._queues[ACTION].popleft()
+
+    def forward(
+        self, batch: dict[str, torch.Tensor], noise=None, time=None
+    ) -> dict[str, torch.Tensor]:
+        """Do a full training forward pass to compute the loss"""
+        if self.config.adapt_to_pi_aloha:
+            batch[OBS_STATE] = self._pi_aloha_decode_state(batch[OBS_STATE])
+            batch[ACTION] = self._pi_aloha_encode_actions_inv(batch[ACTION])
+
+        batch = self.normalize_inputs(batch)
+        batch = self.normalize_targets(batch)
+        images, img_masks = self.prepare_images(batch)
+        state = self.prepare_state(batch)
+        lang_tokens, lang_masks = self.prepare_language(batch)
+
+        actions = self.prepare_action(batch)
+        actions_is_pad = batch.get("actions_id_pad")
+        loss_dict = {}
+        losses = self.model.forward(
+            images, img_masks, lang_tokens, lang_masks, state, actions, noise, time
+        )
+        loss_dict["losses_after_forward"] = losses.clone()
+        if actions_is_pad is not None:
+            in_episode_bound = ~actions_is_pad
+            losses = losses * in_episode_bound.unsqueeze(-1)
+            loss_dict["losses_after_in_ep_bound"] = losses.clone()
+
+        # Remove padding
+        losses = losses[:, :, : self.config.max_action_dim]
+        loss_dict["losses_after_rm_padding"] = losses.clone()
+
+        # For backward pass
+        loss = losses.mean()
+        # For backward pass
+        loss_dict["loss"] = loss.item()
+        return loss, loss_dict
+
+    def prepare_images(self, batch):
+        """Apply DynamicVLA preprocessing to the images, like resizing to 224x224 and padding to keep aspect ratio, and
+        convert pixel range from [0.0, 1.0] to [-1.0, 1.0] as requested by SigLIP.
+        """
+        images = []
+        img_masks = []
+        present_img_keys = [key for key in self.config.image_features if key in batch]
+        missing_img_keys = [
+            key for key in self.config.image_features if key not in batch
+        ]
+
+        if len(present_img_keys) == 0:
+            raise ValueError(
+                "All image features are missing from the batch. At least one expected. "
+                f"(batch: {batch.keys()}) (image_features:{self.config.image_features})"
+            )
+        # Preprocess image features present in the batch
+        for key in present_img_keys:
+            imgs = batch[key][:, None, :, :, :] if batch[key].ndim == 4 else batch[key]
+            b, n, c, h, w = imgs.shape
+            assert n == self.config.n_obs_steps
+            img = imgs.view(b, n * c, h, w)
+            if self.config.resize_imgs_with_padding is not None:
+                img = resize_with_pad(
+                    img, *self.config.resize_imgs_with_padding, pad_value=0
+                )
+            # Normalize from range [0,1] to [-1,1] as expacted by siglip
+            img = img * 2.0 - 1.0
+            if f"{key}_padding_mask" in batch:
+                mask = batch[f"{key}_padding_mask"].bool()
+            else:
+                mask = torch.ones(img.shape[0], dtype=torch.bool, device=img.device)
+
+            images.append(img)
+            img_masks.append(mask)
+
+        # Create image features not present in the batch
+        # as fully 0 padded images.
+        for num_empty_cameras in range(len(missing_img_keys)):
+            if num_empty_cameras >= self.config.empty_cameras:
+                break
+
+            img = torch.ones_like(img) * -1
+            mask = torch.zeros_like(mask)
+            images.append(img)
+            img_masks.append(mask)
+
+        return images, img_masks
+
+    def prepare_language(self, batch) -> tuple[torch.Tensor, torch.Tensor]:
+        """Tokenize the text input"""
+        device = batch[OBS_STATE].device
+        tasks = batch["task"]
+        if isinstance(tasks, str):
+            tasks = [tasks]
+
+        if len(tasks) == 1:
+            tasks = [tasks[0] for _ in range(batch[OBS_STATE].shape[0])]
+
+        tasks = [task if task.endswith("\n") else f"{task}\n" for task in tasks]
+        tokenized_prompt = self.language_tokenizer.__call__(
+            tasks,
+            padding=self.config.pad_language_to,
+            padding_side="right",
+            max_length=self.config.tokenizer_max_length,
+            return_tensors="pt",
+        )
+        lang_tokens = tokenized_prompt["input_ids"].to(device=device)
+        lang_masks = tokenized_prompt["attention_mask"].to(
+            device=device, dtype=torch.bool
+        )
+        return lang_tokens, lang_masks
+
+    def _pi_aloha_decode_state(self, state):
+        # Flip the joints.
+        for motor_idx in [1, 2, 8, 9]:
+            state[:, motor_idx] *= -1
+        # Reverse the gripper transformation that is being applied by the Aloha runtime.
+        for motor_idx in [6, 13]:
+            state[:, motor_idx] = aloha_gripper_to_angular(state[:, motor_idx])
+
+        return state
+
+    def _pi_aloha_encode_actions(self, actions):
+        # Flip the joints.
+        for motor_idx in [1, 2, 8, 9]:
+            actions[:, :, motor_idx] *= -1
+        # Reverse the gripper transformation that is being applied by the Aloha runtime.
+        for motor_idx in [6, 13]:
+            actions[:, :, motor_idx] = aloha_gripper_from_angular(
+                actions[:, :, motor_idx]
+            )
+        return actions
+
+    def _pi_aloha_encode_actions_inv(self, actions):
+        # Flip the joints again.
+        for motor_idx in [1, 2, 8, 9]:
+            actions[:, :, motor_idx] *= -1
+        # Reverse the gripper transformation that is being applied by the Aloha runtime.
+        for motor_idx in [6, 13]:
+            actions[:, :, motor_idx] = aloha_gripper_from_angular_inv(
+                actions[:, :, motor_idx]
+            )
+        return actions
+
+    def prepare_state(self, batch):
+        """Pad state"""
+        state = (
+            batch[OBS_STATE][:, -1, :]
+            if batch[OBS_STATE].ndim > 2
+            else batch[OBS_STATE]
+        )
+        state = pad_vector(state, self.config.max_state_dim)
+        return state
+
+    def prepare_action(self, batch):
+        """Pad action"""
+        actions = pad_vector(batch[ACTION], self.config.max_action_dim)
+        return actions
+
+
+def pad_tensor(tensor, max_len, pad_value=0):
+    """
+    Efficiently pads a torch.Tensor along sequence dimension to match max_len.
+
+    Args:
+        torch.Tensor (torch.Tensor): Shape (B, L, ...) or (B, L).
+        max_len (int): Fixed sequence length.
+        pad_value (int/float): Value for padding.
+
+    Returns:
+        torch.Tensor: Shape (B, max_len, ...) or (B, max_len).
+    """
+    b, d = torch.Tensor.shape[:2]
+
+    # Create a padded torch.Tensor of max_len and copy the existing values
+    padded_tensor = torch.full(
+        (b, max_len, *tensor.shape[2:]),
+        pad_value,
+        dtype=tensor.dtype,
+        device=tensor.device,
+    )
+    padded_tensor[:, :d] = tensor  # Efficient in-place copy
+
+    return padded_tensor
+
+
+class VLAFlowMatching(torch.nn.Module):
+    """
+    ┌──────────────────────────────┐
+    │                 actions      │
+    │                    ▲         │
+    │ ┌─────────┐      ┌─|────┐    │
+    │ |         │────► │      │    │
+    │ |         │ kv   │      │    │
+    │ |         │────► │Action│    │
+    │ |   VLM   │cache │Expert│    |
+    │ │         │────► |      │    │
+    │ │         │      │      │    │
+    │ └▲──▲───▲─┘      └───▲──┘    |
+    │  │  |   |            │       |
+    │  |  |   |          noise     │
+    │  │  │ state                  │
+    │  │ language tokens           │
+    │  image(s)                    │
+    └──────────────────────────────┘
+    """
+
+    def __init__(self, config: DynamicVLAConfig):
+        super().__init__()
+        self.config = config
+
+        if config.temporal_fusion == "conv":
+            self.mults_proj = torch.nn.Sequential(
+                torch.nn.Conv2d(3 * config.n_obs_steps, 3, kernel_size=7, padding=3),
+                torch.nn.GELU(),
+            )
+            vlm_input_channels = 3
+        elif config.temporal_fusion == "attn":
+            vlm_input_channels = 3 * config.n_obs_steps
+        elif config.temporal_fusion == "flat":
+            vlm_input_channels = 3
+        else:
+            raise ValueError(f"Unknown temporal_fusion: {config.temporal_fusion}")
+
+        self.vlm_with_expert: VLMWithExpertModel = self._get_vlm_with_expert(
+            config, config.vlm_model_name, vlm_input_channels
+        )
+        self.state_proj = torch.nn.Linear(
+            config.max_state_dim,
+            self.vlm_with_expert.vlm_config.text_config.hidden_size,
+        )
+        self.action_in_proj = torch.nn.Linear(
+            config.max_action_dim, self.vlm_with_expert.expert_hidden_size
+        )
+        self.action_out_proj = torch.nn.Linear(
+            self.vlm_with_expert.expert_hidden_size, config.max_action_dim
+        )
+        self.action_time_mlp_in = torch.nn.Linear(
+            self.vlm_with_expert.expert_hidden_size * 2,
+            self.vlm_with_expert.expert_hidden_size,
+        )
+        self.action_time_mlp_out = torch.nn.Linear(
+            self.vlm_with_expert.expert_hidden_size,
+            self.vlm_with_expert.expert_hidden_size,
+        )
+        self._set_requires_grad()
+
+    def _get_vlm_with_expert(
+        self,
+        config: DynamicVLAConfig,
+        vlm_model_name: str,
+        vlm_input_channels: int,
+    ):
+        if vlm_model_name.startswith("HuggingFaceTB/SmolVLM2"):
+            vlm_config = AutoConfig.from_pretrained(vlm_model_name)
+            vlm_config.vision_config.num_channels = vlm_input_channels
+            vlm_config.vision_config.patch_size = config.smolvlm_patch_size
+            vlm_config.vision_config.num_attention_heads = (
+                config.smolvlm_attention_heads
+            )
+            vlm_config.vision_config.hidden_size = config.smolvlm_hidden_size
+            vlm_config.vision_config.intermediate_size = (
+                config.smolvlm_intermediate_size
+            )
+            vlm = SmolVLMForConditionalGeneration(config=vlm_config)
+        elif vlm_model_name.startswith("HuggingFaceTB/SmolLM2"):
+            text_config = AutoConfig.from_pretrained(vlm_model_name)
+            vision_config = FastViTConfig(
+                in_channels=vlm_input_channels,
+                position_embeddings=[
+                    None,
+                    None,
+                    None,
+                    {"name": "RepCPE", "spatial_shape": (7, 7)},
+                    {"name": "RepCPE", "spatial_shape": (7, 7)},
+                ],
+                inference_mode=config.fastvlm_inference_mode,
+            )
+            vlm = FastVLMForConditionalGeneration(
+                config=FastVLMConfig(
+                    text_config=text_config,
+                    vision_config=vision_config,
+                )
+            )
+        elif vlm_model_name.startswith("Qwen/Qwen2"):
+            text_config = AutoConfig.from_pretrained(vlm_model_name)
+            vision_config = FastViTConfig(
+                in_channels=vlm_input_channels,
+                position_embeddings=[
+                    None,
+                    None,
+                    None,
+                    {"name": "RepCPE", "spatial_shape": (7, 7)},
+                    {"name": "RepCPE", "spatial_shape": (7, 7)},
+                ],
+                inference_mode=config.fastvlm_inference_mode,
+            )
+            vlm = FastVLMForConditionalGeneration(
+                config=FastVLMConfig(
+                    text_config=text_config,
+                    vision_config=vision_config,
+                )
+            )
+        else:
+            raise ValueError(f"Unknown VLM: {vlm_model_name}")
+
+        return VLMWithExpertModel(
+            model_id=config.vlm_model_name,
+            vlm=vlm,
+            freeze_vision_model=config.freeze_vision_model,
+            freeze_connector=config.freeze_connector,
+            freeze_text_model=config.freeze_text_model,
+            num_vlm_layers=self.config.num_vlm_layers,
+            num_expert_layers=config.num_expert_layers,
+            num_expert_skip_layers=config.num_expert_skip_layers,
+            attention_mode=config.attention_mode,
+            self_attn_every_n_layers=self.config.self_attn_every_n_layers,
+            expert_width_multiplier=self.config.expert_width_multiplier,
+        )
+
+    def _set_requires_grad(self):
+        for params in self.state_proj.parameters():
+            params.requires_grad = self.config.train_state_proj
+
+    def _sample_noise(self, shape, device, dtype=torch.float32):
+        noise = torch.normal(
+            mean=0.0,
+            std=1.0,
+            size=shape,
+            dtype=dtype,
+            device=device,
+        )
+        return noise
+
+    def _sample_time(self, bsize, device, dtype=torch.float32):
+        beta_dist = torch.distributions.Beta(concentration1=1.5, concentration0=1.0)
+        time_beta = beta_dist.sample((bsize,)).to(device=device, dtype=dtype)
+        time = time_beta * 0.999 + 0.001
+        return time
+
+    def _embed_prefix(
+        self, images, img_masks, lang_tokens, lang_masks, state: torch.Tensor = None
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        embs = []
+        pad_masks = []
+        att_masks = []
+        for img, img_mask in zip(images, img_masks, strict=False):
+            bsize = img.size(0)
+            if self.config.temporal_fusion == "conv":
+                # Temporal fusion of image frames with Conv2d
+                img = self.mults_proj(img)
+            elif self.config.temporal_fusion == "flat":
+                # Flatten temporal dimension into batch dimension
+                img = img.view(-1, 3, img.shape[2], img.shape[3])
+
+            img_emb = self.vlm_with_expert.embed_image(img)
+            img_emb_dim = img_emb.size(-1)
+            if self.config.temporal_fusion == "flat":
+                # Reshape back to (batch_size, n_obs_steps * 3, emb_dim)
+                img_emb = img_emb.view(bsize, -1, img_emb_dim)
+
+            # Normalize image embeddings
+            img_emb = img_emb * torch.tensor(
+                img_emb_dim**0.5, dtype=img_emb.dtype, device=img_emb.device
+            )
+            num_img_embs = img_emb.size(1)
+            img_mask = img_mask[:, None].expand(bsize, num_img_embs)
+            embs.append(img_emb)
+            pad_masks.append(img_mask)
+            att_masks += [0] * (num_img_embs)
+
+        lang_emb = self.vlm_with_expert.embed_language_tokens(lang_tokens)
+        # Normalize language embeddings
+        lang_emb_dim = lang_emb.shape[-1]
+        lang_emb = lang_emb * math.sqrt(lang_emb_dim)
+        embs.append(lang_emb)
+        pad_masks.append(lang_masks)
+        num_lang_embs = lang_emb.shape[1]
+        att_masks += [0] * num_lang_embs
+
+        if state is not None:
+            state_emb = self.state_proj(state)
+            state_emb = state_emb[:, None, :] if state_emb.ndim == 2 else state_emb
+            embs.append(state_emb)
+            states_seq_len = state_emb.shape[1]
+            state_mask = torch.ones(
+                state_emb.shape[0],
+                states_seq_len,
+                dtype=torch.bool,
+                device=state_emb.device,
+            )
+            pad_masks.append(state_mask)
+            # Set attention masks so that image and language inputs do not attend to state or actions
+            att_masks += [1] * (states_seq_len)
+
+        embs = torch.cat(embs, dim=1)
+        pad_masks = torch.cat(pad_masks, dim=1)
+        att_masks = torch.tensor(att_masks, dtype=torch.bool, device=pad_masks.device)
+        att_masks = att_masks[None, :]
+
+        seq_len = pad_masks.shape[1]
+        if seq_len < self.config.prefix_length:
+            embs = pad_tensor(embs, self.config.prefix_length, pad_value=0)
+            pad_masks = pad_tensor(pad_masks, self.config.prefix_length, pad_value=0)
+            att_masks = pad_tensor(att_masks, self.config.prefix_length, pad_value=0)
+
+        att_masks = att_masks.expand(bsize, -1)
+        return embs, pad_masks, att_masks
+
+    def _embed_suffix(self, noisy_actions, timestep):
+        embs = []
+        pad_masks = []
+        att_masks = []
+
+        # Fuse timestep + action information using an MLP
+        action_emb = self.action_in_proj(noisy_actions)
+        device = action_emb.device
+        dtype = action_emb.dtype
+        # Embed timestep using sine-cosine positional encoding with sensitivity in the range [0, 1]
+        time_emb = create_sinusoidal_pos_embedding(
+            timestep,
+            self.vlm_with_expert.expert_hidden_size,
+            self.config.min_period,
+            self.config.max_period,
+            device=device,
+        )
+        time_emb = time_emb.type(dtype=dtype)
+
+        time_emb = time_emb[:, None, :].expand_as(action_emb)
+        action_time_emb = torch.cat([action_emb, time_emb], dim=2)
+
+        action_time_emb = self.action_time_mlp_in(action_time_emb)
+        action_time_emb = F.silu(action_time_emb)  # swish == silu
+        action_time_emb = self.action_time_mlp_out(action_time_emb)
+
+        # Add to input tokens
+        embs.append(action_time_emb)
+
+        bsize, action_time_dim = action_time_emb.shape[:2]
+        action_time_mask = torch.ones(
+            bsize, action_time_dim, dtype=torch.bool, device=device
+        )
+        pad_masks.append(action_time_mask)
+
+        # Set attention masks so that image, language and state inputs do not attend to action tokens
+        att_masks += [1] * self.config.chunk_size
+        embs = torch.cat(embs, dim=1)
+        pad_masks = torch.cat(pad_masks, dim=1)
+        att_masks = torch.tensor(att_masks, dtype=embs.dtype, device=embs.device)
+        att_masks = att_masks[None, :].expand(bsize, len(att_masks))
+        return embs, pad_masks, att_masks
+
+    def _get_position_ids(
+        self,
+        prefix_offsets: torch.Tensor | None,
+        pad_masks: torch.Tensor,
+    ) -> torch.Tensor:
+        position_ids = torch.cumsum(pad_masks, dim=1) - 1
+        if prefix_offsets is not None:
+            position_ids += prefix_offsets
+
+        return position_ids
+
+    def forward(
+        self,
+        images,
+        img_masks,
+        lang_tokens,
+        lang_masks,
+        state,
+        actions,
+        noise=None,
+        time=None,
+    ) -> torch.Tensor:
+        """Do a full training forward pass and compute the loss (batch_size x num_steps x num_motors)"""
+        if noise is None:
+            noise = self._sample_noise(actions.shape, actions.device)
+
+        if time is None:
+            time = self._sample_time(actions.shape[0], actions.device)
+
+        time_expanded = time[:, None, None]
+        x_t = time_expanded * noise + (1 - time_expanded) * actions
+        u_t = noise - actions
+        prefix_embs, prefix_pad_masks, prefix_att_masks = self._embed_prefix(
+            images, img_masks, lang_tokens, lang_masks, state=state
+        )
+        suffix_embs, suffix_pad_masks, suffix_att_masks = self._embed_suffix(x_t, time)
+
+        pad_masks = torch.cat([prefix_pad_masks, suffix_pad_masks], dim=1)
+        att_masks = torch.cat([prefix_att_masks, suffix_att_masks], dim=1)
+        att_2d_masks = make_att_2d_masks(pad_masks, att_masks)
+        position_ids = self._get_position_ids(None, pad_masks)
+
+        (_, suffix_out), _ = self.vlm_with_expert(
+            attention_mask=att_2d_masks,
+            position_ids=position_ids,
+            past_key_values=None,
+            inputs_embeds=[prefix_embs, suffix_embs],
+            use_cache=False,
+            fill_kv_cache=False,
+        )
+        suffix_out = suffix_out[:, -self.config.chunk_size :]
+        # Original openpi code, upcast attention output
+        suffix_out = suffix_out.to(dtype=torch.float32)
+        v_t = self.action_out_proj(suffix_out)
+        losses = F.mse_loss(u_t, v_t, reduction="none")
+        return losses
+
+    def sample_vlm_embedding(
+        self, images, img_masks, lang_tokens, lang_masks, state
+    ) -> torch.Tensor:
+        """Do a half inference forward and compute the VLM embedding"""
+
+        prefix_embs, prefix_pad_masks, prefix_att_masks = self._embed_prefix(
+            images, img_masks, lang_tokens, lang_masks, state
+        )
+        prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks)
+        prefix_position_ids = self._get_position_ids(None, prefix_pad_masks)
+        # Compute image and language key value cache
+        _, past_key_values = self.vlm_with_expert.forward(
+            attention_mask=prefix_att_2d_masks,
+            position_ids=prefix_position_ids,
+            past_key_values=None,
+            inputs_embeds=[prefix_embs, None],
+            use_cache=self.config.use_cache,
+            fill_kv_cache=True,
+        )
+        return prefix_pad_masks, past_key_values
+
+    def sample_actions(
+        self, images, img_masks, lang_tokens, lang_masks, state, noise=None
+    ) -> torch.Tensor:
+        """Do a full inference forward and compute the action (batch_size x num_steps x num_motors)"""
+        bsize = state.shape[0]
+        device = state.device
+        if noise is None:
+            actions_shape = (bsize, self.config.chunk_size, self.config.max_action_dim)
+            noise = self._sample_noise(actions_shape, device)
+
+        prefix_pad_masks, past_key_values = self.sample_vlm_embedding(
+            images, img_masks, lang_tokens, lang_masks, state
+        )
+        dt = -1.0 / self.config.num_steps
+        dt = torch.tensor(dt, dtype=torch.float32, device=device)
+
+        x_t = noise
+        time = torch.tensor(1.0, dtype=torch.float32, device=device)
+        while time >= -dt / 2:
+            expanded_time = time.expand(bsize)
+            v_t = self.denoise_step(
+                prefix_pad_masks,
+                past_key_values,
+                x_t,
+                expanded_time,
+            )
+            # Euler step
+            x_t += dt * v_t
+            time += dt
+
+        return x_t
+
+    def denoise_step(
+        self,
+        prefix_pad_masks,
+        past_key_values,
+        x_t,
+        timestep,
+    ):
+        """Apply one denoising step of the noise `x_t` at a given timestep."""
+        suffix_embs, suffix_pad_masks, suffix_att_masks = self._embed_suffix(
+            x_t, timestep
+        )
+
+        suffix_len = suffix_pad_masks.shape[1]
+        batch_size = prefix_pad_masks.shape[0]
+        prefix_len = prefix_pad_masks.shape[1]
+        prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand(
+            batch_size, suffix_len, prefix_len
+        )
+
+        suffix_att_2d_masks = make_att_2d_masks(suffix_pad_masks, suffix_att_masks)
+        full_att_2d_masks = torch.cat([prefix_pad_2d_masks, suffix_att_2d_masks], dim=2)
+        prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None]
+        position_ids = self._get_position_ids(prefix_offsets, suffix_pad_masks)
+
+        outputs_embeds, _ = self.vlm_with_expert.forward(
+            attention_mask=full_att_2d_masks,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=[None, suffix_embs],
+            use_cache=self.config.use_cache,
+            fill_kv_cache=False,
+        )
+        suffix_out = outputs_embeds[1]
+        suffix_out = suffix_out[:, -self.config.chunk_size :]
+        suffix_out = suffix_out.to(dtype=torch.float32)
+        v_t = self.action_out_proj(suffix_out)
+        return v_t

policies/dynamicvla/modeling_fastvlm.py

@@ -0,0 +1,1635 @@
+# -*- coding: utf-8 -*-
+#
+# @File:   modeling_fastvlm.py
+# @Author: Haozhe Xie
+# @Date:   2025-09-10 20:19:11
+# @Last Modified by: Haozhe Xie
+# @Last Modified at: 2026-02-26 10:12:16
+# @Email:  root@haozhexie.com
+
+import functools
+import logging
+import typing
+
+import torch
+import torch.nn.functional as F
+from timm.layers import DropPath, SqueezeExcite
+from transformers import (
+    AutoModel,
+    GenerationMixin,
+    LlamaConfig,
+    LlamaModel,
+    PretrainedConfig,
+    PreTrainedModel,
+    Qwen2Config,
+)
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.modeling_outputs import ModelOutput
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs
+
+
+class KwargsForCausalLM(FlashAttentionKwargs, TransformersKwargs): ...
+
+
+class FastVLMBaseModelOutputWithPast(ModelOutput):
+    last_hidden_state: typing.Optional[torch.FloatTensor] = None
+    past_key_values: typing.Optional[Cache] = None
+    hidden_states: typing.Optional[tuple[torch.FloatTensor]] = None
+    attentions: typing.Optional[tuple[torch.FloatTensor]] = None
+    image_hidden_states: typing.Optional[tuple[torch.FloatTensor]] = None
+
+
+class FastVLMCausalLMOutputWithPast(ModelOutput):
+    loss: typing.Optional[torch.FloatTensor] = None
+    logits: typing.Optional[torch.FloatTensor] = None
+    past_key_values: typing.Optional[Cache] = None
+    hidden_states: typing.Optional[tuple[torch.FloatTensor]] = None
+    attentions: typing.Optional[tuple[torch.FloatTensor]] = None
+    image_hidden_states: typing.Optional[tuple[torch.FloatTensor]] = None
+
+
+class FastViTConfig(PretrainedConfig):
+    model_type = "fastvit"
+    base_config_key = "vision_config"
+
+    def __init__(
+        self,
+        in_channels=3,
+        out_channels=768,
+        image_size=1024,
+        patch_size=64,
+        n_blocks=[2, 12, 24, 4, 2],
+        embed_dims=[96, 192, 384, 768, 1536],
+        mlp_ratios=[4, 4, 4, 4, 4],
+        downsample=[True, True, True, True, True],
+        downsample_patch_size=7,
+        downsample_stride=2,
+        downsample_use_se=None,
+        position_embeddings=None,
+        token_mixers=("repmixer", "repmixer", "repmixer", "attention", "attention"),
+        repmixer_kernel_size=3,
+        use_scale_branch=True,
+        use_layer_scale=True,
+        drop_path_rate=0.0,
+        layer_scale_init_value=1e-5,
+        clsss_ratio=2,
+        inference_mode=False,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.n_blocks = n_blocks
+        self.embed_dims = embed_dims
+        self.mlp_ratios = mlp_ratios
+        self.downsample = downsample
+        self.downsample_patch_size = downsample_patch_size
+        self.downsample_stride = downsample_stride
+        self.downsample_use_se = downsample_use_se
+        self.position_embeddings = position_embeddings
+        self.token_mixers = token_mixers
+        self.repmixer_kernel_size = repmixer_kernel_size
+        self.use_scale_branch = use_scale_branch
+        self.use_layer_scale = use_layer_scale
+        self.drop_path_rate = drop_path_rate
+        self.layer_scale_init_value = layer_scale_init_value
+        self.clsss_ratio = clsss_ratio
+        self.inference_mode = inference_mode
+
+
+class FastVLMConfig(PretrainedConfig):
+
+    model_type = "fastvlm"
+    sub_configs = {"text_config": LlamaConfig | Qwen2Config, "vision_config": FastViTConfig}
+
+    def __init__(
+        self,
+        use_cache=True,
+        tie_word_embeddings=False,
+        max_position_embeddings=None,
+        text_config=None,
+        vision_config=None,
+        image_token_id=128_257,
+        pad_token_id=128_002,
+        **kwargs,
+    ) -> None:
+        # ViTConfig
+        if vision_config is None:
+            self.vision_config = FastViTConfig()
+        elif isinstance(vision_config, dict):
+            self.vision_config = FastViTConfig(**vision_config)
+        elif isinstance(vision_config, FastViTConfig):
+            self.vision_config = vision_config
+        else:
+            raise ValueError("No valid vision_config is provided.")
+
+        # LlamaConfig
+        if text_config is None:
+            self.text_config = LlamaConfig()
+        elif isinstance(text_config, dict):
+            self.text_config = LlamaConfig(**text_config)
+        elif isinstance(text_config, LlamaConfig) or isinstance(text_config, Qwen2Config):
+            self.text_config = text_config
+        else:
+            raise ValueError("No valid text_config is provided.")
+
+        self.use_cache = use_cache
+        self.image_token_id = image_token_id
+        self.tie_word_embeddings = tie_word_embeddings
+        if max_position_embeddings is not None:
+            self.text_config.max_position_embeddings = max_position_embeddings
+
+        super().__init__(
+            **kwargs, pad_token_id=pad_token_id, tie_word_embeddings=tie_word_embeddings
+        )
+
+
+class FastVLMPreTrainedModel(PreTrainedModel):
+    config_class = FastVLMConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _skip_keys_device_placement = "past_key_values"
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+    _supports_cache_class = True
+    _supports_attention_backend = True
+
+    def _init_weights(self, module):
+        std = getattr(
+            self.config,
+            "initializer_range",
+            self.config.get_text_config().initializer_range,
+        )
+
+        if isinstance(module, (torch.nn.Linear, torch.nn.Conv2d)):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, torch.nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, torch.nn.LayerNorm):
+            module.weight.data.fill_(1.0)
+            module.bias.data.zero_()
+
+
+class FastVLMForConditionalGeneration(FastVLMPreTrainedModel, GenerationMixin):
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = FastVLMModel(config)
+        self.lm_head = torch.nn.Linear(
+            config.text_config.hidden_size, config.text_config.vocab_size, bias=False
+        )
+        self.vocab_size = config.text_config.vocab_size
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def enable_input_require_grads(self):
+        """
+        Enables the gradients for the input embeddings. This is useful for fine-tuning
+        adapter weights while keeping the model weights fixed.
+        """
+
+        def make_inputs_require_grads(module, input, output):
+            output.requires_grad_(True)
+
+        self._text_require_grads_hook = (
+            self.get_input_embeddings().register_forward_hook(make_inputs_require_grads)
+        )
+        self._vision_require_grads_hook = (
+            self.model.vision_model.get_input_embeddings().register_forward_hook(
+                make_inputs_require_grads
+            )
+        )
+
+    def disable_input_require_grads(self):
+        self._text_require_grads_hook.remove()
+        self._vision_require_grads_hook.remove()
+
+    def get_input_embeddings(self):
+        return self.model.text_model.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.model.text_model.set_input_embeddings(value)
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def forward(
+        self,
+        input_ids: typing.Optional[torch.LongTensor] = None,
+        attention_mask: typing.Optional[torch.Tensor] = None,
+        position_ids: typing.Optional[torch.LongTensor] = None,
+        past_key_values: typing.Optional[typing.List[torch.FloatTensor]] = None,
+        inputs_embeds: typing.Optional[torch.FloatTensor] = None,
+        pixel_values: typing.Optional[torch.FloatTensor] = None,
+        pixel_attention_mask: typing.Optional[torch.BoolTensor] = None,
+        image_hidden_states: typing.Optional[torch.FloatTensor] = None,
+        labels: typing.Optional[torch.LongTensor] = None,
+        output_attentions: typing.Optional[bool] = None,
+        output_hidden_states: typing.Optional[bool] = None,
+        use_cache: typing.Optional[bool] = None,
+        cache_position: typing.Optional[torch.LongTensor] = None,
+        return_dict: typing.Optional[bool] = None,
+        logits_to_keep: typing.Union[int, torch.Tensor] = 0,
+        **kwargs: Unpack[KwargsForCausalLM],
+    ) -> typing.Union[typing.Tuple, FastVLMCausalLMOutputWithPast]:
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            pixel_values=pixel_values,
+            pixel_attention_mask=pixel_attention_mask,
+            image_hidden_states=image_hidden_states,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            return_dict=return_dict,
+            **kwargs,
+        )
+
+        hidden_states = outputs[0]
+        # Only compute necessary logits, and do not upcast them to float if we are not
+        # computing the loss
+        slice_indices = (
+            slice(-logits_to_keep, None)
+            if isinstance(logits_to_keep, int)
+            else logits_to_keep
+        )
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                vocab_size=self.config.text_config.vocab_size,
+                **kwargs,
+            )
+
+        return FastVLMCausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            image_hidden_states=outputs.image_hidden_states,
+        )
+
+
+class FastVLMModel(FastVLMPreTrainedModel):
+    def __init__(self, config: FastVLMConfig) -> None:
+        super().__init__(config)
+        self.vision_model = FastViT(config.vision_config)
+        self.connector = FastVLMConnector(config)
+        self.text_model: LlamaModel = AutoModel.from_config(config.text_config)
+        self.post_init()
+
+    def enable_input_require_grads(self):
+        """
+        Enables the gradients for the input embeddings.
+
+        This is useful for lora when using gradient checkpointing.
+        c.f. https://github.com/huggingface/peft/issues/1402#issuecomment-1913675032
+
+        Override to set output.requires_grad = True for both the decoder's and vision model's
+        embeddings.
+        """
+
+        def get_lowest_module(module):
+            if len(list(module.children())) == 0:
+                # If the module has no children, it is a leaf module (e.g., Linear, Conv2d, etc.)
+                return module
+            else:
+                # Recursively call the function on each child module
+                return get_lowest_module(list(module.children())[0])
+
+        def make_inputs_require_grads(module, input, output):
+            output.requires_grad_(True)
+
+        self._text_require_grads_hook = (
+            self.get_input_embeddings().register_forward_hook(make_inputs_require_grads)
+        )
+        self._vision_require_grads_hook = get_lowest_module(
+            self.vision_model
+        ).register_forward_hook(make_inputs_require_grads)
+
+    def disable_input_require_grads(self) -> None:
+        self._text_require_grads_hook.remove()
+        self._vision_require_grads_hook.remove()
+
+    def get_input_embeddings(self) -> torch.nn.Module:
+        return self.text_model.get_input_embeddings()
+
+    def set_input_embeddings(self, value: torch.nn.Module) -> None:
+        self.text_model.set_input_embeddings(value)
+
+    def get_image_features(self, pixel_values: torch.FloatTensor):
+        batch_size, num_images, _, _, _ = pixel_values.shape
+        pixel_values = pixel_values.view(
+            batch_size * num_images, *pixel_values.shape[2:]
+        )
+        # Modality projection & resampling
+        image_features = self.vision_model(pixel_values).flatten(2).transpose(1, 2)
+        image_features = self.connector(image_features)
+        return image_features
+
+    def inputs_merger(
+        self,
+        input_ids: torch.LongTensor,
+        inputs_embeds: torch.Tensor,
+        image_hidden_states: torch.Tensor,
+    ):
+        _, patch_size, _ = image_hidden_states.shape
+        image_mask = input_ids == self.config.image_token_id
+        num_image_tokens = image_mask.sum(dim=1)
+
+        if not torch.all(num_image_tokens % patch_size == 0):
+            raise ValueError(
+                "At least one sample has <image> tokens not divisible by patch_size."
+            )
+
+        blocks_per_sample = num_image_tokens // patch_size
+        offsets = torch.nn.functional.pad(
+            blocks_per_sample.cumsum(dim=0), (1, 0), value=0
+        )
+        block_offset = offsets[:-1]
+        row_cum = image_mask.cumsum(dim=-1)
+        chunk_idx = (row_cum - 1) // patch_size
+        local_idx = (row_cum - 1) % patch_size
+        block_idx = block_offset.unsqueeze(1) + chunk_idx
+
+        image_embeds = torch.zeros_like(inputs_embeds)
+        image_embeds[image_mask] = image_hidden_states[
+            block_idx[image_mask], local_idx[image_mask], :
+        ]
+
+        merged_embeds = torch.where(
+            image_mask.unsqueeze(-1), image_embeds, inputs_embeds
+        )
+        return merged_embeds
+
+    def forward(
+        self,
+        input_ids: typing.Optional[torch.LongTensor] = None,
+        attention_mask: typing.Optional[torch.Tensor] = None,
+        position_ids: typing.Optional[torch.LongTensor] = None,
+        past_key_values: typing.Optional[list[torch.FloatTensor]] = None,
+        inputs_embeds: typing.Optional[torch.FloatTensor] = None,
+        pixel_values: typing.Optional[torch.FloatTensor] = None,
+        image_hidden_states: typing.Optional[torch.FloatTensor] = None,
+        output_attentions: typing.Optional[bool] = None,
+        output_hidden_states: typing.Optional[bool] = None,
+        use_cache: typing.Optional[bool] = None,
+        cache_position: typing.Optional[torch.LongTensor] = None,
+        return_dict: typing.Optional[bool] = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> FastVLMBaseModelOutputWithPast:
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        if self.training and self.text_model.gradient_checkpointing and use_cache:
+            logging.warning(
+                "`use_cache=True` is incompatible with gradient checkpointing. "
+                "Setting `use_cache=False`..."
+            )
+            use_cache = False
+
+        # retrieve input_ids and inputs_embeds
+        if input_ids is not None:
+            _, _ = input_ids.shape
+        elif inputs_embeds is not None:
+            _, _, _ = inputs_embeds.shape
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        past_seen_tokens = 0
+        if use_cache:
+            if past_key_values is None:
+                past_key_values = DynamicCache()
+
+            past_seen_tokens = (
+                past_key_values.get_seq_length() if past_key_values is not None else 0
+            )
+
+        if inputs_embeds is not None and input_ids is None and past_seen_tokens == 0:
+            raise ValueError(
+                "When first calling the model, if input_embeds are passed, input_ids "
+                "should not be None."
+            )
+
+        if inputs_embeds is None:
+            inputs_embeds = self.text_model.get_input_embeddings()(input_ids).to(
+                input_ids.device
+            )
+
+        # START VISUAL INPUTS INTEGRATION
+        if pixel_values is not None and image_hidden_states is not None:
+            raise ValueError(
+                "You cannot specify both pixel_values and image_hidden_states at the "
+                "same time"
+            )
+        elif pixel_values is not None:
+            image_hidden_states = self.get_image_features(pixel_values).to(
+                dtype=self.dtype, device=input_ids.device
+            )
+        elif image_hidden_states is not None:
+            image_hidden_states = image_hidden_states.to(
+                dtype=self.dtype, device=input_ids.device
+            )
+
+        if inputs_embeds is not None and image_hidden_states is not None:
+            inputs_embeds = self.inputs_merger(
+                input_ids=input_ids,
+                inputs_embeds=inputs_embeds,
+                image_hidden_states=image_hidden_states,
+            )
+
+        outputs = self.text_model(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            **kwargs,
+        )
+        return FastVLMBaseModelOutputWithPast(
+            last_hidden_state=outputs.last_hidden_state,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            image_hidden_states=image_hidden_states,
+        )
+
+
+class FastViT(torch.nn.Module):
+    """The FastViT model. <https://arxiv.org/pdf/2303.14189.pdf>"""
+
+    def __init__(self, config: FastViTConfig) -> None:
+        super().__init__()
+        n_stages = len(config.n_blocks)
+        if config.position_embeddings is None:
+            config.position_embeddings = [None] * n_stages
+        if config.downsample_use_se is None:
+            config.downsample_use_se = [False] * n_stages
+
+        # Stem
+        self.patch_embedding = torch.nn.Sequential(
+            MobileOneBlock(
+                in_channels=config.in_channels,
+                out_channels=config.embed_dims[0],
+                kernel_size=3,
+                stride=2,
+                padding=1,
+                groups=1,
+                inference_mode=config.inference_mode,
+                use_se=False,
+                num_conv_branches=1,
+                use_scale_branch=config.use_scale_branch,
+            ),
+            MobileOneBlock(
+                in_channels=config.embed_dims[0],
+                out_channels=config.embed_dims[0],
+                kernel_size=3,
+                stride=2,
+                padding=1,
+                groups=config.embed_dims[0],
+                inference_mode=config.inference_mode,
+                use_se=False,
+                num_conv_branches=1,
+                use_scale_branch=config.use_scale_branch,
+            ),
+            MobileOneBlock(
+                in_channels=config.embed_dims[0],
+                out_channels=config.embed_dims[0],
+                kernel_size=1,
+                stride=1,
+                padding=0,
+                groups=1,
+                inference_mode=config.inference_mode,
+                use_se=False,
+                num_conv_branches=1,
+                use_scale_branch=config.use_scale_branch,
+            ),
+        )
+        # Stage Blocks
+        stages = []
+        for i in range(n_stages):
+            # Add position embeddings as requested
+            pe = self._get_position_embedding(config.position_embeddings[i])
+            if pe is not None:
+                stages.append(
+                    pe(
+                        config.embed_dims[i],
+                        config.embed_dims[i],
+                        inference_mode=config.inference_mode,
+                    )
+                )
+
+            stage = self._get_stage_blocks(
+                config.embed_dims[i],
+                i,
+                config.n_blocks,
+                token_mixer_type=config.token_mixers[i],
+                kernel_size=config.repmixer_kernel_size,
+                mlp_ratio=config.mlp_ratios[i],
+                norm_layer=LayerNormChannel,
+                drop_path_rate=config.drop_path_rate,
+                use_layer_scale=config.use_layer_scale,
+                layer_scale_init_value=config.layer_scale_init_value,
+                inference_mode=config.inference_mode,
+            )
+            stages.append(stage)
+            # Patch merging/downsampling between stages
+            if i == n_stages - 1:
+                break
+            if config.downsample[i] or config.embed_dims[i] != config.embed_dims[i + 1]:
+                stages.append(
+                    PatchEmbedding(
+                        patch_size=config.downsample_patch_size,
+                        stride=config.downsample_stride,
+                        in_channels=config.embed_dims[i],
+                        embed_dim=config.embed_dims[i + 1],
+                        inference_mode=config.inference_mode,
+                        use_se=config.downsample_use_se[i + 1],
+                    )
+                )
+
+        self.layers = torch.nn.ModuleList(stages)
+        self.conv_exp = MobileOneBlock(
+            in_channels=config.embed_dims[-1],
+            out_channels=config.embed_dims[-1] * config.clsss_ratio,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            groups=config.embed_dims[-1],
+            inference_mode=config.inference_mode,
+            use_se=True,
+            num_conv_branches=1,
+        )
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m: torch.nn.Module) -> None:
+        """Init. for classification"""
+        if isinstance(m, torch.nn.Linear):
+            torch.nn.init.normal_(m.weight, std=0.02)
+            if isinstance(m, torch.nn.Linear) and m.bias is not None:
+                torch.nn.init.constant_(m.bias, 0)
+
+    def _get_position_embedding(self, config):
+        if config is None:
+            return None
+        elif config["name"] == "RepCPE":
+            return functools.partial(RepCPE, spatial_shape=config["spatial_shape"])
+        else:
+            raise ValueError(f"Position embedding {config['name']} not supported.")
+
+    def _get_stage_blocks(
+        self,
+        dim: int,
+        block_index: int,
+        num_blocks: int,
+        token_mixer_type: str,
+        kernel_size: int = 3,
+        mlp_ratio: float = 4.0,
+        act_layer: torch.nn.Module = torch.nn.GELU,
+        norm_layer: torch.nn.Module = torch.nn.BatchNorm2d,
+        drop_rate: float = 0.0,
+        drop_path_rate: float = 0.0,
+        use_layer_scale: bool = True,
+        layer_scale_init_value: float = 1e-5,
+        inference_mode=False,
+    ) -> torch.nn.Sequential:
+        blocks = []
+        for block_idx in range(num_blocks[block_index]):
+            block_dpr = (
+                drop_path_rate
+                * (block_idx + sum(num_blocks[:block_index]))
+                / (sum(num_blocks) - 1)
+            )
+            if token_mixer_type == "repmixer":
+                blocks.append(
+                    RepMixerBlock(
+                        dim,
+                        kernel_size=kernel_size,
+                        mlp_ratio=mlp_ratio,
+                        act_layer=act_layer,
+                        drop=drop_rate,
+                        drop_path=block_dpr,
+                        use_layer_scale=use_layer_scale,
+                        layer_scale_init_value=layer_scale_init_value,
+                        inference_mode=inference_mode,
+                    )
+                )
+            elif token_mixer_type == "attention":
+                blocks.append(
+                    AttentionBlock(
+                        dim,
+                        mlp_ratio=mlp_ratio,
+                        act_layer=act_layer,
+                        norm_layer=norm_layer,
+                        drop=drop_rate,
+                        drop_path=block_dpr,
+                        use_layer_scale=use_layer_scale,
+                        layer_scale_init_value=layer_scale_init_value,
+                    )
+                )
+            else:
+                raise ValueError(
+                    "Token mixer type: {} not supported".format(token_mixer_type)
+                )
+
+        return torch.nn.Sequential(*blocks)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.patch_embedding(x)
+        for layer in self.layers:
+            x = layer(x)
+
+        return self.conv_exp(x)
+
+
+class FastVLMConnector(torch.nn.Module):
+    def __init__(self, config: FastVLMConfig) -> None:
+        super().__init__()
+
+        vision_out_dim = (
+            config.vision_config.embed_dims[-1] * config.vision_config.clsss_ratio
+        )
+        llm_hidden_size = config.text_config.hidden_size
+        self.connector = torch.nn.Sequential(
+            torch.nn.Linear(vision_out_dim, llm_hidden_size),
+            torch.nn.GELU(),
+            torch.nn.Linear(llm_hidden_size, llm_hidden_size),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.connector(x)
+
+
+# NOTE: The following components are copied from
+# https://huggingface.co/apple/FastVLM-0.5B/blob/main/llava_qwen.py
+#
+#    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.
+
+
+class RepCPE(torch.nn.Module):
+    """Conditional Positional Encodings for Vision Transformers <https://arxiv.org/pdf/2102.10882.pdf>"""
+
+    def __init__(
+        self,
+        in_channels: int,
+        embed_dim: int = 768,
+        spatial_shape: typing.Union[int, typing.Tuple[int, int]] = (7, 7),
+        inference_mode=False,
+    ) -> None:
+        super(RepCPE, self).__init__()
+        if isinstance(spatial_shape, int):
+            spatial_shape = tuple([spatial_shape] * 2)
+        assert isinstance(spatial_shape, typing.Tuple), (
+            f'"spatial_shape" must by a sequence or int, '
+            f"get {type(spatial_shape)} instead."
+        )
+        assert len(spatial_shape) == 2, (
+            f'Length of "spatial_shape" should be 2, '
+            f"got {len(spatial_shape)} instead."
+        )
+
+        self.spatial_shape = spatial_shape
+        self.embed_dim = embed_dim
+        self.in_channels = in_channels
+        self.groups = embed_dim
+
+        if inference_mode:
+            self.reparam_conv = torch.nn.Conv2d(
+                in_channels=self.in_channels,
+                out_channels=self.embed_dim,
+                kernel_size=self.spatial_shape,
+                stride=1,
+                padding=int(self.spatial_shape[0] // 2),
+                groups=self.embed_dim,
+                bias=True,
+            )
+        else:
+            self.pe = torch.nn.Conv2d(
+                in_channels,
+                embed_dim,
+                spatial_shape,
+                1,
+                int(spatial_shape[0] // 2),
+                bias=True,
+                groups=embed_dim,
+            )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if hasattr(self, "reparam_conv"):
+            x = self.reparam_conv(x)
+            return x
+        else:
+            x = self.pe(x) + x
+            return x
+
+    def reparameterize(self) -> None:
+        # Build equivalent Id tensor
+        input_dim = self.in_channels // self.groups
+        kernel_value = torch.zeros(
+            (
+                self.in_channels,
+                input_dim,
+                self.spatial_shape[0],
+                self.spatial_shape[1],
+            ),
+            dtype=self.pe.weight.dtype,
+            device=self.pe.weight.device,
+        )
+        for i in range(self.in_channels):
+            kernel_value[
+                i,
+                i % input_dim,
+                self.spatial_shape[0] // 2,
+                self.spatial_shape[1] // 2,
+            ] = 1
+        id_tensor = kernel_value
+
+        # Reparameterize Id tensor and conv
+        w_final = id_tensor + self.pe.weight
+        b_final = self.pe.bias
+
+        # Introduce reparam conv
+        self.reparam_conv = torch.nn.Conv2d(
+            in_channels=self.in_channels,
+            out_channels=self.embed_dim,
+            kernel_size=self.spatial_shape,
+            stride=1,
+            padding=int(self.spatial_shape[0] // 2),
+            groups=self.embed_dim,
+            bias=True,
+        )
+        self.reparam_conv.weight.data = w_final
+        self.reparam_conv.bias.data = b_final
+
+        self.__delattr__("pe")
+
+
+class LayerNormChannel(torch.nn.Module):
+    """LayerNorm only for Channel Dimension"""
+
+    def __init__(self, num_features, eps=1e-05) -> None:
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(num_features))
+        self.bias = torch.nn.Parameter(torch.zeros(num_features))
+        self.eps = eps
+
+    def forward(self, x) -> torch.Tensor:
+        u = x.mean(1, keepdim=True)
+        s = (x - u).pow(2).mean(1, keepdim=True)
+        x = (x - u) / torch.sqrt(s + self.eps)
+        x = self.weight.unsqueeze(-1).unsqueeze(-1) * x + self.bias.unsqueeze(
+            -1
+        ).unsqueeze(-1)
+
+        return x
+
+
+class MobileOneBlock(torch.nn.Module):
+    """The MobileOne building block. <https://arxiv.org/pdf/2206.04040.pdf>"""
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        padding: int = 0,
+        dilation: int = 1,
+        groups: int = 1,
+        inference_mode: bool = False,
+        use_se: bool = False,
+        use_act: bool = True,
+        use_scale_branch: bool = True,
+        num_conv_branches: int = 1,
+        activation: torch.nn.Module = torch.nn.GELU(),
+    ) -> None:
+        super(MobileOneBlock, self).__init__()
+        self.inference_mode = inference_mode
+        self.groups = groups
+        self.stride = stride
+        self.padding = padding
+        self.dilation = dilation
+        self.kernel_size = kernel_size
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.num_conv_branches = num_conv_branches
+
+        # Check if SE-ReLU is requested
+        if use_se:
+            self.se = SEBlock(out_channels)
+        else:
+            self.se = torch.nn.Identity()
+
+        if use_act:
+            self.activation = activation
+        else:
+            self.activation = torch.nn.Identity()
+
+        if inference_mode:
+            self.reparam_conv = torch.nn.Conv2d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=kernel_size,
+                stride=stride,
+                padding=padding,
+                dilation=dilation,
+                groups=groups,
+                bias=True,
+            )
+        else:
+            # Re-parameterizable skip connection
+            # Fallback, sometimes batchnorm tensors
+            # do not get instantiated correctly on some processes
+            # when using deepspeed + accelerate
+            norm_layer = torch.nn.BatchNorm2d(num_features=in_channels)
+            if norm_layer.weight.shape[0] == 0:
+                norm_layer.weight = torch.nn.Parameter(torch.zeros(in_channels))
+            if norm_layer.bias.shape[0] == 0:
+                norm_layer.bias = torch.nn.Parameter(torch.zeros(in_channels))
+
+            self.rbr_skip = (
+                norm_layer if out_channels == in_channels and stride == 1 else None
+            )
+
+            # Re-parameterizable conv branches
+            if num_conv_branches > 0:
+                rbr_conv = list()
+                for _ in range(self.num_conv_branches):
+                    rbr_conv.append(
+                        self._conv_bn(kernel_size=kernel_size, padding=padding)
+                    )
+                self.rbr_conv = torch.nn.ModuleList(rbr_conv)
+            else:
+                self.rbr_conv = None
+
+            # Re-parameterizable scale branch
+            self.rbr_scale = None
+            if not isinstance(kernel_size, int):
+                kernel_size = kernel_size[0]
+            if (kernel_size > 1) and use_scale_branch:
+                self.rbr_scale = self._conv_bn(kernel_size=1, padding=0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply forward pass."""
+        # Inference mode forward pass.
+        if self.inference_mode:
+            return self.activation(self.se(self.reparam_conv(x)))
+
+        # Multi-branched train-time forward pass.
+        # Skip branch output
+        identity_out = 0
+        if self.rbr_skip is not None:
+            identity_out = self.rbr_skip(x)
+
+        # Scale branch output
+        scale_out = 0
+        if self.rbr_scale is not None:
+            scale_out = self.rbr_scale(x)
+
+        # Other branches
+        out = scale_out + identity_out
+        if self.rbr_conv is not None:
+            for ix in range(self.num_conv_branches):
+                out += self.rbr_conv[ix](x)
+
+        return self.activation(self.se(out))
+
+    def reparameterize(self):
+        """Following works like `RepVGG: Making VGG-style ConvNets Great Again` -
+        https://arxiv.org/pdf/2101.03697.pdf. We re-parameterize multi-branched
+        architecture used at training time to obtain a plain CNN-like structure
+        for inference.
+        """
+        if self.inference_mode:
+            return
+
+        kernel, bias = self._get_kernel_bias()
+        self.reparam_conv = torch.nn.Conv2d(
+            in_channels=self.in_channels,
+            out_channels=self.out_channels,
+            kernel_size=self.kernel_size,
+            stride=self.stride,
+            padding=self.padding,
+            dilation=self.dilation,
+            groups=self.groups,
+            bias=True,
+        )
+        self.reparam_conv.weight.data = kernel
+        self.reparam_conv.bias.data = bias
+
+        # Delete un-used branches
+        self.__delattr__("rbr_conv")
+        self.__delattr__("rbr_scale")
+        if hasattr(self, "rbr_skip"):
+            self.__delattr__("rbr_skip")
+
+        self.inference_mode = True
+
+    def _get_kernel_bias(self) -> typing.Tuple[torch.Tensor, torch.Tensor]:
+        """Method to obtain re-parameterized kernel and bias.
+        Reference: https://github.com/DingXiaoH/RepVGG/blob/main/repvgg.py#L83
+        """
+        # get weights and bias of scale branch
+        kernel_scale = 0
+        bias_scale = 0
+        if self.rbr_scale is not None:
+            kernel_scale, bias_scale = self._fuse_bn_tensor(self.rbr_scale)
+            # Pad scale branch kernel to match conv branch kernel size.
+            pad = self.kernel_size // 2
+            kernel_scale = torch.nn.functional.pad(kernel_scale, [pad, pad, pad, pad])
+
+        # get weights and bias of skip branch
+        kernel_identity = 0
+        bias_identity = 0
+        if self.rbr_skip is not None:
+            kernel_identity, bias_identity = self._fuse_bn_tensor(self.rbr_skip)
+
+        # get weights and bias of conv branches
+        kernel_conv = 0
+        bias_conv = 0
+        if self.rbr_conv is not None:
+            for ix in range(self.num_conv_branches):
+                _kernel, _bias = self._fuse_bn_tensor(self.rbr_conv[ix])
+                kernel_conv += _kernel
+                bias_conv += _bias
+
+        kernel_final = kernel_conv + kernel_scale + kernel_identity
+        bias_final = bias_conv + bias_scale + bias_identity
+        return kernel_final, bias_final
+
+    def _fuse_bn_tensor(
+        self, branch: typing.Union[torch.nn.Sequential, torch.nn.BatchNorm2d]
+    ) -> typing.Tuple[torch.Tensor, torch.Tensor]:
+        """Method to fuse batchnorm layer with preceeding conv layer.
+        Reference: https://github.com/DingXiaoH/RepVGG/blob/main/repvgg.py#L95
+        """
+        if isinstance(branch, torch.nn.Sequential):
+            kernel = branch.conv.weight
+            running_mean = branch.bn.running_mean
+            running_var = branch.bn.running_var
+            gamma = branch.bn.weight
+            beta = branch.bn.bias
+            eps = branch.bn.eps
+        else:
+            assert isinstance(branch, torch.nn.BatchNorm2d)
+            if not hasattr(self, "id_tensor"):
+                input_dim = self.in_channels // self.groups
+
+                kernel_size = self.kernel_size
+                if isinstance(self.kernel_size, int):
+                    kernel_size = (self.kernel_size, self.kernel_size)
+
+                kernel_value = torch.zeros(
+                    (self.in_channels, input_dim, kernel_size[0], kernel_size[1]),
+                    dtype=branch.weight.dtype,
+                    device=branch.weight.device,
+                )
+                for i in range(self.in_channels):
+                    kernel_value[
+                        i, i % input_dim, kernel_size[0] // 2, kernel_size[1] // 2
+                    ] = 1
+                self.id_tensor = kernel_value
+            kernel = self.id_tensor
+            running_mean = branch.running_mean
+            running_var = branch.running_var
+            gamma = branch.weight
+            beta = branch.bias
+            eps = branch.eps
+        std = (running_var + eps).sqrt()
+        t = (gamma / std).reshape(-1, 1, 1, 1)
+        return kernel * t, beta - running_mean * gamma / std
+
+    def _conv_bn(self, kernel_size: int, padding: int) -> torch.nn.Sequential:
+        # Fallback, sometimes batchnorm tensors
+        # do not get instantiated correctly on some processes
+        # when using deepspeed + accelerate
+        norm_layer = torch.nn.BatchNorm2d(num_features=self.out_channels)
+        if norm_layer.weight.shape[0] == 0:
+            norm_layer.weight = torch.nn.Parameter(torch.zeros(self.out_channels))
+        if norm_layer.bias.shape[0] == 0:
+            norm_layer.bias = torch.nn.Parameter(torch.zeros(self.out_channels))
+
+        mod_list = torch.nn.Sequential()
+        mod_list.add_module(
+            "conv",
+            torch.nn.Conv2d(
+                in_channels=self.in_channels,
+                out_channels=self.out_channels,
+                kernel_size=kernel_size,
+                stride=self.stride,
+                padding=padding,
+                groups=self.groups,
+                bias=False,
+            ),
+        )
+        mod_list.add_module("bn", norm_layer)
+        return mod_list
+
+
+class SEBlock(torch.nn.Module):
+    """The Squeeze and Excite module. <https://arxiv.org/pdf/1709.01507.pdf>"""
+
+    def __init__(self, in_channels: int, rd_ratio: float = 0.0625) -> None:
+        super(SEBlock, self).__init__()
+        self.reduce = torch.nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=int(in_channels * rd_ratio),
+            kernel_size=1,
+            stride=1,
+            bias=True,
+        )
+        self.expand = torch.nn.Conv2d(
+            in_channels=int(in_channels * rd_ratio),
+            out_channels=in_channels,
+            kernel_size=1,
+            stride=1,
+            bias=True,
+        )
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        _, c, h, w = inputs.size()
+        x = F.avg_pool2d(inputs, kernel_size=[h, w])
+        # x = F.avg_pool2d(inputs, kernel_size=[16, 16])
+        x = self.reduce(x)
+        x = F.relu(x)
+        x = self.expand(x)
+        x = torch.sigmoid(x)
+        x = x.view(-1, c, 1, 1)
+        return inputs * x
+
+
+class RepMixerBlock(torch.nn.Module):
+    """The Metaformer block with RepMixer as token mixer <https://arxiv.org/pdf/2111.11418.pdf>"""
+
+    def __init__(
+        self,
+        dim: int,
+        kernel_size: int = 3,
+        mlp_ratio: float = 4.0,
+        act_layer: torch.nn.Module = torch.nn.GELU,
+        drop: float = 0.0,
+        drop_path: float = 0.0,
+        use_layer_scale: bool = True,
+        layer_scale_init_value: float = 1e-5,
+        inference_mode: bool = False,
+    ):
+        super().__init__()
+
+        self.token_mixer = RepMixer(
+            dim,
+            kernel_size=kernel_size,
+            use_layer_scale=use_layer_scale,
+            layer_scale_init_value=layer_scale_init_value,
+            inference_mode=inference_mode,
+        )
+        assert mlp_ratio > 0, "MLP ratio should be greater than 0, found: {}".format(
+            mlp_ratio
+        )
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.convffn = ConvFFN(
+            in_channels=dim,
+            hidden_channels=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+        )
+        # Drop Path
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else torch.nn.Identity()
+        # Layer Scale
+        self.use_layer_scale = use_layer_scale
+        if use_layer_scale:
+            self.layer_scale = torch.nn.Parameter(
+                layer_scale_init_value * torch.ones((dim, 1, 1)), requires_grad=True
+            )
+
+    def forward(self, x):
+        if self.use_layer_scale:
+            x = self.token_mixer(x)
+            x = x + self.drop_path(self.layer_scale * self.convffn(x))
+        else:
+            x = self.token_mixer(x)
+            x = x + self.drop_path(self.convffn(x))
+        return x
+
+
+class RepMixer(torch.nn.Module):
+    """Reparameterizable token mixer. <https://arxiv.org/pdf/2303.14189.pdf>"""
+
+    def __init__(
+        self,
+        dim,
+        kernel_size=3,
+        use_layer_scale=True,
+        layer_scale_init_value=1e-5,
+        inference_mode: bool = False,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.kernel_size = kernel_size
+        self.inference_mode = inference_mode
+
+        if inference_mode:
+            self.reparam_conv = torch.nn.Conv2d(
+                in_channels=self.dim,
+                out_channels=self.dim,
+                kernel_size=self.kernel_size,
+                stride=1,
+                padding=self.kernel_size // 2,
+                groups=self.dim,
+                bias=True,
+            )
+        else:
+            self.norm = MobileOneBlock(
+                dim,
+                dim,
+                kernel_size,
+                padding=kernel_size // 2,
+                groups=dim,
+                use_act=False,
+                use_scale_branch=False,
+                num_conv_branches=0,
+            )
+            self.mixer = MobileOneBlock(
+                dim,
+                dim,
+                kernel_size,
+                padding=kernel_size // 2,
+                groups=dim,
+                use_act=False,
+            )
+            self.use_layer_scale = use_layer_scale
+            if use_layer_scale:
+                self.layer_scale = torch.nn.Parameter(
+                    layer_scale_init_value * torch.ones((dim, 1, 1)), requires_grad=True
+                )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if hasattr(self, "reparam_conv"):
+            x = self.reparam_conv(x)
+            return x
+        else:
+            if self.use_layer_scale:
+                x = x + self.layer_scale * (self.mixer(x) - self.norm(x))
+            else:
+                x = x + self.mixer(x) - self.norm(x)
+            return x
+
+    def reparameterize(self) -> None:
+        if self.inference_mode:
+            return
+
+        self.mixer.reparameterize()
+        self.norm.reparameterize()
+
+        if self.use_layer_scale:
+            w = self.mixer.id_tensor + self.layer_scale.unsqueeze(-1) * (
+                self.mixer.reparam_conv.weight - self.norm.reparam_conv.weight
+            )
+            b = torch.squeeze(self.layer_scale) * (
+                self.mixer.reparam_conv.bias - self.norm.reparam_conv.bias
+            )
+        else:
+            w = (
+                self.mixer.id_tensor
+                + self.mixer.reparam_conv.weight
+                - self.norm.reparam_conv.weight
+            )
+            b = self.mixer.reparam_conv.bias - self.norm.reparam_conv.bias
+
+        self.reparam_conv = torch.nn.Conv2d(
+            in_channels=self.dim,
+            out_channels=self.dim,
+            kernel_size=self.kernel_size,
+            stride=1,
+            padding=self.kernel_size // 2,
+            groups=self.dim,
+            bias=True,
+        )
+        self.reparam_conv.weight.data = w
+        self.reparam_conv.bias.data = b
+
+        self.__delattr__("mixer")
+        self.__delattr__("norm")
+        if self.use_layer_scale:
+            self.__delattr__("layer_scale")
+
+
+class ConvFFN(torch.nn.Module):
+    """Convolutional FFN Module."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_channels: typing.Optional[int] = None,
+        out_channels: typing.Optional[int] = None,
+        act_layer: torch.nn.Module = torch.nn.GELU,
+        drop: float = 0.0,
+    ) -> None:
+        super().__init__()
+        out_channels = out_channels or in_channels
+        hidden_channels = hidden_channels or in_channels
+        self.conv = torch.nn.Sequential()
+        self.conv.add_module(
+            "conv",
+            torch.nn.Conv2d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=7,
+                padding=3,
+                groups=in_channels,
+                bias=False,
+            ),
+        )
+
+        # Fallback, sometimes batchnorm tensors
+        # do not get instantiated correctly on some processes
+        # when using deepspeed + accelerate
+        norm_layer = torch.nn.BatchNorm2d(num_features=out_channels)
+        if norm_layer.weight.shape[0] == 0:
+            norm_layer.weight = torch.nn.Parameter(torch.zeros(out_channels))
+        if norm_layer.bias.shape[0] == 0:
+            norm_layer.bias = torch.nn.Parameter(torch.zeros(out_channels))
+
+        self.conv.add_module("bn", norm_layer)
+        self.fc1 = torch.nn.Conv2d(in_channels, hidden_channels, kernel_size=1)
+        self.act = act_layer()
+        self.fc2 = torch.nn.Conv2d(hidden_channels, out_channels, kernel_size=1)
+        self.drop = torch.nn.Dropout(drop)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m: torch.nn.Module) -> None:
+        if isinstance(m, torch.nn.Conv2d):
+            torch.nn.init.normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                torch.nn.init.constant_(m.bias, 0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.conv(x)
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class AttentionBlock(torch.nn.Module):
+    """The metaformer block with MHSA as token mixer. <https://arxiv.org/pdf/2111.11418.pdf>"""
+
+    def __init__(
+        self,
+        dim: int,
+        mlp_ratio: float = 4.0,
+        act_layer: torch.nn.Module = torch.nn.GELU,
+        norm_layer: torch.nn.Module = torch.nn.BatchNorm2d,
+        drop: float = 0.0,
+        drop_path: float = 0.0,
+        use_layer_scale: bool = True,
+        layer_scale_init_value: float = 1e-5,
+    ):
+        super().__init__()
+
+        # Fallback, sometimes batchnorm tensors
+        # do not get instantiated correctly on some processes
+        # when using deepspeed + accelerate
+        norm_layer_ = norm_layer(num_features=dim)
+        if norm_layer_.weight.shape[0] == 0:
+            norm_layer_.weight = torch.nn.Parameter(torch.zeros(dim))
+        if norm_layer_.bias.shape[0] == 0:
+            norm_layer_.bias = torch.nn.Parameter(torch.zeros(dim))
+
+        self.norm = norm_layer_
+        self.token_mixer = MHSA(dim=dim)
+        assert mlp_ratio > 0, "MLP ratio should be greater than 0, found: {}".format(
+            mlp_ratio
+        )
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.convffn = ConvFFN(
+            in_channels=dim,
+            hidden_channels=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+        )
+        # Drop path
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else torch.nn.Identity()
+        # Layer Scale
+        self.use_layer_scale = use_layer_scale
+        if use_layer_scale:
+            self.layer_scale_1 = torch.nn.Parameter(
+                layer_scale_init_value * torch.ones((dim, 1, 1)), requires_grad=True
+            )
+            self.layer_scale_2 = torch.nn.Parameter(
+                layer_scale_init_value * torch.ones((dim, 1, 1)), requires_grad=True
+            )
+
+    def forward(self, x):
+        if self.use_layer_scale:
+            x = x + self.drop_path(self.layer_scale_1 * self.token_mixer(self.norm(x)))
+            x = x + self.drop_path(self.layer_scale_2 * self.convffn(x))
+        else:
+            x = x + self.drop_path(self.token_mixer(self.norm(x)))
+            x = x + self.drop_path(self.convffn(x))
+        return x
+
+
+class MHSA(torch.nn.Module):
+    """Multi-headed Self Attention module. Source modified from:
+    https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        head_dim: int = 32,
+        qkv_bias: bool = False,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+    ) -> None:
+        super().__init__()
+        assert dim % head_dim == 0, "dim should be divisible by head_dim"
+        self.head_dim = head_dim
+        self.num_heads = dim // head_dim
+        self.scale = head_dim**-0.5
+
+        self.qkv = torch.nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = torch.nn.Dropout(attn_drop)
+        self.proj = torch.nn.Linear(dim, dim)
+        self.proj_drop = torch.nn.Dropout(proj_drop)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        shape = x.shape
+        B, C, H, W = shape
+        N = H * W
+        if len(shape) == 4:
+            x = torch.flatten(x, start_dim=2).transpose(-2, -1)  # (B, N, C)
+        qkv = (
+            self.qkv(x)
+            .reshape(B, N, 3, self.num_heads, self.head_dim)
+            .permute(2, 0, 3, 1, 4)
+        )
+        q, k, v = qkv.unbind(0)  # make torchscript happy (cannot use tensor as tuple)
+
+        # trick here to make q@k.t more stable
+        attn = (q * self.scale) @ k.transpose(-2, -1)
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        if len(shape) == 4:
+            x = x.transpose(-2, -1).reshape(B, C, H, W)
+
+        return x
+
+
+class PatchEmbedding(torch.nn.Module):
+    """Convolutional patch embedding layer."""
+
+    def __init__(
+        self,
+        patch_size: int,
+        stride: int,
+        in_channels: int,
+        embed_dim: int,
+        inference_mode: bool = False,
+        use_se: bool = False,
+    ) -> None:
+        super().__init__()
+        block = list()
+        block.append(
+            ReparamLargeKernelConv(
+                in_channels=in_channels,
+                out_channels=embed_dim,
+                kernel_size=patch_size,
+                stride=stride,
+                groups=in_channels,
+                small_kernel=3,
+                inference_mode=inference_mode,
+                use_se=use_se,
+            )
+        )
+        block.append(
+            MobileOneBlock(
+                in_channels=embed_dim,
+                out_channels=embed_dim,
+                kernel_size=1,
+                stride=1,
+                padding=0,
+                groups=1,
+                inference_mode=inference_mode,
+                use_se=False,
+                num_conv_branches=1,
+            )
+        )
+        self.proj = torch.nn.Sequential(*block)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.proj(x)
+        return x
+
+
+class ReparamLargeKernelConv(torch.nn.Module):
+    """The Building Block of RepLKNet <https://arxiv.org/abs/2203.06717>"""
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int,
+        groups: int,
+        small_kernel: int,
+        inference_mode: bool = False,
+        use_se: bool = False,
+        activation: torch.nn.Module = torch.nn.GELU(),
+    ) -> None:
+        super(ReparamLargeKernelConv, self).__init__()
+
+        self.stride = stride
+        self.groups = groups
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.activation = activation
+
+        self.kernel_size = kernel_size
+        self.small_kernel = small_kernel
+        self.padding = kernel_size // 2
+
+        # Check if SE is requested
+        if use_se:
+            self.se = SqueezeExcite(out_channels, rd_ratio=0.25)
+        else:
+            self.se = torch.nn.Identity()
+
+        if inference_mode:
+            self.lkb_reparam = torch.nn.Conv2d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=kernel_size,
+                stride=stride,
+                padding=self.padding,
+                dilation=1,
+                groups=groups,
+                bias=True,
+            )
+        else:
+            self.lkb_origin = self._conv_bn(
+                kernel_size=kernel_size, padding=self.padding
+            )
+            if small_kernel is not None:
+                assert (
+                    small_kernel <= kernel_size
+                ), "The kernel size for re-param cannot be larger than the large kernel!"
+                self.small_conv = self._conv_bn(
+                    kernel_size=small_kernel, padding=small_kernel // 2
+                )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply forward pass."""
+        if hasattr(self, "lkb_reparam"):
+            out = self.lkb_reparam(x)
+        else:
+            out = self.lkb_origin(x)
+            if hasattr(self, "small_conv"):
+                out += self.small_conv(x)
+
+        return self.activation(self.se(out))
+
+    def get_kernel_bias(self) -> typing.Tuple[torch.Tensor, torch.Tensor]:
+        """Method to obtain re-parameterized kernel and bias.
+        Reference: https://github.com/DingXiaoH/RepLKNet-pytorch
+        Returns:
+            Tuple of (kernel, bias) after fusing branches.
+        """
+        eq_k, eq_b = self._fuse_bn(self.lkb_origin.conv, self.lkb_origin.bn)
+        if hasattr(self, "small_conv"):
+            small_k, small_b = self._fuse_bn(self.small_conv.conv, self.small_conv.bn)
+            eq_b += small_b
+            eq_k += torch.nn.functional.pad(
+                small_k, [(self.kernel_size - self.small_kernel) // 2] * 4
+            )
+        return eq_k, eq_b
+
+    def reparameterize(self) -> None:
+        """
+        Following works like `RepVGG: Making VGG-style ConvNets Great Again` -
+        https://arxiv.org/pdf/2101.03697.pdf. We re-parameterize multi-branched
+        architecture used at training time to obtain a plain CNN-like structure
+        for inference.
+        """
+        eq_k, eq_b = self.get_kernel_bias()
+        self.lkb_reparam = torch.nn.Conv2d(
+            in_channels=self.in_channels,
+            out_channels=self.out_channels,
+            kernel_size=self.kernel_size,
+            stride=self.stride,
+            padding=self.padding,
+            dilation=self.lkb_origin.conv.dilation,
+            groups=self.groups,
+            bias=True,
+        )
+
+        self.lkb_reparam.weight.data = eq_k
+        self.lkb_reparam.bias.data = eq_b
+        self.__delattr__("lkb_origin")
+        if hasattr(self, "small_conv"):
+            self.__delattr__("small_conv")
+
+    @staticmethod
+    def _fuse_bn(
+        conv: torch.Tensor, bn: torch.nn.BatchNorm2d
+    ) -> typing.Tuple[torch.Tensor, torch.Tensor]:
+        kernel = conv.weight
+        running_mean = bn.running_mean
+        running_var = bn.running_var
+        gamma = bn.weight
+        beta = bn.bias
+        eps = bn.eps
+        std = (running_var + eps).sqrt()
+        t = (gamma / std).reshape(-1, 1, 1, 1)
+        return kernel * t, beta - running_mean * gamma / std
+
+    def _conv_bn(self, kernel_size: int, padding: int = 0) -> torch.nn.Sequential:
+        # Fallback, sometimes batchnorm tensors
+        # do not get instantiated correctly on some processes
+        # when using deepspeed + accelerate
+        norm_layer = torch.nn.BatchNorm2d(num_features=self.out_channels)
+        if norm_layer.weight.shape[0] == 0:
+            norm_layer.weight = torch.nn.Parameter(torch.zeros(self.out_channels))
+        if norm_layer.bias.shape[0] == 0:
+            norm_layer.bias = torch.nn.Parameter(torch.zeros(self.out_channels))
+
+        mod_list = torch.nn.Sequential()
+        mod_list.add_module(
+            "conv",
+            torch.nn.Conv2d(
+                in_channels=self.in_channels,
+                out_channels=self.out_channels,
+                kernel_size=kernel_size,
+                stride=self.stride,
+                padding=padding,
+                groups=self.groups,
+                bias=False,
+            ),
+        )
+        mod_list.add_module("bn", norm_layer)
+        return mod_list

policies/dynamicvla/modeling_vlm_with_expert.py

@@ -0,0 +1,658 @@
+# -*- coding: utf-8 -*-
+#
+# @File:   modeling_vlm_with_expert.py
+# @Author: Haozhe Xie
+# @Date:   2025-09-16 11:23:15
+# @Last Modified by: Haozhe Xie
+# @Last Modified at: 2026-02-26 11:29:07
+# @Email:  root@haozhexie.com
+
+import collections
+import copy
+import logging
+
+import torch
+from transformers import AutoModel, AutoTokenizer, PretrainedConfig, PreTrainedModel
+
+
+class VLMWithExpertModel(torch.nn.Module):
+    def __init__(
+        self,
+        model_id: str,
+        vlm: PreTrainedModel,
+        freeze_text_model: bool = False,
+        freeze_connector: bool = False,
+        freeze_vision_model: bool = False,
+        attention_mode: str = "self_attn",
+        num_expert_layers: int = -1,
+        num_expert_skip_layers: int = 0,
+        num_vlm_layers: int = -1,
+        self_attn_every_n_layers: int = -1,
+        expert_width_multiplier: float = 0.5,
+    ) -> None:
+        super().__init__()
+        # Tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+        # VLM
+        self.vlm = vlm
+        self.vlm_config = self.vlm.config
+        if hasattr(self.vlm, "lm_head"):
+            del self.vlm.lm_head
+        if num_vlm_layers > 0:
+            logging.info(
+                "Reducing the number of VLM layers from %d to %d ..."
+                % (
+                    len(self._get_text_model(self.get_vlm_model()).layers),
+                    num_vlm_layers,
+                )
+            )
+            del self._get_text_model(self.get_vlm_model()).layers[num_vlm_layers:]
+
+        self.num_vlm_layers = len(self._get_text_model(self.get_vlm_model()).layers)
+        # Action Expert
+        lm_expert_config = self._get_expert_config(
+            self.vlm_config.text_config, num_expert_layers, expert_width_multiplier
+        )
+        self.lm_expert = self._get_expert_model(
+            lm_expert_config,
+            attention_mode,
+            num_expert_skip_layers,
+            self_attn_every_n_layers,
+        )
+        # Remove token embeddings
+        if hasattr(self.lm_expert, "embed_tokens"):
+            del self.lm_expert.embed_tokens
+
+        self.num_attention_heads = self.vlm_config.text_config.num_attention_heads
+        self.num_key_value_heads = self.vlm_config.text_config.num_key_value_heads
+        self.num_vlm_layers = len(self._get_text_model(self.get_vlm_model()).layers)
+        self.num_expert_layers = len(self.lm_expert.layers) - num_expert_skip_layers
+        self.num_expert_skip_layers = num_expert_skip_layers
+        self.self_attn_every_n_layers = self_attn_every_n_layers
+
+        self.freeze_vision_model = freeze_vision_model
+        self.freeze_connector = freeze_connector
+        self.freeze_text_model = freeze_text_model
+        self.attention_mode = attention_mode
+        self.expert_hidden_size = lm_expert_config.hidden_size
+        self._set_requires_grad()
+
+    def _get_vision_model(self, vlm):
+        if hasattr(vlm, "vision_model"):
+            return vlm.vision_model
+        else:
+            raise ValueError("Cannot find vision model in the VLM")
+
+    def _get_text_model(self, vlm):
+        if hasattr(vlm, "text_model"):
+            return vlm.text_model
+        else:
+            raise ValueError("Cannot find text model in the VLM")
+
+    def _get_expert_config(
+        self, text_config: PretrainedConfig, num_layers: int, width_multiplier: float
+    ) -> PretrainedConfig:
+        expert_config = copy.deepcopy(text_config)
+        hidden_size = expert_config.hidden_size
+        expert_config.hidden_size = int(hidden_size * width_multiplier)
+        expert_config.intermediate_size = self._get_intermediate_size(
+            expert_config.hidden_size
+        )
+        num_vlm_layers = len(self._get_text_model(self.get_vlm_model()).layers)
+        expert_config.num_hidden_layers = num_vlm_layers
+        if num_layers > 0:
+            assert (
+                num_vlm_layers % num_layers == 0
+            ), "Number of layers in the VLM %d are not multiple of %d " % (
+                num_vlm_layers,
+                num_layers,
+            )
+            expert_config.num_hidden_layers = num_layers
+
+        return expert_config
+
+    def _get_intermediate_size(
+        self, hidden_dim: int, ffn_dim_multiplier: int = 4, multiple_of: int = 256
+    ) -> int:
+        hidden_dim = int(2 * hidden_dim / 3)
+        hidden_dim = int(ffn_dim_multiplier * hidden_dim)
+        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
+        return hidden_dim
+
+    def _get_expert_model(
+        self,
+        expert_config: PretrainedConfig,
+        attention_mode: str,
+        num_expert_skip_layers: int,
+        self_attn_every_n_layers: int,
+    ) -> PreTrainedModel:
+        text_config = self.vlm_config.text_config
+        if not hasattr(text_config, "head_dim"):
+            text_config.head_dim = (
+                text_config.hidden_size // text_config.num_attention_heads
+            )
+            expert_config.head_dim = text_config.head_dim
+        if not hasattr(expert_config, "attention_bias"):
+            expert_config.attention_bias = False
+
+        expert_model = AutoModel.from_config(expert_config)
+        if "cross" in attention_mode:
+            # Reshape qkv projections to have the same input dimension as the vlm
+            for layer_idx in range(num_expert_skip_layers, len(expert_model.layers)):
+                if (
+                    self_attn_every_n_layers > 0
+                    and layer_idx % self_attn_every_n_layers == 0
+                ):
+                    continue
+
+                # Remove projectors for key (as ROPE has been applied in VLM)
+                expert_model.layers[layer_idx].self_attn.v_proj = torch.nn.Linear(
+                    text_config.num_key_value_heads * text_config.head_dim,
+                    expert_config.num_key_value_heads * expert_config.head_dim,
+                    bias=expert_config.attention_bias,
+                )
+                del expert_model.layers[layer_idx].self_attn.k_proj
+
+        return expert_model
+
+    def get_vlm_model(self) -> PreTrainedModel:
+        return self.vlm.model
+
+    def _set_requires_grad(self) -> None:
+        if self.freeze_vision_model:
+            self._get_vision_model(self.get_vlm_model()).eval()
+            for params in self._get_vision_model(self.get_vlm_model()).parameters():
+                params.requires_grad = False
+        if self.freeze_text_model:
+            self._get_text_model(self.get_vlm_model()).eval()
+            for params in self._get_text_model(self.get_vlm_model()).parameters():
+                params.requires_grad = False
+        if self.freeze_connector and hasattr(self.get_vlm_model(), "connector"):
+            self.get_vlm_model().connector.eval()
+            for params in self.get_vlm_model().connector.parameters():
+                params.requires_grad = False
+
+    def train(self, mode: bool = True) -> None:
+        super().train(mode)
+        if self.freeze_vision_model:
+            self._get_vision_model(self.get_vlm_model()).eval()
+        if self.freeze_connector and hasattr(self.get_vlm_model(), "connector"):
+            self.get_vlm_model().connector.eval()
+        if self.freeze_text_model:
+            self._get_text_model(self.get_vlm_model()).eval()
+
+    def embed_image(self, image: torch.Tensor) -> torch.Tensor:
+        assert len(image.shape) in [
+            4,
+            5,
+        ], f"Image should be [B, C, H, W] or [B, N, C, H, W], got {image.shape}"
+        if len(image.shape) == 4:
+            image = image.unsqueeze(1)  # [B, 1, C, H, W]
+
+        return self.get_vlm_model().get_image_features(image)
+
+    def embed_language_tokens(self, tokens: torch.Tensor) -> torch.Tensor:
+        return self._get_text_model(self.get_vlm_model()).get_input_embeddings()(tokens)
+
+    def _qkv_proj_layer(
+        self,
+        layer: torch.nn.Module,
+        position_ids: torch.Tensor,
+        q_in: torch.Tensor,
+        k_in: torch.Tensor | None = None,
+        v_in: torch.Tensor | None = None,
+    ) -> tuple[
+        dict[int, torch.Tensor], dict[int, torch.Tensor], dict[int, torch.Tensor]
+    ]:
+        assert q_in is not None, "q_in should not be None"
+
+        attn_layer = layer.self_attn
+        q_in = layer.input_layernorm(q_in).to(dtype=attn_layer.q_proj.weight.dtype)
+        if k_in is None and v_in is None:
+            k_in = q_in
+            v_in = q_in
+        else:
+            v_in = v_in.to(dtype=attn_layer.v_proj.weight.dtype)
+            if hasattr(attn_layer, "k_proj"):
+                k_in = k_in.to(dtype=attn_layer.k_proj.weight.dtype)
+
+        v_shape = (*v_in.shape[:-1], -1, attn_layer.head_dim)
+        q_shape = (*q_in.shape[:-1], -1, attn_layer.head_dim)
+        if hasattr(attn_layer, "k_proj"):
+            k_shape = (*k_in.shape[:-1], -1, attn_layer.head_dim)
+
+        q_states, k_states, v_states = {}, {}, {}
+        if position_ids.ndim == 2:  # 1D Rope
+            v_states["t"] = attn_layer.v_proj(v_in).view(v_shape)
+            q_states["t"] = attn_layer.q_proj(q_in).view(q_shape)
+            if hasattr(attn_layer, "q_norm"):
+                q_states["t"] = attn_layer.q_norm(q_states["t"])
+            if hasattr(attn_layer, "k_proj"):
+                k_states["t"] = attn_layer.k_proj(k_in).view(k_shape)
+                if hasattr(attn_layer, "k_norm"):
+                    k_states["t"] = attn_layer.k_norm(k_states["t"])
+        elif position_ids.ndim == 3 and position_ids.shape[2] == 3:  # 3D Rope
+            qtr_head_dim = attn_layer.head_dim // 4
+            # Value
+            v_states["t"] = attn_layer.v_proj(v_in).view(v_shape)
+            # Query
+            _query_states = attn_layer.q_norm(attn_layer.q_proj(q_in).view(q_shape))
+            q_states["t"] = _query_states[..., : qtr_head_dim * 2]
+            q_states["h"] = _query_states[..., qtr_head_dim * 2 : -qtr_head_dim]
+            q_states["w"] = _query_states[..., -qtr_head_dim:]
+            # Key (can be skipped in cross-attention)
+            if hasattr(attn_layer, "k_proj"):
+                _key_states = attn_layer.k_norm(attn_layer.k_proj(k_in).view(k_shape))
+                k_states["t"] = _key_states[..., : qtr_head_dim * 2]
+                k_states["h"] = _key_states[..., qtr_head_dim * 2 : -qtr_head_dim]
+                k_states["w"] = _key_states[..., -qtr_head_dim:]
+        else:
+            raise ValueError(f"Unknown position_ids shape: {position_ids.shape}")
+
+        return q_states, k_states, v_states
+
+    def apply_rope(
+        self,
+        hidden_states: torch.Tensor,
+        positions: torch.Tensor,
+        wavelength: int = 10_000,
+    ):
+        if positions.ndim == 2:  # 1D Rope
+            return self._apply_rope(
+                hidden_states["t"].unsqueeze(0), positions.unsqueeze(0), wavelength
+            ).squeeze(0)
+        elif positions.ndim == 3 and positions.shape[2] == 3:  # 3D Rope
+            roped_states = self._apply_rope(
+                torch.stack(
+                    [
+                        hidden_states["t"],
+                        hidden_states["h"].repeat_interleave(2, dim=-1),
+                        hidden_states["w"].repeat_interleave(2, dim=-1),
+                    ],
+                    dim=0,
+                ),
+                positions.permute(2, 0, 1),
+                wavelength,
+            )
+            half_indexes = torch.arange(
+                0, hidden_states["t"].size(-1), 2, device=roped_states.device
+            )
+            roped_states = roped_states.permute(1, 2, 3, 4, 0)
+            return torch.cat(
+                [
+                    roped_states[..., 0],
+                    roped_states[..., half_indexes, 1],
+                    roped_states[..., half_indexes, 2],
+                ],
+                dim=-1,
+            )
+        else:
+            raise ValueError(f"Unknown position ID shape: {positions.shape}")
+
+    def _apply_rope(
+        self,
+        hidden_states: torch.Tensor,
+        positions: torch.Tensor,
+        max_wavelength: int = 10_000,
+    ):
+        """
+        Applies RoPE positions [B, L, N] to hidden_states [B, L, H, D].
+        """
+        # Cache the sin/cos values for efficiency
+        d_half = hidden_states.size(-1) // 2
+        dtype = hidden_states.dtype
+        x = hidden_states.to(torch.float32)
+
+        freq_exponents = (4.0 / d_half) * torch.arange(
+            d_half, dtype=torch.float32, device=positions.device
+        )
+        timescale = max_wavelength**freq_exponents
+        radians = positions[..., None] / timescale[None, None, None, :]
+        radians = radians[..., None, :]
+        cos = torch.cos(radians)
+        sin = torch.sin(radians)
+
+        x1, x2 = hidden_states.split(d_half, dim=-1)
+        res = torch.empty_like(x)
+        res[..., :d_half] = x1 * cos - x2 * sin
+        res[..., d_half:] = x2 * cos + x1 * sin
+        return res.to(dtype)
+
+    def _self_attn_layer(
+        self,
+        model_layer: list[torch.nn.Module],
+        inputs_embeds: list[torch.Tensor],
+        position_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        head_dim: int,
+        use_cache: bool = True,
+        fill_kv_cache: bool = True,
+        past_key_values: dict[str, torch.Tensor] | None = None,
+    ) -> list[torch.Tensor]:
+        query_states = collections.defaultdict(list)
+        key_states = collections.defaultdict(list)
+        value_states = collections.defaultdict(list)
+        for i, hidden_states in enumerate(inputs_embeds):
+            layer = model_layer[i]
+            if hidden_states is None or layer is None:
+                continue
+
+            q_states, k_states, v_states = self._qkv_proj_layer(
+                layer, position_ids, hidden_states
+            )
+            for dst, src in (
+                (query_states, q_states),
+                (key_states, k_states),
+                (value_states, v_states),
+            ):
+                for k, v in src.items():
+                    dst[k].append(v)
+
+        # B, L, H, D with L sequence length, H number of heads, D head dim
+        # Concatenate on the number of embeddings/tokens
+        for states in (query_states, key_states, value_states):
+            for k, v_list in states.items():
+                states[k] = torch.cat(v_list, dim=1)
+
+        # Both VLM and Expert are empty. May happen during inference.
+        if len(query_states["t"]) == 0:
+            return None, past_key_values
+
+        seq_len = query_states["t"].shape[1]
+        if seq_len < position_ids.shape[1]:
+            _position_ids = position_ids[:, :seq_len]
+            _attention_mask = attention_mask[:, :seq_len, :seq_len]
+        else:
+            _position_ids = position_ids
+            _attention_mask = attention_mask
+
+        # Rotary Position Embedding
+        value_states = value_states["t"]
+        query_states = self.apply_rope(query_states, _position_ids)
+        key_states = self.apply_rope(key_states, _position_ids)
+
+        # KV Cache
+        if use_cache and past_key_values is not None:
+            if fill_kv_cache:
+                past_key_values = {
+                    "key_states": key_states,
+                    "value_states": value_states,
+                }
+            else:  # TODO: some optimization can be done, similar to `StaticCache`
+                key_states = torch.cat(
+                    [past_key_values["key_states"], key_states], dim=1
+                )
+                value_states = torch.cat(
+                    [past_key_values["value_states"], value_states], dim=1
+                )
+
+        # Eager Attention
+        att_output = self._eager_attention(
+            head_dim,
+            query_states,
+            key_states,
+            value_states,
+            _attention_mask,
+        )
+        return [att_output], past_key_values
+
+    def _cross_attn_layer(
+        self,
+        model_layer: list[torch.nn.Module],
+        inputs_embeds: list[torch.Tensor],
+        position_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        head_dim: int,
+        use_cache: bool = True,
+        fill_kv_cache: bool = True,
+        past_key_values: dict[str, torch.Tensor] | None = None,
+    ) -> list[torch.Tensor]:
+        assert len(inputs_embeds) == 2 or (
+            use_cache and past_key_values is not None and not fill_kv_cache
+        )
+
+        att_outputs = []
+        # VLM
+        if len(inputs_embeds) == 2 and not past_key_values:
+            # Prefix attention
+            seq_len = inputs_embeds[0].shape[1]
+            prefix_position_ids, suffix_position_ids = (
+                position_ids[:, :seq_len],
+                position_ids[:, seq_len:],
+            )
+            _query_states, _key_states, value_states = self._qkv_proj_layer(
+                model_layer[0], prefix_position_ids, inputs_embeds[0]
+            )
+            # Rotary Position Embedding
+            value_states = value_states["t"]
+            query_states = self.apply_rope(_query_states, prefix_position_ids)
+            key_states = self.apply_rope(_key_states, prefix_position_ids)
+            # Eager Attention
+            prefix_attention_mask = attention_mask[:, :seq_len, :seq_len]
+            att_output = self._eager_attention(
+                head_dim,
+                query_states,
+                key_states,
+                value_states,
+                prefix_attention_mask,
+            )
+            att_outputs.append(att_output)
+        else:
+            suffix_position_ids = position_ids
+
+        if use_cache and past_key_values is not None:
+            if fill_kv_cache:
+                past_key_values = {
+                    "key_states": key_states,
+                    "value_states": value_states,
+                }
+            else:  # TODO: some optimization can be done, similar to `StaticCache`
+                key_states = past_key_values["key_states"]
+                value_states = past_key_values["value_states"]
+
+        # Expert
+        expert_layer = model_layer[1]
+        if expert_layer is not None:
+            # NOTE: key_states has been ROPEd before. Directly use it here.
+            _query_states, _, _value_states = self._qkv_proj_layer(
+                expert_layer,
+                suffix_position_ids,
+                inputs_embeds[1],
+                None,
+                value_states.view(*value_states.shape[:2], -1),
+            )
+            suffix_position_ids = (
+                suffix_position_ids
+                - torch.min(suffix_position_ids, dim=1, keepdim=True).values
+            )  # start from 0
+            # Rotary Position Embedding
+            query_states = self.apply_rope(_query_states, suffix_position_ids)
+            value_states = _value_states["t"]
+            # Eager Attention
+            suffix_attention_mask = attention_mask[
+                :, -inputs_embeds[1].shape[1] :, : key_states.shape[1] :
+            ]
+            att_output = self._eager_attention(
+                head_dim,
+                query_states,
+                key_states,
+                value_states,
+                suffix_attention_mask,
+            )
+            att_outputs.append(att_output)
+        else:
+            att_outputs.append(None)
+
+        return att_outputs, past_key_values
+
+    def _get_model_layers(
+        self, models: list[PreTrainedModel]
+    ) -> list[list[torch.nn.Module]]:
+        vlm_layers = []
+        expert_layers = []
+        for i in range(self.num_expert_skip_layers):
+            vlm_layers.append(models[0].layers[i])
+            expert_layers.append(None)
+
+        multiple_of = (
+            self.num_vlm_layers - self.num_expert_skip_layers
+        ) // self.num_expert_layers
+        for i in range(self.num_expert_skip_layers, self.num_vlm_layers):
+            if multiple_of > 0 and i > 0 and i % multiple_of != 0:
+                expert_layer = None
+            else:
+                expert_layer_index = i // multiple_of if multiple_of > 0 else i
+                expert_layer = models[1].layers[expert_layer_index]
+
+            vlm_layers.append(models[0].layers[i])
+            expert_layers.append(expert_layer)
+
+        assert len(vlm_layers) == len(expert_layers)
+        return [(vlm_layers[i], expert_layers[i]) for i in range(len(vlm_layers))]
+
+    def forward(
+        self,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: dict[int, dict[str, torch.FloatTensor]] | None = None,
+        inputs_embeds: list[torch.FloatTensor] = None,
+        use_cache: bool | None = None,
+        fill_kv_cache: bool | None = None,
+    ) -> tuple[list[torch.FloatTensor], list[torch.FloatTensor] | None]:
+        models = [self._get_text_model(self.get_vlm_model()), self.lm_expert]
+        model_layers = self._get_model_layers(models)
+
+        # Decoder Layers
+        past_key_values = (
+            {i: {} for i in range(self.num_vlm_layers)}
+            if use_cache and past_key_values is None
+            else past_key_values
+        )
+        for layer_idx in range(self.num_vlm_layers):
+            attn_layer = None
+            if (
+                fill_kv_cache
+                or "cross" not in self.attention_mode
+                or layer_idx < self.num_expert_skip_layers
+                or (
+                    self.self_attn_every_n_layers > 0
+                    and layer_idx % self.self_attn_every_n_layers == 0
+                )
+            ):
+                attn_layer = self._self_attn_layer
+            else:
+                attn_layer = self._cross_attn_layer
+
+            start = 0
+            outputs_embeds = []
+            att_outputs, _past_key_values = attn_layer(
+                model_layers[layer_idx],
+                inputs_embeds,
+                position_ids,
+                attention_mask,
+                self.vlm.config.text_config.head_dim,
+                use_cache,
+                fill_kv_cache,
+                past_key_values[layer_idx] if past_key_values is not None else None,
+            )
+            if att_outputs is None:
+                continue
+            if past_key_values is not None:
+                past_key_values[layer_idx] = _past_key_values
+
+            for i, hidden_states in enumerate(inputs_embeds):
+                layer = model_layers[layer_idx][i]
+                att_output = (
+                    att_outputs[i] if i < len(att_outputs) else att_outputs[0]
+                )  # in case of self_attn
+                if hidden_states is not None:
+                    if layer is None:
+                        outputs_embeds.append(hidden_states)
+                        continue
+
+                    end = start + hidden_states.shape[1]
+                    if att_output.dtype != layer.self_attn.o_proj.weight.dtype:
+                        att_output = att_output.to(layer.self_attn.o_proj.weight.dtype)
+
+                    att_out = att_output[:, start:end]
+                    out_emb = layer.self_attn.o_proj(att_out)
+
+                    out_emb += hidden_states
+                    after_first_residual = out_emb.clone()
+
+                    # Fully Connected Layer in Decoder Layer
+                    out_emb = layer.post_attention_layernorm(out_emb)
+                    out_emb = layer.mlp(out_emb)
+                    out_emb += after_first_residual
+
+                    outputs_embeds.append(out_emb)
+                    start = end if len(att_outputs) == 1 else 0
+                else:
+                    outputs_embeds.append(None)
+
+            inputs_embeds = outputs_embeds
+
+        # Final norm
+        outputs_embeds = []
+        for i, hidden_states in enumerate(inputs_embeds):
+            if hidden_states is not None:
+                out_emb = models[i].norm(hidden_states)
+                outputs_embeds.append(out_emb)
+            else:
+                outputs_embeds.append(None)
+
+        return outputs_embeds, past_key_values
+
+    def _repeat_kv(self, hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+        batch, seq_len, num_key_value_heads, head_dim = hidden_states.shape
+        if n_rep == 1:
+            return hidden_states
+
+        hidden_states = hidden_states[:, :, :, None, :].expand(
+            batch, seq_len, num_key_value_heads, n_rep, head_dim
+        )
+        return hidden_states.reshape(
+            batch, seq_len, num_key_value_heads * n_rep, head_dim
+        )
+
+    def _eager_attention(
+        self,
+        head_dim: int,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attention_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        num_att_heads = self.num_attention_heads
+        num_key_value_heads = self.num_key_value_heads
+        num_key_value_groups = num_att_heads // num_key_value_heads
+
+        batch_size = attention_mask.size(0)
+        key_states = self._repeat_kv(key, num_key_value_groups)
+        value_states = self._repeat_kv(value, num_key_value_groups)
+
+        # Attention here is upcasted to float32 to match the original implementation.
+        att_weights = (head_dim**-0.5) * torch.matmul(
+            query.transpose(1, 2).to(dtype=torch.float32),
+            key_states.permute(0, 2, 3, 1).to(dtype=torch.float32),
+        )
+
+        att_weights = att_weights.to(dtype=torch.float32)
+        masked_att_weights = torch.where(
+            attention_mask[:, None, :, :],
+            att_weights,
+            torch.finfo(att_weights.dtype).min,
+        )
+        probs = torch.nn.functional.softmax(masked_att_weights, dim=-1)
+        probs = probs.to(dtype=value_states.dtype)
+
+        att_output = torch.matmul(probs, value_states.permute(0, 2, 1, 3))
+
+        att_output = att_output.permute(0, 2, 1, 3)
+        # we use -1 because sequence length can change
+        att_output = att_output.reshape(
+            batch_size, -1, num_key_value_heads * num_key_value_groups * head_dim
+        )
+        return att_output

requirements.txt

@@ -0,0 +1,20 @@
+# Do not list: gradio, spaces, huggingface_hub, torch
+# lerobot==0.3.3 is installed --no-deps in app.py (it pins torch<2.8, illegal on ZeroGPU)
+
+accelerate
+draccus==0.10.0
+easydict
+einops
+num2words
+numpy
+opencv-python-headless
+packaging
+pandas
+pillow
+plotly
+safetensors
+scipy
+sentencepiece
+timm
+torchvision
+transformers==5.2.0