Added support for Lingbot-VA

.gitignore

@@ -126,6 +126,7 @@ outputs/
 artifacts/
 data_local/
 credentials/*.secret
+*.bak
 
 CLAUDE.md
 .cursor/

integrations/lingbot_va/README.md

@@ -0,0 +1,90 @@
+<!--
+SPDX-FileCopyrightText: Copyright (c) 2026 Hongyu Zhou
+SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+SPDX-License-Identifier: Apache-2.0
+
+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.
+-->
+
+# flashdreams-lingbot-va
+
+LingBot-VA Image-to-Action-Video (I2AV) integration, packaged as
+a [`flashdreams`](../..) plugin.
+
+This plugin adapts [LingBot-VA](https://github.com/robbyant/lingbot-va) into the
+standard flashdreams runner/pipeline interface, achieving **2.3× speedup** over
+the original repository implementation. Note that the upstream repo wraps models
+with FSDP; compared to the original implementation with FSDP removed, the
+speedup is **1.48×**.
+
+
+## Install
+
+This plugin is a workspace member in the repo-root `pyproject.toml`, which
+is included automatically when you set up the flashdreams environment:
+
+```bash
+uv sync
+```
+
+No separate install step is needed.
+
+## Run
+
+```bash
+uv run flashdreams-run lingbot-va-robotwin-i2av \
+    --input-image-dir assets/example_data/lingbot-va/robotwin \
+    --output-dir outputs/lingbot_va/robotwin_i2av \
+    --checkpoint-root /path/to/lingbot-va-posttrain-robotwin \
+    --num-chunks 10 \
+    --benchmark True
+```
+
+### CLI arguments
+
+| flag | type | default | description |
+| --- | --- | --- | --- |
+| `--checkpoint-root` | str | `robbyant/lingbot-va-posttrain-robotwin` | Local path or HuggingFace repo ID for model weights. Must contain `transformer/`, `vae/`, `text_encoder/`, `tokenizer/` subdirs. |
+| `--input-image-dir` | path | `assets/example_data/lingbot-va/robotwin` | Directory containing three observation camera PNGs (see below). |
+| `--output-dir` | path | `outputs/lingbot_va/robotwin_i2av` | Where to write `demo.mp4`, `actions.npy`, `latents.pt`, and timing JSON. |
+| `--prompt` | str | `"Grab the medium-sized white mug, rotate it, place it on the table, and hook it onto the smooth dark gray rack."` | Text prompt describing the manipulation task. Can also be a path to a `.txt` file. |
+| `--num-chunks` | int | `10` | Number of autoregressive chunks to generate. Each chunk produces `frame_chunk_size` (2) video frames and `action_per_frame × frame_chunk_size` (32) action steps. |
+| `--seed` | int | `42` | Random seed for diffusion sampling. |
+| `--benchmark` | bool | `False` | Print per-chunk and total pipeline timing, and save `timing_flashdreams.json`. |
+| `--compile-network` | bool | `True` | Apply `torch.compile` to the DiT for faster inference. Set `False` for debugging. |
+| `--enable-offload` | bool | `False` | Offload VAE/text-encoder to CPU after use to reduce VRAM (slower). |
+| `--save-video` | bool | `True` | Decode latents and save `demo.mp4`. |
+| `--save-actions` | bool | `True` | Save predicted actions to `actions.npy`. |
+| `--num-inference-steps` | int | `25` | Diffusion steps for video denoising. |
+| `--action-num-inference-steps` | int | `50` | Diffusion steps for action denoising. |
+| `--guidance-scale` | float | `5.0` | Classifier-free guidance scale for video. |
+| `--action-guidance-scale` | float | `1.0` | Classifier-free guidance scale for actions. |
+| `--snr-shift` | float | `5.0` | Flow-match sigma shift for video scheduler. |
+| `--action-snr-shift` | float | `1.0` | Flow-match sigma shift for action scheduler. |
+
+### Input images
+
+The runner expects these files under `--input-image-dir`:
+
+- `observation.images.cam_high.png`
+- `observation.images.cam_left_wrist.png`
+- `observation.images.cam_right_wrist.png`
+
+### Outputs
+
+| file | description |
+| --- | --- |
+| `demo.mp4` | Decoded video (all chunks concatenated). |
+| `actions.npy` | Predicted actions array, shape `(num_chunks × action_per_frame × frame_chunk_size, action_dim)`. |
+| `latents.pt` | Raw latent tensors before VAE decode. |
+| `timing_flashdreams.json` | Per-chunk and total timing (only when `--benchmark True`). |

integrations/lingbot_va/lingbot_va/__init__.py

@@ -0,0 +1,18 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 Hongyu Zhou
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+"""LingBot-VA Robotwin I2AV integration for FlashDreams."""
+
+__all__: list[str] = []

integrations/lingbot_va/lingbot_va/_loaders.py

@@ -0,0 +1,183 @@
+# SPDX-FileCopyrightText: Copyright 2024-2025 The Robbyant Team Authors
+# SPDX-FileCopyrightText: Copyright (c) 2026 Hongyu Zhou
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+"""Model loading utilities adapted from upstream LingBot-VA (wan_va/modules/utils.py
+and wan_va/utils/scheduler.py)."""
+
+from __future__ import annotations
+
+import math
+
+import torch
+from diffusers import AutoencoderKLWan
+from transformers import T5TokenizerFast, UMT5EncoderModel
+
+
+def load_vae(vae_path: str, torch_dtype: torch.dtype, torch_device):
+    vae = AutoencoderKLWan.from_pretrained(vae_path, torch_dtype=torch_dtype)
+    return vae.to(torch_device)
+
+
+def load_text_encoder(text_encoder_path: str, torch_dtype: torch.dtype, torch_device):
+    text_encoder = UMT5EncoderModel.from_pretrained(
+        text_encoder_path, torch_dtype=torch_dtype
+    )
+    return text_encoder.to(torch_device)
+
+
+def load_tokenizer(tokenizer_path: str):
+    return T5TokenizerFast.from_pretrained(tokenizer_path)
+
+
+def patchify(x: torch.Tensor, patch_size: int | None) -> torch.Tensor:
+    if patch_size is None or patch_size == 1:
+        return x
+    batch_size, channels, frames, height, width = x.shape
+    x = x.view(
+        batch_size, channels, frames,
+        height // patch_size, patch_size,
+        width // patch_size, patch_size,
+    )
+    x = x.permute(0, 1, 6, 4, 2, 3, 5).contiguous()
+    x = x.view(
+        batch_size, channels * patch_size * patch_size, frames,
+        height // patch_size, width // patch_size,
+    )
+    return x
+
+
+class WanVAEStreamingWrapper:
+
+    def __init__(self, vae_model):
+        self.vae = vae_model
+        self.encoder = vae_model.encoder
+        self.quant_conv = vae_model.quant_conv
+
+        if hasattr(self.vae, "_cached_conv_counts"):
+            self.enc_conv_num = self.vae._cached_conv_counts["encoder"]
+        else:
+            count = 0
+            for m in self.encoder.modules():
+                if m.__class__.__name__ == "WanCausalConv3d":
+                    count += 1
+            self.enc_conv_num = count
+
+        self.clear_cache()
+
+    def clear_cache(self):
+        self.feat_cache = [None] * self.enc_conv_num
+
+    def encode_chunk(self, x_chunk: torch.Tensor) -> torch.Tensor:
+        if (
+            hasattr(self.vae.config, "patch_size")
+            and self.vae.config.patch_size is not None
+        ):
+            x_chunk = patchify(x_chunk, self.vae.config.patch_size)
+        feat_idx = [0]
+        out = self.encoder(x_chunk, feat_cache=self.feat_cache, feat_idx=feat_idx)
+        enc = self.quant_conv(out)
+        return enc
+
+
+# ---------------------------------------------------------------------------
+# Upstream FlowMatchScheduler (wan_va/utils/scheduler.py)
+# ---------------------------------------------------------------------------
+
+
+class FlowMatchScheduler:
+
+    def __init__(
+        self,
+        num_inference_steps=100,
+        num_train_timesteps=1000,
+        shift=3.0,
+        sigma_max=1.0,
+        sigma_min=0.003 / 1.002,
+        inverse_timesteps=False,
+        extra_one_step=False,
+        reverse_sigmas=False,
+        exponential_shift=False,
+        exponential_shift_mu=None,
+        shift_terminal=None,
+    ):
+        self.num_train_timesteps = num_train_timesteps
+        self.shift = shift
+        self.sigma_max = sigma_max
+        self.sigma_min = sigma_min
+        self.inverse_timesteps = inverse_timesteps
+        self.extra_one_step = extra_one_step
+        self.reverse_sigmas = reverse_sigmas
+        self.exponential_shift = exponential_shift
+        self.exponential_shift_mu = exponential_shift_mu
+        self.shift_terminal = shift_terminal
+        self.set_timesteps(num_inference_steps)
+
+    def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0,
+                      training=False, shift=None, dynamic_shift_len=None):
+        if shift is not None:
+            self.shift = shift
+        sigma_start = self.sigma_min + (self.sigma_max - self.sigma_min) * denoising_strength
+        if self.extra_one_step:
+            self.sigmas = torch.linspace(sigma_start, self.sigma_min,
+                                         num_inference_steps + 1)[:-1]
+        else:
+            self.sigmas = torch.linspace(sigma_start, self.sigma_min,
+                                         num_inference_steps)
+        if self.inverse_timesteps:
+            self.sigmas = torch.flip(self.sigmas, dims=[0])
+        if self.exponential_shift:
+            mu = (self.calculate_shift(dynamic_shift_len)
+                  if dynamic_shift_len is not None else self.exponential_shift_mu)
+            self.sigmas = math.exp(mu) / (math.exp(mu) + (1 / self.sigmas - 1))
+        else:
+            self.sigmas = self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas)
+        if self.shift_terminal is not None:
+            one_minus_z = 1 - self.sigmas
+            scale_factor = one_minus_z[-1] / (1 - self.shift_terminal)
+            self.sigmas = 1 - (one_minus_z / scale_factor)
+        if self.reverse_sigmas:
+            self.sigmas = 1 - self.sigmas
+        self.timesteps = self.sigmas * self.num_train_timesteps
+
+    def step(self, model_output, timestep, sample, to_final=False, **kwargs):
+        if isinstance(timestep, torch.Tensor):
+            timestep = timestep.cpu()
+        timestep_id = torch.argmin((self.timesteps - timestep).abs())
+        sigma = self.sigmas[timestep_id]
+        if to_final or timestep_id + 1 >= len(self.timesteps):
+            sigma_ = 1 if (self.inverse_timesteps or self.reverse_sigmas) else 0
+        else:
+            sigma_ = self.sigmas[timestep_id + 1]
+        prev_sample = sample + model_output * (sigma_ - sigma)
+        return prev_sample
+
+    def add_noise(self, original_samples, noise, timestep, t_dim=2):
+        if isinstance(timestep, torch.Tensor):
+            timestep = timestep.cpu()
+        timestep = timestep[None]
+        timestep_id = torch.argmin((self.timesteps[:, None] - timestep).abs(), dim=0)
+        shape = [1] * noise.ndim
+        shape[t_dim] = timestep_id.shape[0]
+        sigma = self.sigmas[timestep_id].to(original_samples).view(shape)
+        sample = (1 - sigma) * original_samples + sigma * noise
+        return sample
+
+    def calculate_shift(self, image_seq_len, base_seq_len=256,
+                        max_seq_len=8192, base_shift=0.5, max_shift=0.9):
+        m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+        b = base_shift - m * base_seq_len
+        mu = image_seq_len * m + b
+        return mu