add dependency

config.json

@@ -0,0 +1,27 @@
+{
+    "_commit_hash": null,
+    "architectures": [
+      "ViTaminCLIP"
+    ],
+    "auto_map": {
+      "AutoConfig": "configuration_vitamin.ViTaminConfig",
+      "AutoModel": "model.ViTaminCLIP"
+    },
+    "embed_dim": 1152,
+    "vision_cfg": {
+      "timm_model_name": "vitamin_xlarge_384",
+      "timm_model_pretrained": false,
+      "timm_pool": "",
+      "timm_proj": "linear",
+      "timm_drop": 0.0,
+      "timm_drop_path": 0.1,
+      "image_size": 384
+    },
+    "text_cfg": {
+      "context_length": 77,
+      "vocab_size": 49408,
+      "width": 1152,
+      "heads": 16,
+      "layers": 27
+  }
+}

configuration_vitamin.py

@@ -0,0 +1,158 @@
+""" ViTamin
+
+Paper: Designing Scalable Vison Models in the Vision-Language Era
+
+@misc{chen2023designing,
+      title={Designing Scalable Vison Models in the Vision-Language Era},
+      author={Jieneng Chen and Qihang Yu and Xiaohui Shen and Alan Yuille and Liang-Cheih Chen},
+      year={2023},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV}
+}
+
+Based on Apache 2.0 licensed code at https://github.com/Beckschen/ViTamin
+
+by Jieneng Chen 2024
+"""
+
+import copy
+import os
+from collections import OrderedDict
+from typing import TYPE_CHECKING, Any, Mapping, Optional, Union
+
+
+if TYPE_CHECKING:
+    from transformers.processing_utils import ProcessorMixin
+    from transformers.utils import TensorType
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+class ViTaminTextConfig(PretrainedConfig):
+    model_type = "vitamin_text_model"
+
+    def __init__(
+        self,
+        context_length = 77,
+        vocab_size = 49408,
+        width = 1024,
+        heads = 16,
+        layers = 24,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.vocab_size = vocab_size
+        self.context_length = context_length
+        self.width = width
+        self.heads = heads
+        self.layers = layers
+        
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs) -> "PretrainedConfig":
+        config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)
+
+        if 'text_config' in config_dict:
+            config_dict = config_dict['text_config']
+
+        if "model_type" in config_dict and hasattr(cls, "model_type") and config_dict["model_type"] != cls.model_type:
+            logger.warning(
+                f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
+                f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."
+            )
+
+        return cls.from_dict(config_dict, **kwargs)
+
+
+class ViTaminVisionConfig(PretrainedConfig):
+
+    model_type = "vitamin_vision_model"
+
+    def __init__(
+        self,
+        timm_model_name = "vitamin_large",
+        timm_model_pretrained = False,
+        timm_pool = "",
+        timm_proj = "linear",
+        timm_drop = 0.0,
+        timm_drop_path = 0.1,
+        image_size = 256,
+        timm_proj_bias = False,
+        patch_dropout = 0.0,
+        drop_path = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.timm_model_name = timm_model_name
+        self.timm_model_pretrained = timm_model_pretrained
+        self.timm_pool = timm_pool
+        self.timm_proj = timm_proj
+        self.timm_drop = timm_drop
+        self.timm_drop_path = timm_drop_path
+        self.timm_proj_bias = timm_proj_bias
+        self.patch_dropout = patch_dropout
+        self.image_size = image_size
+
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs) -> "PretrainedConfig":
+        config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)
+
+        if 'vision_config' in config_dict:
+            config_dict = config_dict['vision_config']
+
+        if "model_type" in config_dict and hasattr(cls, "model_type") and config_dict["model_type"] != cls.model_type:
+            logger.warning(
+                f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
+                f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."
+            )
+
+        return cls.from_dict(config_dict, **kwargs)
+
+
+
+class ViTaminConfig(PretrainedConfig):
+    model_type = "vitamin"
+    is_composition = True
+
+    def __init__(
+        self, text_config=None, vision_config=None, embed_dim=512,  **kwargs
+    ):
+        super().__init__(**kwargs)
+        if text_config is None:
+            text_config = {}
+            logger.info("`text_config` is `None`. Initializing the `CLIPTextConfig` with default values.")
+
+        if vision_config is None:
+            vision_config = {}
+            logger.info("`vision_config` is `None`. initializing the `CLIPVisionConfig` with default values.")
+        
+        self.embed_dim = embed_dim
+        self.text_config = ViTaminTextConfig(**text_config)
+        self.vision_config = ViTaminVisionConfig(**vision_config)
+        
+    @classmethod
+    def from_text_vision_configs(cls, text_config: ViTaminTextConfig, vision_config: ViTaminVisionConfig, **kwargs):
+        r"""
+        Instantiate a [`CLIPConfig`] (or a derived class) from clip text model configuration and clip vision model
+        configuration.
+        Returns:
+            [`CLIPConfig`]: An instance of a configuration object
+        """
+
+        return cls(text_config=text_config.to_dict(), vision_config=vision_config.to_dict(), **kwargs)
+
+    def to_dict(self):
+        """
+        Serializes this instance to a Python dictionary. Override the default [`~PretrainedConfig.to_dict`].
+        Returns:
+            `Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance,
+        """
+        output = copy.deepcopy(self.__dict__)
+        output["text_config"] = self.text_config.to_dict()
+        output["vision_config"] = self.vision_config.to_dict()
+        output["model_type"] = self.__class__.model_type
+        return output

model.py

@@ -0,0 +1,741 @@
+""" ViTamin
+
+Paper: Designing Scalable Vison Models in the Vision-Language Era
+
+@misc{chen2023designing,
+      title={Designing Scalable Vison Models in the Vision-Language Era},
+      author={Jieneng Chen and Qihang Yu and Xiaohui Shen and Alan Yuille and Liang-Cheih Chen},
+      year={2023},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV}
+}
+
+Based on Apache 2.0 licensed code at https://github.com/Beckschen/ViTamin
+
+by Jieneng Chen 2024
+
+Reference: https://github.com/openai/CLIP. Originally MIT License, Copyright (c) 2021 OpenAI.
+"""
+
+from dataclasses import dataclass
+import logging
+import math
+from typing import Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.utils.checkpoint import checkpoint
+from functools import partial
+from open_clip.hf_model import HFTextEncoder
+from open_clip.modified_resnet import ModifiedResNet
+from open_clip.transformer import LayerNormFp32, LayerNorm, QuickGELU, Attention, VisionTransformer, TextTransformer
+from open_clip.utils import to_2tuple
+import time
+import timm
+from timm.models.vision_transformer import _create_vision_transformer
+from .timm_model import TimmModel 
+from .vitamin import *
+# from .vitamin import HybridEmbed, MbConvStages, VitCfg, VitConvCfg
+from .vitamin import GeGluMlp, ViTamin, HybridEmbed, MbConvStages, VitCfg, VitConvCfg
+from transformers.modeling_utils import PreTrainedModel
+from .configuration_vitamin import ViTaminConfig, ViTaminVisionConfig
+
+@dataclass
+class CLIPVisionCfg:
+    layers: Union[Tuple[int, int, int, int], int] = 12
+    width: int = 768
+    head_width: int = 64
+    mlp_ratio: float = 4.0
+    patch_size: int = 16
+    image_size: Union[Tuple[int, int], int] = 224
+
+    ls_init_value: Optional[float] = None  
+    patch_dropout: float = 0.  
+    input_patchnorm: bool = False  
+    global_average_pool: bool = False 
+    attentional_pool: bool = False  
+    n_queries: int = 256  
+    attn_pooler_heads: int = 8  
+    output_tokens: bool = False
+
+    timm_model_name: str = None  
+    timm_model_pretrained: bool = False  
+    timm_pool: str = 'avg' 
+    timm_proj: str = 'linear'  
+    timm_proj_bias: bool = False 
+    timm_drop: float = 0. 
+    timm_drop_path: Optional[float] = None  
+
+
+@dataclass
+class CLIPTextCfg:
+    context_length: int = 77
+    vocab_size: int = 49408
+    width: int = 512
+    heads: int = 8
+    layers: int = 12
+    ls_init_value: Optional[float] = None  # layer scale initial value
+    hf_model_name: str = None
+    hf_tokenizer_name: str = None
+    hf_model_pretrained: bool = True
+    proj: str = 'mlp'
+    pooler_type: str = 'mean_pooler'
+    embed_cls: bool = False
+    pad_id: int = 0
+    output_tokens: bool = False
+    text_mask: str = 'first' # default first truncate in bpe_tokenizer
+
+
+def get_cast_dtype(precision: str):
+    cast_dtype = None
+    if precision == 'bf16':
+        cast_dtype = torch.bfloat16
+    elif precision == 'fp16':
+        cast_dtype = torch.float16
+    return cast_dtype
+
+
+def get_input_dtype(precision: str):
+    input_dtype = None
+    if precision in ('bf16', 'pure_bf16'):
+        input_dtype = torch.bfloat16
+    elif precision in ('fp16', 'pure_fp16'):
+        input_dtype = torch.float16
+    return input_dtype
+
+
+def _build_vision_tower(
+        embed_dim: int,
+        vision_cfg: CLIPVisionCfg,
+        quick_gelu: bool = False,
+        cast_dtype: Optional[torch.dtype] = None
+):
+    if isinstance(vision_cfg, dict):
+        vision_cfg = CLIPVisionCfg(**vision_cfg)
+
+    act_layer = QuickGELU if quick_gelu else nn.GELU
+
+    if vision_cfg.timm_model_name:
+        visual = TimmModel(
+            vision_cfg.timm_model_name,
+            pretrained=vision_cfg.timm_model_pretrained,
+            pool=vision_cfg.timm_pool,
+            proj=vision_cfg.timm_proj,
+            proj_bias=vision_cfg.timm_proj_bias,
+            drop=vision_cfg.timm_drop,
+            drop_path=vision_cfg.timm_drop_path,
+            patch_drop=vision_cfg.patch_dropout if vision_cfg.patch_dropout > 0 else None,
+            embed_dim=embed_dim,
+            image_size=vision_cfg.image_size,
+        )
+    elif isinstance(vision_cfg.layers, (tuple, list)):
+        vision_heads = vision_cfg.width * 32 // vision_cfg.head_width
+        visual = ModifiedResNet(
+            layers=vision_cfg.layers,
+            output_dim=embed_dim,
+            heads=vision_heads,
+            image_size=vision_cfg.image_size,
+            width=vision_cfg.width,
+        )
+    else:
+        vision_heads = vision_cfg.width // vision_cfg.head_width
+        norm_layer = LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm
+        visual = VisionTransformer(
+            image_size=vision_cfg.image_size,
+            patch_size=vision_cfg.patch_size,
+            width=vision_cfg.width,
+            layers=vision_cfg.layers,
+            heads=vision_heads,
+            mlp_ratio=vision_cfg.mlp_ratio,
+            ls_init_value=vision_cfg.ls_init_value,
+            patch_dropout=vision_cfg.patch_dropout,
+            input_patchnorm=vision_cfg.input_patchnorm,
+            global_average_pool=vision_cfg.global_average_pool,
+            attentional_pool=vision_cfg.attentional_pool,
+            n_queries=vision_cfg.n_queries,
+            attn_pooler_heads=vision_cfg.attn_pooler_heads,
+            output_tokens=vision_cfg.output_tokens,
+            output_dim=embed_dim,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+        )
+
+    return visual
+
+
+def _build_text_tower(
+        embed_dim: int,
+        text_cfg: CLIPTextCfg,
+        quick_gelu: bool = False,
+        cast_dtype: Optional[torch.dtype] = None,
+):
+    if isinstance(text_cfg, dict):
+        text_cfg = CLIPTextCfg(**text_cfg)
+
+    if text_cfg.hf_model_name:
+        text = HFTextEncoder(
+            text_cfg.hf_model_name,
+            output_dim=embed_dim,
+            proj=text_cfg.proj,
+            pooler_type=text_cfg.pooler_type,
+            pretrained=text_cfg.hf_model_pretrained,
+            output_tokens=text_cfg.output_tokens,
+        )
+    else:
+        act_layer = QuickGELU if quick_gelu else nn.GELU
+        norm_layer = LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm
+
+        text = TextTransformer(
+            context_length=text_cfg.context_length,
+            vocab_size=text_cfg.vocab_size,
+            width=text_cfg.width,
+            heads=text_cfg.heads,
+            layers=text_cfg.layers,
+            ls_init_value=text_cfg.ls_init_value,
+            output_dim=embed_dim,
+            embed_cls=text_cfg.embed_cls,
+            output_tokens=text_cfg.output_tokens,
+            pad_id=text_cfg.pad_id,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+        )
+    return text
+
+
+class CLIP(nn.Module):
+    output_dict: torch.jit.Final[bool]
+
+    def __init__(
+            self,
+            embed_dim: int,
+            vision_cfg: CLIPVisionCfg,
+            text_cfg: CLIPTextCfg,
+            quick_gelu: bool = False,
+            cast_dtype: Optional[torch.dtype] = None,
+            output_dict: bool = False,
+    ):
+        super().__init__()
+        self.output_dict = output_dict
+        self.visual = _build_vision_tower(embed_dim, vision_cfg, quick_gelu, cast_dtype)
+
+        text = _build_text_tower(embed_dim, text_cfg, quick_gelu, cast_dtype)
+        self.transformer = text.transformer
+        self.context_length = text.context_length
+        self.vocab_size = text.vocab_size
+        self.token_embedding = text.token_embedding
+        self.positional_embedding = text.positional_embedding
+
+        self.ln_final = text.ln_final
+        self.text_projection = text.text_projection
+        self.register_buffer('attn_mask', text.attn_mask, persistent=False)
+
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+        self.method_lock_text_tower = text.lock
+        self.text_no_grad = False
+
+    def lock_image_tower(self, unlocked_groups=0, freeze_bn_stats=False):
+        # lock image tower as per LiT - https://arxiv.org/abs/2111.07991
+        self.visual.lock(unlocked_groups=unlocked_groups, freeze_bn_stats=freeze_bn_stats)
+
+    def lock_text_tower(self, unlocked_layers: int = 0, freeze_layer_norm: bool = True, unlock_text_proj=False):
+        # added by jieneng
+        self.method_lock_text_tower(unlocked_layers, freeze_layer_norm)
+        self.text_no_grad = True
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True, enable_text=True):
+        self.visual.set_grad_checkpointing(enable)
+        self.transformer.grad_checkpointing = enable_text
+
+    def encode_image(self, image, normalize: bool = False):
+        features = self.visual(image)
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def encode_text(self, text, normalize: bool = False):
+        cast_dtype = self.transformer.get_cast_dtype()
+
+        x = self.token_embedding(text).to(cast_dtype)  # [batch_size, n_ctx, d_model]
+
+        x = x + self.positional_embedding.to(cast_dtype)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x, attn_mask=self.attn_mask)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.ln_final(x)  # [batch_size, n_ctx, transformer.width]
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
+        return F.normalize(x, dim=-1) if normalize else x
+
+    def forward(
+            self,
+            image: Optional[torch.Tensor] = None,
+            text: Optional[torch.Tensor] = None,
+    ):
+        # torch.cuda.synchronize()
+        image_features = self.encode_image(image, normalize=True) if image is not None else None
+
+        if self.text_no_grad:
+            with torch.no_grad():
+                text_features = self.encode_text(text, normalize=True).detach() if text is not None else None
+        else:
+            text_features = self.encode_text(text, normalize=True) if text is not None else None
+
+
+        if self.output_dict:
+            return {
+                "image_features": image_features,
+                "text_features": text_features,
+                "logit_scale": self.logit_scale.exp()
+            }
+        return image_features, text_features, self.logit_scale.exp()
+
+
+# class CustomTextCLIP(nn.Module):
+
+
+class CustomTextCLIP(nn.Module):
+    output_dict: torch.jit.Final[bool]
+
+    def __init__(
+            self,
+            embed_dim: int,
+            vision_cfg: CLIPVisionCfg,
+            text_cfg: CLIPTextCfg,
+            quick_gelu: bool = False,
+            cast_dtype: Optional[torch.dtype] = None,
+            output_dict: bool = False,
+    ):
+        super().__init__()
+        self.output_dict = output_dict
+        self.visual = _build_vision_tower(embed_dim, vision_cfg, quick_gelu, cast_dtype)
+        self.text = _build_text_tower(embed_dim, text_cfg, quick_gelu, cast_dtype)
+        self.context_length = self.text.context_length
+        self.vocab_size = self.text.vocab_size
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+        self.text_no_grad = False
+
+    def lock_image_tower(self, unlocked_groups=0, freeze_bn_stats=False):
+        # lock image tower as per LiT - https://arxiv.org/abs/2111.07991
+        self.visual.lock(unlocked_groups=unlocked_groups, freeze_bn_stats=freeze_bn_stats)
+
+    def lock_text_tower(self, unlocked_layers: int = 0, freeze_layer_norm: bool = True, unlock_text_proj = False):
+        self.text.lock(unlocked_layers, freeze_layer_norm, unlock_text_proj)
+        self.text_no_grad = True
+
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True, enable_text=True):
+        self.visual.set_grad_checkpointing(enable)
+        self.text.set_grad_checkpointing(enable_text)
+
+
+    def encode_image(self, image, normalize: bool = False):
+        features = self.visual(image)
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def encode_text(self, text, normalize: bool = False):
+        features = self.text(text)
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def forward(
+            self,
+            image: Optional[torch.Tensor] = None,
+            text: Optional[torch.Tensor] = None,
+    ):
+        image_features = self.encode_image(image, normalize=True) if image is not None else None
+        # if self.text_no_grad:
+        #     with torch.no_grad():
+        #         text_features = self.encode_text(text, normalize=True).detach() if text is not None else None
+        # else:
+        text_features = self.encode_text(text, normalize=True) if text is not None else None
+        
+        if self.output_dict:
+            return {
+                "image_features": image_features,
+                "text_features": text_features,
+                "logit_scale": self.logit_scale.exp()
+            }
+        return image_features, text_features, self.logit_scale.exp()
+
+
+class ViTaminPreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = ViTaminConfig
+    base_model_prefix = 'vitamin'
+
+
+# hack CLIPVisionModel for llava: https://github.com/huggingface/transformers/blob/9acce7de1cb8229304a467938ebb47727d60cdb2/src/transformers/models/clip/modeling_clip.py#L878
+class ViTaminVisionModel(PreTrainedModel):
+    config_class = ViTaminVisionConfig
+    main_input_name = 'pixel_values'
+
+    def __init__(self, config: ViTaminVisionConfig):
+        super().__init__(config)
+
+        self.visual = _build_vision_tower(config.embed_dim, config)
+        
+    def forward(
+            self,
+            pixel_values: Optional[torch.FloatTensor] = None,
+            select_layer = -2,
+    ):
+        assert len(pixel_values.shape) == 4, f'wrong pixel_values size: {pixel_values.shape}'
+        x = self.visual.trunk.patch_embed.backbone.stem(pixel_values)
+        x = self.visual.trunk.patch_embed.backbone.stages[0](x)
+        x = self.visual.trunk.patch_embed.backbone.stages[1](x)
+        x = self.visual.trunk.patch_embed.backbone.pool(x)
+        x = self.visual.trunk.patch_embed.proj(x)
+        x = x.flatten(2).transpose(1, 2)
+        x = self.visual.trunk.patch_drop(x)
+        x = self.visual.trunk.norm_pre(x)
+        x = self.visual.trunk.blocks[:select_layer+1](x)
+        return x
+
+
+class ViTaminCLIP(ViTaminPreTrainedModel):
+    output_dict: torch.jit.Final[bool]
+    config_class: ViTaminConfig
+
+    def __init__(
+            self,
+            config: ViTaminConfig
+    ):
+        super().__init__(config)
+
+        embed_dim=config.embed_dim #: int,
+        vision_cfg=config.vision_cfg #: CLIPVisionCfg,
+        text_cfg=config.text_cfg #: CLIPTextCfg,
+        quick_gelu=False
+        cast_dtype=None
+        output_dict=False
+        
+        self.config = config
+        self.output_dict = output_dict
+        self.visual = _build_vision_tower(embed_dim, vision_cfg, quick_gelu, cast_dtype)
+        self.text = _build_text_tower(embed_dim, text_cfg, quick_gelu, cast_dtype)
+        self.context_length = self.text.context_length
+        self.vocab_size = self.text.vocab_size
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+        self.text_no_grad = False
+
+    def forward_visual4llava(
+            self,
+            pixel_values: Optional[torch.FloatTensor] = None,
+            select_layer = -2,
+    ):
+        assert len(pixel_values.shape) == 4, f'wrong pixel_values size: {pixel_values.shape}'
+        x = self.visual.trunk.patch_embed.backbone.stem(pixel_values)
+        x = self.visual.trunk.patch_embed.backbone.stages[0](x)
+        x = self.visual.trunk.patch_embed.backbone.stages[1](x)
+        x = self.visual.trunk.patch_embed.backbone.pool(x)
+        x = self.visual.trunk.patch_embed.proj(x)
+        x = x.flatten(2).transpose(1, 2)
+        x = self.visual.trunk.patch_drop(x)
+        x = self.visual.trunk.norm_pre(x)
+        x = self.visual.trunk.blocks[:select_layer+1](x)
+        return x
+
+    def encode_image(self, image, normalize: bool = False):
+        features = self.visual(image)
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def encode_text(self, text, normalize: bool = False):
+        features = self.text(text)
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def forward_pixel(
+            self,
+            image: Optional[torch.Tensor] = None,
+            text: Optional[torch.Tensor] = None,
+    ):
+
+        x = self.visual.trunk.patch_embed.backbone.stem(image)
+        x = self.visual.trunk.patch_embed.backbone.stages[0](x)
+        x = self.visual.trunk.patch_embed.backbone.stages[1](x)
+        x = self.visual.trunk.patch_embed.backbone.pool(x)
+        x = self.visual.trunk.patch_embed.proj(x)
+        x = x.flatten(2).transpose(1, 2)
+        x = self.visual.trunk.patch_drop(x)
+        x = self.visual.trunk.norm_pre(x)
+        x = self.visual.trunk.blocks(x)
+        x = self.visual.trunk.fc_norm(x)
+        x = self.visual.head.proj(x)
+        image_features = F.normalize(x, dim=-1)
+        text_features = self.encode_text(text, normalize=True) if text is not None else None
+        
+        if self.output_dict:
+            return {
+                "image_features": image_features,
+                "text_features": text_features,
+                "logit_scale": self.logit_scale.exp()
+            }
+        return image_features, text_features, self.logit_scale.exp()
+
+    def forward(
+            self,
+            image: Optional[torch.Tensor] = None,
+            text: Optional[torch.Tensor] = None,
+    ):
+        image_features = self.encode_image(image, normalize=True) if image is not None else None
+        # if self.text_no_grad:
+        #     with torch.no_grad():
+        #         text_features = self.encode_text(text, normalize=True).detach() if text is not None else None
+        # else:
+        text_features = self.encode_text(text, normalize=True) if text is not None else None
+        
+        if self.output_dict:
+            return {
+                "image_features": image_features,
+                "text_features": text_features,
+                "logit_scale": self.logit_scale.exp()
+            }
+        return image_features, text_features, self.logit_scale.exp()
+
+def convert_weights_to_lp(model: nn.Module, dtype=torch.float16):
+    """Convert applicable model parameters to low-precision (bf16 or fp16)"""
+
+    def _convert_weights(l):
+        if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
+            l.weight.data = l.weight.data.to(dtype)
+            if l.bias is not None:
+                l.bias.data = l.bias.data.to(dtype)
+
+        if isinstance(l, (nn.MultiheadAttention, Attention)):
+            for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]:
+                tensor = getattr(l, attr)
+                if tensor is not None:
+                    tensor.data = tensor.data.to(dtype)
+
+        if isinstance(l, (CLIP, TextTransformer)):
+            # convert text nn.Parameter projections
+            attr = getattr(l, "text_projection", None)
+            if attr is not None:
+                attr.data = attr.data.to(dtype)
+
+        if isinstance(l, VisionTransformer):
+            # convert vision nn.Parameter projections
+            attr = getattr(l, "proj", None)
+            if attr is not None:
+                attr.data = attr.data.to(dtype)
+
+    model.apply(_convert_weights)
+
+
+convert_weights_to_fp16 = convert_weights_to_lp  # backwards compat
+
+
+# used to maintain checkpoint compatibility
+def convert_to_custom_text_state_dict(state_dict: dict):
+    if 'text_projection' in state_dict:
+        # old format state_dict, move text tower -> .text
+        new_state_dict = {}
+        for k, v in state_dict.items():
+            if any(k.startswith(p) for p in (
+                'text_projection',
+                'positional_embedding',
+                'token_embedding',
+                'transformer',
+                'ln_final',
+            )):
+                k = 'text.' + k
+            new_state_dict[k] = v
+        return new_state_dict
+    return state_dict
+
+
+def build_model_from_openai_state_dict(
+        state_dict: dict,
+        quick_gelu=True,
+        cast_dtype=torch.float16,
+):
+    vit = "visual.proj" in state_dict
+
+    if vit:
+        vision_width = state_dict["visual.conv1.weight"].shape[0]
+        vision_layers = len(
+            [k for k in state_dict.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")])
+        vision_patch_size = state_dict["visual.conv1.weight"].shape[-1]
+        grid_size = round((state_dict["visual.positional_embedding"].shape[0] - 1) ** 0.5)
+        image_size = vision_patch_size * grid_size
+    else:
+        counts: list = [
+            len(set(k.split(".")[2] for k in state_dict if k.startswith(f"visual.layer{b}"))) for b in [1, 2, 3, 4]]
+        vision_layers = tuple(counts)
+        vision_width = state_dict["visual.layer1.0.conv1.weight"].shape[0]
+        output_width = round((state_dict["visual.attnpool.positional_embedding"].shape[0] - 1) ** 0.5)
+        vision_patch_size = None
+        assert output_width ** 2 + 1 == state_dict["visual.attnpool.positional_embedding"].shape[0]
+        image_size = output_width * 32
+
+    embed_dim = state_dict["text_projection"].shape[1]
+    context_length = state_dict["positional_embedding"].shape[0]
+    vocab_size = state_dict["token_embedding.weight"].shape[0]
+    transformer_width = state_dict["ln_final.weight"].shape[0]
+    transformer_heads = transformer_width // 64
+    transformer_layers = len(set(k.split(".")[2] for k in state_dict if k.startswith(f"transformer.resblocks")))
+
+    vision_cfg = CLIPVisionCfg(
+        layers=vision_layers,
+        width=vision_width,
+        patch_size=vision_patch_size,
+        image_size=image_size,
+    )
+    text_cfg = CLIPTextCfg(
+        context_length=context_length,
+        vocab_size=vocab_size,
+        width=transformer_width,
+        heads=transformer_heads,
+        layers=transformer_layers,
+    )
+    model = CLIP(
+        embed_dim,
+        vision_cfg=vision_cfg,
+        text_cfg=text_cfg,
+        quick_gelu=quick_gelu,  # OpenAI models were trained with QuickGELU
+        cast_dtype=cast_dtype,
+    )
+
+    for key in ["input_resolution", "context_length", "vocab_size"]:
+        state_dict.pop(key, None)
+
+    convert_weights_to_fp16(model)  # OpenAI state dicts are partially converted to float16
+    model.load_state_dict(state_dict)
+    return model.eval()
+
+
+def trace_model(model, batch_size=256, device=torch.device('cpu')):
+    model.eval()
+    image_size = model.visual.image_size
+    example_images = torch.ones((batch_size, 3, image_size, image_size), device=device)
+    example_text = torch.zeros((batch_size, model.context_length), dtype=torch.int, device=device)
+    model = torch.jit.trace_module(
+        model,
+        inputs=dict(
+            forward=(example_images, example_text),
+            encode_text=(example_text,),
+            encode_image=(example_images,)
+        ))
+    model.visual.image_size = image_size
+    return model
+
+def resize_pos_embed_timm(state_dict, model, interpolation: str = 'bicubic', antialias: bool = True):
+    # Rescale the grid of position embeddings when loading from state_dict
+    old_pos_embed = state_dict.get('visual.trunk.pos_embed', None) # 1, 196, 1024]
+    if old_pos_embed is None:
+        return
+    
+    grid_size = to_2tuple(model.visual.trunk.patch_embed.grid_size)
+ 
+    
+    if hasattr(model.visual.trunk, 'cls_token') and model.visual.trunk.cls_token is not None:
+        return
+        # extra_tokens?
+        raise NotImplementedError
+
+    new_seq_len = grid_size[0] * grid_size[1]
+    if new_seq_len == old_pos_embed.shape[0]:
+        return
+
+    pos_emb_img = old_pos_embed
+    old_grid_size = to_2tuple(int(math.sqrt(len(pos_emb_img[0]))))
+    old_pos_emb_img = pos_emb_img
+    logging.info('Resizing position embedding grid-size from %s to %s', old_grid_size, grid_size) # Resizing position embedding grid-size from (1, 1) to (21, 21)
+    pos_emb_img = pos_emb_img.reshape(1, old_grid_size[0], old_grid_size[1], -1).permute(0, 3, 1, 2)
+
+    pos_emb_img = F.interpolate(
+        pos_emb_img,
+        size=grid_size,
+        mode=interpolation,
+        antialias=antialias,
+        align_corners=False,
+    )
+    pos_emb_img = pos_emb_img.permute(0, 2, 3, 1).reshape(1, grid_size[0] * grid_size[1], -1)
+    state_dict['visual.trunk.pos_embed'] = pos_emb_img
+
+def resize_pos_embed(state_dict, model, interpolation: str = 'bicubic', antialias: bool = True):
+    # Rescale the grid of position embeddings when loading from state_dict
+    pe_key_name = 'visual.positional_embedding'
+    old_pos_embed = state_dict.get('visual.positional_embedding', None)
+    if old_pos_embed is None:
+        pe_key_name = 'visual.trunk.pos_embed'
+        old_pos_embed = state_dict.get('visual.trunk.pos_embed', None) # 1, 196, 1024]
+        
+    if old_pos_embed is None:
+        return
+    
+    if hasattr(model.visual, 'grid_size'):
+        grid_size = to_2tuple(model.visual.grid_size)
+    elif hasattr(model.visual.trunk.patch_embed, 'grid_size'):
+        grid_size = to_2tuple(model.visual.trunk.patch_embed.grid_size)
+    else:
+        return
+    
+    if hasattr(model.visual.trunk, 'cls_token') and model.visual.trunk.cls_token is not None:
+        extra_tokens = 1  # FIXME detect different token configs (ie no class token, or more)
+    else:
+        extra_tokens = 0
+    new_seq_len = grid_size[0] * grid_size[1] + extra_tokens
+
+    if new_seq_len == old_pos_embed.shape[0]:
+        return
+
+    if extra_tokens:
+        pos_emb_tok, pos_emb_img = old_pos_embed[:extra_tokens], old_pos_embed[extra_tokens:]
+    else:
+        pos_emb_tok, pos_emb_img = None, old_pos_embed
+    old_grid_size = to_2tuple(int(math.sqrt(len(pos_emb_img))))
+    old_pos_emb_img = pos_emb_img
+    logging.info('Resizing position embedding grid-size from %s to %s', old_grid_size, grid_size) # Resizing position embedding grid-size from (1, 1) to (21, 21)
+    pos_emb_img = pos_emb_img.reshape(1, old_grid_size[0], old_grid_size[1], -1).permute(0, 3, 1, 2)
+
+
+    pos_emb_img = F.interpolate(
+        pos_emb_img,
+        size=grid_size,
+        mode=interpolation,
+        antialias=antialias,
+        align_corners=False,
+    )
+    pos_emb_img = pos_emb_img.permute(0, 2, 3, 1).reshape(1, grid_size[0] * grid_size[1], -1)[0]
+    if pos_emb_tok is not None:
+        new_pos_embed = torch.cat([pos_emb_tok, pos_emb_img], dim=0)
+    else:
+        new_pos_embed = pos_emb_img
+    state_dict[pe_key_name] = new_pos_embed
+
+def resize_text_pos_embed(state_dict, model, interpolation: str = 'linear', antialias: bool = False):
+    old_pos_embed = state_dict.get('positional_embedding', None)
+    if old_pos_embed is None:
+        return
+    # FIXME add support for text cls_token
+    model_pos_embed = getattr(model, 'positional_embedding', None)
+    if model_pos_embed is None:
+        model_pos_embed = getattr(model.text, 'positional_embedding', None)
+
+    old_num_pos = old_pos_embed.shape[0]
+    old_width = old_pos_embed.shape[1]
+    num_pos = model_pos_embed.shape[0]
+    width = model_pos_embed.shape[1]
+    assert old_width == width, 'text pos_embed width changed!'
+    if old_num_pos == num_pos:
+        return
+
+    logging.info('Resizing text position embedding num_pos from %s to %s', old_num_pos, num_pos)
+    old_pos_embed = old_pos_embed.reshape(1, old_num_pos, old_width).permute(0, 2, 1)
+    old_pos_embed = F.interpolate(
+        old_pos_embed,
+        size=num_pos,
+        mode=interpolation,
+        antialias=antialias,
+        align_corners=False,
+    )
+    old_pos_embed = old_pos_embed.permute(0, 2, 1)[0]
+    new_pos_embed = old_pos_embed
+
+    state_dict['positional_embedding'] = new_pos_embed

preprocessor_config.json

@@ -0,0 +1,20 @@
+{
+    "crop_size": 384,
+    "do_center_crop": true,
+    "do_normalize": true,
+    "do_resize": true,
+    "feature_extractor_type": "CLIPFeatureExtractor",
+    "image_mean": [
+        0.48145466,
+        0.4578275,
+        0.40821073
+    ],
+    "image_std": [
+        0.26862954,
+        0.26130258,
+        0.27577711
+    ],
+    "resample": 3,
+    "size": 384
+  }
+  

timm_model.py

@@ -0,0 +1,151 @@
+""" timm model adapter
+
+Wraps timm (https://github.com/rwightman/pytorch-image-models) models for use as a vision tower in CLIP model (OpenCLIP).
+"""
+import logging
+from collections import OrderedDict
+
+import torch, sys
+import torch.nn as nn
+import timm
+
+try:
+    import timm
+    from timm.models.layers import Mlp, to_2tuple
+    try:
+        # old timm imports < 0.8.1
+        from timm.models.layers.attention_pool2d import RotAttentionPool2d
+        from timm.models.layers.attention_pool2d import AttentionPool2d as AbsAttentionPool2d
+    except ImportError:
+        # new timm imports >= 0.8.1
+        from timm.layers import RotAttentionPool2d
+        from timm.layers import AttentionPool2d as AbsAttentionPool2d
+except ImportError:
+    timm = None
+from timm.models import create_model
+from open_clip.utils import freeze_batch_norm_2d
+
+from .vitamin import *
+
+class TimmModel(nn.Module):
+    """ timm model adapter
+    """
+
+    def __init__(
+            self,
+            model_name,
+            embed_dim,
+            image_size=224,
+            pool='avg',
+            proj='linear',
+            proj_bias=False,
+            drop=0.,
+            drop_path=None,
+            patch_drop=None,
+            pretrained=False,
+    ):
+        super().__init__()
+        if timm is None:
+            raise RuntimeError("Please `pip install timm` to use timm models.")
+        self.image_size = to_2tuple(image_size)
+
+        # setup kwargs that may not be common across all models
+        timm_kwargs = {}
+        if drop_path is not None:
+            timm_kwargs['drop_path_rate'] = drop_path
+        if patch_drop is not None:
+            timm_kwargs['patch_drop_rate'] = patch_drop
+
+        custom_pool = pool in ('abs_attn', 'rot_attn')
+        if not proj and not custom_pool:
+            # use network classifier head as projection if no proj specified and no custom pooling used
+            self.trunk = timm.create_model(
+                model_name,
+                num_classes=embed_dim,
+                global_pool=pool,
+                pretrained=pretrained,
+                **timm_kwargs,
+            )
+            prev_chs = embed_dim
+        else:
+            self.trunk = timm.create_model(
+                model_name,
+                pretrained=pretrained,
+                **timm_kwargs,
+            )
+            feat_size = self.trunk.default_cfg.get('pool_size', None)
+            feature_ndim = 1 if not feat_size else 2
+            if custom_pool:
+                assert feature_ndim == 2
+                # if attn pooling used, remove both classifier and default pool
+                self.trunk.reset_classifier(0, global_pool='')
+            else:
+                # reset global pool if pool config set, otherwise leave as network default
+                reset_kwargs = dict(global_pool=pool) if pool else {}
+                self.trunk.reset_classifier(0, **reset_kwargs)
+            prev_chs = self.trunk.num_features
+
+        head_layers = OrderedDict()
+
+        # Add custom pooling to head
+        if pool == 'abs_attn':
+            head_layers['pool'] = AbsAttentionPool2d(prev_chs, feat_size=feat_size, out_features=embed_dim)
+            prev_chs = embed_dim
+        elif pool == 'rot_attn':
+            head_layers['pool'] = RotAttentionPool2d(prev_chs, out_features=embed_dim)
+            prev_chs = embed_dim
+
+        # NOTE attention pool ends with a projection layer, so proj should usually be set to '' if such pooling is used
+        if proj == 'linear':
+            head_layers['drop'] = nn.Dropout(drop)
+            head_layers['proj'] = nn.Linear(prev_chs, embed_dim, bias=proj_bias)
+        elif proj == 'mlp':
+            head_layers['mlp'] = Mlp(prev_chs, 2 * embed_dim, embed_dim, drop=(drop, 0), bias=(True, proj_bias))
+        else:
+            assert not proj, f'Unknown projection type {proj}.'
+
+        self.head = nn.Sequential(head_layers)
+
+    def lock(self, unlocked_groups=0, freeze_bn_stats=False):
+        """ lock modules
+        Args:
+            unlocked_groups (int): leave last n layer groups unlocked (default: 0)
+        """
+        if not unlocked_groups:
+            # lock full model
+            for param in self.trunk.parameters():
+                param.requires_grad = False
+            if freeze_bn_stats:
+                freeze_batch_norm_2d(self.trunk)
+        else:
+            # NOTE: partial freeze requires latest timm (master) branch and is subject to change
+            try:
+                # FIXME import here until API stable and in an official release
+                from timm.models.helpers import group_parameters, group_modules
+            except ImportError:
+                raise RuntimeError(
+                    'Please install latest timm `pip install git+https://github.com/rwightman/pytorch-image-models`')
+            matcher = self.trunk.group_matcher()
+            gparams = group_parameters(self.trunk, matcher)
+            max_layer_id = max(gparams.keys())
+            max_layer_id = max_layer_id - unlocked_groups
+            for group_idx in range(max_layer_id + 1):
+                group = gparams[group_idx]
+                for param in group:
+                    self.trunk.get_parameter(param).requires_grad = False
+            if freeze_bn_stats:
+                gmodules = group_modules(self.trunk, matcher, reverse=True)
+                gmodules = {k for k, v in gmodules.items() if v <= max_layer_id}
+                freeze_batch_norm_2d(self.trunk, gmodules)
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        try:
+            self.trunk.set_grad_checkpointing(enable)
+        except Exception as e:
+            logging.warning('grad checkpointing not supported for this timm image tower, continuing without...')
+
+    def forward(self, x):
+        x = self.trunk(x)
+        x = self.head(x)
+        return x

vitamin.py

@@ -0,0 +1,780 @@
+""" ViTamin
+
+Paper: Designing Scalable Vison Models in the Vision-Language Era
+
+@misc{chen2023designing,
+      title={Designing Scalable Vison Models in the Vision-Language Era},
+      author={Jieneng Chen and Qihang Yu and Xiaohui Shen and Alan Yuille and Liang-Cheih Chen},
+      year={2023},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV}
+}
+
+Based on Apache 2.0 licensed code at https://github.com/ViTamin/ViTamin
+
+Modifications and timm support by Jieneng Chen 2023
+
+Adapted from timm codebase, thanks!
+"""
+
+from functools import partial
+from typing import List, Tuple
+from dataclasses import dataclass, replace
+from typing import Callable, Optional, Union, Tuple, List, Sequence
+import math, time
+from torch.jit import Final
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import timm
+from torch.utils.checkpoint import checkpoint
+from timm.models.layers import create_attn, get_norm_layer, get_norm_act_layer, create_conv2d, make_divisible, trunc_normal_tf_
+
+
+from timm.layers import to_2tuple, DropPath, Format
+from timm.layers.norm_act import _create_act
+from timm.models._registry import register_model
+from timm.models._manipulate import named_apply, checkpoint_seq
+from timm.models._builder import build_model_with_cfg
+from timm.models.vision_transformer import get_act_layer, Type, LayerType, Mlp, Block, PatchEmbed, VisionTransformer, checkpoint_filter_fn, get_init_weights_vit, init_weights_vit_timm, _load_weights 
+import logging
+from collections import OrderedDict
+
+
+
+@dataclass
+class VitConvCfg:
+    expand_ratio: float = 4.0
+    expand_output: bool = True  # calculate expansion channels from output (vs input chs)
+    kernel_size: int = 3
+    group_size: int = 1  # 1 == depthwise
+    pre_norm_act: bool = False  # activation after pre-norm
+    stride_mode: str = 'dw'  # stride done via one of 'pool', '1x1', 'dw'
+    pool_type: str = 'avg2'
+    downsample_pool_type: str = 'avg2'
+    act_layer: str = 'gelu' # stem & stage 1234
+    norm_layer: str = ''
+    norm_layer_cl: str = ''
+    norm_eps: Optional[float] = None
+    down_shortcut: Optional[bool] = True
+    mlp: str = 'mlp'
+
+    def __post_init__(self):
+        use_mbconv = True
+        if not self.norm_layer:
+            self.norm_layer = 'batchnorm2d' if use_mbconv else 'layernorm2d'
+        if not self.norm_layer_cl and not use_mbconv:
+            self.norm_layer_cl = 'layernorm'
+        if self.norm_eps is None:
+            self.norm_eps = 1e-5 if use_mbconv else 1e-6
+        self.downsample_pool_type = self.downsample_pool_type or self.pool_type
+
+@dataclass
+class VitCfg:
+    embed_dim: Tuple[Union[int, Tuple[int, ...]], ...] = (96, 192, 384, 768)
+    depths: Tuple[Union[int, Tuple[int, ...]], ...] = (2, 3, 5, 2)
+    stem_width: int = 64
+    conv_cfg: VitConvCfg = VitConvCfg()
+    weight_init: str = 'vit_eff'
+    head_type: str = ""
+    stem_type: str = "stem"
+
+def _init_conv(module, name, scheme=''):
+    if isinstance(module, nn.Conv2d):
+        fan_out = module.kernel_size[0] * module.kernel_size[1] * module.out_channels
+        fan_out //= module.groups
+        nn.init.normal_(module.weight, 0, math.sqrt(2.0 / fan_out))
+        if module.bias is not None:
+            nn.init.zeros_(module.bias)
+
+class Stem(nn.Module):
+    def __init__(
+            self,
+            in_chs: int,
+            out_chs: int,
+            act_layer: str = 'gelu',
+            norm_layer: str = 'layernorm2d',
+            norm_eps: float = 1e-6,
+            bias: bool = True,
+    ):
+        super().__init__()
+        self.grad_checkpointing=False
+        norm_act_layer = partial(get_norm_act_layer(norm_layer, act_layer), eps=norm_eps)
+        self.out_chs = out_chs
+        self.conv1 = create_conv2d(in_chs, out_chs, 3, stride=2, bias=bias)
+        self.norm1 = norm_act_layer(out_chs)
+        self.conv2 = create_conv2d(out_chs, out_chs, 3, stride=1, bias=bias)
+        named_apply(_init_conv, self)
+
+    def forward(self, x):
+        if self.grad_checkpointing:
+            x = checkpoint(self.conv1, x)
+            x = self.norm1(x)
+            x = checkpoint(self.conv2, x)
+        else:
+            x = self.conv1(x)
+            x = self.norm1(x)
+            x = self.conv2(x)
+
+        return x
+
+class Downsample2d(nn.Module):
+    def __init__(
+            self,
+            dim: int,
+            dim_out: int,
+            bias: bool = True,
+    ):
+        super().__init__()
+        self.pool = nn.AvgPool2d(kernel_size=3, stride=2, padding=1, count_include_pad=False)
+        if dim != dim_out:
+            self.expand = nn.Conv2d(dim, dim_out, 1, bias=bias) # 1x1 conv
+        else:
+            self.expand = nn.Identity()
+
+    def forward(self, x):
+        x = self.pool(x)
+        x = self.expand(x)
+        return x
+
+
+class StridedConv(nn.Module):
+    """ downsample 2d as well
+    """
+    def __init__(
+            self, 
+            kernel_size=3, 
+            stride=2, 
+            padding=1,
+            in_chans=3, 
+            embed_dim=768, 
+    ):
+        super().__init__()
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding)
+        norm_layer = partial(get_norm_layer('layernorm2d'), eps=1e-6)
+        self.norm = norm_layer(in_chans) 
+
+    def forward(self, x):
+        x = self.norm(x) 
+        x = self.proj(x)
+        return x
+
+
+class MbConvLNBlock(nn.Module):
+    def __init__(
+            self,
+            in_chs: int,
+            out_chs: int,
+            stride: int = 1,
+            drop_path: float = 0.,
+            kernel_size: int = 3,
+            norm_layer: str = 'layernorm2d',
+            norm_eps: float = 1e-6,
+            act_layer: str = 'gelu',
+            expand_ratio: float = 4.0,
+    ):
+        super(MbConvLNBlock, self).__init__()
+        self.stride, self.in_chs, self.out_chs = stride, in_chs, out_chs
+        mid_chs = make_divisible(out_chs * expand_ratio)
+        prenorm_act_layer = partial(get_norm_act_layer(norm_layer, act_layer), eps=norm_eps)
+
+        if stride == 2: 
+            self.shortcut = Downsample2d(in_chs, out_chs, bias=True)
+        elif in_chs != out_chs:
+            self.shortcut = nn.Conv2d(in_chs, out_chs, 1, bias=True)
+        else:
+            self.shortcut = nn.Identity()
+
+        self.pre_norm = prenorm_act_layer(in_chs, apply_act=False)
+        self.down = nn.Identity()
+        self.conv1_1x1 = create_conv2d(in_chs, mid_chs, 1, stride=1, bias=True)
+        self.act1 = _create_act(act_layer, inplace=True)
+        self.act2 = _create_act(act_layer, inplace=True)
+
+        self.conv2_kxk = create_conv2d(mid_chs, mid_chs, kernel_size, stride=stride, dilation=1, groups=mid_chs, bias=True)
+        self.conv3_1x1 = create_conv2d(mid_chs, out_chs, 1, bias=True)
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+
+    def init_weights(self, scheme=''):
+        named_apply(partial(_init_conv, scheme=scheme), self)
+
+    def forward(self, x):
+        shortcut = self.shortcut(x)
+
+        x = self.pre_norm(x)
+        x = self.down(x) # nn.Identity() 
+
+        # 1x1 expansion conv & act
+        x = self.conv1_1x1(x)
+        x = self.act1(x) 
+
+        # (strided) depthwise 3x3 conv & act
+        x = self.conv2_kxk(x)
+        x = self.act2(x)
+
+        # 1x1 linear projection to output width
+        x = self.conv3_1x1(x)
+        x = self.drop_path(x) + shortcut
+
+        return x
+
+
+class MbConvStages(nn.Module):
+    """ stage 1 and stage 2 of ViTamin: MBConv-LN blocks
+    """
+    def __init__(
+            self,
+            cfg: VitCfg,
+            img_size: Union[int, Tuple[int, int]] = 224, # place holder
+            in_chans: int = 3,
+    ):
+        super().__init__()
+        self.grad_checkpointing = False
+        self.stem = Stem(
+            in_chs=in_chans,
+            out_chs=cfg.stem_width,
+        )
+        stages = []
+        self.num_stages = len(cfg.embed_dim)
+        for s, dim in enumerate(cfg.embed_dim[:2]):
+            blocks = []
+            stage_in_chs = cfg.embed_dim[s-1] if s>0 else cfg.stem_width
+            for d in range(cfg.depths[s]):
+                blocks += [MbConvLNBlock(
+                        in_chs = stage_in_chs if d==0 else dim,
+                        out_chs = dim,
+                        stride = 2 if d == 0 else 1,
+                    )]
+            blocks = nn.Sequential(*blocks)
+            stages += [blocks]
+
+        self.stages = nn.ModuleList(stages)
+        self.pool = StridedConv(
+                        stride=2,
+                        in_chans=cfg.embed_dim[1],
+                        embed_dim=cfg.embed_dim[2]
+                    )
+
+    def forward(self, x):
+        x = self.stem(x)
+        if self.grad_checkpointing and not torch.jit.is_scripting():
+            for stage in self.stages:
+                x = checkpoint_seq(stage, x)
+            x = checkpoint(self.pool, x)
+        else:
+            for stage in self.stages:
+                x = stage(x)
+            x = self.pool(x)
+        
+        return x
+
+class GeGluMlp(nn.Module):
+    def __init__(
+            self, 
+            in_features, 
+            hidden_features,
+            act_layer = None,
+            drop = 0.0,
+    ):
+        super().__init__()
+        norm_layer = partial(get_norm_layer('layernorm'), eps=1e-6)
+        self.norm = norm_layer(in_features)
+        self.act = nn.GELU()
+        self.w0 = nn.Linear(in_features, hidden_features)
+        self.w1 = nn.Linear(in_features, hidden_features)
+        self.w2 = nn.Linear(hidden_features, in_features)
+
+    def forward(self, x):
+        x = self.norm(x)
+        x = self.act(self.w0(x)) * self.w1(x)
+        x = self.w2(x)
+        return x
+
+class HybridEmbed(nn.Module):
+    """ 
+    Extract feature map from stage 1-2, flatten, project to embedding dim.
+    """
+    def __init__(
+            self,
+            backbone,
+            img_size=224,
+            patch_size=1,
+            feature_size=None,
+            in_chans=3,
+            embed_dim=1024,
+            bias=True,
+            dynamic_img_pad=False,
+    ):
+        super().__init__()
+        assert isinstance(backbone, nn.Module)
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.backbone = backbone
+        if feature_size is None:
+            feature_size = img_size[0] // 16
+        feature_size = to_2tuple(feature_size)
+        if hasattr(self.backbone, 'feature_info'):
+            feature_dim = self.backbone.feature_info.channels()[-1]
+        elif hasattr(self.backbone, 'num_features'):
+            feature_dim = self.backbone.num_features
+        else:
+            feature_dim = embed_dim
+        assert feature_size[0] % patch_size[0] == 0 and feature_size[1] % patch_size[1] == 0
+        self.grid_size = (feature_size[0] // patch_size[0], feature_size[1] // patch_size[1])
+        self.num_patches = self.grid_size[0] * self.grid_size[1]
+        self.proj = nn.Identity()
+
+    def forward(self, x):
+        x = self.backbone(x)
+        if isinstance(x, (list, tuple)):
+            x = x[-1]  # last feature if backbone outputs list/tuple of features
+        x = self.proj(x)
+        x = x.flatten(2).transpose(1, 2)
+        return x
+
+class ViTamin(nn.Module):
+    """ hack timm VisionTransformer
+    """
+    dynamic_img_size: Final[bool]
+
+    def __init__(
+            self,
+            img_size: Union[int, Tuple[int, int]] = 224,
+            patch_size: Union[int, Tuple[int, int]] = 16,
+            in_chans: int = 3,
+            num_classes: int = 1000,
+            global_pool = 'token',
+            embed_dim: int = 768,
+            depth: int = 12,
+            num_heads: int = 12,
+            mlp_ratio: float = 4.,
+            qkv_bias: bool = True,
+            qk_norm: bool = False,
+            init_values: Optional[float] = None,
+            class_token: bool = True,
+            no_embed_class: bool = False,
+            reg_tokens: int = 0,
+            pre_norm: bool = False,
+            fc_norm: Optional[bool] = None,
+            dynamic_img_size: bool = False,
+            dynamic_img_pad: bool = False,
+            drop_rate: float = 0.,
+            pos_drop_rate: float = 0.,
+            patch_drop_rate: float = 0.,
+            proj_drop_rate: float = 0.,
+            attn_drop_rate: float = 0.,
+            drop_path_rate: float = 0.,
+            weight_init = '',
+            fix_init: bool = False,
+            embed_layer: Callable = PatchEmbed,
+            norm_layer: Optional[LayerType] = None,
+            act_layer: Optional[LayerType] = None,
+            block_fn: Type[nn.Module] = Block,
+            mlp_layer: Type[nn.Module] = Mlp,
+            is_pos_embed: bool = True
+    ) -> None:
+        """
+        Args:
+            img_size: Input image size.
+            patch_size: Patch size.
+            in_chans: Number of image input channels.
+            num_classes: Mumber of classes for classification head.
+            global_pool: Type of global pooling for final sequence (default: 'token').
+            embed_dim: Transformer embedding dimension.
+            depth: Depth of transformer.
+            num_heads: Number of attention heads.
+            mlp_ratio: Ratio of mlp hidden dim to embedding dim.
+            qkv_bias: Enable bias for qkv projections if True.
+            init_values: Layer-scale init values (layer-scale enabled if not None).
+            class_token: Use class token.
+            no_embed_class: Don't include position embeddings for class (or reg) tokens.
+            reg_tokens: Number of register tokens.
+            fc_norm: Pre head norm after pool (instead of before), if None, enabled when global_pool == 'avg'.
+            drop_rate: Head dropout rate.
+            pos_drop_rate: Position embedding dropout rate.
+            attn_drop_rate: Attention dropout rate.
+            drop_path_rate: Stochastic depth rate.
+            weight_init: Weight initialization scheme.
+            fix_init: Apply weight initialization fix (scaling w/ layer index).
+            embed_layer: Patch embedding layer.
+            norm_layer: Normalization layer.
+            act_layer: MLP activation layer.
+            block_fn: Transformer block layer.
+        """
+        super().__init__()
+        assert global_pool in ('', 'avg', 'token', 'map')
+        assert class_token or global_pool != 'token'
+        use_fc_norm = global_pool == 'avg' if fc_norm is None else fc_norm
+        norm_layer = get_norm_layer(norm_layer) or partial(nn.LayerNorm, eps=1e-6)
+        act_layer = get_act_layer(act_layer) or nn.GELU
+
+        self.num_classes = num_classes
+        self.global_pool = global_pool
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+        self.num_prefix_tokens = 1 if class_token else 0
+        self.num_prefix_tokens += reg_tokens
+        self.num_reg_tokens = reg_tokens
+        self.has_class_token = class_token
+        self.no_embed_class = no_embed_class  # don't embed prefix positions (includes reg)
+        self.dynamic_img_size = dynamic_img_size
+        self.grad_checkpointing = False
+        self.is_pos_embed = is_pos_embed
+        embed_args = {}
+        if dynamic_img_size:
+            # flatten deferred until after pos embed
+            embed_args.update(dict(strict_img_size=False, output_fmt='NHWC'))
+
+        # stage_1_2 = MbConvStages(cfg=VitCfg(
+        #     embed_dim=(160, 320, 1024),
+        #     depths=(2, 4, 1),
+        #     stem_width=160,
+        #     conv_cfg = VitConvCfg(
+        #         norm_layer='layernorm2d',
+        #         norm_eps=1e-6,
+        #     ),
+        #     head_type='1d',
+        #     ), 
+        # )
+        # self.patch_embed = HybridEmbed(
+        #     stage_1_2,
+        #     img_size=img_size,
+        #     patch_size=1,
+        #     in_chans=in_chans,
+        #     embed_dim=embed_dim,
+        #     bias=not pre_norm,
+        #     dynamic_img_pad=dynamic_img_pad,
+        #     **embed_args,)
+        self.patch_embed = embed_layer(
+            img_size=img_size,
+            patch_size=patch_size,
+            in_chans=in_chans,
+            embed_dim=embed_dim,
+            bias=not pre_norm,  # disable bias if pre-norm is used (e.g. CLIP)
+        )
+        
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) if class_token else None
+        self.reg_token = nn.Parameter(torch.zeros(1, reg_tokens, embed_dim)) if reg_tokens else None
+        
+        if self.is_pos_embed:
+            embed_len = num_patches if no_embed_class else num_patches + self.num_prefix_tokens
+            self.pos_embed = nn.Parameter(torch.randn(1, embed_len, embed_dim) * .02)
+        else:
+            self.pos_embed = None
+        
+        self.pos_drop = nn.Dropout(p=pos_drop_rate)
+        if patch_drop_rate > 0:
+            self.patch_drop = PatchDropout(
+                patch_drop_rate,
+                num_prefix_tokens=self.num_prefix_tokens,
+            )
+        else:
+            self.patch_drop = nn.Identity()
+        self.norm_pre = norm_layer(embed_dim) if pre_norm else nn.Identity()
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        self.blocks = nn.Sequential(*[
+            block_fn(
+                dim=embed_dim,
+                num_heads=num_heads,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                qk_norm=qk_norm,
+                init_values=init_values,
+                proj_drop=proj_drop_rate,
+                attn_drop=attn_drop_rate,
+                drop_path=dpr[i],
+                norm_layer=norm_layer,
+                act_layer=act_layer,
+                mlp_layer=mlp_layer,
+            )
+            for i in range(depth)])
+        self.norm = norm_layer(embed_dim) if not use_fc_norm else nn.Identity()
+
+        # Classifier Head
+        if global_pool == 'map':
+            self.attn_pool = AttentionPoolLatent(
+                self.embed_dim,
+                num_heads=num_heads,
+                mlp_ratio=mlp_ratio,
+                norm_layer=norm_layer,
+            )
+        else:
+            self.attn_pool = None
+
+        self.fc_norm = norm_layer(embed_dim) if use_fc_norm else nn.Identity()
+        self.head_drop = nn.Dropout(drop_rate)
+        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+        if weight_init != 'skip':
+            self.init_weights(weight_init)
+        if fix_init:
+            self.fix_init_weight()
+
+    def init_weights(self, mode=''):
+        assert mode in ('jax', 'jax_nlhb', 'moco', '')
+        head_bias = -math.log(self.num_classes) if 'nlhb' in mode else 0.
+        if self.is_pos_embed:
+            trunc_normal_(self.pos_embed, std=.02)
+        if self.cls_token is not None:
+            nn.init.normal_(self.cls_token, std=1e-6)
+        named_apply(get_init_weights_vit(mode, head_bias), self)
+
+    def _init_weights(self, m):
+        # this fn left here for compat with downstream users
+        init_weights_vit_timm(m)
+
+    @torch.jit.ignore()
+    def load_pretrained(self, checkpoint_path, prefix=''):
+        _load_weights(self, checkpoint_path, prefix)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        if self.is_pos_embed:
+            return {'pos_embed', 'cls_token', 'dist_token'}
+        else:
+            return {'cls_token', 'dist_token'}
+
+    @torch.jit.ignore
+    def group_matcher(self, coarse=False):
+        return dict(
+            stem=r'^cls_token|pos_embed|patch_embed',  # stem and embed
+            blocks=[(r'^blocks\.(\d+)', None), (r'^norm', (99999,))]
+        )
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.grad_checkpointing = enable
+        self.patch_embed.backbone.stem.grad_checkpointing = enable # disable https://blog.csdn.net/lhx526080338/article/details/127894671?utm_medium=distribute.pc_relevant.none-task-blog-2~default~baidujs_baidulandingword~default-1-127894671-blog-125562110.235^v38^pc_relevant_anti_t3_base&spm=1001.2101.3001.4242.2&utm_relevant_index=4
+        self.patch_embed.backbone.grad_checkpointing = enable
+
+    @torch.jit.ignore
+    def get_classifier(self):
+        return self.head
+
+    def reset_classifier(self, num_classes: int, global_pool=None):
+        self.num_classes = num_classes
+        if global_pool is not None:
+            assert global_pool in ('', 'avg', 'token')
+            self.global_pool = global_pool
+        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.patch_embed(x)
+        if self.is_pos_embed:
+            x = self._pos_embed(x)
+        x = self.patch_drop(x)
+        x = self.norm_pre(x)
+        if self.grad_checkpointing and not torch.jit.is_scripting():
+            x = checkpoint_seq(self.blocks, x)
+        else:
+            x = self.blocks(x)
+        x = self.norm(x)
+        return x
+
+    def forward_head(self, x: torch.Tensor, pre_logits: bool = False) -> torch.Tensor:
+        if self.attn_pool is not None:
+            x = self.attn_pool(x)
+        elif self.global_pool == 'avg':
+            x = x[:, self.num_prefix_tokens:].mean(dim=1)
+        elif self.global_pool:
+            x = x[:, 0]  # class token
+        x = self.fc_norm(x)
+        x = self.head_drop(x)
+        return x if pre_logits else self.head(x)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.forward_features(x)
+        x = self.forward_head(x)
+        return x
+
+def _create_vision_transformer(variant, pretrained=False, **kwargs):
+    if kwargs.get('features_only', None):
+        raise RuntimeError('features_only not implemented for Vision Transformer models.')
+
+    return build_model_with_cfg(
+        ViTamin, # ViTamin
+        variant,
+        pretrained,
+        pretrained_filter_fn=checkpoint_filter_fn,
+        **kwargs,
+    )
+
+
+def _create_vision_transformer_hybrid(variant, backbone, pretrained=False, **kwargs):
+    embed_layer = partial(HybridEmbed, backbone=backbone)
+    kwargs.setdefault('patch_size', 1)  # default patch size for hybrid models if not set
+    return _create_vision_transformer(variant, pretrained=pretrained, embed_layer=embed_layer, **kwargs)
+
+
+@register_model
+def vitamin_small(pretrained=False, **kwargs) -> VisionTransformer:
+    stage_1_2 = MbConvStages(cfg=VitCfg(
+            embed_dim=(64, 128, 384),
+            depths=(2, 4, 1),
+            stem_width=64,
+            conv_cfg = VitConvCfg(
+                norm_layer='layernorm2d',
+                norm_eps=1e-6,
+            ),
+            head_type='1d',
+        ),
+    )
+    stage3_args = dict(embed_dim=384, depth=14, num_heads=6, mlp_layer=GeGluMlp, mlp_ratio=2., class_token=False, global_pool='avg')
+    model = _create_vision_transformer_hybrid('vitamin_small', backbone=stage_1_2, pretrained=pretrained, **dict(stage3_args, **kwargs))
+    return model
+
+
+@register_model
+def vitamin_base(pretrained=False, **kwargs) -> VisionTransformer:
+    stage_1_2 = MbConvStages(cfg=VitCfg(
+            embed_dim=(128, 256, 768),
+            depths=(2, 4, 1),
+            stem_width=128,
+            conv_cfg = VitConvCfg(
+                norm_layer='layernorm2d',
+                norm_eps=1e-6,
+            ),
+            head_type='1d',
+        ),
+    )
+    stage3_args = dict(embed_dim=768, depth=14, num_heads=12, mlp_layer=GeGluMlp, mlp_ratio=2., class_token=False, global_pool='avg')
+    model = _create_vision_transformer_hybrid('vitamin_base', backbone=stage_1_2, pretrained=pretrained, **dict(stage3_args, **kwargs))
+    return model
+
+
+@register_model
+def vitamin_large(pretrained=False, **kwargs) -> VisionTransformer:
+    stage_1_2 = MbConvStages(cfg=VitCfg(
+        embed_dim=(160, 320, 1024),
+        depths=(2, 4, 1),
+        stem_width=160,
+        conv_cfg = VitConvCfg(
+            norm_layer='layernorm2d',
+            norm_eps=1e-6,
+        ),
+        head_type='1d',
+    ), 
+    )
+    stage3_args = dict(embed_dim=1024, depth=31, num_heads=16, mlp_layer=GeGluMlp, mlp_ratio=2., class_token=False, global_pool='avg')
+    model = _create_vision_transformer_hybrid(
+        'vitamin_large', backbone=stage_1_2, pretrained=pretrained, **dict(stage3_args, **kwargs))
+    return model
+
+
+@register_model
+def vitamin_large_256(pretrained=False, **kwargs) -> VisionTransformer:
+    backbone = MbConvStages(cfg=VitCfg(
+        embed_dim=(160, 320, 1024),
+        depths=(2, 4, 1),
+        stem_width=160,
+        conv_cfg = VitConvCfg(
+            norm_layer='layernorm2d',
+            norm_eps=1e-6,
+        ),
+        head_type='1d',
+    ), 
+    )
+    model_args = dict(img_size=256, embed_dim=1024, depth=31, num_heads=16, mlp_layer=GeGluMlp, mlp_ratio=2., class_token=False, global_pool='avg')
+    model = _create_vision_transformer_hybrid(
+        'vitamin_large_256', backbone=backbone, pretrained=pretrained, **dict(model_args, **kwargs))
+    return model
+
+
+@register_model
+def vitamin_large_336(pretrained=False, **kwargs) -> VisionTransformer:
+    backbone = MbConvStages(cfg=VitCfg(
+        embed_dim=(160, 320, 1024),
+        depths=(2, 4, 1),
+        stem_width=160,
+        conv_cfg = VitConvCfg(
+            norm_layer='layernorm2d',
+            norm_eps=1e-6,
+        ),
+        head_type='1d',
+    ), 
+    )
+    model_args = dict(img_size=336, embed_dim=1024, depth=31, num_heads=16, mlp_layer=GeGluMlp, mlp_ratio=2., class_token=False, global_pool='avg')
+    model = _create_vision_transformer_hybrid(
+        'vitamin_large_336', backbone=backbone, pretrained=pretrained, **dict(model_args, **kwargs))
+    return model
+
+
+@register_model
+def vitamin_large_384(pretrained=False, **kwargs) -> VisionTransformer:
+    backbone = MbConvStages(cfg=VitCfg(
+        embed_dim=(160, 320, 1024),
+        depths=(2, 4, 1),
+        stem_width=160,
+        conv_cfg = VitConvCfg(
+            norm_layer='layernorm2d',
+            norm_eps=1e-6,
+        ),
+        head_type='1d',
+    ), 
+    )
+    model_args = dict(img_size=384, embed_dim=1024, depth=31, num_heads=16, mlp_layer=GeGluMlp, mlp_ratio=2., class_token=False, is_pos_embed=False, global_pool='avg')
+    model = _create_vision_transformer_hybrid(
+        'vitamin_large_384', backbone=backbone, pretrained=pretrained, **dict(model_args, **kwargs))
+    return model
+
+
+@register_model
+def vitamin_xlarge_256(pretrained=False, **kwargs) -> VisionTransformer:
+    backbone = MbConvStages(cfg=VitCfg(
+        embed_dim=(192, 384, 1152),
+        depths=(2, 4, 1),
+        stem_width=192,
+        conv_cfg = VitConvCfg(
+            norm_layer='layernorm2d',
+            norm_eps=1e-6,
+        ),
+        head_type='1d',
+    ), 
+    )
+    model_args = dict(img_size=256, embed_dim=1152, depth=32, num_heads=16, mlp_layer=GeGluMlp, mlp_ratio=2., class_token=False, is_pos_embed=False, global_pool='avg')
+    model = _create_vision_transformer_hybrid(
+        'vitamin_xlarge_256', backbone=backbone, pretrained=pretrained, **dict(model_args, **kwargs))
+    return model
+
+
+@register_model
+def vitamin_xlarge_384(pretrained=False, **kwargs) -> VisionTransformer:
+    backbone = MbConvStages(cfg=VitCfg(
+        embed_dim=(192, 384, 1152),
+        depths=(2, 4, 1),
+        stem_width=192,
+        conv_cfg = VitConvCfg(
+            norm_layer='layernorm2d',
+            norm_eps=1e-6,
+        ),
+        head_type='1d',
+    ), 
+    )
+    model_args = dict(img_size=384, embed_dim=1152, depth=32, num_heads=16, mlp_layer=GeGluMlp, mlp_ratio=2., class_token=False, is_pos_embed=False, global_pool='avg')
+    model = _create_vision_transformer_hybrid(
+        'vitamin_xlarge_384', backbone=backbone, pretrained=pretrained, **dict(model_args, **kwargs))
+    return model
+
+
+def count_params(model: nn.Module):
+    return sum([m.numel() for m in model.parameters()])
+
+
+def count_stage_params(model: nn.Module, prefix='none'):
+    collections = []
+    for name, m in model.named_parameters():
+        print(name)
+        if name.startswith(prefix):
+            collections.append(m.numel())
+    return sum(collections)
+
+
+if __name__ == "__main__":
+    model = timm.create_model('vitamin_large', num_classes=10).cuda()
+    # x = torch.rand([2,3,224,224]).cuda()
+    check_keys(model)