diff --git a/configs/openpi/pi05_libero.json b/configs/openpi/pi05_libero.json new file mode 100644 index 000000000..57adc1adf --- /dev/null +++ b/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 +} diff --git a/lightx2v/infer.py b/lightx2v/infer.py index b9ebaf479..ce963d31d 100755 --- a/lightx2v/infer.py +++ b/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", ) diff --git a/lightx2v/models/networks/openpi/NOTICE.md b/lightx2v/models/networks/openpi/NOTICE.md new file mode 100644 index 000000000..b9dabef6c --- /dev/null +++ b/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. diff --git a/lightx2v/models/networks/openpi/__init__.py b/lightx2v/models/networks/openpi/__init__.py new file mode 100644 index 000000000..4f24af70c --- /dev/null +++ b/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"] diff --git a/lightx2v/models/networks/openpi/config.py b/lightx2v/models/networks/openpi/config.py new file mode 100644 index 000000000..68e4b7775 --- /dev/null +++ b/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}") diff --git a/lightx2v/models/networks/openpi/gemma.py b/lightx2v/models/networks/openpi/gemma.py new file mode 100644 index 000000000..519d342d8 --- /dev/null +++ b/lightx2v/models/networks/openpi/gemma.py @@ -0,0 +1,270 @@ +# Adapted from Physical Intelligence OpenPI (Apache-2.0), commit 15a9616. +# Localized for the LightX2V OpenPI backend; no runtime OpenPI/JAX dependency. + +from typing import Literal + +import torch +from torch import nn +from transformers import GemmaForCausalLM, PaliGemmaForConditionalGeneration +from transformers.models.auto import CONFIG_MAPPING +from transformers.models.gemma import modeling_gemma + + +class PaliGemmaWithExpertModel(nn.Module): + def __init__( + self, + vlm_config, + action_expert_config, + use_adarms=None, + precision: Literal["bfloat16", "float32"] = "bfloat16", + ): + if use_adarms is None: + use_adarms = [False, False] + super().__init__() + + vlm_config_hf = CONFIG_MAPPING["paligemma"]() + vlm_config_hf._vocab_size = 257152 # noqa: SLF001 + vlm_config_hf.image_token_index = 257152 + vlm_config_hf.text_config.hidden_size = vlm_config.width + vlm_config_hf.text_config.intermediate_size = vlm_config.mlp_dim + vlm_config_hf.text_config.num_attention_heads = vlm_config.num_heads + vlm_config_hf.text_config.head_dim = vlm_config.head_dim + vlm_config_hf.text_config.num_hidden_layers = vlm_config.depth + vlm_config_hf.text_config.num_key_value_heads = vlm_config.num_kv_heads + vlm_config_hf.text_config.hidden_activation = "gelu_pytorch_tanh" + vlm_config_hf.text_config.torch_dtype = "float32" + vlm_config_hf.text_config.vocab_size = 257152 + vlm_config_hf.text_config.use_adarms = use_adarms[0] + vlm_config_hf.text_config.adarms_cond_dim = vlm_config.width if use_adarms[0] else None + vlm_config_hf.vision_config.intermediate_size = 4304 + vlm_config_hf.vision_config.projection_dim = 2048 + vlm_config_hf.vision_config.projector_hidden_act = "gelu_fast" + vlm_config_hf.vision_config.torch_dtype = "float32" + + action_expert_config_hf = CONFIG_MAPPING["gemma"]( + head_dim=action_expert_config.head_dim, + hidden_size=action_expert_config.width, + intermediate_size=action_expert_config.mlp_dim, + num_attention_heads=action_expert_config.num_heads, + num_hidden_layers=action_expert_config.depth, + num_key_value_heads=action_expert_config.num_kv_heads, + vocab_size=257152, + hidden_activation="gelu_pytorch_tanh", + torch_dtype="float32", + use_adarms=use_adarms[1], + adarms_cond_dim=action_expert_config.width if use_adarms[1] else None, + ) + + self.paligemma = PaliGemmaForConditionalGeneration(config=vlm_config_hf) + self.gemma_expert = GemmaForCausalLM(config=action_expert_config_hf) + self.gemma_expert.model.embed_tokens = None + + self.to_bfloat16_for_selected_params(precision) + + def to_bfloat16_for_selected_params(self, precision: Literal["bfloat16", "float32"] = "bfloat16"): + if precision == "bfloat16": + self.to(dtype=torch.bfloat16) + elif precision == "float32": + self.to(dtype=torch.float32) + return + else: + raise ValueError(f"Invalid precision: {precision}") + + params_to_keep_float32 = [ + "vision_tower.vision_model.embeddings.patch_embedding.weight", + "vision_tower.vision_model.embeddings.patch_embedding.bias", + "vision_tower.vision_model.embeddings.position_embedding.weight", + "input_layernorm", + "post_attention_layernorm", + "model.norm", + ] + + for name, param in self.named_parameters(): + if any(selector in name for selector in params_to_keep_float32): + param.data = param.data.to(dtype=torch.float32) + + def embed_image(self, image: torch.Tensor): + return self.paligemma.model.get_image_features(image) + + def embed_language_tokens(self, tokens: torch.Tensor): + return self.paligemma.language_model.embed_tokens(tokens) + + def forward( + self, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: list[torch.FloatTensor] | None = None, + inputs_embeds: list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + adarms_cond: list[torch.Tensor] | None = None, + ): + if adarms_cond is None: + adarms_cond = [None, None] + if inputs_embeds[1] is None: + prefix_output = self.paligemma.language_model.forward( + inputs_embeds=inputs_embeds[0], + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + adarms_cond=adarms_cond[0] if adarms_cond is not None else None, + ) + prefix_past_key_values = prefix_output.past_key_values + prefix_output = prefix_output.last_hidden_state + suffix_output = None + elif inputs_embeds[0] is None: + suffix_output = self.gemma_expert.model.forward( + inputs_embeds=inputs_embeds[1], + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + adarms_cond=adarms_cond[1] if adarms_cond is not None else None, + ) + suffix_output = suffix_output.last_hidden_state + prefix_output = None + prefix_past_key_values = None + else: + models = [self.paligemma.language_model, self.gemma_expert.model] + num_layers = self.paligemma.config.text_config.num_hidden_layers + + # Check if gradient checkpointing is enabled for any of the models + use_gradient_checkpointing = (hasattr(self.gemma_expert.model, "gradient_checkpointing") and self.gemma_expert.model.gradient_checkpointing and self.training) or ( + hasattr(self, "gradient_checkpointing") and self.gradient_checkpointing and self.training + ) + + # Force enable gradient checkpointing if we're in training mode and the model supports it + if self.training and hasattr(self.gemma_expert.model, "gradient_checkpointing"): + if not self.gemma_expert.model.gradient_checkpointing: + print("Forcing gradient checkpointing to be enabled for Gemma expert model") + self.gemma_expert.model.gradient_checkpointing = True + use_gradient_checkpointing = True + + # Debug gradient checkpointing status + if hasattr(self, "_debug_gc_printed") and not self._debug_gc_printed: + print(f"Gemma expert model gradient checkpointing: {use_gradient_checkpointing}") + print(f"Model training mode: {self.training}") + print(f"Gemma expert model has gradient_checkpointing attr: {hasattr(self.gemma_expert.model, 'gradient_checkpointing')}") + if hasattr(self.gemma_expert.model, "gradient_checkpointing"): + print(f"Gemma expert model gradient_checkpointing value: {self.gemma_expert.model.gradient_checkpointing}") + self._debug_gc_printed = True + + # Define the complete layer computation function for gradient checkpointing + def compute_layer_complete(layer_idx, inputs_embeds, attention_mask, position_ids, adarms_cond): + models = [self.paligemma.language_model, self.gemma_expert.model] + + query_states = [] + key_states = [] + value_states = [] + gates = [] + for i, hidden_states in enumerate(inputs_embeds): + layer = models[i].layers[layer_idx] + hidden_states, gate = layer.input_layernorm(hidden_states, cond=adarms_cond[i]) # noqa: PLW2901 + gates.append(gate) + + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, layer.self_attn.head_dim) + query_state = layer.self_attn.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + key_state = layer.self_attn.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) + value_state = layer.self_attn.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + query_states.append(query_state) + key_states.append(key_state) + value_states.append(value_state) + + # Concatenate and process attention + query_states = torch.cat(query_states, dim=2) + key_states = torch.cat(key_states, dim=2) + value_states = torch.cat(value_states, dim=2) + + dummy_tensor = torch.zeros( + query_states.shape[0], + query_states.shape[2], + query_states.shape[-1], + device=query_states.device, + dtype=query_states.dtype, + ) + cos, sin = self.paligemma.model.language_model.rotary_emb(dummy_tensor, position_ids) + query_states, key_states = modeling_gemma.apply_rotary_pos_emb(query_states, key_states, cos, sin, unsqueeze_dim=1) + + batch_size = query_states.shape[0] + scaling = self.paligemma.language_model.layers[layer_idx].self_attn.scaling + + # Attention computation + att_output, _ = modeling_gemma.eager_attention_forward( + self.paligemma.language_model.layers[layer_idx].self_attn, + query_states, + key_states, + value_states, + attention_mask, + scaling, + ) + # Get head_dim from the current layer, not from the model + head_dim = self.paligemma.language_model.layers[layer_idx].self_attn.head_dim + att_output = att_output.reshape(batch_size, -1, 1 * 8 * head_dim) + + # Process layer outputs + outputs_embeds = [] + start_pos = 0 + for i, hidden_states in enumerate(inputs_embeds): + layer = models[i].layers[layer_idx] + end_pos = start_pos + 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) + out_emb = layer.self_attn.o_proj(att_output[:, start_pos:end_pos]) + + # first residual + out_emb = modeling_gemma._gated_residual(hidden_states, out_emb, gates[i]) # noqa: SLF001 + after_first_residual = out_emb.clone() + out_emb, gate = layer.post_attention_layernorm(out_emb, cond=adarms_cond[i]) + # Convert to bfloat16 if the next layer (mlp) uses bfloat16 + if layer.mlp.up_proj.weight.dtype == torch.bfloat16: + out_emb = out_emb.to(dtype=torch.bfloat16) + + out_emb = layer.mlp(out_emb) + # second residual + out_emb = modeling_gemma._gated_residual(after_first_residual, out_emb, gate) # noqa: SLF001 + outputs_embeds.append(out_emb) + start_pos = end_pos + + return outputs_embeds + + # Process all layers with gradient checkpointing if enabled + for layer_idx in range(num_layers): + if use_gradient_checkpointing: + inputs_embeds = torch.utils.checkpoint.checkpoint( + compute_layer_complete, + layer_idx, + inputs_embeds, + attention_mask, + position_ids, + adarms_cond, + use_reentrant=False, + preserve_rng_state=False, + ) + else: + inputs_embeds = compute_layer_complete(layer_idx, inputs_embeds, attention_mask, position_ids, adarms_cond) + + # Old code removed - now using compute_layer_complete function above + + # final norm + # Define final norm computation function for gradient checkpointing + def compute_final_norms(inputs_embeds, adarms_cond): + outputs_embeds = [] + for i, hidden_states in enumerate(inputs_embeds): + out_emb, _ = models[i].norm(hidden_states, cond=adarms_cond[i]) + outputs_embeds.append(out_emb) + return outputs_embeds + + # Apply gradient checkpointing to final norm if enabled + if use_gradient_checkpointing: + outputs_embeds = torch.utils.checkpoint.checkpoint(compute_final_norms, inputs_embeds, adarms_cond, use_reentrant=False, preserve_rng_state=False) + else: + outputs_embeds = compute_final_norms(inputs_embeds, adarms_cond) + + prefix_output = outputs_embeds[0] + suffix_output = outputs_embeds[1] + prefix_past_key_values = None + + return [prefix_output, suffix_output], prefix_past_key_values diff --git a/lightx2v/models/networks/openpi/image_tools.py b/lightx2v/models/networks/openpi/image_tools.py new file mode 100644 index 000000000..7d03e2957 --- /dev/null +++ b/lightx2v/models/networks/openpi/image_tools.py @@ -0,0 +1,47 @@ +"""Image helpers copied from OpenPI's torch preprocessing path (Apache-2.0).""" + +import torch +import torch.nn.functional as F # noqa: N812 + + +def resize_with_pad_torch(images: torch.Tensor, height: int, width: int, mode: str = "bilinear") -> torch.Tensor: + """Resize without distortion and pad with black / -1, preserving layout.""" + input_was_unbatched = images.dim() == 3 + channels_last = images.shape[-1] <= 4 + if input_was_unbatched: + images = images.unsqueeze(0) + if channels_last: + images = images.permute(0, 3, 1, 2) + + _, _, current_height, current_width = images.shape + ratio = max(current_width / width, current_height / height) + resized_height = int(current_height / ratio) + resized_width = int(current_width / ratio) + resized = F.interpolate( + images, + size=(resized_height, resized_width), + mode=mode, + align_corners=False if mode == "bilinear" else None, + ) + if images.dtype == torch.uint8: + resized = torch.round(resized).clamp(0, 255).to(torch.uint8) + pad_value = 0 + elif images.dtype == torch.float32: + resized = resized.clamp(-1.0, 1.0) + pad_value = -1.0 + else: + raise ValueError(f"Unsupported image dtype: {images.dtype}") + + pad_h0, remainder_h = divmod(height - resized_height, 2) + pad_w0, remainder_w = divmod(width - resized_width, 2) + resized = F.pad( + resized, + (pad_w0, pad_w0 + remainder_w, pad_h0, pad_h0 + remainder_h), + mode="constant", + value=pad_value, + ) + if channels_last: + resized = resized.permute(0, 2, 3, 1) + if input_was_unbatched: + resized = resized.squeeze(0) + return resized diff --git a/lightx2v/models/networks/openpi/infer/__init__.py b/lightx2v/models/networks/openpi/infer/__init__.py new file mode 100644 index 000000000..370eaf0f9 --- /dev/null +++ b/lightx2v/models/networks/openpi/infer/__init__.py @@ -0,0 +1,5 @@ +from .post_infer import OpenPIPostInfer +from .pre_infer import OpenPIPreInfer +from .transformer_infer import OpenPITransformerInfer + +__all__ = ["OpenPIPostInfer", "OpenPIPreInfer", "OpenPITransformerInfer"] diff --git a/lightx2v/models/networks/openpi/infer/post_infer.py b/lightx2v/models/networks/openpi/infer/post_infer.py new file mode 100644 index 000000000..88a80b46a --- /dev/null +++ b/lightx2v/models/networks/openpi/infer/post_infer.py @@ -0,0 +1,31 @@ +"""Convert padded normalized model actions back to LIBERO's 7-D space.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import torch + +from .pre_infer import load_norm_stats + + +class OpenPIPostInfer: + def __init__(self, norm_stats_path: str | Path, output_action_dim: int = 7): + self.stats = load_norm_stats(norm_stats_path)["actions"] + self.output_action_dim = int(output_action_dim) + + def infer(self, actions: torch.Tensor | np.ndarray) -> np.ndarray: + if isinstance(actions, torch.Tensor): + actions = actions.detach().to(torch.float32).cpu().numpy() + actions = np.asarray(actions, dtype=np.float32) + if actions.ndim == 3: + if actions.shape[0] != 1: + raise ValueError(f"Only batch size 1 is supported by the policy API, got {actions.shape}") + actions = actions[0] + if actions.ndim != 2 or actions.shape[-1] < self.output_action_dim: + raise ValueError(f"Expected [horizon, padded_action_dim], got {actions.shape}") + q01 = self.stats["q01"][: self.output_action_dim] + q99 = self.stats["q99"][: self.output_action_dim] + physical = (actions[:, : self.output_action_dim] + 1.0) / 2.0 * (q99 - q01 + 1e-6) + q01 + return np.asarray(physical, dtype=np.float32) diff --git a/lightx2v/models/networks/openpi/infer/pre_infer.py b/lightx2v/models/networks/openpi/infer/pre_infer.py new file mode 100644 index 000000000..2694bc5ca --- /dev/null +++ b/lightx2v/models/networks/openpi/infer/pre_infer.py @@ -0,0 +1,158 @@ +"""LIBERO input construction for the native PyTorch OpenPI backend.""" + +from __future__ import annotations + +import json +import logging +from pathlib import Path +from typing import Any + +import numpy as np +import sentencepiece +import torch +from PIL import Image + +from ..observation import Observation + +LOGGER = logging.getLogger(__name__) +IMAGE_SIZE = 224 +LIBERO_STATE_DIM = 8 + + +def load_norm_stats(path: str | Path) -> dict[str, dict[str, np.ndarray]]: + with Path(path).open("r", encoding="utf-8") as handle: + payload = json.load(handle) + payload = payload.get("norm_stats", payload) + if "state" not in payload or "actions" not in payload: + raise ValueError(f"Invalid LIBERO norm_stats file: {path}") + # Keep JSON's float64 precision here. Upstream OpenPI constructs NormStats + # with ``np.asarray`` as well; casting the quantiles early changes the + # normalized state by roughly 1e-7. + return {key: {stat: np.asarray(value) for stat, value in stats.items()} for key, stats in payload.items()} + + +def normalize_quantile(value: np.ndarray, stats: dict[str, np.ndarray]) -> np.ndarray: + q01 = stats["q01"][: value.shape[-1]] + q99 = stats["q99"][: value.shape[-1]] + return (value - q01) / (q99 - q01 + 1e-6) * 2.0 - 1.0 + + +def _parse_rgb(image: Any) -> np.ndarray: + if isinstance(image, (str, Path)): + with Image.open(image) as pil: + array = np.asarray(pil.convert("RGB")) + else: + array = np.asarray(image) + if array.ndim != 3: + raise ValueError(f"Expected a 3-D RGB image, got {array.shape}") + if array.shape[0] == 3 and array.shape[-1] != 3: + array = np.moveaxis(array, 0, -1) + if array.shape[-1] != 3: + raise ValueError(f"Expected an HWC RGB image, got {array.shape}") + if np.issubdtype(array.dtype, np.floating): + # This matches LiberoInputs: floating environment frames are [0, 1]. + array = np.clip(array, 0.0, 1.0) + array = np.rint(array * 255.0).astype(np.uint8) + else: + array = np.clip(array, 0, 255).astype(np.uint8, copy=False) + return np.ascontiguousarray(array) + + +def _resize_with_pad(image: np.ndarray, size: int = IMAGE_SIZE) -> np.ndarray: + height, width = image.shape[:2] + if (height, width) == (size, size): + return np.array(image, copy=True) + ratio = max(width / size, height / size) + resized_height = int(height / ratio) + resized_width = int(width / ratio) + resized = Image.fromarray(image, mode="RGB").resize((resized_width, resized_height), resample=Image.BILINEAR) + canvas = Image.new("RGB", (size, size), 0) + canvas.paste(resized, ((size - resized_width) // 2, (size - resized_height) // 2)) + return np.asarray(canvas, dtype=np.uint8).copy() + + +class PaligemmaTokenizer: + """Local-file PaliGemma SentencePiece tokenizer matching OpenPI.""" + + def __init__(self, model_path: str | Path, max_len: int = 200): + self.max_len = int(max_len) + model_path = Path(model_path) + if not model_path.is_file(): + raise FileNotFoundError(f"PaliGemma tokenizer not found: {model_path}") + self.processor = sentencepiece.SentencePieceProcessor(model_proto=model_path.read_bytes()) + + def tokenize(self, prompt: str, state: np.ndarray | None = None) -> tuple[np.ndarray, np.ndarray]: + cleaned = str(prompt).strip().replace("_", " ").replace("\n", " ") + if state is not None: + bins = np.linspace(-1, 1, 257)[:-1] + discrete = np.digitize(state, bins=bins) - 1 + state_text = " ".join(map(str, discrete)) + text = f"Task: {cleaned}, State: {state_text};\nAction: " + tokens = self.processor.encode(text, add_bos=True) + else: + tokens = self.processor.encode(cleaned, add_bos=True) + self.processor.encode("\n") + if len(tokens) > self.max_len: + LOGGER.warning("Prompt uses %d tokens; truncating to %d", len(tokens), self.max_len) + tokens = tokens[: self.max_len] + mask = [True] * len(tokens) + padding = self.max_len - len(tokens) + tokens += [0] * padding + mask += [False] * padding + return np.asarray(tokens, dtype=np.int64), np.asarray(mask, dtype=np.bool_) + + +class OpenPIPreInfer: + """Map logical LIBERO observations to the model's padded torch tensors.""" + + def __init__( + self, + norm_stats_path: str | Path, + tokenizer_path: str | Path, + device: torch.device | str, + action_dim: int = 32, + max_token_len: int = 200, + discrete_state_input: bool = False, + ): + self.norm_stats = load_norm_stats(norm_stats_path) + self.tokenizer = PaligemmaTokenizer(tokenizer_path, max_len=max_token_len) + self.device = torch.device(device) + self.action_dim = int(action_dim) + self.discrete_state_input = bool(discrete_state_input) + + def infer(self, images: dict, state: Any, task_description: str) -> Observation: + try: + agentview = images["agentview"] + wrist = images["wrist"] + except KeyError as exc: + raise KeyError("OpenPI expects images with logical keys 'agentview' and 'wrist'") from exc + + base = _resize_with_pad(_parse_rgb(agentview)) + left_wrist = _resize_with_pad(_parse_rgb(wrist)) + right_wrist = np.zeros_like(base) + + raw_state = np.asarray(state, dtype=np.float32).reshape(-1) + if raw_state.shape != (LIBERO_STATE_DIM,): + raise ValueError(f"pi05_libero expects an 8-D state, got {raw_state.shape}") + normalized_state = normalize_quantile(raw_state, self.norm_stats["state"]) + padded_state = np.pad(normalized_state, (0, self.action_dim - LIBERO_STATE_DIM)) + + tokenizer_state = normalized_state if self.discrete_state_input else None + tokens, token_mask = self.tokenizer.tokenize(task_description, tokenizer_state) + + image_arrays = { + "base_0_rgb": base, + "left_wrist_0_rgb": left_wrist, + "right_wrist_0_rgb": right_wrist, + } + data = { + "image": {key: torch.from_numpy(value).unsqueeze(0).to(self.device) for key, value in image_arrays.items()}, + "image_mask": { + "base_0_rgb": torch.ones(1, dtype=torch.bool, device=self.device), + "left_wrist_0_rgb": torch.ones(1, dtype=torch.bool, device=self.device), + "right_wrist_0_rgb": torch.zeros(1, dtype=torch.bool, device=self.device), + }, + "state": torch.from_numpy(padded_state).unsqueeze(0).to(self.device), + "tokenized_prompt": torch.from_numpy(tokens).unsqueeze(0).to(self.device), + "tokenized_prompt_mask": torch.from_numpy(token_mask).unsqueeze(0).to(self.device), + } + return Observation.from_dict(data) diff --git a/lightx2v/models/networks/openpi/infer/transformer_infer.py b/lightx2v/models/networks/openpi/infer/transformer_infer.py new file mode 100644 index 000000000..7a7182e20 --- /dev/null +++ b/lightx2v/models/networks/openpi/infer/transformer_infer.py @@ -0,0 +1,24 @@ +"""Flow-matching action sampler for OpenPI.""" + +from __future__ import annotations + +import torch + +from ..observation import Observation + + +class OpenPITransformerInfer: + def __init__(self, num_steps: int = 10): + if int(num_steps) <= 0: + raise ValueError("num_steps must be positive") + self.num_steps = int(num_steps) + + @torch.no_grad() + def infer( + self, + model, + observation: Observation, + device: torch.device | str, + noise: torch.Tensor | None = None, + ) -> torch.Tensor: + return model.sample_actions(device, observation, noise=noise, num_steps=self.num_steps) diff --git a/lightx2v/models/networks/openpi/model.py b/lightx2v/models/networks/openpi/model.py new file mode 100644 index 000000000..700f71922 --- /dev/null +++ b/lightx2v/models/networks/openpi/model.py @@ -0,0 +1,193 @@ +"""LightX2V-native wrapper around the official PyTorch pi0.5 architecture.""" + +from __future__ import annotations + +import json +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from torch import nn + +from .config import Pi0Config +from .infer import OpenPIPostInfer, OpenPIPreInfer, OpenPITransformerInfer +from .weights import load_pi05_libero_weights + + +def _read_config(config: Mapping[str, Any] | str | Path) -> tuple[dict[str, Any], Path | None]: + if isinstance(config, (str, Path)): + path = Path(config).expanduser().resolve() + with path.open("r", encoding="utf-8") as handle: + return json.load(handle), path.parent + if isinstance(config, Mapping): + return dict(config), None + try: + return dict(config), None + except (TypeError, ValueError) as exc: + raise TypeError("OpenPI config must be a mapping or a JSON path") from exc + + +def _resolve_path(value: str | Path | None, base_dir: Path | None) -> Path | None: + if value is None or str(value).strip() == "": + return None + path = Path(value).expanduser() + if not path.is_absolute() and base_dir is not None: + path = base_dir / path + return path.resolve() + + +class OpenPIModel(nn.Module): + """Native model/pipeline split into pre-, transformer-, post-infer stages. + + Runtime imports no code from ``/data/liuhongda/openpi`` and requires no + JAX stack. The one-time checkpoint conversion remains an offline step. + """ + + def __init__( + self, + core_model: nn.Module, + pre_infer: OpenPIPreInfer, + transformer_infer: OpenPITransformerInfer, + post_infer: OpenPIPostInfer, + model_config: Pi0Config, + device: torch.device | str, + seed: int | None = 0, + ): + super().__init__() + self.core_model = core_model + self.pre_infer = pre_infer + self.transformer_infer = transformer_infer + self.post_infer = post_infer + self.model_config = model_config + self.device = torch.device(device) + self.seed = None if seed is None or int(seed) < 0 else int(seed) + self._generator: torch.Generator | None = None + self.reset() + + @classmethod + def from_config(cls, config: Mapping[str, Any] | str | Path) -> "OpenPIModel": + values, base_dir = _read_config(config) + model_values = dict(values.get("model", {})) + # Flat LightX2V configs remain supported; nested model values win. + model_values = {**values, **model_values} + model_config = Pi0Config.from_mapping(model_values) + model_config.validate_pi05_libero() + + checkpoint_value = values.get("checkpoint_dir", values.get("model_path")) + checkpoint_path = _resolve_path(checkpoint_value, base_dir) + # A runner/ROS --model_path override must win over the static JSON + # weight_path. With no override, the explicit JSON weight_path is used. + if checkpoint_path is not None: + weight_path = checkpoint_path if checkpoint_path.suffix == ".safetensors" else checkpoint_path / "model.safetensors" + else: + weight_path = _resolve_path(values.get("weight_path"), base_dir) + if weight_path is None: + raise ValueError("OpenPI config requires checkpoint_dir/model_path or weight_path") + checkpoint_dir = weight_path.parent + + norm_stats_path = _resolve_path(values.get("norm_stats_path"), base_dir) + if norm_stats_path is None: + norm_stats_path = checkpoint_dir / "assets/physical-intelligence/libero/norm_stats.json" + tokenizer_path = _resolve_path(values.get("tokenizer_path"), base_dir) + if tokenizer_path is None: + tokenizer_candidates = [ + checkpoint_dir / "assets/paligemma_tokenizer.model", + checkpoint_dir / "paligemma_tokenizer.model", + ] + tokenizer_path = next((path for path in tokenizer_candidates if path.is_file()), tokenizer_candidates[0]) + + device = torch.device(values.get("device", "cuda")) + if device.type == "cuda" and not torch.cuda.is_available(): + raise RuntimeError("OpenPI config requests CUDA, but torch.cuda.is_available() is false") + num_steps = int(values.get("num_inference_steps", values.get("num_flow_steps", values.get("num_steps", 10)))) + output_action_dim = int(values.get("output_action_dim", 7)) + if output_action_dim != 7: + raise ValueError("The released pi05_libero policy must output 7-D LIBERO actions") + + core_model = load_pi05_libero_weights(weight_path, model_config, device) + return cls( + core_model=core_model, + pre_infer=OpenPIPreInfer( + norm_stats_path=norm_stats_path, + tokenizer_path=tokenizer_path, + device=device, + action_dim=model_config.action_dim, + max_token_len=model_config.max_token_len, + discrete_state_input=model_config.discrete_state_input, + ), + transformer_infer=OpenPITransformerInfer(num_steps=num_steps), + post_infer=OpenPIPostInfer(norm_stats_path, output_action_dim=output_action_dim), + model_config=model_config, + device=device, + seed=values.get("seed", 0), + ) + + def _make_generator(self, seed: int) -> torch.Generator: + generator = torch.Generator(device=self.device) + generator.manual_seed(int(seed)) + return generator + + def reset(self) -> None: + self._generator = None if self.seed is None else self._make_generator(self.seed) + + def _sample_noise(self, seed: int | None = None) -> torch.Tensor: + if seed is not None: + generator = self._make_generator(int(seed)) + else: + generator = self._generator + return torch.randn( + (1, self.model_config.action_horizon, self.model_config.action_dim), + dtype=torch.float32, + device=self.device, + generator=generator, + ) + + @torch.no_grad() + def predict_normalized_action_chunk( + self, + images: dict, + state, + task_description: str, + *, + seed: int | None = None, + noise: torch.Tensor | np.ndarray | None = None, + ) -> torch.Tensor: + observation = self.pre_infer.infer(images, state, task_description) + if noise is None: + noise_tensor = self._sample_noise(seed) + else: + noise_tensor = torch.as_tensor(noise, dtype=torch.float32, device=self.device) + if noise_tensor.ndim == 2: + noise_tensor = noise_tensor.unsqueeze(0) + expected = (1, self.model_config.action_horizon, self.model_config.action_dim) + if tuple(noise_tensor.shape) != expected: + raise ValueError(f"Noise must have shape {expected}, got {tuple(noise_tensor.shape)}") + return self.transformer_infer.infer(self.core_model, observation, self.device, noise=noise_tensor) + + @torch.no_grad() + def predict_action_chunk( + self, + images: dict, + state, + task_description: str, + seed: int | None = None, + ) -> np.ndarray: + normalized = self.predict_normalized_action_chunk(images, state, task_description, seed=seed) + actions = self.post_infer.infer(normalized) + expected = (self.model_config.action_horizon, 7) + if actions.shape != expected: + raise RuntimeError(f"OpenPI returned {actions.shape}; expected {expected}") + return actions + + def next_action(self, images: dict, state, task_description: str) -> np.ndarray: + return self.predict_action_chunk(images, state, task_description)[0] + + def forward(self, observation, actions, noise=None, time=None): + """Expose the official training loss for future LightX2V fine-tuning.""" + return self.core_model(observation, actions, noise=noise, time=time) + + def close(self) -> None: + if self.device.type == "cuda": + torch.cuda.empty_cache() diff --git a/lightx2v/models/networks/openpi/observation.py b/lightx2v/models/networks/openpi/observation.py new file mode 100644 index 000000000..a4d70a38a --- /dev/null +++ b/lightx2v/models/networks/openpi/observation.py @@ -0,0 +1,57 @@ +"""Torch-only observation container matching OpenPI's public tensor layout.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + + +@dataclass +class Observation: + # Images are float32 in [-1, 1], normally in BCHW layout for PyTorch. + images: dict[str, torch.Tensor] + image_masks: dict[str, torch.Tensor] + state: torch.Tensor + tokenized_prompt: torch.Tensor | None = None + tokenized_prompt_mask: torch.Tensor | None = None + token_ar_mask: torch.Tensor | None = None + token_loss_mask: torch.Tensor | None = None + + @classmethod + def from_dict(cls, data: dict) -> "Observation": + if ("tokenized_prompt" in data) != ("tokenized_prompt_mask" in data): + raise ValueError("tokenized_prompt and tokenized_prompt_mask must be provided together") + + images: dict[str, torch.Tensor] = {} + for key, value in data["image"].items(): + image = value + if image.dtype == torch.uint8: + if image.ndim != 4 or image.shape[-1] != 3: + raise ValueError(f"uint8 image {key!r} must use BHWC layout, got {tuple(image.shape)}") + image = image.to(torch.float32).permute(0, 3, 1, 2) / 255.0 * 2.0 - 1.0 + images[key] = image + + return cls( + images=images, + image_masks=data["image_mask"], + state=data["state"], + tokenized_prompt=data.get("tokenized_prompt"), + tokenized_prompt_mask=data.get("tokenized_prompt_mask"), + token_ar_mask=data.get("token_ar_mask"), + token_loss_mask=data.get("token_loss_mask"), + ) + + def to(self, device: torch.device | str) -> "Observation": + def move(value): + return None if value is None else value.to(device) + + return Observation( + images={key: value.to(device) for key, value in self.images.items()}, + image_masks={key: value.to(device) for key, value in self.image_masks.items()}, + state=self.state.to(device), + tokenized_prompt=move(self.tokenized_prompt), + tokenized_prompt_mask=move(self.tokenized_prompt_mask), + token_ar_mask=move(self.token_ar_mask), + token_loss_mask=move(self.token_loss_mask), + ) diff --git a/lightx2v/models/networks/openpi/pi0.py b/lightx2v/models/networks/openpi/pi0.py new file mode 100644 index 000000000..359fd7f72 --- /dev/null +++ b/lightx2v/models/networks/openpi/pi0.py @@ -0,0 +1,451 @@ +# Adapted from Physical Intelligence OpenPI (Apache-2.0), commit 15a9616. +# Localized for the LightX2V OpenPI backend; parameter names intentionally unchanged. + +import logging +import math + +import torch +import torch.nn.functional as F # noqa: N812 +from torch import Tensor, nn + +from . import config as _gemma +from . import preprocessing as _preprocessing +from .gemma import PaliGemmaWithExpertModel + + +def get_safe_dtype(target_dtype, device_type): + """Get a safe dtype for the given device type.""" + if device_type == "cpu": + # CPU doesn't support bfloat16, use float32 instead + if target_dtype == torch.bfloat16: + return torch.float32 + if target_dtype == torch.float64: + return torch.float64 + return target_dtype + + +def create_sinusoidal_pos_embedding(time: torch.tensor, dimension: int, min_period: float, max_period: float, device="cpu") -> 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 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] + return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1) + + +def sample_beta(alpha, beta, bsize, device): + alpha_t = torch.as_tensor(alpha, dtype=torch.float32, device=device) + beta_t = torch.as_tensor(beta, dtype=torch.float32, device=device) + dist = torch.distributions.Beta(alpha_t, beta_t) + return dist.sample((bsize,)) + + +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] + return att_2d_masks & pad_2d_masks + + +class PI0Pytorch(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.pi05 = config.pi05 + + paligemma_config = _gemma.get_config(config.paligemma_variant) + action_expert_config = _gemma.get_config(config.action_expert_variant) + + self.paligemma_with_expert = PaliGemmaWithExpertModel( + paligemma_config, + action_expert_config, + use_adarms=[False, True] if self.pi05 else [False, False], + precision=config.dtype, + ) + + self.action_in_proj = nn.Linear(config.action_dim, action_expert_config.width) + self.action_out_proj = nn.Linear(action_expert_config.width, config.action_dim) + + if self.pi05: + self.time_mlp_in = nn.Linear(action_expert_config.width, action_expert_config.width) + self.time_mlp_out = nn.Linear(action_expert_config.width, action_expert_config.width) + else: + self.state_proj = nn.Linear(config.action_dim, action_expert_config.width) + self.action_time_mlp_in = nn.Linear(2 * action_expert_config.width, action_expert_config.width) + self.action_time_mlp_out = nn.Linear(action_expert_config.width, action_expert_config.width) + + torch.set_float32_matmul_precision("high") + if config.pytorch_compile_mode is not None: + self.sample_actions = torch.compile(self.sample_actions, mode=config.pytorch_compile_mode) + + # Initialize gradient checkpointing flag + self.gradient_checkpointing_enabled = False + + msg = "transformers_replace is not installed correctly. Please install it with `uv pip install transformers==4.53.2` and `cp -r ./src/openpi/models_pytorch/transformers_replace/* .venv/lib/python3.11/site-packages/transformers/`." + try: + from transformers.models.siglip import check + + if not check.check_whether_transformers_replace_is_installed_correctly(): + raise ValueError(msg) + except ImportError: + raise ValueError(msg) from None + + def gradient_checkpointing_enable(self): + """Enable gradient checkpointing for memory optimization.""" + self.gradient_checkpointing_enabled = True + self.paligemma_with_expert.paligemma.language_model.gradient_checkpointing = True + self.paligemma_with_expert.paligemma.vision_tower.gradient_checkpointing = True + self.paligemma_with_expert.gemma_expert.model.gradient_checkpointing = True + + logging.info("Enabled gradient checkpointing for PI0Pytorch model") + + def gradient_checkpointing_disable(self): + """Disable gradient checkpointing.""" + self.gradient_checkpointing_enabled = False + self.paligemma_with_expert.paligemma.language_model.gradient_checkpointing = False + self.paligemma_with_expert.paligemma.vision_tower.gradient_checkpointing = False + self.paligemma_with_expert.gemma_expert.model.gradient_checkpointing = False + + logging.info("Disabled gradient checkpointing for PI0Pytorch model") + + def is_gradient_checkpointing_enabled(self): + """Check if gradient checkpointing is enabled.""" + return self.gradient_checkpointing_enabled + + def _apply_checkpoint(self, func, *args, **kwargs): + """Helper method to apply gradient checkpointing if enabled.""" + if self.gradient_checkpointing_enabled and self.training: + return torch.utils.checkpoint.checkpoint(func, *args, use_reentrant=False, preserve_rng_state=False, **kwargs) + return func(*args, **kwargs) + + def _prepare_attention_masks_4d(self, att_2d_masks): + """Helper method to prepare 4D attention masks for transformer.""" + att_2d_masks_4d = att_2d_masks[:, None, :, :] + return torch.where(att_2d_masks_4d, 0.0, -2.3819763e38) + + def _preprocess_observation(self, observation, *, train=True): + """Helper method to preprocess observation.""" + observation = _preprocessing.preprocess_observation_pytorch(observation, train=train) + return ( + list(observation.images.values()), + list(observation.image_masks.values()), + observation.tokenized_prompt, + observation.tokenized_prompt_mask, + observation.state, + ) + + def sample_noise(self, shape, device): + return torch.normal( + mean=0.0, + std=1.0, + size=shape, + dtype=torch.float32, + device=device, + ) + + def sample_time(self, bsize, device): + time_beta = sample_beta(1.5, 1.0, bsize, device) + time = time_beta * 0.999 + 0.001 + return time.to(dtype=torch.float32, device=device) + + def embed_prefix(self, images, img_masks, lang_tokens, lang_masks) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Embed images with SigLIP and language tokens with embedding layer to prepare + for PaliGemma transformer processing. + """ + embs = [] + pad_masks = [] + att_masks = [] + + # Process images + for img, img_mask in zip(images, img_masks, strict=True): + + def image_embed_func(img): + return self.paligemma_with_expert.embed_image(img) + + img_emb = self._apply_checkpoint(image_embed_func, img) + + bsize, num_img_embs = img_emb.shape[:2] + + embs.append(img_emb) + pad_masks.append(img_mask[:, None].expand(bsize, num_img_embs)) + + # Create attention masks so that image tokens attend to each other + att_masks += [0] * num_img_embs + + # Process language tokens + def lang_embed_func(lang_tokens): + lang_emb = self.paligemma_with_expert.embed_language_tokens(lang_tokens) + lang_emb_dim = lang_emb.shape[-1] + return lang_emb * math.sqrt(lang_emb_dim) + + lang_emb = self._apply_checkpoint(lang_embed_func, lang_tokens) + + embs.append(lang_emb) + pad_masks.append(lang_masks) + + # full attention between image and language inputs + num_lang_embs = lang_emb.shape[1] + att_masks += [0] * num_lang_embs + + 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) + + # Get batch size from the first dimension of the concatenated tensors + bsize = pad_masks.shape[0] + att_masks = att_masks[None, :].expand(bsize, len(att_masks)) + + return embs, pad_masks, att_masks + + def embed_suffix(self, state, noisy_actions, timestep): + """Embed state, noisy_actions, timestep to prepare for Expert Gemma processing.""" + embs = [] + pad_masks = [] + att_masks = [] + + if not self.pi05: + if self.state_proj.weight.dtype == torch.float32: + state = state.to(torch.float32) + + # Embed state + def state_proj_func(state): + return self.state_proj(state) + + state_emb = self._apply_checkpoint(state_proj_func, state) + + embs.append(state_emb[:, None, :]) + bsize = state_emb.shape[0] + device = state_emb.device + + state_mask = torch.ones(bsize, 1, dtype=torch.bool, device=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] + + # Embed timestep using sine-cosine positional encoding with sensitivity in the range [0, 1] + time_emb = create_sinusoidal_pos_embedding(timestep, self.action_in_proj.out_features, min_period=4e-3, max_period=4.0, device=timestep.device) + time_emb = time_emb.type(dtype=timestep.dtype) + + # Fuse timestep + action information using an MLP + def action_proj_func(noisy_actions): + return self.action_in_proj(noisy_actions) + + action_emb = self._apply_checkpoint(action_proj_func, noisy_actions) + + if not self.pi05: + time_emb = time_emb[:, None, :].expand_as(action_emb) + action_time_emb = torch.cat([action_emb, time_emb], dim=2) + + # Apply MLP layers + def mlp_func(action_time_emb): + x = self.action_time_mlp_in(action_time_emb) + x = F.silu(x) # swish == silu + return self.action_time_mlp_out(x) + + action_time_emb = self._apply_checkpoint(mlp_func, action_time_emb) + adarms_cond = None + else: + # time MLP (for adaRMS) + def time_mlp_func(time_emb): + x = self.time_mlp_in(time_emb) + x = F.silu(x) # swish == silu + x = self.time_mlp_out(x) + return F.silu(x) + + time_emb = self._apply_checkpoint(time_mlp_func, time_emb) + action_time_emb = action_emb + adarms_cond = 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=timestep.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] + ([0] * (self.config.action_horizon - 1)) + + 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, adarms_cond + + def forward(self, observation, actions, noise=None, time=None) -> Tensor: + """Do a full training forward pass and compute the loss (batch_size x num_steps x num_motors)""" + images, img_masks, lang_tokens, lang_masks, state = self._preprocess_observation(observation, train=True) + + 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) + suffix_embs, suffix_pad_masks, suffix_att_masks, adarms_cond = self.embed_suffix(state, x_t, time) + if self.paligemma_with_expert.paligemma.language_model.layers[0].self_attn.q_proj.weight.dtype == torch.bfloat16: + suffix_embs = suffix_embs.to(dtype=torch.bfloat16) + prefix_embs = prefix_embs.to(dtype=torch.bfloat16) + + 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 = torch.cumsum(pad_masks, dim=1) - 1 + + # Prepare attention masks + att_2d_masks_4d = self._prepare_attention_masks_4d(att_2d_masks) + + # Apply gradient checkpointing if enabled + def forward_func(prefix_embs, suffix_embs, att_2d_masks_4d, position_ids, adarms_cond): + (_, suffix_out), _ = self.paligemma_with_expert.forward( + attention_mask=att_2d_masks_4d, + position_ids=position_ids, + past_key_values=None, + inputs_embeds=[prefix_embs, suffix_embs], + use_cache=False, + adarms_cond=[None, adarms_cond], + ) + return suffix_out + + suffix_out = self._apply_checkpoint(forward_func, prefix_embs, suffix_embs, att_2d_masks_4d, position_ids, adarms_cond) + + suffix_out = suffix_out[:, -self.config.action_horizon :] + suffix_out = suffix_out.to(dtype=torch.float32) + + # Apply gradient checkpointing to final action projection if enabled + def action_out_proj_func(suffix_out): + return self.action_out_proj(suffix_out) + + v_t = self._apply_checkpoint(action_out_proj_func, suffix_out) + + return F.mse_loss(u_t, v_t, reduction="none") + + @torch.no_grad() + def sample_actions(self, device, observation, noise=None, num_steps=10) -> Tensor: + """Do a full inference forward and compute the action (batch_size x num_steps x num_motors)""" + bsize = observation.state.shape[0] + if noise is None: + actions_shape = (bsize, self.config.action_horizon, self.config.action_dim) + noise = self.sample_noise(actions_shape, device) + + images, img_masks, lang_tokens, lang_masks, state = self._preprocess_observation(observation, train=False) + + prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(images, img_masks, lang_tokens, lang_masks) + prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks) + prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1 + + # Compute image and language key value cache + prefix_att_2d_masks_4d = self._prepare_attention_masks_4d(prefix_att_2d_masks) + self.paligemma_with_expert.paligemma.language_model.config._attn_implementation = "eager" # noqa: SLF001 + + _, past_key_values = self.paligemma_with_expert.forward( + attention_mask=prefix_att_2d_masks_4d, + position_ids=prefix_position_ids, + past_key_values=None, + inputs_embeds=[prefix_embs, None], + use_cache=True, + ) + + dt = -1.0 / 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( + state, + prefix_pad_masks, + past_key_values, + x_t, + expanded_time, + ) + + # Euler step - use new tensor assignment instead of in-place operation + x_t = x_t + dt * v_t + time += dt + return x_t + + def denoise_step( + self, + state, + 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, adarms_cond = self.embed_suffix(state, 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 = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1 + + # Prepare attention masks + full_att_2d_masks_4d = self._prepare_attention_masks_4d(full_att_2d_masks) + self.paligemma_with_expert.gemma_expert.model.config._attn_implementation = "eager" # noqa: SLF001 + + outputs_embeds, _ = self.paligemma_with_expert.forward( + attention_mask=full_att_2d_masks_4d, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=[None, suffix_embs], + use_cache=False, + adarms_cond=[None, adarms_cond], + ) + + suffix_out = outputs_embeds[1] + suffix_out = suffix_out[:, -self.config.action_horizon :] + suffix_out = suffix_out.to(dtype=torch.float32) + return self.action_out_proj(suffix_out) diff --git a/lightx2v/models/networks/openpi/preprocessing.py b/lightx2v/models/networks/openpi/preprocessing.py new file mode 100644 index 000000000..6fa6c10ba --- /dev/null +++ b/lightx2v/models/networks/openpi/preprocessing.py @@ -0,0 +1,176 @@ +# Adapted from Physical Intelligence OpenPI (Apache-2.0), commit 15a9616. +# Localized for the LightX2V OpenPI backend. + +import logging +from collections.abc import Sequence + +import torch + +from . import image_tools + +logger = logging.getLogger("openpi") + +# Constants moved from model.py +IMAGE_KEYS = ( + "base_0_rgb", + "left_wrist_0_rgb", + "right_wrist_0_rgb", +) + +IMAGE_RESOLUTION = (224, 224) + + +def preprocess_observation_pytorch( + observation, + *, + train: bool = False, + image_keys: Sequence[str] = IMAGE_KEYS, + image_resolution: tuple[int, int] = IMAGE_RESOLUTION, +): + """Torch.compile-compatible version of preprocess_observation_pytorch with simplified type annotations. + + This function avoids complex type annotations that can cause torch.compile issues. + """ + if not set(image_keys).issubset(observation.images): + raise ValueError(f"images dict missing keys: expected {image_keys}, got {list(observation.images)}") + + batch_shape = observation.state.shape[:-1] + + out_images = {} + for key in image_keys: + image = observation.images[key] + + # TODO: This is a hack to handle both [B, C, H, W] and [B, H, W, C] formats + # Handle both [B, C, H, W] and [B, H, W, C] formats + is_channels_first = image.shape[1] == 3 # Check if channels are in dimension 1 + + if is_channels_first: + # Convert [B, C, H, W] to [B, H, W, C] for processing + image = image.permute(0, 2, 3, 1) + + if image.shape[1:3] != image_resolution: + logger.info(f"Resizing image {key} from {image.shape[1:3]} to {image_resolution}") + image = image_tools.resize_with_pad_torch(image, *image_resolution) + + if train: + # Convert from [-1, 1] to [0, 1] for PyTorch augmentations + image = image / 2.0 + 0.5 + + # Apply PyTorch-based augmentations + if "wrist" not in key: + # Geometric augmentations for non-wrist cameras + height, width = image.shape[1:3] + + # Random crop and resize + crop_height = int(height * 0.95) + crop_width = int(width * 0.95) + + # Random crop + max_h = height - crop_height + max_w = width - crop_width + if max_h > 0 and max_w > 0: + # Use tensor operations instead of .item() for torch.compile compatibility + start_h = torch.randint(0, max_h + 1, (1,), device=image.device) + start_w = torch.randint(0, max_w + 1, (1,), device=image.device) + image = image[:, start_h : start_h + crop_height, start_w : start_w + crop_width, :] + + # Resize back to original size + image = torch.nn.functional.interpolate( + image.permute(0, 3, 1, 2), # [b, h, w, c] -> [b, c, h, w] + size=(height, width), + mode="bilinear", + align_corners=False, + ).permute(0, 2, 3, 1) # [b, c, h, w] -> [b, h, w, c] + + # Random rotation (small angles) + # Use tensor operations instead of .item() for torch.compile compatibility + angle = torch.rand(1, device=image.device) * 10 - 5 # Random angle between -5 and 5 degrees + if torch.abs(angle) > 0.1: # Only rotate if angle is significant + # Convert to radians + angle_rad = angle * torch.pi / 180.0 + + # Create rotation matrix + cos_a = torch.cos(angle_rad) + sin_a = torch.sin(angle_rad) + + # Apply rotation using grid_sample + grid_x = torch.linspace(-1, 1, width, device=image.device) + grid_y = torch.linspace(-1, 1, height, device=image.device) + + # Create meshgrid + grid_y, grid_x = torch.meshgrid(grid_y, grid_x, indexing="ij") + + # Expand to batch dimension + grid_x = grid_x.unsqueeze(0).expand(image.shape[0], -1, -1) + grid_y = grid_y.unsqueeze(0).expand(image.shape[0], -1, -1) + + # Apply rotation transformation + grid_x_rot = grid_x * cos_a - grid_y * sin_a + grid_y_rot = grid_x * sin_a + grid_y * cos_a + + # Stack and reshape for grid_sample + grid = torch.stack([grid_x_rot, grid_y_rot], dim=-1) + + image = torch.nn.functional.grid_sample( + image.permute(0, 3, 1, 2), # [b, h, w, c] -> [b, c, h, w] + grid, + mode="bilinear", + padding_mode="zeros", + align_corners=False, + ).permute(0, 2, 3, 1) # [b, c, h, w] -> [b, h, w, c] + + # Color augmentations for all cameras + # Random brightness + # Use tensor operations instead of .item() for torch.compile compatibility + brightness_factor = 0.7 + torch.rand(1, device=image.device) * 0.6 # Random factor between 0.7 and 1.3 + image = image * brightness_factor + + # Random contrast + # Use tensor operations instead of .item() for torch.compile compatibility + contrast_factor = 0.6 + torch.rand(1, device=image.device) * 0.8 # Random factor between 0.6 and 1.4 + mean = image.mean(dim=[1, 2, 3], keepdim=True) + image = (image - mean) * contrast_factor + mean + + # Random saturation (convert to HSV, modify S, convert back) + # For simplicity, we'll just apply a random scaling to the color channels + # Use tensor operations instead of .item() for torch.compile compatibility + saturation_factor = 0.5 + torch.rand(1, device=image.device) * 1.0 # Random factor between 0.5 and 1.5 + gray = image.mean(dim=-1, keepdim=True) + image = gray + (image - gray) * saturation_factor + + # Clamp values to [0, 1] + image = torch.clamp(image, 0, 1) + + # Back to [-1, 1] + image = image * 2.0 - 1.0 + + # Convert back to [B, C, H, W] format if it was originally channels-first + if is_channels_first: + image = image.permute(0, 3, 1, 2) # [B, H, W, C] -> [B, C, H, W] + + out_images[key] = image + + # obtain mask + out_masks = {} + for key in out_images: + if key not in observation.image_masks: + # do not mask by default + out_masks[key] = torch.ones(batch_shape, dtype=torch.bool, device=observation.state.device) + else: + out_masks[key] = observation.image_masks[key] + + # Create a simple object with the required attributes instead of using the complex Observation class + class SimpleProcessedObservation: + def __init__(self, **kwargs): + for key, value in kwargs.items(): + setattr(self, key, value) + + return SimpleProcessedObservation( + images=out_images, + image_masks=out_masks, + state=observation.state, + tokenized_prompt=observation.tokenized_prompt, + tokenized_prompt_mask=observation.tokenized_prompt_mask, + token_ar_mask=observation.token_ar_mask, + token_loss_mask=observation.token_loss_mask, + ) diff --git a/lightx2v/models/networks/openpi/transformers_replace/models/gemma/configuration_gemma.py b/lightx2v/models/networks/openpi/transformers_replace/models/gemma/configuration_gemma.py new file mode 100644 index 000000000..ae7a9d4b5 --- /dev/null +++ b/lightx2v/models/networks/openpi/transformers_replace/models/gemma/configuration_gemma.py @@ -0,0 +1,174 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/gemma/modular_gemma.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_gemma.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# coding=utf-8 +# Copyright 2024 Google Inc. HuggingFace Inc. team. All rights reserved. +# +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import Optional + +from ...configuration_utils import PretrainedConfig + + +class GemmaConfig(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`GemmaModel`]. It is used to instantiate an Gemma + model according to the specified arguments, defining the model architecture. Instantiating a configuration with the + defaults will yield a similar configuration to that of the Gemma-7B. + e.g. [google/gemma-7b](https://huggingface.co/google/gemma-7b) + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + Args: + vocab_size (`int`, *optional*, defaults to 256000): + Vocabulary size of the Gemma model. Defines the number of different tokens that can be represented by the + `inputs_ids` passed when calling [`GemmaModel`] + hidden_size (`int`, *optional*, defaults to 3072): + Dimension of the hidden representations. + intermediate_size (`int`, *optional*, defaults to 24576): + Dimension of the MLP representations. + num_hidden_layers (`int`, *optional*, defaults to 28): + Number of hidden layers in the Transformer decoder. + num_attention_heads (`int`, *optional*, defaults to 16): + Number of attention heads for each attention layer in the Transformer decoder. + num_key_value_heads (`int`, *optional*, defaults to 16): + This is the number of key_value heads that should be used to implement Grouped Query Attention. If + `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if + `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When + converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed + by meanpooling all the original heads within that group. For more details, check out [this + paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to + `num_attention_heads`. + head_dim (`int`, *optional*, defaults to 256): + The attention head dimension. + hidden_act (`str` or `function`, *optional*, defaults to `"gelu_pytorch_tanh"`): + The legacy activation function. It is overwritten by the `hidden_activation`. + hidden_activation (`str` or `function`, *optional*): + The non-linear activation function (function or string) in the decoder. Will default to `"gelu_pytorch_tanh"` + if not specified. `"gelu_pytorch_tanh"` uses an approximation of the `"gelu"` activation function. + max_position_embeddings (`int`, *optional*, defaults to 8192): + The maximum sequence length that this model might ever be used with. + initializer_range (`float`, *optional*, defaults to 0.02): + The standard deviation of the truncated_normal_initializer for initializing all weight matrices. + rms_norm_eps (`float`, *optional*, defaults to 1e-06): + The epsilon used by the rms normalization layers. + use_cache (`bool`, *optional*, defaults to `True`): + Whether or not the model should return the last key/values attentions (not used by all models). Only + relevant if `config.is_decoder=True`. + pad_token_id (`int`, *optional*, defaults to 0): + Padding token id. + eos_token_id (`int`, *optional*, defaults to 1): + End of stream token id. + bos_token_id (`int`, *optional*, defaults to 2): + Beginning of stream token id. + tie_word_embeddings (`bool`, *optional*, defaults to `True`): + Whether to tie weight embeddings + rope_theta (`float`, *optional*, defaults to 10000.0): + The base period of the RoPE embeddings. + attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`): + Whether to use a bias in the query, key, value and output projection layers during self-attention. + attention_dropout (`float`, *optional*, defaults to 0.0): + The dropout ratio for the attention probabilities. + use_adarms (`bool`, *optional*, defaults to `False`): + Whether to use ADARMS. + adarms_cond_dim (`int`, *optional*, defaults to `None`): + The dimension of the ADARMS condition. + ```python + >>> from transformers import GemmaModel, GemmaConfig + >>> # Initializing a Gemma gemma-7b style configuration + >>> configuration = GemmaConfig() + >>> # Initializing a model from the gemma-7b style configuration + >>> model = GemmaModel(configuration) + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "gemma" + keys_to_ignore_at_inference = ["past_key_values"] + base_model_tp_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + } + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } + + def __init__( + self, + vocab_size=256000, + hidden_size=3072, + intermediate_size=24576, + num_hidden_layers=28, + num_attention_heads=16, + num_key_value_heads=16, + head_dim=256, + hidden_act="gelu_pytorch_tanh", + hidden_activation=None, + max_position_embeddings=8192, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + pad_token_id=0, + eos_token_id=1, + bos_token_id=2, + tie_word_embeddings=True, + rope_theta=10000.0, + attention_bias=False, + attention_dropout=0.0, + use_adarms: bool = False, + adarms_cond_dim: Optional[int] = None, + **kwargs, + ): + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.head_dim = head_dim + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.hidden_activation = hidden_activation + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.rope_theta = rope_theta + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.use_adarms = use_adarms + self.adarms_cond_dim = adarms_cond_dim + + # Set default for adarms_cond_dim if use_adarms is True + if self.use_adarms and self.adarms_cond_dim is None: + self.adarms_cond_dim = self.hidden_size + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + +__all__ = ["GemmaConfig"] diff --git a/lightx2v/models/networks/openpi/transformers_replace/models/gemma/modeling_gemma.py b/lightx2v/models/networks/openpi/transformers_replace/models/gemma/modeling_gemma.py new file mode 100644 index 000000000..29e3a2749 --- /dev/null +++ b/lightx2v/models/networks/openpi/transformers_replace/models/gemma/modeling_gemma.py @@ -0,0 +1,840 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/gemma/modular_gemma.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_gemma.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# coding=utf-8 +# Copyright 2024 Google Inc. HuggingFace Inc. team. All rights reserved. +# +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import Callable, Optional, Union + +import torch +from torch import nn + +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache +from ...generation import GenerationMixin +from ...masking_utils import create_causal_mask +from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import ( + BaseModelOutputWithPast, + CausalLMOutputWithPast, + SequenceClassifierOutputWithPast, + TokenClassifierOutput, +) +from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import LossKwargs, auto_docstring, can_return_tuple, logging +from .configuration_gemma import GemmaConfig + +logger = logging.get_logger(__name__) + + +class GemmaRMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6, cond_dim: Optional[int] = None): + super().__init__() + self.eps = eps + self.dim = dim + self.cond_dim = cond_dim + + # Dense layer for adaptive normalization (if cond_dim is provided) + if cond_dim is not None: + # self.dense = nn.Linear(cond_dim, dim * 3, bias=True, dtype=torch.bfloat16) + self.dense = nn.Linear(cond_dim, dim * 3, bias=True) + # Initialize with zeros (matches source implementation) + nn.init.zeros_(self.dense.weight) + else: + self.weight = nn.Parameter(torch.zeros(dim, dtype=torch.bfloat16)) + self.dense = None + + def _norm(self, x): + # Compute variance in float32 (like the source implementation) + var = torch.mean(torch.square(x.float()), dim=-1, keepdim=True) + # Compute normalization in float32 + normed_inputs = x * torch.rsqrt(var + self.eps) + return normed_inputs + + def forward(self, x, cond=None): + dtype = x.dtype # original dtype, could be half-precision + normed_inputs = self._norm(x) + + if cond is None or self.dense is None: + # regular RMSNorm + # scale by learned parameter in float32 (matches source implementation) + normed_inputs = normed_inputs * (1.0 + self.weight.float()) + return normed_inputs.to(dtype), None # return in original dtype with None gate + + # adaptive RMSNorm (if cond is provided and dense layer exists) + if cond.shape[-1] != self.cond_dim: + raise ValueError(f"Expected cond dimension {self.cond_dim}, got {cond.shape[-1]}") + + # self.dense.to(dtype=torch.bfloat16).to(dtype=torch.float32) + modulation = self.dense(cond) + # Reshape modulation to broadcast properly: [batch, 1, features] for [batch, seq, features] + if len(x.shape) == 3: # [batch, seq, features] + modulation = modulation.unsqueeze(1) + + scale, shift, gate = torch.chunk(modulation, 3, dim=-1) + + # Apply adaptive normalization: use model weight dtype to ensure compatibility + # model_dtype = self.dense.weight.dtype # Use the model's dtype (bfloat16) + # scale = scale.to(model_dtype) + # shift = shift.to(model_dtype) + # gate = gate.to(model_dtype) + # normed_inputs = normed_inputs.to(model_dtype) # Convert normed_inputs to model dtype + + normed_inputs = normed_inputs * (1 + scale.to(torch.float32)) + shift.to(torch.float32) + + return normed_inputs.to(dtype), gate.to(dtype) + + def extra_repr(self): + repr_str = f"{tuple(self.weight.shape)}, eps={self.eps}" + if self.dense is not None: + repr_str += f", adaptive=True, cond_dim={self.cond_dim}" + return repr_str + + +class GemmaMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +class GemmaRotaryEmbedding(nn.Module): + def __init__(self, config: GemmaConfig, device=None): + super().__init__() + # BC: "rope_type" was originally "type" + if hasattr(config, "rope_scaling") and config.rope_scaling is not None: + self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type")) + else: + self.rope_type = "default" + self.max_seq_len_cached = config.max_position_embeddings + self.original_max_seq_len = config.max_position_embeddings + + self.config = config + self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] + + inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.original_inv_freq = self.inv_freq + + @torch.no_grad() + @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) + def forward(self, x, position_ids): + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) + position_ids_expanded = position_ids[:, None, :].float() + + device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): # Force float32 + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + position_ids (`torch.Tensor`, *optional*): + Deprecated and unused. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def _gated_residual(x, y, gate): + """ + Applies gated residual connection with optional gate parameter. + + Args: + x: Input tensor (residual) + y: Output tensor to be added + gate: Optional gate tensor to modulate the addition + + Returns: + x + y if gate is None, otherwise x + y * gate + """ + if x is None and y is None: + return None + if x is None or y is None: + return x if x is not None else y + if gate is None: + return x + y + return x + y * gate + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: Optional[torch.Tensor], + scaling: float, + dropout: float = 0.0, + **kwargs, +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +class GemmaAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: GemmaConfig, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = True + + self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias) + self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias) + self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias) + self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor], + past_key_value: Optional[Cache] = None, + cache_position: Optional[torch.LongTensor] = None, + use_cache: bool = False, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + # Use cache if provided + if past_key_value is not None: + if use_cache: + # sin and cos are specific to RoPE models; cache_position needed for the static cache + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + else: + key_states = torch.cat([past_key_value[self.layer_idx][0], key_states], dim=2) + value_states = torch.cat([past_key_value[self.layer_idx][1], value_states], dim=2) + + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class GemmaDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: GemmaConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = GemmaAttention(config=config, layer_idx=layer_idx) + + self.mlp = GemmaMLP(config) + cond_dim = getattr(config, "adarms_cond_dim", None) if getattr(config, "use_adarms", False) else None + self.input_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps, cond_dim=cond_dim) + self.post_attention_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps, cond_dim=cond_dim) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + adarms_cond: Optional[torch.Tensor] = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: + residual = hidden_states + hidden_states, gate = self.input_layernorm(hidden_states, adarms_cond) + + # Self Attention + hidden_states, self_attn_weights = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = _gated_residual(residual, hidden_states, gate) + + # Fully Connected + residual = hidden_states + hidden_states, gate = self.post_attention_layernorm(hidden_states, adarms_cond) + hidden_states = self.mlp(hidden_states) + hidden_states = _gated_residual(residual, hidden_states, gate) + + outputs = (hidden_states,) + if output_attentions: + outputs += (self_attn_weights,) + + return outputs + + +@auto_docstring +class GemmaPreTrainedModel(PreTrainedModel): + config_class = GemmaConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["GemmaDecoderLayer"] + _skip_keys_device_placement = ["past_key_values"] + _supports_flash_attn_3 = True + _supports_flash_attn_2 = True + _supports_sdpa = True + _supports_flex_attn = True + _supports_cache_class = True + _supports_quantized_cache = True + _supports_static_cache = True + _supports_attention_backend = True + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, 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, GemmaRMSNorm): + if hasattr(module, "weight"): + module.weight.data.fill_(1.0) + + +@auto_docstring +class GemmaModel(GemmaPreTrainedModel): + def __init__(self, config: GemmaConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList([GemmaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) + + cond_dim = getattr(config, "adarms_cond_dim", None) if getattr(config, "use_adarms", False) else None + self.norm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps, cond_dim=cond_dim) + self.rotary_emb = GemmaRotaryEmbedding(config=config) + self.gradient_checkpointing = False + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + adarms_cond: Optional[torch.Tensor] = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> BaseModelOutputWithPast: + """ + adarms_cond (`torch.Tensor` of shape `(batch_size, cond_dim)`, *optional*): + Condition for ADARMS. + """ + 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 + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if self.gradient_checkpointing and self.training and use_cache: + logger.warning_once("`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`.") + use_cache = False + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if use_cache and past_key_values is None: + past_key_values = DynamicCache() + + if cache_position is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + cache_position = torch.arange(past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device) + + if position_ids is None: + position_ids = cache_position.unsqueeze(0) + + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + position_ids=position_ids, + ) + + # embed positions + hidden_states = inputs_embeds + # Convert to bfloat16 if the first layer uses bfloat16 + if len(self.layers) > 0 and self.layers[0].self_attn.q_proj.weight.dtype == torch.bfloat16: + hidden_states = hidden_states.to(torch.bfloat16) + + # create position embeddings to be shared across the decoder layers + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + # normalized + # Gemma downcasts the below to float16, causing sqrt(3072)=55.4256 to become 55.5 + # See https://github.com/huggingface/transformers/pull/29402 + normalizer = torch.tensor(self.config.hidden_size**0.5, dtype=hidden_states.dtype) + # hidden_states = hidden_states * normalizer + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + + for decoder_layer in self.layers[: self.config.num_hidden_layers]: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + layer_outputs = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + adarms_cond=adarms_cond, + **kwargs, + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + hidden_states, _ = self.norm(hidden_states, adarms_cond) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values if use_cache else None, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + +class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs): ... + + +@auto_docstring +class GemmaForCausalLM(GemmaPreTrainedModel, GenerationMixin): + _tied_weights_keys = ["lm_head.weight"] + _tp_plan = {"lm_head": "colwise_rep"} + _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + + def __init__(self, config): + super().__init__(config) + self.model = GemmaModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + logits_to_keep: Union[int, torch.Tensor] = 0, + adarms_cond: Optional[torch.Tensor] = None, + **kwargs: Unpack[KwargsForCausalLM], + ) -> CausalLMOutputWithPast: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + adarms_cond (`torch.Tensor` of shape `(batch_size, cond_dim)`, *optional*): + Condition for ADARMS. + + Example: + + ```python + >>> from transformers import AutoTokenizer, GemmaForCausalLM + + >>> model = GemmaForCausalLM.from_pretrained("google/gemma-7b") + >>> tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b") + + >>> prompt = "What is your favorite condiment?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "What is your favorite condiment?" + ```""" + 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 + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs: BaseModelOutputWithPast = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + cache_position=cache_position, + adarms_cond=adarms_cond, + **kwargs, + ) + + hidden_states = outputs.last_hidden_state + # 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.vocab_size, **kwargs) + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + The Gemma Model transformer with a sequence classification head on top (linear layer). + + [`GemmaForSequenceClassification`] uses the last token in order to do the classification, as other causal models + (e.g. GPT-2) do. + + Since it does classification on the last token, it requires to know the position of the last token. If a + `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If + no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the + padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in + each row of the batch). + """ +) +class GemmaForSequenceClassification(GemmaPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.model = GemmaModel(config) + self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + adarms_cond: Optional[torch.Tensor] = None, + ) -> SequenceClassifierOutputWithPast: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + + adarms_cond (`torch.Tensor` of shape `(batch_size, cond_dim)`, *optional*): + Condition for ADARMS. + """ + + transformer_outputs: BaseModelOutputWithPast = self.model( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + adarms_cond=adarms_cond, + ) + hidden_states = transformer_outputs.last_hidden_state + logits = self.score(hidden_states) + + if input_ids is not None: + batch_size = input_ids.shape[0] + else: + batch_size = inputs_embeds.shape[0] + + if self.config.pad_token_id is None and batch_size != 1: + raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.") + if self.config.pad_token_id is None: + last_non_pad_token = -1 + elif input_ids is not None: + # To handle both left- and right- padding, we take the rightmost token that is not equal to pad_token_id + non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32) + token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32) + last_non_pad_token = (token_indices * non_pad_mask).argmax(-1) + else: + last_non_pad_token = -1 + logger.warning_once(f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be unexpected if using padding tokens in conjunction with `inputs_embeds.`") + + pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token] + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, pooled_logits=pooled_logits, config=self.config) + + return SequenceClassifierOutputWithPast( + loss=loss, + logits=pooled_logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + +@auto_docstring +class GemmaForTokenClassification(GemmaPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.model = GemmaModel(config) + if getattr(config, "classifier_dropout", None) is not None: + classifier_dropout = config.classifier_dropout + elif getattr(config, "hidden_dropout", None) is not None: + classifier_dropout = config.hidden_dropout + else: + classifier_dropout = 0.1 + self.dropout = nn.Dropout(classifier_dropout) + self.score = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + adarms_cond: Optional[torch.Tensor] = None, + ) -> TokenClassifierOutput: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + + adarms_cond (`torch.Tensor` of shape `(batch_size, cond_dim)`, *optional*): + Condition for ADARMS. + """ + + outputs: BaseModelOutputWithPast = self.model( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + adarms_cond=adarms_cond, + ) + sequence_output = outputs.last_hidden_state + sequence_output = self.dropout(sequence_output) + logits = self.score(sequence_output) + + loss = None + if labels is not None: + loss = self.loss_function(logits, labels, self.config) + + return TokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +__all__ = [ + "GemmaModel", + "GemmaForCausalLM", + "GemmaForSequenceClassification", + "GemmaForTokenClassification", + "GemmaPreTrainedModel", +] diff --git a/lightx2v/models/networks/openpi/transformers_replace/models/paligemma/modeling_paligemma.py b/lightx2v/models/networks/openpi/transformers_replace/models/paligemma/modeling_paligemma.py new file mode 100644 index 000000000..1838525bf --- /dev/null +++ b/lightx2v/models/networks/openpi/transformers_replace/models/paligemma/modeling_paligemma.py @@ -0,0 +1,591 @@ +# coding=utf-8 +# Copyright 2024 the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch PaliGemmamodel.""" + +from dataclasses import dataclass +from typing import Optional, Union + +import torch +import torch.utils.checkpoint +from torch import nn + +from ...cache_utils import Cache, HybridCache, StaticCache +from ...generation import GenerationMixin +from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_outputs import BaseModelOutputWithPast +from ...modeling_utils import PreTrainedModel +from ...processing_utils import Unpack +from ...utils import LossKwargs, ModelOutput, auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging +from ..auto import AutoModel +from .configuration_paligemma import PaliGemmaConfig + +logger = logging.get_logger(__name__) + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for Paligemma outputs, with hidden states and attentions. + """ +) +class PaligemmaModelOutputWithPast(BaseModelOutputWithPast): + r""" + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) + + Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + image_hidden_states (`torch.FloatTensor`, *optional*): + A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`. + image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state. + """ + + image_hidden_states: Optional[torch.FloatTensor] = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for PaliGemma causal language model (or autoregressive) outputs. + """ +) +class PaliGemmaCausalLMOutputWithPast(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Language modeling loss (for next-token prediction). + logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.text_config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) + + Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + image_hidden_states (`torch.FloatTensor`, *optional*): + A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`. + image_hidden_states of the model produced by the vision encoder after projecting last hidden state. + """ + + loss: Optional[torch.FloatTensor] = None + logits: Optional[torch.FloatTensor] = None + past_key_values: Optional[Union[list[torch.FloatTensor], Cache]] = None + hidden_states: Optional[tuple[torch.FloatTensor]] = None + attentions: Optional[tuple[torch.FloatTensor]] = None + image_hidden_states: Optional[torch.FloatTensor] = None + + +class PaliGemmaMultiModalProjector(nn.Module): + def __init__(self, config: PaliGemmaConfig): + super().__init__() + self.linear = nn.Linear(config.vision_config.hidden_size, config.vision_config.projection_dim, bias=True) + + def forward(self, image_features): + hidden_states = self.linear(image_features) + + return hidden_states + + +@auto_docstring +class PaliGemmaPreTrainedModel(PreTrainedModel): + config_class = PaliGemmaConfig + base_model_prefix = "" + supports_gradient_checkpointing = True + _no_split_modules = ["PaliGemmaMultiModalProjector"] + _skip_keys_device_placement = "past_key_values" + _supports_cache_class = True + _supports_quantized_cache = True + _supports_static_cache = True + _supports_flash_attn_2 = True + _supports_sdpa = True + _supports_flex_attn = True + _supports_attention_backend = True + + def _init_weights(self, module): + # important: this ported version of PaliGemmaisn't meant for training from scratch - only + # inference and fine-tuning + std = getattr(self.config, "initializer_range", self.config.get_text_config().initializer_range) + + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + + +@auto_docstring( + custom_intro=""" + The Base Paligemma model which consists of a vision backbone and a language model withou language modeling head., + """ +) +class PaliGemmaModel(PaliGemmaPreTrainedModel): + _checkpoint_conversion_mapping = {"language_model.model": "language_model"} + # we are filtering the logits/labels so we shouldn't divide the loss based on num_items_in_batch + accepts_loss_kwargs = False + + def __init__(self, config: PaliGemmaConfig): + super().__init__(config) + self.vision_tower = AutoModel.from_config(config=config.vision_config) + self.multi_modal_projector = PaliGemmaMultiModalProjector(config) + self.vocab_size = config.text_config.vocab_size + + language_model = AutoModel.from_config(config=config.text_config) + self.language_model = language_model + + self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1 + self.post_init() + + # Copied from transformers.models.llava.modeling_llava.LlavaModel.get_input_embeddings with Llava->PaliGemma + def get_input_embeddings(self): + return self.language_model.get_input_embeddings() + + # Copied from transformers.models.llava.modeling_llava.LlavaModel.set_input_embeddings with Llava->PaliGemma + def set_input_embeddings(self, value): + self.language_model.set_input_embeddings(value) + + def set_decoder(self, decoder): + self.language_model = decoder + + def get_decoder(self): + return self.language_model + + def _update_causal_mask( + self, + attention_mask, + token_type_ids=None, + past_key_values=None, + cache_position=None, + input_tensor=None, + is_training: Optional[bool] = None, + ): + if self.config.text_config._attn_implementation == "flash_attention_2": + if attention_mask is not None and 0.0 in attention_mask: + return attention_mask + return None + is_training = is_training if is_training is not None else self.training + using_static_cache = isinstance(past_key_values, StaticCache) + min_dtype = torch.finfo(self.dtype).min + if input_tensor is None: + input_tensor = attention_mask + + inputs_lead_dim, sequence_length = input_tensor.shape[:2] + if using_static_cache: + target_length = past_key_values.get_max_cache_shape() + elif isinstance(past_key_values, HybridCache): + target_length = past_key_values.get_max_cache_shape() + else: + target_length = attention_mask.shape[-1] if isinstance(attention_mask, torch.Tensor) else cache_position[0] + sequence_length + 1 + + if attention_mask is not None and attention_mask.dim() == 4: + # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. + return attention_mask + + causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=self.dtype, device=cache_position.device) + # Causal diagonal mask only if training, otherwise attend to the whole prefix. Training-specific attn for prefix is handled below + if sequence_length != 1: + if is_training: + causal_mask = torch.triu(causal_mask, diagonal=1) + else: + causal_mask[:, :sequence_length] = 0.0 + + causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) + causal_mask = causal_mask[None, None, :, :].expand(inputs_lead_dim, 1, -1, -1) + if attention_mask is not None: + causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit + mask_length = attention_mask.shape[-1] + + # First unmask prefix tokens during training + if is_training: + if token_type_ids is None: + raise ValueError("Token type ids must be provided during training") + causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(token_type_ids[:, None, None, :].to(causal_mask.device) == 0, 0) + + # Then apply padding mask (will mask pad tokens) + padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(causal_mask.device) + padding_mask = padding_mask == 0 + causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(padding_mask, min_dtype) + + return causal_mask + + def get_image_features(self, pixel_values: torch.FloatTensor): + """ + Obtains image last hidden states from the vision tower and apply multimodal projection. + + Args: + pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`) + The tensors corresponding to the input images. + Returns: + image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). + """ + image_outputs = self.vision_tower(pixel_values) + selected_image_feature = image_outputs.last_hidden_state + image_features = self.multi_modal_projector(selected_image_feature) + return image_features + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor = None, + pixel_values: torch.FloatTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Union[list[torch.FloatTensor], Cache]] = None, + token_type_ids: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> Union[tuple, PaligemmaModelOutputWithPast]: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.text_config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.text_config.vocab_size]`. + + Example: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, PaliGemmaForConditionalGeneration + + >>> model = PaliGemmaForConditionalGeneration.from_pretrained("google/paligemma2-3b-mix-224") + >>> processor = AutoProcessor.from_pretrained("google/paligemma2-3b-mix-224") + + >>> prompt = "Where is the cat standing?" + >>> url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> inputs = processor(images=image, text=prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(**inputs,) + >>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Where is the cat standing?\nsnow" + ```""" + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + 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 + + is_training = token_type_ids is not None and labels is not None + + # Replace image id woth PAD if the image token if OOV, to avoid index-errors + if input_ids is not None and self.config.image_token_id >= self.vocab_size: + special_image_mask = input_ids == self.config.image_token_id + llm_input_ids = input_ids.clone() + llm_input_ids[special_image_mask] = 0 + else: + llm_input_ids = input_ids + + if inputs_embeds is None: + inputs_embeds = self.get_input_embeddings()(llm_input_ids) + + if cache_position is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + cache_position = torch.arange(past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device) + + if position_ids is None: + position_ids = cache_position.unsqueeze(0) + 1 # Paligemma positions are 1-indexed + + # Merge text and images + if pixel_values is not None: + image_features = self.get_image_features(pixel_values) + + if input_ids is None: + special_image_mask = inputs_embeds == self.get_input_embeddings()(torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)) + else: + special_image_mask = (input_ids == self.config.image_token_id).unsqueeze(-1) + special_image_mask = special_image_mask.expand_as(inputs_embeds).to(inputs_embeds.device) + + if not is_torchdynamo_compiling() and inputs_embeds[special_image_mask].numel() != image_features.numel(): + image_tokens_in_text = (special_image_mask).sum(dim=1).sum(dim=0)[0] + raise ValueError( + f"Number of images does not match number of special image tokens in the input text. " + f"Got {image_tokens_in_text} image tokens in the text but {image_features.shape[0] * image_features.shape[1]} " + "tokens from image embeddings." + ) + image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) + inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features) + + causal_mask = self._update_causal_mask(attention_mask, token_type_ids, past_key_values, cache_position, inputs_embeds, is_training) + outputs = self.language_model( + attention_mask=causal_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=True, + cache_position=cache_position, + **kwargs, + ) + + return PaligemmaModelOutputWithPast( + 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_features if pixel_values is not None else None, + ) + + +class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs): ... + + +@auto_docstring( + custom_intro=""" + The Base Paligemma model which consists of a vision backbone and a language model without language modeling head., + """ +) +class PaliGemmaForConditionalGeneration(PaliGemmaPreTrainedModel, GenerationMixin): + _checkpoint_conversion_mapping = { + "^language_model.model": "model.language_model", + "^vision_tower": "model.vision_tower", + "^multi_modal_projector": "model.multi_modal_projector", + "^language_model.lm_head": "lm_head", + } + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config: PaliGemmaConfig): + super().__init__(config) + self.model = PaliGemmaModel(config) + self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False) + self.post_init() + + def get_input_embeddings(self): + return self.model.get_input_embeddings() + + def set_input_embeddings(self, value): + self.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 set_decoder(self, decoder): + self.model.set_decoder(decoder) + + def get_decoder(self): + return self.model.get_decoder() + + def get_image_features(self, pixel_values): + return self.model.get_image_features(pixel_values) + + # Make modules available throught conditional class for BC + @property + def language_model(self): + return self.model.language_model + + @property + def vision_tower(self): + return self.model.vision_tower + + @property + def multi_modal_projector(self): + return self.model.multi_modal_projector + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor = None, + pixel_values: torch.FloatTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Union[list[torch.FloatTensor], Cache]] = None, + token_type_ids: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + logits_to_keep: Union[int, torch.Tensor] = 0, + **kwargs: Unpack[KwargsForCausalLM], + ) -> Union[tuple, PaliGemmaCausalLMOutputWithPast]: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.text_config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.text_config.vocab_size]`. + + Example: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, PaliGemmaForConditionalGeneration + + >>> model = PaliGemmaForConditionalGeneration.from_pretrained("google/paligemma2-3b-mix-224") + >>> processor = AutoProcessor.from_pretrained("google/paligemma2-3b-mix-224") + + >>> prompt = "Where is the cat standing?" + >>> url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> inputs = processor(images=image, text=prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(**inputs,) + >>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Where is the cat standing?\nsnow" + ```""" + 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 + + outputs = self.model( + input_ids=input_ids, + pixel_values=pixel_values, + token_type_ids=token_type_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + labels=labels, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=True, + cache_position=cache_position, + **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 PaliGemmaCausalLMOutputWithPast( + 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, + ) + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + inputs_embeds=None, + cache_position=None, + position_ids=None, + pixel_values=None, + attention_mask=None, + token_type_ids=None, + use_cache=True, + logits_to_keep=None, + labels=None, + **kwargs, + ): + # Overwritten -- custom `position_ids` and `pixel_values` handling + model_inputs = super().prepare_inputs_for_generation( + input_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + cache_position=cache_position, + use_cache=use_cache, + logits_to_keep=logits_to_keep, + token_type_ids=token_type_ids, + **kwargs, + ) + + # position_ids in Paligemma are 1-indexed + if model_inputs.get("position_ids") is not None: + model_inputs["position_ids"] += 1 + # If we're in cached decoding stage, pixel values should be None because input ids do not contain special image token anymore + # Otherwise we need pixel values to be passed to model. NOTE: use_cache=False needs pixel_values always + if cache_position[0] == 0: + model_inputs["pixel_values"] = pixel_values + is_training = token_type_ids is not None and labels is not None + if cache_position[0] == 0 and isinstance(past_key_values, HybridCache): + input_tensor = inputs_embeds if inputs_embeds is not None else input_ids + causal_mask = self.model._update_causal_mask(attention_mask, token_type_ids, past_key_values, cache_position, input_tensor, is_training) + model_inputs["attention_mask"] = causal_mask + + return model_inputs + + @staticmethod + # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._prepare_4d_causal_attention_mask_with_cache_position + def _prepare_4d_causal_attention_mask_with_cache_position( + attention_mask: torch.Tensor, + sequence_length: int, + target_length: int, + dtype: torch.dtype, + cache_position: torch.Tensor, + batch_size: int, + **kwargs, + ): + """ + Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape + `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. + + Args: + attention_mask (`torch.Tensor`): + A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape + `(batch_size, 1, query_length, key_value_length)`. + sequence_length (`int`): + The sequence length being processed. + target_length (`int`): + The target length: when generating with static cache, the mask should be as long as the static cache, + to account for the 0 padding, the part of the cache that is not filled yet. + dtype (`torch.dtype`): + The dtype to use for the 4D attention mask. + cache_position (`torch.Tensor`): + Indices depicting the position of the input sequence tokens in the sequence. + batch_size (`torch.Tensor`): + Batch size. + """ + if attention_mask is not None and attention_mask.dim() == 4: + # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. + causal_mask = attention_mask + else: + min_dtype = torch.finfo(dtype).min + causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device) + if sequence_length != 1: + causal_mask = torch.triu(causal_mask, diagonal=1) + causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) + causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) + if attention_mask is not None: + causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit + mask_length = attention_mask.shape[-1] + padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(causal_mask.device) + padding_mask = padding_mask == 0 + causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(padding_mask, min_dtype) + + return causal_mask + + +__all__ = ["PaliGemmaForConditionalGeneration", "PaliGemmaPreTrainedModel", "PaliGemmaModel"] diff --git a/lightx2v/models/networks/openpi/transformers_replace/models/siglip/check.py b/lightx2v/models/networks/openpi/transformers_replace/models/siglip/check.py new file mode 100644 index 000000000..d899dc1b9 --- /dev/null +++ b/lightx2v/models/networks/openpi/transformers_replace/models/siglip/check.py @@ -0,0 +1,5 @@ +import transformers + + +def check_whether_transformers_replace_is_installed_correctly(): + return transformers.__version__ == "4.53.2" diff --git a/lightx2v/models/networks/openpi/transformers_replace/models/siglip/modeling_siglip.py b/lightx2v/models/networks/openpi/transformers_replace/models/siglip/modeling_siglip.py new file mode 100644 index 000000000..bcae34ab6 --- /dev/null +++ b/lightx2v/models/networks/openpi/transformers_replace/models/siglip/modeling_siglip.py @@ -0,0 +1,1196 @@ +# coding=utf-8 +# Copyright 2024 Google AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch Siglip model.""" + +import math +import warnings +from dataclasses import dataclass +from typing import Any, Callable, Optional, Union + +import numpy as np +import torch +import torch.utils.checkpoint +from torch import nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss +from torch.nn.init import _calculate_fan_in_and_fan_out + +from ...activations import ACT2FN +from ...modeling_attn_mask_utils import _prepare_4d_attention_mask +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, ImageClassifierOutput +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...utils import ModelOutput, auto_docstring, can_return_tuple, logging, torch_int +from .configuration_siglip import SiglipConfig, SiglipTextConfig, SiglipVisionConfig + +logger = logging.get_logger(__name__) + + +def _trunc_normal_(tensor, mean, std, a, b): + # Cut & paste from PyTorch official master until it's in a few official releases - RW + # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf + def norm_cdf(x): + # Computes standard normal cumulative distribution function + return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0 + + if (mean < a - 2 * std) or (mean > b + 2 * std): + warnings.warn( + "mean is more than 2 std from [a, b] in nn.init.trunc_normal_. The distribution of values may be incorrect.", + stacklevel=2, + ) + + # Values are generated by using a truncated uniform distribution and + # then using the inverse CDF for the normal distribution. + # Get upper and lower cdf values + lower_cdf = norm_cdf((a - mean) / std) + upper_cdf = norm_cdf((b - mean) / std) + + # Uniformly fill tensor with values from [lower_cdf, upper_cdf], then + # translate to [2 * lower_cdf - 1, 2 * upper_cdf - 1]. + tensor.uniform_(2 * lower_cdf - 1, 2 * upper_cdf - 1) + + # Use inverse cdf transform for normal distribution to get truncated + # standard normal + tensor.erfinv_() + + # Transform to proper mean, std + tensor.mul_(std * math.sqrt(2.0)) + tensor.add_(mean) + + # Clamp to ensure it's in the proper range + tensor.clamp_(min=a, max=b) + + +def trunc_normal_tf_(tensor: torch.Tensor, mean: float = 0.0, std: float = 1.0, a: float = -2.0, b: float = 2.0) -> torch.Tensor: + """Fills the input Tensor with values drawn from a truncated + normal distribution. The values are effectively drawn from the + normal distribution :math:`\\mathcal{N}(\text{mean}, \text{std}^2)` + with values outside :math:`[a, b]` redrawn until they are within + the bounds. The method used for generating the random values works + best when :math:`a \\leq \text{mean} \\leq b`. + + NOTE: this 'tf' variant behaves closer to Tensorflow / JAX impl where the + bounds [a, b] are applied when sampling the normal distribution with mean=0, std=1.0 + and the result is subsequently scaled and shifted by the mean and std args. + + Args: + tensor: an n-dimensional `torch.Tensor` + mean: the mean of the normal distribution + std: the standard deviation of the normal distribution + a: the minimum cutoff value + b: the maximum cutoff value + """ + with torch.no_grad(): + _trunc_normal_(tensor, 0, 1.0, a, b) + tensor.mul_(std).add_(mean) + + +def variance_scaling_(tensor, scale=1.0, mode="fan_in", distribution="normal"): + fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor) + if mode == "fan_in": + denom = fan_in + elif mode == "fan_out": + denom = fan_out + elif mode == "fan_avg": + denom = (fan_in + fan_out) / 2 + + variance = scale / denom + + if distribution == "truncated_normal": + # constant is stddev of standard normal truncated to (-2, 2) + trunc_normal_tf_(tensor, std=math.sqrt(variance) / 0.87962566103423978) + elif distribution == "normal": + with torch.no_grad(): + tensor.normal_(std=math.sqrt(variance)) + elif distribution == "uniform": + bound = math.sqrt(3 * variance) + with torch.no_grad(): + tensor.uniform_(-bound, bound) + else: + raise ValueError(f"invalid distribution {distribution}") + + +def lecun_normal_(tensor): + variance_scaling_(tensor, mode="fan_in", distribution="truncated_normal") + + +def default_flax_embed_init(tensor): + variance_scaling_(tensor, mode="fan_in", distribution="normal") + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for vision model's outputs that also contains image embeddings of the pooling of the last hidden states. + """ +) +# Copied from transformers.models.clip.modeling_clip.CLIPVisionModelOutput with CLIP->Siglip +class SiglipVisionModelOutput(ModelOutput): + r""" + image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`): + The image embeddings obtained by applying the projection layer to the pooler_output. + """ + + image_embeds: Optional[torch.FloatTensor] = None + last_hidden_state: Optional[torch.FloatTensor] = None + hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None + attentions: Optional[tuple[torch.FloatTensor, ...]] = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for text model's outputs that also contains a pooling of the last hidden states. + """ +) +# Copied from transformers.models.clip.modeling_clip.CLIPTextModelOutput with CLIP->Siglip +class SiglipTextModelOutput(ModelOutput): + r""" + text_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`): + The text embeddings obtained by applying the projection layer to the pooler_output. + """ + + text_embeds: Optional[torch.FloatTensor] = None + last_hidden_state: Optional[torch.FloatTensor] = None + hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None + attentions: Optional[tuple[torch.FloatTensor, ...]] = None + + +@dataclass +@auto_docstring +# Copied from transformers.models.clip.modeling_clip.CLIPOutput with CLIP->Siglip +class SiglipOutput(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `return_loss` is `True`): + Contrastive loss for image-text similarity. + logits_per_image (`torch.FloatTensor` of shape `(image_batch_size, text_batch_size)`): + The scaled dot product scores between `image_embeds` and `text_embeds`. This represents the image-text + similarity scores. + logits_per_text (`torch.FloatTensor` of shape `(text_batch_size, image_batch_size)`): + The scaled dot product scores between `text_embeds` and `image_embeds`. This represents the text-image + similarity scores. + text_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`): + The text embeddings obtained by applying the projection layer to the pooled output of [`SiglipTextModel`]. + image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`): + The image embeddings obtained by applying the projection layer to the pooled output of [`SiglipVisionModel`]. + text_model_output (`BaseModelOutputWithPooling`): + The output of the [`SiglipTextModel`]. + vision_model_output (`BaseModelOutputWithPooling`): + The output of the [`SiglipVisionModel`]. + """ + + loss: Optional[torch.FloatTensor] = None + logits_per_image: Optional[torch.FloatTensor] = None + logits_per_text: Optional[torch.FloatTensor] = None + text_embeds: Optional[torch.FloatTensor] = None + image_embeds: Optional[torch.FloatTensor] = None + text_model_output: BaseModelOutputWithPooling = None + vision_model_output: BaseModelOutputWithPooling = None + + def to_tuple(self) -> tuple[Any]: + return tuple(self[k] if k not in ["text_model_output", "vision_model_output"] else getattr(self, k).to_tuple() for k in self.keys()) + + +class SiglipVisionEmbeddings(nn.Module): + def __init__(self, config: SiglipVisionConfig): + super().__init__() + self.config = config + self.embed_dim = config.hidden_size + self.image_size = config.image_size + self.patch_size = config.patch_size + + self.patch_embedding = nn.Conv2d( + in_channels=config.num_channels, + out_channels=self.embed_dim, + kernel_size=self.patch_size, + stride=self.patch_size, + padding="valid", + ) + + self.num_patches = (self.image_size // self.patch_size) ** 2 + self.num_positions = self.num_patches + self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim) + self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False) + + def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor: + """ + This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution + images. This method is also adapted to support torch.jit tracing and no class embeddings. + + Adapted from: + - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and + - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211 + """ + + num_patches = embeddings.shape[1] + num_positions = self.position_embedding.weight.shape[0] + + # always interpolate when tracing to ensure the exported model works for dynamic input shapes + if not torch.jit.is_tracing() and num_patches == num_positions and height == width: + return self.position_embedding(self.position_ids) + + patch_pos_embed = self.position_embedding.weight.unsqueeze(0) + + dim = embeddings.shape[-1] + + new_height = height // self.patch_size + new_width = width // self.patch_size + + sqrt_num_positions = torch_int(num_positions**0.5) + patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim) + patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2) + + patch_pos_embed = nn.functional.interpolate( + patch_pos_embed, + size=(new_height, new_width), + mode="bicubic", + align_corners=False, + ) + + patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) + return patch_pos_embed + + def forward(self, pixel_values: torch.FloatTensor, interpolate_pos_encoding=False) -> torch.Tensor: + _, _, height, width = pixel_values.shape + target_dtype = self.patch_embedding.weight.dtype + patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype)) # shape = [*, width, grid, grid] + embeddings = patch_embeds.flatten(2).transpose(1, 2) + + if interpolate_pos_encoding: + embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width) + else: + embeddings = embeddings + self.position_embedding(self.position_ids) + return embeddings + + +# Copied from transformers.models.clip.modeling_clip.CLIPTextEmbeddings with CLIP->Siglip +class SiglipTextEmbeddings(nn.Module): + def __init__(self, config: SiglipTextConfig): + super().__init__() + embed_dim = config.hidden_size + + self.token_embedding = nn.Embedding(config.vocab_size, embed_dim) + self.position_embedding = nn.Embedding(config.max_position_embeddings, embed_dim) + + # position_ids (1, len position emb) is contiguous in memory and exported when serialized + self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False) + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + ) -> torch.Tensor: + seq_length = input_ids.shape[-1] if input_ids is not None else inputs_embeds.shape[-2] + max_position_embedding = self.position_embedding.weight.shape[0] + + if seq_length > max_position_embedding: + raise ValueError(f"Sequence length must be less than max_position_embeddings (got `sequence length`: {seq_length} and max_position_embeddings: {max_position_embedding}") + + if position_ids is None: + position_ids = self.position_ids[:, :seq_length] + + if inputs_embeds is None: + inputs_embeds = self.token_embedding(input_ids) + + position_embeddings = self.position_embedding(position_ids) + embeddings = inputs_embeds + position_embeddings + + return embeddings + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: Optional[torch.Tensor], + scaling: float, + dropout: float = 0.0, + **kwargs, +): + attn_weights = torch.matmul(query, key.transpose(-1, -2)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + + attn_output = torch.matmul(attn_weights, value) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +class SiglipAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config): + super().__init__() + self.config = config + self.embed_dim = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.embed_dim // self.num_heads + if self.head_dim * self.num_heads != self.embed_dim: + raise ValueError(f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`: {self.num_heads}).") + self.scale = self.head_dim**-0.5 + self.dropout = config.attention_dropout + self.is_causal = False + + self.k_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.v_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.q_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.out_proj = nn.Linear(self.embed_dim, self.embed_dim) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = False, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """Input shape: Batch x Time x Channel""" + + batch_size, seq_length, embed_dim = hidden_states.shape + + queries = self.q_proj(hidden_states) + keys = self.k_proj(hidden_states) + values = self.v_proj(hidden_states) + + queries = queries.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) + keys = keys.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) + values = values.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) + + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + if self.config._attn_implementation == "sdpa" and output_attentions: + logger.warning_once( + "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to " + 'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.' + ) + else: + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + + attn_output, attn_weights = attention_interface( + self, + queries, + keys, + values, + attention_mask, + is_causal=self.is_causal, + scaling=self.scale, + dropout=0.0 if not self.training else self.dropout, + ) + + attn_output = attn_output.reshape(batch_size, seq_length, embed_dim).contiguous() + attn_output = self.out_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights + + +# Copied from transformers.models.clip.modeling_clip.CLIPMLP with CLIP->Siglip +class SiglipMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.activation_fn = ACT2FN[config.hidden_act] + self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size) + self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.fc1(hidden_states) + hidden_states = self.activation_fn(hidden_states) + hidden_states = self.fc2(hidden_states) + return hidden_states + + +class SiglipEncoderLayer(GradientCheckpointingLayer): + def __init__(self, config: Union[SiglipVisionConfig, SiglipTextConfig]): + super().__init__() + self.embed_dim = config.hidden_size + self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + self.self_attn = SiglipAttention(config) + self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + self.mlp = SiglipMLP(config) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + output_attentions: Optional[bool] = False, + ) -> tuple[torch.FloatTensor]: + """ + Args: + hidden_states (`torch.FloatTensor`): + Input to the layer of shape `(batch, seq_len, embed_dim)`. + attention_mask (`torch.FloatTensor`): + Attention mask of shape `(batch, 1, q_len, k_v_seq_len)` where padding elements are indicated by very large negative values. + output_attentions (`bool`, *optional*, defaults to `False`): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + """ + residual = hidden_states + + hidden_states = self.layer_norm1(hidden_states) + hidden_states, attn_weights = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + output_attentions=output_attentions, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.layer_norm2(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (attn_weights,) + + return outputs + + +@auto_docstring +class SiglipPreTrainedModel(PreTrainedModel): + config_class = SiglipConfig + base_model_prefix = "siglip" + supports_gradient_checkpointing = True + + _no_split_modules = [ + "SiglipTextEmbeddings", + "SiglipEncoderLayer", + "SiglipVisionEmbeddings", + "SiglipEncoderLayer", + "SiglipMultiheadAttentionPoolingHead", + ] + _supports_flash_attn_2 = True + _supports_sdpa = True + _supports_flex_attn = True + _supports_attention_backend = True + + def _init_weights(self, module): + """Initialize the weights""" + if isinstance(module, SiglipVisionEmbeddings): + width = self.config.vision_config.hidden_size if isinstance(self.config, SiglipConfig) else self.config.hidden_size + nn.init.normal_(module.position_embedding.weight, std=1 / np.sqrt(width)) + elif isinstance(module, nn.Embedding): + default_flax_embed_init(module.weight) + elif isinstance(module, SiglipAttention): + nn.init.xavier_uniform_(module.q_proj.weight) + nn.init.xavier_uniform_(module.k_proj.weight) + nn.init.xavier_uniform_(module.v_proj.weight) + nn.init.xavier_uniform_(module.out_proj.weight) + nn.init.zeros_(module.q_proj.bias) + nn.init.zeros_(module.k_proj.bias) + nn.init.zeros_(module.v_proj.bias) + nn.init.zeros_(module.out_proj.bias) + elif isinstance(module, SiglipMLP): + nn.init.xavier_uniform_(module.fc1.weight) + nn.init.xavier_uniform_(module.fc2.weight) + nn.init.normal_(module.fc1.bias, std=1e-6) + nn.init.normal_(module.fc2.bias, std=1e-6) + elif isinstance(module, SiglipMultiheadAttentionPoolingHead): + nn.init.xavier_uniform_(module.probe.data) + nn.init.xavier_uniform_(module.attention.in_proj_weight.data) + nn.init.zeros_(module.attention.in_proj_bias.data) + elif isinstance(module, SiglipModel): + logit_scale_init = torch.log(torch.tensor(1.0)) + module.logit_scale.data.fill_(logit_scale_init) + module.logit_bias.data.zero_() + elif isinstance(module, SiglipForImageClassification): + nn.init.normal_( + module.classifier.weight, + std=self.config.vision_config.hidden_size**-0.5 * self.config.initializer_factor, + ) + elif isinstance(module, (nn.Linear, nn.Conv2d)): + lecun_normal_(module.weight) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.LayerNorm): + module.bias.data.zero_() + module.weight.data.fill_(1.0) + + +# Copied from transformers.models.altclip.modeling_altclip.AltCLIPEncoder with AltCLIP->Siglip +class SiglipEncoder(nn.Module): + """ + Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a + [`SiglipEncoderLayer`]. + + Args: + config: SiglipConfig + """ + + def __init__(self, config: SiglipConfig): + super().__init__() + self.config = config + self.layers = nn.ModuleList([SiglipEncoderLayer(config) for _ in range(config.num_hidden_layers)]) + self.gradient_checkpointing = False + + # Ignore copy + @can_return_tuple + def forward( + self, + inputs_embeds, + attention_mask: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + ) -> BaseModelOutput: + r""" + Args: + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. + This is useful if you want more control over how to convert `input_ids` indices into associated vectors + than the model's internal embedding lookup matrix. + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors + for more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + """ + 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 + + encoder_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + + hidden_states = inputs_embeds + for encoder_layer in self.layers: + if output_hidden_states: + encoder_states = encoder_states + (hidden_states,) + + layer_outputs = encoder_layer( + hidden_states, + attention_mask, + output_attentions=output_attentions, + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_attentions = all_attentions + (layer_outputs[1],) + + if output_hidden_states: + encoder_states = encoder_states + (hidden_states,) + + return BaseModelOutput( + last_hidden_state=hidden_states, + hidden_states=encoder_states, + attentions=all_attentions, + ) + + +class SiglipTextTransformer(nn.Module): + def __init__(self, config: SiglipTextConfig): + super().__init__() + self.config = config + embed_dim = config.hidden_size + self.embeddings = SiglipTextEmbeddings(config) + self.encoder = SiglipEncoder(config) + self.final_layer_norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + + self.head = nn.Linear(embed_dim, config.projection_size) + self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2" + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + ) -> BaseModelOutputWithPooling: + 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 + + if input_ids is None: + raise ValueError("You have to specify input_ids") + + input_shape = input_ids.size() + input_ids = input_ids.view(-1, input_shape[-1]) + + hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids) + + # note: SigLIP's text model does not use a causal mask, unlike the original CLIP model. + # expand attention_mask + if attention_mask is not None and not self._use_flash_attention_2: + # [batch_size, seq_len] -> [batch_size, 1, tgt_seq_len, src_seq_len] + attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype) + + encoder_outputs: BaseModelOutput = self.encoder( + inputs_embeds=hidden_states, + attention_mask=attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + ) + + last_hidden_state = encoder_outputs.last_hidden_state + last_hidden_state = self.final_layer_norm(last_hidden_state) + + # Assuming "sticky" EOS tokenization, last token is always EOS. + pooled_output = last_hidden_state[:, -1, :] + pooled_output = self.head(pooled_output) + + return BaseModelOutputWithPooling( + last_hidden_state=last_hidden_state, + pooler_output=pooled_output, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + The text model from SigLIP without any head or projection on top. + """ +) +class SiglipTextModel(SiglipPreTrainedModel): + config_class = SiglipTextConfig + + def __init__(self, config: SiglipTextConfig): + super().__init__(config) + self.text_model = SiglipTextTransformer(config) + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self) -> nn.Module: + return self.text_model.embeddings.token_embedding + + def set_input_embeddings(self, value): + self.text_model.embeddings.token_embedding = value + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + ) -> BaseModelOutputWithPooling: + r""" + Examples: + + ```python + >>> from transformers import AutoTokenizer, SiglipTextModel + + >>> model = SiglipTextModel.from_pretrained("google/siglip-base-patch16-224") + >>> tokenizer = AutoTokenizer.from_pretrained("google/siglip-base-patch16-224") + + >>> # important: make sure to set padding="max_length" as that's how the model was trained + >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding="max_length", return_tensors="pt") + + >>> outputs = model(**inputs) + >>> last_hidden_state = outputs.last_hidden_state + >>> pooled_output = outputs.pooler_output # pooled (EOS token) states + ```""" + + return self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + ) + + +class SiglipVisionTransformer(nn.Module): + def __init__(self, config: SiglipVisionConfig): + super().__init__() + self.config = config + embed_dim = config.hidden_size + + self.embeddings = SiglipVisionEmbeddings(config) + self.encoder = SiglipEncoder(config) + self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + self.use_head = True if not hasattr(config, "vision_use_head") else config.vision_use_head + if self.use_head: + self.head = SiglipMultiheadAttentionPoolingHead(config) + + @can_return_tuple + @auto_docstring + def forward( + self, + pixel_values, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + interpolate_pos_encoding: Optional[bool] = False, + ) -> BaseModelOutputWithPooling: + 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 + + hidden_states = self.embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding) + # Convert to bfloat16 if the encoder uses bfloat16 + if len(self.encoder.layers) > 0 and self.encoder.layers[0].self_attn.q_proj.weight.dtype == torch.bfloat16: + hidden_states = hidden_states.to(torch.bfloat16) + + encoder_outputs: BaseModelOutput = self.encoder( + inputs_embeds=hidden_states, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + ) + + last_hidden_state = encoder_outputs.last_hidden_state + last_hidden_state = self.post_layernorm(last_hidden_state) + + pooler_output = self.head(last_hidden_state) if self.use_head else None + + return BaseModelOutputWithPooling( + last_hidden_state=last_hidden_state, + pooler_output=pooler_output, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + ) + + +class SiglipMultiheadAttentionPoolingHead(nn.Module): + """Multihead Attention Pooling.""" + + def __init__(self, config: SiglipVisionConfig): + super().__init__() + + self.probe = nn.Parameter(torch.randn(1, 1, config.hidden_size)) + self.attention = torch.nn.MultiheadAttention(config.hidden_size, config.num_attention_heads, batch_first=True) + self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.mlp = SiglipMLP(config) + + def forward(self, hidden_state): + batch_size = hidden_state.shape[0] + probe = self.probe.repeat(batch_size, 1, 1) + + hidden_state = self.attention(probe, hidden_state, hidden_state)[0] + + residual = hidden_state + hidden_state = self.layernorm(hidden_state) + hidden_state = residual + self.mlp(hidden_state) + + return hidden_state[:, 0] + + +@auto_docstring( + custom_intro=""" + The vision model from SigLIP without any head or projection on top. + """ +) +class SiglipVisionModel(SiglipPreTrainedModel): + config_class = SiglipVisionConfig + main_input_name = "pixel_values" + + def __init__(self, config: SiglipVisionConfig): + super().__init__(config) + + self.vision_model = SiglipVisionTransformer(config) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self) -> nn.Module: + return self.vision_model.embeddings.patch_embedding + + @can_return_tuple + @auto_docstring + def forward( + self, + pixel_values, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + interpolate_pos_encoding: bool = False, + ) -> BaseModelOutputWithPooling: + r""" + Examples: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, SiglipVisionModel + + >>> model = SiglipVisionModel.from_pretrained("google/siglip-base-patch16-224") + >>> processor = AutoProcessor.from_pretrained("google/siglip-base-patch16-224") + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> inputs = processor(images=image, return_tensors="pt") + + >>> outputs = model(**inputs) + >>> last_hidden_state = outputs.last_hidden_state + >>> pooled_output = outputs.pooler_output # pooled features + ```""" + + return self.vision_model( + pixel_values=pixel_values, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + interpolate_pos_encoding=interpolate_pos_encoding, + ) + + +@auto_docstring +class SiglipModel(SiglipPreTrainedModel): + config_class = SiglipConfig + + def __init__(self, config: SiglipConfig): + super().__init__(config) + + if not isinstance(config.text_config, SiglipTextConfig): + raise TypeError(f"config.text_config is expected to be of type SiglipTextConfig but is of type {type(config.text_config)}.") + + if not isinstance(config.vision_config, SiglipVisionConfig): + raise TypeError(f"config.vision_config is expected to be of type SiglipVisionConfig but is of type {type(config.vision_config)}.") + + text_config = config.text_config + vision_config = config.vision_config + + # First, initialize the text and vision models with proper attention implementation + text_model = SiglipTextModel._from_config(text_config) + vision_model = SiglipVisionModel._from_config(vision_config) + + # Second, get the text and vision submodules (for backward compatibility) + self.text_model = text_model.text_model + self.vision_model = vision_model.vision_model + + self.logit_scale = nn.Parameter(torch.randn(1)) + self.logit_bias = nn.Parameter(torch.randn(1)) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def get_text_features( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + ) -> torch.FloatTensor: + r""" + Returns: + text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by + applying the projection layer to the pooled output of [`SiglipTextModel`]. + + Examples: + + ```python + >>> from transformers import AutoTokenizer, AutoModel + >>> import torch + + >>> model = AutoModel.from_pretrained("google/siglip-base-patch16-224") + >>> tokenizer = AutoTokenizer.from_pretrained("google/siglip-base-patch16-224") + + >>> # important: make sure to set padding="max_length" as that's how the model was trained + >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding="max_length", return_tensors="pt") + >>> with torch.no_grad(): + ... text_features = model.get_text_features(**inputs) + ```""" + # Use SigLIP model's config for some fields (if specified) instead of those of vision & text components. + 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 + + text_outputs: BaseModelOutputWithPooling = self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + ) + + pooled_output = text_outputs.pooler_output + + return pooled_output + + @auto_docstring + def get_image_features( + self, + pixel_values: Optional[torch.FloatTensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + interpolate_pos_encoding: bool = False, + ) -> torch.FloatTensor: + r""" + Returns: + image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by + applying the projection layer to the pooled output of [`SiglipVisionModel`]. + + Examples: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, AutoModel + >>> import torch + + >>> model = AutoModel.from_pretrained("google/siglip-base-patch16-224") + >>> processor = AutoProcessor.from_pretrained("google/siglip-base-patch16-224") + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> inputs = processor(images=image, return_tensors="pt") + + >>> with torch.no_grad(): + ... image_features = model.get_image_features(**inputs) + ```""" + # Use SiglipModel's config for some fields (if specified) instead of those of vision & text components. + 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 + + vision_outputs: BaseModelOutputWithPooling = self.vision_model( + pixel_values=pixel_values, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + interpolate_pos_encoding=interpolate_pos_encoding, + ) + + pooled_output = vision_outputs.pooler_output + + return pooled_output + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + pixel_values: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + return_loss: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + interpolate_pos_encoding: bool = False, + ) -> SiglipOutput: + r""" + return_loss (`bool`, *optional*): + Whether or not to return the contrastive loss. + + Examples: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, AutoModel + >>> import torch + + >>> model = AutoModel.from_pretrained("google/siglip-base-patch16-224") + >>> processor = AutoProcessor.from_pretrained("google/siglip-base-patch16-224") + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> texts = ["a photo of 2 cats", "a photo of 2 dogs"] + >>> # important: we pass `padding=max_length` since the model was trained with this + >>> inputs = processor(text=texts, images=image, padding="max_length", return_tensors="pt") + + >>> with torch.no_grad(): + ... outputs = model(**inputs) + + >>> logits_per_image = outputs.logits_per_image + >>> probs = torch.sigmoid(logits_per_image) # these are the probabilities + >>> print(f"{probs[0][0]:.1%} that image 0 is '{texts[0]}'") + 31.9% that image 0 is 'a photo of 2 cats' + ```""" + # Use SigLIP model's config for some fields (if specified) instead of those of vision & text components. + 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 + + vision_outputs: BaseModelOutputWithPooling = self.vision_model( + pixel_values=pixel_values, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + interpolate_pos_encoding=interpolate_pos_encoding, + ) + + text_outputs: BaseModelOutputWithPooling = self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + ) + + image_embeds = vision_outputs.pooler_output + text_embeds = text_outputs.pooler_output + + # normalized features + image_embeds = image_embeds / image_embeds.norm(p=2, dim=-1, keepdim=True) + text_embeds = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True) + + # cosine similarity as logits + logits_per_text = torch.matmul(text_embeds, image_embeds.t().to(text_embeds.device)) + + logit_scale, logit_bias = self.logit_scale.to(text_embeds.device), self.logit_bias.to(text_embeds.device) + logits_per_text = logits_per_text * logit_scale.exp() + logit_bias + + logits_per_image = logits_per_text.t() + + loss = None + if return_loss: + # Adapted from https://github.com/google-research/big_vision/blob/01edb81a4716f93a48be43b3a4af14e29cdb3a7f/big_vision/trainers/proj/image_text/siglip.py#L287 + eye = torch.eye(logits_per_text.size(0), device=logits_per_text.device) + m1_diag1 = -torch.ones_like(logits_per_text) + 2 * eye + loglik = torch.nn.functional.logsigmoid(m1_diag1 * logits_per_text) + nll = -torch.sum(loglik, dim=-1) + loss = nll.mean() + + return SiglipOutput( + loss=loss, + logits_per_image=logits_per_image, + logits_per_text=logits_per_text, + text_embeds=text_embeds, + image_embeds=image_embeds, + text_model_output=text_outputs, + vision_model_output=vision_outputs, + ) + + +@auto_docstring( + custom_intro=""" + SigLIP vision encoder with an image classification head on top (a linear layer on top of the pooled final hidden states of + the patch tokens) e.g. for ImageNet. + """ +) +class SiglipForImageClassification(SiglipPreTrainedModel): + main_input_name = "pixel_values" + + def __init__(self, config: SiglipConfig) -> None: + super().__init__(config) + + self.num_labels = config.num_labels + + # Create the vision model with proper attention + # and take only vision_model submodule (for backward compatibility) + vision_model = SiglipVisionModel._from_config(config.vision_config) + self.vision_model = vision_model.vision_model + + # Classifier head + self.classifier = nn.Linear(config.vision_config.hidden_size, config.num_labels) if config.num_labels > 0 else nn.Identity() + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + @auto_docstring + def forward( + self, + pixel_values: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + interpolate_pos_encoding: bool = False, + ) -> ImageClassifierOutput: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the image classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + + Examples: + + ```python + >>> from transformers import AutoImageProcessor, SiglipForImageClassification + >>> import torch + >>> from PIL import Image + >>> import requests + + >>> torch.manual_seed(3) # doctest: +IGNORE_RESULT + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> # note: we are loading a `SiglipModel` from the hub here, + >>> # so the head will be randomly initialized, hence the predictions will be random if seed is not set above. + >>> image_processor = AutoImageProcessor.from_pretrained("google/siglip-base-patch16-224") + >>> model = SiglipForImageClassification.from_pretrained("google/siglip-base-patch16-224") + + >>> inputs = image_processor(images=image, return_tensors="pt") + >>> outputs = model(**inputs) + >>> logits = outputs.logits + >>> # model predicts one of the two classes + >>> predicted_class_idx = logits.argmax(-1).item() + >>> print("Predicted class:", model.config.id2label[predicted_class_idx]) + Predicted class: LABEL_1 + ```""" + 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 + + outputs: BaseModelOutputWithPooling = self.vision_model( + pixel_values, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + interpolate_pos_encoding=interpolate_pos_encoding, + ) + + sequence_output = outputs.last_hidden_state + + # average pool the patch tokens + sequence_output = torch.mean(sequence_output, dim=1) + # apply classifier + logits = self.classifier(sequence_output) + + loss = None + if labels is not None: + # move labels to correct device to enable model parallelism + labels = labels.to(logits.device) + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(logits, labels) + + return ImageClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +__all__ = [ + "SiglipModel", + "SiglipPreTrainedModel", + "SiglipTextModel", + "SiglipVisionModel", + "SiglipForImageClassification", +] diff --git a/lightx2v/models/networks/openpi/weights/__init__.py b/lightx2v/models/networks/openpi/weights/__init__.py new file mode 100644 index 000000000..62ed9f1ce --- /dev/null +++ b/lightx2v/models/networks/openpi/weights/__init__.py @@ -0,0 +1,3 @@ +from .loader import load_pi05_libero_weights, validate_transformers_runtime + +__all__ = ["load_pi05_libero_weights", "validate_transformers_runtime"] diff --git a/lightx2v/models/networks/openpi/weights/loader.py b/lightx2v/models/networks/openpi/weights/loader.py new file mode 100644 index 000000000..c8b9e2374 --- /dev/null +++ b/lightx2v/models/networks/openpi/weights/loader.py @@ -0,0 +1,66 @@ +"""Strict SafeTensors loader for the converted pi0.5-LIBERO checkpoint.""" + +from __future__ import annotations + +import logging +from pathlib import Path + +import torch +from safetensors.torch import load_model + +from ..config import Pi0Config + +LOGGER = logging.getLogger(__name__) + + +def validate_transformers_runtime() -> None: + """Fail early unless the official patched Transformers runtime is active.""" + import transformers + + if transformers.__version__ != "4.53.2": + raise RuntimeError( + "OpenPI requires its private patched transformers==4.53.2 runtime; " + f"the current process imported transformers=={transformers.__version__}. " + "Launch with scripts/openpi/run_libero_*.sh or prepend OPENPI_PYTHON_RUNTIME to PYTHONPATH." + ) + try: + from transformers.models.siglip import check + except ImportError as exc: + raise RuntimeError("OpenPI Transformers patches are missing (siglip/check.py not found)") from exc + if not check.check_whether_transformers_replace_is_installed_correctly(): + raise RuntimeError("OpenPI Transformers 4.53.2 is present but the official replacement patches are missing") + + +def load_pi05_libero_weights( + weight_path: str | Path, + config: Pi0Config, + device: torch.device | str, + *, + training: bool = False, +): + """Build the exact official parameter tree and load it with strict key checks.""" + validate_transformers_runtime() + config.validate_pi05_libero() + weight_path = Path(weight_path).expanduser().resolve() + if not weight_path.is_file(): + raise FileNotFoundError(f"Converted OpenPI SafeTensors file not found: {weight_path}") + + # Import after validating the process-local dependency layer. This module + # preserves the official parameter names, so strict SafeTensors loading is + # meaningful and does not need key rewriting. + from ..pi0 import PI0Pytorch + + model = PI0Pytorch(config) + missing, unexpected = load_model(model, weight_path, strict=True, device="cpu") + if missing or unexpected: # Defensive: strict=True normally raises first. + raise RuntimeError(f"OpenPI weight mismatch: missing={missing}, unexpected={unexpected}") + + # Match policy_config.create_trained_policy in upstream OpenPI: most + # parameters are BF16, while numerically sensitive norms/vision embeddings + # are restored to FP32. + model.paligemma_with_expert.to_bfloat16_for_selected_params(config.dtype) + model.to(torch.device(device)) + model.train(training) + parameter_count = sum(parameter.numel() for parameter in model.parameters()) + LOGGER.info("Loaded pi05_libero PyTorch weights strictly: %.3fB parameters", parameter_count / 1e9) + return model diff --git a/lightx2v/models/runners/openpi/__init__.py b/lightx2v/models/runners/openpi/__init__.py new file mode 100644 index 000000000..1c99a799c --- /dev/null +++ b/lightx2v/models/runners/openpi/__init__.py @@ -0,0 +1,5 @@ +"""OpenPI policy and offline runner integration.""" + +from .openpi_runner import OpenPIPolicy, OpenPIRunner + +__all__ = ["OpenPIPolicy", "OpenPIRunner"] diff --git a/lightx2v/models/runners/openpi/libero_rollout.py b/lightx2v/models/runners/openpi/libero_rollout.py new file mode 100644 index 000000000..b4c2e5dbc --- /dev/null +++ b/lightx2v/models/runners/openpi/libero_rollout.py @@ -0,0 +1,292 @@ +"""Run the native PyTorch OpenPI policy in a local LIBERO simulator. + +This is the non-ROS closed-loop evaluation path. OpenPI predicts robot actions, +not pixels, so an actual simulator rollout is required to produce a video of +the robot executing those actions. +""" + +from __future__ import annotations + +import argparse +import json +import logging +import math +import os +import sys +import time +from pathlib import Path + +import imageio.v2 as imageio +import numpy as np +import torch + +from .openpi_runner import OpenPIPolicy + +LOGGER = logging.getLogger(__name__) + +LIBERO_BENCHMARKS = ("libero_spatial", "libero_object", "libero_goal", "libero_10", "libero_90") +MAX_STEPS_BY_BENCHMARK = { + "libero_spatial": 220, + "libero_object": 280, + "libero_goal": 300, + "libero_10": 520, + "libero_90": 400, +} +LIBERO_DUMMY_ACTION = np.asarray([0.0] * 6 + [-1.0], dtype=np.float32) + + +def _configure_libero(libero_root: Path, config_dir: Path): + """Load the local LIBERO checkout without importing OpenPI runtime code.""" + libero_root = libero_root.expanduser().resolve() + benchmark_root = libero_root / "libero" / "libero" + required = ( + benchmark_root / "bddl_files", + benchmark_root / "init_files", + benchmark_root / "assets", + ) + missing = [str(path) for path in required if not path.is_dir()] + if missing: + raise FileNotFoundError(f"LIBERO checkout is incomplete under {libero_root}: {missing}") + + config_dir = config_dir.expanduser().resolve() + config_dir.mkdir(parents=True, exist_ok=True) + (config_dir / "config.yaml").write_text( + "\n".join( + ( + f"benchmark_root: {benchmark_root}", + f"bddl_files: {benchmark_root / 'bddl_files'}", + f"init_states: {benchmark_root / 'init_files'}", + f"datasets: {libero_root / 'libero' / 'datasets'}", + f"assets: {benchmark_root / 'assets'}", + "", + ) + ), + encoding="utf-8", + ) + os.environ["LIBERO_CONFIG_PATH"] = str(config_dir) + root_text = str(libero_root) + if root_text not in sys.path: + sys.path.insert(0, root_text) + + from libero.libero import benchmark, get_libero_path + from libero.libero.envs import OffScreenRenderEnv + + return benchmark, get_libero_path, OffScreenRenderEnv + + +def _quat_to_axis_angle(quaternion) -> np.ndarray: + quat = np.asarray(quaternion, dtype=np.float32).copy() + if quat.shape != (4,): + raise ValueError(f"Expected LIBERO quaternion shape (4,), got {quat.shape}") + quat[3] = np.clip(quat[3], -1.0, 1.0) + denominator = math.sqrt(max(0.0, 1.0 - float(quat[3]) ** 2)) + if math.isclose(denominator, 0.0): + return np.zeros(3, dtype=np.float32) + return (quat[:3] * (2.0 * math.acos(float(quat[3])) / denominator)).astype(np.float32) + + +def _rotate_rgb(observation: dict, key: str) -> np.ndarray: + if key not in observation: + raise KeyError(f"LIBERO observation is missing {key!r}") + image = np.asarray(observation[key]) + if image.ndim != 3 or image.shape[-1] != 3: + raise ValueError(f"LIBERO image {key!r} must be HWC RGB, got {image.shape}") + # LIBERO renders both policy cameras rotated by 180 degrees relative to the + # released pi05_libero training observations. + return np.ascontiguousarray(image[::-1, ::-1], dtype=np.uint8) + + +def _state_from_observation(observation: dict) -> np.ndarray: + state = np.concatenate( + ( + np.asarray(observation["robot0_eef_pos"], dtype=np.float32), + _quat_to_axis_angle(observation["robot0_eef_quat"]), + np.asarray(observation["robot0_gripper_qpos"], dtype=np.float32), + ) + ).astype(np.float32) + if state.shape != (8,) or not np.isfinite(state).all(): + raise ValueError(f"Expected finite 8-D LIBERO state, got {state.shape}") + return state + + +def _load_policy_config(args: argparse.Namespace) -> dict: + config_path = args.config_json.expanduser().resolve() + with config_path.open("r", encoding="utf-8") as handle: + config = json.load(handle) + config.update( + { + "model_cls": "openpi", + "task": "i2va", + "model_path": str(args.model_path.expanduser().resolve()), + "config_json": str(config_path), + "seed": args.seed, + "actions_per_plan": args.actions_per_plan, + } + ) + return config + + +def _validate_output_path(path: Path, suffix: str, label: str) -> Path: + path = path.expanduser().resolve() + if path.suffix.lower() != suffix: + raise ValueError(f"{label} must end in {suffix}: {path}") + path.parent.mkdir(parents=True, exist_ok=True) + return path + + +def _write_video(path: Path, frames: list[np.ndarray], fps: int) -> None: + if not frames: + raise RuntimeError("LIBERO rollout produced no video frames") + temporary = path.with_name(f".{path.stem}.tmp.mp4") + imageio.mimwrite(temporary, frames, fps=fps) + temporary.replace(path) + + +def run_rollout(args: argparse.Namespace) -> dict: + if args.actions_per_plan < 1: + raise ValueError("--actions-per-plan must be positive") + if args.num_steps_wait < 0: + raise ValueError("--num-steps-wait must be non-negative") + if args.fps < 1: + raise ValueError("--fps must be positive") + if args.render_size < 1: + raise ValueError("--render-size must be positive") + max_steps = args.max_steps if args.max_steps is not None else MAX_STEPS_BY_BENCHMARK[args.benchmark] + if max_steps < 1: + raise ValueError("--max-steps must be positive") + + save_video_path = _validate_output_path(args.save_video_path, ".mp4", "save_video_path") + save_action_path = _validate_output_path(args.save_action_path, ".npy", "save_action_path") + save_metrics_path = _validate_output_path(args.save_metrics_path, ".json", "save_metrics_path") + + benchmark_module, get_libero_path, env_type = _configure_libero(args.libero_root, args.libero_config_dir) + factories = benchmark_module.get_benchmark_dict() + task_suite = factories[args.benchmark]() + task_count = task_suite.get_num_tasks() + if not 0 <= args.task_id < task_count: + raise ValueError(f"task_id must be in [0, {task_count}), got {args.task_id}") + task = task_suite.get_task(args.task_id) + # Torch >= 2.6 defaults torch.load to weights_only=True. LIBERO init-state + # files are trusted local tensors rather than model weights, so load them + # explicitly with the legacy behavior used by the existing ROS simulator. + init_states_path = Path(get_libero_path("init_states")) / task.problem_folder / task.init_states_file + initial_states = torch.load(init_states_path, map_location="cpu", weights_only=False) + if not 0 <= args.init_state_id < len(initial_states): + raise ValueError(f"init_state_id must be in [0, {len(initial_states)}), got {args.init_state_id}") + + task_description = args.task_description.strip() or str(task.language) + bddl_path = Path(get_libero_path("bddl_files")) / task.problem_folder / task.bddl_file + env = env_type( + bddl_file_name=str(bddl_path), + camera_heights=args.render_size, + camera_widths=args.render_size, + ) + env.seed(args.seed) + + policy = None + frames: list[np.ndarray] = [] + executed_actions: list[np.ndarray] = [] + success = False + started = time.perf_counter() + + try: + LOGGER.info("Loading local PyTorch OpenPI policy") + policy = OpenPIPolicy.from_config(_load_policy_config(args)) + LOGGER.info("Starting %s task %d, init state %d: %s", args.benchmark, args.task_id, args.init_state_id, task_description) + env.reset() + observation = env.set_init_state(initial_states[args.init_state_id]) + policy.reset() + + for _ in range(args.num_steps_wait): + observation, _reward, done, _info = env.step(LIBERO_DUMMY_ACTION.tolist()) + if done: + success = True + break + + frames.append(_rotate_rgb(observation, "agentview_image")) + for step in range(0 if success else max_steps): + images = { + "agentview": _rotate_rgb(observation, "agentview_image"), + "wrist": _rotate_rgb(observation, "robot0_eye_in_hand_image"), + } + action = np.asarray( + policy.next_action(images=images, state=_state_from_observation(observation), task_description=task_description), + dtype=np.float32, + ).reshape(-1) + if action.shape != (7,) or not np.isfinite(action).all(): + raise ValueError(f"OpenPI returned an invalid LIBERO action: shape={action.shape}") + executed_actions.append(action.copy()) + observation, _reward, done, _info = env.step(action.tolist()) + frames.append(_rotate_rgb(observation, "agentview_image")) + success = bool(done) + if success or (step + 1) % 10 == 0: + LOGGER.info("Rollout step %d/%d, success=%s", step + 1, max_steps, success) + if success: + break + finally: + env.close() + if policy is not None: + policy.close() + + action_array = np.asarray(executed_actions, dtype=np.float32).reshape(-1, 7) + np.save(save_action_path, action_array) + _write_video(save_video_path, frames, args.fps) + + metrics = { + "benchmark": args.benchmark, + "task_id": args.task_id, + "task_name": str(task.name), + "task_description": task_description, + "bddl_file": str(bddl_path), + "init_states_file": str(init_states_path), + "init_state_id": args.init_state_id, + "init_state_count": len(initial_states), + "seed": args.seed, + "success": success, + "policy_steps": int(action_array.shape[0]), + "warmup_steps": args.num_steps_wait, + "actions_per_plan": args.actions_per_plan, + "video_frames": len(frames), + "video_fps": args.fps, + "elapsed_seconds": time.perf_counter() - started, + "video_path": str(save_video_path), + "action_path": str(save_action_path), + } + save_metrics_path.write_text(json.dumps(metrics, indent=2) + "\n", encoding="utf-8") + LOGGER.info("Saved rollout video: %s", save_video_path) + LOGGER.info("Saved executed actions %s: %s", action_array.shape, save_action_path) + LOGGER.info("Saved rollout metrics (success=%s): %s", success, save_metrics_path) + return metrics + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="Run a local pi05_libero PyTorch rollout and record MP4 video") + parser.add_argument("--model-path", type=Path, required=True) + parser.add_argument("--config-json", type=Path, required=True) + parser.add_argument("--libero-root", type=Path, required=True) + parser.add_argument("--libero-config-dir", type=Path, required=True) + parser.add_argument("--benchmark", choices=LIBERO_BENCHMARKS, default="libero_spatial") + parser.add_argument("--task-id", type=int, default=0) + parser.add_argument("--init-state-id", type=int, default=0) + parser.add_argument("--seed", type=int, default=7) + parser.add_argument("--task-description", default="") + parser.add_argument("--actions-per-plan", type=int, default=5) + parser.add_argument("--num-steps-wait", type=int, default=10) + parser.add_argument("--max-steps", type=int) + parser.add_argument("--render-size", type=int, default=256) + parser.add_argument("--fps", type=int, default=10) + parser.add_argument("--save-video-path", type=Path, required=True) + parser.add_argument("--save-action-path", type=Path, required=True) + parser.add_argument("--save-metrics-path", type=Path, required=True) + return parser + + +def main() -> None: + logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s") + args = build_parser().parse_args() + metrics = run_rollout(args) + print(json.dumps(metrics, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/lightx2v/models/runners/openpi/openpi_runner.py b/lightx2v/models/runners/openpi/openpi_runner.py new file mode 100644 index 000000000..ff905d34c --- /dev/null +++ b/lightx2v/models/runners/openpi/openpi_runner.py @@ -0,0 +1,352 @@ +import json +import os +import subprocess +import sys +from collections import deque +from pathlib import Path +from typing import Any + +import numpy as np +from PIL import Image +from loguru import logger + +from lightx2v.models.runners.base_runner import BaseRunner +from lightx2v.utils.registry_factory import RUNNER_REGISTER + +AGENTVIEW_IMAGE_NAME = "agentview_image.png" +WRIST_IMAGE_NAME = "wrist_image.png" + + +class OpenPIPolicy: + """Thin stateful policy wrapper around the native LightX2V OpenPI model.""" + + def __init__(self, config): + self.config = config + self.action_horizon = int(config.get("action_horizon", 10)) + self.output_action_dim = int(config.get("output_action_dim", 7)) + self.actions_per_plan = int(config.get("actions_per_plan", 5)) + if self.action_horizon <= 0: + raise ValueError(f"OpenPI requires a positive action_horizon, got {self.action_horizon}.") + if self.output_action_dim <= 0: + raise ValueError(f"OpenPI requires a positive output_action_dim, got {self.output_action_dim}.") + if not 1 <= self.actions_per_plan <= self.action_horizon: + raise ValueError(f"OpenPI actions_per_plan must be in [1, {self.action_horizon}], got {self.actions_per_plan}.") + + # Keep OpenPI/Transformers imports out of LightX2V's process startup. + # OpenPIRunner prepends the isolated patched runtime only in the local + # worker before this constructor imports the native network. + from lightx2v.models.networks.openpi import OpenPIModel + + self.model = OpenPIModel.from_config(config) + self.pending_actions = deque() + + @classmethod + def from_config(cls, config): + return cls(config) + + def predict_action_chunk(self, images, state, task_description, seed=None): + actions = self.model.predict_action_chunk( + images=images, + state=state, + task_description=task_description, + seed=seed, + ) + actions = np.asarray(actions, dtype=np.float32) + expected_shape = (self.action_horizon, self.output_action_dim) + if actions.shape != expected_shape: + raise ValueError(f"OpenPI expected action chunk shape {expected_shape}, got {actions.shape}.") + if not np.isfinite(actions).all(): + raise ValueError("OpenPI produced non-finite actions.") + return np.ascontiguousarray(actions) + + def next_action(self, images, state, task_description): + if not self.pending_actions: + # With seed=None the model advances its config-seeded generator + # between replans instead of replaying identical flow noise. + action_chunk = self.predict_action_chunk(images, state, task_description) + self.pending_actions.extend(action.copy() for action in action_chunk[: self.actions_per_plan]) + if not self.pending_actions: + raise RuntimeError("OpenPI produced an empty action chunk.") + return self.pending_actions.popleft() + + def reset(self): + self.pending_actions.clear() + self.model.reset() + + def close(self): + self.pending_actions.clear() + self.model.close() + + +@RUNNER_REGISTER("openpi") +class OpenPIRunner(BaseRunner): + """LightX2V entry point for local π0.5-LIBERO inference.""" + + def init_modules(self): + if self.config.get("task") != "i2va": + raise ValueError(f"OpenPI currently supports only task='i2va', got {self.config.get('task')!r}.") + self._worker_process = os.environ.get("OPENPI_WORKER_PROCESS") == "1" + if self._worker_process: + logger.info("Loading native PyTorch OpenPI policy in the isolated worker...") + self.policy = OpenPIPolicy.from_config(self.config) + logger.info("Native PyTorch OpenPI policy loaded.") + else: + self.run_mode = str(os.environ.get("OPENPI_RUN_MODE", self.config.get("openpi_run_mode", "rollout"))).strip() + if self.run_mode not in {"rollout", "single_observation"}: + raise ValueError(f"Unsupported OpenPI run mode {self.run_mode!r}; expected 'rollout' or 'single_observation'.") + logger.info("OpenPI selected by lightx2v.infer; dispatch mode={}", self.run_mode) + self.config.lock() + + @staticmethod + def _require_path(value, label, suffix=None): + normalized = str(value or "").strip() + if not normalized: + raise ValueError(f"OpenPI requires {label}.") + path = Path(normalized).expanduser().resolve() + if suffix is not None and path.suffix.lower() != suffix: + raise ValueError(f"OpenPI {label} must end in {suffix}: {path}") + return path + + def _worker_environment(self): + runtime_value = os.environ.get("OPENPI_TRANSFORMERS_RUNTIME_PATH", self.config.get("transformers_runtime_path", "")) + runtime_path = self._require_path(runtime_value, "transformers_runtime_path") + if not (runtime_path / "transformers").is_dir(): + raise FileNotFoundError(f"OpenPI patched Transformers package is missing: {runtime_path / 'transformers'}") + + project_root = Path(__file__).resolve().parents[4] + child_env = os.environ.copy() + child_env["OPENPI_WORKER_PROCESS"] = "1" + child_env["USE_FLAX"] = "0" + python_paths = (str(runtime_path), str(project_root), child_env.get("PYTHONPATH", "")) + child_env["PYTHONPATH"] = os.pathsep.join(path for path in python_paths if path) + return child_env + + def _single_observation_command(self, input_info): + task_description = str(input_info.prompt or "").strip() + if not task_description: + raise ValueError("OpenPI single_observation mode requires prompt as the LIBERO task_description.") + return [ + "-m", + "lightx2v.models.runners.openpi.single_observation", + "--model-path", + str(self._require_path(self.config.get("model_path"), "model_path")), + "--config-json", + str(self._require_path(self.config.get("config_json"), "config_json", ".json")), + "--seed", + str(input_info.seed), + "--task-description", + task_description, + "--image-path", + str(input_info.image_path or ""), + "--state-path", + str(input_info.state_path or ""), + "--save-action-path", + str(self._require_path(input_info.save_action_path, "save_action_path", ".npy")), + ] + + def _rollout_command(self, input_info): + video_path = self._require_path(input_info.save_result_path, "save_result_path", ".mp4") + action_path = self._require_path(input_info.save_action_path, "save_action_path", ".npy") + metrics_value = os.environ.get("OPENPI_SAVE_METRICS_PATH", "").strip() + metrics_path = self._require_path(metrics_value, "save_metrics_path", ".json") if metrics_value else video_path.with_suffix(".metrics.json") + libero_root = os.environ.get("OPENPI_LIBERO_ROOT", "/data/liuhongda/openpi/third_party/libero") + libero_config_dir = os.environ.get("OPENPI_LIBERO_CONFIG_DIR", "/data/liuhongda/openpi_data/runtime_configs/lightx2v_openpi_libero") + command = [ + "-m", + "lightx2v.models.runners.openpi.libero_rollout", + "--model-path", + str(self._require_path(self.config.get("model_path"), "model_path")), + "--config-json", + str(self._require_path(self.config.get("config_json"), "config_json", ".json")), + "--libero-root", + str(self._require_path(libero_root, "libero_root")), + "--libero-config-dir", + str(Path(libero_config_dir).expanduser().resolve()), + "--benchmark", + os.environ.get("LIBERO_BENCHMARK", "libero_spatial"), + "--task-id", + os.environ.get("LIBERO_TASK_ID", "0"), + "--init-state-id", + os.environ.get("LIBERO_INIT_STATE_ID", "0"), + "--seed", + str(input_info.seed), + "--actions-per-plan", + os.environ.get("OPENPI_ACTIONS_PER_PLAN", str(self.config.get("actions_per_plan", 5))), + "--num-steps-wait", + os.environ.get("OPENPI_NUM_STEPS_WAIT", str(self.config.get("num_steps_wait", 10))), + "--render-size", + os.environ.get("OPENPI_RENDER_SIZE", "256"), + "--fps", + os.environ.get("OPENPI_VIDEO_FPS", "10"), + "--save-video-path", + str(video_path), + "--save-action-path", + str(action_path), + "--save-metrics-path", + str(metrics_path), + ] + task_description = str(input_info.prompt or "").strip() + if task_description: + command.extend(("--task-description", task_description)) + max_steps = os.environ.get("OPENPI_MAX_STEPS", "").strip() + if max_steps: + command.extend(("--max-steps", max_steps)) + return command + + def _run_isolated_worker(self, input_info): + child_env = self._worker_environment() + python_bin = os.environ.get("OPENPI_PYTHON", sys.executable) + if self.run_mode == "rollout": + arguments = self._rollout_command(input_info) + else: + arguments = self._single_observation_command(input_info) + command = [python_bin, *arguments] + logger.info("Starting synchronous local OpenPI {} worker", self.run_mode) + subprocess.run(command, env=child_env, check=True) + logger.info("OpenPI {} worker completed", self.run_mode) + + if input_info.return_result_tensor: + action_path = self._require_path(input_info.save_action_path, "save_action_path", ".npy") + return {"actions": np.load(action_path)} + return {"actions": None} + + @staticmethod + def _load_rgb(path): + image_path = Path(path).expanduser().resolve() + if not image_path.is_file(): + raise FileNotFoundError(f"OpenPI image does not exist: {image_path}") + image = np.asarray(Image.open(image_path).convert("RGB"), dtype=np.uint8) + if image.ndim != 3 or image.shape[-1] != 3: + raise ValueError(f"OpenPI expected an HWC RGB image, got {image.shape} from {image_path}.") + # The shared LIBERO simulator already performs the official 180-degree + # rotation. Neither the runner nor the network flips these inputs again. + return np.ascontiguousarray(image) + + def _load_image_pair(self): + policy_image = getattr(self.input_info, "policy_image", None) + if policy_image is not None: + if not isinstance(policy_image, dict): + raise TypeError("OpenPI policy_image must be a dict with logical keys 'agentview' and 'wrist'.") + missing = [key for key in ("agentview", "wrist") if key not in policy_image] + if missing: + raise KeyError(f"OpenPI policy_image is missing camera keys: {missing}") + images = {key: np.asarray(policy_image[key], dtype=np.uint8) for key in ("agentview", "wrist")} + for key, image in images.items(): + if image.ndim != 3 or image.shape[-1] != 3: + raise ValueError(f"OpenPI camera '{key}' must be HWC RGB, got {image.shape}.") + images[key] = np.ascontiguousarray(image) + return images + + image_path = str(getattr(self.input_info, "image_path", "") or "").strip() + if not image_path: + raise ValueError("OpenPI requires image_path.") + expanded = os.path.expanduser(image_path) + if os.path.isdir(expanded): + agentview_path = os.path.join(expanded, AGENTVIEW_IMAGE_NAME) + wrist_path = os.path.join(expanded, WRIST_IMAGE_NAME) + else: + paths = [item.strip() for item in expanded.split(",") if item.strip()] + if len(paths) != 2: + raise ValueError("OpenPI image_path must be a directory containing agentview_image.png and wrist_image.png, or two comma-separated image paths in that order.") + agentview_path, wrist_path = paths + return { + "agentview": self._load_rgb(agentview_path), + "wrist": self._load_rgb(wrist_path), + } + + @staticmethod + def _unwrap_state_payload(payload: Any): + if isinstance(payload, dict): + for key in ("state", "qpos", "robot_state", "observation.state", "observation/state"): + if key in payload: + return payload[key] + raise KeyError("OpenPI state mapping must contain one of: state, qpos, robot_state, observation.state, observation/state.") + return payload + + @staticmethod + def _parse_text_state(text): + normalized = str(text).strip() + if not normalized: + return np.empty((0,), dtype=np.float32) + try: + payload = json.loads(normalized) + except json.JSONDecodeError: + payload = None + if payload is not None: + return OpenPIRunner._unwrap_state_payload(payload) + normalized = normalized.translate(str.maketrans({",": " ", "[": " ", "]": " ", "(": " ", ")": " ", ";": " "})) + return np.fromstring(normalized, sep=" ", dtype=np.float32) + + def _load_state(self): + policy_state = getattr(self.input_info, "policy_state", None) + if policy_state is not None: + payload = policy_state + else: + state_source = str(getattr(self.input_info, "state_path", "") or "").strip() + if not state_source: + raise ValueError("OpenPI requires state_path containing the 8-dimensional LIBERO state.") + state_path = Path(state_source).expanduser() + if state_path.is_file(): + suffix = state_path.suffix.lower() + if suffix == ".npy": + payload = np.load(state_path, allow_pickle=True) + if isinstance(payload, np.ndarray) and payload.shape == () and isinstance(payload.item(), dict): + payload = payload.item() + elif suffix == ".npz": + with np.load(state_path, allow_pickle=True) as archive: + payload = {key: archive[key] for key in archive.files} + else: + payload = self._parse_text_state(state_path.read_text(encoding="utf-8")) + else: + # Inline comma/space-separated state text is convenient for + # smoke tests while retaining state_path CLI compatibility. + payload = self._parse_text_state(state_source) + + payload = self._unwrap_state_payload(payload) + state = np.asarray(payload, dtype=np.float32).reshape(-1) + expected_dim = int(self.config.get("state_dim", 8)) + if state.size != expected_dim: + raise ValueError(f"OpenPI LIBERO state must contain {expected_dim} floats, got {state.size}.") + if not np.isfinite(state).all(): + raise ValueError("OpenPI state contains non-finite values.") + return np.ascontiguousarray(state) + + def _resolve_action_output_path(self): + save_action_path = str(getattr(self.input_info, "save_action_path", "") or "").strip() + if not save_action_path: + raise ValueError("OpenPI offline inference requires save_action_path.") + output_path = Path(save_action_path).expanduser().resolve() + if output_path.suffix.lower() != ".npy": + raise ValueError(f"OpenPI save_action_path must end in .npy, got: {output_path}") + return output_path + + def run_pipeline(self, input_info): + if not self._worker_process: + return self._run_isolated_worker(input_info) + + self.input_info = input_info + task_description = str(self.input_info.prompt or "").strip() + if not task_description: + raise ValueError("OpenPI requires prompt as the LIBERO task_description.") + + images = self._load_image_pair() + state = self._load_state() + actions = self.policy.predict_action_chunk( + images=images, + state=state, + task_description=task_description, + seed=self.input_info.seed, + ) + + output_path = self._resolve_action_output_path() + output_path.parent.mkdir(parents=True, exist_ok=True) + np.save(output_path, actions) + logger.info("Saved OpenPI action chunk {} to {}", actions.shape, output_path) + + if self.input_info.return_result_tensor: + return {"actions": actions} + return {"actions": None} + + def end_run(self): + if hasattr(self, "policy"): + self.policy.close() diff --git a/lightx2v/models/runners/openpi/single_observation.py b/lightx2v/models/runners/openpi/single_observation.py new file mode 100644 index 000000000..e2172f621 --- /dev/null +++ b/lightx2v/models/runners/openpi/single_observation.py @@ -0,0 +1,73 @@ +"""Run one OpenPI image/state observation without importing the shared CLI.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +from lightx2v.utils.input_info import init_empty_input_info +from lightx2v.utils.lockable_dict import LockableDict + +from .openpi_runner import OpenPIRunner + + +def _load_config(args: argparse.Namespace) -> LockableDict: + config_path = args.config_json.expanduser().resolve() + if not config_path.is_file(): + raise FileNotFoundError(f"OpenPI config JSON does not exist: {config_path}") + + model_path = args.model_path.expanduser().resolve() + if not (model_path / "model.safetensors").is_file(): + raise FileNotFoundError(f"OpenPI PyTorch checkpoint is incomplete: {model_path}") + + with config_path.open("r", encoding="utf-8") as handle: + config = json.load(handle) + config.update( + { + "model_cls": "openpi", + "task": "i2va", + "model_path": str(model_path), + "config_json": str(config_path), + "seed": args.seed, + "warmup": False, + } + ) + return LockableDict(config) + + +def run_single_observation(args: argparse.Namespace) -> Path: + runner = OpenPIRunner(_load_config(args)) + try: + runner.init_modules() + input_info = init_empty_input_info("i2va") + input_info.seed = args.seed + input_info.prompt = args.task_description + input_info.image_path = str(args.image_path) + input_info.state_path = str(args.state_path) + input_info.save_action_path = str(args.save_action_path) + runner.run_pipeline(input_info) + finally: + runner.end_run() + return args.save_action_path.expanduser().resolve() + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="Run one local pi05-LIBERO image/state observation") + parser.add_argument("--model-path", type=Path, required=True) + parser.add_argument("--config-json", type=Path, required=True) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--task-description", required=True) + parser.add_argument("--image-path", type=Path, required=True) + parser.add_argument("--state-path", type=Path, required=True) + parser.add_argument("--save-action-path", type=Path, required=True) + return parser + + +def main() -> None: + output_path = run_single_observation(build_parser().parse_args()) + print(f"Saved OpenPI action chunk: {output_path}") + + +if __name__ == "__main__": + main() diff --git a/lightx2v_ros/src/inference/inference/openpi_node/__init__.py b/lightx2v_ros/src/inference/inference/openpi_node/__init__.py new file mode 100644 index 000000000..67bb3945f --- /dev/null +++ b/lightx2v_ros/src/inference/inference/openpi_node/__init__.py @@ -0,0 +1 @@ +"""ROS2 inference node for the OpenPI policy.""" diff --git a/lightx2v_ros/src/inference/inference/openpi_node/main.py b/lightx2v_ros/src/inference/inference/openpi_node/main.py new file mode 100644 index 000000000..dda3b1549 --- /dev/null +++ b/lightx2v_ros/src/inference/inference/openpi_node/main.py @@ -0,0 +1,216 @@ +import numpy as np +import rclpy +from common.contract import get_contract +from rclpy.node import Node +from sensor_msgs.msg import Image +from std_msgs.msg import Bool, Float32MultiArray, Int32, String + +from lightx2v.models.runners.openpi.openpi_runner import OpenPIPolicy +from lightx2v.utils.set_config import auto_calc_config, get_default_config + + +class OpenPINode(Node): + """Thin ROS bridge between the shared LIBERO simulator and OpenPI.""" + + def __init__(self): + super().__init__("openpi_node") + + self.declare_parameter("env", "libero") + self.declare_parameter("config_json", "") + self.declare_parameter("model_path", "") + self.declare_parameter("seed", 7) + self.declare_parameter("actions_per_plan", -1) + self.declare_parameter("num_steps_wait", -1) + + env = str(self.get_parameter("env").value).strip().lower() + self.contract = get_contract(env) + if self.contract.name != "libero": + raise ValueError("OpenPI ROS integration currently supports only LIBERO.") + + self.get_logger().info("[libero] loading OpenPI policy") + self.policy_config = self.build_policy_config() + self.policy = OpenPIPolicy.from_config(self.policy_config) + self.get_logger().info("[libero] OpenPI policy loaded") + + self.images = {camera: None for camera in self.contract.policy_input_cameras} + self.state = None + self.task_description = None + self.success = False + self.episode_index = 0 + self.episode_observation_count = 0 + self.last_processed_observation = -1 + + configured_wait = int(self.get_parameter("num_steps_wait").value) + self.num_steps_wait = configured_wait if configured_wait >= 0 else int(self.policy_config.get("num_steps_wait", 10)) + self.dummy_action = np.zeros(self.contract.action_dim, dtype=np.float32) + self.dummy_action[-1] = -1.0 + + self.action_pub = self.create_publisher(Float32MultiArray, self.contract.action_topic, 10) + self._camera_subs = [] + for camera in self.contract.policy_input_cameras: + self._camera_subs.append(self.create_subscription(Image, self.contract.camera_topic(camera), self._make_image_cb(camera), 10)) + self.create_subscription(Float32MultiArray, self.contract.state_topic, self.on_state, 10) + self.create_subscription(String, self.contract.task_topic, self.on_task, 10) + self.create_subscription(Bool, self.contract.success_topic, self.on_success, 10) + self.create_subscription(Int32, self.contract.episode_topic, self.on_episode, 10) + self.create_subscription(Int32, self.contract.observation_ready_topic, self.on_observation_ready, 10) + + self.get_logger().info( + "[libero] openpi_node ready: " + f"cameras={list(self.contract.policy_input_cameras)} action_dim={self.contract.action_dim} " + f"state_dim={self.contract.state_dim} actions_per_plan={self.policy_config.get('actions_per_plan', 5)} " + f"num_steps_wait={self.num_steps_wait}" + ) + + def build_policy_config(self): + config_json = str(self.get_parameter("config_json").value).strip() + if not config_json: + raise ValueError("OpenPI ROS node requires `config_json`.") + model_path = str(self.get_parameter("model_path").value).strip() + if not model_path: + raise ValueError("OpenPI ROS node requires `model_path`.") + + seed = int(self.get_parameter("seed").value) + config = get_default_config() + config.update( + { + "model_cls": "openpi", + "task": "i2va", + "model_path": model_path, + "config_json": config_json, + "seed": seed, + } + ) + config = auto_calc_config(config) + + # Explicit ROS parameters are runtime overrides and must win over JSON. + config["model_cls"] = "openpi" + config["task"] = "i2va" + config["model_path"] = model_path + config["seed"] = seed + actions_per_plan = int(self.get_parameter("actions_per_plan").value) + if actions_per_plan > 0: + config["actions_per_plan"] = actions_per_plan + + expected_dims = { + "state_dim": self.contract.state_dim, + "output_action_dim": self.contract.action_dim, + } + for key, expected in expected_dims.items(): + if key in config and int(config[key]) != expected: + raise ValueError(f"OpenPI config `{key}`={config[key]} does not match LIBERO contract ({expected}).") + return config + + def _make_image_cb(self, camera): + def _callback(msg): + # The shared LIBERO simulator already applies the official OpenPI + # 180-degree observation rotation. Do not flip the image again here. + self.images[camera] = image_msg_to_rgb(msg) + + return _callback + + def on_state(self, msg): + state = np.asarray(msg.data, dtype=np.float32).reshape(-1) + if state.size != self.contract.state_dim: + self.get_logger().error(f"expected state length {self.contract.state_dim}, got {state.size}") + return + self.state = state + + def on_task(self, msg): + self.task_description = str(msg.data).strip() + + def on_success(self, msg): + self.success = bool(msg.data) + + def on_episode(self, msg): + episode = int(msg.data) + if episode == self.episode_index: + return + self.episode_index = episode + self.success = False + self.episode_observation_count = 0 + self.last_processed_observation = -1 + self.policy.reset() + self.get_logger().info(f"new episode {episode}; cleared OpenPI action queue") + + def missing_inputs(self): + missing = [camera for camera in self.contract.policy_input_cameras if self.images.get(camera) is None] + if self.state is None: + missing.append("state") + if not self.task_description: + missing.append("task_description") + return missing + + def on_observation_ready(self, msg): + observation_index = int(msg.data) + if observation_index <= self.last_processed_observation: + return + if self.success: + self.last_processed_observation = observation_index + return + + missing = self.missing_inputs() + if missing: + self.get_logger().warning(f"observation {observation_index} waiting for: {missing}") + return + + if self.episode_observation_count < self.num_steps_wait: + action = self.dummy_action.copy() + self.get_logger().info(f"observation {observation_index}: publishing LIBERO warmup action ({self.episode_observation_count + 1}/{self.num_steps_wait})") + else: + self.get_logger().info(f"observation {observation_index}: running/consuming OpenPI action chunk") + action = self.policy.next_action( + images={camera: self.images[camera] for camera in self.contract.policy_input_cameras}, + state=self.state, + task_description=self.task_description, + ) + + self.publish_action(action) + self.episode_observation_count += 1 + self.last_processed_observation = observation_index + + def publish_action(self, action): + action = np.asarray(action, dtype=np.float32).reshape(-1) + if action.size != self.contract.action_dim: + raise ValueError(f"expected action length {self.contract.action_dim}, got {action.size}") + if not np.isfinite(action).all(): + raise ValueError("OpenPI produced a non-finite action.") + msg = Float32MultiArray() + msg.data = action.tolist() + self.action_pub.publish(msg) + + def destroy_node(self): + if hasattr(self, "policy"): + self.policy.close() + super().destroy_node() + + +def image_msg_to_rgb(msg): + encoding = msg.encoding.lower() + if encoding not in {"rgb8", "bgr8"}: + raise ValueError(f"unsupported image encoding: {msg.encoding}") + row = np.frombuffer(msg.data, dtype=np.uint8).reshape(msg.height, msg.step) + image = row[:, : msg.width * 3].reshape(msg.height, msg.width, 3) + if encoding == "bgr8": + image = image[:, :, ::-1] + return np.ascontiguousarray(image.copy()) + + +def main(args=None): + rclpy.init(args=args) + node = OpenPINode() + try: + rclpy.spin(node) + except KeyboardInterrupt: + pass + except Exception: + if rclpy.ok(): + raise + finally: + node.destroy_node() + if rclpy.ok(): + rclpy.shutdown() + + +if __name__ == "__main__": + main() diff --git a/lightx2v_ros/src/inference/setup.py b/lightx2v_ros/src/inference/setup.py index ce367aec1..95ca933aa 100644 --- a/lightx2v_ros/src/inference/setup.py +++ b/lightx2v_ros/src/inference/setup.py @@ -21,6 +21,7 @@ "cosmos3_node = inference.cosmos3_node.main:main", "fastwam_node = inference.fastwam_node.main:main", "lingbot_va_node = inference.lingbot_va_node.main:main", + "openpi_node = inference.openpi_node.main:main", ], }, ) diff --git a/scripts/openpi/INTEGRATION_REPORT.md b/scripts/openpi/INTEGRATION_REPORT.md new file mode 100644 index 000000000..99d119a36 --- /dev/null +++ b/scripts/openpi/INTEGRATION_REPORT.md @@ -0,0 +1,264 @@ +# OpenPI pi0.5-LIBERO PyTorch 转换与 LightX2V 接入报告 + +## 结论 + +本次接入采用以下单一路径： + +```text +官方 pi05_libero Orbax/JAX checkpoint + -> 官方 convert_jax_model_to_pytorch.py + -> model.safetensors + -> LightX2V 原生 PyTorch network + -> python -m lightx2v.infer + -> registry OpenPIRunner + -> 同步本地隔离 worker + -> 非 ROS 本地 LIBERO 闭环或静态 i2va +``` + +LightX2V policy/network 运行时不 import OpenPI 的模型代码、JAX、Flax 或 Orbax； +本地闭环只使用 `/data/liuhongda/openpi/third_party/libero` 中的 LIBERO simulator。 +整个调用链不启动 OpenPI server、网页或 viewer。这里的 worker 只是由 +`OpenPIRunner` 同步等待的本地进程，用于隔离 Python 依赖，不是常驻服务或异步 +推理后端。ROS 闭环仍走后文所述的 `openpi_node` 链路。 + +## 权重转换结果 + +- 本地 OpenPI 源码版本：`15a9616a00943ada6c20a0f158e3adb39df2ccac` +- 转换入口：OpenPI 官方 `examples/convert_jax_model_to_pytorch.py` +- 转换配置：`pi05_libero`，输出精度 `bfloat16` +- JAX 输入： + `/data/liuhongda/openpi_data/openpi-assets/checkpoints/pi05_libero` +- PyTorch 输出： + `/data/liuhongda/openpi_data/openpi-assets/checkpoints/pi05_libero_pytorch` +- `model.safetensors`：7,233,650,408 bytes，812 个 tensor keys +- 逻辑 state dict 为 813 keys；其中 1 个是 SafeTensors metadata 记录的合法 tied-weight alias +- key/shape/dtype manifest SHA256： + `ee81d609ff73d395731f9f3df3b0caefcbd17f83d2ff153d26166e1bd024e20d` +- SHA256： + `812e78eb87ddcf6acd87acaf1561a0adcda7ed035b5733b4371064c2f1a16a77` +- 参数总量：3,616,757,520 +- 严格加载结果：`missing=[]`、`unexpected=[]` + +完整权重目录： + +```text +pi05_libero_pytorch/ +├── model.safetensors +├── config.json +├── SHA256SUMS +└── assets/ + ├── paligemma_tokenizer.model + └── physical-intelligence/libero/norm_stats.json +``` + +`assets` 是显式补齐的：官方转换脚本查找 `checkpoint_dir.parent/assets`，而本地发布 +权重把它放在 `checkpoint_dir/assets`。 + +## 模型结构 + +- vision prefix：SigLIP，hidden 1152，27 层，16 heads +- language prefix：Gemma 2B，width 2048，18 层，8 heads，1 KV head +- action expert：Gemma 300M，width 1024，18 层，8 heads，1 KV head +- π0.5 AdaRMS：由 timestep MLP 条件控制 action expert +- action：内部 padding 到 32 维，horizon 为 10 +- sampler：10 步 Euler flow matching +- LIBERO 输出：反归一化后截取前 7 维，得到 `(10, 7)` + +关键参数树保持官方名字不变，例如： + +```text +paligemma_with_expert.paligemma.* +paligemma_with_expert.gemma_expert.* +action_in_proj.* +action_out_proj.* +time_mlp_in.* +time_mlp_out.* +``` + +## LightX2V 文件组织 + +```text +configs/openpi/pi05_libero.json +lightx2v/models/networks/openpi/ +├── config.py +├── gemma.py +├── image_tools.py +├── model.py +├── observation.py +├── pi0.py +├── preprocessing.py +├── infer/{pre_infer.py,transformer_infer.py,post_infer.py} +├── weights/loader.py +└── transformers_replace/... +lightx2v/models/runners/openpi/openpi_runner.py +lightx2v/models/runners/openpi/libero_rollout.py +lightx2v/models/runners/openpi/single_observation.py +lightx2v_ros/src/inference/inference/openpi_node/main.py +scripts/openpi/ +├── convert_pi05_libero_to_pytorch.sh +├── setup_pytorch_runtime.sh +├── prepare_libero_sample.py +├── validate_pytorch_parity.py +├── run_libero_i2va.sh +├── run_libero_task_i2va.sh +└── run_libero_ros_i2va.sh +``` + +网络按 LightX2V 的 `networks / infer / weights / runners` 边界组织。没有强行继承 +视频生成用的 `BaseTransformerModel`，因为那会改变官方 SafeTensors 参数树。 + +## 输入到输出 + +```text +agentview RGB + wrist RGB + state(8) + task description + -> 两张有效图 + 一张全零、mask=false 的右腕占位图 + -> resize-with-pad 224x224 + -> state q01/q99 quantile normalize + pad 到 32 + -> PaliGemma SentencePiece tokenize + pad 到 200 + -> pi0.5 sample_actions，输出 (1,10,32) + -> actions q01/q99 unnormalize + -> slice 前 7 维，输出 float32 (10,7) +``` + +与官方 OpenPI 的 image、mask、state、token 和 token mask 预处理逐元素完全一致。 + +## 环境隔离 + +base 未被修改： + +```text +/opt/conda/bin/python +torch 2.8.0+cu128 +transformers 5.14.1 +``` + +`run_libero_i2va.sh` 首先使用这个 base 环境进入公共入口： + +```text +python -m lightx2v.infer + -> RUNNER_REGISTER["openpi"] + -> OpenPIRunner +``` + +公共入口会 eager import 其他 LightX2V runner，其中 Motus、HiDream 等模型需要 +Transformers 5.14.1 提供的 Qwen3-VL API。另一方面，OpenPI 官方 PyTorch 实现和 +五个 replacement 文件严格绑定 Transformers 4.53.2。因此不能在启动公共入口前 +把 OpenPI 的 4.53.2 放到全局 `PYTHONPATH`，也不能在同一个 Python 进程里动态 +替换已经 import 的 Transformers。 + +`OpenPIRunner` 在 registry 正常完成选择后，才同步启动本地 worker。只有 worker +进程的 `PYTHONPATH` 前置： + +```text +/data/liuhongda/openpi_data/python_deps/openpi_pytorch_runtime +``` + +其中包含 `transformers 4.53.2`、`huggingface-hub 0.32.3`、 +`tokenizers 0.21.1` 和官方五个 Transformers replacement 文件。闭环 rollout 和 +静态单观测模式都从 `lightx2v.infer` 进入 `OpenPIRunner`；隔离 worker 内部再分别 +执行对应的 OpenPI 本地逻辑。`lightx2v.infer` 仅按已有模型风格增加 runner import +和 `model_cls` choice，没有基于 `sys.argv` 的 OpenPI 特殊分支。 + +worker 设置 `USE_FLAX=0`，防止私有 Transformers 因 base 中存在 Flax 而自动 +加载 JAX。实测导入 patched Transformers/SigLIP 后 `jax_loaded=False`、 +`flax_loaded=False`。公共进程不会 import OpenPI network，worker 退出码会由 +`OpenPIRunner` 检查；worker 失败会使公共推理命令失败，所以脚本仍然是同步、 +可失败感知的单条本地调用链。 + +## 启动方式 + +转换和私有运行时： + +```bash +cd /data/liuhongda/LightX2V +bash scripts/openpi/convert_pi05_libero_to_pytorch.sh +bash scripts/openpi/setup_pytorch_runtime.sh +``` + +本地 LIBERO 闭环推理并录制机器人执行视频： + +```bash +bash scripts/openpi/run_libero_i2va.sh +``` + +该脚本的公共调用链为： + +```text +run_libero_i2va.sh + -> python -m lightx2v.infer --model_cls openpi --task i2va ... + -> RUNNER_REGISTER["openpi"] + -> OpenPIRunner + -> 同步本地隔离 worker +``` + +默认输出 MP4、实际执行动作轨迹和成功指标。切换任务： + +```bash +LIBERO_BENCHMARK=libero_goal \ +LIBERO_TASK_ID=3 \ +LIBERO_INIT_STATE_ID=0 \ +OPENPI_SAVE_VIDEO_PATH=/absolute/output/rollout.mp4 \ +OPENPI_SAVE_ACTION_PATH=/absolute/output/rollout.actions.npy \ +bash scripts/openpi/run_libero_i2va.sh +``` + +单帧静态 smoke inference 仍可用： + +```bash +OPENPI_RUN_MODE=single_observation \ +OPENPI_IMAGE_PATH=/absolute/input_dir \ +OPENPI_STATE_PATH=/absolute/input_dir/state.npy \ +OPENPI_TASK_DESCRIPTION="pick up the black bowl and place it on the plate" \ +OPENPI_SAVE_ACTION_PATH=/absolute/output/actions.npy \ +bash scripts/openpi/run_libero_i2va.sh +``` + +ROS 闭环： + +```bash +LIBERO_BENCHMARK=libero_10 \ +LIBERO_TASK_ID=5 \ +LIBERO_INIT_STATE_ID=0 \ +bash scripts/openpi/run_libero_ros_i2va.sh +``` + +ROS 复用 LightX2V 的共享 LIBERO simulator contract。simulator 已执行官方要求的 +180 度图像旋转，所以 OpenPI node、runner 和 network 均不再次 flip。 + +## 已完成验证 + +- 官方转换命令成功完成 +- SafeTensors 六组关键参数前缀存在 +- SafeTensors 的 key/shape/dtype manifest、参数量和 checksum 全部通过 +- 官方模型严格加载 36.17 亿参数，无 missing/unexpected keys +- LightX2V 本地模型严格加载成功 +- 官方预处理与 LightX2V 预处理逐元素一致 +- H200 上使用 base Python 完成本地真实前向 +- H200 上完成 `libero_spatial/task0/init0` 本地闭环：68 个 policy steps 后 + BDDL success +- 视频： + `/data/liuhongda/LightX2V/save_results/output_openpi_pi05_libero.mp4` +- 视频属性：H.264 MP4、256x256、10 FPS、69 帧、6.9 秒 +- 实际执行动作：float32、shape `(68, 7)`、全部 finite +- 指标： + `/data/liuhongda/LightX2V/save_results/output_openpi_pi05_libero.metrics.json` +- 官方 OpenPI PyTorch 与 LightX2V 固定输入/固定 noise： + max abs error `2.613627914094252e-08`，`atol=1e-6` 通过 +- 对齐报告： + `/data/liuhongda/openpi_data/results/pi05_libero_pytorch_parity.json` +- Python AST/compile、JSON、shell、`git diff --check` 均通过 + +非 ROS 本地 MuJoCo 闭环已经实际跑通。当前机器没有可 source 的 ROS2/colcon 环境， +因此 ROS 节点完成了静态契约检查，但没有执行 ROS 版本的闭环 rollout。 + +## 训练边界 + +`PI0Pytorch.forward()`、训练 loss 和 gradient-checkpointing 接口已随官方 PyTorch +架构保留在 `networks/openpi`，可作为后续 LightX2V kernel/offload/量化优化入口。 +本次两个目标 runner 是离线推理和 ROS 推理；数据加载、optimizer、DDP 与 checkpoint +保存仍以 OpenPI README 中的 `scripts/train_pytorch.py pi05_libero ...` 为已验证基线， +没有在本次改动中伪造一套未验证的 LightX2V 训练调度器。 + +如果后续在 LightX2V 内补训练，需要同时补 LIBERO dataset/collator、action 的 +quantile normalize + 32 维 padding、optimizer/DDP 和 checkpoint 保存约定；应保存 +`core_model.state_dict()`，避免外层 `OpenPIModel` 自动增加 `core_model.` 前缀。 diff --git a/scripts/openpi/README.md b/scripts/openpi/README.md new file mode 100644 index 000000000..44c661ef6 --- /dev/null +++ b/scripts/openpi/README.md @@ -0,0 +1,215 @@ +# OpenPI pi0.5-LIBERO in LightX2V + +This integration runs the released `pi05_libero` policy as a native local +PyTorch model. It does not start an OpenPI policy server, web page, or viewer. + +## 1. Convert the released checkpoint + +```bash +cd /data/liuhongda/LightX2V +bash scripts/openpi/convert_pi05_libero_to_pytorch.sh +``` + +The default output is: + +```text +/data/liuhongda/openpi_data/openpi-assets/checkpoints/pi05_libero_pytorch/ +├── model.safetensors +├── config.json +└── assets/ + ├── paligemma_tokenizer.model + └── physical-intelligence/libero/norm_stats.json +``` + +## 2. Prepare the isolated PyTorch dependency layer + +```bash +bash scripts/openpi/setup_pytorch_runtime.sh +``` + +This creates +`/data/liuhongda/openpi_data/python_deps/openpi_pytorch_runtime` with patched +Transformers 4.53.2. It does not change packages installed in the base +environment. The shared `lightx2v.infer` process continues to use the base +Transformers 5.14.1 installation. Only the synchronous OpenPI worker prepends +the private runtime to its `PYTHONPATH`; that worker also sets `USE_FLAX=0` so +Transformers does not auto-import the base environment's JAX/Flax packages. + +The split is required because the official OpenPI PyTorch implementation uses +five replacement files tied to Transformers 4.53.2, while other LightX2V +runners imported by the shared entry require newer Transformers APIs such as +Qwen3-VL. The two versions are therefore never imported into the same Python +process. + +## 3. Local closed-loop LIBERO rollout video + +OpenPI predicts actions rather than video pixels. The default launcher now +runs those actions in the local LIBERO MuJoCo simulator and records the real +agent-view rollout: + +```bash +cd /data/liuhongda/LightX2V +bash scripts/openpi/run_libero_i2va.sh +``` + +The launcher follows the same public inference entry used by other LightX2V +models: + +```text +scripts/openpi/run_libero_i2va.sh + -> python -m lightx2v.infer + -> RUNNER_REGISTER["openpi"] + -> OpenPIRunner + -> synchronous local OpenPI/LIBERO worker + -> MP4 + executed actions + metrics +``` + +The shared entry and registry run in the base environment. `OpenPIRunner` +waits for the local worker and propagates a worker failure back to the command. +The worker is dependency isolation only: it is not a policy server, web +service, viewer, or asynchronous background process. + +Default outputs: + +```text +save_results/output_openpi_pi05_libero.mp4 +save_results/output_openpi_pi05_libero.actions.npy +save_results/output_openpi_pi05_libero.metrics.json +``` + +Select another LIBERO task or output path with environment variables: + +```bash +LIBERO_BENCHMARK=libero_goal \ +LIBERO_TASK_ID=3 \ +LIBERO_INIT_STATE_ID=0 \ +OPENPI_SAVE_VIDEO_PATH=/absolute/output/rollout.mp4 \ +OPENPI_SAVE_ACTION_PATH=/absolute/output/rollout.actions.npy \ +OPENPI_SAVE_METRICS_PATH=/absolute/output/rollout.metrics.json \ +bash scripts/openpi/run_libero_i2va.sh +``` + +The rollout performs 10 dummy stabilization steps, replans every 5 policy +steps, records the correctly oriented `agentview_image` at 10 FPS, and stops +when LIBERO reports BDDL task success or the suite-specific step cap is hit. + +## 4. Run a sample by LIBERO benchmark and task ID + +Use the task launcher when selecting examples from the local LIBERO checkout: + +```text +/data/liuhongda/openpi/third_party/libero/libero/libero/bddl_files +/data/liuhongda/openpi/third_party/libero/libero/libero/init_files +``` + +The first argument is the benchmark (task suite), the second is the zero-based +task ID, and the optional third argument is the zero-based initialization-state +ID (default: `0`). Every bundled `.pruned_init` contains 50 states, so valid +initialization-state IDs are `0-49`: + +```bash +cd /data/liuhongda/LightX2V +bash scripts/openpi/run_libero_task_i2va.sh libero_goal 3 0 +``` + +Supported benchmark and task-ID ranges: + +| Benchmark | Task IDs | +| --- | ---: | +| `libero_spatial` | 0-9 | +| `libero_object` | 0-9 | +| `libero_goal` | 0-9 | +| `libero_10` | 0-9 | +| `libero_90` | 0-89 | + +LIBERO's benchmark map (rather than alphabetical filename order) resolves the +selected pair to its exact `.bddl` and `.pruned_init` files. For the example +above, results are organized as: + +```text +save_results/openpi_libero_tasks/ +└── libero_goal_task_3/ + ├── init_state_0.mp4 + ├── init_state_0.actions.npy + └── init_state_0.metrics.json +``` + +The metrics JSON records the task name and description, success result, exact +BDDL/init-state file paths, selected initialization-state ID, and rollout +statistics. Running a different initialization state of the same task places +another set of `init_state_N.*` files in the same task directory. + +Change the result root or cap the rollout during a smoke test with: + +```bash +OPENPI_LIBERO_RESULT_ROOT=/absolute/output/root \ +OPENPI_MAX_STEPS=20 \ +bash scripts/openpi/run_libero_task_i2va.sh libero_10 5 2 +``` + +Use `bash scripts/openpi/run_libero_task_i2va.sh --help` to display the command +summary. Other model/runtime overrides accepted by `run_libero_i2va.sh` remain +available, including `OPENPI_SEED`, `OPENPI_RENDER_SIZE`, and +`OPENPI_VIDEO_FPS`. + +## 5. Static image/state-to-action smoke inference + +Place the two RGB frames and state in one input directory: + +```text +INPUT_DIR/ +├── agentview_image.png +├── wrist_image.png +└── state.npy # float32 shape (8,) +``` + +A reproducible sample can be extracted from the already-downloaded raw +LIBERO dataset with: + +```bash +python scripts/openpi/prepare_libero_sample.py +``` + +Run: + +```bash +OPENPI_RUN_MODE=single_observation \ +OPENPI_IMAGE_PATH=/absolute/path/to/INPUT_DIR \ +OPENPI_TASK_DESCRIPTION="pick up the black bowl and place it on the plate" \ +OPENPI_SAVE_ACTION_PATH=/absolute/path/to/actions.npy \ +bash scripts/openpi/run_libero_i2va.sh +``` + +This mode also enters through `python -m lightx2v.infer`, resolves +`OpenPIRunner` through the shared registry, and then runs the synchronous local +worker under the private Transformers 4.53.2 dependency layer. There is no +OpenPI-specific `sys.argv` branch in `lightx2v.infer`. + +The result is a float32 NumPy array with shape `(10, 7)` at the exact +`OPENPI_SAVE_ACTION_PATH`. + +The default closed-loop rollout was exercised on an H200. The generated +artifacts are: + +```text +/data/liuhongda/LightX2V/save_results/output_openpi_pi05_libero.mp4 +/data/liuhongda/LightX2V/save_results/output_openpi_pi05_libero.actions.npy +/data/liuhongda/LightX2V/save_results/output_openpi_pi05_libero.metrics.json +/data/liuhongda/openpi_data/results/pi05_libero_pytorch_parity.json +``` + +The parity report uses identical LIBERO input and fixed flow noise for upstream +OpenPI PyTorch and the LightX2V-native network. + +## 6. Closed-loop ROS LIBERO rollout + +```bash +LIBERO_BENCHMARK=libero_10 \ +LIBERO_TASK_ID=5 \ +LIBERO_INIT_STATE_ID=0 \ +bash scripts/openpi/run_libero_ros_i2va.sh +``` + +The ROS script builds the existing LightX2V workspace, starts the shared local +LIBERO simulator, and starts `openpi_node`. The simulator already performs the +official 180-degree image rotation; the runner and policy do not flip again. diff --git a/scripts/openpi/convert_pi05_libero_to_pytorch.sh b/scripts/openpi/convert_pi05_libero_to_pytorch.sh new file mode 100755 index 000000000..8c77fda34 --- /dev/null +++ b/scripts/openpi/convert_pi05_libero_to_pytorch.sh @@ -0,0 +1,137 @@ +#!/usr/bin/env bash + +set -euo pipefail + +openpi_path="${OPENPI_PATH:-/data/liuhongda/openpi}" +source_checkpoint="${OPENPI_JAX_CHECKPOINT:-/data/liuhongda/openpi_data/openpi-assets/checkpoints/pi05_libero}" +output_checkpoint="${OPENPI_PYTORCH_CHECKPOINT:-/data/liuhongda/openpi_data/openpi-assets/checkpoints/pi05_libero_pytorch}" +tokenizer_path="${OPENPI_TOKENIZER_PATH:-/data/liuhongda/openpi_data/big_vision/paligemma_tokenizer.model}" +python_bin="${OPENPI_CONVERT_PYTHON:-${openpi_path}/.venv/bin/python}" +transformers_dir="${openpi_path}/.venv/lib/python3.11/site-packages/transformers" +patch_dir="${openpi_path}/src/openpi/models_pytorch/transformers_replace" + +for required in \ + "${python_bin}" \ + "${openpi_path}/examples/convert_jax_model_to_pytorch.py" \ + "${source_checkpoint}/params/_METADATA" \ + "${source_checkpoint}/assets/physical-intelligence/libero/norm_stats.json" \ + "${tokenizer_path}"; do + if [[ ! -e "${required}" ]]; then + echo "Required conversion input is missing: ${required}" >&2 + exit 1 + fi +done + +if [[ -f "${output_checkpoint}/model.safetensors" && "${OPENPI_FORCE_CONVERT:-0}" != "1" ]]; then + if [[ -f "${output_checkpoint}/SHA256SUMS" ]]; then + ( + cd "${output_checkpoint}" + sha256sum -c SHA256SUMS + ) + else + echo "Existing checkpoint has no SHA256SUMS; validating its tensor manifest before creating one." >&2 + fi + echo "Converted checkpoint already exists: ${output_checkpoint}/model.safetensors" + echo "Set OPENPI_FORCE_CONVERT=1 only when you intentionally want to overwrite it." +else + # The official PyTorch implementation needs five OpenPI replacements on + # top of transformers==4.53.2. This changes only OpenPI's own .venv, not + # the user's base environment. + cp -a "${patch_dir}/." "${transformers_dir}/" + ( + cd "${openpi_path}" + "${python_bin}" -u examples/convert_jax_model_to_pytorch.py \ + --checkpoint-dir "${source_checkpoint}" \ + --config-name pi05_libero \ + --output-path "${output_checkpoint}" \ + --precision bfloat16 + ) +fi + +# Upstream's converter looks for checkpoint_dir.parent/assets, while the +# released local layout stores assets under checkpoint_dir/assets. +mkdir -p "${output_checkpoint}/assets" +cp -a "${source_checkpoint}/assets/." "${output_checkpoint}/assets/" +cp -a "${tokenizer_path}" "${output_checkpoint}/assets/paligemma_tokenizer.model" + +"${python_bin}" - "${output_checkpoint}" <<'PY' +import hashlib +import json +import math +from pathlib import Path +import sys + +from safetensors import safe_open + +root = Path(sys.argv[1]) +required = ( + root / "model.safetensors", + root / "config.json", + root / "assets/physical-intelligence/libero/norm_stats.json", + root / "assets/paligemma_tokenizer.model", +) +missing = [str(path) for path in required if not path.is_file()] +if missing: + raise FileNotFoundError(f"Converted checkpoint is incomplete: {missing}") +with safe_open(root / "model.safetensors", framework="pt", device="cpu") as handle: + keys = list(handle.keys()) + manifest_rows = [] + parameter_count = 0 + dtypes = set() + for key in keys: + tensor_slice = handle.get_slice(key) + shape = tuple(tensor_slice.get_shape()) + dtype = tensor_slice.get_dtype() + manifest_rows.append(f"{key}|{dtype}|{shape}") + parameter_count += math.prod(shape) + dtypes.add(dtype) + +manifest_sha256 = hashlib.sha256(("\n".join(manifest_rows) + "\n").encode()).hexdigest() +expected_manifest_sha256 = "ee81d609ff73d395731f9f3df3b0caefcbd17f83d2ff153d26166e1bd024e20d" +if len(keys) != 812 or parameter_count != 3_616_757_520 or dtypes != {"BF16"}: + raise RuntimeError( + "Converted pi05_libero tensor inventory mismatch: " + f"keys={len(keys)}, parameters={parameter_count}, dtypes={sorted(dtypes)}" + ) +if manifest_sha256 != expected_manifest_sha256: + raise RuntimeError( + "Converted pi05_libero key/shape/dtype manifest mismatch: " + f"expected {expected_manifest_sha256}, got {manifest_sha256}" + ) + +required_prefixes = ( + "paligemma_with_expert.paligemma.", + "paligemma_with_expert.gemma_expert.", + "action_in_proj.", + "action_out_proj.", + "time_mlp_in.", + "time_mlp_out.", +) +missing_prefixes = [prefix for prefix in required_prefixes if not any(key.startswith(prefix) for key in keys)] +if missing_prefixes: + raise RuntimeError(f"SafeTensors is missing OpenPI parameter groups: {missing_prefixes}") +print( + json.dumps( + { + "checkpoint": str(root), + "tensor_keys": len(keys), + "parameters": parameter_count, + "tensor_manifest_sha256": manifest_sha256, + "bytes": (root / "model.safetensors").stat().st_size, + }, + indent=2, + ) +) +PY + +( + cd "${output_checkpoint}" + sha256sum \ + model.safetensors \ + config.json \ + assets/paligemma_tokenizer.model \ + assets/physical-intelligence/libero/norm_stats.json \ + > SHA256SUMS +) + +echo "pi05_libero PyTorch checkpoint is ready: ${output_checkpoint}" diff --git a/scripts/openpi/prepare_libero_sample.py b/scripts/openpi/prepare_libero_sample.py new file mode 100644 index 000000000..d58b947a0 --- /dev/null +++ b/scripts/openpi/prepare_libero_sample.py @@ -0,0 +1,63 @@ +"""Extract one reproducible offline OpenPI input from a downloaded LIBERO HDF5. + +This is a data-preparation utility, not part of model runtime. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import h5py +import numpy as np +from PIL import Image + +DEFAULT_SOURCE = Path( + "/data/liuhongda/openpi_data/raw/huggingface/yifengzhu-hf/LIBERO-datasets/libero_spatial/pick_up_the_black_bowl_between_the_plate_and_the_ramekin_and_place_it_on_the_plate_demo.hdf5" +) +DEFAULT_OUTPUT = Path("/data/liuhongda/openpi_data/examples/pi05_libero/libero_spatial_task0_demo0_step0") +DEFAULT_PROMPT = "pick up the black bowl between the plate and the ramekin and place it on the plate" + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--source", type=Path, default=DEFAULT_SOURCE) + parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT) + parser.add_argument("--demo", default="demo_0") + parser.add_argument("--step", type=int, default=0) + args = parser.parse_args() + + source = args.source.expanduser().resolve() + output = args.output.expanduser().resolve() + if not source.is_file(): + raise FileNotFoundError(source) + + with h5py.File(source, "r") as handle: + demo = handle[f"data/{args.demo}"] + step = int(args.step) + agentview = np.asarray(demo["obs/agentview_rgb"][step], dtype=np.uint8) + wrist = np.asarray(demo["obs/eye_in_hand_rgb"][step], dtype=np.uint8) + state = np.concatenate([demo["obs/ee_pos"][step], demo["obs/ee_ori"][step], demo["obs/gripper_states"][step]]).astype(np.float32) + reference_actions = np.asarray(demo["actions"][step : step + 10], dtype=np.float32) + + if state.shape != (8,): + raise ValueError(f"Expected 8-D LIBERO state, got {state.shape}") + output.mkdir(parents=True, exist_ok=True) + Image.fromarray(agentview, mode="RGB").save(output / "agentview_image.png") + Image.fromarray(wrist, mode="RGB").save(output / "wrist_image.png") + np.save(output / "state.npy", state) + np.save(output / "reference_actions.npy", reference_actions) + metadata = { + "source": str(source), + "demo": args.demo, + "step": int(args.step), + "task_description": DEFAULT_PROMPT, + "state_layout": ["eef_pos[3]", "eef_axis_angle[3]", "gripper_qpos[2]"], + } + (output / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n", encoding="utf-8") + print(output) + + +if __name__ == "__main__": + main() diff --git a/scripts/openpi/run_libero_i2va.sh b/scripts/openpi/run_libero_i2va.sh new file mode 100755 index 000000000..a7c9b0613 --- /dev/null +++ b/scripts/openpi/run_libero_i2va.sh @@ -0,0 +1,83 @@ +#!/usr/bin/env bash + +set -euo pipefail + +lightx2v_path="${LIGHTX2V_PATH:-/data/liuhongda/LightX2V}" +model_path="${OPENPI_MODEL_PATH:-/data/liuhongda/openpi_data/openpi-assets/checkpoints/pi05_libero_pytorch}" +config_json="${OPENPI_CONFIG:-${lightx2v_path}/configs/openpi/pi05_libero.json}" +transformers_runtime_path="${OPENPI_TRANSFORMERS_RUNTIME_PATH:-/data/liuhongda/openpi_data/python_deps/openpi_pytorch_runtime}" +python_bin="${OPENPI_PYTHON:-/opt/conda/bin/python}" +run_mode="${OPENPI_RUN_MODE:-rollout}" + +image_path="${OPENPI_IMAGE_PATH:-/data/liuhongda/openpi_data/examples/pi05_libero/libero_spatial_task0_demo0_step0}" +state_path="${OPENPI_STATE_PATH:-${image_path}/state.npy}" +task_description="${OPENPI_TASK_DESCRIPTION:-}" +save_action_path="${OPENPI_SAVE_ACTION_PATH:-${lightx2v_path}/save_results/output_openpi_pi05_libero.actions.npy}" +save_video_path="${OPENPI_SAVE_VIDEO_PATH:-${OPENPI_SAVE_RESULT_PATH:-${lightx2v_path}/save_results/output_openpi_pi05_libero.mp4}}" +save_metrics_path="${OPENPI_SAVE_METRICS_PATH:-${lightx2v_path}/save_results/output_openpi_pi05_libero.metrics.json}" + +libero_root="${OPENPI_LIBERO_ROOT:-/data/liuhongda/openpi/third_party/libero}" +libero_config_dir="${OPENPI_LIBERO_CONFIG_DIR:-/data/liuhongda/openpi_data/runtime_configs/lightx2v_openpi_libero}" + +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}" +export PYTHONPATH="${PYTHONPATH:-}" +export USE_FLAX=0 +export MUJOCO_GL="${MUJOCO_GL:-egl}" +export PYOPENGL_PLATFORM="${PYOPENGL_PLATFORM:-egl}" +export MUJOCO_EGL_DEVICE_ID="${OPENPI_EGL_DEVICE_ID:-0}" +export NUMBA_CACHE_DIR="${OPENPI_NUMBA_CACHE_DIR:-${NUMBA_CACHE_DIR:-/tmp/lightx2v-openpi-numba-cache}}" +export MPLCONFIGDIR="${OPENPI_MPLCONFIG_DIR:-${MPLCONFIGDIR:-/tmp/lightx2v-openpi-matplotlib-cache}}" +mkdir -p "${NUMBA_CACHE_DIR}" "${MPLCONFIGDIR}" + +for required in "${python_bin}" "${model_path}/model.safetensors" "${config_json}" "${transformers_runtime_path}/transformers"; do + if [[ ! -e "${required}" ]]; then + echo "Required OpenPI runtime input is missing: ${required}" >&2 + exit 1 + fi +done + +source "${lightx2v_path}/scripts/base/base.sh" + +# The common LightX2V process stays on the base environment. OpenPIRunner adds +# the isolated patched Transformers runtime only to its local worker process. +export PYTHONPATH="${lightx2v_path}${PYTHONPATH:+:${PYTHONPATH}}" +export OPENPI_PYTHON="${python_bin}" +export OPENPI_RUN_MODE="${run_mode}" +export OPENPI_TRANSFORMERS_RUNTIME_PATH="${transformers_runtime_path}" +export OPENPI_LIBERO_ROOT="${libero_root}" +export OPENPI_LIBERO_CONFIG_DIR="${libero_config_dir}" +export OPENPI_SAVE_METRICS_PATH="${save_metrics_path}" +unset OPENPI_WORKER_PROCESS + +case "${run_mode}" in + rollout) + seed="${OPENPI_SEED:-7}" + if [[ ! -d "${libero_root}/libero/libero/bddl_files" ]]; then + echo "LIBERO checkout is incomplete: ${libero_root}" >&2 + exit 1 + fi + ;; + single_observation) + seed="${OPENPI_SEED:-0}" + if [[ -z "${task_description}" ]]; then + task_description="pick up the black bowl between the plate and the ramekin and place it on the plate" + fi + ;; + *) + echo "Unsupported OPENPI_RUN_MODE=${run_mode}; expected rollout or single_observation." >&2 + exit 2 + ;; +esac + +echo "OpenPI call chain: lightx2v.infer -> OpenPIRunner -> local ${run_mode} worker" +"${python_bin}" -m lightx2v.infer \ + --model_cls openpi \ + --task i2va \ + --model_path "${model_path}" \ + --config_json "${config_json}" \ + --seed "${seed}" \ + --prompt "${task_description}" \ + --image_path "${image_path}" \ + --state_path "${state_path}" \ + --save_result_path "${save_video_path}" \ + --save_action_path "${save_action_path}" diff --git a/scripts/openpi/run_libero_ros_i2va.sh b/scripts/openpi/run_libero_ros_i2va.sh new file mode 100755 index 000000000..d2758206e --- /dev/null +++ b/scripts/openpi/run_libero_ros_i2va.sh @@ -0,0 +1,66 @@ +#!/usr/bin/env bash + +set -euo pipefail + +lightx2v_path="${LIGHTX2V_PATH:-/data/liuhongda/LightX2V}" +ros_workspace="${lightx2v_path}/lightx2v_ros" +model_path="${OPENPI_LIBERO_MODEL_PATH:-/data/liuhongda/openpi_data/openpi-assets/checkpoints/pi05_libero_pytorch}" +config_json="${OPENPI_LIBERO_CONFIG:-${lightx2v_path}/configs/openpi/pi05_libero.json}" +transformers_runtime_path="${OPENPI_TRANSFORMERS_RUNTIME_PATH:-/data/liuhongda/openpi_data/python_deps/openpi_pytorch_runtime}" +lyrical_setup="${HOME:-}/ros2_lyrical/install/setup.sh" +if [[ -n "${ROS_SETUP:-}" ]]; then + ros_setup="${ROS_SETUP}" +elif [[ -f "${lyrical_setup}" ]]; then + ros_setup="${lyrical_setup}" +else + ros_setup="/opt/ros/jazzy/setup.bash" +fi + +if [[ ! -f "${ros_setup}" ]]; then + echo "ROS setup not found: ${ros_setup}. Set ROS_SETUP to your ROS2 setup script." >&2 + exit 1 +fi + +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}" +export PYTHONPATH="${PYTHONPATH:-}" +export USE_FLAX=0 + +set +u +source "${ros_setup}" +set -u +source "${lightx2v_path}/scripts/base/base.sh" + +# Keep OpenPI's patched Transformers 4.53.2 isolated from the base +# installation while making it visible to the inference ROS process. +export PYTHONPATH="${transformers_runtime_path}:${lightx2v_path}${PYTHONPATH:+:${PYTHONPATH}}" + +cd "${ros_workspace}" +colcon build --symlink-install --packages-select common simulator inference +set +u +source "${ros_workspace}/install/setup.bash" +set -u + +simulator_pid="" +cleanup() { + if [[ -n "${simulator_pid}" ]]; then + kill "${simulator_pid}" 2>/dev/null || true + wait "${simulator_pid}" 2>/dev/null || true + fi +} +trap cleanup EXIT INT TERM + +ros2 run simulator libero_node --ros-args \ + -p autostart:=true \ + -p "benchmark:=${LIBERO_BENCHMARK:-libero_10}" \ + -p "task_id:=${LIBERO_TASK_ID:-5}" \ + -p "init_state_id:=${LIBERO_INIT_STATE_ID:-0}" \ + -p "seed:=${LIBERO_SEED:-0}" & +simulator_pid=$! + +ros2 run inference openpi_node --ros-args \ + -p env:=libero \ + -p "model_path:=${model_path}" \ + -p "config_json:=${config_json}" \ + -p "seed:=${OPENPI_SEED:-7}" \ + -p "actions_per_plan:=${OPENPI_ACTIONS_PER_PLAN:-5}" \ + -p "num_steps_wait:=${OPENPI_NUM_STEPS_WAIT:-10}" diff --git a/scripts/openpi/run_libero_task_i2va.sh b/scripts/openpi/run_libero_task_i2va.sh new file mode 100755 index 000000000..ee29693ce --- /dev/null +++ b/scripts/openpi/run_libero_task_i2va.sh @@ -0,0 +1,113 @@ +#!/usr/bin/env bash + +set -euo pipefail + +usage() { + cat <<'EOF' +Run one task from the local LIBERO BDDL and initialization-state samples. + +Usage: + bash scripts/openpi/run_libero_task_i2va.sh [init_state_id] + +Benchmarks and zero-based task-id ranges: + libero_spatial 0-9 + libero_object 0-9 + libero_goal 0-9 + libero_10 0-9 + libero_90 0-89 + +Example: + bash scripts/openpi/run_libero_task_i2va.sh libero_goal 3 0 + +Each task's pruned init file contains 50 states; init_state_id is 0-49. + +Optional environment variables: + OPENPI_LIBERO_RESULT_ROOT Result root directory + OPENPI_MAX_STEPS Override the suite-specific rollout step limit + OPENPI_SEED Policy and simulator seed (default: 7) + OPENPI_VIDEO_FPS Output video FPS (default: 10) + OPENPI_RENDER_SIZE Square render resolution (default: 256) +EOF +} + +if [[ "${1:-}" == "-h" || "${1:-}" == "--help" ]]; then + usage + exit 0 +fi + +if [[ $# -lt 2 || $# -gt 3 ]]; then + usage >&2 + exit 2 +fi + +benchmark="$1" +task_id="$2" +init_state_id="${3:-0}" + +case "${benchmark}" in + libero_spatial|libero_object|libero_goal|libero_10) + task_count=10 + ;; + libero_90) + task_count=90 + ;; + *) + echo "Unsupported LIBERO benchmark: ${benchmark}" >&2 + usage >&2 + exit 2 + ;; +esac + +if [[ ! "${task_id}" =~ ^[0-9]+$ ]]; then + echo "task_id must be a non-negative integer: ${task_id}" >&2 + exit 2 +fi +if [[ ! "${init_state_id}" =~ ^[0-9]+$ ]]; then + echo "init_state_id must be a non-negative integer: ${init_state_id}" >&2 + exit 2 +fi + +task_id=$((10#${task_id})) +init_state_id=$((10#${init_state_id})) +if (( task_id >= task_count )); then + echo "task_id ${task_id} is out of range for ${benchmark}; expected 0-$((task_count - 1))." >&2 + exit 2 +fi +if (( init_state_id >= 50 )); then + echo "init_state_id ${init_state_id} is out of range; expected 0-49." >&2 + exit 2 +fi + +lightx2v_path="${LIGHTX2V_PATH:-/data/liuhongda/LightX2V}" +libero_root="${OPENPI_LIBERO_ROOT:-/data/liuhongda/openpi/third_party/libero}" +bddl_dir="${libero_root}/libero/libero/bddl_files/${benchmark}" +init_dir="${libero_root}/libero/libero/init_files/${benchmark}" + +for required_dir in "${bddl_dir}" "${init_dir}"; do + if [[ ! -d "${required_dir}" ]]; then + echo "Required LIBERO sample directory is missing: ${required_dir}" >&2 + exit 1 + fi +done + +result_root="${OPENPI_LIBERO_RESULT_ROOT:-${lightx2v_path}/save_results/openpi_libero_tasks}" +task_output_dir="${result_root}/${benchmark}_task_${task_id}" +output_prefix="init_state_${init_state_id}" +mkdir -p "${task_output_dir}" + +echo "LIBERO benchmark : ${benchmark}" +echo "Task ID : ${task_id} (zero-based)" +echo "Init-state ID : ${init_state_id} (zero-based)" +echo "BDDL directory : ${bddl_dir}" +echo "Init directory : ${init_dir}" +echo "Output directory : ${task_output_dir}" + +LIBERO_BENCHMARK="${benchmark}" \ +LIBERO_TASK_ID="${task_id}" \ +LIBERO_INIT_STATE_ID="${init_state_id}" \ +OPENPI_SAVE_VIDEO_PATH="${task_output_dir}/${output_prefix}.mp4" \ +OPENPI_SAVE_ACTION_PATH="${task_output_dir}/${output_prefix}.actions.npy" \ +OPENPI_SAVE_METRICS_PATH="${task_output_dir}/${output_prefix}.metrics.json" \ +bash "${lightx2v_path}/scripts/openpi/run_libero_i2va.sh" + +echo "Saved LIBERO task results under: ${task_output_dir}" diff --git a/scripts/openpi/setup_pytorch_runtime.sh b/scripts/openpi/setup_pytorch_runtime.sh new file mode 100755 index 000000000..a9b285702 --- /dev/null +++ b/scripts/openpi/setup_pytorch_runtime.sh @@ -0,0 +1,201 @@ +#!/usr/bin/env bash + +set -euo pipefail + +script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +lightx2v_path="${LIGHTX2V_PATH:-$(cd -- "${script_dir}/../.." && pwd)}" +openpi_path="${OPENPI_PATH:-/data/liuhongda/openpi}" +runtime_path="${OPENPI_PYTORCH_RUNTIME_PATH:-/data/liuhongda/openpi_data/python_deps/openpi_pytorch_runtime}" +python_bin="${OPENPI_RUNTIME_PYTHON:-/opt/conda/bin/python}" +openpi_venv_python="${OPENPI_VENV_PYTHON:-${openpi_path}/.venv/bin/python}" +patch_source="${lightx2v_path}/lightx2v/models/networks/openpi/transformers_replace" + +transformers_version="4.53.2" +huggingface_hub_version="0.32.3" +tokenizers_version="0.21.1" + +if [[ ! -x "${python_bin}" ]]; then + echo "Python executable not found: ${python_bin}" >&2 + exit 1 +fi + +runtime_path="$(${python_bin} -c 'import os, sys; print(os.path.abspath(os.path.expanduser(sys.argv[1])))' "${runtime_path}")" +case "${runtime_path}" in + /|/data|/data/liuhongda|/data/liuhongda/openpi_data|/data/liuhongda/openpi_data/python_deps) + echo "Refusing to use an unsafe runtime target: ${runtime_path}" >&2 + exit 1 + ;; +esac + +patch_files=( + "models/gemma/configuration_gemma.py" + "models/gemma/modeling_gemma.py" + "models/paligemma/modeling_paligemma.py" + "models/siglip/check.py" + "models/siglip/modeling_siglip.py" +) + +for relative_path in "${patch_files[@]}"; do + if [[ ! -f "${patch_source}/${relative_path}" ]]; then + echo "Required OpenPI Transformers patch is missing: ${patch_source}/${relative_path}" >&2 + exit 1 + fi +done + +overlay_openpi_patches() { + local target_root="$1" + if [[ ! -d "${target_root}/transformers" ]]; then + echo "Transformers package is missing under ${target_root}" >&2 + return 1 + fi + cp -a "${patch_source}/." "${target_root}/transformers/" +} + +runtime_has_exact_layout() { + local target_root="$1" + [[ -d "${target_root}/transformers" ]] \ + && [[ -d "${target_root}/transformers-${transformers_version}.dist-info" ]] \ + && [[ -d "${target_root}/huggingface_hub" ]] \ + && [[ -d "${target_root}/huggingface_hub-${huggingface_hub_version}.dist-info" ]] \ + && [[ -d "${target_root}/tokenizers" ]] \ + && [[ -d "${target_root}/tokenizers-${tokenizers_version}.dist-info" ]] +} + +guard_runtime() { + local target_root="$1" + USE_FLAX=0 PYTHONPATH="${target_root}${PYTHONPATH:+:${PYTHONPATH}}" \ + "${python_bin}" - "${target_root}" "${patch_source}" <<'PY' +import filecmp +import importlib.metadata +import inspect +from pathlib import Path +import sys + +runtime_root = Path(sys.argv[1]).resolve() +patch_root = Path(sys.argv[2]).resolve() + +expected_versions = { + "transformers": "4.53.2", + "huggingface-hub": "0.32.3", + "tokenizers": "0.21.1", +} +for distribution, expected in expected_versions.items(): + actual = importlib.metadata.version(distribution) + if actual != expected: + raise RuntimeError(f"{distribution} version mismatch: expected {expected}, got {actual}") + +import huggingface_hub +import tokenizers +import transformers + +for module in (transformers, huggingface_hub, tokenizers): + module_path = Path(module.__file__).resolve() + if not module_path.is_relative_to(runtime_root): + raise RuntimeError(f"{module.__name__} was imported outside the private runtime: {module_path}") + +patch_files = ( + "models/gemma/configuration_gemma.py", + "models/gemma/modeling_gemma.py", + "models/paligemma/modeling_paligemma.py", + "models/siglip/check.py", + "models/siglip/modeling_siglip.py", +) +for relative_path in patch_files: + source = patch_root / relative_path + installed = runtime_root / "transformers" / relative_path + if not installed.is_file() or not filecmp.cmp(source, installed, shallow=False): + raise RuntimeError(f"OpenPI Transformers patch mismatch: {installed}") + +from transformers.models.gemma.configuration_gemma import GemmaConfig +from transformers.models.gemma.modeling_gemma import GemmaRMSNorm, _gated_residual +from transformers.models.paligemma import modeling_paligemma # noqa: F401 +from transformers.models.siglip import check, modeling_siglip # noqa: F401 + +if not check.check_whether_transformers_replace_is_installed_correctly(): + raise RuntimeError("OpenPI Transformers replacement guard returned false") + +config = GemmaConfig(use_adarms=True, adarms_cond_dim=8) +if not config.use_adarms or config.adarms_cond_dim != 8: + raise RuntimeError("Patched GemmaConfig does not expose AdaRMS configuration") +if "cond" not in inspect.signature(GemmaRMSNorm.forward).parameters: + raise RuntimeError("Patched GemmaRMSNorm does not accept AdaRMS conditioning") +if not callable(_gated_residual): + raise RuntimeError("Patched Gemma gated residual helper is unavailable") + +print(f"private transformers={transformers.__version__} ({transformers.__file__})") +print(f"private huggingface-hub={huggingface_hub.__version__} ({huggingface_hub.__file__})") +print(f"private tokenizers={tokenizers.__version__} ({tokenizers.__file__})") +print("OpenPI Transformers replacement guard: OK") +PY +} + +if runtime_has_exact_layout "${runtime_path}"; then + echo "Refreshing OpenPI patches in existing private runtime: ${runtime_path}" + overlay_openpi_patches "${runtime_path}" + if guard_runtime "${runtime_path}"; then + echo "OpenPI private PyTorch runtime is ready: ${runtime_path}" + exit 0 + fi + echo "Existing private runtime failed validation; rebuilding it." >&2 +fi + +runtime_parent="$(dirname -- "${runtime_path}")" +mkdir -p "${runtime_parent}" +stage_dir="$(mktemp -d "${runtime_parent}/.openpi_pytorch_runtime.XXXXXX")" + +cleanup() { + if [[ -n "${stage_dir:-}" && -d "${stage_dir}" ]]; then + rm -rf -- "${stage_dir}" + fi +} +trap cleanup EXIT + +copied_from_openpi_venv=false +if [[ -x "${openpi_venv_python}" ]]; then + source_site="$(${openpi_venv_python} -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')" + if [[ -d "${source_site}/transformers" \ + && -d "${source_site}/transformers-${transformers_version}.dist-info" \ + && -d "${source_site}/huggingface_hub" \ + && -d "${source_site}/huggingface_hub-${huggingface_hub_version}.dist-info" \ + && -d "${source_site}/tokenizers" \ + && -d "${source_site}/tokenizers-${tokenizers_version}.dist-info" ]]; then + echo "Building private runtime from local OpenPI environment: ${source_site}" + cp -a \ + "${source_site}/transformers" \ + "${source_site}/transformers-${transformers_version}.dist-info" \ + "${source_site}/huggingface_hub" \ + "${source_site}/huggingface_hub-${huggingface_hub_version}.dist-info" \ + "${source_site}/tokenizers" \ + "${source_site}/tokenizers-${tokenizers_version}.dist-info" \ + "${stage_dir}/" + copied_from_openpi_venv=true + fi +fi + +if [[ "${copied_from_openpi_venv}" != true ]]; then + echo "Exact local packages were not found; using pip --target fallback." + "${python_bin}" -m pip install \ + --disable-pip-version-check \ + --no-input \ + --no-deps \ + --target "${stage_dir}" \ + "transformers==${transformers_version}" \ + "huggingface-hub==${huggingface_hub_version}" \ + "tokenizers==${tokenizers_version}" +fi + +overlay_openpi_patches "${stage_dir}" +guard_runtime "${stage_dir}" + +if [[ -e "${runtime_path}" || -L "${runtime_path}" ]]; then + backup_path="${runtime_path}.invalid.$(date -u +%Y%m%dT%H%M%SZ).$$" + mv -- "${runtime_path}" "${backup_path}" + echo "Previous invalid runtime was preserved at: ${backup_path}" +fi +mv -- "${stage_dir}" "${runtime_path}" +stage_dir="" + +guard_runtime "${runtime_path}" +echo "OpenPI private PyTorch runtime is ready: ${runtime_path}" +echo "Use it without changing base packages:" +echo " PYTHONPATH=${runtime_path}\${PYTHONPATH:+:\$PYTHONPATH} ${python_bin} " diff --git a/scripts/openpi/validate_pytorch_parity.py b/scripts/openpi/validate_pytorch_parity.py new file mode 100644 index 000000000..5c1f1667b --- /dev/null +++ b/scripts/openpi/validate_pytorch_parity.py @@ -0,0 +1,124 @@ +"""Compare upstream OpenPI PyTorch and LightX2V with identical input/noise. + +Run this with OpenPI's patched conversion environment. It is a validation tool; +the deployed LightX2V runtime itself does not import OpenPI or JAX. +""" + +from __future__ import annotations + +import argparse +import dataclasses +import gc +import json +import sys +import types +from pathlib import Path + +import h5py +import numpy as np +import torch + + +def load_sample(path: Path) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + with h5py.File(path, "r") as handle: + demo = handle["data/demo_0"] + image = np.asarray(demo["obs/agentview_rgb"][0], dtype=np.uint8) + wrist = np.asarray(demo["obs/eye_in_hand_rgb"][0], dtype=np.uint8) + state = np.concatenate([demo["obs/ee_pos"][0], demo["obs/ee_ori"][0], demo["obs/gripper_states"][0]]).astype(np.float32) + return image, wrist, state + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--checkpoint", + type=Path, + default=Path("/data/liuhongda/openpi_data/openpi-assets/checkpoints/pi05_libero_pytorch"), + ) + parser.add_argument( + "--config", + type=Path, + default=Path("/data/liuhongda/LightX2V/configs/openpi/pi05_libero.json"), + ) + parser.add_argument( + "--sample", + type=Path, + default=Path( + "/data/liuhongda/openpi_data/raw/huggingface/yifengzhu-hf/LIBERO-datasets/libero_spatial/pick_up_the_black_bowl_between_the_plate_and_the_ramekin_and_place_it_on_the_plate_demo.hdf5" + ), + ) + parser.add_argument( + "--output", + type=Path, + default=Path("/data/liuhongda/openpi_data/results/pi05_libero_pytorch_parity.json"), + ) + parser.add_argument("--device", default="cuda") + args = parser.parse_args() + + image, wrist, state = load_sample(args.sample) + prompt = "pick up the black bowl between the plate and the ramekin and place it on the plate" + noise = np.random.default_rng(0).standard_normal((10, 32)).astype(np.float32) + + from openpi.policies import policy_config + from openpi.training import config as training_config + + upstream_config = training_config.get_config("pi05_libero") + upstream_config = dataclasses.replace( + upstream_config, + model=dataclasses.replace(upstream_config.model, pytorch_compile_mode=None), + ) + upstream_policy = policy_config.create_trained_policy( + upstream_config, + args.checkpoint, + pytorch_device=args.device, + ) + raw_input = { + "observation/image": image, + "observation/wrist_image": wrist, + "observation/state": state, + "prompt": prompt, + } + upstream_actions = np.asarray(upstream_policy.infer(raw_input, noise=noise)["actions"]) + del upstream_policy + gc.collect() + if args.device.startswith("cuda"): + torch.cuda.empty_cache() + + # Import only the LightX2V network family; bypass LightX2V's top-level + # platform initialization because this tool already owns the torch device. + package = types.ModuleType("lightx2v") + package.__path__ = ["/data/liuhongda/LightX2V/lightx2v"] + sys.modules["lightx2v"] = package + from lightx2v.models.networks.openpi import OpenPIModel + + with args.config.open("r", encoding="utf-8") as handle: + local_config = json.load(handle) + local_config["device"] = args.device + model = OpenPIModel.from_config(local_config) + normalized = model.predict_normalized_action_chunk( + {"agentview": image, "wrist": wrist}, + state, + prompt, + noise=noise, + ) + local_actions = model.post_infer.infer(normalized) + + difference = np.asarray(local_actions, dtype=np.float64) - np.asarray(upstream_actions, dtype=np.float64) + report = { + "upstream_shape": list(upstream_actions.shape), + "lightx2v_shape": list(local_actions.shape), + "max_abs_error": float(np.max(np.abs(difference))), + "mean_abs_error": float(np.mean(np.abs(difference))), + "allclose_atol_1e-6": bool(np.allclose(local_actions, upstream_actions, rtol=0.0, atol=1e-6)), + "upstream_actions": upstream_actions.tolist(), + "lightx2v_actions": local_actions.tolist(), + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8") + print(json.dumps({key: value for key, value in report.items() if not key.endswith("actions")}, indent=2)) + if not report["allclose_atol_1e-6"]: + raise SystemExit("OpenPI PyTorch parity check failed") + + +if __name__ == "__main__": + main()