Initial EVA-CLIP HF upload

__init__.py

@@ -0,0 +1,6 @@
+from .configuration_eva_clip import EvaClipVisionConfig
+from .modeling_eva_clip import EvaClipVisionModel
+
+__all__ = ["EvaClipVisionConfig", "EvaClipVisionModel"]
+
+

config.json

@@ -0,0 +1,38 @@
+{
+  "_name_or_path": "eva-clip-E-14-plus",
+  "architectures": [
+    "EvaClipVisionModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_eva_clip.EvaClipVisionConfig",
+    "AutoModel": "modeling_eva_clip.EvaClipVisionModel"
+  },
+  "model_type": "eva-clip-vision",
+  "vision_tower": "eva-clip-E-14-plus",
+  "embed_dim": 1024,
+  "vision_cfg": {
+    "eva_model_name": "eva-clip-E-14-plus",
+    "image_size": 448,
+    "patch_size": 14,
+    "width": 1792,
+    "layers": 64,
+    "mlp_ratio": 8.571428571428571,
+    "head_width": 112,
+    "drop_path_rate": 0.0,
+    "qkv_bias": true,
+    "xattn": true,
+    "postnorm": true,
+    "global_average_pool": false,
+    "patch_dropout": 0.0,
+    "rope": false,
+    "pt_hw_seq_len": 32,
+    "intp_freq": false,
+    "naiveswiglu": false,
+    "subln": false,
+    "fusedLN": false
+  },
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.44.0"
+}
+
+

configuration_eva_clip.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+from copy import deepcopy
+from typing import Any, Dict, Optional
+
+from transformers.configuration_utils import PretrainedConfig
+
+DEFAULT_MODEL_NAME = "eva-clip-E-14-plus"
+
+DEFAULT_VISION_CFG: Dict[str, Any] = {
+    "eva_model_name": DEFAULT_MODEL_NAME,
+    "image_size": 448,
+    "patch_size": 14,
+    "width": 1792,
+    "layers": 64,
+    "mlp_ratio": 8.571428571428571,
+    "head_width": 112,
+    "drop_path_rate": 0.0,
+    "qkv_bias": True,
+    "xattn": True,
+    "postnorm": True,
+    "global_average_pool": False,
+    "patch_dropout": 0.0,
+    "rope": False,
+    "pt_hw_seq_len": 32,
+    "intp_freq": False,
+    "naiveswiglu": False,
+    "subln": False,
+    "fusedLN": False,
+}
+
+
+class EvaClipVisionConfig(PretrainedConfig):
+    model_type = "eva-clip-vision"
+
+    def __init__(
+        self,
+        *,
+        vision_tower: str = DEFAULT_MODEL_NAME,
+        embed_dim: int = 1024,
+        vision_cfg: Optional[Dict[str, Any]] = None,
+        vision_tower_pretrained: Optional[str] = None,
+        projection_dim: Optional[int] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+
+        cfg = deepcopy(DEFAULT_VISION_CFG)
+        if vision_cfg is not None:
+            cfg.update(vision_cfg)
+
+        self.vision_tower = vision_tower
+        self.embed_dim = embed_dim
+        self.vision_cfg = cfg
+        self.vision_tower_pretrained = vision_tower_pretrained
+        self.projection_dim = projection_dim if projection_dim is not None else embed_dim
+
+
+__all__ = ["EvaClipVisionConfig", "DEFAULT_VISION_CFG", "DEFAULT_MODEL_NAME"]
+
+

legacy_eva_clip/eva_clip_encoder.py

@@ -0,0 +1,75 @@
+import torch
+import torch.nn as nn
+
+from .eva_clip_processors import EvaClipImageTrainProcessor
+from .eva_vit import EVAEncoderWrapper
+from .factory import list_models, add_model_config, get_model_config
+
+
+
+class EvaClipVisionTower(nn.Module):
+    def __init__(self, vision_tower, args, delay_load=False):
+        super().__init__()
+
+        self.is_loaded = False
+        self.vision_tower_name = vision_tower
+        self.vision_tower_pretrained = args.vision_tower_pretrained
+        self.config = get_model_config(vision_tower)
+
+
+        if not delay_load:
+            print(f"Loading EVA ViT: {self.vision_tower_name}")
+            self.load_model()
+        elif getattr(args, "unfreeze_mm_vision_tower", False):
+            # TODO: better detector is needed.
+            print(f"The checkpoint seems to contain `vision_tower` weights: `unfreeze_mm_vision_tower`: True.")
+            self.load_model()
+        elif hasattr(args, "mm_tunable_parts") and "mm_vision_tower" in args.mm_tunable_parts:
+            print(f"The checkpoint seems to contain `vision_tower` weights: `mm_tunable_parts` contains `mm_vision_tower`.")
+            self.load_model()
+        else:
+            self.cfg_only = self.config
+
+
+    def load_model(self, device_map=None):
+        print(f"Pretrained: {self.vision_tower_pretrained}")
+        self.image_processor = EvaClipImageTrainProcessor(self.config["vision_cfg"]["image_size"])
+        self.vision_tower = EVAEncoderWrapper(self.vision_tower_pretrained, self.config)
+        print(f"Loaded image processor: {self.image_processor}")
+        self.vision_tower.requires_grad_(False)
+        self.is_loaded = True
+
+    def forward(self, images):
+        if type(images) is list:
+            image_features = []
+            for image in images:
+                image_feature = self.vision_tower(image.to(device=self.device, dtype=self.dtype).unsqueeze(0)).to(image.dtype)
+                image_features.append(image_feature)
+        else:
+            image_features = self.vision_tower(images.to(device=self.device, dtype=self.dtype)).to(images.dtype)
+
+        return image_features
+
+    @property
+    def dtype(self):
+        return self.vision_tower.dtype
+
+    @property
+    def device(self):
+        return self.vision_tower.device
+
+    @property
+    def hidden_size(self):
+        return self.config["vision_cfg"]["width"]
+
+    @property
+    def num_patches(self):
+        return (self.config["vision_cfg"]["image_size"] // self.config["vision_cfg"]["patch_size"]) ** 2
+
+    @property
+    def num_patches_per_side(self):
+        return self.config["vision_cfg"]["image_size"] // self.config["vision_cfg"]["patch_size"]
+
+    @property
+    def image_size(self):
+        return self.config["vision_cfg"]["image_size"]

legacy_eva_clip/eva_clip_processors.py

@@ -0,0 +1,74 @@
+"""
+# Adapted from https://github.com/baaivision/EVA/tree/master/EVA-CLIP
+"""
+
+from torchvision import transforms
+from torchvision.transforms.functional import InterpolationMode
+from transformers.image_processing_utils import BatchFeature
+from PIL import Image
+from transformers.image_transforms import convert_to_rgb
+import numpy as np
+
+class BaseProcessor:
+    def __init__(self):
+        self.transform = lambda x: x
+        return
+
+    def __call__(self, item):
+        return self.transform(item)
+
+
+class EvaClipImageBaseProcessor(BaseProcessor):
+    def __init__(self, mean=None, std=None):
+        self.mean = (0.48145466, 0.4578275, 0.40821073) if mean is None else mean
+        self.std = (0.26862954, 0.26130258, 0.27577711) if std is None else std
+        self.normalize = transforms.Normalize(self.mean, self.std)
+
+
+    @property
+    def image_mean(self):
+        return self.mean
+
+
+class EvaClipImageTrainProcessor(EvaClipImageBaseProcessor):
+    def __init__(self, image_size=224, mean=None, std=None, min_scale=0.5, max_scale=1.0):
+        super().__init__(mean=mean, std=std)
+
+        self.transform = transforms.Compose(
+            [
+                convert_to_rgb,
+                transforms.Resize(
+                    image_size,
+                    interpolation=InterpolationMode.BICUBIC,
+                ),
+                transforms.CenterCrop(image_size),
+                transforms.ToTensor(),
+                self.normalize,
+            ]
+        )
+
+
+        self.image_size = image_size
+
+    def preprocess(self, images, return_tensors):
+        if isinstance(images, Image.Image):
+            images = [images]
+        else:
+            assert isinstance(images, list)
+        
+        transformed_images = [self.transform(image).numpy() for image in images]
+        data = {"pixel_values": transformed_images}
+
+
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+    def __call__(self, item):
+        return self.transform(item)
+
+    @property
+    def crop_size(self):
+        return {"height": self.image_size, "width": self.image_size}
+
+    @property
+    def size(self):
+        return {"shortest_edge": self.image_size}

legacy_eva_clip/eva_vit.py

@@ -0,0 +1,815 @@
+"""
+# Adapted from https://github.com/baaivision/EVA/tree/master/EVA-CLIP
+"""
+
+from math import pi
+import os
+from pathlib import Path
+from typing import Optional
+import importlib.util
+
+import torch
+from torch import nn
+from einops import rearrange, repeat
+import logging
+
+_hub_spec = importlib.util.find_spec("huggingface_hub")
+if _hub_spec is not None:
+    from huggingface_hub import snapshot_download  # type: ignore
+else:  # pragma: no cover - optional dependency
+    snapshot_download = None
+
+_DEFAULT_EVA_REPO = "jiuhai/eva_clip_vision_tower"
+XFORMERS_AVAILABLE = False  # populated later once xops import resolves
+_XFORMERS_WARNING_EMITTED = False
+
+
+def _resolve_vision_checkpoint_path(vision_tower_pretrained: Optional[str]) -> str:
+    """Determine where to load the EVA visual tower weights from."""
+
+    candidate_files: list[Path] = []
+    if vision_tower_pretrained:
+        supplied = Path(vision_tower_pretrained)
+        if supplied.is_file():
+            candidate_files.append(supplied)
+        elif supplied.is_dir():
+            for filename in (
+                "pytorch_model.bin",
+                "model.safetensors",
+                "visual.pth",
+                "visual.bin",
+            ):
+                candidate = supplied / filename
+                if candidate.exists():
+                    candidate_files.append(candidate)
+                    break
+
+    if candidate_files:
+        return str(candidate_files[0])
+
+    if snapshot_download is None:
+        raise FileNotFoundError(
+            "EVA vision weights not found locally and huggingface_hub is unavailable. "
+            "Provide --vision-tower-path pointing to a directory containing the checkpoint."
+        )
+
+    cache_dir = Path(snapshot_download(repo_id=_DEFAULT_EVA_REPO))
+    default_file = cache_dir / "pytorch_model.bin"
+    if not default_file.exists():
+        raise FileNotFoundError(f"Default EVA checkpoint not found at {default_file}.")
+    return str(default_file)
+
+
+def broadcat(tensors, dim=-1):
+    num_tensors = len(tensors)
+    shape_lens = set(list(map(lambda t: len(t.shape), tensors)))
+    assert len(shape_lens) == 1, "tensors must all have the same number of dimensions"
+    shape_len = list(shape_lens)[0]
+    dim = (dim + shape_len) if dim < 0 else dim
+    dims = list(zip(*map(lambda t: list(t.shape), tensors)))
+    expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim]
+    assert all([*map(lambda t: len(set(t[1])) <= 2, expandable_dims)]), "invalid dimensions for broadcastable concatentation"
+    max_dims = list(map(lambda t: (t[0], max(t[1])), expandable_dims))
+    expanded_dims = list(map(lambda t: (t[0], (t[1],) * num_tensors), max_dims))
+    expanded_dims.insert(dim, (dim, dims[dim]))
+    expandable_shapes = list(zip(*map(lambda t: t[1], expanded_dims)))
+    tensors = list(map(lambda t: t[0].expand(*t[1]), zip(tensors, expandable_shapes)))
+    return torch.cat(tensors, dim=dim)
+
+
+def rotate_half(x):
+    x = rearrange(x, "... (d r) -> ... d r", r=2)
+    x1, x2 = x.unbind(dim=-1)
+    x = torch.stack((-x2, x1), dim=-1)
+    return rearrange(x, "... d r -> ... (d r)")
+
+
+class VisionRotaryEmbeddingFast(nn.Module):
+    def __init__(self, dim, pt_seq_len, ft_seq_len=None, custom_freqs=None, freqs_for="lang", theta=10000, max_freq=10, num_freqs=1, patch_dropout=0.0):
+        super().__init__()
+        if custom_freqs:
+            freqs = custom_freqs
+        elif freqs_for == "lang":
+            freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+        elif freqs_for == "pixel":
+            freqs = torch.linspace(1.0, max_freq / 2, dim // 2) * pi
+        elif freqs_for == "constant":
+            freqs = torch.ones(num_freqs).float()
+        else:
+            raise ValueError(f"unknown modality {freqs_for}")
+
+        if ft_seq_len is None:
+            ft_seq_len = pt_seq_len
+        t = torch.arange(ft_seq_len) / ft_seq_len * pt_seq_len
+
+        freqs = torch.einsum("..., f -> ... f", t, freqs)
+        freqs = repeat(freqs, "... n -> ... (n r)", r=2)
+        freqs = broadcat((freqs[:, None, :], freqs[None, :, :]), dim=-1)
+
+        freqs_cos = freqs.cos().view(-1, freqs.shape[-1])
+        freqs_sin = freqs.sin().view(-1, freqs.shape[-1])
+
+        self.patch_dropout = patch_dropout
+
+        self.register_buffer("freqs_cos", freqs_cos)
+        self.register_buffer("freqs_sin", freqs_sin)
+
+        logging.info(f"Shape of rope freq: {self.freqs_cos.shape}")
+
+    def forward(self, t, patch_indices_keep=None):
+        if patch_indices_keep is not None:
+            batch = t.size()[0]
+            batch_indices = torch.arange(batch)
+            batch_indices = batch_indices[..., None]
+
+            freqs_cos = repeat(self.freqs_cos, "i j -> n i m j", n=t.shape[0], m=t.shape[1])
+            freqs_sin = repeat(self.freqs_sin, "i j -> n i m j", n=t.shape[0], m=t.shape[1])
+
+            freqs_cos = freqs_cos[batch_indices, patch_indices_keep]
+            freqs_cos = rearrange(freqs_cos, "n i m j -> n m i j")
+            freqs_sin = freqs_sin[batch_indices, patch_indices_keep]
+            freqs_sin = rearrange(freqs_sin, "n i m j -> n m i j")
+
+            return t * freqs_cos + rotate_half(t) * freqs_sin
+
+        return t * self.freqs_cos + rotate_half(t) * self.freqs_sin
+
+
+class LayerNorm(nn.LayerNorm):
+    """Subclass torch's LayerNorm (with cast back to input dtype)."""
+
+    def forward(self, x: torch.Tensor):
+        orig_type = x.dtype
+        x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        return x.to(orig_type)
+
+
+class PatchDropout(nn.Module):
+    """
+    https://arxiv.org/abs/2212.00794
+    """
+
+    def __init__(self, prob, exclude_first_token=True):
+        super().__init__()
+        assert 0 <= prob < 1.
+        self.prob = prob
+        self.exclude_first_token = exclude_first_token  # exclude CLS token
+        print(f"os.getenv('RoPE')={os.getenv('RoPE')}")
+
+    def forward(self, x):
+        if not self.training or self.prob == 0.:
+            return x
+
+        if self.exclude_first_token:
+            cls_tokens, x = x[:, :1], x[:, 1:]
+        else:
+            cls_tokens = torch.jit.annotate(torch.Tensor, x[:, :1])
+
+        batch = x.size()[0]
+        num_tokens = x.size()[1]
+
+        batch_indices = torch.arange(batch)
+        batch_indices = batch_indices[..., None]
+
+        keep_prob = 1 - self.prob
+        num_patches_keep = max(1, int(num_tokens * keep_prob))
+
+        rand = torch.randn(batch, num_tokens)
+        patch_indices_keep = rand.topk(num_patches_keep, dim=-1).indices
+
+        x = x[batch_indices, patch_indices_keep]
+
+        if self.exclude_first_token:
+            x = torch.cat((cls_tokens, x), dim=1)
+
+        if self.training and os.getenv('RoPE') == '1':
+            return x, patch_indices_keep
+
+        return x
+
+
+# --------------------------------------------------------
+# Adapted from  https://github.com/microsoft/unilm/tree/master/beit
+# --------------------------------------------------------
+import math
+import os
+import torch.nn as nn
+import torch.nn.functional as F
+
+try:
+    from timm.models.layers import drop_path, to_2tuple, trunc_normal_
+except:
+    from timm.layers import drop_path, to_2tuple, trunc_normal_
+
+if os.getenv("ENV_TYPE") == "deepspeed":
+    try:
+        from deepspeed.runtime.activation_checkpointing.checkpointing import checkpoint
+    except:
+        from torch.utils.checkpoint import checkpoint
+else:
+    from torch.utils.checkpoint import checkpoint
+
+try:
+    import xformers.ops as xops
+    XFORMERS_AVAILABLE = True
+except ImportError:
+    xops = None
+    # print("Please 'pip install xformers'")
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+    
+    def extra_repr(self) -> str:
+        return 'p={}'.format(self.drop_prob)
+
+
+class Mlp(nn.Module):
+    def __init__(
+        self, 
+        in_features, 
+        hidden_features=None, 
+        out_features=None, 
+        act_layer=nn.GELU, 
+        norm_layer=nn.LayerNorm, 
+        drop=0.,
+        subln=False,
+
+        ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+
+        self.ffn_ln = norm_layer(hidden_features) if subln else nn.Identity()
+
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        # x = self.drop(x)
+        # commit this for the orignal BERT implement 
+        x = self.ffn_ln(x)
+
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.SiLU, drop=0., 
+                norm_layer=nn.LayerNorm, subln=False):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+
+        self.w1 = nn.Linear(in_features, hidden_features)
+        self.w2 = nn.Linear(in_features, hidden_features)
+
+        self.act = act_layer()
+        self.ffn_ln = norm_layer(hidden_features) if subln else nn.Identity()
+        self.w3 = nn.Linear(hidden_features, out_features)
+        
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x1 = self.w1(x)
+        x2 = self.w2(x)
+        hidden = self.act(x1) * x2
+        x = self.ffn_ln(hidden)
+        x = self.w3(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(
+            self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.,
+            proj_drop=0., window_size=None, attn_head_dim=None, xattn=False, rope=None, subln=False, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        if attn_head_dim is not None:
+            head_dim = attn_head_dim
+        all_head_dim = head_dim * self.num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.subln = subln
+        if self.subln:
+            self.q_proj = nn.Linear(dim, all_head_dim, bias=False)
+            self.k_proj = nn.Linear(dim, all_head_dim, bias=False)
+            self.v_proj = nn.Linear(dim, all_head_dim, bias=False)
+        else:
+            if qkv_bias:
+                self.qkv = nn.Linear(dim, all_head_dim * 3, bias=True)
+            else:
+                self.qkv = nn.Linear(dim, all_head_dim * 3, bias=False)
+
+        # if qkv_bias:
+        #     self.q_bias = nn.Parameter(torch.zeros(all_head_dim))
+        #     self.v_bias = nn.Parameter(torch.zeros(all_head_dim))
+        # else:
+        #     self.q_bias = None
+        #     self.v_bias = None
+
+        self.window_size = None
+        self.relative_position_bias_table = None
+        self.relative_position_index = None
+
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.inner_attn_ln = norm_layer(all_head_dim) if subln else nn.Identity()
+        # self.proj = nn.Linear(all_head_dim, all_head_dim)
+        self.proj = nn.Linear(all_head_dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.xattn = xattn and XFORMERS_AVAILABLE
+        self.xattn_drop = attn_drop
+        global _XFORMERS_WARNING_EMITTED
+        if xattn and not self.xattn and not _XFORMERS_WARNING_EMITTED:
+            logging.warning("xformers.ops not available; falling back to standard attention for EVA-CLIP vision tower.")
+            _XFORMERS_WARNING_EMITTED = True
+
+        self.rope = rope
+
+    def forward(self, x, rel_pos_bias=None, attn_mask=None):
+        B, N, C = x.shape
+        if self.subln: 
+            q = F.linear(input=x, weight=self.q_proj.weight, bias=self.q_bias)
+            k = F.linear(input=x, weight=self.k_proj.weight, bias=None)
+            v = F.linear(input=x, weight=self.v_proj.weight, bias=self.v_bias)
+
+            q = q.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)     # B, num_heads, N, C
+            k = k.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)  
+            v = v.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3) 
+        else: 
+
+            # qkv_bias = None
+            # if self.q_bias is not None:
+            #     qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))
+            
+            # qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
+
+            qkv = self.qkv(x)
+
+            qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)   # 3, B, num_heads, N, C
+            q, k, v = qkv[0], qkv[1], qkv[2]
+
+        if self.rope:
+            q_t = q[:, :, 1:, :]
+            ro_q_t = self.rope(q_t)
+            q = torch.cat((q[:, :, :1, :], ro_q_t), -2).type_as(v)
+
+            k_t = k[:, :, 1:, :]
+            ro_k_t = self.rope(k_t)
+            k = torch.cat((k[:, :, :1, :], ro_k_t), -2).type_as(v)
+
+        if self.xattn:
+            q = q.permute(0, 2, 1, 3)   # B, num_heads, N, C -> B, N, num_heads, C
+            k = k.permute(0, 2, 1, 3)
+            v = v.permute(0, 2, 1, 3)
+
+            x = xops.memory_efficient_attention(
+                q, k, v,
+                p=self.xattn_drop,
+                scale=self.scale,
+                )
+            x = x.reshape(B, N, -1)
+            x = self.inner_attn_ln(x)
+            x = self.proj(x)
+            x = self.proj_drop(x)
+        else:
+            q = q * self.scale
+            attn = (q @ k.transpose(-2, -1))
+
+            if self.relative_position_bias_table is not None:
+                relative_position_bias = \
+                    self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
+                        self.window_size[0] * self.window_size[1] + 1,
+                        self.window_size[0] * self.window_size[1] + 1, -1)  # Wh*Ww,Wh*Ww,nH
+                relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
+                attn = attn + relative_position_bias.unsqueeze(0).type_as(attn)
+
+            if rel_pos_bias is not None:
+                attn = attn + rel_pos_bias.type_as(attn)
+
+            if attn_mask is not None:
+                attn_mask = attn_mask.bool()
+                attn = attn.masked_fill(~attn_mask[:, None, None, :], float("-inf"))
+            
+            attn = attn.softmax(dim=-1)
+            attn = self.attn_drop(attn)
+
+            x = (attn @ v).transpose(1, 2).reshape(B, N, -1)
+            x = self.inner_attn_ln(x)
+            x = self.proj(x)
+            x = self.proj_drop(x)
+        return x
+
+
+class Block(nn.Module):
+
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., init_values=None, act_layer=nn.GELU, norm_layer=nn.LayerNorm,
+                 window_size=None, attn_head_dim=None, xattn=False, rope=None, postnorm=False,
+                 subln=False, naiveswiglu=False):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
+            attn_drop=attn_drop, proj_drop=drop, window_size=window_size, attn_head_dim=attn_head_dim,
+            xattn=xattn, rope=rope, subln=subln, norm_layer=norm_layer)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+
+        if naiveswiglu:
+            self.mlp = SwiGLU(
+                in_features=dim, 
+                hidden_features=mlp_hidden_dim, 
+                subln=subln,
+                norm_layer=norm_layer,
+            )
+        else:
+            self.mlp = Mlp(
+                in_features=dim, 
+                hidden_features=mlp_hidden_dim, 
+                act_layer=act_layer,
+                subln=subln,
+                drop=drop
+            )
+
+        if init_values is not None and init_values > 0:
+            self.gamma_1 = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True)
+            self.gamma_2 = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True)
+        else:
+            self.gamma_1, self.gamma_2 = None, None
+
+        self.postnorm = postnorm
+
+    def forward(self, x, rel_pos_bias=None, attn_mask=None):
+        if self.gamma_1 is None:
+            if self.postnorm:
+                x = x + self.drop_path(self.norm1(self.attn(x, rel_pos_bias=rel_pos_bias, attn_mask=attn_mask)))
+                x = x + self.drop_path(self.norm2(self.mlp(x)))
+            else:
+                x = x + self.drop_path(self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias, attn_mask=attn_mask))
+                x = x + self.drop_path(self.mlp(self.norm2(x)))
+        else:
+            if self.postnorm:
+                x = x + self.drop_path(self.gamma_1 * self.norm1(self.attn(x, rel_pos_bias=rel_pos_bias, attn_mask=attn_mask)))
+                x = x + self.drop_path(self.gamma_2 * self.norm2(self.mlp(x)))
+            else:
+                x = x + self.drop_path(self.gamma_1 * self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias, attn_mask=attn_mask))
+                x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """ Image to Patch Embedding
+    """
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
+        self.patch_shape = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+
+    def forward(self, x, **kwargs):
+        B, C, H, W = x.shape
+        # FIXME look at relaxing size constraints
+        assert H == self.img_size[0] and W == self.img_size[1], \
+            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        x = self.proj(x).flatten(2).transpose(1, 2)
+        return x
+
+
+class RelativePositionBias(nn.Module):
+
+    def __init__(self, window_size, num_heads):
+        super().__init__()
+        self.window_size = window_size
+        self.num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3
+        self.relative_position_bias_table = nn.Parameter(torch.zeros(self.num_relative_distance, num_heads))  # 2*Wh-1 * 2*Ww-1, nH
+        # cls to token & token 2 cls & cls to cls
+
+        # get pair-wise relative position index for each token inside the window
+        coords_h = torch.arange(window_size[0])
+        coords_w = torch.arange(window_size[1])
+        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
+        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
+        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
+        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
+        relative_coords[:, :, 0] += window_size[0] - 1  # shift to start from 0
+        relative_coords[:, :, 1] += window_size[1] - 1
+        relative_coords[:, :, 0] *= 2 * window_size[1] - 1
+        relative_position_index = torch.zeros(size=(window_size[0] * window_size[1] + 1,) * 2, dtype=relative_coords.dtype)
+        relative_position_index[1:, 1:] = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
+        relative_position_index[0, 0:] = self.num_relative_distance - 3
+        relative_position_index[0:, 0] = self.num_relative_distance - 2
+        relative_position_index[0, 0] = self.num_relative_distance - 1
+
+        self.register_buffer("relative_position_index", relative_position_index)
+
+    def forward(self):
+        relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(self.window_size[0] * self.window_size[1] + 1, self.window_size[0] * self.window_size[1] + 1, -1)  # Wh*Ww,Wh*Ww,nH
+        return relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
+
+
+class EVAVisionTransformer(nn.Module):
+    """Vision Transformer with support for patch or hybrid CNN input stage"""
+
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,
+                 num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,
+                 drop_path_rate=0., norm_layer=nn.LayerNorm, init_values=None, patch_dropout=0.,
+                 use_abs_pos_emb=True, use_rel_pos_bias=False, use_shared_rel_pos_bias=False, rope=False,
+                 use_mean_pooling=True, init_scale=0.001, grad_checkpointing=False, xattn=False, postnorm=False,
+                 pt_hw_seq_len=16, intp_freq=False, naiveswiglu=False, subln=False,
+                 ):
+        super().__init__()
+        self.image_size = img_size
+        # self.num_classes = num_classes
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+
+        self.patch_embed = PatchEmbed(
+            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        if use_abs_pos_emb:
+            self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
+        else:
+            self.pos_embed = None
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        self.rel_pos_bias = None
+        self.rope = None
+
+        self.naiveswiglu = naiveswiglu
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        self.use_rel_pos_bias = use_rel_pos_bias
+        self.blocks = nn.ModuleList([
+            Block(
+                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer,
+                init_values=init_values, window_size=self.patch_embed.patch_shape if use_rel_pos_bias else None,
+                xattn=xattn, rope=self.rope, postnorm=postnorm, subln=subln, naiveswiglu=naiveswiglu)
+            for i in range(depth)])
+
+        # setting a patch_dropout of 0. would mean it is disabled and this function would be the identity fn
+        self.patch_dropout = PatchDropout(patch_dropout) if patch_dropout > 0. else nn.Identity()
+
+        self.grad_checkpointing = grad_checkpointing
+
+    def fix_init_weight(self):
+        def rescale(param, layer_id):
+            param.div_(math.sqrt(2.0 * layer_id))
+
+        for layer_id, layer in enumerate(self.blocks):
+            rescale(layer.attn.proj.weight.data, layer_id + 1)
+            if self.naiveswiglu:
+                rescale(layer.mlp.w3.weight.data, layer_id + 1)
+            else:
+                rescale(layer.mlp.fc2.weight.data, layer_id + 1)
+
+    def get_cast_dtype(self) -> torch.dtype:
+        return self.blocks[0].mlp.fc2.weight.dtype
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def get_num_layers(self):
+        return len(self.blocks)
+
+    def lock(self, unlocked_groups=0, freeze_bn_stats=False):
+        assert unlocked_groups == 0, "partial locking not currently supported for this model"
+        for param in self.parameters():
+            param.requires_grad = False
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.grad_checkpointing = enable
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {"pos_embed", "cls_token"}
+
+    def get_classifier(self):
+        return self.head
+
+    def reset_classifier(self, num_classes, global_pool=""):
+        self.num_classes = num_classes
+        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+    def forward_features(self, x):
+        x = self.patch_embed(x)
+        batch_size, seq_len, _ = x.size()
+
+        cls_tokens = self.cls_token.expand(batch_size, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
+        x = torch.cat((cls_tokens, x), dim=1)
+        if self.pos_embed is not None:
+            x = x + self.pos_embed
+        x = self.pos_drop(x)
+
+        # a patch_dropout of 0. would mean it is disabled and this function would do nothing but return what was passed in
+        if os.getenv('RoPE') == '1':
+            if self.training and not isinstance(self.patch_dropout, nn.Identity):
+                x, patch_indices_keep = self.patch_dropout(x)
+                self.rope.forward = partial(self.rope.forward, patch_indices_keep=patch_indices_keep)
+            else:
+                self.rope.forward = partial(self.rope.forward, patch_indices_keep=None)
+                x = self.patch_dropout(x)
+        else:
+            x = self.patch_dropout(x)
+
+        rel_pos_bias = None
+
+        for blk in self.blocks:
+            if self.grad_checkpointing:
+                x = checkpoint(blk, x, (rel_pos_bias,))
+            else:
+                x = blk(x, rel_pos_bias=rel_pos_bias)
+
+        return x
+
+    def forward(self, x, return_all_features=False):
+
+        """
+        :return:
+            forward_features function returns raw features of ViT,
+            forward with return_all_features returns normalized features of ViT
+        :param x:
+        :param return_all_features:
+        """
+
+        features = self.forward_features(x)  # [B, n_patch, C]
+        return features
+
+
+def load_state_dict(checkpoint_path: str, map_location: str = "cpu", model_key: str = "model|module|state_dict", is_openai: bool = False, skip_list: list = []):
+    if is_openai:
+        model = torch.jit.load(checkpoint_path, map_location="cpu").eval()
+        state_dict = model.state_dict()
+        for key in ["input_resolution", "context_length", "vocab_size"]:
+            state_dict.pop(key, None)
+    else:
+        checkpoint = torch.load(checkpoint_path, map_location=map_location)
+        for mk in model_key.split("|"):
+            if isinstance(checkpoint, dict) and mk in checkpoint:
+                state_dict = checkpoint[mk]
+                break
+            else:
+                state_dict = checkpoint
+        if next(iter(state_dict.items()))[0].startswith("module"):
+            state_dict = {k[7:]: v for k, v in state_dict.items()}
+
+    for k in skip_list:
+        if k in list(state_dict.keys()):
+            logging.info(f"Removing key {k} from pretrained checkpoint")
+            del state_dict[k]
+
+    if os.getenv("RoPE") == "1":
+        for k in list(state_dict.keys()):
+            if "freqs_cos" in k or "freqs_sin" in k:
+                del state_dict[k]
+    return state_dict
+
+
+def load_clip_visual_state_dict(checkpoint_path: str, map_location: str = "cpu", is_openai: bool = False, skip_list: list = []):
+    state_dict = load_state_dict(checkpoint_path, map_location=map_location, is_openai=is_openai, skip_list=skip_list)
+    # for k in list(state_dict.keys()):
+    #     if not k.startswith("visual."):
+    #         del state_dict[k]
+    # for k in list(state_dict.keys()):
+    #     if k.startswith("visual."):
+    #         new_k = k[7:]
+    #         state_dict[new_k] = state_dict[k]
+    #         del state_dict[k]
+    return state_dict
+
+
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+
+try:
+    from apex.normalization import FusedLayerNorm
+except:
+    FusedLayerNorm = LayerNorm
+    # print("Please build and install Nvidia apex package with option '--cuda_ext' according to https://github.com/NVIDIA/apex#from-source .")
+
+
+@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  # layer scale initial value
+    patch_dropout: float = 0.0  # what fraction of patches to dropout during training (0 would mean disabled and no patches dropped) - 0.5 to 0.75 recommended in the paper for optimal results
+    global_average_pool: bool = False  # whether to global average pool the last embedding layer, instead of using CLS token (https://arxiv.org/abs/2205.01580)
+    drop_path_rate: Optional[float] = None  # drop path rate
+    timm_model_name: str = None  # a valid model name overrides layers, width, patch_size
+    timm_model_pretrained: bool = False  # use (imagenet) pretrained weights for named model
+    timm_pool: str = "avg"  # feature pooling for timm model ('abs_attn', 'rot_attn', 'avg', '')
+    timm_proj: str = "linear"  # linear projection for timm model output ('linear', 'mlp', '')
+    timm_proj_bias: bool = False  # enable bias final projection
+    eva_model_name: str = None  # a valid eva model name overrides layers, width, patch_size
+    qkv_bias: bool = True
+    fusedLN: bool = False
+    xattn: bool = False
+    postnorm: bool = False
+    rope: bool = False
+    pt_hw_seq_len: int = 16  # 224/14
+    intp_freq: bool = False
+    naiveswiglu: bool = False
+    subln: bool = False
+
+
+def create_norm_layer_factory(use_fused_ln, eps=1e-6):
+    # Otherwise, use the standard LayerNorm
+    return lambda num_features: nn.LayerNorm(num_features, eps=eps)
+
+
+def _build_vision_tower(vision_tower_path: str, embed_dim: int, vision_cfg: CLIPVisionCfg, **kwargs):
+    if isinstance(vision_cfg, dict):
+        vision_cfg = CLIPVisionCfg(**vision_cfg)
+
+    if vision_cfg.eva_model_name:
+        vision_heads = vision_cfg.width // vision_cfg.head_width
+        # Determine the appropriate norm layer factory based on the configuration
+        norm_layer_factory = create_norm_layer_factory(vision_cfg.fusedLN, eps=1e-6)
+
+        # breakpoint()
+        visual = EVAVisionTransformer(
+            img_size=vision_cfg.image_size,
+            patch_size=vision_cfg.patch_size,
+            num_classes=embed_dim,
+            use_mean_pooling=vision_cfg.global_average_pool,  # False
+            init_values=vision_cfg.ls_init_value,
+            patch_dropout=vision_cfg.patch_dropout,
+            embed_dim=vision_cfg.width,
+            depth=vision_cfg.layers,
+            num_heads=vision_heads,
+            mlp_ratio=vision_cfg.mlp_ratio,
+            qkv_bias=vision_cfg.qkv_bias,
+            drop_path_rate=vision_cfg.drop_path_rate,
+            norm_layer=norm_layer_factory,
+            xattn=vision_cfg.xattn,
+            rope=vision_cfg.rope,
+            postnorm=vision_cfg.postnorm,
+            pt_hw_seq_len=vision_cfg.pt_hw_seq_len,  # 224/14
+            intp_freq=vision_cfg.intp_freq,
+            naiveswiglu=vision_cfg.naiveswiglu,
+            subln=vision_cfg.subln,
+        )
+        # breakpoint()
+        state_dict = load_clip_visual_state_dict(vision_tower_path)
+        incompatible_keys = visual.load_state_dict(state_dict, strict=False)
+        print("EVA-CLIP incompatible_keys:", incompatible_keys)
+
+    return visual
+
+
+class EVAEncoderWrapper(nn.Module):
+    def __init__(self, vision_tower_pretrained, config):
+        super(EVAEncoderWrapper, self).__init__()
+        self.config = config
+        vision_tower_path = _resolve_vision_checkpoint_path(vision_tower_pretrained)
+        self.config["vision_tower_path"] = vision_tower_path
+        self.vision_tower_path = vision_tower_path
+        self.model = _build_vision_tower(**self.config)
+
+    def forward(self, image, **kwargs):
+        encode = self.model(image, return_all_features=True)[:, 1:, :]  # remove the CLS token
+        return encode
+
+    @property
+    def dtype(self):
+        return list(self.parameters())[-1].dtype
+
+    @property
+    def device(self):
+        return list(self.parameters())[-1].device

legacy_eva_clip/factory.py

@@ -0,0 +1,59 @@
+import json
+import logging
+import os
+import pathlib
+import re
+from copy import deepcopy
+from pathlib import Path
+from typing import Optional, Tuple, Union, Dict, Any
+import torch
+
+_MODEL_CONFIG_PATHS = [Path(__file__).parent / f"model_configs/"]
+_MODEL_CONFIGS = {}  # directory (model_name: config) of model architecture configs
+
+
+def _natural_key(string_):
+    return [int(s) if s.isdigit() else s for s in re.split(r"(\d+)", string_.lower())]
+
+
+def _rescan_model_configs():
+    global _MODEL_CONFIGS
+
+    config_ext = (".json",)
+    config_files = []
+    for config_path in _MODEL_CONFIG_PATHS:
+        if config_path.is_file() and config_path.suffix in config_ext:
+            config_files.append(config_path)
+        elif config_path.is_dir():
+            for ext in config_ext:
+                config_files.extend(config_path.glob(f"*{ext}"))
+    for cf in config_files:
+        with open(cf, "r", encoding="utf8") as f:
+            model_cfg = json.load(f)
+            if all(a in model_cfg for a in ("embed_dim", "vision_cfg", "text_cfg")):
+                _MODEL_CONFIGS[cf.stem] = model_cfg
+
+    _MODEL_CONFIGS = dict(sorted(_MODEL_CONFIGS.items(), key=lambda x: _natural_key(x[0])))
+
+
+_rescan_model_configs()  # initial populate of model config registry
+
+
+def list_models():
+    """enumerate available model architectures based on config files"""
+    return list(_MODEL_CONFIGS.keys())
+
+
+def add_model_config(path):
+    """add model config path or file and update registry"""
+    if not isinstance(path, Path):
+        path = Path(path)
+    _MODEL_CONFIG_PATHS.append(path)
+    _rescan_model_configs()
+
+
+def get_model_config(model_name):
+    if model_name in _MODEL_CONFIGS:
+        return deepcopy(_MODEL_CONFIGS[model_name])
+    else:
+        return None

legacy_eva_clip/model_configs/EVA-CLIP-18B.json

@@ -0,0 +1,27 @@
+{
+    "embed_dim": 1536,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 48,
+        "width": 5120,
+        "head_width": 128,
+        "mlp_ratio": 5,
+        "patch_size": 14,
+        "eva_model_name": "eva-clip-18b-14-x",
+        "drop_path_rate": 0,
+        "qkv_bias": false,
+        "xattn": true,
+        "postnorm": true,
+        "fusedLN": false,
+        "use_rms_norm": true
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 1280,
+        "heads": 20,
+        "layers": 32,
+        "xattn": false,
+        "fusedLN": false
+    }
+}

legacy_eva_clip/model_configs/EVA-CLIP-8B-plus.json

@@ -0,0 +1,27 @@
+{
+    "embed_dim": 1280,
+    "vision_cfg": {
+        "image_size": 448,
+        "layers": 32,
+        "width": 4096,
+        "head_width": 128,
+        "mlp_ratio": 5,
+        "patch_size": 14,
+        "eva_model_name": "eva-clip-8b-14-plus-x",
+        "drop_path_rate": 0,
+        "qkv_bias": false,
+        "xattn": true,
+        "postnorm": false,
+        "fusedLN": false,
+        "use_rms_norm": true
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 1280,
+        "heads": 20,
+        "layers": 32,
+        "xattn": false,
+        "fusedLN": false
+    }
+}

legacy_eva_clip/model_configs/EVA-CLIP-8B.json

@@ -0,0 +1,27 @@
+{
+    "embed_dim": 1280,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 32,
+        "width": 4096,
+        "head_width": 128,
+        "mlp_ratio": 5,
+        "patch_size": 14,
+        "eva_model_name": "eva-clip-8b-14-x",
+        "drop_path_rate": 0,
+        "qkv_bias": false,
+        "xattn": true,
+        "postnorm": false,
+        "fusedLN": false,
+        "use_rms_norm": true
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 1280,
+        "heads": 20,
+        "layers": 32,
+        "xattn": false,
+        "fusedLN": false
+    }
+}

legacy_eva_clip/model_configs/EVA01-CLIP-B-16.json

@@ -0,0 +1,19 @@
+{
+    "embed_dim": 512,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 12,
+        "width": 768,
+        "patch_size": 16,
+        "eva_model_name": "eva-clip-b-16",
+        "ls_init_value": 0.1,
+        "drop_path_rate": 0.0
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}

legacy_eva_clip/model_configs/EVA01-CLIP-g-14-plus.json

@@ -0,0 +1,24 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 40,
+        "width": 1408,
+        "head_width": 88,
+        "mlp_ratio": 4.3637,
+        "patch_size": 14,
+        "eva_model_name": "eva-clip-g-14-x",
+        "drop_path_rate": 0,
+        "xattn": true,
+        "fusedLN": true
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 1024,
+        "heads": 16,
+        "layers": 24,
+        "xattn": false,
+        "fusedLN": true
+    }
+}

legacy_eva_clip/model_configs/EVA01-CLIP-g-14.json

@@ -0,0 +1,24 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 40,
+        "width": 1408,
+        "head_width": 88,
+        "mlp_ratio": 4.3637,
+        "patch_size": 14,
+        "eva_model_name": "eva-clip-g-14-x",
+        "drop_path_rate": 0.4,
+        "xattn": true,
+        "fusedLN": true
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 768,
+        "heads": 12,
+        "layers": 12,
+        "xattn": false,
+        "fusedLN": true
+    }
+}

legacy_eva_clip/model_configs/EVA02-CLIP-B-16.json

@@ -0,0 +1,29 @@
+{
+    "embed_dim": 512,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 12,
+        "width": 768,
+        "head_width": 64,
+        "patch_size": 16,
+        "mlp_ratio": 2.6667,
+        "eva_model_name": "eva-clip-b-16-X",
+        "drop_path_rate": 0.0,
+        "xattn": true,
+        "fusedLN": true,
+        "rope": true,
+        "pt_hw_seq_len": 16,
+        "intp_freq": true,
+        "naiveswiglu": true,
+        "subln": true
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12,
+        "xattn": true,
+        "fusedLN": true
+    }
+}

legacy_eva_clip/model_configs/EVA02-CLIP-L-14-336.json

@@ -0,0 +1,29 @@
+{
+    "embed_dim": 768,
+    "vision_cfg": {
+        "image_size": 336,
+        "layers": 24,
+        "width": 1024,
+        "drop_path_rate": 0,
+        "head_width": 64,
+        "mlp_ratio": 2.6667,
+        "patch_size": 14,
+        "eva_model_name": "eva-clip-l-14-336",
+        "xattn": true,
+        "fusedLN": true,
+        "rope": true,
+        "pt_hw_seq_len": 16,
+        "intp_freq": true,
+        "naiveswiglu": true,
+        "subln": true
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 768,
+        "heads": 12,
+        "layers": 12,
+        "xattn": false,
+        "fusedLN": true
+    }
+}

legacy_eva_clip/model_configs/EVA02-CLIP-L-14.json

@@ -0,0 +1,29 @@
+{
+    "embed_dim": 768,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 24,
+        "width": 1024,
+        "drop_path_rate": 0,
+        "head_width": 64,
+        "mlp_ratio": 2.6667,
+        "patch_size": 14,
+        "eva_model_name": "eva-clip-l-14",
+        "xattn": true,
+        "fusedLN": true,
+        "rope": true,
+        "pt_hw_seq_len": 16,
+        "intp_freq": true,
+        "naiveswiglu": true,
+        "subln": true
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 768,
+        "heads": 12,
+        "layers": 12,
+        "xattn": false,
+        "fusedLN": true
+    }
+}

legacy_eva_clip/model_configs/EVA02-CLIP-bigE-14-plus.json

@@ -0,0 +1,28 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 64,
+        "width": 1792,
+        "head_width": 112,
+        "mlp_ratio": 8.571428571428571,
+        "patch_size": 14,
+        "eva_model_name": "eva-clip-4b-14-x",
+        "drop_path_rate": 0,
+        "xattn": true,
+        "postnorm": true,
+        "fusedLN": true
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 1280,
+        "heads": 20,
+        "layers": 32,
+        "xattn": false,
+        "fusedLN": true
+    }
+}
+
+
+

legacy_eva_clip/model_configs/EVA02-CLIP-bigE-14.json

@@ -0,0 +1,25 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 64,
+        "width": 1792,
+        "head_width": 112,
+        "mlp_ratio": 8.571428571428571,
+        "patch_size": 14,
+        "eva_model_name": "eva-clip-4b-14-x",
+        "drop_path_rate": 0,
+        "xattn": true,
+        "postnorm": true,
+        "fusedLN": true
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 1024,
+        "heads": 16,
+        "layers": 24,
+        "xattn": false,
+        "fusedLN": true
+    }
+}

legacy_eva_clip/model_configs/Internal-EVA02-CLIP-10B-14-448.json

@@ -0,0 +1,25 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 448,
+        "layers": 77,
+        "width": 2304,
+        "head_width": 144,
+        "mlp_ratio": 10.9722,
+        "patch_size": 14,
+        "eva_model_name": "eva-clip-10b-14-x",
+        "drop_path_rate": 0,
+        "xattn": true,
+        "postnorm": false,
+        "fusedLN": true
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 1280,
+        "heads": 20,
+        "layers": 32,
+        "xattn": false,
+        "fusedLN": true
+    }
+}

legacy_eva_clip/model_configs/Internal-EVA02-CLIP-10B-14.json

@@ -0,0 +1,25 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 77,
+        "width": 2304,
+        "head_width": 144,
+        "mlp_ratio": 10.9722,
+        "patch_size": 14,
+        "eva_model_name": "eva-clip-10b-14-x",
+        "drop_path_rate": 0,
+        "xattn": true,
+        "postnorm": false,
+        "fusedLN": true
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 1280,
+        "heads": 20,
+        "layers": 32,
+        "xattn": false,
+        "fusedLN": true
+    }
+}

legacy_eva_clip/model_configs/eva-clip-E-14-plus.json

@@ -0,0 +1,27 @@
+{
+  "embed_dim": 1024,
+  "vision_cfg": {
+      "drop_path_rate": 0,
+      "eva_model_name": "eva-clip-E-14-plus",
+      "head_width": 112,
+      "image_size": 448,
+      "layers": 64,
+      "mlp_ratio": 8.571428571428571,
+      "patch_size": 14,
+      "postnorm": true,
+      "qkv_bias": true,
+      "width": 1792,
+      "xattn": true
+    },
+  "text_cfg": {
+      "context_length": 77,
+      "vocab_size": 49408,
+      "width": 1280,
+      "heads": 20,
+      "layers": 32,
+      "xattn": false,
+      "fusedLN": true
+  }
+}
+
+

legacy_eva_clip/visual.py

@@ -0,0 +1,450 @@
+# --------------------------------------------------------
+# Adapted from  https://github.com/microsoft/unilm/tree/master/beit
+# --------------------------------------------------------
+
+import os
+from functools import partial
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint
+
+try:
+    from timm.models.layers import drop_path, to_2tuple
+except:
+    from timm.layers import drop_path, to_2tuple
+
+try:
+    import xformers.ops as xops
+except ImportError:
+    xops = None
+    print("Please 'pip install xformers'")
+
+
+class PatchDropout(nn.Module):
+    """
+    https://arxiv.org/abs/2212.00794
+    """
+
+    def __init__(self, prob, exclude_first_token=True):
+        super().__init__()
+        assert 0 <= prob < 1.
+        self.prob = prob
+        self.exclude_first_token = exclude_first_token  # exclude CLS token
+        print(f"os.getenv('RoPE')={os.getenv('RoPE')}")
+
+    def forward(self, x):
+        if not self.training or self.prob == 0.:
+            return x
+
+        if self.exclude_first_token:
+            cls_tokens, x = x[:, :1], x[:, 1:]
+        else:
+            cls_tokens = torch.jit.annotate(torch.Tensor, x[:, :1])
+
+        batch = x.size()[0]
+        num_tokens = x.size()[1]
+
+        batch_indices = torch.arange(batch)
+        batch_indices = batch_indices[..., None]
+
+        keep_prob = 1 - self.prob
+        num_patches_keep = max(1, int(num_tokens * keep_prob))
+
+        rand = torch.randn(batch, num_tokens)
+        patch_indices_keep = rand.topk(num_patches_keep, dim=-1).indices
+
+        x = x[batch_indices, patch_indices_keep]
+
+        if self.exclude_first_token:
+            x = torch.cat((cls_tokens, x), dim=1)
+
+        if self.training and os.getenv('RoPE') == '1':
+            return x, patch_indices_keep
+
+        return x
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+    
+    def extra_repr(self) -> str:
+        return 'p={}'.format(self.drop_prob)
+
+
+class Mlp(nn.Module):
+    def __init__(
+        self, 
+        in_features, 
+        hidden_features=None, 
+        out_features=None, 
+        act_layer=nn.GELU, 
+        norm_layer=nn.LayerNorm, 
+        drop=0.,
+        subln=False,
+
+        ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+
+        self.ffn_ln = norm_layer(hidden_features) if subln else nn.Identity()
+
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        # x = self.drop(x)
+        # commit this for the orignal BERT implement 
+        x = self.ffn_ln(x)
+
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+class SwiGLU(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.SiLU, drop=0., 
+                norm_layer=nn.LayerNorm, subln=False):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+
+        self.w1 = nn.Linear(in_features, hidden_features)
+        self.w2 = nn.Linear(in_features, hidden_features)
+
+        self.act = act_layer()
+        self.ffn_ln = norm_layer(hidden_features) if subln else nn.Identity()
+        self.w3 = nn.Linear(hidden_features, out_features)
+        
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x1 = self.w1(x)
+        x2 = self.w2(x)
+        hidden = self.act(x1) * x2
+        x = self.ffn_ln(hidden)
+        x = self.w3(x)
+        x = self.drop(x)
+        return x
+
+class Attention(nn.Module):
+    def __init__(
+            self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.,
+            proj_drop=0., window_size=None, attn_head_dim=None, xattn=False, rope=None, subln=False, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        if attn_head_dim is not None:
+            head_dim = attn_head_dim
+        all_head_dim = head_dim * self.num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.subln = subln
+        if self.subln:
+            self.q_proj = nn.Linear(dim, all_head_dim, bias=False)
+            self.k_proj = nn.Linear(dim, all_head_dim, bias=False)
+            self.v_proj = nn.Linear(dim, all_head_dim, bias=False)
+        else:
+            if qkv_bias:
+                self.qkv = nn.Linear(dim, all_head_dim * 3, bias=True)
+            else:
+                self.qkv = nn.Linear(dim, all_head_dim * 3, bias=False)
+
+        # if qkv_bias:
+        #     self.q_bias = nn.Parameter(torch.zeros(all_head_dim))
+        #     self.v_bias = nn.Parameter(torch.zeros(all_head_dim))
+        # else:
+        #     self.q_bias = None
+        #     self.v_bias = None
+
+        self.window_size = None
+        self.relative_position_bias_table = None
+        self.relative_position_index = None
+
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.inner_attn_ln = norm_layer(all_head_dim) if subln else nn.Identity()
+        # self.proj = nn.Linear(all_head_dim, all_head_dim)
+        self.proj = nn.Linear(all_head_dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.xattn = xattn
+        self.xattn_drop = attn_drop
+
+        self.rope = rope
+
+    def forward(self, x, rel_pos_bias=None, attn_mask=None):
+        B, N, C = x.shape
+        if self.subln: 
+            q = F.linear(input=x, weight=self.q_proj.weight, bias=self.q_bias)
+            k = F.linear(input=x, weight=self.k_proj.weight, bias=None)
+            v = F.linear(input=x, weight=self.v_proj.weight, bias=self.v_bias)
+
+            q = q.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)     # B, num_heads, N, C
+            k = k.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)  
+            v = v.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3) 
+        else: 
+
+            # qkv_bias = None
+            # if self.q_bias is not None:
+            #     qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))
+            
+            # qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
+
+            qkv = self.qkv(x)
+
+            qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)   # 3, B, num_heads, N, C
+            q, k, v = qkv[0], qkv[1], qkv[2]
+
+        if self.rope:
+            q_t = q[:, :, 1:, :]
+            ro_q_t = self.rope(q_t)
+            q = torch.cat((q[:, :, :1, :], ro_q_t), -2).type_as(v)
+
+            k_t = k[:, :, 1:, :]
+            ro_k_t = self.rope(k_t)
+            k = torch.cat((k[:, :, :1, :], ro_k_t), -2).type_as(v)
+
+        if self.xattn:
+            q = q.permute(0, 2, 1, 3)   # B, num_heads, N, C -> B, N, num_heads, C
+            k = k.permute(0, 2, 1, 3)
+            v = v.permute(0, 2, 1, 3)
+
+            x = xops.memory_efficient_attention(
+                q, k, v,
+                p=self.xattn_drop,
+                scale=self.scale,
+                )
+            x = x.reshape(B, N, -1)
+            x = self.inner_attn_ln(x)
+            x = self.proj(x)
+            x = self.proj_drop(x)
+        else:
+            q = q * self.scale
+            attn = (q @ k.transpose(-2, -1))
+
+            if self.relative_position_bias_table is not None:
+                relative_position_bias = \
+                    self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
+                        self.window_size[0] * self.window_size[1] + 1,
+                        self.window_size[0] * self.window_size[1] + 1, -1)  # Wh*Ww,Wh*Ww,nH
+                relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
+                attn = attn + relative_position_bias.unsqueeze(0).type_as(attn)
+
+            if rel_pos_bias is not None:
+                attn = attn + rel_pos_bias.type_as(attn)
+
+            if attn_mask is not None:
+                attn_mask = attn_mask.bool()
+                attn = attn.masked_fill(~attn_mask[:, None, None, :], float("-inf"))
+            
+            attn = attn.softmax(dim=-1)
+            attn = self.attn_drop(attn)
+
+            x = (attn @ v).transpose(1, 2).reshape(B, N, -1)
+            x = self.inner_attn_ln(x)
+            x = self.proj(x)
+            x = self.proj_drop(x)
+        return x
+
+
+class Block(nn.Module):
+
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., init_values=None, act_layer=nn.GELU, norm_layer=nn.LayerNorm,
+                 window_size=None, attn_head_dim=None, xattn=False, rope=None, postnorm=False,
+                 subln=False, naiveswiglu=False):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
+            attn_drop=attn_drop, proj_drop=drop, window_size=window_size, attn_head_dim=attn_head_dim,
+            xattn=xattn, rope=rope, subln=subln, norm_layer=norm_layer)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+
+        if naiveswiglu:
+            self.mlp = SwiGLU(
+                in_features=dim, 
+                hidden_features=mlp_hidden_dim, 
+                subln=subln,
+                norm_layer=norm_layer,
+            )
+        else:
+            self.mlp = Mlp(
+                in_features=dim, 
+                hidden_features=mlp_hidden_dim, 
+                act_layer=act_layer,
+                subln=subln,
+                drop=drop
+            )
+
+        if init_values is not None and init_values > 0:
+            self.gamma_1 = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True)
+            self.gamma_2 = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True)
+        else:
+            self.gamma_1, self.gamma_2 = None, None
+
+        self.postnorm = postnorm
+
+    def forward(self, x, rel_pos_bias=None, attn_mask=None):
+        if self.gamma_1 is None:
+            if self.postnorm:
+                x = x + self.drop_path(self.norm1(self.attn(x, rel_pos_bias=rel_pos_bias, attn_mask=attn_mask)))
+                x = x + self.drop_path(self.norm2(self.mlp(x)))
+            else:
+                x = x + self.drop_path(self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias, attn_mask=attn_mask))
+                x = x + self.drop_path(self.mlp(self.norm2(x)))
+        else:
+            if self.postnorm:
+                x = x + self.drop_path(self.gamma_1 * self.norm1(self.attn(x, rel_pos_bias=rel_pos_bias, attn_mask=attn_mask)))
+                x = x + self.drop_path(self.gamma_2 * self.norm2(self.mlp(x)))
+            else:
+                x = x + self.drop_path(self.gamma_1 * self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias, attn_mask=attn_mask))
+                x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """ Image to Patch Embedding
+    """
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
+        self.patch_shape = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+
+    def forward(self, x, **kwargs):
+        B, C, H, W = x.shape
+        # FIXME look at relaxing size constraints
+        assert H == self.img_size[0] and W == self.img_size[1], \
+            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        x = self.proj(x).flatten(2).transpose(1, 2)
+        return x
+
+
+class EVAVisionTransformer(nn.Module):
+    """ Vision Transformer with support for patch or hybrid CNN input stage
+    """
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,
+                 num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,
+                 drop_path_rate=0., norm_layer=nn.LayerNorm, init_values=None, patch_dropout=0.,
+                 use_abs_pos_emb=True, use_rel_pos_bias=False, use_shared_rel_pos_bias=False, rope=False,
+                 use_mean_pooling=True, init_scale=0.001, grad_checkpointing=False, xattn=False, postnorm=False,
+                 pt_hw_seq_len=16, intp_freq=False, naiveswiglu=False, subln=False,
+                 ):
+        super().__init__()
+        self.image_size = img_size
+        # self.num_classes = num_classes
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+
+        self.patch_embed = PatchEmbed(
+            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        if use_abs_pos_emb:
+            self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
+        else:
+            self.pos_embed = None
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        self.rel_pos_bias = None
+        self.rope = None
+
+        self.naiveswiglu = naiveswiglu
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        self.use_rel_pos_bias = use_rel_pos_bias
+        self.blocks = nn.ModuleList([
+            Block(
+                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer,
+                init_values=init_values, window_size=self.patch_embed.patch_shape if use_rel_pos_bias else None,
+                xattn=xattn, rope=self.rope, postnorm=postnorm, subln=subln, naiveswiglu=naiveswiglu)
+            for i in range(depth)])
+
+        # setting a patch_dropout of 0. would mean it is disabled and this function would be the identity fn
+        self.patch_dropout = PatchDropout(patch_dropout) if patch_dropout > 0. else nn.Identity()
+
+        self.grad_checkpointing = grad_checkpointing
+
+
+    def get_num_layers(self):
+        return len(self.blocks)
+    
+    def lock(self, unlocked_groups=0, freeze_bn_stats=False):
+        assert unlocked_groups == 0, 'partial locking not currently supported for this model'
+        for param in self.parameters():
+            param.requires_grad = False
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.grad_checkpointing = enable
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'pos_embed', 'cls_token'}
+
+
+    def forward_features(self, x):
+        x = self.patch_embed(x)
+        batch_size, seq_len, _ = x.size()
+
+        cls_tokens = self.cls_token.expand(batch_size, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
+        x = torch.cat((cls_tokens, x), dim=1)
+        if self.pos_embed is not None:
+            x = x + self.pos_embed
+        x = self.pos_drop(x)
+
+        # a patch_dropout of 0. would mean it is disabled and this function would do nothing but return what was passed in
+        if os.getenv('RoPE') == '1':
+            if self.training and not isinstance(self.patch_dropout, nn.Identity):
+                x, patch_indices_keep = self.patch_dropout(x)
+                self.rope.forward = partial(self.rope.forward, patch_indices_keep=patch_indices_keep)
+            else:
+                self.rope.forward = partial(self.rope.forward, patch_indices_keep=None)
+                x = self.patch_dropout(x)
+        else:
+            x = self.patch_dropout(x)
+
+        rel_pos_bias = None
+
+        for blk in self.blocks:
+            if self.grad_checkpointing:
+                x = checkpoint(blk, x, (rel_pos_bias,))
+            else:
+                x = blk(x, rel_pos_bias=rel_pos_bias)
+
+        return x
+
+    def forward(self, x):
+
+        """
+        :return:
+            forward_features function returns raw features of ViT,
+            forward with return_all_features returns normalized features of ViT
+        :param x:
+        :param return_all_features:
+        """
+
+        features = self.forward_features(x)  # [B, n_patch, C]
+
+        return features

legacy_eva_vit.py

@@ -0,0 +1,815 @@
+"""
+# Adapted from https://github.com/baaivision/EVA/tree/master/EVA-CLIP
+"""
+
+from math import pi
+import os
+from pathlib import Path
+from typing import Optional
+import importlib.util
+
+import torch
+from torch import nn
+from einops import rearrange, repeat
+import logging
+
+_hub_spec = importlib.util.find_spec("huggingface_hub")
+if _hub_spec is not None:
+    from huggingface_hub import snapshot_download  # type: ignore
+else:  # pragma: no cover - optional dependency
+    snapshot_download = None
+
+_DEFAULT_EVA_REPO = "jiuhai/eva_clip_vision_tower"
+XFORMERS_AVAILABLE = False  # populated later once xops import resolves
+_XFORMERS_WARNING_EMITTED = False
+
+
+def _resolve_vision_checkpoint_path(vision_tower_pretrained: Optional[str]) -> str:
+    """Determine where to load the EVA visual tower weights from."""
+
+    candidate_files: list[Path] = []
+    if vision_tower_pretrained:
+        supplied = Path(vision_tower_pretrained)
+        if supplied.is_file():
+            candidate_files.append(supplied)
+        elif supplied.is_dir():
+            for filename in (
+                "pytorch_model.bin",
+                "model.safetensors",
+                "visual.pth",
+                "visual.bin",
+            ):
+                candidate = supplied / filename
+                if candidate.exists():
+                    candidate_files.append(candidate)
+                    break
+
+    if candidate_files:
+        return str(candidate_files[0])
+
+    if snapshot_download is None:
+        raise FileNotFoundError(
+            "EVA vision weights not found locally and huggingface_hub is unavailable. "
+            "Provide --vision-tower-path pointing to a directory containing the checkpoint."
+        )
+
+    cache_dir = Path(snapshot_download(repo_id=_DEFAULT_EVA_REPO))
+    default_file = cache_dir / "pytorch_model.bin"
+    if not default_file.exists():
+        raise FileNotFoundError(f"Default EVA checkpoint not found at {default_file}.")
+    return str(default_file)
+
+
+def broadcat(tensors, dim=-1):
+    num_tensors = len(tensors)
+    shape_lens = set(list(map(lambda t: len(t.shape), tensors)))
+    assert len(shape_lens) == 1, "tensors must all have the same number of dimensions"
+    shape_len = list(shape_lens)[0]
+    dim = (dim + shape_len) if dim < 0 else dim
+    dims = list(zip(*map(lambda t: list(t.shape), tensors)))
+    expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim]
+    assert all([*map(lambda t: len(set(t[1])) <= 2, expandable_dims)]), "invalid dimensions for broadcastable concatentation"
+    max_dims = list(map(lambda t: (t[0], max(t[1])), expandable_dims))
+    expanded_dims = list(map(lambda t: (t[0], (t[1],) * num_tensors), max_dims))
+    expanded_dims.insert(dim, (dim, dims[dim]))
+    expandable_shapes = list(zip(*map(lambda t: t[1], expanded_dims)))
+    tensors = list(map(lambda t: t[0].expand(*t[1]), zip(tensors, expandable_shapes)))
+    return torch.cat(tensors, dim=dim)
+
+
+def rotate_half(x):
+    x = rearrange(x, "... (d r) -> ... d r", r=2)
+    x1, x2 = x.unbind(dim=-1)
+    x = torch.stack((-x2, x1), dim=-1)
+    return rearrange(x, "... d r -> ... (d r)")
+
+
+class VisionRotaryEmbeddingFast(nn.Module):
+    def __init__(self, dim, pt_seq_len, ft_seq_len=None, custom_freqs=None, freqs_for="lang", theta=10000, max_freq=10, num_freqs=1, patch_dropout=0.0):
+        super().__init__()
+        if custom_freqs:
+            freqs = custom_freqs
+        elif freqs_for == "lang":
+            freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+        elif freqs_for == "pixel":
+            freqs = torch.linspace(1.0, max_freq / 2, dim // 2) * pi
+        elif freqs_for == "constant":
+            freqs = torch.ones(num_freqs).float()
+        else:
+            raise ValueError(f"unknown modality {freqs_for}")
+
+        if ft_seq_len is None:
+            ft_seq_len = pt_seq_len
+        t = torch.arange(ft_seq_len) / ft_seq_len * pt_seq_len
+
+        freqs = torch.einsum("..., f -> ... f", t, freqs)
+        freqs = repeat(freqs, "... n -> ... (n r)", r=2)
+        freqs = broadcat((freqs[:, None, :], freqs[None, :, :]), dim=-1)
+
+        freqs_cos = freqs.cos().view(-1, freqs.shape[-1])
+        freqs_sin = freqs.sin().view(-1, freqs.shape[-1])
+
+        self.patch_dropout = patch_dropout
+
+        self.register_buffer("freqs_cos", freqs_cos)
+        self.register_buffer("freqs_sin", freqs_sin)
+
+        logging.info(f"Shape of rope freq: {self.freqs_cos.shape}")
+
+    def forward(self, t, patch_indices_keep=None):
+        if patch_indices_keep is not None:
+            batch = t.size()[0]
+            batch_indices = torch.arange(batch)
+            batch_indices = batch_indices[..., None]
+
+            freqs_cos = repeat(self.freqs_cos, "i j -> n i m j", n=t.shape[0], m=t.shape[1])
+            freqs_sin = repeat(self.freqs_sin, "i j -> n i m j", n=t.shape[0], m=t.shape[1])
+
+            freqs_cos = freqs_cos[batch_indices, patch_indices_keep]
+            freqs_cos = rearrange(freqs_cos, "n i m j -> n m i j")
+            freqs_sin = freqs_sin[batch_indices, patch_indices_keep]
+            freqs_sin = rearrange(freqs_sin, "n i m j -> n m i j")
+
+            return t * freqs_cos + rotate_half(t) * freqs_sin
+
+        return t * self.freqs_cos + rotate_half(t) * self.freqs_sin
+
+
+class LayerNorm(nn.LayerNorm):
+    """Subclass torch's LayerNorm (with cast back to input dtype)."""
+
+    def forward(self, x: torch.Tensor):
+        orig_type = x.dtype
+        x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        return x.to(orig_type)
+
+
+class PatchDropout(nn.Module):
+    """
+    https://arxiv.org/abs/2212.00794
+    """
+
+    def __init__(self, prob, exclude_first_token=True):
+        super().__init__()
+        assert 0 <= prob < 1.
+        self.prob = prob
+        self.exclude_first_token = exclude_first_token  # exclude CLS token
+        print(f"os.getenv('RoPE')={os.getenv('RoPE')}")
+
+    def forward(self, x):
+        if not self.training or self.prob == 0.:
+            return x
+
+        if self.exclude_first_token:
+            cls_tokens, x = x[:, :1], x[:, 1:]
+        else:
+            cls_tokens = torch.jit.annotate(torch.Tensor, x[:, :1])
+
+        batch = x.size()[0]
+        num_tokens = x.size()[1]
+
+        batch_indices = torch.arange(batch)
+        batch_indices = batch_indices[..., None]
+
+        keep_prob = 1 - self.prob
+        num_patches_keep = max(1, int(num_tokens * keep_prob))
+
+        rand = torch.randn(batch, num_tokens)
+        patch_indices_keep = rand.topk(num_patches_keep, dim=-1).indices
+
+        x = x[batch_indices, patch_indices_keep]
+
+        if self.exclude_first_token:
+            x = torch.cat((cls_tokens, x), dim=1)
+
+        if self.training and os.getenv('RoPE') == '1':
+            return x, patch_indices_keep
+
+        return x
+
+
+# --------------------------------------------------------
+# Adapted from  https://github.com/microsoft/unilm/tree/master/beit
+# --------------------------------------------------------
+import math
+import os
+import torch.nn as nn
+import torch.nn.functional as F
+
+try:
+    from timm.models.layers import drop_path, to_2tuple, trunc_normal_
+except:
+    from timm.layers import drop_path, to_2tuple, trunc_normal_
+
+if os.getenv("ENV_TYPE") == "deepspeed":
+    try:
+        from deepspeed.runtime.activation_checkpointing.checkpointing import checkpoint
+    except:
+        from torch.utils.checkpoint import checkpoint
+else:
+    from torch.utils.checkpoint import checkpoint
+
+try:
+    import xformers.ops as xops
+    XFORMERS_AVAILABLE = True
+except ImportError:
+    xops = None
+    # print("Please 'pip install xformers'")
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+    
+    def extra_repr(self) -> str:
+        return 'p={}'.format(self.drop_prob)
+
+
+class Mlp(nn.Module):
+    def __init__(
+        self, 
+        in_features, 
+        hidden_features=None, 
+        out_features=None, 
+        act_layer=nn.GELU, 
+        norm_layer=nn.LayerNorm, 
+        drop=0.,
+        subln=False,
+
+        ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+
+        self.ffn_ln = norm_layer(hidden_features) if subln else nn.Identity()
+
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        # x = self.drop(x)
+        # commit this for the orignal BERT implement 
+        x = self.ffn_ln(x)
+
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.SiLU, drop=0., 
+                norm_layer=nn.LayerNorm, subln=False):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+
+        self.w1 = nn.Linear(in_features, hidden_features)
+        self.w2 = nn.Linear(in_features, hidden_features)
+
+        self.act = act_layer()
+        self.ffn_ln = norm_layer(hidden_features) if subln else nn.Identity()
+        self.w3 = nn.Linear(hidden_features, out_features)
+        
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x1 = self.w1(x)
+        x2 = self.w2(x)
+        hidden = self.act(x1) * x2
+        x = self.ffn_ln(hidden)
+        x = self.w3(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(
+            self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.,
+            proj_drop=0., window_size=None, attn_head_dim=None, xattn=False, rope=None, subln=False, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        if attn_head_dim is not None:
+            head_dim = attn_head_dim
+        all_head_dim = head_dim * self.num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.subln = subln
+        if self.subln:
+            self.q_proj = nn.Linear(dim, all_head_dim, bias=False)
+            self.k_proj = nn.Linear(dim, all_head_dim, bias=False)
+            self.v_proj = nn.Linear(dim, all_head_dim, bias=False)
+        else:
+            if qkv_bias:
+                self.qkv = nn.Linear(dim, all_head_dim * 3, bias=True)
+            else:
+                self.qkv = nn.Linear(dim, all_head_dim * 3, bias=False)
+
+        # if qkv_bias:
+        #     self.q_bias = nn.Parameter(torch.zeros(all_head_dim))
+        #     self.v_bias = nn.Parameter(torch.zeros(all_head_dim))
+        # else:
+        #     self.q_bias = None
+        #     self.v_bias = None
+
+        self.window_size = None
+        self.relative_position_bias_table = None
+        self.relative_position_index = None
+
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.inner_attn_ln = norm_layer(all_head_dim) if subln else nn.Identity()
+        # self.proj = nn.Linear(all_head_dim, all_head_dim)
+        self.proj = nn.Linear(all_head_dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.xattn = xattn and XFORMERS_AVAILABLE
+        self.xattn_drop = attn_drop
+        global _XFORMERS_WARNING_EMITTED
+        if xattn and not self.xattn and not _XFORMERS_WARNING_EMITTED:
+            logging.warning("xformers.ops not available; falling back to standard attention for EVA-CLIP vision tower.")
+            _XFORMERS_WARNING_EMITTED = True
+
+        self.rope = rope
+
+    def forward(self, x, rel_pos_bias=None, attn_mask=None):
+        B, N, C = x.shape
+        if self.subln: 
+            q = F.linear(input=x, weight=self.q_proj.weight, bias=self.q_bias)
+            k = F.linear(input=x, weight=self.k_proj.weight, bias=None)
+            v = F.linear(input=x, weight=self.v_proj.weight, bias=self.v_bias)
+
+            q = q.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)     # B, num_heads, N, C
+            k = k.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)  
+            v = v.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3) 
+        else: 
+
+            # qkv_bias = None
+            # if self.q_bias is not None:
+            #     qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))
+            
+            # qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
+
+            qkv = self.qkv(x)
+
+            qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)   # 3, B, num_heads, N, C
+            q, k, v = qkv[0], qkv[1], qkv[2]
+
+        if self.rope:
+            q_t = q[:, :, 1:, :]
+            ro_q_t = self.rope(q_t)
+            q = torch.cat((q[:, :, :1, :], ro_q_t), -2).type_as(v)
+
+            k_t = k[:, :, 1:, :]
+            ro_k_t = self.rope(k_t)
+            k = torch.cat((k[:, :, :1, :], ro_k_t), -2).type_as(v)
+
+        if self.xattn:
+            q = q.permute(0, 2, 1, 3)   # B, num_heads, N, C -> B, N, num_heads, C
+            k = k.permute(0, 2, 1, 3)
+            v = v.permute(0, 2, 1, 3)
+
+            x = xops.memory_efficient_attention(
+                q, k, v,
+                p=self.xattn_drop,
+                scale=self.scale,
+                )
+            x = x.reshape(B, N, -1)
+            x = self.inner_attn_ln(x)
+            x = self.proj(x)
+            x = self.proj_drop(x)
+        else:
+            q = q * self.scale
+            attn = (q @ k.transpose(-2, -1))
+
+            if self.relative_position_bias_table is not None:
+                relative_position_bias = \
+                    self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
+                        self.window_size[0] * self.window_size[1] + 1,
+                        self.window_size[0] * self.window_size[1] + 1, -1)  # Wh*Ww,Wh*Ww,nH
+                relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
+                attn = attn + relative_position_bias.unsqueeze(0).type_as(attn)
+
+            if rel_pos_bias is not None:
+                attn = attn + rel_pos_bias.type_as(attn)
+
+            if attn_mask is not None:
+                attn_mask = attn_mask.bool()
+                attn = attn.masked_fill(~attn_mask[:, None, None, :], float("-inf"))
+            
+            attn = attn.softmax(dim=-1)
+            attn = self.attn_drop(attn)
+
+            x = (attn @ v).transpose(1, 2).reshape(B, N, -1)
+            x = self.inner_attn_ln(x)
+            x = self.proj(x)
+            x = self.proj_drop(x)
+        return x
+
+
+class Block(nn.Module):
+
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., init_values=None, act_layer=nn.GELU, norm_layer=nn.LayerNorm,
+                 window_size=None, attn_head_dim=None, xattn=False, rope=None, postnorm=False,
+                 subln=False, naiveswiglu=False):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
+            attn_drop=attn_drop, proj_drop=drop, window_size=window_size, attn_head_dim=attn_head_dim,
+            xattn=xattn, rope=rope, subln=subln, norm_layer=norm_layer)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+
+        if naiveswiglu:
+            self.mlp = SwiGLU(
+                in_features=dim, 
+                hidden_features=mlp_hidden_dim, 
+                subln=subln,
+                norm_layer=norm_layer,
+            )
+        else:
+            self.mlp = Mlp(
+                in_features=dim, 
+                hidden_features=mlp_hidden_dim, 
+                act_layer=act_layer,
+                subln=subln,
+                drop=drop
+            )
+
+        if init_values is not None and init_values > 0:
+            self.gamma_1 = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True)
+            self.gamma_2 = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True)
+        else:
+            self.gamma_1, self.gamma_2 = None, None
+
+        self.postnorm = postnorm
+
+    def forward(self, x, rel_pos_bias=None, attn_mask=None):
+        if self.gamma_1 is None:
+            if self.postnorm:
+                x = x + self.drop_path(self.norm1(self.attn(x, rel_pos_bias=rel_pos_bias, attn_mask=attn_mask)))
+                x = x + self.drop_path(self.norm2(self.mlp(x)))
+            else:
+                x = x + self.drop_path(self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias, attn_mask=attn_mask))
+                x = x + self.drop_path(self.mlp(self.norm2(x)))
+        else:
+            if self.postnorm:
+                x = x + self.drop_path(self.gamma_1 * self.norm1(self.attn(x, rel_pos_bias=rel_pos_bias, attn_mask=attn_mask)))
+                x = x + self.drop_path(self.gamma_2 * self.norm2(self.mlp(x)))
+            else:
+                x = x + self.drop_path(self.gamma_1 * self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias, attn_mask=attn_mask))
+                x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """ Image to Patch Embedding
+    """
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
+        self.patch_shape = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+
+    def forward(self, x, **kwargs):
+        B, C, H, W = x.shape
+        # FIXME look at relaxing size constraints
+        assert H == self.img_size[0] and W == self.img_size[1], \
+            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        x = self.proj(x).flatten(2).transpose(1, 2)
+        return x
+
+
+class RelativePositionBias(nn.Module):
+
+    def __init__(self, window_size, num_heads):
+        super().__init__()
+        self.window_size = window_size
+        self.num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3
+        self.relative_position_bias_table = nn.Parameter(torch.zeros(self.num_relative_distance, num_heads))  # 2*Wh-1 * 2*Ww-1, nH
+        # cls to token & token 2 cls & cls to cls
+
+        # get pair-wise relative position index for each token inside the window
+        coords_h = torch.arange(window_size[0])
+        coords_w = torch.arange(window_size[1])
+        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
+        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
+        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
+        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
+        relative_coords[:, :, 0] += window_size[0] - 1  # shift to start from 0
+        relative_coords[:, :, 1] += window_size[1] - 1
+        relative_coords[:, :, 0] *= 2 * window_size[1] - 1
+        relative_position_index = torch.zeros(size=(window_size[0] * window_size[1] + 1,) * 2, dtype=relative_coords.dtype)
+        relative_position_index[1:, 1:] = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
+        relative_position_index[0, 0:] = self.num_relative_distance - 3
+        relative_position_index[0:, 0] = self.num_relative_distance - 2
+        relative_position_index[0, 0] = self.num_relative_distance - 1
+
+        self.register_buffer("relative_position_index", relative_position_index)
+
+    def forward(self):
+        relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(self.window_size[0] * self.window_size[1] + 1, self.window_size[0] * self.window_size[1] + 1, -1)  # Wh*Ww,Wh*Ww,nH
+        return relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
+
+
+class EVAVisionTransformer(nn.Module):
+    """Vision Transformer with support for patch or hybrid CNN input stage"""
+
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,
+                 num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,
+                 drop_path_rate=0., norm_layer=nn.LayerNorm, init_values=None, patch_dropout=0.,
+                 use_abs_pos_emb=True, use_rel_pos_bias=False, use_shared_rel_pos_bias=False, rope=False,
+                 use_mean_pooling=True, init_scale=0.001, grad_checkpointing=False, xattn=False, postnorm=False,
+                 pt_hw_seq_len=16, intp_freq=False, naiveswiglu=False, subln=False,
+                 ):
+        super().__init__()
+        self.image_size = img_size
+        # self.num_classes = num_classes
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+
+        self.patch_embed = PatchEmbed(
+            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        if use_abs_pos_emb:
+            self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
+        else:
+            self.pos_embed = None
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        self.rel_pos_bias = None
+        self.rope = None
+
+        self.naiveswiglu = naiveswiglu
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        self.use_rel_pos_bias = use_rel_pos_bias
+        self.blocks = nn.ModuleList([
+            Block(
+                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer,
+                init_values=init_values, window_size=self.patch_embed.patch_shape if use_rel_pos_bias else None,
+                xattn=xattn, rope=self.rope, postnorm=postnorm, subln=subln, naiveswiglu=naiveswiglu)
+            for i in range(depth)])
+
+        # setting a patch_dropout of 0. would mean it is disabled and this function would be the identity fn
+        self.patch_dropout = PatchDropout(patch_dropout) if patch_dropout > 0. else nn.Identity()
+
+        self.grad_checkpointing = grad_checkpointing
+
+    def fix_init_weight(self):
+        def rescale(param, layer_id):
+            param.div_(math.sqrt(2.0 * layer_id))
+
+        for layer_id, layer in enumerate(self.blocks):
+            rescale(layer.attn.proj.weight.data, layer_id + 1)
+            if self.naiveswiglu:
+                rescale(layer.mlp.w3.weight.data, layer_id + 1)
+            else:
+                rescale(layer.mlp.fc2.weight.data, layer_id + 1)
+
+    def get_cast_dtype(self) -> torch.dtype:
+        return self.blocks[0].mlp.fc2.weight.dtype
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def get_num_layers(self):
+        return len(self.blocks)
+
+    def lock(self, unlocked_groups=0, freeze_bn_stats=False):
+        assert unlocked_groups == 0, "partial locking not currently supported for this model"
+        for param in self.parameters():
+            param.requires_grad = False
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.grad_checkpointing = enable
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {"pos_embed", "cls_token"}
+
+    def get_classifier(self):
+        return self.head
+
+    def reset_classifier(self, num_classes, global_pool=""):
+        self.num_classes = num_classes
+        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+    def forward_features(self, x):
+        x = self.patch_embed(x)
+        batch_size, seq_len, _ = x.size()
+
+        cls_tokens = self.cls_token.expand(batch_size, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
+        x = torch.cat((cls_tokens, x), dim=1)
+        if self.pos_embed is not None:
+            x = x + self.pos_embed
+        x = self.pos_drop(x)
+
+        # a patch_dropout of 0. would mean it is disabled and this function would do nothing but return what was passed in
+        if os.getenv('RoPE') == '1':
+            if self.training and not isinstance(self.patch_dropout, nn.Identity):
+                x, patch_indices_keep = self.patch_dropout(x)
+                self.rope.forward = partial(self.rope.forward, patch_indices_keep=patch_indices_keep)
+            else:
+                self.rope.forward = partial(self.rope.forward, patch_indices_keep=None)
+                x = self.patch_dropout(x)
+        else:
+            x = self.patch_dropout(x)
+
+        rel_pos_bias = None
+
+        for blk in self.blocks:
+            if self.grad_checkpointing:
+                x = checkpoint(blk, x, (rel_pos_bias,))
+            else:
+                x = blk(x, rel_pos_bias=rel_pos_bias)
+
+        return x
+
+    def forward(self, x, return_all_features=False):
+
+        """
+        :return:
+            forward_features function returns raw features of ViT,
+            forward with return_all_features returns normalized features of ViT
+        :param x:
+        :param return_all_features:
+        """
+
+        features = self.forward_features(x)  # [B, n_patch, C]
+        return features
+
+
+def load_state_dict(checkpoint_path: str, map_location: str = "cpu", model_key: str = "model|module|state_dict", is_openai: bool = False, skip_list: list = []):
+    if is_openai:
+        model = torch.jit.load(checkpoint_path, map_location="cpu").eval()
+        state_dict = model.state_dict()
+        for key in ["input_resolution", "context_length", "vocab_size"]:
+            state_dict.pop(key, None)
+    else:
+        checkpoint = torch.load(checkpoint_path, map_location=map_location)
+        for mk in model_key.split("|"):
+            if isinstance(checkpoint, dict) and mk in checkpoint:
+                state_dict = checkpoint[mk]
+                break
+            else:
+                state_dict = checkpoint
+        if next(iter(state_dict.items()))[0].startswith("module"):
+            state_dict = {k[7:]: v for k, v in state_dict.items()}
+
+    for k in skip_list:
+        if k in list(state_dict.keys()):
+            logging.info(f"Removing key {k} from pretrained checkpoint")
+            del state_dict[k]
+
+    if os.getenv("RoPE") == "1":
+        for k in list(state_dict.keys()):
+            if "freqs_cos" in k or "freqs_sin" in k:
+                del state_dict[k]
+    return state_dict
+
+
+def load_clip_visual_state_dict(checkpoint_path: str, map_location: str = "cpu", is_openai: bool = False, skip_list: list = []):
+    state_dict = load_state_dict(checkpoint_path, map_location=map_location, is_openai=is_openai, skip_list=skip_list)
+    # for k in list(state_dict.keys()):
+    #     if not k.startswith("visual."):
+    #         del state_dict[k]
+    # for k in list(state_dict.keys()):
+    #     if k.startswith("visual."):
+    #         new_k = k[7:]
+    #         state_dict[new_k] = state_dict[k]
+    #         del state_dict[k]
+    return state_dict
+
+
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+
+try:
+    from apex.normalization import FusedLayerNorm
+except:
+    FusedLayerNorm = LayerNorm
+    # print("Please build and install Nvidia apex package with option '--cuda_ext' according to https://github.com/NVIDIA/apex#from-source .")
+
+
+@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  # layer scale initial value
+    patch_dropout: float = 0.0  # what fraction of patches to dropout during training (0 would mean disabled and no patches dropped) - 0.5 to 0.75 recommended in the paper for optimal results
+    global_average_pool: bool = False  # whether to global average pool the last embedding layer, instead of using CLS token (https://arxiv.org/abs/2205.01580)
+    drop_path_rate: Optional[float] = None  # drop path rate
+    timm_model_name: str = None  # a valid model name overrides layers, width, patch_size
+    timm_model_pretrained: bool = False  # use (imagenet) pretrained weights for named model
+    timm_pool: str = "avg"  # feature pooling for timm model ('abs_attn', 'rot_attn', 'avg', '')
+    timm_proj: str = "linear"  # linear projection for timm model output ('linear', 'mlp', '')
+    timm_proj_bias: bool = False  # enable bias final projection
+    eva_model_name: str = None  # a valid eva model name overrides layers, width, patch_size
+    qkv_bias: bool = True
+    fusedLN: bool = False
+    xattn: bool = False
+    postnorm: bool = False
+    rope: bool = False
+    pt_hw_seq_len: int = 16  # 224/14
+    intp_freq: bool = False
+    naiveswiglu: bool = False
+    subln: bool = False
+
+
+def create_norm_layer_factory(use_fused_ln, eps=1e-6):
+    # Otherwise, use the standard LayerNorm
+    return lambda num_features: nn.LayerNorm(num_features, eps=eps)
+
+
+def _build_vision_tower(vision_tower_path: str, embed_dim: int, vision_cfg: CLIPVisionCfg, **kwargs):
+    if isinstance(vision_cfg, dict):
+        vision_cfg = CLIPVisionCfg(**vision_cfg)
+
+    if vision_cfg.eva_model_name:
+        vision_heads = vision_cfg.width // vision_cfg.head_width
+        # Determine the appropriate norm layer factory based on the configuration
+        norm_layer_factory = create_norm_layer_factory(vision_cfg.fusedLN, eps=1e-6)
+
+        # breakpoint()
+        visual = EVAVisionTransformer(
+            img_size=vision_cfg.image_size,
+            patch_size=vision_cfg.patch_size,
+            num_classes=embed_dim,
+            use_mean_pooling=vision_cfg.global_average_pool,  # False
+            init_values=vision_cfg.ls_init_value,
+            patch_dropout=vision_cfg.patch_dropout,
+            embed_dim=vision_cfg.width,
+            depth=vision_cfg.layers,
+            num_heads=vision_heads,
+            mlp_ratio=vision_cfg.mlp_ratio,
+            qkv_bias=vision_cfg.qkv_bias,
+            drop_path_rate=vision_cfg.drop_path_rate,
+            norm_layer=norm_layer_factory,
+            xattn=vision_cfg.xattn,
+            rope=vision_cfg.rope,
+            postnorm=vision_cfg.postnorm,
+            pt_hw_seq_len=vision_cfg.pt_hw_seq_len,  # 224/14
+            intp_freq=vision_cfg.intp_freq,
+            naiveswiglu=vision_cfg.naiveswiglu,
+            subln=vision_cfg.subln,
+        )
+        # breakpoint()
+        state_dict = load_clip_visual_state_dict(vision_tower_path)
+        incompatible_keys = visual.load_state_dict(state_dict, strict=False)
+        print("EVA-CLIP incompatible_keys:", incompatible_keys)
+
+    return visual
+
+
+class EVAEncoderWrapper(nn.Module):
+    def __init__(self, vision_tower_pretrained, config):
+        super(EVAEncoderWrapper, self).__init__()
+        self.config = config
+        vision_tower_path = _resolve_vision_checkpoint_path(vision_tower_pretrained)
+        self.config["vision_tower_path"] = vision_tower_path
+        self.vision_tower_path = vision_tower_path
+        self.model = _build_vision_tower(**self.config)
+
+    def forward(self, image, **kwargs):
+        encode = self.model(image, return_all_features=True)[:, 1:, :]  # remove the CLS token
+        return encode
+
+    @property
+    def dtype(self):
+        return list(self.parameters())[-1].dtype
+
+    @property
+    def device(self):
+        return list(self.parameters())[-1].device

model-00001-of-00009.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d84ba46c42ae442da367bbf89fd91b99f031abaae7702e97af614830a00b7dc6
+size 2075135128

model-00002-of-00009.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bd471ff827bd24f908cff1a974848cd0d59601745dd7fab6ae37a67a6af2e1aa
+size 2063504224

model-00003-of-00009.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:946270642a7cca24ea645c93045607423862753c1ed94e85d5d49ffc1cab8caa
+size 2122242576

model-00004-of-00009.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:891c19b33d0b85c8c3c1a6c679781947152e5ada83f368fdb670928db165ef46
+size 2063544008

model-00005-of-00009.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9e7de9bce95a2a80881f20a1cc191c1a77ef28cbc8d210b7c97538bb946e6c56
+size 2063504264

model-00006-of-00009.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:21ccc9ffbffc60786ae93e1d412273cbd9d95a1e3faaa45536f7f7dd763364cd
+size 2122242576

model-00007-of-00009.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d00a946cb9f25eb2ae2e82390c29a74f2f264f98cd0bb5ecd7e7404ddfa35605
+size 2063544008

model-00008-of-00009.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:65e7dba74ace0cf4fa91135decf5628306528dacca00dfc6670d22cca78c4f69
+size 2063504264

model-00009-of-00009.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1a2a80b0493b0a1da9b766beed290defea69acc068a0ae21ad04486844930667
+size 763686520

model.safetensors.index.json

@@ -0,0 +1,779 @@
+{
+  "metadata": {
+    "total_size": 17400833024
+  },
+  "weight_map": {
+    "cls_token": "model-00001-of-00009.safetensors",
+    "pos_embed": "model-00001-of-00009.safetensors",
+    "patch_embed.proj.weight": "model-00001-of-00009.safetensors",
+    "patch_embed.proj.bias": "model-00001-of-00009.safetensors",
+    "blocks.0.norm1.weight": "model-00001-of-00009.safetensors",
+    "blocks.0.norm1.bias": "model-00001-of-00009.safetensors",
+    "blocks.0.attn.qkv.weight": "model-00001-of-00009.safetensors",
+    "blocks.0.attn.qkv.bias": "model-00001-of-00009.safetensors",
+    "blocks.0.attn.proj.weight": "model-00001-of-00009.safetensors",
+    "blocks.0.attn.proj.bias": "model-00001-of-00009.safetensors",
+    "blocks.0.norm2.weight": "model-00001-of-00009.safetensors",
+    "blocks.0.norm2.bias": "model-00001-of-00009.safetensors",
+    "blocks.0.mlp.fc1.weight": "model-00001-of-00009.safetensors",
+    "blocks.0.mlp.fc1.bias": "model-00001-of-00009.safetensors",
+    "blocks.0.mlp.fc2.weight": "model-00001-of-00009.safetensors",
+    "blocks.0.mlp.fc2.bias": "model-00001-of-00009.safetensors",
+    "blocks.1.norm1.weight": "model-00001-of-00009.safetensors",
+    "blocks.1.norm1.bias": "model-00001-of-00009.safetensors",
+    "blocks.1.attn.qkv.weight": "model-00001-of-00009.safetensors",
+    "blocks.1.attn.qkv.bias": "model-00001-of-00009.safetensors",
+    "blocks.1.attn.proj.weight": "model-00001-of-00009.safetensors",
+    "blocks.1.attn.proj.bias": "model-00001-of-00009.safetensors",
+    "blocks.1.norm2.weight": "model-00001-of-00009.safetensors",
+    "blocks.1.norm2.bias": "model-00001-of-00009.safetensors",
+    "blocks.1.mlp.fc1.weight": "model-00001-of-00009.safetensors",
+    "blocks.1.mlp.fc1.bias": "model-00001-of-00009.safetensors",
+    "blocks.1.mlp.fc2.weight": "model-00001-of-00009.safetensors",
+    "blocks.1.mlp.fc2.bias": "model-00001-of-00009.safetensors",
+    "blocks.2.norm1.weight": "model-00001-of-00009.safetensors",
+    "blocks.2.norm1.bias": "model-00001-of-00009.safetensors",
+    "blocks.2.attn.qkv.weight": "model-00001-of-00009.safetensors",
+    "blocks.2.attn.qkv.bias": "model-00001-of-00009.safetensors",
+    "blocks.2.attn.proj.weight": "model-00001-of-00009.safetensors",
+    "blocks.2.attn.proj.bias": "model-00001-of-00009.safetensors",
+    "blocks.2.norm2.weight": "model-00001-of-00009.safetensors",
+    "blocks.2.norm2.bias": "model-00001-of-00009.safetensors",
+    "blocks.2.mlp.fc1.weight": "model-00001-of-00009.safetensors",
+    "blocks.2.mlp.fc1.bias": "model-00001-of-00009.safetensors",
+    "blocks.2.mlp.fc2.weight": "model-00001-of-00009.safetensors",
+    "blocks.2.mlp.fc2.bias": "model-00001-of-00009.safetensors",
+    "blocks.3.norm1.weight": "model-00001-of-00009.safetensors",
+    "blocks.3.norm1.bias": "model-00001-of-00009.safetensors",
+    "blocks.3.attn.qkv.weight": "model-00001-of-00009.safetensors",
+    "blocks.3.attn.qkv.bias": "model-00001-of-00009.safetensors",
+    "blocks.3.attn.proj.weight": "model-00001-of-00009.safetensors",
+    "blocks.3.attn.proj.bias": "model-00001-of-00009.safetensors",
+    "blocks.3.norm2.weight": "model-00001-of-00009.safetensors",
+    "blocks.3.norm2.bias": "model-00001-of-00009.safetensors",
+    "blocks.3.mlp.fc1.weight": "model-00001-of-00009.safetensors",
+    "blocks.3.mlp.fc1.bias": "model-00001-of-00009.safetensors",
+    "blocks.3.mlp.fc2.weight": "model-00001-of-00009.safetensors",
+    "blocks.3.mlp.fc2.bias": "model-00001-of-00009.safetensors",
+    "blocks.4.norm1.weight": "model-00001-of-00009.safetensors",
+    "blocks.4.norm1.bias": "model-00001-of-00009.safetensors",
+    "blocks.4.attn.qkv.weight": "model-00001-of-00009.safetensors",
+    "blocks.4.attn.qkv.bias": "model-00001-of-00009.safetensors",
+    "blocks.4.attn.proj.weight": "model-00001-of-00009.safetensors",
+    "blocks.4.attn.proj.bias": "model-00001-of-00009.safetensors",
+    "blocks.4.norm2.weight": "model-00001-of-00009.safetensors",
+    "blocks.4.norm2.bias": "model-00001-of-00009.safetensors",
+    "blocks.4.mlp.fc1.weight": "model-00001-of-00009.safetensors",
+    "blocks.4.mlp.fc1.bias": "model-00001-of-00009.safetensors",
+    "blocks.4.mlp.fc2.weight": "model-00001-of-00009.safetensors",
+    "blocks.4.mlp.fc2.bias": "model-00001-of-00009.safetensors",
+    "blocks.5.norm1.weight": "model-00001-of-00009.safetensors",
+    "blocks.5.norm1.bias": "model-00001-of-00009.safetensors",
+    "blocks.5.attn.qkv.weight": "model-00001-of-00009.safetensors",
+    "blocks.5.attn.qkv.bias": "model-00001-of-00009.safetensors",
+    "blocks.5.attn.proj.weight": "model-00001-of-00009.safetensors",
+    "blocks.5.attn.proj.bias": "model-00001-of-00009.safetensors",
+    "blocks.5.norm2.weight": "model-00001-of-00009.safetensors",
+    "blocks.5.norm2.bias": "model-00001-of-00009.safetensors",
+    "blocks.5.mlp.fc1.weight": "model-00001-of-00009.safetensors",
+    "blocks.5.mlp.fc1.bias": "model-00001-of-00009.safetensors",
+    "blocks.5.mlp.fc2.weight": "model-00001-of-00009.safetensors",
+    "blocks.5.mlp.fc2.bias": "model-00001-of-00009.safetensors",
+    "blocks.6.norm1.weight": "model-00001-of-00009.safetensors",
+    "blocks.6.norm1.bias": "model-00001-of-00009.safetensors",
+    "blocks.6.attn.qkv.weight": "model-00001-of-00009.safetensors",
+    "blocks.6.attn.qkv.bias": "model-00001-of-00009.safetensors",
+    "blocks.6.attn.proj.weight": "model-00001-of-00009.safetensors",
+    "blocks.6.attn.proj.bias": "model-00001-of-00009.safetensors",
+    "blocks.6.norm2.weight": "model-00001-of-00009.safetensors",
+    "blocks.6.norm2.bias": "model-00001-of-00009.safetensors",
+    "blocks.6.mlp.fc1.weight": "model-00001-of-00009.safetensors",
+    "blocks.6.mlp.fc1.bias": "model-00001-of-00009.safetensors",
+    "blocks.6.mlp.fc2.weight": "model-00001-of-00009.safetensors",
+    "blocks.6.mlp.fc2.bias": "model-00001-of-00009.safetensors",
+    "blocks.7.norm1.weight": "model-00001-of-00009.safetensors",
+    "blocks.7.norm1.bias": "model-00001-of-00009.safetensors",
+    "blocks.7.attn.qkv.weight": "model-00001-of-00009.safetensors",
+    "blocks.7.attn.qkv.bias": "model-00001-of-00009.safetensors",
+    "blocks.7.attn.proj.weight": "model-00001-of-00009.safetensors",
+    "blocks.7.attn.proj.bias": "model-00001-of-00009.safetensors",
+    "blocks.7.norm2.weight": "model-00001-of-00009.safetensors",
+    "blocks.7.norm2.bias": "model-00001-of-00009.safetensors",
+    "blocks.7.mlp.fc1.weight": "model-00001-of-00009.safetensors",
+    "blocks.7.mlp.fc1.bias": "model-00001-of-00009.safetensors",
+    "blocks.7.mlp.fc2.weight": "model-00002-of-00009.safetensors",
+    "blocks.7.mlp.fc2.bias": "model-00002-of-00009.safetensors",
+    "blocks.8.norm1.weight": "model-00002-of-00009.safetensors",
+    "blocks.8.norm1.bias": "model-00002-of-00009.safetensors",
+    "blocks.8.attn.qkv.weight": "model-00002-of-00009.safetensors",
+    "blocks.8.attn.qkv.bias": "model-00002-of-00009.safetensors",
+    "blocks.8.attn.proj.weight": "model-00002-of-00009.safetensors",
+    "blocks.8.attn.proj.bias": "model-00002-of-00009.safetensors",
+    "blocks.8.norm2.weight": "model-00002-of-00009.safetensors",
+    "blocks.8.norm2.bias": "model-00002-of-00009.safetensors",
+    "blocks.8.mlp.fc1.weight": "model-00002-of-00009.safetensors",
+    "blocks.8.mlp.fc1.bias": "model-00002-of-00009.safetensors",
+    "blocks.8.mlp.fc2.weight": "model-00002-of-00009.safetensors",
+    "blocks.8.mlp.fc2.bias": "model-00002-of-00009.safetensors",
+    "blocks.9.norm1.weight": "model-00002-of-00009.safetensors",
+    "blocks.9.norm1.bias": "model-00002-of-00009.safetensors",
+    "blocks.9.attn.qkv.weight": "model-00002-of-00009.safetensors",
+    "blocks.9.attn.qkv.bias": "model-00002-of-00009.safetensors",
+    "blocks.9.attn.proj.weight": "model-00002-of-00009.safetensors",
+    "blocks.9.attn.proj.bias": "model-00002-of-00009.safetensors",
+    "blocks.9.norm2.weight": "model-00002-of-00009.safetensors",
+    "blocks.9.norm2.bias": "model-00002-of-00009.safetensors",
+    "blocks.9.mlp.fc1.weight": "model-00002-of-00009.safetensors",
+    "blocks.9.mlp.fc1.bias": "model-00002-of-00009.safetensors",
+    "blocks.9.mlp.fc2.weight": "model-00002-of-00009.safetensors",
+    "blocks.9.mlp.fc2.bias": "model-00002-of-00009.safetensors",
+    "blocks.10.norm1.weight": "model-00002-of-00009.safetensors",
+    "blocks.10.norm1.bias": "model-00002-of-00009.safetensors",
+    "blocks.10.attn.qkv.weight": "model-00002-of-00009.safetensors",
+    "blocks.10.attn.qkv.bias": "model-00002-of-00009.safetensors",
+    "blocks.10.attn.proj.weight": "model-00002-of-00009.safetensors",
+    "blocks.10.attn.proj.bias": "model-00002-of-00009.safetensors",
+    "blocks.10.norm2.weight": "model-00002-of-00009.safetensors",
+    "blocks.10.norm2.bias": "model-00002-of-00009.safetensors",
+    "blocks.10.mlp.fc1.weight": "model-00002-of-00009.safetensors",
+    "blocks.10.mlp.fc1.bias": "model-00002-of-00009.safetensors",
+    "blocks.10.mlp.fc2.weight": "model-00002-of-00009.safetensors",
+    "blocks.10.mlp.fc2.bias": "model-00002-of-00009.safetensors",
+    "blocks.11.norm1.weight": "model-00002-of-00009.safetensors",
+    "blocks.11.norm1.bias": "model-00002-of-00009.safetensors",
+    "blocks.11.attn.qkv.weight": "model-00002-of-00009.safetensors",
+    "blocks.11.attn.qkv.bias": "model-00002-of-00009.safetensors",
+    "blocks.11.attn.proj.weight": "model-00002-of-00009.safetensors",
+    "blocks.11.attn.proj.bias": "model-00002-of-00009.safetensors",
+    "blocks.11.norm2.weight": "model-00002-of-00009.safetensors",
+    "blocks.11.norm2.bias": "model-00002-of-00009.safetensors",
+    "blocks.11.mlp.fc1.weight": "model-00002-of-00009.safetensors",
+    "blocks.11.mlp.fc1.bias": "model-00002-of-00009.safetensors",
+    "blocks.11.mlp.fc2.weight": "model-00002-of-00009.safetensors",
+    "blocks.11.mlp.fc2.bias": "model-00002-of-00009.safetensors",
+    "blocks.12.norm1.weight": "model-00002-of-00009.safetensors",
+    "blocks.12.norm1.bias": "model-00002-of-00009.safetensors",
+    "blocks.12.attn.qkv.weight": "model-00002-of-00009.safetensors",
+    "blocks.12.attn.qkv.bias": "model-00002-of-00009.safetensors",
+    "blocks.12.attn.proj.weight": "model-00002-of-00009.safetensors",
+    "blocks.12.attn.proj.bias": "model-00002-of-00009.safetensors",
+    "blocks.12.norm2.weight": "model-00002-of-00009.safetensors",
+    "blocks.12.norm2.bias": "model-00002-of-00009.safetensors",
+    "blocks.12.mlp.fc1.weight": "model-00002-of-00009.safetensors",
+    "blocks.12.mlp.fc1.bias": "model-00002-of-00009.safetensors",
+    "blocks.12.mlp.fc2.weight": "model-00002-of-00009.safetensors",
+    "blocks.12.mlp.fc2.bias": "model-00002-of-00009.safetensors",
+    "blocks.13.norm1.weight": "model-00002-of-00009.safetensors",
+    "blocks.13.norm1.bias": "model-00002-of-00009.safetensors",
+    "blocks.13.attn.qkv.weight": "model-00002-of-00009.safetensors",
+    "blocks.13.attn.qkv.bias": "model-00002-of-00009.safetensors",
+    "blocks.13.attn.proj.weight": "model-00002-of-00009.safetensors",
+    "blocks.13.attn.proj.bias": "model-00002-of-00009.safetensors",
+    "blocks.13.norm2.weight": "model-00002-of-00009.safetensors",
+    "blocks.13.norm2.bias": "model-00002-of-00009.safetensors",
+    "blocks.13.mlp.fc1.weight": "model-00002-of-00009.safetensors",
+    "blocks.13.mlp.fc1.bias": "model-00002-of-00009.safetensors",
+    "blocks.13.mlp.fc2.weight": "model-00002-of-00009.safetensors",
+    "blocks.13.mlp.fc2.bias": "model-00002-of-00009.safetensors",
+    "blocks.14.norm1.weight": "model-00002-of-00009.safetensors",
+    "blocks.14.norm1.bias": "model-00002-of-00009.safetensors",
+    "blocks.14.attn.qkv.weight": "model-00002-of-00009.safetensors",
+    "blocks.14.attn.qkv.bias": "model-00002-of-00009.safetensors",
+    "blocks.14.attn.proj.weight": "model-00002-of-00009.safetensors",
+    "blocks.14.attn.proj.bias": "model-00002-of-00009.safetensors",
+    "blocks.14.norm2.weight": "model-00002-of-00009.safetensors",
+    "blocks.14.norm2.bias": "model-00002-of-00009.safetensors",
+    "blocks.14.mlp.fc1.weight": "model-00002-of-00009.safetensors",
+    "blocks.14.mlp.fc1.bias": "model-00002-of-00009.safetensors",
+    "blocks.14.mlp.fc2.weight": "model-00002-of-00009.safetensors",
+    "blocks.14.mlp.fc2.bias": "model-00002-of-00009.safetensors",
+    "blocks.15.norm1.weight": "model-00002-of-00009.safetensors",
+    "blocks.15.norm1.bias": "model-00002-of-00009.safetensors",
+    "blocks.15.attn.qkv.weight": "model-00002-of-00009.safetensors",
+    "blocks.15.attn.qkv.bias": "model-00002-of-00009.safetensors",
+    "blocks.15.attn.proj.weight": "model-00002-of-00009.safetensors",
+    "blocks.15.attn.proj.bias": "model-00002-of-00009.safetensors",
+    "blocks.15.norm2.weight": "model-00002-of-00009.safetensors",
+    "blocks.15.norm2.bias": "model-00002-of-00009.safetensors",
+    "blocks.15.mlp.fc1.weight": "model-00003-of-00009.safetensors",
+    "blocks.15.mlp.fc1.bias": "model-00003-of-00009.safetensors",
+    "blocks.15.mlp.fc2.weight": "model-00003-of-00009.safetensors",
+    "blocks.15.mlp.fc2.bias": "model-00003-of-00009.safetensors",
+    "blocks.16.norm1.weight": "model-00003-of-00009.safetensors",
+    "blocks.16.norm1.bias": "model-00003-of-00009.safetensors",
+    "blocks.16.attn.qkv.weight": "model-00003-of-00009.safetensors",
+    "blocks.16.attn.qkv.bias": "model-00003-of-00009.safetensors",
+    "blocks.16.attn.proj.weight": "model-00003-of-00009.safetensors",
+    "blocks.16.attn.proj.bias": "model-00003-of-00009.safetensors",
+    "blocks.16.norm2.weight": "model-00003-of-00009.safetensors",
+    "blocks.16.norm2.bias": "model-00003-of-00009.safetensors",
+    "blocks.16.mlp.fc1.weight": "model-00003-of-00009.safetensors",
+    "blocks.16.mlp.fc1.bias": "model-00003-of-00009.safetensors",
+    "blocks.16.mlp.fc2.weight": "model-00003-of-00009.safetensors",
+    "blocks.16.mlp.fc2.bias": "model-00003-of-00009.safetensors",
+    "blocks.17.norm1.weight": "model-00003-of-00009.safetensors",
+    "blocks.17.norm1.bias": "model-00003-of-00009.safetensors",
+    "blocks.17.attn.qkv.weight": "model-00003-of-00009.safetensors",
+    "blocks.17.attn.qkv.bias": "model-00003-of-00009.safetensors",
+    "blocks.17.attn.proj.weight": "model-00003-of-00009.safetensors",
+    "blocks.17.attn.proj.bias": "model-00003-of-00009.safetensors",
+    "blocks.17.norm2.weight": "model-00003-of-00009.safetensors",
+    "blocks.17.norm2.bias": "model-00003-of-00009.safetensors",
+    "blocks.17.mlp.fc1.weight": "model-00003-of-00009.safetensors",
+    "blocks.17.mlp.fc1.bias": "model-00003-of-00009.safetensors",
+    "blocks.17.mlp.fc2.weight": "model-00003-of-00009.safetensors",
+    "blocks.17.mlp.fc2.bias": "model-00003-of-00009.safetensors",
+    "blocks.18.norm1.weight": "model-00003-of-00009.safetensors",
+    "blocks.18.norm1.bias": "model-00003-of-00009.safetensors",
+    "blocks.18.attn.qkv.weight": "model-00003-of-00009.safetensors",
+    "blocks.18.attn.qkv.bias": "model-00003-of-00009.safetensors",
+    "blocks.18.attn.proj.weight": "model-00003-of-00009.safetensors",
+    "blocks.18.attn.proj.bias": "model-00003-of-00009.safetensors",
+    "blocks.18.norm2.weight": "model-00003-of-00009.safetensors",
+    "blocks.18.norm2.bias": "model-00003-of-00009.safetensors",
+    "blocks.18.mlp.fc1.weight": "model-00003-of-00009.safetensors",
+    "blocks.18.mlp.fc1.bias": "model-00003-of-00009.safetensors",
+    "blocks.18.mlp.fc2.weight": "model-00003-of-00009.safetensors",
+    "blocks.18.mlp.fc2.bias": "model-00003-of-00009.safetensors",
+    "blocks.19.norm1.weight": "model-00003-of-00009.safetensors",
+    "blocks.19.norm1.bias": "model-00003-of-00009.safetensors",
+    "blocks.19.attn.qkv.weight": "model-00003-of-00009.safetensors",
+    "blocks.19.attn.qkv.bias": "model-00003-of-00009.safetensors",
+    "blocks.19.attn.proj.weight": "model-00003-of-00009.safetensors",
+    "blocks.19.attn.proj.bias": "model-00003-of-00009.safetensors",
+    "blocks.19.norm2.weight": "model-00003-of-00009.safetensors",
+    "blocks.19.norm2.bias": "model-00003-of-00009.safetensors",
+    "blocks.19.mlp.fc1.weight": "model-00003-of-00009.safetensors",
+    "blocks.19.mlp.fc1.bias": "model-00003-of-00009.safetensors",
+    "blocks.19.mlp.fc2.weight": "model-00003-of-00009.safetensors",
+    "blocks.19.mlp.fc2.bias": "model-00003-of-00009.safetensors",
+    "blocks.20.norm1.weight": "model-00003-of-00009.safetensors",
+    "blocks.20.norm1.bias": "model-00003-of-00009.safetensors",
+    "blocks.20.attn.qkv.weight": "model-00003-of-00009.safetensors",
+    "blocks.20.attn.qkv.bias": "model-00003-of-00009.safetensors",
+    "blocks.20.attn.proj.weight": "model-00003-of-00009.safetensors",
+    "blocks.20.attn.proj.bias": "model-00003-of-00009.safetensors",
+    "blocks.20.norm2.weight": "model-00003-of-00009.safetensors",
+    "blocks.20.norm2.bias": "model-00003-of-00009.safetensors",
+    "blocks.20.mlp.fc1.weight": "model-00003-of-00009.safetensors",
+    "blocks.20.mlp.fc1.bias": "model-00003-of-00009.safetensors",
+    "blocks.20.mlp.fc2.weight": "model-00003-of-00009.safetensors",
+    "blocks.20.mlp.fc2.bias": "model-00003-of-00009.safetensors",
+    "blocks.21.norm1.weight": "model-00003-of-00009.safetensors",
+    "blocks.21.norm1.bias": "model-00003-of-00009.safetensors",
+    "blocks.21.attn.qkv.weight": "model-00003-of-00009.safetensors",
+    "blocks.21.attn.qkv.bias": "model-00003-of-00009.safetensors",
+    "blocks.21.attn.proj.weight": "model-00003-of-00009.safetensors",
+    "blocks.21.attn.proj.bias": "model-00003-of-00009.safetensors",
+    "blocks.21.norm2.weight": "model-00003-of-00009.safetensors",
+    "blocks.21.norm2.bias": "model-00003-of-00009.safetensors",
+    "blocks.21.mlp.fc1.weight": "model-00003-of-00009.safetensors",
+    "blocks.21.mlp.fc1.bias": "model-00003-of-00009.safetensors",
+    "blocks.21.mlp.fc2.weight": "model-00003-of-00009.safetensors",
+    "blocks.21.mlp.fc2.bias": "model-00003-of-00009.safetensors",
+    "blocks.22.norm1.weight": "model-00003-of-00009.safetensors",
+    "blocks.22.norm1.bias": "model-00003-of-00009.safetensors",
+    "blocks.22.attn.qkv.weight": "model-00003-of-00009.safetensors",
+    "blocks.22.attn.qkv.bias": "model-00003-of-00009.safetensors",
+    "blocks.22.attn.proj.weight": "model-00003-of-00009.safetensors",
+    "blocks.22.attn.proj.bias": "model-00003-of-00009.safetensors",
+    "blocks.22.norm2.weight": "model-00003-of-00009.safetensors",
+    "blocks.22.norm2.bias": "model-00003-of-00009.safetensors",
+    "blocks.22.mlp.fc1.weight": "model-00003-of-00009.safetensors",
+    "blocks.22.mlp.fc1.bias": "model-00003-of-00009.safetensors",
+    "blocks.22.mlp.fc2.weight": "model-00003-of-00009.safetensors",
+    "blocks.22.mlp.fc2.bias": "model-00003-of-00009.safetensors",
+    "blocks.23.norm1.weight": "model-00003-of-00009.safetensors",
+    "blocks.23.norm1.bias": "model-00003-of-00009.safetensors",
+    "blocks.23.attn.qkv.weight": "model-00004-of-00009.safetensors",
+    "blocks.23.attn.qkv.bias": "model-00004-of-00009.safetensors",
+    "blocks.23.attn.proj.weight": "model-00004-of-00009.safetensors",
+    "blocks.23.attn.proj.bias": "model-00004-of-00009.safetensors",
+    "blocks.23.norm2.weight": "model-00004-of-00009.safetensors",
+    "blocks.23.norm2.bias": "model-00004-of-00009.safetensors",
+    "blocks.23.mlp.fc1.weight": "model-00004-of-00009.safetensors",
+    "blocks.23.mlp.fc1.bias": "model-00004-of-00009.safetensors",
+    "blocks.23.mlp.fc2.weight": "model-00004-of-00009.safetensors",
+    "blocks.23.mlp.fc2.bias": "model-00004-of-00009.safetensors",
+    "blocks.24.norm1.weight": "model-00004-of-00009.safetensors",
+    "blocks.24.norm1.bias": "model-00004-of-00009.safetensors",
+    "blocks.24.attn.qkv.weight": "model-00004-of-00009.safetensors",
+    "blocks.24.attn.qkv.bias": "model-00004-of-00009.safetensors",
+    "blocks.24.attn.proj.weight": "model-00004-of-00009.safetensors",
+    "blocks.24.attn.proj.bias": "model-00004-of-00009.safetensors",
+    "blocks.24.norm2.weight": "model-00004-of-00009.safetensors",
+    "blocks.24.norm2.bias": "model-00004-of-00009.safetensors",
+    "blocks.24.mlp.fc1.weight": "model-00004-of-00009.safetensors",
+    "blocks.24.mlp.fc1.bias": "model-00004-of-00009.safetensors",
+    "blocks.24.mlp.fc2.weight": "model-00004-of-00009.safetensors",
+    "blocks.24.mlp.fc2.bias": "model-00004-of-00009.safetensors",
+    "blocks.25.norm1.weight": "model-00004-of-00009.safetensors",
+    "blocks.25.norm1.bias": "model-00004-of-00009.safetensors",
+    "blocks.25.attn.qkv.weight": "model-00004-of-00009.safetensors",
+    "blocks.25.attn.qkv.bias": "model-00004-of-00009.safetensors",
+    "blocks.25.attn.proj.weight": "model-00004-of-00009.safetensors",
+    "blocks.25.attn.proj.bias": "model-00004-of-00009.safetensors",
+    "blocks.25.norm2.weight": "model-00004-of-00009.safetensors",
+    "blocks.25.norm2.bias": "model-00004-of-00009.safetensors",
+    "blocks.25.mlp.fc1.weight": "model-00004-of-00009.safetensors",
+    "blocks.25.mlp.fc1.bias": "model-00004-of-00009.safetensors",
+    "blocks.25.mlp.fc2.weight": "model-00004-of-00009.safetensors",
+    "blocks.25.mlp.fc2.bias": "model-00004-of-00009.safetensors",
+    "blocks.26.norm1.weight": "model-00004-of-00009.safetensors",
+    "blocks.26.norm1.bias": "model-00004-of-00009.safetensors",
+    "blocks.26.attn.qkv.weight": "model-00004-of-00009.safetensors",
+    "blocks.26.attn.qkv.bias": "model-00004-of-00009.safetensors",
+    "blocks.26.attn.proj.weight": "model-00004-of-00009.safetensors",
+    "blocks.26.attn.proj.bias": "model-00004-of-00009.safetensors",
+    "blocks.26.norm2.weight": "model-00004-of-00009.safetensors",
+    "blocks.26.norm2.bias": "model-00004-of-00009.safetensors",
+    "blocks.26.mlp.fc1.weight": "model-00004-of-00009.safetensors",
+    "blocks.26.mlp.fc1.bias": "model-00004-of-00009.safetensors",
+    "blocks.26.mlp.fc2.weight": "model-00004-of-00009.safetensors",
+    "blocks.26.mlp.fc2.bias": "model-00004-of-00009.safetensors",
+    "blocks.27.norm1.weight": "model-00004-of-00009.safetensors",
+    "blocks.27.norm1.bias": "model-00004-of-00009.safetensors",
+    "blocks.27.attn.qkv.weight": "model-00004-of-00009.safetensors",
+    "blocks.27.attn.qkv.bias": "model-00004-of-00009.safetensors",
+    "blocks.27.attn.proj.weight": "model-00004-of-00009.safetensors",
+    "blocks.27.attn.proj.bias": "model-00004-of-00009.safetensors",
+    "blocks.27.norm2.weight": "model-00004-of-00009.safetensors",
+    "blocks.27.norm2.bias": "model-00004-of-00009.safetensors",
+    "blocks.27.mlp.fc1.weight": "model-00004-of-00009.safetensors",
+    "blocks.27.mlp.fc1.bias": "model-00004-of-00009.safetensors",
+    "blocks.27.mlp.fc2.weight": "model-00004-of-00009.safetensors",
+    "blocks.27.mlp.fc2.bias": "model-00004-of-00009.safetensors",
+    "blocks.28.norm1.weight": "model-00004-of-00009.safetensors",
+    "blocks.28.norm1.bias": "model-00004-of-00009.safetensors",
+    "blocks.28.attn.qkv.weight": "model-00004-of-00009.safetensors",
+    "blocks.28.attn.qkv.bias": "model-00004-of-00009.safetensors",
+    "blocks.28.attn.proj.weight": "model-00004-of-00009.safetensors",
+    "blocks.28.attn.proj.bias": "model-00004-of-00009.safetensors",
+    "blocks.28.norm2.weight": "model-00004-of-00009.safetensors",
+    "blocks.28.norm2.bias": "model-00004-of-00009.safetensors",
+    "blocks.28.mlp.fc1.weight": "model-00004-of-00009.safetensors",
+    "blocks.28.mlp.fc1.bias": "model-00004-of-00009.safetensors",
+    "blocks.28.mlp.fc2.weight": "model-00004-of-00009.safetensors",
+    "blocks.28.mlp.fc2.bias": "model-00004-of-00009.safetensors",
+    "blocks.29.norm1.weight": "model-00004-of-00009.safetensors",
+    "blocks.29.norm1.bias": "model-00004-of-00009.safetensors",
+    "blocks.29.attn.qkv.weight": "model-00004-of-00009.safetensors",
+    "blocks.29.attn.qkv.bias": "model-00004-of-00009.safetensors",
+    "blocks.29.attn.proj.weight": "model-00004-of-00009.safetensors",
+    "blocks.29.attn.proj.bias": "model-00004-of-00009.safetensors",
+    "blocks.29.norm2.weight": "model-00004-of-00009.safetensors",
+    "blocks.29.norm2.bias": "model-00004-of-00009.safetensors",
+    "blocks.29.mlp.fc1.weight": "model-00004-of-00009.safetensors",
+    "blocks.29.mlp.fc1.bias": "model-00004-of-00009.safetensors",
+    "blocks.29.mlp.fc2.weight": "model-00004-of-00009.safetensors",
+    "blocks.29.mlp.fc2.bias": "model-00004-of-00009.safetensors",
+    "blocks.30.norm1.weight": "model-00004-of-00009.safetensors",
+    "blocks.30.norm1.bias": "model-00004-of-00009.safetensors",
+    "blocks.30.attn.qkv.weight": "model-00004-of-00009.safetensors",
+    "blocks.30.attn.qkv.bias": "model-00004-of-00009.safetensors",
+    "blocks.30.attn.proj.weight": "model-00004-of-00009.safetensors",
+    "blocks.30.attn.proj.bias": "model-00004-of-00009.safetensors",
+    "blocks.30.norm2.weight": "model-00004-of-00009.safetensors",
+    "blocks.30.norm2.bias": "model-00004-of-00009.safetensors",
+    "blocks.30.mlp.fc1.weight": "model-00004-of-00009.safetensors",
+    "blocks.30.mlp.fc1.bias": "model-00004-of-00009.safetensors",
+    "blocks.30.mlp.fc2.weight": "model-00005-of-00009.safetensors",
+    "blocks.30.mlp.fc2.bias": "model-00005-of-00009.safetensors",
+    "blocks.31.norm1.weight": "model-00005-of-00009.safetensors",
+    "blocks.31.norm1.bias": "model-00005-of-00009.safetensors",
+    "blocks.31.attn.qkv.weight": "model-00005-of-00009.safetensors",
+    "blocks.31.attn.qkv.bias": "model-00005-of-00009.safetensors",
+    "blocks.31.attn.proj.weight": "model-00005-of-00009.safetensors",
+    "blocks.31.attn.proj.bias": "model-00005-of-00009.safetensors",
+    "blocks.31.norm2.weight": "model-00005-of-00009.safetensors",
+    "blocks.31.norm2.bias": "model-00005-of-00009.safetensors",
+    "blocks.31.mlp.fc1.weight": "model-00005-of-00009.safetensors",
+    "blocks.31.mlp.fc1.bias": "model-00005-of-00009.safetensors",
+    "blocks.31.mlp.fc2.weight": "model-00005-of-00009.safetensors",
+    "blocks.31.mlp.fc2.bias": "model-00005-of-00009.safetensors",
+    "blocks.32.norm1.weight": "model-00005-of-00009.safetensors",
+    "blocks.32.norm1.bias": "model-00005-of-00009.safetensors",
+    "blocks.32.attn.qkv.weight": "model-00005-of-00009.safetensors",
+    "blocks.32.attn.qkv.bias": "model-00005-of-00009.safetensors",
+    "blocks.32.attn.proj.weight": "model-00005-of-00009.safetensors",
+    "blocks.32.attn.proj.bias": "model-00005-of-00009.safetensors",
+    "blocks.32.norm2.weight": "model-00005-of-00009.safetensors",
+    "blocks.32.norm2.bias": "model-00005-of-00009.safetensors",
+    "blocks.32.mlp.fc1.weight": "model-00005-of-00009.safetensors",
+    "blocks.32.mlp.fc1.bias": "model-00005-of-00009.safetensors",
+    "blocks.32.mlp.fc2.weight": "model-00005-of-00009.safetensors",
+    "blocks.32.mlp.fc2.bias": "model-00005-of-00009.safetensors",
+    "blocks.33.norm1.weight": "model-00005-of-00009.safetensors",
+    "blocks.33.norm1.bias": "model-00005-of-00009.safetensors",
+    "blocks.33.attn.qkv.weight": "model-00005-of-00009.safetensors",
+    "blocks.33.attn.qkv.bias": "model-00005-of-00009.safetensors",
+    "blocks.33.attn.proj.weight": "model-00005-of-00009.safetensors",
+    "blocks.33.attn.proj.bias": "model-00005-of-00009.safetensors",
+    "blocks.33.norm2.weight": "model-00005-of-00009.safetensors",
+    "blocks.33.norm2.bias": "model-00005-of-00009.safetensors",
+    "blocks.33.mlp.fc1.weight": "model-00005-of-00009.safetensors",
+    "blocks.33.mlp.fc1.bias": "model-00005-of-00009.safetensors",
+    "blocks.33.mlp.fc2.weight": "model-00005-of-00009.safetensors",
+    "blocks.33.mlp.fc2.bias": "model-00005-of-00009.safetensors",
+    "blocks.34.norm1.weight": "model-00005-of-00009.safetensors",
+    "blocks.34.norm1.bias": "model-00005-of-00009.safetensors",
+    "blocks.34.attn.qkv.weight": "model-00005-of-00009.safetensors",
+    "blocks.34.attn.qkv.bias": "model-00005-of-00009.safetensors",
+    "blocks.34.attn.proj.weight": "model-00005-of-00009.safetensors",
+    "blocks.34.attn.proj.bias": "model-00005-of-00009.safetensors",
+    "blocks.34.norm2.weight": "model-00005-of-00009.safetensors",
+    "blocks.34.norm2.bias": "model-00005-of-00009.safetensors",
+    "blocks.34.mlp.fc1.weight": "model-00005-of-00009.safetensors",
+    "blocks.34.mlp.fc1.bias": "model-00005-of-00009.safetensors",
+    "blocks.34.mlp.fc2.weight": "model-00005-of-00009.safetensors",
+    "blocks.34.mlp.fc2.bias": "model-00005-of-00009.safetensors",
+    "blocks.35.norm1.weight": "model-00005-of-00009.safetensors",
+    "blocks.35.norm1.bias": "model-00005-of-00009.safetensors",
+    "blocks.35.attn.qkv.weight": "model-00005-of-00009.safetensors",
+    "blocks.35.attn.qkv.bias": "model-00005-of-00009.safetensors",
+    "blocks.35.attn.proj.weight": "model-00005-of-00009.safetensors",
+    "blocks.35.attn.proj.bias": "model-00005-of-00009.safetensors",
+    "blocks.35.norm2.weight": "model-00005-of-00009.safetensors",
+    "blocks.35.norm2.bias": "model-00005-of-00009.safetensors",
+    "blocks.35.mlp.fc1.weight": "model-00005-of-00009.safetensors",
+    "blocks.35.mlp.fc1.bias": "model-00005-of-00009.safetensors",
+    "blocks.35.mlp.fc2.weight": "model-00005-of-00009.safetensors",
+    "blocks.35.mlp.fc2.bias": "model-00005-of-00009.safetensors",
+    "blocks.36.norm1.weight": "model-00005-of-00009.safetensors",
+    "blocks.36.norm1.bias": "model-00005-of-00009.safetensors",
+    "blocks.36.attn.qkv.weight": "model-00005-of-00009.safetensors",
+    "blocks.36.attn.qkv.bias": "model-00005-of-00009.safetensors",
+    "blocks.36.attn.proj.weight": "model-00005-of-00009.safetensors",
+    "blocks.36.attn.proj.bias": "model-00005-of-00009.safetensors",
+    "blocks.36.norm2.weight": "model-00005-of-00009.safetensors",
+    "blocks.36.norm2.bias": "model-00005-of-00009.safetensors",
+    "blocks.36.mlp.fc1.weight": "model-00005-of-00009.safetensors",
+    "blocks.36.mlp.fc1.bias": "model-00005-of-00009.safetensors",
+    "blocks.36.mlp.fc2.weight": "model-00005-of-00009.safetensors",
+    "blocks.36.mlp.fc2.bias": "model-00005-of-00009.safetensors",
+    "blocks.37.norm1.weight": "model-00005-of-00009.safetensors",
+    "blocks.37.norm1.bias": "model-00005-of-00009.safetensors",
+    "blocks.37.attn.qkv.weight": "model-00005-of-00009.safetensors",
+    "blocks.37.attn.qkv.bias": "model-00005-of-00009.safetensors",
+    "blocks.37.attn.proj.weight": "model-00005-of-00009.safetensors",
+    "blocks.37.attn.proj.bias": "model-00005-of-00009.safetensors",
+    "blocks.37.norm2.weight": "model-00005-of-00009.safetensors",
+    "blocks.37.norm2.bias": "model-00005-of-00009.safetensors",
+    "blocks.37.mlp.fc1.weight": "model-00005-of-00009.safetensors",
+    "blocks.37.mlp.fc1.bias": "model-00005-of-00009.safetensors",
+    "blocks.37.mlp.fc2.weight": "model-00005-of-00009.safetensors",
+    "blocks.37.mlp.fc2.bias": "model-00005-of-00009.safetensors",
+    "blocks.38.norm1.weight": "model-00005-of-00009.safetensors",
+    "blocks.38.norm1.bias": "model-00005-of-00009.safetensors",
+    "blocks.38.attn.qkv.weight": "model-00005-of-00009.safetensors",
+    "blocks.38.attn.qkv.bias": "model-00005-of-00009.safetensors",
+    "blocks.38.attn.proj.weight": "model-00005-of-00009.safetensors",
+    "blocks.38.attn.proj.bias": "model-00005-of-00009.safetensors",
+    "blocks.38.norm2.weight": "model-00005-of-00009.safetensors",
+    "blocks.38.norm2.bias": "model-00005-of-00009.safetensors",
+    "blocks.38.mlp.fc1.weight": "model-00006-of-00009.safetensors",
+    "blocks.38.mlp.fc1.bias": "model-00006-of-00009.safetensors",
+    "blocks.38.mlp.fc2.weight": "model-00006-of-00009.safetensors",
+    "blocks.38.mlp.fc2.bias": "model-00006-of-00009.safetensors",
+    "blocks.39.norm1.weight": "model-00006-of-00009.safetensors",
+    "blocks.39.norm1.bias": "model-00006-of-00009.safetensors",
+    "blocks.39.attn.qkv.weight": "model-00006-of-00009.safetensors",
+    "blocks.39.attn.qkv.bias": "model-00006-of-00009.safetensors",
+    "blocks.39.attn.proj.weight": "model-00006-of-00009.safetensors",
+    "blocks.39.attn.proj.bias": "model-00006-of-00009.safetensors",
+    "blocks.39.norm2.weight": "model-00006-of-00009.safetensors",
+    "blocks.39.norm2.bias": "model-00006-of-00009.safetensors",
+    "blocks.39.mlp.fc1.weight": "model-00006-of-00009.safetensors",
+    "blocks.39.mlp.fc1.bias": "model-00006-of-00009.safetensors",
+    "blocks.39.mlp.fc2.weight": "model-00006-of-00009.safetensors",
+    "blocks.39.mlp.fc2.bias": "model-00006-of-00009.safetensors",
+    "blocks.40.norm1.weight": "model-00006-of-00009.safetensors",
+    "blocks.40.norm1.bias": "model-00006-of-00009.safetensors",
+    "blocks.40.attn.qkv.weight": "model-00006-of-00009.safetensors",
+    "blocks.40.attn.qkv.bias": "model-00006-of-00009.safetensors",
+    "blocks.40.attn.proj.weight": "model-00006-of-00009.safetensors",
+    "blocks.40.attn.proj.bias": "model-00006-of-00009.safetensors",
+    "blocks.40.norm2.weight": "model-00006-of-00009.safetensors",
+    "blocks.40.norm2.bias": "model-00006-of-00009.safetensors",
+    "blocks.40.mlp.fc1.weight": "model-00006-of-00009.safetensors",
+    "blocks.40.mlp.fc1.bias": "model-00006-of-00009.safetensors",
+    "blocks.40.mlp.fc2.weight": "model-00006-of-00009.safetensors",
+    "blocks.40.mlp.fc2.bias": "model-00006-of-00009.safetensors",
+    "blocks.41.norm1.weight": "model-00006-of-00009.safetensors",
+    "blocks.41.norm1.bias": "model-00006-of-00009.safetensors",
+    "blocks.41.attn.qkv.weight": "model-00006-of-00009.safetensors",
+    "blocks.41.attn.qkv.bias": "model-00006-of-00009.safetensors",
+    "blocks.41.attn.proj.weight": "model-00006-of-00009.safetensors",
+    "blocks.41.attn.proj.bias": "model-00006-of-00009.safetensors",
+    "blocks.41.norm2.weight": "model-00006-of-00009.safetensors",
+    "blocks.41.norm2.bias": "model-00006-of-00009.safetensors",
+    "blocks.41.mlp.fc1.weight": "model-00006-of-00009.safetensors",
+    "blocks.41.mlp.fc1.bias": "model-00006-of-00009.safetensors",
+    "blocks.41.mlp.fc2.weight": "model-00006-of-00009.safetensors",
+    "blocks.41.mlp.fc2.bias": "model-00006-of-00009.safetensors",
+    "blocks.42.norm1.weight": "model-00006-of-00009.safetensors",
+    "blocks.42.norm1.bias": "model-00006-of-00009.safetensors",
+    "blocks.42.attn.qkv.weight": "model-00006-of-00009.safetensors",
+    "blocks.42.attn.qkv.bias": "model-00006-of-00009.safetensors",
+    "blocks.42.attn.proj.weight": "model-00006-of-00009.safetensors",
+    "blocks.42.attn.proj.bias": "model-00006-of-00009.safetensors",
+    "blocks.42.norm2.weight": "model-00006-of-00009.safetensors",
+    "blocks.42.norm2.bias": "model-00006-of-00009.safetensors",
+    "blocks.42.mlp.fc1.weight": "model-00006-of-00009.safetensors",
+    "blocks.42.mlp.fc1.bias": "model-00006-of-00009.safetensors",
+    "blocks.42.mlp.fc2.weight": "model-00006-of-00009.safetensors",
+    "blocks.42.mlp.fc2.bias": "model-00006-of-00009.safetensors",
+    "blocks.43.norm1.weight": "model-00006-of-00009.safetensors",
+    "blocks.43.norm1.bias": "model-00006-of-00009.safetensors",
+    "blocks.43.attn.qkv.weight": "model-00006-of-00009.safetensors",
+    "blocks.43.attn.qkv.bias": "model-00006-of-00009.safetensors",
+    "blocks.43.attn.proj.weight": "model-00006-of-00009.safetensors",
+    "blocks.43.attn.proj.bias": "model-00006-of-00009.safetensors",
+    "blocks.43.norm2.weight": "model-00006-of-00009.safetensors",
+    "blocks.43.norm2.bias": "model-00006-of-00009.safetensors",
+    "blocks.43.mlp.fc1.weight": "model-00006-of-00009.safetensors",
+    "blocks.43.mlp.fc1.bias": "model-00006-of-00009.safetensors",
+    "blocks.43.mlp.fc2.weight": "model-00006-of-00009.safetensors",
+    "blocks.43.mlp.fc2.bias": "model-00006-of-00009.safetensors",
+    "blocks.44.norm1.weight": "model-00006-of-00009.safetensors",
+    "blocks.44.norm1.bias": "model-00006-of-00009.safetensors",
+    "blocks.44.attn.qkv.weight": "model-00006-of-00009.safetensors",
+    "blocks.44.attn.qkv.bias": "model-00006-of-00009.safetensors",
+    "blocks.44.attn.proj.weight": "model-00006-of-00009.safetensors",
+    "blocks.44.attn.proj.bias": "model-00006-of-00009.safetensors",
+    "blocks.44.norm2.weight": "model-00006-of-00009.safetensors",
+    "blocks.44.norm2.bias": "model-00006-of-00009.safetensors",
+    "blocks.44.mlp.fc1.weight": "model-00006-of-00009.safetensors",
+    "blocks.44.mlp.fc1.bias": "model-00006-of-00009.safetensors",
+    "blocks.44.mlp.fc2.weight": "model-00006-of-00009.safetensors",
+    "blocks.44.mlp.fc2.bias": "model-00006-of-00009.safetensors",
+    "blocks.45.norm1.weight": "model-00006-of-00009.safetensors",
+    "blocks.45.norm1.bias": "model-00006-of-00009.safetensors",
+    "blocks.45.attn.qkv.weight": "model-00006-of-00009.safetensors",
+    "blocks.45.attn.qkv.bias": "model-00006-of-00009.safetensors",
+    "blocks.45.attn.proj.weight": "model-00006-of-00009.safetensors",
+    "blocks.45.attn.proj.bias": "model-00006-of-00009.safetensors",
+    "blocks.45.norm2.weight": "model-00006-of-00009.safetensors",
+    "blocks.45.norm2.bias": "model-00006-of-00009.safetensors",
+    "blocks.45.mlp.fc1.weight": "model-00006-of-00009.safetensors",
+    "blocks.45.mlp.fc1.bias": "model-00006-of-00009.safetensors",
+    "blocks.45.mlp.fc2.weight": "model-00006-of-00009.safetensors",
+    "blocks.45.mlp.fc2.bias": "model-00006-of-00009.safetensors",
+    "blocks.46.norm1.weight": "model-00006-of-00009.safetensors",
+    "blocks.46.norm1.bias": "model-00006-of-00009.safetensors",
+    "blocks.46.attn.qkv.weight": "model-00007-of-00009.safetensors",
+    "blocks.46.attn.qkv.bias": "model-00007-of-00009.safetensors",
+    "blocks.46.attn.proj.weight": "model-00007-of-00009.safetensors",
+    "blocks.46.attn.proj.bias": "model-00007-of-00009.safetensors",
+    "blocks.46.norm2.weight": "model-00007-of-00009.safetensors",
+    "blocks.46.norm2.bias": "model-00007-of-00009.safetensors",
+    "blocks.46.mlp.fc1.weight": "model-00007-of-00009.safetensors",
+    "blocks.46.mlp.fc1.bias": "model-00007-of-00009.safetensors",
+    "blocks.46.mlp.fc2.weight": "model-00007-of-00009.safetensors",
+    "blocks.46.mlp.fc2.bias": "model-00007-of-00009.safetensors",
+    "blocks.47.norm1.weight": "model-00007-of-00009.safetensors",
+    "blocks.47.norm1.bias": "model-00007-of-00009.safetensors",
+    "blocks.47.attn.qkv.weight": "model-00007-of-00009.safetensors",
+    "blocks.47.attn.qkv.bias": "model-00007-of-00009.safetensors",
+    "blocks.47.attn.proj.weight": "model-00007-of-00009.safetensors",
+    "blocks.47.attn.proj.bias": "model-00007-of-00009.safetensors",
+    "blocks.47.norm2.weight": "model-00007-of-00009.safetensors",
+    "blocks.47.norm2.bias": "model-00007-of-00009.safetensors",
+    "blocks.47.mlp.fc1.weight": "model-00007-of-00009.safetensors",
+    "blocks.47.mlp.fc1.bias": "model-00007-of-00009.safetensors",
+    "blocks.47.mlp.fc2.weight": "model-00007-of-00009.safetensors",
+    "blocks.47.mlp.fc2.bias": "model-00007-of-00009.safetensors",
+    "blocks.48.norm1.weight": "model-00007-of-00009.safetensors",
+    "blocks.48.norm1.bias": "model-00007-of-00009.safetensors",
+    "blocks.48.attn.qkv.weight": "model-00007-of-00009.safetensors",
+    "blocks.48.attn.qkv.bias": "model-00007-of-00009.safetensors",
+    "blocks.48.attn.proj.weight": "model-00007-of-00009.safetensors",
+    "blocks.48.attn.proj.bias": "model-00007-of-00009.safetensors",
+    "blocks.48.norm2.weight": "model-00007-of-00009.safetensors",
+    "blocks.48.norm2.bias": "model-00007-of-00009.safetensors",
+    "blocks.48.mlp.fc1.weight": "model-00007-of-00009.safetensors",
+    "blocks.48.mlp.fc1.bias": "model-00007-of-00009.safetensors",
+    "blocks.48.mlp.fc2.weight": "model-00007-of-00009.safetensors",
+    "blocks.48.mlp.fc2.bias": "model-00007-of-00009.safetensors",
+    "blocks.49.norm1.weight": "model-00007-of-00009.safetensors",
+    "blocks.49.norm1.bias": "model-00007-of-00009.safetensors",
+    "blocks.49.attn.qkv.weight": "model-00007-of-00009.safetensors",
+    "blocks.49.attn.qkv.bias": "model-00007-of-00009.safetensors",
+    "blocks.49.attn.proj.weight": "model-00007-of-00009.safetensors",
+    "blocks.49.attn.proj.bias": "model-00007-of-00009.safetensors",
+    "blocks.49.norm2.weight": "model-00007-of-00009.safetensors",
+    "blocks.49.norm2.bias": "model-00007-of-00009.safetensors",
+    "blocks.49.mlp.fc1.weight": "model-00007-of-00009.safetensors",
+    "blocks.49.mlp.fc1.bias": "model-00007-of-00009.safetensors",
+    "blocks.49.mlp.fc2.weight": "model-00007-of-00009.safetensors",
+    "blocks.49.mlp.fc2.bias": "model-00007-of-00009.safetensors",
+    "blocks.50.norm1.weight": "model-00007-of-00009.safetensors",
+    "blocks.50.norm1.bias": "model-00007-of-00009.safetensors",
+    "blocks.50.attn.qkv.weight": "model-00007-of-00009.safetensors",
+    "blocks.50.attn.qkv.bias": "model-00007-of-00009.safetensors",
+    "blocks.50.attn.proj.weight": "model-00007-of-00009.safetensors",
+    "blocks.50.attn.proj.bias": "model-00007-of-00009.safetensors",
+    "blocks.50.norm2.weight": "model-00007-of-00009.safetensors",
+    "blocks.50.norm2.bias": "model-00007-of-00009.safetensors",
+    "blocks.50.mlp.fc1.weight": "model-00007-of-00009.safetensors",
+    "blocks.50.mlp.fc1.bias": "model-00007-of-00009.safetensors",
+    "blocks.50.mlp.fc2.weight": "model-00007-of-00009.safetensors",
+    "blocks.50.mlp.fc2.bias": "model-00007-of-00009.safetensors",
+    "blocks.51.norm1.weight": "model-00007-of-00009.safetensors",
+    "blocks.51.norm1.bias": "model-00007-of-00009.safetensors",
+    "blocks.51.attn.qkv.weight": "model-00007-of-00009.safetensors",
+    "blocks.51.attn.qkv.bias": "model-00007-of-00009.safetensors",
+    "blocks.51.attn.proj.weight": "model-00007-of-00009.safetensors",
+    "blocks.51.attn.proj.bias": "model-00007-of-00009.safetensors",
+    "blocks.51.norm2.weight": "model-00007-of-00009.safetensors",
+    "blocks.51.norm2.bias": "model-00007-of-00009.safetensors",
+    "blocks.51.mlp.fc1.weight": "model-00007-of-00009.safetensors",
+    "blocks.51.mlp.fc1.bias": "model-00007-of-00009.safetensors",
+    "blocks.51.mlp.fc2.weight": "model-00007-of-00009.safetensors",
+    "blocks.51.mlp.fc2.bias": "model-00007-of-00009.safetensors",
+    "blocks.52.norm1.weight": "model-00007-of-00009.safetensors",
+    "blocks.52.norm1.bias": "model-00007-of-00009.safetensors",
+    "blocks.52.attn.qkv.weight": "model-00007-of-00009.safetensors",
+    "blocks.52.attn.qkv.bias": "model-00007-of-00009.safetensors",
+    "blocks.52.attn.proj.weight": "model-00007-of-00009.safetensors",
+    "blocks.52.attn.proj.bias": "model-00007-of-00009.safetensors",
+    "blocks.52.norm2.weight": "model-00007-of-00009.safetensors",
+    "blocks.52.norm2.bias": "model-00007-of-00009.safetensors",
+    "blocks.52.mlp.fc1.weight": "model-00007-of-00009.safetensors",
+    "blocks.52.mlp.fc1.bias": "model-00007-of-00009.safetensors",
+    "blocks.52.mlp.fc2.weight": "model-00007-of-00009.safetensors",
+    "blocks.52.mlp.fc2.bias": "model-00007-of-00009.safetensors",
+    "blocks.53.norm1.weight": "model-00007-of-00009.safetensors",
+    "blocks.53.norm1.bias": "model-00007-of-00009.safetensors",
+    "blocks.53.attn.qkv.weight": "model-00007-of-00009.safetensors",
+    "blocks.53.attn.qkv.bias": "model-00007-of-00009.safetensors",
+    "blocks.53.attn.proj.weight": "model-00007-of-00009.safetensors",
+    "blocks.53.attn.proj.bias": "model-00007-of-00009.safetensors",
+    "blocks.53.norm2.weight": "model-00007-of-00009.safetensors",
+    "blocks.53.norm2.bias": "model-00007-of-00009.safetensors",
+    "blocks.53.mlp.fc1.weight": "model-00007-of-00009.safetensors",
+    "blocks.53.mlp.fc1.bias": "model-00007-of-00009.safetensors",
+    "blocks.53.mlp.fc2.weight": "model-00008-of-00009.safetensors",
+    "blocks.53.mlp.fc2.bias": "model-00008-of-00009.safetensors",
+    "blocks.54.norm1.weight": "model-00008-of-00009.safetensors",
+    "blocks.54.norm1.bias": "model-00008-of-00009.safetensors",
+    "blocks.54.attn.qkv.weight": "model-00008-of-00009.safetensors",
+    "blocks.54.attn.qkv.bias": "model-00008-of-00009.safetensors",
+    "blocks.54.attn.proj.weight": "model-00008-of-00009.safetensors",
+    "blocks.54.attn.proj.bias": "model-00008-of-00009.safetensors",
+    "blocks.54.norm2.weight": "model-00008-of-00009.safetensors",
+    "blocks.54.norm2.bias": "model-00008-of-00009.safetensors",
+    "blocks.54.mlp.fc1.weight": "model-00008-of-00009.safetensors",
+    "blocks.54.mlp.fc1.bias": "model-00008-of-00009.safetensors",
+    "blocks.54.mlp.fc2.weight": "model-00008-of-00009.safetensors",
+    "blocks.54.mlp.fc2.bias": "model-00008-of-00009.safetensors",
+    "blocks.55.norm1.weight": "model-00008-of-00009.safetensors",
+    "blocks.55.norm1.bias": "model-00008-of-00009.safetensors",
+    "blocks.55.attn.qkv.weight": "model-00008-of-00009.safetensors",
+    "blocks.55.attn.qkv.bias": "model-00008-of-00009.safetensors",
+    "blocks.55.attn.proj.weight": "model-00008-of-00009.safetensors",
+    "blocks.55.attn.proj.bias": "model-00008-of-00009.safetensors",
+    "blocks.55.norm2.weight": "model-00008-of-00009.safetensors",
+    "blocks.55.norm2.bias": "model-00008-of-00009.safetensors",
+    "blocks.55.mlp.fc1.weight": "model-00008-of-00009.safetensors",
+    "blocks.55.mlp.fc1.bias": "model-00008-of-00009.safetensors",
+    "blocks.55.mlp.fc2.weight": "model-00008-of-00009.safetensors",
+    "blocks.55.mlp.fc2.bias": "model-00008-of-00009.safetensors",
+    "blocks.56.norm1.weight": "model-00008-of-00009.safetensors",
+    "blocks.56.norm1.bias": "model-00008-of-00009.safetensors",
+    "blocks.56.attn.qkv.weight": "model-00008-of-00009.safetensors",
+    "blocks.56.attn.qkv.bias": "model-00008-of-00009.safetensors",
+    "blocks.56.attn.proj.weight": "model-00008-of-00009.safetensors",
+    "blocks.56.attn.proj.bias": "model-00008-of-00009.safetensors",
+    "blocks.56.norm2.weight": "model-00008-of-00009.safetensors",
+    "blocks.56.norm2.bias": "model-00008-of-00009.safetensors",
+    "blocks.56.mlp.fc1.weight": "model-00008-of-00009.safetensors",
+    "blocks.56.mlp.fc1.bias": "model-00008-of-00009.safetensors",
+    "blocks.56.mlp.fc2.weight": "model-00008-of-00009.safetensors",
+    "blocks.56.mlp.fc2.bias": "model-00008-of-00009.safetensors",
+    "blocks.57.norm1.weight": "model-00008-of-00009.safetensors",
+    "blocks.57.norm1.bias": "model-00008-of-00009.safetensors",
+    "blocks.57.attn.qkv.weight": "model-00008-of-00009.safetensors",
+    "blocks.57.attn.qkv.bias": "model-00008-of-00009.safetensors",
+    "blocks.57.attn.proj.weight": "model-00008-of-00009.safetensors",
+    "blocks.57.attn.proj.bias": "model-00008-of-00009.safetensors",
+    "blocks.57.norm2.weight": "model-00008-of-00009.safetensors",
+    "blocks.57.norm2.bias": "model-00008-of-00009.safetensors",
+    "blocks.57.mlp.fc1.weight": "model-00008-of-00009.safetensors",
+    "blocks.57.mlp.fc1.bias": "model-00008-of-00009.safetensors",
+    "blocks.57.mlp.fc2.weight": "model-00008-of-00009.safetensors",
+    "blocks.57.mlp.fc2.bias": "model-00008-of-00009.safetensors",
+    "blocks.58.norm1.weight": "model-00008-of-00009.safetensors",
+    "blocks.58.norm1.bias": "model-00008-of-00009.safetensors",
+    "blocks.58.attn.qkv.weight": "model-00008-of-00009.safetensors",
+    "blocks.58.attn.qkv.bias": "model-00008-of-00009.safetensors",
+    "blocks.58.attn.proj.weight": "model-00008-of-00009.safetensors",
+    "blocks.58.attn.proj.bias": "model-00008-of-00009.safetensors",
+    "blocks.58.norm2.weight": "model-00008-of-00009.safetensors",
+    "blocks.58.norm2.bias": "model-00008-of-00009.safetensors",
+    "blocks.58.mlp.fc1.weight": "model-00008-of-00009.safetensors",
+    "blocks.58.mlp.fc1.bias": "model-00008-of-00009.safetensors",
+    "blocks.58.mlp.fc2.weight": "model-00008-of-00009.safetensors",
+    "blocks.58.mlp.fc2.bias": "model-00008-of-00009.safetensors",
+    "blocks.59.norm1.weight": "model-00008-of-00009.safetensors",
+    "blocks.59.norm1.bias": "model-00008-of-00009.safetensors",
+    "blocks.59.attn.qkv.weight": "model-00008-of-00009.safetensors",
+    "blocks.59.attn.qkv.bias": "model-00008-of-00009.safetensors",
+    "blocks.59.attn.proj.weight": "model-00008-of-00009.safetensors",
+    "blocks.59.attn.proj.bias": "model-00008-of-00009.safetensors",
+    "blocks.59.norm2.weight": "model-00008-of-00009.safetensors",
+    "blocks.59.norm2.bias": "model-00008-of-00009.safetensors",
+    "blocks.59.mlp.fc1.weight": "model-00008-of-00009.safetensors",
+    "blocks.59.mlp.fc1.bias": "model-00008-of-00009.safetensors",
+    "blocks.59.mlp.fc2.weight": "model-00008-of-00009.safetensors",
+    "blocks.59.mlp.fc2.bias": "model-00008-of-00009.safetensors",
+    "blocks.60.norm1.weight": "model-00008-of-00009.safetensors",
+    "blocks.60.norm1.bias": "model-00008-of-00009.safetensors",
+    "blocks.60.attn.qkv.weight": "model-00008-of-00009.safetensors",
+    "blocks.60.attn.qkv.bias": "model-00008-of-00009.safetensors",
+    "blocks.60.attn.proj.weight": "model-00008-of-00009.safetensors",
+    "blocks.60.attn.proj.bias": "model-00008-of-00009.safetensors",
+    "blocks.60.norm2.weight": "model-00008-of-00009.safetensors",
+    "blocks.60.norm2.bias": "model-00008-of-00009.safetensors",
+    "blocks.60.mlp.fc1.weight": "model-00008-of-00009.safetensors",
+    "blocks.60.mlp.fc1.bias": "model-00008-of-00009.safetensors",
+    "blocks.60.mlp.fc2.weight": "model-00008-of-00009.safetensors",
+    "blocks.60.mlp.fc2.bias": "model-00008-of-00009.safetensors",
+    "blocks.61.norm1.weight": "model-00008-of-00009.safetensors",
+    "blocks.61.norm1.bias": "model-00008-of-00009.safetensors",
+    "blocks.61.attn.qkv.weight": "model-00008-of-00009.safetensors",
+    "blocks.61.attn.qkv.bias": "model-00008-of-00009.safetensors",
+    "blocks.61.attn.proj.weight": "model-00008-of-00009.safetensors",
+    "blocks.61.attn.proj.bias": "model-00008-of-00009.safetensors",
+    "blocks.61.norm2.weight": "model-00008-of-00009.safetensors",
+    "blocks.61.norm2.bias": "model-00008-of-00009.safetensors",
+    "blocks.61.mlp.fc1.weight": "model-00009-of-00009.safetensors",
+    "blocks.61.mlp.fc1.bias": "model-00009-of-00009.safetensors",
+    "blocks.61.mlp.fc2.weight": "model-00009-of-00009.safetensors",
+    "blocks.61.mlp.fc2.bias": "model-00009-of-00009.safetensors",
+    "blocks.62.norm1.weight": "model-00009-of-00009.safetensors",
+    "blocks.62.norm1.bias": "model-00009-of-00009.safetensors",
+    "blocks.62.attn.qkv.weight": "model-00009-of-00009.safetensors",
+    "blocks.62.attn.qkv.bias": "model-00009-of-00009.safetensors",
+    "blocks.62.attn.proj.weight": "model-00009-of-00009.safetensors",
+    "blocks.62.attn.proj.bias": "model-00009-of-00009.safetensors",
+    "blocks.62.norm2.weight": "model-00009-of-00009.safetensors",
+    "blocks.62.norm2.bias": "model-00009-of-00009.safetensors",
+    "blocks.62.mlp.fc1.weight": "model-00009-of-00009.safetensors",
+    "blocks.62.mlp.fc1.bias": "model-00009-of-00009.safetensors",
+    "blocks.62.mlp.fc2.weight": "model-00009-of-00009.safetensors",
+    "blocks.62.mlp.fc2.bias": "model-00009-of-00009.safetensors",
+    "blocks.63.norm1.weight": "model-00009-of-00009.safetensors",
+    "blocks.63.norm1.bias": "model-00009-of-00009.safetensors",
+    "blocks.63.attn.qkv.weight": "model-00009-of-00009.safetensors",
+    "blocks.63.attn.qkv.bias": "model-00009-of-00009.safetensors",
+    "blocks.63.attn.proj.weight": "model-00009-of-00009.safetensors",
+    "blocks.63.attn.proj.bias": "model-00009-of-00009.safetensors",
+    "blocks.63.norm2.weight": "model-00009-of-00009.safetensors",
+    "blocks.63.norm2.bias": "model-00009-of-00009.safetensors",
+    "blocks.63.mlp.fc1.weight": "model-00009-of-00009.safetensors",
+    "blocks.63.mlp.fc1.bias": "model-00009-of-00009.safetensors",
+    "blocks.63.mlp.fc2.weight": "model-00009-of-00009.safetensors",
+    "blocks.63.mlp.fc2.bias": "model-00009-of-00009.safetensors"
+  }
+}

modeling_eva_clip.py

@@ -0,0 +1,119 @@
+from __future__ import annotations
+
+from collections import OrderedDict
+from copy import deepcopy
+from typing import Optional, Tuple
+
+import torch
+from torch import nn
+from transformers.modeling_utils import PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+from transformers.utils import logging
+
+from .configuration_eva_clip import DEFAULT_VISION_CFG, EvaClipVisionConfig
+from .legacy_eva_vit import (
+        CLIPVisionCfg,
+        EVAVisionTransformer,
+        create_norm_layer_factory,
+    )
+
+logger = logging.get_logger(__name__)
+
+
+if not hasattr(nn.Module, "get_parameter_or_buffer"):  # pragma: no cover - compatibility shim
+    def _get_parameter_or_buffer(self, target: str):
+        try:
+            return self.get_parameter(target)
+        except AttributeError:
+            pass
+        try:
+            return self.get_buffer(target)
+        except AttributeError:
+            pass
+        raise AttributeError(f"`{target}` is neither a parameter nor a buffer.")
+
+    nn.Module.get_parameter_or_buffer = _get_parameter_or_buffer  # type: ignore[attr-defined]
+
+
+class EvaClipVisionModelOutput(ModelOutput):
+    last_hidden_state: torch.FloatTensor
+
+
+class EvaClipVisionModel(PreTrainedModel):
+    config_class = EvaClipVisionConfig
+    base_model_prefix = "vision_model"
+    main_input_name = "pixel_values"
+    supports_gradient_checkpointing = True
+
+    def __init__(self, config: EvaClipVisionConfig) -> None:
+        super().__init__(config)
+        vision_cfg = self._build_vision_cfg(config)
+        self._vision_cfg_dict = vision_cfg
+        norm_layer_factory = create_norm_layer_factory(vision_cfg.fusedLN, eps=1e-6)
+
+        vision_heads = vision_cfg.width // vision_cfg.head_width
+        self.vision_model = EVAVisionTransformer(
+            img_size=vision_cfg.image_size,
+            patch_size=vision_cfg.patch_size,
+            in_chans=3,
+            num_classes=config.embed_dim,
+            embed_dim=vision_cfg.width,
+            depth=vision_cfg.layers,
+            num_heads=vision_heads,
+            mlp_ratio=vision_cfg.mlp_ratio,
+            qkv_bias=vision_cfg.qkv_bias,
+            drop_path_rate=vision_cfg.drop_path_rate or 0.0,
+            norm_layer=norm_layer_factory,
+            xattn=vision_cfg.xattn,
+            rope=vision_cfg.rope,
+            postnorm=vision_cfg.postnorm,
+            pt_hw_seq_len=vision_cfg.pt_hw_seq_len,
+            intp_freq=vision_cfg.intp_freq,
+            naiveswiglu=vision_cfg.naiveswiglu,
+            subln=vision_cfg.subln,
+            patch_dropout=vision_cfg.patch_dropout,
+            use_mean_pooling=vision_cfg.global_average_pool,
+        )
+
+        self.image_size = vision_cfg.image_size
+        self.patch_size = vision_cfg.patch_size
+
+    def _build_vision_cfg(self, config: EvaClipVisionConfig) -> CLIPVisionCfg:
+        merged = deepcopy(DEFAULT_VISION_CFG)
+        merged.update(config.vision_cfg or {})
+        if "pt_hw_seq_len" not in merged and merged.get("image_size") and merged.get("patch_size"):
+            merged["pt_hw_seq_len"] = merged["image_size"] // merged["patch_size"]
+        return CLIPVisionCfg(**merged)
+
+    def forward(
+        self,
+        pixel_values: torch.FloatTensor,
+        return_dict: bool = True,
+        output_hidden_states: bool = False,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.FloatTensor, ...] | EvaClipVisionModelOutput:
+        del output_hidden_states, output_attentions  # not supported
+        hidden_states = self.vision_model(pixel_values, return_all_features=True)
+        if hidden_states.shape[1] > 0:
+            hidden_states = hidden_states[:, 1:, :]
+
+        if not return_dict:
+            return (hidden_states,)
+        return EvaClipVisionModelOutput(last_hidden_state=hidden_states)
+
+    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs: Optional[dict] = None) -> None:
+        del gradient_checkpointing_kwargs
+        if hasattr(self.vision_model, "set_grad_checkpointing"):
+            self.vision_model.set_grad_checkpointing(True)
+
+    # The EVA checkpoints omit the top-level module prefix. Mirror that convention here.
+    def load_state_dict(self, state_dict, strict: bool = True):
+        if any(k.startswith("vision_model.") for k in state_dict):
+            processed = state_dict
+        else:
+            processed = OrderedDict((f"vision_model.{k}", v) for k, v in state_dict.items())
+        return super().load_state_dict(processed, strict=strict)
+
+
+__all__ = ["EvaClipVisionModel", "EvaClipVisionModelOutput"]
+