File size: 3,617 Bytes
7f5d27f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | from __future__ import annotations
from collections.abc import Mapping
import torch
from transformers import AutoConfig, AutoModel, AutoModelForImageTextToText, PretrainedConfig, PreTrainedModel
try:
from .action_head import MiniCPMV_VLA_ActionHead
from .configuration_minicpm_vla import MiniCPMVLAConfig
except ImportError: # pragma: no cover - fallback for non-package loaders
from action_head import MiniCPMV_VLA_ActionHead
from configuration_minicpm_vla import MiniCPMVLAConfig
def _torch_dtype(name: str) -> torch.dtype:
try:
return getattr(torch, name)
except AttributeError as exc:
raise ValueError(f"Unsupported torch dtype: {name}") from exc
def _build_vlm_config(config_dict: dict) -> PretrainedConfig:
values = dict(config_dict)
model_type = values.pop("model_type")
return AutoConfig.for_model(model_type, **values)
class MiniCPMV_VLA(PreTrainedModel):
"""Complete VLA model loadable with ``AutoModel.from_pretrained``."""
config_class = MiniCPMVLAConfig
base_model_prefix = ""
_tied_weights_keys = {
"vlm.lm_head.weight": "vlm.model.language_model.embed_tokens.weight",
}
_no_split_modules = ["BasicTransformerBlock"]
def __init__(self, config: MiniCPMVLAConfig):
super().__init__(config)
if config.vlm_config is None:
raise ValueError("config.vlm_config is required")
vlm_config = _build_vlm_config(config.vlm_config)
self.vlm_dtype = _torch_dtype(config.vlm_dtype)
self.action_head_dtype = _torch_dtype(config.action_head_dtype)
self.vlm = AutoModelForImageTextToText.from_config(vlm_config).to(self.vlm_dtype)
self.vlm.config.hidden_size = self.vlm.config.text_config.hidden_size
self.action_head = MiniCPMV_VLA_ActionHead(
action_dim=config.action_dim,
state_dim=config.state_dim,
action_horizon=config.action_horizon,
num_inference_timesteps=config.num_inference_timesteps,
max_num_embodiments=config.max_num_embodiments,
).to(self.action_head_dtype)
self.post_init()
def get_input_embeddings(self):
return self.vlm.get_input_embeddings()
def set_input_embeddings(self, value):
return self.vlm.set_input_embeddings(value)
def _vlm_forward(self, data: Mapping[str, object]):
inputs = {key: value for key, value in data.items() if value is not None}
if "language_attention_mask" in inputs and "attention_mask" not in inputs:
inputs["attention_mask"] = inputs.pop("language_attention_mask")
inputs.setdefault("output_hidden_states", True)
inputs.setdefault("return_dict", True)
inputs.setdefault("use_cache", False)
inputs.setdefault("logits_to_keep", 1)
with torch.autocast("cuda", dtype=self.vlm_dtype, enabled=torch.cuda.is_available()):
try:
return self.vlm(**inputs)
except TypeError as exc:
if "logits_to_keep" not in str(exc):
raise
inputs.pop("logits_to_keep")
return self.vlm(**inputs)
@torch.no_grad()
def predict_action(
self,
state: torch.Tensor,
embodiment_id: torch.Tensor,
**vlm_inputs,
) -> torch.Tensor:
vl_embs = self._vlm_forward(vlm_inputs).hidden_states[-1].to(self.action_head_dtype)
return self.action_head.predict_action(vl_embs, state, embodiment_id)
AutoConfig.register("minicpm_vla", MiniCPMVLAConfig)
AutoModel.register(MiniCPMVLAConfig, MiniCPMV_VLA)
|