# ------------------------------------------------------------------------------ # 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 import logging import traceback from typing import Any, Dict import numpy as np import torch import torch.nn.functional as F from fastapi import FastAPI from fastapi.responses import JSONResponse from PIL import Image import uvicorn import json_numpy import cv2 from transformers import PreTrainedModel from .modeling_florence2 import Florence2ForConditionalGeneration from .transformer import SoftPromptedTransformer from .action_hub import build_action_space from .configuration_xvla import XVLAConfig class XVLA(PreTrainedModel): """ 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, ) 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, num_skills=getattr(config, 'num_skills', 0), ) # Skill classifier (auxiliary head for inference-time skill prediction) if getattr(config, 'num_skills', 0) > 0: self.skill_classifier = torch.nn.Linear(projection_dim, config.num_skills) self.skill_classifier_weight = getattr(config, 'skill_classifier_weight', 0.1) # Class weights for weighted CE (sqrt-inverse frequency) skill_class_weights = getattr(config, 'skill_class_weights', None) if skill_class_weights is not None: self.register_buffer('_skill_class_weights', torch.tensor(skill_class_weights, dtype=torch.float32)) else: self._skill_class_weights = None # Progress head: continuous regression (0→1) for position within current skill # Conditioned on VLM features + skill embedding if getattr(config, 'use_progress_head', False): skill_emb_dim = 64 self.progress_skill_emb = torch.nn.Embedding( max(getattr(config, 'num_skills', 0), 1), skill_emb_dim) self.progress_head = torch.nn.Sequential( torch.nn.Linear(projection_dim + skill_emb_dim, 1), torch.nn.Sigmoid(), ) self.progress_head_weight = getattr(config, 'progress_head_weight', 0.1) # Deferred FastAPI app self.app: FastAPI | None = None # ============================= 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] skill_id: torch.LongTensor | None = None, skill_progress: torch.Tensor | None = None, **kwargs, ) -> 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. 4) Auxiliary skill classifier + progress head losses. """ enc = self.forward_vlm(input_ids, image_input, image_mask) # Auxiliary heads (computed in fp32, detached from VLM) skill_cls_loss = None progress_loss = None if (hasattr(self, 'skill_classifier') and skill_id is not None) or \ (hasattr(self, 'progress_head') and skill_progress is not None): with torch.amp.autocast('cuda', enabled=False): vlm_pooled = enc["vlm_features"].detach().mean(dim=1).float() if hasattr(self, 'skill_classifier') and skill_id is not None: skill_logits = self.skill_classifier(vlm_pooled) skill_cls_loss = F.cross_entropy(skill_logits, skill_id, weight=self._skill_class_weights) if hasattr(self, 'progress_head') and skill_progress is not None and skill_id is not None: skill_emb = self.progress_skill_emb(skill_id) # [B, 64] progress_input = torch.cat([vlm_pooled, skill_emb], dim=-1) pred_progress = self.progress_head(progress_input).squeeze(-1) # [B] progress_loss = F.mse_loss(pred_progress, skill_progress) 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, skill_id=skill_id, **enc, ) loss_dict = self.action_space.compute_loss(pred_action, action) if skill_cls_loss is not None: loss_dict["skill_cls_loss"] = skill_cls_loss * self.skill_classifier_weight if progress_loss is not None: loss_dict["progress_loss"] = progress_loss * self.progress_head_weight return loss_dict # ================================= 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, skill_id: torch.LongTensor | None = None, **kwargs, ) -> torch.Tensor: """ Iterative denoising (linear schedule). Applies action_space.postprocess at the end (e.g., sigmoid on gripper). If skill_id is None and a skill_classifier exists, auto-predicts skill. Also predicts progress (0→1) if progress_head exists. """ self.eval() enc = self.forward_vlm(input_ids, image_input, image_mask) # Auto-predict skill when not provided if skill_id is None and hasattr(self, 'skill_classifier'): with torch.amp.autocast('cuda', enabled=False): vlm_pooled = enc["vlm_features"].mean(dim=1).float() skill_logits = self.skill_classifier(vlm_pooled) skill_id = skill_logits.argmax(dim=-1) else: vlm_pooled = None # Predict progress for monitoring (conditioned on skill) if hasattr(self, 'progress_head') and skill_id is not None: with torch.amp.autocast('cuda', enabled=False): if vlm_pooled is None: vlm_pooled = enc["vlm_features"].mean(dim=1).float() skill_emb = self.progress_skill_emb(skill_id) progress_input = torch.cat([vlm_pooled, skill_emb], dim=-1) self._last_predicted_progress = self.progress_head(progress_input).squeeze(-1) 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, skill_id=skill_id, **enc, ) return self.action_space.postprocess(action) # =============================== FastAPI service ============================= def _build_app(self, processor): """ Minimal FastAPI app for XVLA inference. Args: processor: callable(images, text) -> Dict[str, torch.Tensor] expected keys: "input_ids", "image_input", "image_mask" """ if self.app is not None: return app = FastAPI() @app.post("/act") def act(payload: Dict[str, Any]): try: self.eval() # Decode up to 3 image inputs images = [] for key in ("image0", "image1", "image2"): if key not in payload: continue v = json_numpy.loads(payload[key]) if isinstance(v, np.ndarray): if v.ndim == 1: # encoded bytes v = cv2.imdecode(v, cv2.IMREAD_COLOR) images.append(Image.fromarray(v)) elif isinstance(v, (list, tuple)): images.append(Image.fromarray(np.array(v))) elif isinstance(v, str): images.append(Image.open(v)) if not images: return JSONResponse({"error": "No valid images found."}, status_code=400) # Multimodal preprocessing by processor inputs = processor(images, payload["language_instruction"]) if not {"input_ids", "image_input", "image_mask"}.issubset(inputs): return JSONResponse({"error": "Processor returned incomplete inputs."}, status_code=400) # Build proprio/domain tensors proprio = torch.as_tensor(np.asarray(json_numpy.loads(payload["proprio"]))) domain_id = torch.tensor([int(payload["domain_id"])], dtype=torch.long) # Align to model's device/dtype device = next(self.parameters()).device dtype = next(self.parameters()).dtype def to_model(t: torch.Tensor) -> torch.Tensor: if not isinstance(t, torch.Tensor): t = torch.as_tensor(t) # cast floats to model dtype, keep integral/bool as-is return t.to(device=device, dtype=dtype) if t.is_floating_point() else t.to(device=device) inputs = {k: to_model(v) for k, v in inputs.items()} inputs.update({ "proprio": to_model(proprio.unsqueeze(0)), "domain_id": domain_id.to(device), }) # Inference steps = int(payload.get("steps", 10)) action = self.generate_actions(**inputs, steps=steps).squeeze(0).float().cpu().numpy() return JSONResponse({"action": action.tolist()}) except Exception: logging.error(traceback.format_exc()) return JSONResponse({"error": "Request failed"}, status_code=400) self.app = app def run(self, processor, host: str = "0.0.0.0", port: int = 8000): """ Launch the FastAPI service. """ self._build_app(processor) assert self.app is not None uvicorn.run(self.app, host=host, port=port)