feat(openpi): add native PyTorch LIBERO integration

configs/openpi/pi05_libero.json

@@ -0,0 +1,28 @@
+{
+  "model_type": "pi05",
+  "pi05": true,
+  "discrete_state_input": false,
+  "weight_path": "/data/liuhongda/openpi_data/openpi-assets/checkpoints/pi05_libero_pytorch/model.safetensors",
+  "norm_stats_path": "/data/liuhongda/openpi_data/openpi-assets/checkpoints/pi05_libero_pytorch/assets/physical-intelligence/libero/norm_stats.json",
+  "tokenizer_path": "/data/liuhongda/openpi_data/openpi-assets/checkpoints/pi05_libero_pytorch/assets/paligemma_tokenizer.model",
+  "transformers_runtime_path": "/data/liuhongda/openpi_data/python_deps/openpi_pytorch_runtime",
+  "paligemma_variant": "gemma_2b",
+  "action_expert_variant": "gemma_300m",
+  "action_dim": 32,
+  "output_action_dim": 7,
+  "state_dim": 8,
+  "image_size": 224,
+  "action_horizon": 10,
+  "max_token_len": 200,
+  "num_flow_steps": 10,
+  "num_inference_steps": 10,
+  "actions_per_plan": 5,
+  "num_steps_wait": 10,
+  "use_quantile_norm": true,
+  "device": "cuda",
+  "dtype": "bfloat16",
+  "compile_mode": null,
+  "pytorch_compile_mode": null,
+  "extra_delta_transform": false,
+  "seed": 0
+}

lightx2v/infer.py

@@ -25,6 +25,7 @@
 from lightx2v.models.runners.minimax_h3.minimax_h3_runner import MiniMaxH3Runner  # noqa: F401
 from lightx2v.models.runners.motus.motus_runner import MotusRunner  # noqa: F401
 from lightx2v.models.runners.neopp.neopp_runner import NeoppRunner  # noqa: F401
+from lightx2v.models.runners.openpi.openpi_runner import OpenPIRunner  # noqa: F401
 from lightx2v.models.runners.qwen_image.qwen_image_runner import QwenImageRunner  # noqa: F401
 from lightx2v.models.runners.seedvr.seedvr_runner import SeedVRRunner  # noqa: F401
 from lightx2v.models.runners.swiftvr.swiftvr_runner import SwiftVRRunner  # noqa: F401
@@ -145,6 +146,7 @@ def main():
             "infinitetalk",
             "fastwam",
             "lingbot_video",
+            "openpi",
         ],
         default="wan2.1",
     )

lightx2v/models/networks/openpi/NOTICE.md

@@ -0,0 +1,10 @@
+# OpenPI attribution
+
+The `pi0.py`, `gemma.py`, `preprocessing.py`, and `transformers_replace/`
+implementation in this directory is adapted from Physical Intelligence's
+OpenPI project at commit `15a9616a00943ada6c20a0f158e3adb39df2ccac`.
+
+OpenPI and the copied Hugging Face Transformers source files are distributed
+under the Apache License 2.0. The localization changes replace OpenPI/JAX
+imports with LightX2V-local, PyTorch-only modules while intentionally retaining
+the official model parameter names for SafeTensors compatibility.

lightx2v/models/networks/openpi/__init__.py

@@ -0,0 +1,8 @@
+"""Native PyTorch OpenPI network family for LightX2V."""
+
+from .config import Pi0Config
+from .model import OpenPIModel
+from .observation import Observation
+from .pi0 import PI0Pytorch
+
+__all__ = ["Observation", "OpenPIModel", "PI0Pytorch", "Pi0Config"]

lightx2v/models/networks/openpi/config.py

@@ -0,0 +1,88 @@
+"""Configuration objects for the native PyTorch pi0.5 backend.
+
+The numerical values mirror Physical Intelligence's released ``pi05_libero``
+configuration.  This module is intentionally pure Python: importing it does
+not require JAX, Flax, Orbax, or the OpenPI source tree.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal
+
+
+@dataclass(frozen=True)
+class GemmaConfig:
+    width: int
+    depth: int
+    mlp_dim: int
+    num_heads: int
+    num_kv_heads: int
+    head_dim: int
+
+
+GemmaVariant = Literal["dummy", "gemma_300m", "gemma_2b"]
+
+
+def get_config(variant: GemmaVariant) -> GemmaConfig:
+    """Return the Gemma dimensions used by the official OpenPI model."""
+    if variant == "dummy":
+        return GemmaConfig(width=64, depth=4, mlp_dim=128, num_heads=8, num_kv_heads=1, head_dim=16)
+    if variant == "gemma_300m":
+        return GemmaConfig(width=1024, depth=18, mlp_dim=4096, num_heads=8, num_kv_heads=1, head_dim=256)
+    if variant == "gemma_2b":
+        return GemmaConfig(width=2048, depth=18, mlp_dim=16384, num_heads=8, num_kv_heads=1, head_dim=256)
+    raise ValueError(f"Unsupported OpenPI Gemma variant: {variant!r}")
+
+
+@dataclass(frozen=True)
+class Pi0Config:
+    """Minimal model config consumed by :class:`PI0Pytorch`.
+
+    Defaults are the released pi0.5-LIBERO checkpoint.  ``action_dim`` is the
+    padded internal dimension; the LIBERO environment still consumes 7-D
+    actions.
+    """
+
+    action_dim: int = 32
+    action_horizon: int = 10
+    max_token_len: int = 200
+    dtype: Literal["bfloat16", "float32"] = "bfloat16"
+    paligemma_variant: GemmaVariant = "gemma_2b"
+    action_expert_variant: GemmaVariant = "gemma_300m"
+    pi05: bool = True
+    discrete_state_input: bool = False
+    pytorch_compile_mode: str | None = None
+
+    @classmethod
+    def from_mapping(cls, config: dict) -> "Pi0Config":
+        precision = config.get("dtype", config.get("precision", "bfloat16"))
+        return cls(
+            action_dim=int(config.get("action_dim", 32)),
+            action_horizon=int(config.get("action_horizon", 10)),
+            max_token_len=int(config.get("max_token_len", 200)),
+            dtype=str(precision),
+            paligemma_variant=str(config.get("paligemma_variant", "gemma_2b")),
+            action_expert_variant=str(config.get("action_expert_variant", "gemma_300m")),
+            pi05=bool(config.get("pi05", True)),
+            discrete_state_input=bool(config.get("discrete_state_input", False)),
+            pytorch_compile_mode=config.get("pytorch_compile_mode", config.get("compile_mode")),
+        )
+
+    def validate_pi05_libero(self) -> None:
+        expected = {
+            "pi05": True,
+            "paligemma_variant": "gemma_2b",
+            "action_expert_variant": "gemma_300m",
+            "action_dim": 32,
+            "action_horizon": 10,
+            "max_token_len": 200,
+            "discrete_state_input": False,
+        }
+        actual = {name: getattr(self, name) for name in expected}
+        wrong = {name: (actual[name], value) for name, value in expected.items() if actual[name] != value}
+        if wrong:
+            details = ", ".join(f"{name}={got!r} (expected {want!r})" for name, (got, want) in wrong.items())
+            raise ValueError(f"Configuration does not match the released pi05_libero checkpoint: {details}")
+        if self.dtype not in {"bfloat16", "float32"}:
+            raise ValueError(f"Unsupported OpenPI dtype: {self.dtype!r}")