""" 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 from transformers import PretrainedConfig from transformers.models.auto import CONFIG_MAPPING # === Utilities for Mapping Prismatic names to HF names === # fmt: off VALID_TEXT_TOKEN_GATE_TEXT_POOLING_MODES = {"mean", "mlp", "cross_attention", "contrastive_alignment_score"} VALID_GAZING_MODES = {"mlp", "self_attention"} VALID_LAYER_GATE_MODES = {"none", "threshold"} 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 PrismaticConfig(PretrainedConfig): model_type: str = "prismatic" is_composition: bool = False def __init__( self, vision_backbone_id: str = "siglip-vit-so400m", llm_backbone_id: str = "vicuna-v15-7b", arch_specifier: str = "no-align+gelu-mlp", use_fused_vision_backbone: Optional[bool] = None, 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, use_text_token_gate: bool = False, text_token_gate_use_text_summary: bool = True, text_token_gate_use_vision_tokens: bool = True, text_token_gate_hidden_dim: int = 512, text_token_gate_text_pooling_mode: str = "mean", text_token_gate_text_pool_hidden_dim: int = 128, text_token_gate_cross_attention_dim: int = 256, text_token_gate_cross_attention_heads: int = 1, text_token_gate_mlp_depth: int = 1, contrastive_visual_tau: float = 0.1, contrastive_text_tau: float = 1.0, gazing_mode: str = "mlp", self_attn_dim: int = 512, self_attn_heads: int = 8, self_attn_layers: int = 1, layer_gate_mode: str = "none", layer_gate_threshold: float = 0.15, layer_gate_strength: float = 0.5, text_token_gate_budget: float = 0.5, text_token_gate_budget_loss_weight: float = 0.01, text_token_gate_linear_mean_penalty_weight: float = 0.0, text_token_gate_pool_instruction_only: bool = False, text_token_gate_filter_stopwords: bool = False, **kwargs: str, ) -> None: # Deprecated: nonzero text_token_gate_linear_mean_penalty_weight now enables the penalty directly. kwargs.pop("text_token_gate_use_linear_mean_penalty", None) legacy_weighted_text_pooling = kwargs.pop("text_token_gate_use_weighted_text_pooling", None) if legacy_weighted_text_pooling is not None and text_token_gate_text_pooling_mode == "mean": text_token_gate_text_pooling_mode = "mlp" if legacy_weighted_text_pooling else "mean" 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 = }") if text_token_gate_text_pooling_mode not in VALID_TEXT_TOKEN_GATE_TEXT_POOLING_MODES: raise ValueError( "`text_token_gate_text_pooling_mode` must be one of " f"{sorted(VALID_TEXT_TOKEN_GATE_TEXT_POOLING_MODES)}, got {text_token_gate_text_pooling_mode!r}" ) if text_token_gate_cross_attention_dim <= 0: raise ValueError( "`text_token_gate_cross_attention_dim` must be positive, " f"got {text_token_gate_cross_attention_dim}" ) if text_token_gate_cross_attention_heads <= 0: raise ValueError( "`text_token_gate_cross_attention_heads` must be positive, " f"got {text_token_gate_cross_attention_heads}" ) if text_token_gate_cross_attention_dim % text_token_gate_cross_attention_heads != 0: raise ValueError( "`text_token_gate_cross_attention_dim` must be divisible by " "`text_token_gate_cross_attention_heads`; got " f"{text_token_gate_cross_attention_dim} and {text_token_gate_cross_attention_heads}" ) if text_token_gate_mlp_depth <= 0: raise ValueError( "`text_token_gate_mlp_depth` must be positive, " f"got {text_token_gate_mlp_depth}" ) if contrastive_visual_tau <= 0: raise ValueError(f"`contrastive_visual_tau` must be positive, got {contrastive_visual_tau}") if contrastive_text_tau <= 0: raise ValueError(f"`contrastive_text_tau` must be positive, got {contrastive_text_tau}") if gazing_mode not in VALID_GAZING_MODES: raise ValueError(f"`gazing_mode` must be one of {sorted(VALID_GAZING_MODES)}, got {gazing_mode!r}") if self_attn_dim <= 0: raise ValueError(f"`self_attn_dim` must be positive, got {self_attn_dim}") if self_attn_heads <= 0: raise ValueError(f"`self_attn_heads` must be positive, got {self_attn_heads}") if self_attn_dim % self_attn_heads != 0: raise ValueError( "`self_attn_dim` must be divisible by `self_attn_heads`; " f"got {self_attn_dim} and {self_attn_heads}" ) if self_attn_layers <= 0: raise ValueError(f"`self_attn_layers` must be positive, got {self_attn_layers}") if layer_gate_mode not in VALID_LAYER_GATE_MODES: raise ValueError( f"`layer_gate_mode` must be one of {sorted(VALID_LAYER_GATE_MODES)}, got {layer_gate_mode!r}" ) if not 0.0 <= layer_gate_threshold <= 1.0: raise ValueError(f"`layer_gate_threshold` must be in [0, 1], got {layer_gate_threshold}") if not 0.0 <= layer_gate_strength <= 1.0: raise ValueError(f"`layer_gate_strength` must be in [0, 1], got {layer_gate_strength}") # 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 self.use_text_token_gate = use_text_token_gate self.text_token_gate_use_text_summary = text_token_gate_use_text_summary self.text_token_gate_use_vision_tokens = text_token_gate_use_vision_tokens self.text_token_gate_hidden_dim = text_token_gate_hidden_dim self.text_token_gate_text_pooling_mode = text_token_gate_text_pooling_mode self.text_token_gate_text_pool_hidden_dim = text_token_gate_text_pool_hidden_dim self.text_token_gate_cross_attention_dim = text_token_gate_cross_attention_dim self.text_token_gate_cross_attention_heads = text_token_gate_cross_attention_heads self.text_token_gate_mlp_depth = text_token_gate_mlp_depth self.contrastive_visual_tau = contrastive_visual_tau self.contrastive_text_tau = contrastive_text_tau self.gazing_mode = gazing_mode self.self_attn_dim = self_attn_dim self.self_attn_heads = self_attn_heads self.self_attn_layers = self_attn_layers self.layer_gate_mode = layer_gate_mode self.layer_gate_threshold = layer_gate_threshold self.layer_gate_strength = layer_gate_strength self.text_token_gate_budget = text_token_gate_budget self.text_token_gate_budget_loss_weight = text_token_gate_budget_loss_weight self.text_token_gate_linear_mean_penalty_weight = text_token_gate_linear_mean_penalty_weight self.text_token_gate_pool_instruction_only = text_token_gate_pool_instruction_only self.text_token_gate_filter_stopwords = text_token_gate_filter_stopwords # [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) 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, use_text_token_gate: bool = False, text_token_gate_use_text_summary: bool = True, text_token_gate_use_vision_tokens: bool = True, text_token_gate_hidden_dim: int = 512, text_token_gate_text_pooling_mode: str = "mean", text_token_gate_text_pool_hidden_dim: int = 128, text_token_gate_cross_attention_dim: int = 256, text_token_gate_cross_attention_heads: int = 1, text_token_gate_mlp_depth: int = 1, contrastive_visual_tau: float = 0.1, contrastive_text_tau: float = 1.0, gazing_mode: str = "mlp", self_attn_dim: int = 512, self_attn_heads: int = 8, self_attn_layers: int = 1, layer_gate_mode: str = "none", layer_gate_threshold: float = 0.15, layer_gate_strength: float = 0.5, text_token_gate_budget: float = 0.5, text_token_gate_budget_loss_weight: float = 0.01, text_token_gate_linear_mean_penalty_weight: float = 0.0, text_token_gate_pool_instruction_only: bool = False, text_token_gate_filter_stopwords: bool = False, **kwargs: str, ) -> None: self.norm_stats, self.n_action_bins = norm_stats, n_action_bins self.use_text_token_gate = use_text_token_gate self.text_token_gate_use_text_summary = text_token_gate_use_text_summary self.text_token_gate_use_vision_tokens = text_token_gate_use_vision_tokens self.text_token_gate_hidden_dim = text_token_gate_hidden_dim self.text_token_gate_text_pooling_mode = text_token_gate_text_pooling_mode self.text_token_gate_text_pool_hidden_dim = text_token_gate_text_pool_hidden_dim self.text_token_gate_cross_attention_dim = text_token_gate_cross_attention_dim self.text_token_gate_cross_attention_heads = text_token_gate_cross_attention_heads self.text_token_gate_mlp_depth = text_token_gate_mlp_depth self.contrastive_visual_tau = contrastive_visual_tau self.contrastive_text_tau = contrastive_text_tau self.gazing_mode = gazing_mode self.self_attn_dim = self_attn_dim self.self_attn_heads = self_attn_heads self.self_attn_layers = self_attn_layers self.layer_gate_mode = layer_gate_mode self.layer_gate_threshold = layer_gate_threshold self.layer_gate_strength = layer_gate_strength self.text_token_gate_budget = text_token_gate_budget self.text_token_gate_budget_loss_weight = text_token_gate_budget_loss_weight self.text_token_gate_linear_mean_penalty_weight = text_token_gate_linear_mean_penalty_weight self.text_token_gate_pool_instruction_only = text_token_gate_pool_instruction_only self.text_token_gate_filter_stopwords = text_token_gate_filter_stopwords super().__init__( use_text_token_gate=use_text_token_gate, text_token_gate_use_text_summary=text_token_gate_use_text_summary, text_token_gate_use_vision_tokens=text_token_gate_use_vision_tokens, text_token_gate_hidden_dim=text_token_gate_hidden_dim, text_token_gate_text_pooling_mode=text_token_gate_text_pooling_mode, text_token_gate_text_pool_hidden_dim=text_token_gate_text_pool_hidden_dim, text_token_gate_cross_attention_dim=text_token_gate_cross_attention_dim, text_token_gate_cross_attention_heads=text_token_gate_cross_attention_heads, text_token_gate_mlp_depth=text_token_gate_mlp_depth, contrastive_visual_tau=contrastive_visual_tau, contrastive_text_tau=contrastive_text_tau, gazing_mode=gazing_mode, self_attn_dim=self_attn_dim, self_attn_heads=self_attn_heads, self_attn_layers=self_attn_layers, layer_gate_mode=layer_gate_mode, layer_gate_threshold=layer_gate_threshold, layer_gate_strength=layer_gate_strength, text_token_gate_budget=text_token_gate_budget, text_token_gate_budget_loss_weight=text_token_gate_budget_loss_weight, text_token_gate_linear_mean_penalty_weight=text_token_gate_linear_mean_penalty_weight, text_token_gate_pool_instruction_only=text_token_gate_pool_instruction_only, text_token_gate_filter_stopwords=text_token_gate_filter_stopwords, **kwargs, )