""" configuration_prismatic.py HuggingFace-style configuration definition for Prismatic VLMs, inheriting from `transformers.PretrainedConfig`. Default configuration specifies `siglip-224px+7b`. """ from typing import Any, Dict, List, Optional import transformers from transformers import PretrainedConfig from transformers.models.auto import CONFIG_MAPPING import numpy as np # === Utilities for Mapping Prismatic names to HF names === # fmt: off VISION_BACKBONE_TO_RESOLUTION: Dict[str, List[int]] = { "clip-vit-l": [224], "siglip-vit-so400m": [224], "dinov2-vit-l": [224], "in1k-vit-l": [224], "clip-vit-l-336px": [336], "siglip-vit-so400m-384px": [384], "dinoclip-vit-l-336px": [336, 336], "dinosiglip-vit-so-224px": [224, 224], "dinosiglip-vit-so-384px": [384, 384], } VISION_BACKBONE_TO_TIMM_ID: Dict[str, List[str]] = { "clip-vit-l": ["vit_large_patch14_clip_224.openai"], "clip-vit-l-336px": ["vit_large_patch14_clip_336.openai"], "dinov2-vit-l": ["vit_large_patch14_reg4_dinov2.lvd142m"], "in1k-vit-l": ["vit_large_patch16_224.augreg_in21k_ft_in1k"], "siglip-vit-so400m": ["vit_so400m_patch14_siglip_224"], "siglip-vit-so400m-384px": ["vit_so400m_patch14_siglip_384"], "dinoclip-vit-l-336px": ["vit_large_patch14_reg4_dinov2.lvd142m", "vit_large_patch14_clip_336.openai"], "dinosiglip-vit-so-224px": ["vit_large_patch14_reg4_dinov2.lvd142m", "vit_so400m_patch14_siglip_224"], "dinosiglip-vit-so-384px": ["vit_large_patch14_reg4_dinov2.lvd142m", "vit_so400m_patch14_siglip_384"], } TIMM_OVERRIDE_ACT_LAYER: Dict[str, List[Optional[str]]] = { "clip-vit-l": ["quick_gelu"], "clip-vit-l-336px": ["quick_gelu"], "dinov2-vit-l": [None], "in1k-vit-l": [None], "siglip-vit-so400m": [None], "siglip-vit-so400m-384px": [None], "dinoclip-vit-l-336px": [None, "quick_gelu"], "dinosiglip-vit-so-224px": [None, None], "dinosiglip-vit-so-384px": [None, None] } LLM_BACKBONE_TO_HF_PATH = { "llama2-7b-pure": "meta-llama/Llama-2-7b-hf", "llama2-13b-pure": "meta-llama/Llama-2-13b-hf", "llama2-7b-chat": "meta-llama/Llama-2-7b-chat-hf", "llama2-13b-chat": "meta-llama/Llama-2-13b-chat-hf", "vicuna-v15-7b": "lmsys/vicuna-7b-v1.5", "vicuna-v15-13b": "lmsys/vicuna-13b-v1.5", "mistral-v0.1-7b-pure": "mistralai/Mistral-7B-v0.1", "mistral-v0.1-7b-instruct": "mistralai/Mistral-7B-Instruct-v0.1", "phi-2-3b": "microsoft/phi-2", } LLM_BACKBONE_TO_HF_METACLASS = { "llama2-7b-pure": "llama", "llama2-13b-pure": "llama", "llama2-7b-chat": "llama", "llama2-13b-chat": "llama", "vicuna-v15-7b": "llama", "vicuna-v15-13b": "llama", "mistral-v0.1-7b-pure": "mistral", "mistral-v0.1-7b-instruct": "mistral", "phi-2-3b": "phi", } VALID_VISION_BACKBONES = set(VISION_BACKBONE_TO_RESOLUTION.keys()) VALID_LLM_BACKBONES = set(LLM_BACKBONE_TO_HF_PATH) # fmt: on class WaypointTokenizer: """ Wraps base LLM/VLM tokenizer and overloads least used token as a control token NOTE: By default, assumes a BPE-style tokenizer akin to the LlamaTokenizer, where *the least used tokens* appear at the end of the vocabulary! TODO: Adding new token vs overloading? When I call `tokenizer.add_token()` vocab stays the same """ def __init__(self, tokenizer: transformers.PreTrainedTokenizerBase, num_tokens: int = 10) -> None: self.tokenizer = tokenizer self.num_tokens = num_tokens def __call__(self, *_) -> str: """Get the text token for control""" return self.tokenizer.decode(self.control_token_ids) @property def control_token_ids(self) -> np.ndarray: # Assumes we're overwriting the final tokens of the vocabulary (least used tokens) return np.arange(self.num_tokens) + int(self.tokenizer.vocab_size - self.num_tokens) @property def num_control_tokens(self) -> int: return self.num_tokens class PrismaticConfig(PretrainedConfig): model_type: str = "prismatic" is_composition: bool = False def __init__( self, vision_backbone_id: str = "dinosiglip-vit-so-224px", llm_backbone_id: str = "llama2-7b-pure", arch_specifier: str = "no-align+gelu-mlp", ## TODO: check use_fused_vision_backbone: Optional[bool] = None, ## TODO: check image_resize_strategy: str = "letterbox", text_config: Optional[Dict[str, Any]] = None, llm_max_length: int = 2048, pad_token_id: int = 32000, pad_to_multiple_of: int = 64, output_projector_states: bool = False, **kwargs: str, ) -> None: if vision_backbone_id not in VALID_VISION_BACKBONES: raise ValueError(f"Vision backbone `{vision_backbone_id}` not in {VALID_VISION_BACKBONES = }") if llm_backbone_id not in VALID_LLM_BACKBONES: raise ValueError(f"LLM backbone `{llm_backbone_id}` not in {VALID_LLM_BACKBONES = }") # Set Prismatic Configuration Fields self.vision_backbone_id = vision_backbone_id self.llm_backbone_id = llm_backbone_id self.arch_specifier = arch_specifier self.output_projector_states = output_projector_states # [Contract] All vision backbone parameters are lists =>> supports fused backbones with different preprocessing self.use_fused_vision_backbone = ( use_fused_vision_backbone if use_fused_vision_backbone is not None else any(self.vision_backbone_id.startswith(v) for v in ["dinoclip", "dinosiglip"]) ) self.timm_model_ids = VISION_BACKBONE_TO_TIMM_ID[self.vision_backbone_id] self.timm_override_act_layers = TIMM_OVERRIDE_ACT_LAYER[self.vision_backbone_id] self.image_sizes = VISION_BACKBONE_TO_RESOLUTION[self.vision_backbone_id] self.image_resize_strategy = image_resize_strategy self.hf_llm_id = LLM_BACKBONE_TO_HF_PATH[self.llm_backbone_id] self.llm_max_length = llm_max_length self.pad_token_id, self.pad_to_multiple_of = pad_token_id, pad_to_multiple_of # [IMPORTANT] HF Utilities actually look for a `text_config` field... we need to use that specific naming! self.text_config = ( CONFIG_MAPPING[LLM_BACKBONE_TO_HF_METACLASS[self.llm_backbone_id]](**text_config) if text_config is not None else CONFIG_MAPPING[LLM_BACKBONE_TO_HF_METACLASS[self.llm_backbone_id]]() ) # Dispatch **kwargs to super() =>> note that `pad_token_id` collides, so we pass it in here as well... super().__init__(pad_token_id=pad_token_id, **kwargs) # Here we need trajectory_vla config, with # prismatic_config fields and then the waypointer fields class TrajectoryVLAConfig(PretrainedConfig): def __init__( self, prismatic_config = {}, token_size: int = 1024, # Timestep token size cheat: bool = False, # If True, cheat and use action tokens; Works only with OpenVLA checkpoint num_timesteps: int = 20, # Number of prediction time steps rotation_components: int = 9, # Number of rotation componens: euler -> 3, quaternion -> 4, rotmat -> 9 num_timestep_tokens : int = 3, seperate_control_proj: bool = True, # If True, project control components separately timestep_proj_config: Dict[str, Any] = {}, token_proj_config: Dict[str, Any] = {}, transformer_config: Dict[str, Any] = {}, # prismatic_config: PrismaticConfig, # waypointer_config: Dict[str, Any], # **kwargs: str, ): # super().__init__(**prismatic_config) self.prismatic_config = PrismaticConfig(**prismatic_config) self.token_size = token_size self.cheat = cheat self.num_timesteps = num_timesteps self.rotation_components = rotation_components self.seperate_control_proj = seperate_control_proj self.timestep_proj_config = timestep_proj_config self.token_proj_config = token_proj_config self.transformer_config = transformer_config # self.num_timestep_tokens = num_timestep_tokens @property def control_components(self) -> int: # Number of control dimensions: 3 translation, N rotation, 1 gripper return 3 + self.rotation_components + 1 @property def num_timestep_tokens(self) -> int: return self.timestep_proj_config['num_tokens'] # class WaypointerConfig(ConfigurableModuleConfig): # token_size: int = 1024 # Timestep token size # cheat: bool # If True, cheat and use action tokens; Works only with OpenVLA checkpoint # timestep_proj_config: AutoConfig # Timestep tokens # token_proj_config: TokenProjectorConfig # LLM output tokens projection and packing # transformer_config: AutoConfig # Transformer config # # Output configurations # num_timesteps: int = 20 # Number of prediction time steps # rotation_components: int = 3 # Number of rotation componens: euler -> 3, quaternion -> 4, rotmat -> 9 # separate_control_proj: bool = True # If True, project control components separately # @property # def control_components(self) -> int: # # Number of control dimensions: 3 translation, N rotation, 1 gripper # return 3 + self.rotation_components + 1 # @property # def num_timestep_tokens(self) -> int: # return self.timestep_proj_config.num_tokens class OpenVLAConfig(PrismaticConfig): model_type: str = "openvla" def __init__( self, norm_stats: Optional[Dict[str, Dict[str, Dict[str, Dict[str, List[float]]]]]] = None, n_action_bins: int = 256, **kwargs: str, ) -> None: self.norm_stats, self.n_action_bins = norm_stats, n_action_bins super().__init__(**kwargs) if __name__ == "__main__" : # yaml_file = 'barrel/pipes/vlams/configs/waypoints/waypointer_multistep_fractal.yaml' prismatic_config = PrismaticConfig() print(prismatic_config) prismatic_config_dict = { "vision_backbone_id":"dinosiglip-vit-so-224px", # "llm_backbone_id":"llama2-7b-pure",meta-llama/Llama-2-7b-hf "llm_backbone_id": "meta-llama/Llama-2-7b-hf", "arch_specifier": "no-align+gelu-mlp", ## TODO: check "use_fused_vision_backbone" :None, ## TODO: check "image_resize_strategy" : "letterbox", "text_config" : None, "llm_max_length" : 2048, "pad_token_id" :32000, "pad_to_multiple_of" : 64, "output_projector_states" : False, } token_proj_config = { "vit_tokens_layers": [2176, 1024], "control_tokens_layers": [4096, 2048, 1024], "image_tokens_mode": 'vit', } timestep_proj_config = { "pos_embed_scale": 1.0, "proj_layers": [1024], "time_delta_sec": 0.1, "num_tokens":3 } TrajectoryVlaConfig = { "prismatic_config":prismatic_config_dict, "token_size": 1024, "cheat": False, "num_timesteps": 20, "rotation_components": 3, "seperate_control_proj": True, "timestep_proj_config": {}, "token_proj_config": {}, "transformer_config": {}, } TrajectoryVLAConfig = TrajectoryVLAConfig( **TrajectoryVlaConfig) print(TrajectoryVLAConfig) class WaypointTokenizer: """ Wraps base LLM/VLM tokenizer and overloads least used token as a control token NOTE: By default, assumes a BPE-style tokenizer akin to the LlamaTokenizer, where *the least used tokens* appear at the end of the vocabulary! TODO: Adding new token vs overloading? When I call `tokenizer.add_token()` vocab stays the same """ def __init__(self, tokenizer: transformers.PreTrainedTokenizerBase, num_tokens: int = 10) -> None: self.tokenizer = tokenizer self.num_tokens = num_tokens def __call__(self, *_) -> str: """Get the text token for control""" return self.tokenizer.decode(self.control_token_ids) @property def control_token_ids(self) -> np.ndarray: # Assumes we're overwriting the final tokens of the vocabulary (least used tokens) return np.arange(self.num_tokens) + int(self.tokenizer.vocab_size - self.num_tokens) @property def num_control_tokens(self) -> int: return self.num_tokens