modeling_xvla.py

# ------------------------------------------------------------------------------
# Copyright 2025 2toINF (https://github.com/2toINF)
#
# 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 __future__ import annotations



from typing import Any, Dict, List
import torch

import numpy as np
from PIL import Image
from fastapi import FastAPI
import cv2

from transformers import PreTrainedModel
from .server import ModelServer
from .modeling_florence2 import Florence2ForConditionalGeneration
from .transformer import SoftPromptedTransformer
from .action_hub import build_action_space
from .configuration_xvla import XVLAConfig


class XVLA(PreTrainedModel, ModelServer):
    """
    XVLA: HuggingFace-compatible Vision-Language-Action policy.

    Components:
      • Florence2 encoder-only backbone (vision-language)
      • SoftPromptedTransformer (temporal/action head)
      • Action space (pre/post-processing + loss)
    """
    config_class = XVLAConfig
    base_model_prefix = "xvla"
    supports_gradient_checkpointing = True

    def __init__(self, config: XVLAConfig, *args, **kwargs):
        super().__init__(config, *args, **kwargs)

        # Core settings
        self.num_actions: int = config.num_actions
        self.use_proprio: bool = config.use_proprio
        self.action_mode: str = config.action_mode.lower()
        # Action space (dimensions + hooks)
        if config.action_mode.lower() == "auto":
            self.action_space = build_action_space(
                config.action_mode.lower(),
                real_dim=config.real_action_dim,
                max_dim=config.max_action_dim,
                idx_for_delta=config.idx_for_delta,
                idx_for_mask_proprio=config.idx_for_mask_proprio
            )
        else:
            self.action_space = build_action_space(config.action_mode.lower())
        dim_action = self.action_space.dim_action
        dim_proprio = getattr(self.action_space, "dim_proprio", dim_action)

        # Florence2 backbone (encoder only)
        self.vlm = Florence2ForConditionalGeneration(config.florence_config).to(torch.float32)
        if hasattr(self.vlm, "language_model"):
            lm = self.vlm.language_model
            if hasattr(lm, "model") and hasattr(lm.model, "decoder"):
                del lm.model.decoder
            if hasattr(lm, "lm_head"):
                del lm.lm_head

        projection_dim = getattr(self.vlm.config, "projection_dim", None)
        if projection_dim is None:
            raise ValueError("Florence2 config must provide `projection_dim` for multimodal fusion.")

        # Temporal/action head
        self.transformer = SoftPromptedTransformer(
            hidden_size=config.hidden_size,
            multi_modal_input_size=projection_dim,
            depth=config.depth,
            num_heads=config.num_heads,
            mlp_ratio=config.mlp_ratio,
            num_domains=config.num_domains,
            dim_action=dim_action,
            dim_propio=dim_proprio,
            len_soft_prompts=config.len_soft_prompts,
            dim_time=config.dim_time,
            max_len_seq=config.max_len_seq,
            use_hetero_proj=config.use_hetero_proj,
        )

        # Deferred FastAPI app
        self.app: FastAPI | None = None

    # ========================== pretrained loading ================================
    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        """
        Load pretrained XVLA, automatically handling action-head dimension
        mismatches.

        * Shape-compatible parameters are loaded normally.
        * Mismatched parameters are logged and explicitly re-initialised
          (Xavier-uniform for weight, zeros for bias — matching
          ``DomainAwareLinear.__init__``).
        """
        import os
        import json
        import logging
        from collections import OrderedDict

        logger = logging.getLogger(__name__)

        config = kwargs.pop("config", None)
        torch_dtype = kwargs.pop("torch_dtype", None)

        if config is None:
            config = cls.config_class.from_pretrained(
                pretrained_model_name_or_path, **kwargs
            )

        model = cls(config, *model_args)
        if torch_dtype is not None:
            model = model.to(torch_dtype)

        pretrained_state = cls._load_pretrained_state_dict(
            pretrained_model_name_or_path
        )
        model_state = model.state_dict()

        to_load = OrderedDict()
        mismatched = []

        for key, param in pretrained_state.items():
            if key not in model_state:
                continue
            if param.shape == model_state[key].shape:
                to_load[key] = param
            else:
                mismatched.append(
                    (key, tuple(param.shape), tuple(model_state[key].shape))
                )

        model.load_state_dict(to_load, strict=False)

        if mismatched:
            logger.warning(
                "=== Mismatched pretrained keys (reinitialized) ===\n"
                + "\n".join(
                    f"  {k}: pretrained {ps} -> current {cs}"
                    for k, ps, cs in mismatched
                )
            )
            for key, _, _ in mismatched:
                parts = key.split(".")
                module = model
                for part in parts[:-1]:
                    module = getattr(module, part)
                param = getattr(module, parts[-1])
                with torch.no_grad():
                    if "bias" in key:
                        torch.nn.init.zeros_(param)
                    elif param.dim() >= 2:
                        torch.nn.init.xavier_uniform_(param)
                    else:
                        torch.nn.init.zeros_(param)
            logger.warning(
                "Above %d parameter(s) have been re-initialised.",
                len(mismatched),
            )

        return model

    @staticmethod
    def _load_pretrained_state_dict(model_path: str) -> dict:
        """Load state dict from a local checkpoint (file or directory).

        Supports single-file, directory, and sharded safetensors / bin.
        """
        import os
        import json
        from collections import OrderedDict

        def _load_safetensors(path):
            from safetensors.torch import load_file
            return load_file(path)

        def _load_bin(path):
            return torch.load(path, map_location="cpu")

        if os.path.isfile(model_path):
            if model_path.endswith(".safetensors"):
                return _load_safetensors(model_path)
            return _load_bin(model_path)

        for fname, loader in [
            ("model.safetensors", _load_safetensors),
            ("pytorch_model.bin", _load_bin),
        ]:
            fpath = os.path.join(model_path, fname)
            if os.path.isfile(fpath):
                return loader(fpath)

        for index_name, loader in [
            ("model.safetensors.index.json", _load_safetensors),
            ("pytorch_model.bin.index.json", _load_bin),
        ]:
            index_path = os.path.join(model_path, index_name)
            if os.path.isfile(index_path):
                with open(index_path) as f:
                    weight_map = json.load(f)["weight_map"]
                state_dict = OrderedDict()
                for shard_file in dict.fromkeys(weight_map.values()):
                    state_dict.update(
                        loader(os.path.join(model_path, shard_file))
                    )
                return state_dict

        raise FileNotFoundError(
            f"No checkpoint found at '{model_path}'. Expected "
            f"model.safetensors, pytorch_model.bin, or sharded index files."
        )

    # ============================= Florence2 encoder =============================
    def forward_vlm(
        self,
        input_ids: torch.LongTensor,        # [B, L]
        pixel_values: torch.FloatTensor,    # [B, V, C, H, W]
        image_mask: torch.Tensor,           # [B, V] (bool or 0/1)
    ) -> Dict[str, torch.Tensor]:
        """
        Encode text + multi-view images via Florence2 encoder.

        Returns:
          { "vlm_features": [B, T_enc, D], "aux_visual_inputs": [B, (V-1)*N, D] }
        """
        B, V = pixel_values.shape[:2]
        flat_mask = image_mask.view(-1).to(torch.bool)         # [B*V]
        flat_images = pixel_values.flatten(0, 1)                # [B*V, C, H, W]

        num_valid = int(flat_mask.sum().item())
        if num_valid == 0:
            raise ValueError("At least one image view must be valid per batch.")

        valid_images = flat_images[flat_mask]                   # [#valid, C, H, W]
        valid_feats = self.vlm._encode_image(valid_images)      # [#valid, N, D]
        N, D = valid_feats.shape[1:]

        image_features = valid_feats.new_zeros((B * V, N, D))
        image_features[flat_mask] = valid_feats
        image_features = image_features.view(B, V, N, D)        # [B, V, N, D]

        inputs_embeds = self.vlm.get_input_embeddings()(input_ids)  # [B, L, D]

        merged_embeds, attention_mask = self.vlm._merge_input_ids_with_image_features(
            image_features[:, 0],  # first view: [B, N, D]
            inputs_embeds,         # [B, L, D]
        )

        enc_out = self.vlm.language_model.model.encoder(
            attention_mask=attention_mask,
            inputs_embeds=merged_embeds,
        )[0]  # [B, T_enc, D]

        aux_visual_inputs = image_features[:, 1:].reshape(B, -1, D)  # remaining views flattened
        return {"vlm_features": enc_out, "aux_visual_inputs": aux_visual_inputs}

    # ================================= training =================================
    def forward(
        self,
        input_ids: torch.LongTensor,
        image_input: torch.FloatTensor,
        image_mask: torch.Tensor,
        domain_id: torch.LongTensor,
        proprio: torch.Tensor,
        action: torch.Tensor,  # [B, T=num_actions, D=dim_action]
    ) -> Dict[str, torch.Tensor]:
        """
        1) Encode multimodal inputs.
        2) Diffusion-style noisy mixture of actions: x_t = t*noise + (1-t)*gt.
        3) Space-specific preprocessing, prediction, and supervised loss.
        """
        action, proprio = self.action_space.prepare_for_training(action, proprio)
        enc = self.forward_vlm(input_ids, image_input, image_mask)

        B = input_ids.shape[0]
        t = (torch.rand(1, device=input_ids.device)
             + torch.arange(B, device=input_ids.device) / B) % (1 - 1e-5)

        action_noisy = torch.randn_like(action) * t.view(-1, 1, 1) + action * (1 - t).view(-1, 1, 1)
        proprio_m, action_noisy_m = self.action_space.preprocess(proprio, action_noisy)

        pred_action = self.transformer(
            domain_id=domain_id,
            action_with_noise=action_noisy_m,
            t=t,
            proprio=proprio_m,
            **enc,
        )
        return self.action_space.compute_loss(pred_action, action)

    # ================================= inference =================================
    @torch.no_grad()
    def generate_actions(
        self,
        input_ids: torch.LongTensor,
        image_input: torch.FloatTensor,
        image_mask: torch.Tensor,
        domain_id: torch.LongTensor,
        proprio: torch.Tensor,
        steps: int = 10,
    ) -> torch.Tensor:
        """
        Iterative denoising (linear schedule).
        Applies action_space.postprocess at the end (e.g., sigmoid on gripper).
        """
        self.eval()
        enc = self.forward_vlm(input_ids, image_input, image_mask)

        B = input_ids.shape[0]
        D = self.action_space.dim_action

        x1 = torch.randn(B, self.num_actions, D, device=proprio.device, dtype=proprio.dtype)
        action = torch.zeros_like(x1)

        steps = max(1, int(steps))
        for i in range(steps, 0, -1):
            t = torch.full((B,), i / steps, device=proprio.device, dtype=proprio.dtype)
            x_t = x1 * t.view(-1, 1, 1) + action * (1 - t).view(-1, 1, 1)
            proprio_m, x_t_m = self.action_space.preprocess(proprio, x_t)
            action = self.transformer(
                domain_id=domain_id,
                action_with_noise=x_t_m,
                proprio=proprio_m,
                t=t,
                **enc,
            )
        return self.action_space.postprocess(action, proprio=proprio)

    # =============================== FastAPI service =============================


    def inference_api(self, payload: Dict[str, Any] | List[Dict[str, Any]], **kwargs) -> np.ndarray:
        """
        XVLA inference supporting:
        - Single sample: payload is a dict of scalars/arrays.
        - Grouped batch: payload is a list of dicts with same-length fields.

        payload contents:
        - "language_instruction": str or List[str], optional
        - "image0", "image1", ... : np.ndarray (H, W, C) or encoded buffer, required
        - "proprio": np.ndarray (D,) or (B, D), required
        - "domain_id": int / List[int] if batch > 1, required
        - "steps": int, optional, default=10
        - "batch_size": int, optional, default=1
        
        Returns:
        - (T, D) for single sample
        - (B, T, D) for grouped batch
        """
        # -------------------------
        # 1) Normalize payload -> List[Dict[str, Any]]
        # -------------------------
        processor = kwargs.get("processor")
        if isinstance(payload, dict):
            batch_payloads: List[Dict[str, Any]] = [payload]
        batch_size = len(batch_payloads)
        device = next(self.parameters()).device
        dtype = next(self.parameters()).dtype
        # -------------------------
        # 2) Utilities
        # -------------------------
        def move_to_device(x: Any) -> torch.Tensor:
            """Convert to tensor and move to model device/dtype."""
            tensor = x if isinstance(x, torch.Tensor) else torch.as_tensor(x)
            if tensor.is_floating_point():
                return tensor.to(device=device, dtype=dtype)
            return tensor.to(device=device)

        def decode_image_list(sample: Dict[str, Any]) -> List[Image.Image]:
            """Decode image0/image1/... from np.ndarray into PIL Images."""
            images: List[Image.Image] = []
            idx = 0
            while f"image{idx}" in sample:
                arr = sample[f"image{idx}"]
                if not isinstance(arr, np.ndarray): raise ValueError(f"image{idx} must be np.ndarray, got {type(arr)}")
                if arr.ndim == 1:  # encoded buffer
                    arr = cv2.imdecode(arr, cv2.IMREAD_COLOR)
                    if arr is None: raise ValueError(f"cv2.imdecode failed for image{idx}")
                    arr = cv2.cvtColor(arr, cv2.COLOR_BGR2RGB)
                images.append(Image.fromarray(arr))
                idx += 1
            if not images:
                raise ValueError("Missing images: expected keys image0, image1, ...")
            return images
        # -------------------------
        # 3) Per-sample preprocessing + strict collation (no padding)
        # -------------------------
        language_batch: List[str] = []
        images_batch: List[List[Image.Image]] = []
        proprio_batch: List[torch.Tensor] = []
        domain_id_list: List[int] = []
        denoiseing_steps = batch_payloads[0].get("steps", 10)
        
        for sample in batch_payloads:
            images_batch.append(decode_image_list(sample))
            language_batch.append(sample.get("language_instruction", ""))
            proprio_batch.append(move_to_device(sample["proprio"]))
            domain_id_list.append(int(sample.get("domain_id", 0)))
        model_inputs = processor(
            images=images_batch,
            language_instruction=language_batch,
        )
        model_inputs = {k: move_to_device(v) for k, v in model_inputs.items()}
        model_inputs.update(
            proprio=torch.stack(proprio_batch, dim=0),  # (B, state_dim)
            domain_id=torch.tensor(domain_id_list, dtype=torch.long, device=device),  # (B,)
            steps=denoiseing_steps,  # one scalar for whole batch
        )
        # -------------------------
        # 4) Inference
        # -------------------------
        self.eval()
        with torch.inference_mode():
            actions = self.generate_actions(**model_inputs)  # expected: (B, T, D)
        actions_np = actions.float().cpu().numpy()
        return actions_np[0] if batch_size == 1 else actions_np