configuration_prismatic.py

"""
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,
        )