Delete modeling_prismatic.py.back.20260315_133310

modeling_prismatic.py.back.20260315_133310

@@ -1,1498 +0,0 @@
-"""
-modeling_prismatic.py
-
-Core HuggingFace-style PrismaticPreTrainedModel and PrismaticForConditionalGeneration class definitions.
-Inherits from the default `transformers.PretrainedModel`. Meant to be standalone and self-contained,
-but exactly replicate the logic in `prismatic.models.vlms.prismatic.py`.
-"""
-
-import logging
-from dataclasses import dataclass
-from functools import partial
-import math
-import random
-from typing import Any, Callable, ClassVar, Dict, List, Optional, Tuple, Union
-import numpy as np
-import timm
-import tokenizers
-import torch
-import torch.nn as nn
-import transformers
-from timm.models.vision_transformer import LayerScale
-from transformers import AutoModelForCausalLM, PretrainedConfig, PreTrainedModel
-from transformers.modeling_outputs import ModelOutput
-
-from prismatic.training.train_utils import (
-    get_current_action_mask,
-    get_next_actions_mask,
-)
-from prismatic.vla.constants import (
-    ACTION_DIM,
-    ACTION_PROPRIO_NORMALIZATION_TYPE,
-    ACTION_TOKEN_BEGIN_IDX,
-    IGNORE_INDEX,
-    NUM_ACTIONS_CHUNK,
-    STOP_INDEX,
-    NormalizationType,
-    NUM_TOKENS
-)
-from .configuration_prismatic import OpenVLAConfig, PrismaticConfig
-
-
-
-# Set up logger
-logger = logging.getLogger(__name__)
-
-
-# === Utility Functions for Monkey-Patching ===
-def unpack_tuple(fn: Callable[[Any], Tuple[Any]]) -> Callable[[Any], Any]:
-    def wrapper(*args: Any, **kwargs: Any) -> Any:
-        result = fn(*args, **kwargs)
-        return result[0] if isinstance(result, tuple) else result
-
-    return wrapper
-
-
-
-# HF Transformers overwrites parameters with names containing `gamma`; we're going to patch VisionBackbone.LayerScale.
-#   =>> TIMM :: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L109
-#   =>> Transformers :: https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L3960
-def _ls_new_forward(self, x: torch.Tensor) -> torch.Tensor:
-    return x.mul_(self.scale_factor) if self.inplace else x * self.scale_factor
-
-
-
-def ls_apply_patch(ls_module: LayerScale):
-    ls_module.scale_factor = nn.Parameter(ls_module.gamma.clone())
-    ls_module.forward = _ls_new_forward.__get__(ls_module, LayerScale)
-    del ls_module.gamma
-
-
-
-# === Prismatic Vision Backbone (nn.Module) Definitions (w/ Fused Backbone Support) ===
-class PrismaticVisionBackbone(nn.Module):
-    """
-    Vision backbone for Prismatic models that handles image feature extraction.
-
-    Supports both single backbone (e.g., SigLIP) and fused backbone (e.g., SigLIP + DINOv2) configurations.
-    For fused backbones, features from both models are concatenated along the feature dimension.
-    """
-
-    def __init__(
-        self,
-        use_fused_vision_backbone: bool,
-        image_sizes: List[int],
-        timm_model_ids: List[str],
-        timm_override_act_layers: List[Optional[str]],
-    ) -> None:
-        """
-        Initialize the vision backbone.
-
-        Args:
-            use_fused_vision_backbone: Whether to use two backbones and fuse their features
-            image_sizes: List of image sizes for each backbone
-            timm_model_ids: List of TIMM model IDs to use for each backbone
-            timm_override_act_layers: List of activation layer overrides for each backbone
-        """
-        super().__init__()
-        self.use_fused_vision_backbone = use_fused_vision_backbone
-        self.num_images_in_input = 1  # Default value, can be overridden later
-
-        # Validate number of (fused) vision backbones
-        if len(timm_model_ids) > 2:
-            raise ValueError("Prismatic models only support up to 2 (fused) vision backbones!")
-
-        # Create primary featurizer
-        self.featurizer = self._create_featurizer(
-            model_id=timm_model_ids[0], img_size=image_sizes[0], act_layer=timm_override_act_layers[0]
-        )
-        self.embed_dim = self.featurizer.embed_dim
-
-        # Create secondary featurizer if using fused backbone
-        if self.use_fused_vision_backbone:
-            self.fused_featurizer = self._create_featurizer(
-                model_id=timm_model_ids[1], img_size=image_sizes[1], act_layer=timm_override_act_layers[1]
-            )
-            self.embed_dim += self.fused_featurizer.embed_dim
-
-        # Patch LayerScale modules for HF compatibility
-        self._patch_layer_scales()
-
-
-    def _create_featurizer(self, model_id: str, img_size: int, act_layer: Optional[str]) -> nn.Module:
-        """
-        Create a TIMM-based featurizer model with appropriate configurations.
-
-        Args:
-            model_id: The TIMM model ID to load
-            img_size: Input image size for the model
-            act_layer: Override for the activation layer type
-
-        Returns:
-            A configured featurizer model
-        """
-        featurizer = timm.create_model(
-            model_id,
-            pretrained=False,
-            num_classes=0,
-            img_size=img_size,
-            act_layer=act_layer,
-        )
-
-        # Monkey-patch the forward function to extract the second-to-last layer features
-        num_blocks = len(featurizer.blocks)
-        featurizer.forward = unpack_tuple(partial(featurizer.get_intermediate_layers, n={num_blocks - 2}))
-
-        return featurizer
-
-
-    def _patch_layer_scales(self) -> None:
-        """
-        Patch all LayerScale modules to be compatible with HF's parameter naming.
-
-        HF Transformers overwrites parameters with names containing 'gamma',
-        so we need to rename and modify the forward method.
-        """
-        # Patch primary featurizer
-        for module in self.featurizer.modules():
-            if isinstance(module, LayerScale):
-                ls_apply_patch(module)
-
-        # Patch secondary featurizer if it exists
-        if self.use_fused_vision_backbone:
-            for module in self.fused_featurizer.modules():
-                if isinstance(module, LayerScale):
-                    ls_apply_patch(module)
-
-
-    def get_num_patches(self) -> int:
-        """
-        Returns the number of vision patches output by the vision backbone.
-
-        Returns:
-            Number of patches per image
-        """
-        return self.featurizer.patch_embed.num_patches
-
-
-    def get_num_images_in_input(self) -> int:
-        """
-        Returns the number of input images for the vision backbone.
-
-        Returns:
-            Number of images expected in the input
-        """
-        return self.num_images_in_input
-
-
-    def set_num_images_in_input(self, num_images_in_input: int) -> None:
-        """
-        Sets the number of input images for the vision backbone.
-
-        Args:
-            num_images_in_input: Number of images to expect in the input
-        """
-        self.num_images_in_input = num_images_in_input
-
-
-    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
-        """
-        Implements the forward pass for the vision backbone.
-
-        If `self.use_fused_vision_backbone == True`, uses both SigLIP and DINOv2 transformers to extract visual features
-        (otherwise uses SigLIP only). Allows multi-image inputs (but only for fused vision backbone).
-
-        Args:
-            pixel_values (torch.Tensor): Pixels for input image(s), (B, C, H, W).
-        """
-        if self.num_images_in_input == 1:
-            if not self.use_fused_vision_backbone:
-                return self.featurizer(pixel_values)
-
-            # Split `pixel_values :: [bsz, 2 * 3, resolution, resolution]` =>> featurize =>> channel stack
-            img, img_fused = torch.split(pixel_values, [3, 3], dim=1)
-            patches, patches_fused = self.featurizer(img), self.fused_featurizer(img_fused)
-
-            return torch.cat([patches, patches_fused], dim=2)
-
-        else:
-            assert self.use_fused_vision_backbone, "Multi-image inputs require using fused backbone!"
-
-            # Split `pixel_values` into individual images (each with 6 channels: 3 for SigLIP + 3 for DINOv2)
-            images = torch.split(pixel_values, [6] * self.num_images_in_input, dim=1)
-
-            # Process each image and collect patches
-            all_patches = []
-            for img in images:
-                # Split each image further into two stacks of channels (each with 3 channels)
-                img_regular, img_fused = torch.split(img, [3, 3], dim=1)
-
-                # Get patches from both SigLIP and DINOv2 vision transformers
-                patches = self.featurizer(img_regular)
-                patches_fused = self.fused_featurizer(img_fused)
-
-                # Concatenate SigLIP and DINOv2 patches along the hidden dimension
-                combined_patches = torch.cat([patches, patches_fused], dim=2)
-                all_patches.append(combined_patches)
-
-            # Concatenate all patches along the patch dimension
-            return torch.cat(all_patches, dim=1)
-
-
-
-# === Prismatic Projector (nn.Module) Definitions ===
-class PrismaticProjector(nn.Module):
-    def __init__(self, use_fused_vision_backbone: bool, vision_dim: int, llm_dim: int) -> None:
-        super().__init__()
-        self.use_fused_vision_backbone = use_fused_vision_backbone
-        self.vision_dim, self.llm_dim = vision_dim, llm_dim
-
-        # Switch on `use_fused_vision_backbone` =>> use slightly different MLPs and projection factors!
-        if not self.use_fused_vision_backbone:
-            self.fc1 = nn.Linear(self.vision_dim, self.llm_dim, bias=True)
-            self.fc2 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
-            self.act_fn1 = nn.GELU()
-        else:
-            initial_projection_dim = 4 * vision_dim
-            self.fc1 = nn.Linear(self.vision_dim, initial_projection_dim, bias=True)
-            self.fc2 = nn.Linear(initial_projection_dim, self.llm_dim, bias=True)
-            self.fc3 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
-            self.act_fn1 = nn.GELU()
-            self.act_fn2 = nn.GELU()
-
-    def forward(self, img_patches: torch.Tensor) -> torch.Tensor:
-        if not self.use_fused_vision_backbone:
-            projected_features = self.fc1(img_patches)
-            projected_features = self.act_fn1(projected_features)
-            projected_features = self.fc2(projected_features)
-        else:
-            projected_features = self.fc1(img_patches)
-            projected_features = self.act_fn1(projected_features)
-            projected_features = self.fc2(projected_features)
-            projected_features = self.act_fn2(projected_features)
-            projected_features = self.fc3(projected_features)
-
-        return projected_features
-
-
-
-# === Main HF Class Definitions ===
-@dataclass
-class PrismaticCausalLMOutputWithPast(ModelOutput):
-    """Base class for Prismatic casual (visually-conditioned) language model outputs; also exposes visual features."""
-
-    loss: Optional[torch.FloatTensor] = None
-    logits: torch.FloatTensor = None
-    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
-    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
-    attentions: Optional[Tuple[torch.FloatTensor]] = None
-
-    # Additions for VLMs
-    projector_features: Optional[torch.FloatTensor] = None
-    language_embeddings: Optional[torch.FloatTensor] = None
-
-
-
-class PrismaticPreTrainedModel(PreTrainedModel):
-    config_class: PretrainedConfig = PrismaticConfig
-    base_model_prefix: str = "model"
-    supports_gradient_checkpointing: bool = True
-
-    _no_split_modules: ClassVar[List[str]] = ["PrismaticProjector"]
-    _skip_keys_device_placement: str = "past_key_values"
-    _supports_flash_attn_2: bool = True
-
-    def _init_weights(self, module: nn.Module) -> None:
-        # Important :: this HF ported version is *not* meant for training from scratch; only inference and fine-tuning!
-        #   => As such, this init_weights code is not correct; if training VLMs from scratch, use the main codebase at
-        #      https://github.com/TRI-ML/prismatic-vlms
-        std = (
-            self.config.initializer_range
-            if hasattr(self.config, "initializer_range")
-            else self.config.text_config.initializer_range
-        )
-
-        if hasattr(module, "class_embedding"):
-            module.class_embedding.data.normal_(mean=0.0, std=std)
-
-        if isinstance(module, (nn.Linear, nn.Conv2d)):
-            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_()
-
-    @property
-    def _supports_sdpa(self) -> bool:
-        """Check LLM supports SDPA Attention"""
-        return self.language_model._supports_sdpa
-
-
-
-class PrismaticForConditionalGeneration(PrismaticPreTrainedModel):
-    def __init__(self, config: PrismaticConfig) -> None:
-        super().__init__(config)
-
-        # [Validation] Lightweight Validate on `config` Fields + Dependency Versions
-        if config.use_fused_vision_backbone is None:
-            raise ValueError("Missing config field `use_fused_vision_backbone`")
-
-        if timm.__version__ not in {"0.9.10", "0.9.11", "0.9.12", "0.9.16"}:
-            raise NotImplementedError(
-                "TIMM Version must be >= 0.9.10 and < 1.0.0 (breaking); please raise a GitHub Issue "
-                "if you urgently need support for latest TIMM versions."
-            )
-
-        if (transformers.__version__ != "4.40.1") or (tokenizers.__version__ != "0.19.1"):
-            logger.warning(
-                f"Expected `transformers==4.40.1` and `tokenizers==0.19.1` but got "
-                f"`transformers=={transformers.__version__}` and `tokenizers=={tokenizers.__version__}`; "
-                f"there might be inference-time regressions due to dependency changes. If in doubt, please"
-                f"use the above versions."
-            )
-        
-        # Instantiate PrismaticVisionBackbone (w/ Potential Fused Backbone)
-        self.vision_backbone = PrismaticVisionBackbone(
-            config.use_fused_vision_backbone, config.image_sizes, config.timm_model_ids, config.timm_override_act_layers
-        )
-
-        # Create Multimodal Projector
-        self.projector = PrismaticProjector(
-            config.use_fused_vision_backbone,
-            vision_dim=self.vision_backbone.embed_dim,
-            llm_dim=config.text_config.hidden_size,
-        )
-
-        # Instantiate LLM Backbone
-        self.language_model = AutoModelForCausalLM.from_config(
-            config.text_config, attn_implementation=config._attn_implementation
-        )
-
-        self.vocab_size = config.text_config.vocab_size
-        self.pad_token_id = config.pad_token_id
-        self.llm_dim = config.text_config.hidden_size
-
-        # Action query token
-        self.action_queries = nn.Embedding(NUM_TOKENS, self.llm_dim)
-        self.action_queries.weight.data.zero_()
-
-        # Uniform wrist dropout / language conditioning flags (persisted via config)
-        self.uniform_vision_dropout_enabled = getattr(config, "uniform_vision_dropout_enabled", False)
-        self.uniform_vision_dropout_ratio = getattr(config, "uniform_vision_dropout_ratio", 0.0)
-        self.language_conditioning_enabled = getattr(config, "language_conditioning_enabled", False)
-
-        # HF Boilerplate =>> initializes weights via `_init_weights()` and sets gradient checkpointing
-        self.post_init()
-
-    # === `PreTrainedModel` Boilerplate ===
-    def get_input_embeddings(self) -> nn.Module:
-        return self.language_model.get_input_embeddings()
-    def set_version(self, version: str):
-        self.version = version
-        return self.version
-
-
-    def set_input_embeddings(self, value: nn.Module) -> None:
-        self.language_model.set_input_embeddings(value)
-
-    def get_output_embeddings(self) -> nn.Module:
-        return self.language_model.get_output_embeddings()
-
-    def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
-        self.language_model.set_output_embeddings(new_embeddings)
-
-    def get_decoder(self) -> nn.Module:
-        return self.language_model.get_decoder()
-
-    def set_decoder(self, decoder: nn.Module) -> None:
-        self.language_model.set_decoder(decoder)
-
-    def tie_weights(self) -> None:
-        self.language_model.tie_weights()  # Note: `Llama-2` and `Mistral` don't tie weights (no-op)
-
-    def resize_token_embeddings(
-        self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None
-    ) -> nn.Embedding:
-        updated_embeddings = self.language_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
-
-        # Update config/instance variables
-        self.config.text_config.vocab_size = updated_embeddings.num_embeddings
-        self.vocab_size = updated_embeddings.num_embeddings
-
-        return updated_embeddings
-
-    def set_uniform_vision_dropout(self, enabled: bool, dropout_ratio: float) -> None:
-        """Enable or disable uniform wrist patch dropout during multimodal masking."""
-        self.uniform_vision_dropout_enabled = enabled
-        self.uniform_vision_dropout_ratio = dropout_ratio
-        setattr(self.config, "uniform_vision_dropout_enabled", enabled)
-        setattr(self.config, "uniform_vision_dropout_ratio", float(dropout_ratio))
-
-    def set_language_conditioning(self, enabled: bool) -> None:
-        """Toggle language conditioning inputs for the action head."""
-        self.language_conditioning_enabled = enabled
-        setattr(self.config, "language_conditioning_enabled", enabled)
-
-    def _replace_input_embeddings(self, input_embeddings, all_actions_mask, noisy_action_features):
-        """
-        Replace embeddings in input_embeddings at positions where all_actions_mask is True
-        with embeddings from noisy_action_features, using vectorized operations.
-
-        Args:
-            input_embeddings: Tensor of shape (B, S, D)
-            all_actions_mask: Boolean tensor of shape (B, S)
-            noisy_action_features: Tensor of shape (B, K, D) where K is the number of True values in mask per sample
-
-        Returns:
-            Modified input_embeddings tensor
-        """
-        # Clone input to avoid modifying the original tensor
-        new_input_embeddings = input_embeddings.clone()
-
-        # Create a tensor with the same shape of input_embeddings to hold the noisy action features
-        repositioned_noisy_action_features = torch.zeros_like(input_embeddings)
-
-        # Create batch indices for splicing
-        batch_indices = torch.arange(input_embeddings.shape[0], device=input_embeddings.device)
-        batch_indices = batch_indices.unsqueeze(1).expand(-1, noisy_action_features.shape[1])
-
-        # Get indices where mask is True for each sample
-        masked_indices = torch.stack([torch.where(mask)[0] for mask in all_actions_mask])
-
-        # Move the noisy action features into their correct positions
-        # print(noisy_action_features.size())
-        
-        repositioned_noisy_action_features[batch_indices, masked_indices] = noisy_action_features
-
-        # Combine original input embeddings and noisy action embeddings using the mask
-        new_input_embeddings = torch.where(
-            all_actions_mask.unsqueeze(-1), repositioned_noisy_action_features, new_input_embeddings
-        )
-
-        return new_input_embeddings
-    
-    def _process_action_masks(self, labels):
-        """Helper to get action masks from labels"""
-        current_action_mask = get_current_action_mask(labels)
-        next_actions_mask = get_next_actions_mask(labels)
-        all_actions_mask = current_action_mask | next_actions_mask  # (B, seq_len)
-        return all_actions_mask
-
-    def _process_vision_features(self, pixel_values, language_embeddings=None, use_film=False):
-        """Process vision features with optional FiLM conditioning"""
-        if use_film:
-            # FiLM: Infuse language inputs into visual features
-            patch_features = self.vision_backbone(pixel_values, language_embeddings)  # (bsz, 256 * num_images, D)
-        else:
-            patch_features = self.vision_backbone(pixel_values)  # (bsz, 256 * num_images, D)
-
-        # Project patch embeddings into language embedding space
-        return self.projector(patch_features)
-
-    def _process_proprio_features(self, projected_patch_embeddings, proprio, proprio_projector):
-        """Process proprioceptive features and append to vision features"""
-        if proprio_projector is not None and proprio is not None:
-            # projected_patch_embeddings: (bsz, num_patches * num_images, llm_dim)
-            # proprio: (bsz, proprio_dim) or (propro_dim,)
-            proprio = proprio.reshape(projected_patch_embeddings.shape[0], -1)  # (bsz, proprio_dim)
-            proprio_features = proprio_projector(proprio)  # (bsz, llm_dim)
-            proprio_features = proprio_features.unsqueeze(dim=1)  # (bsz, 1, llm_dim)
-            # For simplicity, just append proprio token to the end of projected vision patch tokens
-            return torch.cat((projected_patch_embeddings, proprio_features), dim=1)
-        return projected_patch_embeddings
-
-    def _get_wrist_enabled_mask(
-        self,
-        modality_mask: Optional[torch.Tensor],
-        batch_size: int,
-        device: torch.device,
-    ) -> Optional[torch.Tensor]:
-        num_wrist_views = max(self.vision_backbone.get_num_images_in_input() - 1, 0)
-        if num_wrist_views == 0:
-            return None
-
-        wrist_enabled = torch.ones((batch_size, num_wrist_views), dtype=torch.bool, device=device)
-        if modality_mask is None:
-            return wrist_enabled
-
-        if modality_mask.dim() == 1:
-            modality_mask = modality_mask.unsqueeze(0)
-
-        if modality_mask.shape[1] >= 2:
-            vision_enabled = modality_mask[:, 1].bool().unsqueeze(1)
-            wrist_enabled = vision_enabled.repeat(1, num_wrist_views)
-
-        if modality_mask.shape[1] > 4:
-            provided_wrist_flags = min(num_wrist_views, modality_mask.shape[1] - 4)
-            wrist_enabled[:, :provided_wrist_flags] = modality_mask[:, 4 : 4 + provided_wrist_flags].bool()
-
-        return wrist_enabled
-
-    def _get_strict_mae_enabled_mask(
-        self,
-        modality_mask: Optional[torch.Tensor],
-        batch_size: int,
-        device: torch.device,
-    ) -> torch.Tensor:
-        """
-        Return per-sample boolean flag for strict MAE-style masking.
-        Layout convention (dataset-side modality_mask):
-            [language, vision, query, proprio, wrist_flag_0, wrist_flag_1, strict_mae_flag]
-        """
-        strict_enabled = torch.zeros((batch_size,), dtype=torch.bool, device=device)
-        if modality_mask is None:
-            return strict_enabled
-
-        if modality_mask.dim() == 1:
-            modality_mask = modality_mask.unsqueeze(0)
-
-        strict_flag_idx = 6
-        if modality_mask.shape[1] > strict_flag_idx:
-            strict_enabled = modality_mask[:, strict_flag_idx].bool()
-        return strict_enabled
-
-    def _build_multimodal_attention(
-        self,
-        input_embeddings,
-        projected_patch_embeddings,
-        attention_mask,
-        modality_mask=None,
-        action_mask: Optional[torch.Tensor] = None,
-        use_query_mask: bool = False,
-        query_mask_ratio: float = 0.0,
-    ):
-        """Build multimodal embeddings and attention mask, with optional intra-query masking."""
-
-        projected_patch_attention_mask = None
-        if attention_mask is not None:
-            device = attention_mask.device
-            dtype_mask = input_embeddings.dtype
-            attention_keep = attention_mask.bool().clone()
-            # Preserve the original token keep mask before query masking mutates it.
-            base_attention_keep = attention_keep.clone()
-
-            batch_size = attention_keep.shape[0]
-            vision_enabled = torch.ones(batch_size, dtype=torch.bool, device=device)
-            query_enabled = torch.ones(batch_size, dtype=torch.bool, device=device)
-            proprio_enabled = torch.ones(batch_size, dtype=torch.bool, device=device)
-            wrist_enabled = None
-            strict_mae_enabled = torch.zeros(batch_size, dtype=torch.bool, device=device)
-
-            if modality_mask is not None:
-                modality_mask = modality_mask.to(device)
-                if modality_mask.dim() == 1:
-                    modality_mask = modality_mask.unsqueeze(0)
-                num_modalities = modality_mask.shape[1]
-                if num_modalities >= 2:
-                    vision_enabled = modality_mask[:, 1].bool()
-                if num_modalities >= 3:
-                    query_enabled = modality_mask[:, 2].bool()
-                if num_modalities >= 4:
-                    proprio_enabled = modality_mask[:, 3].bool()
-                wrist_enabled = self._get_wrist_enabled_mask(modality_mask, batch_size, device)
-                strict_mae_enabled = self._get_strict_mae_enabled_mask(modality_mask, batch_size, device)
-            else:
-                wrist_enabled = self._get_wrist_enabled_mask(None, batch_size, device)
-
-            if strict_mae_enabled.any():
-                strict_token_mask = (~base_attention_keep) & strict_mae_enabled.unsqueeze(1)
-                if strict_token_mask.any():
-                    input_embeddings = input_embeddings.masked_fill(strict_token_mask.unsqueeze(-1), 0)
-
-            wrist_patch_dropout_mask = None
-            if (
-                self.training
-                and self.uniform_vision_dropout_enabled
-                and self.vision_backbone.get_num_images_in_input() > 1
-            ):
-                wrist_patch_dropout_mask = self._build_uniform_wrist_dropout_mask(
-                    projected_patch_embeddings,
-                    vision_enabled,
-                    wrist_enabled,
-                    dtype=projected_patch_embeddings.dtype,
-                    device=device,
-                )
-
-            if (
-                self.training
-                and use_query_mask
-                and query_mask_ratio > 0.0
-                and action_mask is not None
-            ):
-                action_mask_bool = action_mask.to(device=device).bool()
-                for b in range(batch_size):
-                    if query_enabled[b]:
-                        continue
-                    query_indices = torch.where(action_mask_bool[b] & attention_keep[b])[0]
-                    if query_indices.numel() == 0:
-                        continue
-                    keep = torch.rand(query_indices.numel(), device=device) > query_mask_ratio
-                    if not keep.any():
-                        keep[torch.randint(query_indices.numel(), (1,), device=device)] = True
-                    attention_keep[b, query_indices] = keep
-
-                    keep_float = keep.to(input_embeddings.dtype)
-                    input_embeddings[b, query_indices] *= keep_float.unsqueeze(-1)
-                    kept_tokens = keep_float.sum()
-                    if kept_tokens > 0:
-                        scale = query_indices.numel() / kept_tokens
-                        input_embeddings[b, query_indices] *= scale
-
-            attention_mask = attention_keep.to(dtype_mask)
-
-            projected_patch_attention_mask = torch.ones(
-                (projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
-                dtype=dtype_mask,
-                device=device,
-            )
-            if modality_mask is not None:
-                patches_per_image = self.vision_backbone.get_num_patches()
-                num_vision_tokens = patches_per_image * self.vision_backbone.get_num_images_in_input()
-                wrist_start = patches_per_image
-                wrist_end = num_vision_tokens
-                if wrist_end > wrist_start:
-                    if wrist_enabled is not None:
-                        for wrist_idx in range(wrist_enabled.shape[1]):
-                            start_idx = patches_per_image * (wrist_idx + 1)
-                            end_idx = min(start_idx + patches_per_image, num_vision_tokens)
-                            if end_idx <= start_idx:
-                                continue
-                            wrist_mask = wrist_enabled[:, wrist_idx].to(projected_patch_attention_mask.dtype).unsqueeze(1)
-                            projected_patch_attention_mask[:, start_idx:end_idx] *= wrist_mask
-                    else:
-                        vision_mask = vision_enabled.to(projected_patch_attention_mask.dtype).unsqueeze(1)
-                        projected_patch_attention_mask[:, wrist_start:wrist_end] *= vision_mask
-                if projected_patch_attention_mask.shape[1] > num_vision_tokens:
-                    proprio_mask = proprio_enabled.to(projected_patch_attention_mask.dtype).unsqueeze(1)
-                    projected_patch_attention_mask[:, num_vision_tokens:] *= proprio_mask
-
-                if strict_mae_enabled.any() and wrist_enabled is not None:
-                    strict_wrist_keep = torch.ones_like(projected_patch_attention_mask)
-                    for wrist_idx in range(wrist_enabled.shape[1]):
-                        start_idx = patches_per_image * (wrist_idx + 1)
-                        end_idx = min(start_idx + patches_per_image, num_vision_tokens)
-                        if end_idx <= start_idx:
-                            continue
-                        unmasked_or_disabled = (wrist_enabled[:, wrist_idx] | (~strict_mae_enabled)).to(
-                            strict_wrist_keep.dtype
-                        ).unsqueeze(1)
-                        strict_wrist_keep[:, start_idx:end_idx] *= unmasked_or_disabled
-                    projected_patch_embeddings = projected_patch_embeddings * strict_wrist_keep.unsqueeze(-1)
-
-            if wrist_patch_dropout_mask is not None:
-                projected_patch_attention_mask = projected_patch_attention_mask * wrist_patch_dropout_mask
-                projected_patch_embeddings = projected_patch_embeddings * wrist_patch_dropout_mask.unsqueeze(-1)
-
-        multimodal_embeddings = torch.cat(
-            [input_embeddings[:, :1, :], projected_patch_embeddings, input_embeddings[:, 1:, :]], dim=1
-        )
-
-        multimodal_attention_mask = None
-        if attention_mask is not None:
-            multimodal_attention_mask = torch.cat(
-                [attention_mask[:, :1], projected_patch_attention_mask, attention_mask[:, 1:]], dim=1
-            )
-
-        return multimodal_embeddings, multimodal_attention_mask
-
-    def _build_multimodal_labels(self, labels, projected_patch_embeddings):
-        """Build multimodal labels with IGNORE_INDEX for patch embeddings"""
-        if labels is not None:
-            projected_patch_labels = torch.full(
-                (projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
-                fill_value=IGNORE_INDEX,
-                dtype=labels.dtype,
-                device=labels.device,
-            )
-            return torch.cat([labels[:, :1], projected_patch_labels, labels[:, 1:]], dim=1)
-        return None
-
-    def _build_uniform_wrist_dropout_mask(
-        self,
-        projected_patch_embeddings: torch.Tensor,
-        vision_enabled: torch.Tensor,
-        wrist_enabled: Optional[torch.Tensor],
-        dtype: torch.dtype,
-        device: torch.device,
-    ) -> Optional[torch.Tensor]:
-        """Build DropPath-style mask for wrist patches ensuring even coverage."""
-
-        batch_size, total_tokens = projected_patch_embeddings.shape[:2]
-        patches_per_image = self.vision_backbone.get_num_patches()
-        num_images = self.vision_backbone.get_num_images_in_input()
-        if patches_per_image == 0 or num_images <= 1:
-            return None
-
-        mask = torch.ones((batch_size, total_tokens), device=device, dtype=dtype)
-        grid_size = max(int(math.sqrt(patches_per_image)), 1)
-        wrist_ranges = [
-            (img_idx * patches_per_image, (img_idx + 1) * patches_per_image)
-            for img_idx in range(1, num_images)
-        ]
-
-        dropout_ratio = max(0.0, min(float(self.uniform_vision_dropout_ratio), 1.0))
-
-        for b in range(batch_size):
-            if not vision_enabled[b]:
-                continue
-            for wrist_idx, (start_idx, end_idx) in enumerate(wrist_ranges):
-                if wrist_enabled is not None and wrist_idx < wrist_enabled.shape[1] and not wrist_enabled[b, wrist_idx]:
-                    continue
-                num_patches = end_idx - start_idx
-                if num_patches <= 0:
-                    continue
-                dropout_tokens = int(round(num_patches * dropout_ratio))
-                dropout_tokens = min(dropout_tokens, max(num_patches - 1, 0))
-
-                keep_mask = torch.ones(num_patches, device=device, dtype=dtype)
-                if dropout_tokens > 0:
-                    selected_indices = self._sample_uniform_patch_indices(num_patches, dropout_tokens, grid_size)
-                    if selected_indices:
-                        keep_mask[selected_indices] = 0.0
-                kept = keep_mask.sum()
-                if kept > 0 and kept < num_patches:
-                    keep_mask = keep_mask * (num_patches / kept)
-                mask[b, start_idx:end_idx] = keep_mask
-
-        if torch.all(mask == 1):
-            return None
-        return mask
-
-    def _sample_uniform_patch_indices(
-        self,
-        patches_per_image: int,
-        num_to_drop: int,
-        grid_size: int,
-    ) -> List[int]:
-        if num_to_drop <= 0:
-            return []
-
-        selected: List[int] = []
-        attempts = 0
-        max_attempts = max(num_to_drop * 20, 1)
-
-        while len(selected) < num_to_drop and attempts < max_attempts:
-            candidate = random.randrange(patches_per_image)
-            row, col = divmod(candidate, grid_size)
-            if any(abs(row - r0) <= 1 and abs(col - c0) <= 1 for idx0 in selected for r0, c0 in [divmod(idx0, grid_size)]):
-                attempts += 1
-                continue
-            selected.append(candidate)
-
-        if len(selected) < num_to_drop:
-            remaining = [idx for idx in range(patches_per_image) if idx not in selected]
-            if remaining:
-                needed = min(num_to_drop - len(selected), len(remaining))
-                selected.extend(random.sample(remaining, needed))
-
-        return selected[:num_to_drop]
-
-    # === Core Prismatic VLM `forward()` Logic ===
-    def forward(
-        self,
-        input_ids: Optional[torch.LongTensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        pixel_values: Optional[torch.FloatTensor] = None,
-        labels: Optional[torch.LongTensor] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        past_key_values: Optional[List[torch.FloatTensor]] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        output_projector_features: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-        proprio=None,
-        proprio_projector=None,
-        noisy_actions=None,
-        noisy_action_projector=None,
-        diffusion_timestep_embeddings=None,
-        use_film: bool = False,
-        modality_mask: Optional[torch.Tensor] = None,
-        use_query_mask: bool = False,
-        query_mask_ratio: float = 0.0,
-    ) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
-        """Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
-        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
-        )
-        output_projector_features = output_projector_features if output_projector_features is not None else False
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        # Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
-        use_cache = use_cache and not self.training
-
-        # Instantiate Placeholder for Projector Features / Language Conditioning
-        projected_patch_embeddings = None
-        language_embeddings = None
-
-        # === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
-        if input_ids.shape[1] == 1:
-            assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
-            assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
-            assert labels is None, "Unexpected key `labels` provided during cached generation!"
-
-            language_model_output = self.language_model(
-                input_ids=input_ids,
-                attention_mask=None,
-                position_ids=None,
-                past_key_values=past_key_values,
-                inputs_embeds=None,
-                labels=None,
-                use_cache=use_cache,
-                output_attentions=output_attentions,
-                output_hidden_states=output_hidden_states,
-                return_dict=return_dict,
-            )
-
-        # === Handle Unimodal Forward ===
-        elif pixel_values is None:
-            assert (input_ids is not None) and (inputs_embeds is None), "Missing `input_ids` in language-only forward!"
-            assert past_key_values is None, "Unexpected key `past_key_values` provided during language-only forward!"
-
-            language_model_output = self.language_model(
-                input_ids=input_ids,
-                attention_mask=attention_mask,
-                position_ids=None,
-                past_key_values=None,
-                inputs_embeds=None,
-                labels=labels,
-                use_cache=use_cache,
-                output_attentions=output_attentions,
-                output_hidden_states=output_hidden_states,
-                return_dict=return_dict,
-            )
-
-        # === Handle Multimodal Forward ===
-        elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
-            assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
-
-            # Get input embeddings (from language model embeddings)
-            input_embeddings = self.get_input_embeddings()(input_ids)  # (B, seq_len, D)
-            if labels is not None and labels.device != input_embeddings.device:
-                labels = labels.to(input_embeddings.device)
-
-            
-            # Extract action masks
-            all_actions_mask = self._process_action_masks(labels)
-
-            # Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
-            
-            # print(input_embeddings[~all_actions_mask].size())
-            language_embeddings = input_embeddings[~all_actions_mask].reshape(
-                input_embeddings.shape[0], -1, input_embeddings.shape[2]
-            )  # (B, lang_seq_len, llm_dim)
-
-            # Get visual features
-            projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
-
-            if modality_mask is not None:
-                modality_mask = modality_mask.to(projected_patch_embeddings.device)
-                if modality_mask.dim() == 1:
-                    modality_mask = modality_mask.unsqueeze(0)
-                modality_mask = modality_mask.to(projected_patch_embeddings.dtype)
-
-            #NOTE Add proprioceptive state if provided
-            # if proprio_projector is not None:
-            #     projected_patch_embeddings = self._process_proprio_features(
-            #         projected_patch_embeddings,
-            #         proprio,
-            #         proprio_projector,
-            #     )
-
-            # Process action embeddings - standard action queries
-            action_queries = self.action_queries.weight  # (NUM_TOKENS, h)
-            action_queries = action_queries.unsqueeze(0).repeat(input_embeddings.shape[0], 1, 1)  # (B, NUM_TOKENS, h)
-
-            # Replace action token embeddings with action queries
-            all_actions_mask = self._process_action_masks(labels)
-            input_embeddings = self._replace_input_embeddings(
-                input_embeddings, all_actions_mask, action_queries
-            )
-
-            # Build multimodal embeddings & attention mask
-            multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
-                input_embeddings,
-                projected_patch_embeddings,
-                attention_mask,
-                modality_mask=modality_mask,
-                action_mask=all_actions_mask,
-                use_query_mask=use_query_mask,
-                query_mask_ratio=query_mask_ratio,
-            )
-
-            # Build labels for multimodal sequence if needed
-            multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
-
-            # Dispatch to language model
-            language_model_output = self.language_model(
-                input_ids=None,
-                attention_mask=multimodal_attention_mask,
-                position_ids=None,
-                past_key_values=None,
-                inputs_embeds=multimodal_embeddings,
-                labels=None,
-                use_cache=use_cache,
-                output_attentions=output_attentions,
-                output_hidden_states=output_hidden_states,
-                return_dict=return_dict,
-                ) 
-
-        # === Otherwise =>> Assume Invalid! ===
-        elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
-            raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
-
-        else:
-            raise ValueError(
-                "Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
-                f"=> `input_ids` = {input_ids is not None}\n"
-                f"=> `attention_mask` = {attention_mask is not None}\n"
-                f"=> `pixel_values` = {pixel_values is not None}\n"
-                f"=> `labels` = {labels is not None}\n"
-                f"=> `input_embeds` = {inputs_embeds is not None}\n"
-                f"=> `past_key_values` = {past_key_values is not None}\n"
-                f"=> `use_cache` = {use_cache}"
-            )
-
-        # Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
-        if not return_dict:
-            if output_projector_features and (projected_patch_embeddings is not None):
-                return *language_model_output, projected_patch_embeddings
-
-            return language_model_output
-
-        return PrismaticCausalLMOutputWithPast(
-            loss=language_model_output.loss,
-            logits=language_model_output.logits,
-            past_key_values=language_model_output.past_key_values,
-            hidden_states=language_model_output.hidden_states,
-            attentions=language_model_output.attentions,
-            projector_features=projected_patch_embeddings,
-            language_embeddings=language_embeddings,
-            )
-
-
-# === GenerationMixin Methods ===
-class OpenVLAForActionPrediction(PrismaticForConditionalGeneration):
-    config_class: PretrainedConfig = OpenVLAConfig
-
-    def __init__(self, config: OpenVLAConfig) -> None:
-        super().__init__(config)
-        self.norm_stats = config.norm_stats
-        
-
-        # Compute action bins
-        self.bins = np.linspace(-1, 1, config.n_action_bins)
-        self.bin_centers = (self.bins[:-1] + self.bins[1:]) / 2.0
-
-        # Compute vocab size for de-tokenization -- revert added "multiple of"
-        self.vocab_size = self.config.text_config.vocab_size - self.config.pad_to_multiple_of
-
-    def _prepare_input_for_action_prediction(self, input_ids, attention_mask):
-        """Prepares input for action prediction by adding necessary tokens"""
-        # Add (ACTION_DIM * NUM_ACTIONS_CHUNK) placeholder tokens to input_ids to simulate action tokens
-        placeholder_action_token_ids = (
-            torch.ones((input_ids.shape[0], NUM_TOKENS)).to(input_ids.device).to(input_ids.dtype)
-        )
-        input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1)
-
-        # Add stop token to sequence (needed in non-causal bi-directional self-attention, as it appears at train time)
-        stop_token_id = torch.ones((input_ids.shape[0], 1)).to(input_ids.device).to(input_ids.dtype) * STOP_INDEX
-        input_ids = torch.cat([input_ids, stop_token_id], dim=-1)
-
-        # Extend the attention mask to fit the new shape of input
-        # Note: Only batch size == 1 supported right now
-        mask_extension = (
-            torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1]))
-            .to(attention_mask.device)
-            .to(attention_mask.dtype)
-        )
-        attention_mask = torch.cat([attention_mask, mask_extension], dim=-1)
-
-        return input_ids, attention_mask
-
-    def _prepare_labels_for_action_prediction(self, labels, input_ids):
-        """Creates labels tensor for action prediction if not provided"""
-        # Extend labels tensor with fake action labels
-        ARBITRARY_ACTION_TOKEN_IDX = ACTION_TOKEN_BEGIN_IDX + 1
-        labels_extension = (
-            torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype)
-            * ARBITRARY_ACTION_TOKEN_IDX
-        )
-        labels = torch.cat([labels, labels_extension], dim=-1)
-
-        # Replace last label token with stop token
-        labels[:, -1] = STOP_INDEX
-
-        return labels
-
-    def _unnormalize_actions(self, normalized_actions, unnorm_key=None):
-        """Unnormalize actions using dataset statistics"""
-        action_norm_stats = self.get_action_stats(unnorm_key)
-
-        if ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS:
-            mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["min"], dtype=bool))
-            action_high, action_low = np.array(action_norm_stats["max"]), np.array(action_norm_stats["min"])
-        elif ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS_Q99:
-            mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["q01"], dtype=bool))
-            action_high, action_low = np.array(action_norm_stats["q99"]), np.array(action_norm_stats["q01"])
-        else:
-            raise ValueError("Unsupported action/proprio normalization type detected!")
-
-        actions = np.where(
-            mask,
-            0.5 * (normalized_actions + 1) * (action_high - action_low + 1e-8) + action_low,
-            normalized_actions,
-        )
-
-        return actions
-
-
-    def _regression_or_discrete_prediction(
-        self,
-        input_embeddings,
-        all_actions_mask,
-        projected_patch_embeddings,
-        attention_mask,
-        labels,
-        NUM_PATCHES,
-        NUM_PROMPT_TOKENS,
-        action_head=None,
-        proprio=None,
-        proprio_projector=None,
-    ):
-        """Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
-
-        # Standard action queries
-        action_queries = self.action_queries.weight  # (num_tokens, h)
-        action_queries = action_queries.unsqueeze(0).repeat(input_embeddings.shape[0], 1, 1)  # (B, num_tokens, h)
-
-        # Replace action token embeddings with action queries
-        input_embeddings = self._replace_input_embeddings(input_embeddings.clone(), all_actions_mask, action_queries)
-
-        # Build multimodal embeddings and attention mask
-        multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
-            input_embeddings, projected_patch_embeddings, attention_mask
-        )
-
-        # Forward pass through language model
-        language_model_output = self.language_model(
-            input_ids=None,
-            attention_mask=multimodal_attention_mask,
-            position_ids=None,
-            past_key_values=None,
-            inputs_embeds=multimodal_embeddings,
-            labels=None,
-            use_cache=None,
-            output_attentions=False,
-            output_hidden_states=True,
-            return_dict=True,
-        )
-
-        # Extract hidden states for action tokens
-        multi_layer_hidden_states = []
-       
-        for item in language_model_output.hidden_states[0:]:
-            # last_hidden_states = output.hidden_states[-1]  # (B, seq_len, D)
-            # Get hidden states for text portion of prompt+response (after the vision patches)
-            text_hidden_states = item
-            # Get hidden states for action portion of response
-            actions_hidden_states = text_hidden_states[:, NUM_PATCHES+ NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + NUM_TOKENS, :,].reshape(1, 1, NUM_TOKENS , -1).to(torch.bfloat16)
-            
-            batch_size = item.shape[0]
-            task_latten_states = item[:, :NUM_PATCHES].reshape(batch_size, 1, NUM_PATCHES , -1)
-            all_hidden_states = torch.cat((task_latten_states, actions_hidden_states),2)
-            multi_layer_hidden_states.append(all_hidden_states)
-            
-        multi_layer_hidden_states = torch.cat(multi_layer_hidden_states, dim = 1)
-        
-
-        # Handle different prediction methods
-        if action_head is not None:
-            # L1 regression prediction
-            normalized_actions = action_head.predict_action(
-                multi_layer_hidden_states,
-                proprio=proprio,
-                proprio_projector=proprio_projector,
-            )
-            normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
-            normalized_actions = normalized_actions.float().cpu().detach().numpy()
-        else:
-            # Discrete token-based prediction
-            predicted_action_token_ids = (
-                language_model_output.logits[
-                    :,
-                    NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
-                ]
-                .argmax(dim=2)
-                .cpu()
-                .numpy()
-            )
-            discretized_actions = self.vocab_size - predicted_action_token_ids
-            discretized_actions = np.clip(discretized_actions - 1, a_min=0, a_max=self.bin_centers.shape[0] - 1)
-            normalized_actions = self.bin_centers[discretized_actions]
-            normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
-
-        return normalized_actions, actions_hidden_states
-
-    def _unwrap_action_head(self, action_head: Optional[nn.Module]) -> Optional[nn.Module]:
-        """Return the underlying nn.Module if wrapped (e.g., DDP)."""
-        if action_head is None:
-            return None
-        return action_head.module if hasattr(action_head, "module") else action_head
-
-    def _gather_action_hidden_states(
-        self,
-        hidden_states: Tuple[torch.Tensor, ...],
-        num_patches: int,
-        num_prompt_tokens: int,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """Extract task and action hidden states across layers for downstream heads."""
-        multi_layer_hidden_states = []
-        actions_hidden_states = None
-        for item in hidden_states:
-            batch_size = item.shape[0]
-            actions_hidden_states = item[
-                :,
-                num_patches + num_prompt_tokens : num_patches + num_prompt_tokens + NUM_TOKENS,
-                :,
-            ].reshape(batch_size, 1, NUM_TOKENS, -1)
-            task_latten_states = item[:, :num_patches].reshape(batch_size, 1, num_patches, -1)
-            all_hidden_states = torch.cat((task_latten_states, actions_hidden_states), 2)
-            multi_layer_hidden_states.append(all_hidden_states)
-        return torch.cat(multi_layer_hidden_states, dim=1), actions_hidden_states
-
-    def _compute_action_hidden_states(
-        self,
-        input_embeddings,
-        all_actions_mask,
-        projected_patch_embeddings,
-        attention_mask,
-        labels,
-        NUM_PATCHES,
-        NUM_PROMPT_TOKENS,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """Build multimodal sequence and collect hidden states for flow / mip inference."""
-        # Standard action queries
-        action_queries = self.action_queries.weight
-        action_queries = action_queries.unsqueeze(0).repeat(input_embeddings.shape[0], 1, 1)
-
-        input_embeddings = self._replace_input_embeddings(
-            input_embeddings.clone(), all_actions_mask, action_queries
-        )
-
-        multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
-            input_embeddings, projected_patch_embeddings, attention_mask
-        )
-
-        language_model_output = self.language_model(
-            input_ids=None,
-            attention_mask=multimodal_attention_mask,
-            position_ids=None,
-            past_key_values=None,
-            inputs_embeds=multimodal_embeddings,
-            labels=None,
-            use_cache=None,
-            output_attentions=False,
-            output_hidden_states=True,
-            return_dict=True,
-        )
-
-        multi_layer_hidden_states, actions_hidden_states = self._gather_action_hidden_states(
-            language_model_output.hidden_states,
-            NUM_PATCHES,
-            NUM_PROMPT_TOKENS,
-        )
-        return multi_layer_hidden_states, actions_hidden_states
-
-    def _predict_flow_actions(
-        self,
-        base_hidden_states: torch.Tensor,
-        action_head: nn.Module,
-        proprio=None,
-        proprio_projector=None,
-        flow_num_steps: int = 1,
-        flow_sample_mode: str = "stochastic",
-    ) -> torch.Tensor:
-        """Run ODE-style integration using the learned velocity field for flow inference.
-
-        Args:
-            base_hidden_states: Base hidden states for action prediction
-            action_head: Action head module
-            proprio: Proprioceptive features
-            proprio_projector: Proprioceptive projector
-            flow_num_steps: Number of ODE solver steps
-            flow_sample_mode: Sampling mode (stochastic or deterministic)
-        """
-        batch_size = base_hidden_states.shape[0]
-        device = base_hidden_states.device
-        dtype = base_hidden_states.dtype
-
-        num_steps = max(1, int(flow_num_steps))
-        sample_mode = flow_sample_mode
-
-        # Initialize action state
-        if sample_mode == "stochastic":
-            act_s = torch.randn(
-                batch_size, NUM_ACTIONS_CHUNK, ACTION_DIM, device=device, dtype=dtype
-            )
-        else:
-            act_s = torch.zeros(
-                batch_size, NUM_ACTIONS_CHUNK, ACTION_DIM, device=device, dtype=dtype
-            )
-
-        # ODE integration
-        t_schedule = torch.linspace(0.0, 1.0, steps=num_steps + 1, device=device, dtype=dtype)
-        for i in range(num_steps):
-            s_val = t_schedule[i]
-            t_val = t_schedule[i + 1]
-            timestep = s_val.expand(batch_size, 1)
-
-            # Predict velocity at current state
-            # action_init: (B, NUM_ACTIONS_CHUNK, ACTION_DIM)
-            # timestep: (B, 1)
-            velocity = action_head.predict_action(
-                base_hidden_states,
-                proprio=proprio,
-                proprio_projector=proprio_projector,
-                action_init=act_s,
-                timestep=timestep,
-                phase="Inference",
-            )
-            if velocity.dim() == 2:
-                velocity = velocity.view(batch_size, NUM_ACTIONS_CHUNK, ACTION_DIM)
-
-            # Euler step: act_s = act_s + velocity * dt
-            act_s = act_s + velocity * (t_val - s_val)
-
-        return act_s
-
-    def _predict_mip_actions(
-        self,
-        base_hidden_states: torch.Tensor,
-        action_head: nn.Module,
-        proprio=None,
-        proprio_projector=None,
-        flow_t_two_step: float = 0.9,
-    ) -> torch.Tensor:
-        """Two-step deterministic MIP inference.
-
-        Args:
-            base_hidden_states: Base hidden states for action prediction
-            action_head: Action head module
-            proprio: Proprioceptive features
-            proprio_projector: Proprioceptive projector
-            flow_t_two_step: Two-step midpoint timestep
-        """
-        batch_size = base_hidden_states.shape[0]
-        device = base_hidden_states.device
-        dtype = base_hidden_states.dtype
-
-        t_two_step = float(flow_t_two_step)
-
-        # Initialize states
-        zero_actions = torch.zeros(
-            batch_size, NUM_ACTIONS_CHUNK, ACTION_DIM, device=device, dtype=dtype
-        )
-        t_zero = torch.zeros(batch_size, 1, device=device, dtype=dtype)
-        t_mid = torch.full((batch_size, 1), t_two_step, device=device, dtype=dtype)
-
-        # First step: predict from zero
-        # action_init: (B, NUM_ACTIONS_CHUNK, ACTION_DIM)
-        # timestep: (B, 1)
-        pred_first = action_head.predict_action(
-            base_hidden_states,
-            proprio=proprio,
-            proprio_projector=proprio_projector,
-            action_init=zero_actions,
-            timestep=t_zero,
-            phase="Inference",
-        )
-        if pred_first.dim() == 2:
-            pred_first = pred_first.view(batch_size, NUM_ACTIONS_CHUNK, ACTION_DIM)
-        # print(f"pred_first: {pred_first}")
-        # Second step: predict from pred_first
-        pred_second = action_head.predict_action(
-            base_hidden_states,
-            proprio=proprio,
-            proprio_projector=proprio_projector,
-            action_init=pred_first,
-            timestep=t_mid,
-            phase="Inference",
-        )
-        if pred_second.dim() == 2:
-            pred_second = pred_second.view(batch_size, NUM_ACTIONS_CHUNK, ACTION_DIM)
-
-        return pred_second
-
-
-    def predict_action(
-        self,
-        input_ids: Optional[torch.LongTensor] = None,
-        unnorm_key: Optional[str] = None,
-        proprio=None,
-        proprio_projector=None,
-        action_head=None,
-        noisy_action_projector=None,
-        use_film: bool = False,
-        action_head_config=None,
-        **kwargs: str,
-    ) -> np.ndarray:
-        """Predict actions from input sequence, with options for different prediction methods.
-
-        Args:
-            input_ids: Input token ids
-            unnorm_key: Key for unnormalization statistics
-            proprio: Proprioceptive features
-            proprio_projector: Projector for proprioceptive features
-            action_head: Optional head for L1 regression or diffusion-based prediction
-            noisy_action_projector: Projector for noisy actions in diffusion-based prediction
-            use_film: Whether to use FiLM conditioning
-            action_head_config: Configuration for action head (includes flow/mip settings)
-            **kwargs: Additional arguments including pixel_values and attention_mask
-
-        Returns:
-            Tuple of (unnormalized_actions, action_hidden_states)
-        """
-
-        pixel_values = kwargs["pixel_values"] # [1, 12, 224, 224]
-        attention_mask = kwargs["attention_mask"] # 
-
-        # Create fake labels tensor (needed for action mask)
-        labels = input_ids.clone()
-        labels[:] = IGNORE_INDEX
-
-        # Get number of tokens in prompt (excluding the start token)
-        NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1  # Subtract action tokens and stop token
-
-        # Prepare inputs by adding necessary tokens
-        input_ids, attention_mask = self._prepare_input_for_action_prediction(input_ids, attention_mask)
-
-        # Update labels tensor for action mask computation later
-        labels = self._prepare_labels_for_action_prediction(labels, input_ids)
-
-        # Get input embeddings and action masks
-        input_embeddings = self.get_input_embeddings()(input_ids)
-        all_actions_mask = self._process_action_masks(labels)
-
-        # Extract language embeddings
-        language_embeddings = input_embeddings[~all_actions_mask].reshape(
-            input_embeddings.shape[0], -1, input_embeddings.shape[2]
-        )
-
-        # Process vision features
-        projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
-
-        # Add proprioceptive features if provided
-        use_proprio = proprio_projector is not None and proprio is not None
-        if use_proprio:
-            proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
-        #     projected_patch_embeddings = self._process_proprio_features(
-        #         projected_patch_embeddings, proprio, proprio_projector
-        #     )
-
-        # Calculate number of patches (including proprio token if present)
-        NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
-        # if use_proprio:
-        #     NUM_PATCHES += 1
-
-        actions_hidden_states = None
-        normalized_actions = None
-        loss_type = getattr(action_head_config, "loss_type", None)
-        unwrapped_action_head = self._unwrap_action_head(action_head)
-
-        if loss_type in ("flow", "mip"):
-            if unwrapped_action_head is None:
-                raise ValueError(f"`action_head` must be provided for {loss_type} inference.")
-
-            # Get flow/mip settings from action_head_config if provided, otherwise fall back to self.config
-            if action_head_config is not None:
-                flow_num_steps = getattr(action_head_config, "flow_num_steps", 1)
-                flow_sample_mode = getattr(action_head_config, "flow_sample_mode", "stochastic")
-                flow_t_two_step = getattr(action_head_config, "flow_t_two_step", 0.9)
-            else:
-                flow_num_steps = getattr(self.config, "flow_num_steps", 1)
-                flow_sample_mode = getattr(self.config, "flow_sample_mode", "stochastic")
-                flow_t_two_step = getattr(self.config, "flow_t_two_step", 0.9)
-
-            multi_layer_hidden_states, actions_hidden_states = self._compute_action_hidden_states(
-                input_embeddings,
-                all_actions_mask,
-                projected_patch_embeddings,
-                attention_mask,
-                labels,
-                NUM_PATCHES,
-                NUM_PROMPT_TOKENS,
-            )
-
-            if loss_type == "flow":
-                normalized_actions = self._predict_flow_actions(
-                    multi_layer_hidden_states,
-                    unwrapped_action_head,
-                    proprio=proprio,
-                    proprio_projector=proprio_projector,
-                    flow_num_steps=flow_num_steps,
-                    flow_sample_mode=flow_sample_mode,
-                )
-            else:
-                normalized_actions = self._predict_mip_actions(
-                    multi_layer_hidden_states,
-                    unwrapped_action_head,
-                    proprio=proprio,
-                    proprio_projector=proprio_projector,
-                    flow_t_two_step=flow_t_two_step,
-                )
-        else:
-            # Run regression or discrete token-based prediction
-            normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
-                input_embeddings,
-                all_actions_mask,
-                projected_patch_embeddings,
-                attention_mask,
-                labels,
-                NUM_PATCHES,
-                NUM_PROMPT_TOKENS,
-                action_head=action_head,
-                proprio=proprio,
-                proprio_projector=proprio_projector,
-            )
-
-        if isinstance(normalized_actions, torch.Tensor):
-            normalized_actions = normalized_actions.reshape(normalized_actions.shape[0], -1, ACTION_DIM)
-            normalized_actions = normalized_actions[0].float().cpu().detach().numpy()
-
-        # Unnormalize predicted actions
-        actions = self._unnormalize_actions(normalized_actions, unnorm_key)
-
-        return actions, actions_hidden_states
-
-
-
-    @staticmethod
-    def _check_unnorm_key(norm_stats: Dict[str, Dict[str, Any]], unnorm_key: Optional[str]) -> str:
-        """Validate and resolve the unnormalization key for action statistics"""
-        if unnorm_key is None:
-            assert len(norm_stats) == 1, (
-                f"Your model was trained on more than one dataset, "
-                f"please pass a `unnorm_key` from the following options to choose the statistics "
-                f"used for un-normalizing actions: {norm_stats.keys()}"
-            )
-            unnorm_key = next(iter(norm_stats.keys()))
-
-        assert unnorm_key in norm_stats, (
-            f"The `unnorm_key` you chose is not in the set of available dataset statistics, "
-            f"please choose from: {norm_stats.keys()}"
-        )
-        return unnorm_key
-
-    def get_action_dim(self, unnorm_key: Optional[str] = None) -> int:
-        """Get the dimensionality of the policy's action space."""
-        unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
-        return len(self.norm_stats[unnorm_key]["action"]["min"])
-
-    def get_action_stats(self, unnorm_key: Optional[str] = None) -> Dict[str, Any]:
-        """Get all the logged statistics for the given dataset."""
-        unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
-        return self.norm_stats[unnorm_key]["action"]