Add model

.gitattributes

@@ -17,10 +17,6 @@
 *.ot filter=lfs diff=lfs merge=lfs -text
 *.parquet filter=lfs diff=lfs merge=lfs -text
 *.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
 saved_model/**/* filter=lfs diff=lfs merge=lfs -text
@@ -33,3 +29,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,12 +1,11 @@
 ---
-title: Ovsam
+title: Open-Vocabulary SAM
 emoji: 📚
 colorFrom: green
 colorTo: red
 sdk: gradio
 sdk_version: 4.13.0
-app_file: app.py
+app_file: main.py
 pinned: false
+python_version: 3.10
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app/configs/sam_r50x16_fpn.py

@@ -0,0 +1,81 @@
+from mmcv.ops import RoIAlign
+from mmdet.models import FPN, SingleRoIExtractor
+
+from app.models.model import SAMSegmentor
+from app.models.openclip_backbone import OpenCLIPBackbone
+from app.models.ovsam_head import OVSAMHead
+from app.models.sam_pe import SAMPromptEncoder
+from app.models.transformer_neck import MultiLayerTransformerNeck
+
+model = dict(
+    type=SAMSegmentor,
+    data_preprocessor=None,
+    enable_backbone=True,
+    backbone=dict(
+        type=OpenCLIPBackbone,
+        model_name='RN50x16',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='openai'
+        )
+    ),
+    neck=dict(
+        type=MultiLayerTransformerNeck,
+        input_size=(1024, 1024),
+        in_channels=[384, 768, 1536, 3072],
+        strides=[4, 8, 16, 32],
+        layer_ids=(0, 1, 2, 3),
+        embed_channels=1280,
+        out_channels=256,
+        fix=True,
+        init_cfg=dict(
+            type='Pretrained',
+            checkpoint='./models/sam2clip_vith_rn50.pth',
+            prefix='neck_student',
+        )
+    ),
+    prompt_encoder=dict(
+        type=SAMPromptEncoder,
+        model_name='vit_h',
+        fix=True,
+        init_cfg=dict(
+            type='sam_pretrain',
+            checkpoint='vit_h'
+        )
+    ),
+    fpn_neck=dict(
+        type=FPN,
+        in_channels=[384, 768, 1536, 3072],
+        out_channels=256,
+        num_outs=4,
+        init_cfg=dict(
+            type='Pretrained',
+            checkpoint='./models/R50x16_fpn_lvis_norare_v3det.pth',
+            prefix='fpn_neck',
+        ),
+    ),
+    mask_decoder=dict(
+        type=OVSAMHead,
+        model_name='vit_h',
+        with_label_token=True,
+        gen_box=True,
+        ov_classifier_name='RN50x16_LVISV1Dataset',
+        roi_extractor=dict(
+            type=SingleRoIExtractor,
+            roi_layer=dict(type=RoIAlign, output_size=7, sampling_ratio=0),
+            out_channels=256,
+            featmap_strides=[4, 8, 16, 32]
+        ),
+        fix=False,
+        init_cfg=dict(
+            type='Pretrained',
+            checkpoint='./models/ovsam_R50x16_lvisnorare.pth',
+            prefix='mask_decoder',
+        ),
+        load_roi_conv=dict(
+            checkpoint='./models/R50x16_fpn_lvis_norare_v3det.pth',
+            prefix='roi_conv',
+        )
+    )
+)

app/configs/sam_vith.py

@@ -0,0 +1,38 @@
+from app.models.last_layer import LastLayerNeck
+from app.models.model import SAMSegmentor
+from app.models.sam_backbone import SAMBackbone
+from app.models.sam_mask_decoder import SAMMaskDecoder
+from app.models.sam_pe import SAMPromptEncoder
+
+model = dict(
+    type=SAMSegmentor,
+    data_preprocessor=None,
+    backbone=dict(
+        type=SAMBackbone,
+        model_name='vit_h',
+        fix=True,
+        init_cfg=dict(
+            type='sam_pretrain',
+            checkpoint='vit_h'
+        )
+    ),
+    neck=dict(type=LastLayerNeck),
+    prompt_encoder=dict(
+        type=SAMPromptEncoder,
+        model_name='vit_h',
+        fix=True,
+        init_cfg=dict(
+            type='sam_pretrain',
+            checkpoint='vit_h'
+        )
+    ),
+    mask_decoder=dict(
+        type=SAMMaskDecoder,
+        model_name='vit_h',
+        fix=True,
+        init_cfg=dict(
+            type='sam_pretrain',
+            checkpoint='vit_h'
+        )
+    )
+)

app/models/last_layer.py

@@ -0,0 +1,20 @@
+from typing import Tuple
+
+from mmengine.model import BaseModule
+from torch import Tensor
+
+from mmdet.registry import MODELS
+
+
+@MODELS.register_module()
+class LastLayerNeck(BaseModule):
+    r"""Last Layer Neck
+
+    Return the last layer feature of the backbone.
+    """
+
+    def __init__(self) -> None:
+        super().__init__(init_cfg=None)
+
+    def forward(self, inputs: Tuple[Tensor]) -> Tensor:
+        return inputs[-1]

app/models/model.py

@@ -0,0 +1,92 @@
+import torch.nn.functional as F
+from mmengine.model import BaseModel
+
+from mmdet.registry import MODELS
+from mmdet.utils import ConfigType, OptConfigType, OptMultiConfig
+
+
+@MODELS.register_module()
+class SAMSegmentor(BaseModel):
+    MASK_THRESHOLD = 0.5
+
+    def __init__(
+            self,
+            backbone: ConfigType,
+            neck: ConfigType,
+            prompt_encoder: ConfigType,
+            mask_decoder: ConfigType,
+            data_preprocessor: OptConfigType = None,
+            fpn_neck: OptConfigType = None,
+            init_cfg: OptMultiConfig = None,
+            use_clip_feat: bool = False,
+            use_head_feat: bool = False,
+            use_gt_prompt: bool = False,
+            use_point: bool = False,
+            enable_backbone: bool = False,
+    ) -> None:
+        super().__init__(data_preprocessor=data_preprocessor, init_cfg=init_cfg)
+
+        self.backbone = MODELS.build(backbone)
+        self.neck = MODELS.build(neck)
+        self.pe = MODELS.build(prompt_encoder)
+        self.mask_decoder = MODELS.build(mask_decoder)
+        if fpn_neck is not None:
+            self.fpn_neck = MODELS.build(fpn_neck)
+        else:
+            self.fpn_neck = None
+
+        self.use_clip_feat = use_clip_feat
+        self.use_head_feat = use_head_feat
+        self.use_gt_prompt = use_gt_prompt
+        self.use_point = use_point
+
+        self.enable_backbone = enable_backbone
+
+    def extract_feat(self, inputs):
+        backbone_feat = self.backbone(inputs)
+        neck_feat = self.neck(backbone_feat)
+        if self.fpn_neck is not None:
+            fpn_feat = self.fpn_neck(backbone_feat)
+        else:
+            fpn_feat = None
+
+        return dict(
+            backbone_feat=backbone_feat,
+            neck_feat=neck_feat,
+            fpn_feat=fpn_feat
+        )
+
+    def extract_masks(self, feat_cache, prompts):
+        sparse_embed, dense_embed = self.pe(
+            prompts,
+            image_size=(1024, 1024),
+            with_points='point_coords' in prompts,
+            with_bboxes='bboxes' in prompts,
+        )
+
+        kwargs = dict()
+        if self.enable_backbone:
+            kwargs['backbone_feats'] = feat_cache['backbone_feat']
+            kwargs['backbone'] = self.backbone
+            kwargs['fpn_feats'] = feat_cache['fpn_feat']
+        low_res_masks, iou_predictions, cls_pred = self.mask_decoder(
+            image_embeddings=feat_cache['neck_feat'],
+            image_pe=self.pe.get_dense_pe(),
+            sparse_prompt_embeddings=sparse_embed,
+            dense_prompt_embeddings=dense_embed,
+            multi_mask_output=False,
+            **kwargs
+        )
+        masks = F.interpolate(
+            low_res_masks,
+            scale_factor=4.,
+            mode='bilinear',
+            align_corners=False,
+        )
+
+        masks = masks.sigmoid()
+        cls_pred = cls_pred.softmax(-1)[..., :-1]
+        return masks.detach().cpu().numpy(), cls_pred.detach().cpu()
+
+    def forward(self, inputs):
+        return inputs

app/models/openclip_backbone.py

@@ -0,0 +1,292 @@
+from typing import Optional, List
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+from mmdet.registry import MODELS
+
+from mmengine.model import BaseModule
+from mmengine.dist import get_dist_info
+from mmengine.logging import MMLogger
+
+import ext.open_clip as open_clip
+from utils.load_checkpoint import load_checkpoint_with_prefix
+
+
+@MODELS.register_module()
+class OpenCLIPBackbone(BaseModule):
+    """OpenCLIPBackbone,
+    Please refer to:
+    https://github.com/mlfoundations/open_clip/tree/5f7892b672b21e6853d0f6c11b18dda9bcf36c8d#pretrained-model-interface
+    for the supported models and checkpoints.
+    """
+    STAGES = 4
+
+    def __init__(
+            self,
+            img_size: int = 1024,
+            model_name: str = '',
+            fix: bool = True,
+            fix_layers: Optional[List] = None,
+            init_cfg=None,
+    ):
+        assert init_cfg is not None and init_cfg['type'] in ['clip_pretrain', 'image_pretrain', 'Pretrained'], \
+            f"{init_cfg['type']} is not supported."
+        pretrained = init_cfg['checkpoint']
+        super().__init__(init_cfg=None)
+        self.init_cfg = init_cfg
+        self.logger = MMLogger.get_current_instance()
+        rank, world_size = get_dist_info()
+
+        if world_size > 1:
+            if rank == 0:
+                if init_cfg['type'] == 'clip_pretrain':
+                    _ = open_clip.create_model_from_pretrained(model_name, pretrained=pretrained,
+                                                               return_transform=False, logger=self.logger)
+                elif init_cfg['type'] == 'image_pretrain':
+                    _ = open_clip.create_model(model_name, pretrained_image=True, logger=self.logger)
+
+            else:
+                pass
+            dist.barrier()
+
+        # Get the clip model
+        if init_cfg['type'] == 'clip_pretrain':
+            clip_model = open_clip.create_model_from_pretrained(model_name, pretrained=pretrained,
+                                                                return_transform=False, logger=self.logger)
+        elif init_cfg['type'] == 'image_pretrain':
+            clip_model = open_clip.create_model(model_name, pretrained_image=True, logger=self.logger)
+        elif init_cfg['type'] == 'Pretrained':
+            clip_model = open_clip.create_model(model_name, pretrained_image=False, logger=self.logger)
+        else:
+            raise NotImplementedError
+
+        self.out_indices = (0, 1, 2, 3)
+        model_name_lower = model_name.lower()
+        if 'convnext_' in model_name_lower:
+            model_type = 'convnext'
+            if '_base' in model_name_lower:
+                output_channels = [128, 256, 512, 1024]
+                feat_size = 0
+            elif '_large' in model_name_lower:
+                output_channels = [192, 384, 768, 1536]
+                feat_size = 0
+            elif '_xxlarge' in model_name_lower:
+                output_channels = [384, 768, 1536, 3072]
+                feat_size = 0
+            else:
+                raise NotImplementedError(f"{model_name} not supported yet.")
+        elif 'rn' in model_name_lower:
+            model_type = 'resnet'
+            if model_name_lower.replace('-quickgelu', '') in ['rn50', 'rn101']:
+                output_channels = [256, 512, 1024, 2048]
+                feat_size = 7
+            elif model_name_lower == 'rn50x4':
+                output_channels = [320, 640, 1280, 2560]
+                feat_size = 9
+            elif model_name_lower == 'rn50x16':
+                output_channels = [384, 768, 1536, 3072]
+                feat_size = 12
+            elif model_name_lower == 'rn50x64':
+                output_channels = [512, 1024, 2048, 4096]
+                feat_size = 14
+            else:
+                raise NotImplementedError(f"{model_name} not supported yet.")
+        else:
+            raise NotImplementedError(f"{model_name} not supported yet.")
+
+        self.model_name = model_name
+        self.fix = fix
+        self.model_type = model_type
+        self.output_channels = output_channels
+        self.feat_size = feat_size
+
+        # Get the visual model
+        if self.model_type == 'resnet':
+            self.stem = nn.Sequential(*[
+                clip_model.visual.conv1, clip_model.visual.bn1, clip_model.visual.act1,
+                clip_model.visual.conv2, clip_model.visual.bn2, clip_model.visual.act2,
+                clip_model.visual.conv3, clip_model.visual.bn3, clip_model.visual.act3,
+            ])
+        elif self.model_type == 'convnext':
+            self.stem = clip_model.visual.trunk.stem
+        else:
+            raise ValueError
+
+        if self.model_type == 'resnet':
+            self.avgpool = clip_model.visual.avgpool
+        elif self.model_type == 'convnext':
+            self.avgpool = nn.Identity()
+        else:
+            raise ValueError
+
+        self.res_layers = []
+        for i in range(self.STAGES):
+            if self.model_type == 'resnet':
+                layer_name = f'layer{i + 1}'
+                layer = getattr(clip_model.visual, layer_name)
+            elif self.model_type == 'convnext':
+                layer_name = f'layer{i + 1}'
+                layer = clip_model.visual.trunk.stages[i]
+            else:
+                raise ValueError
+            self.add_module(layer_name, layer)
+            self.res_layers.append(layer_name)
+
+        if self.model_type == 'resnet':
+            self.norm_pre = nn.Identity()
+        elif self.model_type == 'convnext':
+            self.norm_pre = clip_model.visual.trunk.norm_pre
+
+        if self.model_type == 'resnet':
+            self.head = clip_model.visual.attnpool
+        elif self.model_type == 'convnext':
+            self.head = nn.Sequential(*[
+                clip_model.visual.trunk.head,
+                clip_model.visual.head,
+            ])
+
+        if self.init_cfg['type'] == 'Pretrained':
+            checkpoint_path = pretrained
+            state_dict = load_checkpoint_with_prefix(checkpoint_path, prefix=self.init_cfg['prefix'])
+            self.load_state_dict(state_dict, strict=True)
+
+        self.fix_layers = fix_layers
+
+        if not self.fix:
+            self.train()
+            for name, param in self.norm_pre.named_parameters():
+                param.requires_grad = False
+            for name, param in self.head.named_parameters():
+                param.requires_grad = False
+            if self.fix_layers is not None:
+                for i, layer_name in enumerate(self.res_layers):
+                    if i in self.fix_layers:
+                        res_layer = getattr(self, layer_name)
+                        for name, param in res_layer.named_parameters():
+                            param.requires_grad = False
+
+        if self.fix:
+            self.train(mode=False)
+            for name, param in self.named_parameters():
+                param.requires_grad = False
+
+    def init_weights(self):
+        self.logger.info(f"Init Config for {self.model_name}")
+        self.logger.info(self.init_cfg)
+
+    def train(self: torch.nn.Module, mode: bool = True) -> torch.nn.Module:
+        if not isinstance(mode, bool):
+            raise ValueError("training mode is expected to be boolean")
+        if self.fix:
+            super().train(mode=False)
+        else:
+            super().train(mode=mode)
+            if self.fix_layers is not None:
+                for i, layer_name in enumerate(self.res_layers):
+                    if i in self.fix_layers:
+                        res_layer = getattr(self, layer_name)
+                        res_layer.train(mode=False)
+        return self
+
+    def forward_func(self, x):
+        x = self.stem(x)
+        x = self.avgpool(x)
+        outs = []
+        for i, layer_name in enumerate(self.res_layers):
+            res_layer = getattr(self, layer_name)
+            x = res_layer(x).contiguous()
+            if i in self.out_indices:
+                outs.append(x)
+        return tuple(outs)
+
+    def get_clip_feature(self, backbone_feat):
+        if self.model_type == 'resnet':
+            return backbone_feat
+        elif self.model_type == 'convnext':
+            return self.norm_pre(backbone_feat)
+        raise NotImplementedError
+
+    def forward_feat(self, features):
+        if self.model_type == 'convnext':
+            batch, num_query, channel = features.shape
+            features = features.reshape(batch * num_query, channel, 1, 1)
+            features = self.head(features)
+            return features.view(batch, num_query, features.shape[-1])
+        elif self.model_type == 'resnet':
+            num_query, channel, seven, seven = features.shape
+            features = self.head(features)
+            return features
+
+    def forward(self, x):
+        if self.fix:
+            with torch.no_grad():
+                outs = self.forward_func(x)
+        else:
+            outs = self.forward_func(x)
+        return outs
+
+    def get_text_model(self):
+        return OpenCLIPBackboneText(
+            self.model_name,
+            init_cfg=self.init_cfg
+        )
+
+
+@MODELS.register_module()
+class OpenCLIPBackboneText(BaseModule):
+    def __init__(
+            self,
+            model_name: str = '',
+            init_cfg=None,
+    ):
+        assert init_cfg is not None and init_cfg['type'] == 'clip_pretrain', f"{init_cfg['type']} is not supported."
+        pretrained = init_cfg['checkpoint']
+        super().__init__(init_cfg=None)
+        self.init_cfg = init_cfg
+        self.logger = MMLogger.get_current_instance()
+        rank, world_size = get_dist_info()
+
+        if world_size > 1:
+            if rank == 0:
+                _ = open_clip.create_model_from_pretrained(model_name, pretrained=pretrained, return_transform=False,
+                                                           logger=self.logger)
+            else:
+                pass
+            dist.barrier()
+
+        # Get the clip model
+        clip_model = open_clip.create_model_from_pretrained(model_name, pretrained=pretrained, return_transform=False,
+                                                            logger=self.logger)
+
+        # Get the textual model
+        self.text_tokenizer = open_clip.get_tokenizer(model_name)
+        self.text_transformer = clip_model.transformer
+        self.text_token_embedding = clip_model.token_embedding
+        self.text_pe = clip_model.positional_embedding
+        self.text_ln_final = clip_model.ln_final
+        self.text_proj = clip_model.text_projection
+
+        self.register_buffer('text_attn_mask', clip_model.attn_mask)
+
+        self.param_dtype = torch.float32
+        self.model_name = model_name
+
+    def init_weights(self):
+        self.logger.info(f"Init Config for {self.model_name}")
+        self.logger.info(self.init_cfg)
+
+    # Copied from
+    # https://github.com/openai/CLIP/blob/d50d76daa670286dd6cacf3bcd80b5e4823fc8e1/clip/model.py#L343
+    @torch.no_grad()
+    def forward(self, text):
+        text_tokens = self.text_tokenizer(text).to(device=self.text_proj.device)
+        x = self.text_token_embedding(text_tokens).to(self.param_dtype)
+        x = x + self.text_pe.to(self.param_dtype)
+        x = x.permute(1, 0, 2)
+        x = self.text_transformer(x, attn_mask=self.text_attn_mask)
+        x = x.permute(1, 0, 2)
+        x = self.text_ln_final(x)  # [batch_size, n_ctx, transformer.width]
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        x = x[torch.arange(x.shape[0]), text_tokens.argmax(dim=-1)] @ self.text_proj
+        return x

app/models/ovsam_head.py

@@ -0,0 +1,226 @@
+import copy
+import os
+from typing import Literal, Tuple, List, Optional
+
+import torch
+from mmcv.cnn import ConvModule
+from mmdet.structures.bbox import bbox2roi
+from mmdet.structures.mask import mask2bbox
+from torch import nn
+import torch.nn.functional as F
+from mmengine import MMLogger
+from mmengine.model import BaseModule
+from mmdet.registry import MODELS
+
+from ext.sam import MaskDecoder
+from ext.sam.mask_decoder import MLP as SAMMLP
+from ext.meta.sam_meta import meta_dict, checkpoint_dict
+from utils.load_checkpoint import load_checkpoint_with_prefix
+
+
+@MODELS.register_module()
+class OVSAMHead(BaseModule):
+
+    def __init__(
+            self,
+            model_name: Literal['vit_h', 'vit_l', 'vit_b'] = 'vit_h',
+            with_label_token: bool = False,
+            ov_classifier_name: Optional[str] = None,
+            logit: Optional[float] = None,
+            roi_extractor=None,
+            fix: bool = True,
+            init_cfg=None,
+            cur_mask=1,
+            roi_extractor_single=None,
+            load_roi_conv=None,
+            gen_box=False,
+    ):
+        assert init_cfg is not None and \
+               init_cfg['type'] in ['sam_pretrain', 'Pretrained'], f"{init_cfg['type']} is not supported."
+        pretrained = init_cfg['checkpoint']
+        super().__init__(init_cfg=None)
+        self.init_cfg = init_cfg
+        self.logger = MMLogger.get_current_instance()
+        if roi_extractor_single is not None:
+            self.roi_extractor_single = MODELS.build(roi_extractor_single)
+            self.roi_merge_proj = nn.Linear(768 * 2, 768)
+        else:
+            self.roi_extractor_single = None
+            self.roi_merge_proj = None
+
+        mask_decoder = MaskDecoder(
+            num_multimask_outputs=cur_mask - 1,
+            transformer_dim=meta_dict[model_name]['prompt_embed_dim'],
+            iou_head_depth=3,
+            iou_head_hidden_dim=256,
+            with_iou=False
+        )
+
+        if self.init_cfg['type'] == 'sam_pretrain':
+            raise NotImplementedError
+
+        self.mask_decoder = mask_decoder
+
+        self.with_label_token = with_label_token
+
+        if self.with_label_token:
+            ov_path = os.path.join(os.path.expanduser('./models/'), f"{ov_classifier_name}.pth")
+            cls_embed = torch.load(ov_path)
+            cls_embed_norm = cls_embed.norm(p=2, dim=-1)
+            assert torch.allclose(cls_embed_norm, torch.ones_like(cls_embed_norm))
+
+            _dim = cls_embed.size(2)
+            _prototypes = cls_embed.size(1)
+            back_token = torch.zeros(1, _dim, dtype=torch.float32, device='cpu')
+            cls_embed = torch.cat([
+                cls_embed, back_token.repeat(_prototypes, 1)[None]
+            ], dim=0)
+            self.register_buffer('cls_embed', cls_embed.permute(2, 0, 1).contiguous(), persistent=False)
+
+            if logit is None:
+                logit_scale = torch.tensor(4.6052, dtype=torch.float32)
+            else:
+                logit_scale = torch.tensor(logit, dtype=torch.float32)
+            self.register_buffer('logit_scale', logit_scale, persistent=False)
+
+            transformer_dim = self.mask_decoder.mask_tokens.weight.shape[1]
+            self.label_token = nn.Embedding(1, transformer_dim)
+            self.label_mlp = SAMMLP(transformer_dim, transformer_dim, _dim, 3)
+
+        self.gen_box = gen_box
+
+        if roi_extractor is not None:
+            self.roi = MODELS.build(roi_extractor)
+            self.roi_conv = nn.Sequential(*[
+                ConvModule(in_channels=self.roi.out_channels, out_channels=_dim, kernel_size=1, bias=False)
+            ])
+        else:
+            self.roi = None
+
+        if self.init_cfg['type'] == 'Pretrained':
+            checkpoint_path = pretrained
+            state_dict = load_checkpoint_with_prefix(checkpoint_path, prefix=self.init_cfg['prefix'])
+            self.load_state_dict(state_dict, strict=True)
+        if roi_extractor is not None and load_roi_conv is not None:
+            checkpoint_path = load_roi_conv['checkpoint']
+            state_dict = load_checkpoint_with_prefix(checkpoint_path, prefix=load_roi_conv['prefix'])
+            self.roi_conv.load_state_dict(state_dict, strict=True)
+
+        self.fix = fix
+
+        if self.fix:
+            self.train(mode=False)
+            for name, param in self.named_parameters():
+                param.requires_grad = False
+
+    def init_weights(self):
+        self.logger.info(f"Init Config for {self.__class__.__name__}")
+        self.logger.info(self.init_cfg)
+
+    def forward_logit(self, cls_embd):
+        cls_pred = torch.einsum('bnc,ckp->bnkp', F.normalize(cls_embd, dim=-1), self.cls_embed)
+        cls_pred = cls_pred.max(-1).values
+        cls_pred = self.logit_scale.exp() * cls_pred
+        return cls_pred
+
+    def predict_masks(
+            self,
+            image_embeddings: torch.Tensor,
+            image_pe: torch.Tensor,
+            sparse_prompt_embeddings: torch.Tensor,
+            dense_prompt_embeddings: torch.Tensor,
+            fpn_feats: List[torch.Tensor],
+            roi_list: Optional[List[torch.Tensor]],
+            backbone_feature: torch.Tensor,
+            backbone=None
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Predicts masks. See 'forward' for more details."""
+        num_instances = int(sparse_prompt_embeddings.size(0))
+        # Concatenate output tokens
+        output_tokens = torch.cat([
+            self.label_token.weight,
+            self.mask_decoder.mask_tokens.weight], dim=0
+        )
+        output_tokens = output_tokens.unsqueeze(0).expand(num_instances, -1, -1)
+        queries = torch.cat((output_tokens, sparse_prompt_embeddings), dim=1)
+
+        # image_embeddings = torch.repeat_interleave(image_embeddings, num_instances, dim=0)
+        image_embeddings = image_embeddings + dense_prompt_embeddings
+        pos_img = torch.repeat_interleave(image_pe, num_instances, dim=0)
+        b, c, h, w = image_embeddings.shape
+
+        # Run the transformer
+        queries, mask_feats = self.mask_decoder.transformer(image_embeddings, pos_img, queries)
+        label_query = queries[:, 0, :]
+        mask_embeds = queries[:, 1:(1 + self.mask_decoder.num_mask_tokens), :]
+
+        # Upscale mask embeddings and predict masks using the mask tokens
+        mask_feats = mask_feats.transpose(1, 2).view(b, c, h, w)
+        mask_feats = self.mask_decoder.output_upscaling(mask_feats)
+        mask_queries_list: List[torch.Tensor] = []
+        for i in range(self.mask_decoder.num_mask_tokens):
+            mask_queries_list.append(self.mask_decoder.output_hypernetworks_mlps[i](mask_embeds[:, i, :]))
+        mask_queries = torch.stack(mask_queries_list, dim=1)
+        b, c, h, w = mask_feats.shape
+        masks = (mask_queries @ mask_feats.view(b, c, h * w)).view(b, -1, h, w)
+
+        # Generate class labels
+        if self.with_label_token:
+            cls_embed_list = []
+            assert self.mask_decoder.num_mask_tokens == 1
+            for i in range(self.mask_decoder.num_mask_tokens):
+                cls_embed_list.append(self.label_mlp(label_query))
+            cls_embed = torch.stack(cls_embed_list, dim=1)
+            if self.gen_box:
+                bboxes = mask2bbox(masks.sigmoid()[:, 0] > 0.5) * 4
+                roi_list = bbox2roi([bboxes])
+            roi_feats = self.roi(fpn_feats, roi_list)
+            roi_feats = self.roi_conv(roi_feats)
+            roi_feats = roi_feats.mean(dim=-1).mean(dim=-1)
+            if self.roi_extractor_single:
+                roi_feats_clip = self.roi_extractor_single(
+                    backbone.get_clip_feature(backbone_feature[-1:]), roi_list
+                )
+                roi_feats_clip = backbone.forward_feat(roi_feats_clip)
+                roi_feats = self.roi_merge_proj(torch.cat([roi_feats, roi_feats_clip], dim=-1))
+            roi_feats = roi_feats[:, None] + 0 * cls_embed
+            cls_pred = self.forward_logit(roi_feats)
+        else:
+            cls_pred = None
+        return masks, None, cls_pred
+
+    def forward(
+            self,
+            image_embeddings: torch.Tensor,
+            image_pe: torch.Tensor,
+            sparse_prompt_embeddings: torch.Tensor,
+            dense_prompt_embeddings: torch.Tensor,
+            multi_mask_output: bool,
+            data_samples=None,
+            fpn_feats=None,
+            backbone_feats=None,
+            backbone=None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor]:
+        num_prompts = len(sparse_prompt_embeddings)
+        image_embeddings = torch.repeat_interleave(image_embeddings, num_prompts, dim=0)
+
+        masks, _, cls_pred = self.predict_masks(
+            image_embeddings=image_embeddings,
+            image_pe=image_pe,
+            sparse_prompt_embeddings=sparse_prompt_embeddings,
+            dense_prompt_embeddings=dense_prompt_embeddings,
+            fpn_feats=fpn_feats,
+            roi_list=None,
+            backbone_feature=backbone_feats,
+            backbone=backbone,
+        )
+
+        # Select the correct mask or masks for output
+        if multi_mask_output:
+            mask_slice = slice(1, None)
+        else:
+            mask_slice = slice(0, 1)
+        masks = masks[:, mask_slice, :, :]
+
+        # Prepare output
+        return masks, None, cls_pred

app/models/sam_backbone.py

@@ -0,0 +1,113 @@
+from functools import partial
+from typing import Literal
+
+import torch
+import torch.nn as nn
+from mmdet.registry import MODELS
+
+from mmengine.model import BaseModule
+from mmengine.logging import MMLogger
+
+from ext.sam import ImageEncoderViT
+from ext.meta.sam_meta import meta_dict, checkpoint_dict
+from utils.load_checkpoint import load_checkpoint_with_prefix
+
+
+@MODELS.register_module()
+class SAMBackbone(BaseModule):
+
+    def __init__(
+            self,
+            model_name: Literal['vit_h', 'vit_l', 'vit_b'] = 'vit_h',
+            fix: bool = True,
+            init_cfg=None,
+    ):
+        assert init_cfg is not None and init_cfg['type'] in \
+               ['sam_pretrain', 'Pretrained'], f"{init_cfg['type']} is not supported."
+        pretrained = init_cfg['checkpoint']
+        super().__init__(init_cfg=None)
+        self.init_cfg = init_cfg
+        self.logger = MMLogger.get_current_instance()
+
+        backbone_meta = meta_dict[model_name]
+
+        backbone = ImageEncoderViT(
+            depth=backbone_meta['encoder_depth'],
+            embed_dim=backbone_meta['encoder_embed_dim'],
+            num_heads=backbone_meta['encoder_num_heads'],
+            patch_size=backbone_meta['vit_patch_size'],
+            img_size=backbone_meta['image_size'],
+            global_attn_indexes=backbone_meta['encoder_global_attn_indexes'],
+            out_chans=backbone_meta['prompt_embed_dim'],
+            norm_layer=partial(torch.nn.LayerNorm, eps=1e-6),
+            qkv_bias=True,
+            use_rel_pos=True,
+            mlp_ratio=4,
+            window_size=14,
+        )
+        if self.init_cfg['type'] == 'sam_pretrain':
+            checkpoint_path = checkpoint_dict[pretrained]
+            state_dict = load_checkpoint_with_prefix(checkpoint_path, prefix='image_encoder')
+            backbone.load_state_dict(state_dict, strict=True)
+
+        self.stem = backbone.patch_embed
+        self.pos_embed = backbone.pos_embed
+
+        self.res_layers = []
+        last_pos = 0
+        for idx, cur_pos in enumerate(backbone_meta['encoder_global_attn_indexes']):
+            blocks = backbone.blocks[last_pos:cur_pos + 1]
+            layer_name = f'layer{idx + 1}'
+            self.add_module(layer_name, nn.Sequential(*blocks))
+            self.res_layers.append(layer_name)
+            last_pos = cur_pos + 1
+
+        self.out_proj = backbone.neck
+
+        if self.init_cfg['type'] == 'Pretrained':
+            checkpoint_path = pretrained
+            state_dict = load_checkpoint_with_prefix(checkpoint_path, prefix=self.init_cfg['prefix'])
+            self.load_state_dict(state_dict, strict=True)
+
+        self.model_name = model_name
+        self.fix = fix
+        self.model_type = 'vit'
+        self.output_channels = None
+        self.out_indices = (0, 1, 2, 3)
+        if self.fix:
+            self.train(mode=False)
+            for name, param in self.named_parameters():
+                param.requires_grad = False
+
+    def init_weights(self):
+        self.logger.info(f"Init Config for {self.model_name}")
+        self.logger.info(self.init_cfg)
+
+    def train(self: torch.nn.Module, mode: bool = True) -> torch.nn.Module:
+        if not isinstance(mode, bool):
+            raise ValueError("training mode is expected to be boolean")
+        if self.fix:
+            super().train(mode=False)
+        else:
+            super().train(mode=mode)
+        return self
+
+    def forward_func(self, x):
+        x = self.stem(x)
+        x = x + self.pos_embed
+        outs = []
+        for i, layer_name in enumerate(self.res_layers):
+            res_layer = getattr(self, layer_name)
+            x = res_layer(x)
+            if i in self.out_indices:
+                outs.append(x.permute(0, 3, 1, 2).contiguous())
+        outs[-1] = self.out_proj(outs[-1])
+        return tuple(outs)
+
+    def forward(self, x):
+        if self.fix:
+            with torch.no_grad():
+                outs = self.forward_func(x)
+        else:
+            outs = self.forward_func(x)
+        return outs

app/models/sam_mask_decoder.py

@@ -0,0 +1,140 @@
+from typing import Literal, Tuple, List
+
+import torch
+import torch.nn.functional as F
+from mmdet.structures import SampleList
+from mmengine import MMLogger
+from mmengine.model import BaseModule
+from mmdet.registry import MODELS
+
+from ext.sam import MaskDecoder
+from ext.meta.sam_meta import meta_dict, checkpoint_dict
+from utils.load_checkpoint import load_checkpoint_with_prefix
+
+
+@MODELS.register_module()
+class SAMMaskDecoder(BaseModule):
+
+    def __init__(
+            self,
+            model_name: Literal['vit_h', 'vit_l', 'vit_b'] = 'vit_h',
+            fix: bool = True,
+            init_cfg=None,
+    ):
+        assert init_cfg is not None and \
+               init_cfg['type'] in ['sam_pretrain', 'Pretrained'], f"{init_cfg['type']} is not supported."
+        pretrained = init_cfg['checkpoint']
+        super().__init__(init_cfg=None)
+        self.init_cfg = init_cfg
+        self.logger = MMLogger.get_current_instance()
+
+        mask_decoder = MaskDecoder(
+            num_multimask_outputs=3,
+            transformer_dim=meta_dict[model_name]['prompt_embed_dim'],
+            iou_head_depth=3,
+            iou_head_hidden_dim=256,
+        )
+        if self.init_cfg['type'] == 'sam_pretrain':
+            checkpoint_path = checkpoint_dict[pretrained]
+            state_dict = load_checkpoint_with_prefix(checkpoint_path, prefix='mask_decoder')
+            mask_decoder.load_state_dict(state_dict, strict=True)
+
+        self.mask_decoder = mask_decoder
+        if self.init_cfg['type'] == 'Pretrained':
+            checkpoint_path = pretrained
+            state_dict = load_checkpoint_with_prefix(checkpoint_path, prefix=self.init_cfg['prefix'])
+            self.load_state_dict(state_dict, strict=True)
+
+        self.fix = fix
+        if self.fix:
+            self.train(mode=False)
+            for name, param in self.named_parameters():
+                param.requires_grad = False
+
+    def init_weights(self):
+        self.logger.info(f"Init Config for {self.__class__.__name__}")
+        self.logger.info(self.init_cfg)
+
+    def forward_logit(self, cls_embd):
+        cls_pred = torch.einsum('bnc,ckp->bnkp', F.normalize(cls_embd, dim=-1), self.cls_embed)
+        cls_pred = cls_pred.max(-1).values
+        cls_pred = self.logit_scale.exp() * cls_pred
+        return cls_pred
+
+    def predict_masks(
+            self,
+            image_embeddings: torch.Tensor,
+            image_pe: torch.Tensor,
+            sparse_prompt_embeddings: torch.Tensor,
+            dense_prompt_embeddings: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Predicts masks. See 'forward' for more details."""
+        num_instances = int(sparse_prompt_embeddings.shape[0])
+        # Concatenate output tokens
+        output_tokens = torch.cat([self.mask_decoder.iou_token.weight, self.mask_decoder.mask_tokens.weight], dim=0)
+        output_tokens = output_tokens.unsqueeze(0).expand(num_instances, -1, -1)
+        queries = torch.cat((output_tokens, sparse_prompt_embeddings), dim=1)
+
+        # image_embeddings = torch.repeat_interleave(image_embeddings, num_instances, dim=0)
+        image_embeddings = image_embeddings + dense_prompt_embeddings
+        pos_img = torch.repeat_interleave(image_pe, num_instances, dim=0)
+        b, c, h, w = image_embeddings.shape
+
+        # Run the transformer
+        queries, mask_feats = self.mask_decoder.transformer(image_embeddings, pos_img, queries)
+        iou_query = queries[:, 0, :]
+        mask_embeds = queries[:, 1:(1 + self.mask_decoder.num_mask_tokens), :]
+
+        # Upscale mask embeddings and predict masks using the mask tokens
+        mask_feats = mask_feats.transpose(1, 2).view(b, c, h, w)
+        mask_feats = self.mask_decoder.output_upscaling(mask_feats)
+        mask_queries_list: List[torch.Tensor] = []
+        for i in range(self.mask_decoder.num_mask_tokens):
+            mask_queries_list.append(self.mask_decoder.output_hypernetworks_mlps[i](mask_embeds[:, i, :]))
+        mask_queries = torch.stack(mask_queries_list, dim=1)
+        b, c, h, w = mask_feats.shape
+        masks = (mask_queries @ mask_feats.view(b, c, h * w)).view(b, -1, h, w)
+
+        # Generate mask quality predictions
+        iou_pred = self.mask_decoder.iou_prediction_head(iou_query)
+
+        return masks, iou_pred, None
+
+    def forward(
+            self,
+            image_embeddings: torch.Tensor,
+            image_pe: torch.Tensor,
+            sparse_prompt_embeddings: torch.Tensor,
+            dense_prompt_embeddings: torch.Tensor,
+            multi_mask_output: bool,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        num_prompts = len(sparse_prompt_embeddings)
+        image_embeddings = torch.repeat_interleave(image_embeddings, num_prompts, dim=0)
+        masks, iou_pred, cls_pred = self.predict_masks(
+            image_embeddings=image_embeddings,
+            image_pe=image_pe,
+            sparse_prompt_embeddings=sparse_prompt_embeddings,
+            dense_prompt_embeddings=dense_prompt_embeddings,
+        )
+
+        # Select the correct mask or masks for output
+        if multi_mask_output:
+            mask_slice = slice(1, None)
+        else:
+            mask_slice = slice(0, 1)
+        masks = masks[:, mask_slice, :, :]
+        iou_pred = iou_pred[:, mask_slice]
+
+        # Prepare output
+        return masks, iou_pred, cls_pred
+
+    def forward_train(
+            self,
+            image_embeddings: torch.Tensor,
+            image_pe: torch.Tensor,
+            sparse_prompt_embeddings: torch.Tensor,
+            dense_prompt_embeddings: torch.Tensor,
+            batch_ind_list: List[int],
+            data_samples: SampleList,
+    ):
+        raise NotImplementedError

app/models/sam_pe.py

@@ -0,0 +1,152 @@
+from typing import Tuple, Literal
+
+import torch
+from mmengine import MMLogger
+
+from mmdet.registry import MODELS
+from mmengine.model import BaseModule
+from mmengine.structures import InstanceData
+
+from ext.sam import PromptEncoder
+from ext.meta.sam_meta import meta_dict, checkpoint_dict
+from utils.load_checkpoint import load_checkpoint_with_prefix
+
+
+@MODELS.register_module()
+class SAMPromptEncoder(BaseModule):
+
+    def __init__(
+            self,
+            model_name: Literal['vit_h', 'vit_l', 'vit_b'] = 'vit_h',
+            fix: bool = True,
+            init_cfg=None,
+    ):
+        assert init_cfg is not None and init_cfg['type'] == 'sam_pretrain', f"{init_cfg['type']} is not supported."
+        pretrained = init_cfg['checkpoint']
+        super().__init__(init_cfg=None)
+        self.init_cfg = init_cfg
+        self.logger = MMLogger.get_current_instance()
+
+        backbone_meta = meta_dict[model_name]
+        checkpoint_path = checkpoint_dict[pretrained]
+
+        prompt_encoder = PromptEncoder(
+            embed_dim=256,
+            image_embedding_size=(backbone_meta['image_embedding_size'], backbone_meta['image_embedding_size']),
+            input_image_size=(backbone_meta['image_size'], backbone_meta['image_size']),
+            mask_in_chans=16,
+        )
+        state_dict = load_checkpoint_with_prefix(checkpoint_path, prefix='prompt_encoder')
+        prompt_encoder.load_state_dict(state_dict, strict=True)
+
+        # meta
+        self.embed_dim = prompt_encoder.embed_dim
+        self.input_image_size = prompt_encoder.input_image_size
+        self.image_embedding_size = prompt_encoder.image_embedding_size
+        self.num_point_embeddings = 4
+        self.mask_input_size = prompt_encoder.mask_input_size
+
+        # positional encoding
+        self.pe_layer = prompt_encoder.pe_layer
+
+        # mask encoding
+        self.mask_downscaling = prompt_encoder.mask_downscaling
+        self.no_mask_embed = prompt_encoder.no_mask_embed
+
+        # point encoding
+        self.point_embeddings = prompt_encoder.point_embeddings
+        self.not_a_point_embed = prompt_encoder.not_a_point_embed
+
+        self.fix = fix
+        if self.fix:
+            self.train(mode=False)
+            for name, param in self.named_parameters():
+                param.requires_grad = False
+
+    @property
+    def device(self):
+        return self.no_mask_embed.weight.device
+
+    def init_weights(self):
+        self.logger.info(f"Init Config for {self.__class__.__name__}")
+        self.logger.info(self.init_cfg)
+
+    def train(self: torch.nn.Module, mode: bool = True) -> torch.nn.Module:
+        if not isinstance(mode, bool):
+            raise ValueError("training mode is expected to be boolean")
+        if self.fix:
+            super().train(mode=False)
+        else:
+            super().train(mode=mode)
+        return self
+
+    def _embed_boxes(self, bboxes: torch.Tensor, image_size: Tuple[int, int]) -> torch.Tensor:
+        """Embeds box prompts."""
+        bboxes = bboxes + 0.5  # Shift to center of pixel
+        coords = bboxes.reshape(-1, 2, 2)
+        corner_embedding = self.pe_layer.forward_with_coords(coords, image_size)
+        corner_embedding[:, 0, :] += self.point_embeddings[2].weight
+        corner_embedding[:, 1, :] += self.point_embeddings[3].weight
+        return corner_embedding
+
+    def _embed_masks(self, masks: torch.Tensor) -> torch.Tensor:
+        """Embeds mask inputs."""
+        mask_embedding = self.mask_downscaling(masks)
+        return mask_embedding
+
+    def get_dense_pe(self) -> torch.Tensor:
+        return self.pe_layer(self.image_embedding_size).unsqueeze(0)
+
+    def _embed_points(
+        self,
+        points: torch.Tensor,
+        labels: torch.Tensor,
+        pad: bool,
+    ) -> torch.Tensor:
+        """Embeds point prompts."""
+        points = points + 0.5  # Shift to center of pixel
+        if pad:
+            padding_point = torch.zeros((points.shape[0], 1, 2), device=points.device)
+            padding_label = -torch.ones((labels.shape[0], 1), device=labels.device)
+            points = torch.cat([points, padding_point], dim=1)
+            labels = torch.cat([labels, padding_label], dim=1)
+        point_embedding = self.pe_layer.forward_with_coords(points, self.input_image_size)
+        point_embedding[labels == -1] = 0.0
+        point_embedding[labels == -1] += self.not_a_point_embed.weight
+        point_embedding[labels == 0] += self.point_embeddings[0].weight
+        point_embedding[labels == 1] += self.point_embeddings[1].weight
+        return point_embedding
+
+    def forward(
+            self,
+            instances: InstanceData,
+            image_size: Tuple[int, int],
+            with_points: bool = False,
+            with_bboxes: bool = False,
+            with_masks: bool = False,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        assert with_points or with_bboxes or with_masks
+        bs = len(instances)
+        sparse_embeddings = torch.empty((bs, 0, self.embed_dim), device=self.device)
+        if with_points:
+            assert 'point_coords' in instances
+            coords = instances.point_coords
+            labels = torch.ones_like(coords)[:, :, 0]
+            point_embeddings = self._embed_points(coords, labels, pad=not with_bboxes)
+            sparse_embeddings = torch.cat([sparse_embeddings, point_embeddings], dim=1)
+
+        if with_bboxes:
+            assert 'bboxes' in instances
+            box_embeddings = self._embed_boxes(
+                instances.bboxes, image_size=image_size
+            )
+            sparse_embeddings = torch.cat([sparse_embeddings, box_embeddings], dim=1)
+
+        if with_masks:
+            assert 'masks' in instances
+            dense_embeddings = self._embed_masks(instances.masks.masks)
+        else:
+            dense_embeddings = self.no_mask_embed.weight.reshape(1, -1, 1, 1).expand(
+                bs, -1, self.image_embedding_size[0], self.image_embedding_size[1]
+            )
+        return sparse_embeddings, dense_embeddings

app/models/transformer_neck.py

@@ -0,0 +1,158 @@
+from functools import partial
+from typing import Tuple, List, Optional
+
+import torch
+from torch import Tensor, nn
+
+from mmengine.model import BaseModule, normal_init
+from mmdet.registry import MODELS
+from mmdet.models.layers import PatchEmbed
+
+from ext.meta.sam_meta import checkpoint_dict
+from ext.sam.common import LayerNorm2d
+from ext.sam.image_encoder import Block
+
+from utils.load_checkpoint import load_checkpoint_with_prefix
+
+
+@MODELS.register_module()
+class MultiLayerTransformerNeck(BaseModule):
+    STRIDE = 16
+
+    def __init__(
+            self,
+            input_size: Tuple[int, int],
+            in_channels: List[int],
+            embed_channels: int,
+            out_channels: int,
+            layer_ids: Tuple[int] = (0, 1, 2, 3),
+            strides: Tuple[int] = (4, 8, 16, 32),
+            embedding_path: Optional[str] = None,
+            fix=False,
+            init_cfg=None
+    ) -> None:
+        super().__init__(init_cfg=None)
+
+        self.transformer_size = (input_size[0] // self.STRIDE, input_size[1] // self.STRIDE)
+        self.layer_ids = layer_ids
+
+        self.patch_embeds = nn.ModuleList()
+        for idx, in_ch in enumerate(in_channels):
+            if idx in layer_ids:
+                if strides[idx] > self.STRIDE:
+                    patch_embed = PatchEmbed(
+                        conv_type=nn.ConvTranspose2d,
+                        in_channels=in_ch,
+                        embed_dims=embed_channels,
+                        kernel_size=strides[idx] // self.STRIDE,
+                        stride=strides[idx] // self.STRIDE,
+                        input_size=(input_size[0] // strides[idx], input_size[1] // strides[idx])
+                    )
+                else:
+                    patch_embed = PatchEmbed(
+                        in_channels=in_ch,
+                        embed_dims=embed_channels,
+                        kernel_size=self.STRIDE // strides[idx],
+                        stride=self.STRIDE // strides[idx],
+                        input_size=(input_size[0] // strides[idx], input_size[1] // strides[idx])
+                    )
+                self.patch_embeds.append(patch_embed)
+            else:
+                self.patch_embeds.append(nn.Identity())
+
+        if embedding_path is not None:
+            assert embedding_path.startswith('sam_')
+            embedding_ckpt = embedding_path.split('_', maxsplit=1)[1]
+            path = checkpoint_dict[embedding_ckpt]
+            state_dict = load_checkpoint_with_prefix(path, prefix='image_encoder')
+            pos_embed = state_dict['pos_embed']
+        else:
+            # For loading from checkpoint
+            pos_embed = torch.zeros(1, input_size[0] // self.STRIDE, input_size[1] // self.STRIDE, embed_channels)
+
+        self.register_buffer('pos_embed', pos_embed)
+
+        self.level_encoding = nn.Embedding(len(layer_ids), embed_channels)
+
+        depth = 5
+        global_attn_indexes = [4]
+        window_size = 14
+
+        self.blocks = nn.ModuleList()
+        for i in range(depth):
+            block = Block(
+                dim=embed_channels,
+                num_heads=16,
+                mlp_ratio=4,
+                qkv_bias=True,
+                norm_layer=partial(torch.nn.LayerNorm, eps=1e-6),
+                act_layer=nn.GELU,
+                use_rel_pos=True,
+                rel_pos_zero_init=True,
+                window_size=window_size if i not in global_attn_indexes else 0,
+                input_size=self.transformer_size,
+            )
+            self.blocks.append(block)
+
+        self.neck = nn.Sequential(
+            nn.Conv2d(
+                embed_channels,
+                out_channels,
+                kernel_size=1,
+                bias=False,
+            ),
+            LayerNorm2d(out_channels),
+            nn.Conv2d(
+                out_channels,
+                out_channels,
+                kernel_size=3,
+                padding=1,
+                bias=False,
+            ),
+            LayerNorm2d(out_channels),
+        )
+
+        self.fix = fix
+        if self.fix:
+            self.train(mode=False)
+            for name, param in self.named_parameters():
+                param.requires_grad = False
+
+        if init_cfg is not None:
+            assert init_cfg['type'] == 'Pretrained'
+            checkpoint_path = init_cfg['checkpoint']
+            state_dict = load_checkpoint_with_prefix(checkpoint_path, prefix=init_cfg['prefix'])
+            self.load_state_dict(state_dict, strict=True)
+            self._is_init = True
+
+    def init_weights(self):
+        normal_init(self.level_encoding, mean=0, std=1)
+
+    def train(self: torch.nn.Module, mode: bool = True) -> torch.nn.Module:
+        if not isinstance(mode, bool):
+            raise ValueError("training mode is expected to be boolean")
+        if self.fix:
+            super().train(mode=False)
+        else:
+            super().train(mode=mode)
+        return self
+
+    def forward(self, inputs: Tuple[Tensor]) -> Tensor:
+        input_embeddings = []
+        level_cnt = 0
+        for idx, feat in enumerate(inputs):
+            if idx not in self.layer_ids:
+                continue
+            feat, size = self.patch_embeds[idx](feat)
+            feat = feat.unflatten(1, size)
+            feat = feat + self.level_encoding.weight[level_cnt]
+            input_embeddings.append(feat)
+            level_cnt += 1
+
+        feat = sum(input_embeddings)
+        feat = feat + self.pos_embed
+        for block in self.blocks:
+            feat = block(feat)
+        feat = feat.permute(0, 3, 1, 2).contiguous()
+        feat = self.neck(feat)
+        return feat

ext/class_names/lvis_list.py

@@ -0,0 +1,242 @@
+LVIS_CLASSES = ('aerosol_can', 'air_conditioner', 'airplane', 'alarm_clock',
+                'alcohol', 'alligator', 'almond', 'ambulance', 'amplifier', 'anklet',
+                'antenna', 'apple', 'applesauce', 'apricot', 'apron', 'aquarium',
+                'arctic_(type_of_shoe)', 'armband', 'armchair', 'armoire', 'armor',
+                'artichoke', 'trash_can', 'ashtray', 'asparagus', 'atomizer',
+                'avocado', 'award', 'awning', 'ax', 'baboon', 'baby_buggy',
+                'basketball_backboard', 'backpack', 'handbag', 'suitcase', 'bagel',
+                'bagpipe', 'baguet', 'bait', 'ball', 'ballet_skirt', 'balloon',
+                'bamboo', 'banana', 'Band_Aid', 'bandage', 'bandanna', 'banjo',
+                'banner', 'barbell', 'barge', 'barrel', 'barrette', 'barrow',
+                'baseball_base', 'baseball', 'baseball_bat', 'baseball_cap',
+                'baseball_glove', 'basket', 'basketball', 'bass_horn', 'bat_(animal)',
+                'bath_mat', 'bath_towel', 'bathrobe', 'bathtub', 'batter_(food)',
+                'battery', 'beachball', 'bead', 'bean_curd', 'beanbag', 'beanie',
+                'bear', 'bed', 'bedpan', 'bedspread', 'cow', 'beef_(food)', 'beeper',
+                'beer_bottle', 'beer_can', 'beetle', 'bell', 'bell_pepper', 'belt',
+                'belt_buckle', 'bench', 'beret', 'bib', 'Bible', 'bicycle', 'visor',
+                'billboard', 'binder', 'binoculars', 'bird', 'birdfeeder', 'birdbath',
+                'birdcage', 'birdhouse', 'birthday_cake', 'birthday_card',
+                'pirate_flag', 'black_sheep', 'blackberry', 'blackboard', 'blanket',
+                'blazer', 'blender', 'blimp', 'blinker', 'blouse', 'blueberry',
+                'gameboard', 'boat', 'bob', 'bobbin', 'bobby_pin', 'boiled_egg',
+                'bolo_tie', 'deadbolt', 'bolt', 'bonnet', 'book', 'bookcase',
+                'booklet', 'bookmark', 'boom_microphone', 'boot', 'bottle',
+                'bottle_opener', 'bouquet', 'bow_(weapon)',
+                'bow_(decorative_ribbons)', 'bow-tie', 'bowl', 'pipe_bowl',
+                'bowler_hat', 'bowling_ball', 'box', 'boxing_glove', 'suspenders',
+                'bracelet', 'brass_plaque', 'brassiere', 'bread-bin', 'bread',
+                'breechcloth', 'bridal_gown', 'briefcase', 'broccoli', 'broach',
+                'broom', 'brownie', 'brussels_sprouts', 'bubble_gum', 'bucket',
+                'horse_buggy', 'bull', 'bulldog', 'bulldozer', 'bullet_train',
+                'bulletin_board', 'bulletproof_vest', 'bullhorn', 'bun', 'bunk_bed',
+                'buoy', 'burrito', 'bus_(vehicle)', 'business_card', 'butter',
+                'butterfly', 'button', 'cab_(taxi)', 'cabana', 'cabin_car', 'cabinet',
+                'locker', 'cake', 'calculator', 'calendar', 'calf', 'camcorder',
+                'camel', 'camera', 'camera_lens', 'camper_(vehicle)', 'can',
+                'can_opener', 'candle', 'candle_holder', 'candy_bar', 'candy_cane',
+                'walking_cane', 'canister', 'canoe', 'cantaloup', 'canteen',
+                'cap_(headwear)', 'bottle_cap', 'cape', 'cappuccino',
+                'car_(automobile)', 'railcar_(part_of_a_train)', 'elevator_car',
+                'car_battery', 'identity_card', 'card', 'cardigan', 'cargo_ship',
+                'carnation', 'horse_carriage', 'carrot', 'tote_bag', 'cart', 'carton',
+                'cash_register', 'casserole', 'cassette', 'cast', 'cat',
+                'cauliflower', 'cayenne_(spice)', 'CD_player', 'celery',
+                'cellular_telephone', 'chain_mail', 'chair', 'chaise_longue',
+                'chalice', 'chandelier', 'chap', 'checkbook', 'checkerboard',
+                'cherry', 'chessboard', 'chicken_(animal)', 'chickpea',
+                'chili_(vegetable)', 'chime', 'chinaware', 'crisp_(potato_chip)',
+                'poker_chip', 'chocolate_bar', 'chocolate_cake', 'chocolate_milk',
+                'chocolate_mousse', 'choker', 'chopping_board', 'chopstick',
+                'Christmas_tree', 'slide', 'cider', 'cigar_box', 'cigarette',
+                'cigarette_case', 'cistern', 'clarinet', 'clasp', 'cleansing_agent',
+                'cleat_(for_securing_rope)', 'clementine', 'clip', 'clipboard',
+                'clippers_(for_plants)', 'cloak', 'clock', 'clock_tower',
+                'clothes_hamper', 'clothespin', 'clutch_bag', 'coaster', 'coat',
+                'coat_hanger', 'coatrack', 'cock', 'cockroach', 'cocoa_(beverage)',
+                'coconut', 'coffee_maker', 'coffee_table', 'coffeepot', 'coil',
+                'coin', 'colander', 'coleslaw', 'coloring_material',
+                'combination_lock', 'pacifier', 'comic_book', 'compass',
+                'computer_keyboard', 'condiment', 'cone', 'control',
+                'convertible_(automobile)', 'sofa_bed', 'cooker', 'cookie',
+                'cooking_utensil', 'cooler_(for_food)', 'cork_(bottle_plug)',
+                'corkboard', 'corkscrew', 'edible_corn', 'cornbread', 'cornet',
+                'cornice', 'cornmeal', 'corset', 'costume', 'cougar', 'coverall',
+                'cowbell', 'cowboy_hat', 'crab_(animal)', 'crabmeat', 'cracker',
+                'crape', 'crate', 'crayon', 'cream_pitcher', 'crescent_roll', 'crib',
+                'crock_pot', 'crossbar', 'crouton', 'crow', 'crowbar', 'crown',
+                'crucifix', 'cruise_ship', 'police_cruiser', 'crumb', 'crutch',
+                'cub_(animal)', 'cube', 'cucumber', 'cufflink', 'cup', 'trophy_cup',
+                'cupboard', 'cupcake', 'hair_curler', 'curling_iron', 'curtain',
+                'cushion', 'cylinder', 'cymbal', 'dagger', 'dalmatian', 'dartboard',
+                'date_(fruit)', 'deck_chair', 'deer', 'dental_floss', 'desk',
+                'detergent', 'diaper', 'diary', 'die', 'dinghy', 'dining_table',
+                'tux', 'dish', 'dish_antenna', 'dishrag', 'dishtowel', 'dishwasher',
+                'dishwasher_detergent', 'dispenser', 'diving_board', 'Dixie_cup',
+                'dog', 'dog_collar', 'doll', 'dollar', 'dollhouse', 'dolphin',
+                'domestic_ass', 'doorknob', 'doormat', 'doughnut', 'dove',
+                'dragonfly', 'drawer', 'underdrawers', 'dress', 'dress_hat',
+                'dress_suit', 'dresser', 'drill', 'drone', 'dropper',
+                'drum_(musical_instrument)', 'drumstick', 'duck', 'duckling',
+                'duct_tape', 'duffel_bag', 'dumbbell', 'dumpster', 'dustpan', 'eagle',
+                'earphone', 'earplug', 'earring', 'easel', 'eclair', 'eel', 'egg',
+                'egg_roll', 'egg_yolk', 'eggbeater', 'eggplant', 'electric_chair',
+                'refrigerator', 'elephant', 'elk', 'envelope', 'eraser', 'escargot',
+                'eyepatch', 'falcon', 'fan', 'faucet', 'fedora', 'ferret',
+                'Ferris_wheel', 'ferry', 'fig_(fruit)', 'fighter_jet', 'figurine',
+                'file_cabinet', 'file_(tool)', 'fire_alarm', 'fire_engine',
+                'fire_extinguisher', 'fire_hose', 'fireplace', 'fireplug',
+                'first-aid_kit', 'fish', 'fish_(food)', 'fishbowl', 'fishing_rod',
+                'flag', 'flagpole', 'flamingo', 'flannel', 'flap', 'flash',
+                'flashlight', 'fleece', 'flip-flop_(sandal)', 'flipper_(footwear)',
+                'flower_arrangement', 'flute_glass', 'foal', 'folding_chair',
+                'food_processor', 'football_(American)', 'football_helmet',
+                'footstool', 'fork', 'forklift', 'freight_car', 'French_toast',
+                'freshener', 'frisbee', 'frog', 'fruit_juice', 'frying_pan', 'fudge',
+                'funnel', 'futon', 'gag', 'garbage', 'garbage_truck', 'garden_hose',
+                'gargle', 'gargoyle', 'garlic', 'gasmask', 'gazelle', 'gelatin',
+                'gemstone', 'generator', 'giant_panda', 'gift_wrap', 'ginger',
+                'giraffe', 'cincture', 'glass_(drink_container)', 'globe', 'glove',
+                'goat', 'goggles', 'goldfish', 'golf_club', 'golfcart',
+                'gondola_(boat)', 'goose', 'gorilla', 'gourd', 'grape', 'grater',
+                'gravestone', 'gravy_boat', 'green_bean', 'green_onion', 'griddle',
+                'grill', 'grits', 'grizzly', 'grocery_bag', 'guitar', 'gull', 'gun',
+                'hairbrush', 'hairnet', 'hairpin', 'halter_top', 'ham', 'hamburger',
+                'hammer', 'hammock', 'hamper', 'hamster', 'hair_dryer', 'hand_glass',
+                'hand_towel', 'handcart', 'handcuff', 'handkerchief', 'handle',
+                'handsaw', 'hardback_book', 'harmonium', 'hat', 'hatbox', 'veil',
+                'headband', 'headboard', 'headlight', 'headscarf', 'headset',
+                'headstall_(for_horses)', 'heart', 'heater', 'helicopter', 'helmet',
+                'heron', 'highchair', 'hinge', 'hippopotamus', 'hockey_stick', 'hog',
+                'home_plate_(baseball)', 'honey', 'fume_hood', 'hook', 'hookah',
+                'hornet', 'horse', 'hose', 'hot-air_balloon', 'hotplate', 'hot_sauce',
+                'hourglass', 'houseboat', 'hummingbird', 'hummus', 'polar_bear',
+                'icecream', 'popsicle', 'ice_maker', 'ice_pack', 'ice_skate',
+                'igniter', 'inhaler', 'iPod', 'iron_(for_clothing)', 'ironing_board',
+                'jacket', 'jam', 'jar', 'jean', 'jeep', 'jelly_bean', 'jersey',
+                'jet_plane', 'jewel', 'jewelry', 'joystick', 'jumpsuit', 'kayak',
+                'keg', 'kennel', 'kettle', 'key', 'keycard', 'kilt', 'kimono',
+                'kitchen_sink', 'kitchen_table', 'kite', 'kitten', 'kiwi_fruit',
+                'knee_pad', 'knife', 'knitting_needle', 'knob', 'knocker_(on_a_door)',
+                'koala', 'lab_coat', 'ladder', 'ladle', 'ladybug', 'lamb_(animal)',
+                'lamb-chop', 'lamp', 'lamppost', 'lampshade', 'lantern', 'lanyard',
+                'laptop_computer', 'lasagna', 'latch', 'lawn_mower', 'leather',
+                'legging_(clothing)', 'Lego', 'legume', 'lemon', 'lemonade',
+                'lettuce', 'license_plate', 'life_buoy', 'life_jacket', 'lightbulb',
+                'lightning_rod', 'lime', 'limousine', 'lion', 'lip_balm', 'liquor',
+                'lizard', 'log', 'lollipop', 'speaker_(stereo_equipment)', 'loveseat',
+                'machine_gun', 'magazine', 'magnet', 'mail_slot', 'mailbox_(at_home)',
+                'mallard', 'mallet', 'mammoth', 'manatee', 'mandarin_orange',
+                'manger', 'manhole', 'map', 'marker', 'martini', 'mascot',
+                'mashed_potato', 'masher', 'mask', 'mast', 'mat_(gym_equipment)',
+                'matchbox', 'mattress', 'measuring_cup', 'measuring_stick',
+                'meatball', 'medicine', 'melon', 'microphone', 'microscope',
+                'microwave_oven', 'milestone', 'milk', 'milk_can', 'milkshake',
+                'minivan', 'mint_candy', 'mirror', 'mitten', 'mixer_(kitchen_tool)',
+                'money', 'monitor_(computer_equipment) computer_monitor', 'monkey',
+                'motor', 'motor_scooter', 'motor_vehicle', 'motorcycle',
+                'mound_(baseball)', 'mouse_(computer_equipment)', 'mousepad',
+                'muffin', 'mug', 'mushroom', 'music_stool', 'musical_instrument',
+                'nailfile', 'napkin', 'neckerchief', 'necklace', 'necktie', 'needle',
+                'nest', 'newspaper', 'newsstand', 'nightshirt',
+                'nosebag_(for_animals)', 'noseband_(for_animals)', 'notebook',
+                'notepad', 'nut', 'nutcracker', 'oar', 'octopus_(food)',
+                'octopus_(animal)', 'oil_lamp', 'olive_oil', 'omelet', 'onion',
+                'orange_(fruit)', 'orange_juice', 'ostrich', 'ottoman', 'oven',
+                'overalls_(clothing)', 'owl', 'packet', 'inkpad', 'pad', 'paddle',
+                'padlock', 'paintbrush', 'painting', 'pajamas', 'palette',
+                'pan_(for_cooking)', 'pan_(metal_container)', 'pancake', 'pantyhose',
+                'papaya', 'paper_plate', 'paper_towel', 'paperback_book',
+                'paperweight', 'parachute', 'parakeet', 'parasail_(sports)',
+                'parasol', 'parchment', 'parka', 'parking_meter', 'parrot',
+                'passenger_car_(part_of_a_train)', 'passenger_ship', 'passport',
+                'pastry', 'patty_(food)', 'pea_(food)', 'peach', 'peanut_butter',
+                'pear', 'peeler_(tool_for_fruit_and_vegetables)', 'wooden_leg',
+                'pegboard', 'pelican', 'pen', 'pencil', 'pencil_box',
+                'pencil_sharpener', 'pendulum', 'penguin', 'pennant', 'penny_(coin)',
+                'pepper', 'pepper_mill', 'perfume', 'persimmon', 'person', 'pet',
+                'pew_(church_bench)', 'phonebook', 'phonograph_record', 'piano',
+                'pickle', 'pickup_truck', 'pie', 'pigeon', 'piggy_bank', 'pillow',
+                'pin_(non_jewelry)', 'pineapple', 'pinecone', 'ping-pong_ball',
+                'pinwheel', 'tobacco_pipe', 'pipe', 'pistol', 'pita_(bread)',
+                'pitcher_(vessel_for_liquid)', 'pitchfork', 'pizza', 'place_mat',
+                'plate', 'platter', 'playpen', 'pliers', 'plow_(farm_equipment)',
+                'plume', 'pocket_watch', 'pocketknife', 'poker_(fire_stirring_tool)',
+                'pole', 'polo_shirt', 'poncho', 'pony', 'pool_table', 'pop_(soda)',
+                'postbox_(public)', 'postcard', 'poster', 'pot', 'flowerpot',
+                'potato', 'potholder', 'pottery', 'pouch', 'power_shovel', 'prawn',
+                'pretzel', 'printer', 'projectile_(weapon)', 'projector', 'propeller',
+                'prune', 'pudding', 'puffer_(fish)', 'puffin', 'pug-dog', 'pumpkin',
+                'puncher', 'puppet', 'puppy', 'quesadilla', 'quiche', 'quilt',
+                'rabbit', 'race_car', 'racket', 'radar', 'radiator', 'radio_receiver',
+                'radish', 'raft', 'rag_doll', 'raincoat', 'ram_(animal)', 'raspberry',
+                'rat', 'razorblade', 'reamer_(juicer)', 'rearview_mirror', 'receipt',
+                'recliner', 'record_player', 'reflector', 'remote_control',
+                'rhinoceros', 'rib_(food)', 'rifle', 'ring', 'river_boat', 'road_map',
+                'robe', 'rocking_chair', 'rodent', 'roller_skate', 'Rollerblade',
+                'rolling_pin', 'root_beer', 'router_(computer_equipment)',
+                'rubber_band', 'runner_(carpet)', 'plastic_bag',
+                'saddle_(on_an_animal)', 'saddle_blanket', 'saddlebag', 'safety_pin',
+                'sail', 'salad', 'salad_plate', 'salami', 'salmon_(fish)',
+                'salmon_(food)', 'salsa', 'saltshaker', 'sandal_(type_of_shoe)',
+                'sandwich', 'satchel', 'saucepan', 'saucer', 'sausage', 'sawhorse',
+                'saxophone', 'scale_(measuring_instrument)', 'scarecrow', 'scarf',
+                'school_bus', 'scissors', 'scoreboard', 'scraper', 'screwdriver',
+                'scrubbing_brush', 'sculpture', 'seabird', 'seahorse', 'seaplane',
+                'seashell', 'sewing_machine', 'shaker', 'shampoo', 'shark',
+                'sharpener', 'Sharpie', 'shaver_(electric)', 'shaving_cream', 'shawl',
+                'shears', 'sheep', 'shepherd_dog', 'sherbert', 'shield', 'shirt',
+                'shoe', 'shopping_bag', 'shopping_cart', 'short_pants', 'shot_glass',
+                'shoulder_bag', 'shovel', 'shower_head', 'shower_cap',
+                'shower_curtain', 'shredder_(for_paper)', 'signboard', 'silo', 'sink',
+                'skateboard', 'skewer', 'ski', 'ski_boot', 'ski_parka', 'ski_pole',
+                'skirt', 'skullcap', 'sled', 'sleeping_bag', 'sling_(bandage)',
+                'slipper_(footwear)', 'smoothie', 'snake', 'snowboard', 'snowman',
+                'snowmobile', 'soap', 'soccer_ball', 'sock', 'sofa', 'softball',
+                'solar_array', 'sombrero', 'soup', 'soup_bowl', 'soupspoon',
+                'sour_cream', 'soya_milk', 'space_shuttle', 'sparkler_(fireworks)',
+                'spatula', 'spear', 'spectacles', 'spice_rack', 'spider', 'crawfish',
+                'sponge', 'spoon', 'sportswear', 'spotlight', 'squid_(food)',
+                'squirrel', 'stagecoach', 'stapler_(stapling_machine)', 'starfish',
+                'statue_(sculpture)', 'steak_(food)', 'steak_knife', 'steering_wheel',
+                'stepladder', 'step_stool', 'stereo_(sound_system)', 'stew',
+                'stirrer', 'stirrup', 'stool', 'stop_sign', 'brake_light', 'stove',
+                'strainer', 'strap', 'straw_(for_drinking)', 'strawberry',
+                'street_sign', 'streetlight', 'string_cheese', 'stylus', 'subwoofer',
+                'sugar_bowl', 'sugarcane_(plant)', 'suit_(clothing)', 'sunflower',
+                'sunglasses', 'sunhat', 'surfboard', 'sushi', 'mop', 'sweat_pants',
+                'sweatband', 'sweater', 'sweatshirt', 'sweet_potato', 'swimsuit',
+                'sword', 'syringe', 'Tabasco_sauce', 'table-tennis_table', 'table',
+                'table_lamp', 'tablecloth', 'tachometer', 'taco', 'tag', 'taillight',
+                'tambourine', 'army_tank', 'tank_(storage_vessel)',
+                'tank_top_(clothing)', 'tape_(sticky_cloth_or_paper)', 'tape_measure',
+                'tapestry', 'tarp', 'tartan', 'tassel', 'tea_bag', 'teacup',
+                'teakettle', 'teapot', 'teddy_bear', 'telephone', 'telephone_booth',
+                'telephone_pole', 'telephoto_lens', 'television_camera',
+                'television_set', 'tennis_ball', 'tennis_racket', 'tequila',
+                'thermometer', 'thermos_bottle', 'thermostat', 'thimble', 'thread',
+                'thumbtack', 'tiara', 'tiger', 'tights_(clothing)', 'timer',
+                'tinfoil', 'tinsel', 'tissue_paper', 'toast_(food)', 'toaster',
+                'toaster_oven', 'toilet', 'toilet_tissue', 'tomato', 'tongs',
+                'toolbox', 'toothbrush', 'toothpaste', 'toothpick', 'cover',
+                'tortilla', 'tow_truck', 'towel', 'towel_rack', 'toy',
+                'tractor_(farm_equipment)', 'traffic_light', 'dirt_bike',
+                'trailer_truck', 'train_(railroad_vehicle)', 'trampoline', 'tray',
+                'trench_coat', 'triangle_(musical_instrument)', 'tricycle', 'tripod',
+                'trousers', 'truck', 'truffle_(chocolate)', 'trunk', 'vat', 'turban',
+                'turkey_(food)', 'turnip', 'turtle', 'turtleneck_(clothing)',
+                'typewriter', 'umbrella', 'underwear', 'unicycle', 'urinal', 'urn',
+                'vacuum_cleaner', 'vase', 'vending_machine', 'vent', 'vest',
+                'videotape', 'vinegar', 'violin', 'vodka', 'volleyball', 'vulture',
+                'waffle', 'waffle_iron', 'wagon', 'wagon_wheel', 'walking_stick',
+                'wall_clock', 'wall_socket', 'wallet', 'walrus', 'wardrobe',
+                'washbasin', 'automatic_washer', 'watch', 'water_bottle',
+                'water_cooler', 'water_faucet', 'water_heater', 'water_jug',
+                'water_gun', 'water_scooter', 'water_ski', 'water_tower',
+                'watering_can', 'watermelon', 'weathervane', 'webcam', 'wedding_cake',
+                'wedding_ring', 'wet_suit', 'wheel', 'wheelchair', 'whipped_cream',
+                'whistle', 'wig', 'wind_chime', 'windmill', 'window_box_(for_plants)',
+                'windshield_wiper', 'windsock', 'wine_bottle', 'wine_bucket',
+                'wineglass', 'blinder_(for_horses)', 'wok', 'wolf', 'wooden_spoon',
+                'wreath', 'wrench', 'wristband', 'wristlet', 'yacht', 'yogurt',
+                'yoke_(animal_equipment)', 'zebra', 'zucchini')

ext/meta/sam_meta.py

@@ -0,0 +1,41 @@
+meta_dict = {
+    'vit_h': dict(
+        encoder_embed_dim=1280,
+        encoder_depth=32,
+        encoder_num_heads=16,
+        encoder_global_attn_indexes=[7, 15, 23, 31],
+        # common
+        prompt_embed_dim=256,
+        image_size=1024,
+        vit_patch_size=16,
+        image_embedding_size=64
+    ),
+    'vit_l': dict(
+        encoder_embed_dim=1024,
+        encoder_depth=24,
+        encoder_num_heads=16,
+        encoder_global_attn_indexes=[5, 11, 17, 23],
+        # common
+        prompt_embed_dim=256,
+        image_size=1024,
+        vit_patch_size=16,
+        image_embedding_size=64
+    ),
+    'vit_b': dict(
+        encoder_embed_dim=768,
+        encoder_depth=12,
+        encoder_num_heads=12,
+        encoder_global_attn_indexes=[2, 5, 8, 11],
+        # common
+        prompt_embed_dim=256,
+        image_size=1024,
+        vit_patch_size=16,
+        image_embedding_size=64
+    )
+}
+
+checkpoint_dict = {
+    'vit_h': 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth',
+    'vit_l': 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth',
+    'vit_b': 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth',
+}

ext/open_clip/__init__.py

@@ -0,0 +1,15 @@
+from .coca_model import CoCa
+from .constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
+from .factory import create_model, create_model_and_transforms, create_model_from_pretrained, get_tokenizer, create_loss
+from .factory import list_models, add_model_config, get_model_config, load_checkpoint
+from .loss import ClipLoss, DistillClipLoss, CoCaLoss
+from .model import CLIP, CustomTextCLIP, CLIPTextCfg, CLIPVisionCfg, \
+    convert_weights_to_lp, convert_weights_to_fp16, trace_model, get_cast_dtype, get_input_dtype
+from .openai import load_openai_model, list_openai_models
+from .pretrained import list_pretrained, list_pretrained_models_by_tag, list_pretrained_tags_by_model, \
+    get_pretrained_url, download_pretrained_from_url, is_pretrained_cfg, get_pretrained_cfg, download_pretrained
+from .push_to_hf_hub import push_pretrained_to_hf_hub, push_to_hf_hub
+from .tokenizer import SimpleTokenizer, tokenize, decode
+from .transform import image_transform, AugmentationCfg
+from .zero_shot_classifier import build_zero_shot_classifier, build_zero_shot_classifier_legacy
+from .zero_shot_metadata import OPENAI_IMAGENET_TEMPLATES, SIMPLE_IMAGENET_TEMPLATES, IMAGENET_CLASSNAMES

ext/open_clip/bpe_simple_vocab_16e6.txt.gz

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

ext/open_clip/coca_model.py

@@ -0,0 +1,458 @@
+from typing import Optional
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+import numpy as np
+from dataclasses import dataclass
+
+from .transformer import (
+    LayerNormFp32,
+    LayerNorm,
+    QuickGELU,
+    MultimodalTransformer,
+)
+from .model import CLIPTextCfg, CLIPVisionCfg, _build_vision_tower, _build_text_tower
+
+try:
+    from transformers import (
+        BeamSearchScorer,
+        LogitsProcessorList,
+        TopPLogitsWarper,
+        TopKLogitsWarper,
+        RepetitionPenaltyLogitsProcessor,
+        MinLengthLogitsProcessor,
+        MaxLengthCriteria,
+        StoppingCriteriaList
+    )
+
+    GENERATION_TYPES = {
+        "top_k": TopKLogitsWarper,
+        "top_p": TopPLogitsWarper,
+        "beam_search": "beam_search"
+    }
+    _has_transformers = True
+except ImportError as e:
+    GENERATION_TYPES = {
+        "top_k": None,
+        "top_p": None,
+        "beam_search": "beam_search"
+    }
+    _has_transformers = False
+
+
+@dataclass
+class MultimodalCfg(CLIPTextCfg):
+    mlp_ratio: int = 4
+    dim_head: int = 64
+    heads: int = 8
+    n_queries: int = 256
+    attn_pooler_heads: int = 8
+
+
+def _build_text_decoder_tower(
+        embed_dim,
+        multimodal_cfg,
+        quick_gelu: bool = False,
+        cast_dtype: Optional[torch.dtype] = None,
+):
+    multimodal_cfg = MultimodalCfg(**multimodal_cfg) if isinstance(multimodal_cfg, dict) else multimodal_cfg
+    act_layer = QuickGELU if quick_gelu else nn.GELU
+    norm_layer = (
+        LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm
+    )
+
+    decoder = MultimodalTransformer(
+        context_length=multimodal_cfg.context_length,
+        width=multimodal_cfg.width,
+        heads=multimodal_cfg.heads,
+        layers=multimodal_cfg.layers,
+        ls_init_value=multimodal_cfg.ls_init_value,
+        output_dim=embed_dim,
+        act_layer=act_layer,
+        norm_layer=norm_layer,
+    )
+
+    return decoder
+
+
+class CoCa(nn.Module):
+    def __init__(
+            self,
+            embed_dim,
+            multimodal_cfg: MultimodalCfg,
+            text_cfg: CLIPTextCfg,
+            vision_cfg: CLIPVisionCfg,
+            quick_gelu: bool = False,
+            cast_dtype: Optional[torch.dtype] = None,
+            pad_id: int = 0,
+    ):
+        super().__init__()
+        multimodal_cfg = MultimodalCfg(**multimodal_cfg) if isinstance(multimodal_cfg, dict) else multimodal_cfg
+        text_cfg = CLIPTextCfg(**text_cfg) if isinstance(text_cfg, dict) else text_cfg
+        vision_cfg = CLIPVisionCfg(**vision_cfg) if isinstance(vision_cfg, dict) else vision_cfg
+
+        self.text = _build_text_tower(
+            embed_dim=embed_dim,
+            text_cfg=text_cfg,
+            quick_gelu=quick_gelu,
+            cast_dtype=cast_dtype,
+        )
+
+        vocab_size = (
+            text_cfg.vocab_size  # for hf models
+            if hasattr(text_cfg, "hf_model_name") and text_cfg.hf_model_name is not None
+            else text_cfg.vocab_size
+        )
+
+        self.visual = _build_vision_tower(
+            embed_dim=embed_dim,
+            vision_cfg=vision_cfg,
+            quick_gelu=quick_gelu,
+            cast_dtype=cast_dtype,
+        )
+
+        self.text_decoder = _build_text_decoder_tower(
+            vocab_size,
+            multimodal_cfg=multimodal_cfg,
+            quick_gelu=quick_gelu,
+            cast_dtype=cast_dtype,
+        )
+
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+        self.pad_id = pad_id
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.visual.set_grad_checkpointing(enable)
+        self.text.set_grad_checkpointing(enable)
+        self.text_decoder.set_grad_checkpointing(enable)
+
+    def _encode_image(self, images, normalize=True):
+        image_latent, tokens_embs = self.visual(images)
+        image_latent = F.normalize(image_latent, dim=-1) if normalize else image_latent
+        return image_latent, tokens_embs
+
+    def _encode_text(self, text, normalize=True, embed_cls=True):
+        text = text[:, :-1] if embed_cls else text # make space for CLS token
+        text_latent, token_emb = self.text(text)
+        text_latent = F.normalize(text_latent, dim=-1) if normalize else text_latent
+        return text_latent, token_emb
+
+    def encode_image(self, images, normalize=True):
+        image_latent, _ = self._encode_image(images, normalize=normalize)
+        return image_latent
+
+    def encode_text(self, text, normalize=True, embed_cls=True):
+        text_latent, _ = self._encode_text(text, normalize=normalize, embed_cls=embed_cls)
+        return text_latent
+
+    def forward(self, image, text, embed_cls=True, image_latent=None, image_embs=None):
+        text_latent, token_embs = self._encode_text(text, embed_cls=embed_cls)
+        if image_latent is None or image_embs is None:
+            image_latent, image_embs = self._encode_image(image)
+
+        # TODO: add assertion to avoid bugs?
+        labels = text[:, -token_embs.shape[1]:]
+
+        logits = self.text_decoder(image_embs, token_embs)
+        return {
+            "image_features": image_latent,
+            "text_features": text_latent,
+            "logits": logits,
+            "labels": labels,
+            "logit_scale": self.logit_scale.exp()
+        }
+
+    def generate(
+        self,
+        image,
+        text=None,
+        seq_len=30,
+        max_seq_len=77,
+        temperature=1.,
+        generation_type="beam_search",
+        top_p=0.1,  # keep tokens in the 1 - top_p quantile
+        top_k=1,  # keeps the top_k most probable tokens
+        pad_token_id=None,
+        eos_token_id=None,
+        sot_token_id=None,
+        num_beams=6,
+        num_beam_groups=3,
+        min_seq_len=5,
+        stopping_criteria=None,
+        repetition_penalty=1.0,
+        fixed_output_length=False # if True output.shape == (batch_size, seq_len)
+    ):
+        # taking many ideas and components from HuggingFace GenerationMixin
+        # https://huggingface.co/docs/transformers/main/en/main_classes/text_generation
+        assert _has_transformers, "Please install transformers for generate functionality. `pip install transformers`."
+        assert seq_len > min_seq_len, "seq_len must be larger than min_seq_len"
+
+        with torch.no_grad():
+            sot_token_id = 49406 if sot_token_id is None else sot_token_id
+            eos_token_id = 49407 if eos_token_id is None else eos_token_id
+            pad_token_id = self.pad_id if pad_token_id is None else pad_token_id
+            logit_processor = LogitsProcessorList(
+                [
+                    MinLengthLogitsProcessor(min_seq_len, eos_token_id),
+                    RepetitionPenaltyLogitsProcessor(repetition_penalty),
+                ]
+            )
+
+            if stopping_criteria is None:
+                stopping_criteria = [MaxLengthCriteria(max_length=seq_len)]
+
+            stopping_criteria = StoppingCriteriaList(
+                stopping_criteria
+            )
+
+            device = image.device
+
+            if generation_type == "beam_search":
+                output = self._generate_beamsearch(
+                    image_inputs = image,
+                    pad_token_id=pad_token_id,
+                    eos_token_id=eos_token_id,
+                    sot_token_id=sot_token_id,
+                    num_beams=num_beams,
+                    num_beam_groups=num_beam_groups,
+                    min_seq_len=min_seq_len,
+                    stopping_criteria=stopping_criteria,
+                    logit_processor=logit_processor,
+                )
+                if fixed_output_length and output.shape[1] < seq_len:
+                    return torch.cat(
+                        (output, torch.ones(output.shape[0], seq_len-output.shape[1], device=device, dtype=output.dtype) * self.pad_id),
+                        dim=1
+                    )
+                return output
+
+            elif generation_type == "top_p":
+                logit_warper = GENERATION_TYPES[generation_type](top_p)
+            elif generation_type == "top_k":
+                logit_warper = GENERATION_TYPES[generation_type](top_k)
+            else:
+                raise ValueError(
+                    f"generation_type has to be one of "
+                    f"{'| ' + ' | '.join(list(GENERATION_TYPES.keys())) + ' |'}."
+                )
+
+            image_latent, image_embs = self._encode_image(image)
+
+            if text is None:
+                text = torch.ones((image.shape[0], 1), device=device, dtype=torch.long) * sot_token_id
+
+            was_training = self.training
+            num_dims = len(text.shape)
+
+            if num_dims == 1:
+                text = text[None, :]
+
+            cur_len = text.shape[1]
+            self.eval()
+            out = text
+
+            while True:
+                x = out[:, -max_seq_len:]
+                cur_len = x.shape[1]
+                logits = self(image, x, image_latent=image_latent, image_embs=image_embs, embed_cls=False)["logits"][:, -1]
+                mask = (out[:, -1] == eos_token_id) | (out[:, -1] == pad_token_id)
+                sample = torch.ones((out.shape[0], 1), device=device, dtype=torch.long) * pad_token_id
+
+                if mask.all():
+                    if not fixed_output_length:
+                        break
+                else:
+                    logits = logits[~mask, :]
+                    filtered_logits = logit_processor(x[~mask, :], logits)
+                    filtered_logits = logit_warper(x[~mask, :], filtered_logits)
+                    probs = F.softmax(filtered_logits / temperature, dim=-1)
+
+                    if (cur_len + 1 == seq_len):
+                        sample[~mask, :] = torch.ones((sum(~mask), 1), device=device, dtype=torch.long) * eos_token_id
+                    else:
+                        sample[~mask, :] = torch.multinomial(probs, 1)
+
+                out = torch.cat((out, sample), dim=-1)
+
+                cur_len += 1
+
+                if stopping_criteria(out, None):
+                    break
+
+            if num_dims == 1:
+                out = out.squeeze(0)
+
+            self.train(was_training)
+            return out
+
+    def _generate_beamsearch(
+            self,
+            image_inputs,
+            pad_token_id=None,
+            eos_token_id=None,
+            sot_token_id=None,
+            num_beams=6,
+            num_beam_groups=3,
+            min_seq_len=5,
+            stopping_criteria=None,
+            logit_processor=None,
+            logit_warper=None,
+    ):
+        device = image_inputs.device
+        batch_size = image_inputs.shape[0]
+        image_inputs = torch.repeat_interleave(image_inputs, num_beams, dim=0)
+        image_latent, image_embs = self._encode_image(image_inputs)
+
+        input_ids = torch.ones((batch_size * num_beams, 1), device=device, dtype=torch.long)
+        input_ids = input_ids * sot_token_id
+        beam_scorer = BeamSearchScorer(
+            batch_size=batch_size,
+            num_beams=num_beams,
+            device=device,
+            num_beam_groups=num_beam_groups,
+        )
+        # instantiate logits processors
+        logits_processor = (
+            LogitsProcessorList([MinLengthLogitsProcessor(min_seq_len, eos_token_id=eos_token_id)])
+            if logit_processor is None
+            else logit_processor
+        )
+
+        batch_size = len(beam_scorer._beam_hyps)
+        num_beams = beam_scorer.num_beams
+        num_beam_groups = beam_scorer.num_beam_groups
+        num_sub_beams = num_beams // num_beam_groups
+        batch_beam_size, cur_len = input_ids.shape
+        beam_indices = None
+
+        if num_beams * batch_size != batch_beam_size:
+            raise ValueError(
+                f"Batch dimension of `input_ids` should be {num_beams * batch_size}, but is {batch_beam_size}."
+            )
+
+        beam_scores = torch.full((batch_size, num_beams), -1e9, dtype=torch.float, device=device)
+        # initialise score of first beam of each group with 0 and the rest with 1e-9. This ensures that the beams in
+        # the same group don't produce same tokens everytime.
+        beam_scores[:, ::num_sub_beams] = 0
+        beam_scores = beam_scores.view((batch_size * num_beams,))
+
+        while True:
+
+            # predicted tokens in cur_len step
+            current_tokens = torch.zeros(batch_size * num_beams, dtype=input_ids.dtype, device=device)
+
+            # indices which will form the beams in the next time step
+            reordering_indices = torch.zeros(batch_size * num_beams, dtype=torch.long, device=device)
+
+            # do one decoder step on all beams of all sentences in batch
+            model_inputs = prepare_inputs_for_generation(input_ids=input_ids, image_inputs=image_inputs)
+            outputs = self(
+                model_inputs['images'],
+                model_inputs['text'],
+                embed_cls=False,
+                image_latent=image_latent,
+                image_embs=image_embs
+            )
+
+            for beam_group_idx in range(num_beam_groups):
+                group_start_idx = beam_group_idx * num_sub_beams
+                group_end_idx = min(group_start_idx + num_sub_beams, num_beams)
+                group_size = group_end_idx - group_start_idx
+
+                # indices of beams of current group among all sentences in batch
+                batch_group_indices = []
+
+                for batch_idx in range(batch_size):
+                    batch_group_indices.extend(
+                        [batch_idx * num_beams + idx for idx in range(group_start_idx, group_end_idx)]
+                    )
+                group_input_ids = input_ids[batch_group_indices]
+
+                # select outputs of beams of currentg group only
+                next_token_logits = outputs['logits'][batch_group_indices, -1, :]
+                vocab_size = next_token_logits.shape[-1]
+
+                next_token_scores_processed = logits_processor(
+                    group_input_ids, next_token_logits, current_tokens=current_tokens, beam_group_idx=beam_group_idx
+                )
+                next_token_scores = next_token_scores_processed + beam_scores[batch_group_indices].unsqueeze(-1)
+                next_token_scores = next_token_scores.expand_as(next_token_scores_processed)
+
+                # reshape for beam search
+                next_token_scores = next_token_scores.view(batch_size, group_size * vocab_size)
+
+                next_token_scores, next_tokens = torch.topk(
+                    next_token_scores, 2 * group_size, dim=1, largest=True, sorted=True
+                )
+
+                next_indices = torch.div(next_tokens, vocab_size, rounding_mode="floor")
+                next_tokens = next_tokens % vocab_size
+
+                # stateless
+                process_beam_indices = sum(beam_indices, ()) if beam_indices is not None else None
+                beam_outputs = beam_scorer.process(
+                    group_input_ids,
+                    next_token_scores,
+                    next_tokens,
+                    next_indices,
+                    pad_token_id=pad_token_id,
+                    eos_token_id=eos_token_id,
+                    beam_indices=process_beam_indices,
+                )
+                beam_scores[batch_group_indices] = beam_outputs["next_beam_scores"]
+                beam_next_tokens = beam_outputs["next_beam_tokens"]
+                beam_idx = beam_outputs["next_beam_indices"]
+
+                input_ids[batch_group_indices] = group_input_ids[beam_idx]
+                group_input_ids = torch.cat([group_input_ids[beam_idx, :], beam_next_tokens.unsqueeze(-1)], dim=-1)
+                current_tokens[batch_group_indices] = group_input_ids[:, -1]
+
+                # (beam_idx // group_size) -> batch_idx
+                # (beam_idx % group_size) -> offset of idx inside the group
+                reordering_indices[batch_group_indices] = (
+                    num_beams * torch.div(beam_idx, group_size, rounding_mode="floor") + group_start_idx + (beam_idx % group_size)
+                )
+
+            input_ids = torch.cat([input_ids, current_tokens.unsqueeze(-1)], dim=-1)
+
+            # increase cur_len
+            cur_len = cur_len + 1
+            if beam_scorer.is_done or stopping_criteria(input_ids, None):
+                break
+
+        final_beam_indices = sum(beam_indices, ()) if beam_indices is not None else None
+        sequence_outputs = beam_scorer.finalize(
+            input_ids,
+            beam_scores,
+            next_tokens,
+            next_indices,
+            pad_token_id=pad_token_id,
+            eos_token_id=eos_token_id,
+            max_length=stopping_criteria.max_length,
+            beam_indices=final_beam_indices,
+        )
+        return sequence_outputs['sequences']
+
+
+def prepare_inputs_for_generation(input_ids, image_inputs, past=None, **kwargs):
+    if past:
+        input_ids = input_ids[:, -1].unsqueeze(-1)
+
+    attention_mask = kwargs.get("attention_mask", None)
+    position_ids = kwargs.get("position_ids", None)
+
+    if attention_mask is not None and position_ids is None:
+        # create position_ids on the fly for batch generation
+        position_ids = attention_mask.long().cumsum(-1) - 1
+        position_ids.masked_fill_(attention_mask == 0, 1)
+    else:
+        position_ids = None
+    return {
+        "text": input_ids,
+        "images": image_inputs,
+        "past_key_values": past,
+        "position_ids": position_ids,
+        "attention_mask": attention_mask,
+    }

ext/open_clip/constants.py

@@ -0,0 +1,2 @@
+OPENAI_DATASET_MEAN = (0.48145466, 0.4578275, 0.40821073)
+OPENAI_DATASET_STD = (0.26862954, 0.26130258, 0.27577711)

ext/open_clip/factory.py

@@ -0,0 +1,387 @@
+import json
+import logging
+import os
+import pathlib
+import re
+from copy import deepcopy
+from pathlib import Path
+from typing import Any, Dict, Optional, Tuple, Union
+
+import torch
+
+from .constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
+from .model import CLIP, CustomTextCLIP, convert_weights_to_lp, convert_to_custom_text_state_dict,\
+    resize_pos_embed, get_cast_dtype
+from .coca_model import CoCa
+from .loss import ClipLoss, DistillClipLoss, CoCaLoss
+from .openai import load_openai_model
+from .pretrained import is_pretrained_cfg, get_pretrained_cfg, download_pretrained,\
+    list_pretrained_tags_by_model, download_pretrained_from_hf
+from .transform import image_transform, AugmentationCfg
+from .tokenizer import HFTokenizer, tokenize
+
+
+HF_HUB_PREFIX = 'hf-hub:'
+_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') 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 = {k: v for k, v in 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
+
+
+def get_tokenizer(model_name):
+    if model_name.startswith(HF_HUB_PREFIX):
+        tokenizer = HFTokenizer(model_name[len(HF_HUB_PREFIX):])
+    else:
+        config = get_model_config(model_name)
+        tokenizer = HFTokenizer(
+            config['text_cfg']['hf_tokenizer_name']) if 'hf_tokenizer_name' in config['text_cfg'] else tokenize
+    return tokenizer
+
+
+def load_state_dict(checkpoint_path: str, map_location='cpu'):
+    checkpoint = torch.load(checkpoint_path, map_location=map_location)
+    if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
+        state_dict = checkpoint['state_dict']
+    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()}
+    return state_dict
+
+
+def load_checkpoint(model, checkpoint_path, strict=True):
+    state_dict = load_state_dict(checkpoint_path)
+    # detect old format and make compatible with new format
+    if 'positional_embedding' in state_dict and not hasattr(model, 'positional_embedding'):
+        state_dict = convert_to_custom_text_state_dict(state_dict)
+    resize_pos_embed(state_dict, model)
+    incompatible_keys = model.load_state_dict(state_dict, strict=strict)
+    return incompatible_keys
+
+
+def create_model(
+        model_name: str,
+        pretrained: Optional[str] = None,
+        precision: str = 'fp32',
+        device: Union[str, torch.device] = 'cpu',
+        jit: bool = False,
+        force_quick_gelu: bool = False,
+        force_custom_text: bool = False,
+        force_patch_dropout: Optional[float] = None,
+        force_image_size: Optional[Union[int, Tuple[int, int]]] = None,
+        pretrained_image: bool = False,
+        pretrained_hf: bool = True,
+        cache_dir: Optional[str] = None,
+        output_dict: Optional[bool] = None,
+        require_pretrained: bool = False,
+        logger: logging.Logger = logging,
+):
+    has_hf_hub_prefix = model_name.startswith(HF_HUB_PREFIX)
+    if has_hf_hub_prefix:
+        model_id = model_name[len(HF_HUB_PREFIX):]
+        checkpoint_path = download_pretrained_from_hf(model_id, cache_dir=cache_dir)
+        config_path = download_pretrained_from_hf(model_id, filename='open_clip_config.json', cache_dir=cache_dir)
+
+        with open(config_path, 'r', encoding='utf-8') as f:
+            config = json.load(f)
+        pretrained_cfg = config['preprocess_cfg']
+        model_cfg = config['model_cfg']
+    else:
+        model_name = model_name.replace('/', '-')  # for callers using old naming with / in ViT names
+        checkpoint_path = None
+        pretrained_cfg = {}
+        model_cfg = None
+
+    if isinstance(device, str):
+        device = torch.device(device)
+
+    if pretrained and pretrained.lower() == 'openai':
+        logger.info(f'Loading pretrained {model_name} from OpenAI.')
+        model = load_openai_model(
+            model_name,
+            precision=precision,
+            device=device,
+            cache_dir=cache_dir,
+        )
+    else:
+        model_cfg = model_cfg or get_model_config(model_name)
+        if model_cfg is not None:
+            logger.info(f'Loaded {model_name} model config.')
+        else:
+            logger.error(f'Model config for {model_name} not found; available models {list_models()}.')
+            raise RuntimeError(f'Model config for {model_name} not found.')
+
+        if force_quick_gelu:
+            # override for use of QuickGELU on non-OpenAI transformer models
+            model_cfg["quick_gelu"] = True
+
+        if force_patch_dropout is not None:
+            # override the default patch dropout value
+            model_cfg["vision_cfg"]["patch_dropout"] = force_patch_dropout
+
+        if force_image_size is not None:
+            # override model config's image size
+            model_cfg["vision_cfg"]["image_size"] = force_image_size
+
+        is_timm_model = 'timm_model_name' in model_cfg.get('vision_cfg', {})
+        if pretrained_image:
+            if is_timm_model:
+                # pretrained weight loading for timm models set via vision_cfg
+                model_cfg['vision_cfg']['timm_model_pretrained'] = True
+            else:
+                assert False, 'pretrained image towers currently only supported for timm models'
+
+        # cast_dtype set for fp16 and bf16 (manual mixed-precision), not set for 'amp' or 'pure' modes
+        cast_dtype = get_cast_dtype(precision)
+        is_hf_model = 'hf_model_name' in model_cfg.get('text_cfg', {})
+        custom_text = model_cfg.pop('custom_text', False) or force_custom_text or is_hf_model
+
+        if custom_text:
+            if is_hf_model:
+                model_cfg['text_cfg']['hf_model_pretrained'] = pretrained_hf
+            if "coca" in model_name:
+                model = CoCa(**model_cfg, cast_dtype=cast_dtype)
+            else:
+                model = CustomTextCLIP(**model_cfg, cast_dtype=cast_dtype)
+        else:
+            model = CLIP(**model_cfg, cast_dtype=cast_dtype)
+
+        if precision in ("fp16", "bf16"):
+            dtype = torch.float16 if 'fp16' in precision else torch.bfloat16
+            # manual mixed precision that matches original OpenAI behaviour
+            if is_timm_model:
+                # FIXME this is a bit janky, create timm based model in low-precision and
+                # then cast only LayerNormFp32 instances back to float32 so they don't break.
+                # Why? The convert_weights_to_lp fn only works with native models.
+                model.to(device=device, dtype=dtype)
+                from .transformer import LayerNormFp32
+                def _convert_ln(m):
+                    if isinstance(m, LayerNormFp32):
+                        m.weight.data = m.weight.data.to(torch.float32)
+                        m.bias.data = m.bias.data.to(torch.float32)
+                model.apply(_convert_ln)
+            else:
+                model.to(device=device)
+                convert_weights_to_lp(model, dtype=dtype)
+        elif precision in ("pure_fp16", "pure_bf16"):
+            dtype = torch.float16 if 'fp16' in precision else torch.bfloat16
+            model.to(device=device, dtype=dtype)
+        else:
+            model.to(device=device)
+
+        pretrained_loaded = False
+        if pretrained:
+            checkpoint_path = ''
+            pretrained_cfg = get_pretrained_cfg(model_name, pretrained)
+            if pretrained_cfg:
+                checkpoint_path = download_pretrained(pretrained_cfg, cache_dir=cache_dir)
+            elif os.path.exists(pretrained):
+                checkpoint_path = pretrained
+
+            if checkpoint_path:
+                logger.info(f'Loading pretrained {model_name} weights ({pretrained}).')
+                load_checkpoint(model, checkpoint_path)
+            else:
+                error_str = (
+                    f'Pretrained weights ({pretrained}) not found for model {model_name}.'
+                    f'Available pretrained tags ({list_pretrained_tags_by_model(model_name)}.')
+                logger.warning(error_str)
+                raise RuntimeError(error_str)
+            pretrained_loaded = True
+        elif has_hf_hub_prefix:
+            logger.info(f'Loading pretrained {model_name} weights ({pretrained}).')
+            load_checkpoint(model, checkpoint_path)
+            pretrained_loaded = True
+
+        if require_pretrained and not pretrained_loaded:
+            # callers of create_model_from_pretrained always expect pretrained weights
+            raise RuntimeError(
+                f'Pretrained weights were required for (model: {model_name}, pretrained: {pretrained}) but not loaded.')
+
+        # set image / mean metadata from pretrained_cfg if available, or use default
+        model.visual.image_mean = pretrained_cfg.get('mean', None) or OPENAI_DATASET_MEAN
+        model.visual.image_std = pretrained_cfg.get('std', None) or OPENAI_DATASET_STD
+
+    if output_dict and hasattr(model, "output_dict"):
+        model.output_dict = True
+
+    if jit:
+        model = torch.jit.script(model)
+
+    return model
+
+
+def create_loss(args):
+    if args.distill:
+        return DistillClipLoss(
+            local_loss=args.local_loss,
+            gather_with_grad=args.gather_with_grad,
+            cache_labels=True,
+            rank=args.rank,
+            world_size=args.world_size,
+            use_horovod=args.horovod,
+        )
+    elif "coca" in args.model.lower():
+        return CoCaLoss(
+            caption_loss_weight=args.coca_caption_loss_weight,
+            clip_loss_weight=args.coca_contrastive_loss_weight,
+            local_loss=args.local_loss,
+            gather_with_grad=args.gather_with_grad,
+            cache_labels=True,
+            rank=args.rank,
+            world_size=args.world_size,
+            use_horovod=args.horovod,
+        )
+    return ClipLoss(
+        local_loss=args.local_loss,
+        gather_with_grad=args.gather_with_grad,
+        cache_labels=True,
+        rank=args.rank,
+        world_size=args.world_size,
+        use_horovod=args.horovod,
+    )
+
+
+def create_model_and_transforms(
+        model_name: str,
+        pretrained: Optional[str] = None,
+        precision: str = 'fp32',
+        device: Union[str, torch.device] = 'cpu',
+        jit: bool = False,
+        force_quick_gelu: bool = False,
+        force_custom_text: bool = False,
+        force_patch_dropout: Optional[float] = None,
+        force_image_size: Optional[Union[int, Tuple[int, int]]] = None,
+        pretrained_image: bool = False,
+        pretrained_hf: bool = True,
+        image_mean: Optional[Tuple[float, ...]] = None,
+        image_std: Optional[Tuple[float, ...]] = None,
+        aug_cfg: Optional[Union[Dict[str, Any], AugmentationCfg]] = None,
+        cache_dir: Optional[str] = None,
+        output_dict: Optional[bool] = None,
+        logger: logging.Logger = logging,
+):
+    model = create_model(
+        model_name,
+        pretrained,
+        precision=precision,
+        device=device,
+        jit=jit,
+        force_quick_gelu=force_quick_gelu,
+        force_custom_text=force_custom_text,
+        force_patch_dropout=force_patch_dropout,
+        force_image_size=force_image_size,
+        pretrained_image=pretrained_image,
+        pretrained_hf=pretrained_hf,
+        cache_dir=cache_dir,
+        output_dict=output_dict,
+        logger=logger,
+    )
+
+    image_mean = image_mean or getattr(model.visual, 'image_mean', None)
+    image_std = image_std or getattr(model.visual, 'image_std', None)
+    preprocess_train = image_transform(
+        model.visual.image_size,
+        is_train=True,
+        mean=image_mean,
+        std=image_std,
+        aug_cfg=aug_cfg,
+    )
+    preprocess_val = image_transform(
+        model.visual.image_size,
+        is_train=False,
+        mean=image_mean,
+        std=image_std,
+    )
+
+    return model, preprocess_train, preprocess_val
+
+
+def create_model_from_pretrained(
+        model_name: str,
+        pretrained: Optional[str] = None,
+        precision: str = 'fp32',
+        device: Union[str, torch.device] = 'cpu',
+        jit: bool = False,
+        force_quick_gelu: bool = False,
+        force_custom_text: bool = False,
+        force_image_size: Optional[Union[int, Tuple[int, int]]] = None,
+        return_transform: bool = True,
+        image_mean: Optional[Tuple[float, ...]] = None,
+        image_std: Optional[Tuple[float, ...]] = None,
+        cache_dir: Optional[str] = None,
+        logger: logging.Logger = logging,
+):
+    model = create_model(
+        model_name,
+        pretrained,
+        precision=precision,
+        device=device,
+        jit=jit,
+        force_quick_gelu=force_quick_gelu,
+        force_custom_text=force_custom_text,
+        force_image_size=force_image_size,
+        cache_dir=cache_dir,
+        require_pretrained=True,
+        logger=logger,
+    )
+
+    if not return_transform:
+        return model
+
+    image_mean = image_mean or getattr(model.visual, 'image_mean', None)
+    image_std = image_std or getattr(model.visual, 'image_std', None)
+    preprocess = image_transform(
+        model.visual.image_size,
+        is_train=False,
+        mean=image_mean,
+        std=image_std,
+    )
+
+    return model, preprocess

ext/open_clip/hf_configs.py

@@ -0,0 +1,56 @@
+# HF architecture dict:
+arch_dict = {
+    # https://huggingface.co/docs/transformers/model_doc/roberta#roberta
+    "roberta": {
+        "config_names": {
+            "context_length": "max_position_embeddings",
+            "vocab_size": "vocab_size",
+            "width": "hidden_size",
+            "heads": "num_attention_heads",
+            "layers": "num_hidden_layers",
+            "layer_attr": "layer",
+            "token_embeddings_attr": "embeddings"
+        },
+        "pooler": "mean_pooler",
+    },
+    # https://huggingface.co/docs/transformers/model_doc/xlm-roberta#transformers.XLMRobertaConfig
+    "xlm-roberta": {
+        "config_names": {
+            "context_length": "max_position_embeddings",
+            "vocab_size": "vocab_size",
+            "width": "hidden_size",
+            "heads": "num_attention_heads",
+            "layers": "num_hidden_layers",
+            "layer_attr": "layer",
+            "token_embeddings_attr": "embeddings"
+        },
+        "pooler": "mean_pooler",
+    },
+    # https://huggingface.co/docs/transformers/model_doc/mt5#mt5
+    "mt5": {
+        "config_names": {
+            # unlimited seqlen
+            # https://github.com/google-research/text-to-text-transfer-transformer/issues/273
+            # https://github.com/huggingface/transformers/blob/v4.24.0/src/transformers/models/t5/modeling_t5.py#L374
+            "context_length": "",
+            "vocab_size": "vocab_size",
+            "width": "d_model",
+            "heads": "num_heads",
+            "layers": "num_layers",
+            "layer_attr": "block",
+            "token_embeddings_attr": "embed_tokens"
+        },
+        "pooler": "mean_pooler",
+    },
+    # https://huggingface.co/docs/transformers/model_doc/bert
+    "bert": {
+        "config_names": {
+            "context_length": "max_position_embeddings",
+            "vocab_size": "vocab_size",
+            "width": "hidden_size",
+            "heads": "num_attention_heads",
+            "layers": "num_hidden_layers",
+        },
+        "pooler": "cls_pooler",
+    },
+}

ext/open_clip/hf_model.py

@@ -0,0 +1,193 @@
+""" huggingface model adapter
+
+Wraps HuggingFace transformers (https://github.com/huggingface/transformers) models for use as a text tower in CLIP model.
+"""
+import re
+
+import torch
+import torch.nn as nn
+from torch import TensorType
+
+try:
+    import transformers
+    from transformers import AutoModel, AutoTokenizer, AutoConfig, PretrainedConfig
+    from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, \
+        BaseModelOutputWithPoolingAndCrossAttentions
+except ImportError as e:
+    transformers = None
+
+
+    class BaseModelOutput:
+        pass
+
+
+    class PretrainedConfig:
+        pass
+
+from .hf_configs import arch_dict
+
+
+# utils
+def _camel2snake(s):
+    return re.sub(r'(?<!^)(?=[A-Z])', '_', s).lower()
+
+
+# TODO: ?last - for gpt-like models
+_POOLERS = {}
+
+
+def register_pooler(cls):
+    """Decorator registering pooler class"""
+    _POOLERS[_camel2snake(cls.__name__)] = cls
+    return cls
+
+
+@register_pooler
+class MeanPooler(nn.Module):
+    """Mean pooling"""
+
+    def forward(self, x: BaseModelOutput, attention_mask: TensorType):
+        masked_output = x.last_hidden_state * attention_mask.unsqueeze(-1)
+        return masked_output.sum(dim=1) / attention_mask.sum(-1, keepdim=True)
+
+
+@register_pooler
+class MaxPooler(nn.Module):
+    """Max pooling"""
+
+    def forward(self, x: BaseModelOutput, attention_mask: TensorType):
+        masked_output = x.last_hidden_state.masked_fill(attention_mask.unsqueeze(-1), -torch.inf)
+        return masked_output.max(1).values
+
+
+@register_pooler
+class ClsPooler(nn.Module):
+    """CLS token pooling"""
+
+    def __init__(self, use_pooler_output=True):
+        super().__init__()
+        self.cls_token_position = 0
+        self.use_pooler_output = use_pooler_output
+
+    def forward(self, x: BaseModelOutput, attention_mask: TensorType):
+        if (self.use_pooler_output and
+            isinstance(x, (BaseModelOutputWithPooling, BaseModelOutputWithPoolingAndCrossAttentions)) and
+            (x.pooler_output is not None)
+        ):
+            return x.pooler_output
+
+        return x.last_hidden_state[:, self.cls_token_position, :]
+
+
+@register_pooler
+class ClsLastHiddenStatePooler(nn.Module):
+    """CLS token pooling
+    NOTE: this is equivalent to ClsPooler above with use_pooler_output=False
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.cls_token_position = 0
+
+    def forward(self, x: BaseModelOutput, attention_mask: TensorType):
+        return x.last_hidden_state[:, self.cls_token_position, :]
+
+
+class HFTextEncoder(nn.Module):
+    """HuggingFace model adapter"""
+    output_tokens: torch.jit.Final[bool]
+
+    def __init__(
+            self,
+            model_name_or_path: str,
+            output_dim: int,
+            config: PretrainedConfig = None,
+            pooler_type: str = None,
+            proj: str = None,
+            pretrained: bool = True,
+            output_tokens: bool = False,
+    ):
+        super().__init__()
+        self.output_tokens = output_tokens
+        self.output_dim = output_dim
+
+        # TODO: find better way to get this information
+        uses_transformer_pooler = (pooler_type == "cls_pooler")
+
+        if transformers is None:
+            raise RuntimeError("Please `pip install transformers` to use pre-trained HuggingFace models")
+        if config is None:
+            self.config = AutoConfig.from_pretrained(model_name_or_path)
+            create_func, model_args = (AutoModel.from_pretrained, model_name_or_path) if pretrained else (
+                AutoModel.from_config, self.config)
+            # TODO: do all model configs have this attribute? PretrainedConfig does so yes??
+            if hasattr(self.config, "is_encoder_decoder") and self.config.is_encoder_decoder:
+                self.transformer = create_func(model_args)
+                self.transformer = self.transformer.encoder
+            else:
+                self.transformer = create_func(model_args, add_pooling_layer=uses_transformer_pooler)
+        else:
+            self.config = config
+            self.transformer = AutoModel.from_config(config)
+        if pooler_type is None:  # get default arch pooler
+            pooler_type = (arch_dict[self.config.model_type]["pooler"])
+
+        # FIXME downstream users of OpenCLIP models use these attr, need to verify valid across all models
+        self.vocab_size = getattr(self.config, 'vocab_size', 0)
+        self.context_length = getattr(self.config, 'max_position_embeddings', 0)
+
+        self.pooler = _POOLERS[pooler_type]()
+
+        d_model = getattr(self.config, arch_dict[self.config.model_type]["config_names"]["width"])
+        if (d_model == output_dim) and (proj is None):  # do we always need a proj?
+            self.proj = nn.Identity()
+        elif proj == 'linear':
+            self.proj = nn.Linear(d_model, output_dim, bias=False)
+        elif proj == 'mlp':
+            hidden_size = (d_model + output_dim) // 2
+            self.proj = nn.Sequential(
+                nn.Linear(d_model, hidden_size, bias=False),
+                nn.GELU(),
+                nn.Linear(hidden_size, output_dim, bias=False),
+            )
+
+    def forward(self, x: TensorType):
+        attn_mask = (x != self.config.pad_token_id).long()
+        out = self.transformer(input_ids=x, attention_mask=attn_mask)
+        pooled_out = self.pooler(out, attn_mask)
+        projected = self.proj(pooled_out)
+
+        seq_len = out.last_hidden_state.shape[1]
+        tokens = (
+            out.last_hidden_state[:, torch.arange(seq_len) != self.pooler.cls_token_position, :] 
+            if type(self.pooler) == ClsPooler 
+            else out.last_hidden_state
+        )
+        
+        if self.output_tokens:
+            return projected, tokens
+        return projected
+
+    def lock(self, unlocked_layers: int = 0, freeze_layer_norm: bool = True):
+        if not unlocked_layers:  # full freezing
+            for n, p in self.transformer.named_parameters():
+                p.requires_grad = (not freeze_layer_norm) if "LayerNorm" in n.split(".") else False
+            return
+
+        encoder = self.transformer.encoder if hasattr(self.transformer, 'encoder') else self.transformer
+        layer_list = getattr(encoder, arch_dict[self.config.model_type]["config_names"]["layer_attr"])
+        print(f"Unlocking {unlocked_layers}/{len(layer_list) + 1} layers of hf model")
+        embeddings = getattr(
+            self.transformer, arch_dict[self.config.model_type]["config_names"]["token_embeddings_attr"])
+        modules = [embeddings, *layer_list][:-unlocked_layers]
+        # freeze layers
+        for module in modules:
+            for n, p in module.named_parameters():
+                p.requires_grad = (not freeze_layer_norm) if "LayerNorm" in n.split(".") else False
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.transformer.gradient_checkpointing_enable()
+
+    def init_parameters(self):
+        pass

ext/open_clip/loss.py

@@ -0,0 +1,216 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+try:
+    import torch.distributed.nn
+    from torch import distributed as dist
+
+    has_distributed = True
+except ImportError:
+    has_distributed = False
+
+try:
+    import horovod.torch as hvd
+except ImportError:
+    hvd = None
+
+
+def gather_features(
+        image_features,
+        text_features,
+        local_loss=False,
+        gather_with_grad=False,
+        rank=0,
+        world_size=1,
+        use_horovod=False
+):
+    assert has_distributed, 'torch.distributed did not import correctly, please use a PyTorch version with support.'
+    if use_horovod:
+        assert hvd is not None, 'Please install horovod'
+        if gather_with_grad:
+            all_image_features = hvd.allgather(image_features)
+            all_text_features = hvd.allgather(text_features)
+        else:
+            with torch.no_grad():
+                all_image_features = hvd.allgather(image_features)
+                all_text_features = hvd.allgather(text_features)
+            if not local_loss:
+                # ensure grads for local rank when all_* features don't have a gradient
+                gathered_image_features = list(all_image_features.chunk(world_size, dim=0))
+                gathered_text_features = list(all_text_features.chunk(world_size, dim=0))
+                gathered_image_features[rank] = image_features
+                gathered_text_features[rank] = text_features
+                all_image_features = torch.cat(gathered_image_features, dim=0)
+                all_text_features = torch.cat(gathered_text_features, dim=0)
+    else:
+        # We gather tensors from all gpus
+        if gather_with_grad:
+            all_image_features = torch.cat(torch.distributed.nn.all_gather(image_features), dim=0)
+            all_text_features = torch.cat(torch.distributed.nn.all_gather(text_features), dim=0)
+        else:
+            gathered_image_features = [torch.zeros_like(image_features) for _ in range(world_size)]
+            gathered_text_features = [torch.zeros_like(text_features) for _ in range(world_size)]
+            dist.all_gather(gathered_image_features, image_features)
+            dist.all_gather(gathered_text_features, text_features)
+            if not local_loss:
+                # ensure grads for local rank when all_* features don't have a gradient
+                gathered_image_features[rank] = image_features
+                gathered_text_features[rank] = text_features
+            all_image_features = torch.cat(gathered_image_features, dim=0)
+            all_text_features = torch.cat(gathered_text_features, dim=0)
+
+    return all_image_features, all_text_features
+
+
+class ClipLoss(nn.Module):
+
+    def __init__(
+            self,
+            local_loss=False,
+            gather_with_grad=False,
+            cache_labels=False,
+            rank=0,
+            world_size=1,
+            use_horovod=False,
+    ):
+        super().__init__()
+        self.local_loss = local_loss
+        self.gather_with_grad = gather_with_grad
+        self.cache_labels = cache_labels
+        self.rank = rank
+        self.world_size = world_size
+        self.use_horovod = use_horovod
+
+        # cache state
+        self.prev_num_logits = 0
+        self.labels = {}
+
+    def get_ground_truth(self, device, num_logits) -> torch.Tensor:
+        # calculated ground-truth and cache if enabled
+        if self.prev_num_logits != num_logits or device not in self.labels:
+            labels = torch.arange(num_logits, device=device, dtype=torch.long)
+            if self.world_size > 1 and self.local_loss:
+                labels = labels + num_logits * self.rank
+            if self.cache_labels:
+                self.labels[device] = labels
+                self.prev_num_logits = num_logits
+        else:
+            labels = self.labels[device]
+        return labels
+
+    def get_logits(self, image_features, text_features, logit_scale):
+        if self.world_size > 1:
+            all_image_features, all_text_features = gather_features(
+                image_features, text_features,
+                self.local_loss, self.gather_with_grad, self.rank, self.world_size, self.use_horovod)
+
+            if self.local_loss:
+                logits_per_image = logit_scale * image_features @ all_text_features.T
+                logits_per_text = logit_scale * text_features @ all_image_features.T
+            else:
+                logits_per_image = logit_scale * all_image_features @ all_text_features.T
+                logits_per_text = logits_per_image.T
+        else:
+            logits_per_image = logit_scale * image_features @ text_features.T
+            logits_per_text = logit_scale * text_features @ image_features.T
+        
+        return logits_per_image, logits_per_text
+
+    def forward(self, image_features, text_features, logit_scale, output_dict=False):
+        device = image_features.device
+        logits_per_image, logits_per_text = self.get_logits(image_features, text_features, logit_scale)
+
+        labels = self.get_ground_truth(device, logits_per_image.shape[0])
+
+        total_loss = (
+            F.cross_entropy(logits_per_image, labels) +
+            F.cross_entropy(logits_per_text, labels)
+        ) / 2
+
+        return {"contrastive_loss": total_loss} if output_dict else total_loss
+
+
+class CoCaLoss(ClipLoss):
+    def __init__(
+            self,
+            caption_loss_weight,
+            clip_loss_weight,
+            pad_id=0,  # pad_token for open_clip custom tokenizer
+            local_loss=False,
+            gather_with_grad=False,
+            cache_labels=False,
+            rank=0,
+            world_size=1,
+            use_horovod=False,
+    ):
+        super().__init__(
+            local_loss=local_loss,
+            gather_with_grad=gather_with_grad,
+            cache_labels=cache_labels,
+            rank=rank,
+            world_size=world_size,
+            use_horovod=use_horovod
+        )
+
+        self.clip_loss_weight = clip_loss_weight
+        self.caption_loss_weight = caption_loss_weight
+        self.caption_loss = nn.CrossEntropyLoss(ignore_index=pad_id)
+
+    def forward(self, image_features, text_features, logits, labels, logit_scale, output_dict=False):
+        
+        clip_loss = torch.tensor(0)
+        
+        if self.clip_loss_weight:
+            clip_loss = super().forward(image_features, text_features, logit_scale)
+            clip_loss = self.clip_loss_weight * clip_loss
+
+        caption_loss = self.caption_loss(
+            logits.permute(0, 2, 1),
+            labels,
+        )
+        caption_loss = caption_loss * self.caption_loss_weight
+
+        if output_dict:
+            return {"contrastive_loss": clip_loss, "caption_loss": caption_loss}
+
+        return clip_loss, caption_loss
+
+
+class DistillClipLoss(ClipLoss):
+
+    def dist_loss(self, teacher_logits, student_logits):
+        return -(teacher_logits.softmax(dim=1) * student_logits.log_softmax(dim=1)).sum(dim=1).mean(dim=0)
+
+    def forward(
+            self,
+            image_features,
+            text_features,
+            logit_scale,
+            dist_image_features,
+            dist_text_features,
+            dist_logit_scale,
+            output_dict=False,
+    ):
+        logits_per_image, logits_per_text = \
+            self.get_logits(image_features, text_features, logit_scale)
+
+        dist_logits_per_image, dist_logits_per_text = \
+            self.get_logits(dist_image_features, dist_text_features, dist_logit_scale)
+
+        labels = self.get_ground_truth(image_features.device, logits_per_image.shape[0])
+
+        contrastive_loss = (
+            F.cross_entropy(logits_per_image, labels) +
+            F.cross_entropy(logits_per_text, labels)
+        ) / 2
+
+        distill_loss = (
+            self.dist_loss(dist_logits_per_image, logits_per_image) +
+            self.dist_loss(dist_logits_per_text, logits_per_text)
+        ) / 2
+
+        if output_dict:
+            return {"contrastive_loss": contrastive_loss, "distill_loss": distill_loss}
+
+        return contrastive_loss, distill_loss

ext/open_clip/model.py

@@ -0,0 +1,473 @@
+""" CLIP Model
+
+Adapted from https://github.com/openai/CLIP. Originally MIT License, Copyright (c) 2021 OpenAI.
+"""
+from dataclasses import dataclass
+import logging
+import math
+from typing import Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.utils.checkpoint import checkpoint
+
+from .hf_model import HFTextEncoder
+from .modified_resnet import ModifiedResNet
+from .timm_model import TimmModel
+from .transformer import LayerNormFp32, LayerNorm, QuickGELU, Attention, VisionTransformer, TextTransformer
+from .utils import to_2tuple
+
+
+@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.  # 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
+    input_patchnorm: bool = False  # whether to use dual patchnorm - would only apply the input layernorm on each patch, as post-layernorm already exist in original clip vit design
+    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)
+    attentional_pool: bool = False  # whether to use attentional pooler in the last embedding layer
+    n_queries: int = 256  # n_queries for attentional pooler
+    attn_pooler_heads: int = 8  # n heads for attentional_pooling
+    output_tokens: bool = False
+
+    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
+    timm_drop: float = 0.  # head dropout
+    timm_drop_path: Optional[float] = None  # backbone stochastic depth
+
+
+@dataclass
+class CLIPTextCfg:
+    context_length: int = 77
+    vocab_size: int = 49408
+    width: int = 512
+    heads: int = 8
+    layers: int = 12
+    ls_init_value: Optional[float] = None  # layer scale initial value
+    hf_model_name: str = None
+    hf_tokenizer_name: str = None
+    hf_model_pretrained: bool = True
+    proj: str = 'mlp'
+    pooler_type: str = 'mean_pooler'
+    embed_cls: bool = False
+    pad_id: int = 0
+    output_tokens: bool = False
+
+
+def get_cast_dtype(precision: str):
+    cast_dtype = None
+    if precision == 'bf16':
+        cast_dtype = torch.bfloat16
+    elif precision == 'fp16':
+        cast_dtype = torch.float16
+    return cast_dtype
+
+
+def get_input_dtype(precision: str):
+    input_dtype = None
+    if precision in ('bf16', 'pure_bf16'):
+        input_dtype = torch.bfloat16
+    elif precision in ('fp16', 'pure_fp16'):
+        input_dtype = torch.float16
+    return input_dtype
+
+
+def _build_vision_tower(
+        embed_dim: int,
+        vision_cfg: CLIPVisionCfg,
+        quick_gelu: bool = False,
+        cast_dtype: Optional[torch.dtype] = None
+):
+    if isinstance(vision_cfg, dict):
+        vision_cfg = CLIPVisionCfg(**vision_cfg)
+
+    # OpenAI models are pretrained w/ QuickGELU but native nn.GELU is both faster and more
+    # memory efficient in recent PyTorch releases (>= 1.10).
+    # NOTE: timm models always use native GELU regardless of quick_gelu flag.
+    act_layer = QuickGELU if quick_gelu else nn.GELU
+
+    if vision_cfg.timm_model_name:
+        visual = TimmModel(
+            vision_cfg.timm_model_name,
+            pretrained=vision_cfg.timm_model_pretrained,
+            pool=vision_cfg.timm_pool,
+            proj=vision_cfg.timm_proj,
+            proj_bias=vision_cfg.timm_proj_bias,
+            drop=vision_cfg.timm_drop,
+            drop_path=vision_cfg.timm_drop_path,
+            patch_drop=vision_cfg.patch_dropout if vision_cfg.patch_dropout > 0 else None,
+            embed_dim=embed_dim,
+            image_size=vision_cfg.image_size,
+        )
+    elif isinstance(vision_cfg.layers, (tuple, list)):
+        vision_heads = vision_cfg.width * 32 // vision_cfg.head_width
+        visual = ModifiedResNet(
+            layers=vision_cfg.layers,
+            output_dim=embed_dim,
+            heads=vision_heads,
+            image_size=vision_cfg.image_size,
+            width=vision_cfg.width,
+        )
+    else:
+        vision_heads = vision_cfg.width // vision_cfg.head_width
+        norm_layer = LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm
+        visual = VisionTransformer(
+            image_size=vision_cfg.image_size,
+            patch_size=vision_cfg.patch_size,
+            width=vision_cfg.width,
+            layers=vision_cfg.layers,
+            heads=vision_heads,
+            mlp_ratio=vision_cfg.mlp_ratio,
+            ls_init_value=vision_cfg.ls_init_value,
+            patch_dropout=vision_cfg.patch_dropout,
+            input_patchnorm=vision_cfg.input_patchnorm,
+            global_average_pool=vision_cfg.global_average_pool,
+            attentional_pool=vision_cfg.attentional_pool,
+            n_queries=vision_cfg.n_queries,
+            attn_pooler_heads=vision_cfg.attn_pooler_heads,
+            output_tokens=vision_cfg.output_tokens,
+            output_dim=embed_dim,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+        )
+
+    return visual
+
+
+def _build_text_tower(
+        embed_dim: int,
+        text_cfg: CLIPTextCfg,
+        quick_gelu: bool = False,
+        cast_dtype: Optional[torch.dtype] = None,
+):
+    if isinstance(text_cfg, dict):
+        text_cfg = CLIPTextCfg(**text_cfg)
+
+    if text_cfg.hf_model_name:
+        text = HFTextEncoder(
+            text_cfg.hf_model_name,
+            output_dim=embed_dim,
+            proj=text_cfg.proj,
+            pooler_type=text_cfg.pooler_type,
+            pretrained=text_cfg.hf_model_pretrained,
+            output_tokens=text_cfg.output_tokens,
+        )
+    else:
+        act_layer = QuickGELU if quick_gelu else nn.GELU
+        norm_layer = LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm
+
+        text = TextTransformer(
+            context_length=text_cfg.context_length,
+            vocab_size=text_cfg.vocab_size,
+            width=text_cfg.width,
+            heads=text_cfg.heads,
+            layers=text_cfg.layers,
+            ls_init_value=text_cfg.ls_init_value,
+            output_dim=embed_dim,
+            embed_cls=text_cfg.embed_cls,
+            output_tokens=text_cfg.output_tokens,
+            pad_id=text_cfg.pad_id,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+        )
+    return text
+
+
+class CLIP(nn.Module):
+    output_dict: torch.jit.Final[bool]
+
+    def __init__(
+            self,
+            embed_dim: int,
+            vision_cfg: CLIPVisionCfg,
+            text_cfg: CLIPTextCfg,
+            quick_gelu: bool = False,
+            cast_dtype: Optional[torch.dtype] = None,
+            output_dict: bool = False,
+    ):
+        super().__init__()
+        self.output_dict = output_dict
+        self.visual = _build_vision_tower(embed_dim, vision_cfg, quick_gelu, cast_dtype)
+
+        text = _build_text_tower(embed_dim, text_cfg, quick_gelu, cast_dtype)
+        self.transformer = text.transformer
+        self.context_length = text.context_length
+        self.vocab_size = text.vocab_size
+        self.token_embedding = text.token_embedding
+        self.positional_embedding = text.positional_embedding
+        self.ln_final = text.ln_final
+        self.text_projection = text.text_projection
+        self.register_buffer('attn_mask', text.attn_mask, persistent=False)
+
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+    def lock_image_tower(self, unlocked_groups=0, freeze_bn_stats=False):
+        # lock image tower as per LiT - https://arxiv.org/abs/2111.07991
+        self.visual.lock(unlocked_groups=unlocked_groups, freeze_bn_stats=freeze_bn_stats)
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.visual.set_grad_checkpointing(enable)
+        self.transformer.grad_checkpointing = enable
+
+    def encode_image(self, image, normalize: bool = False):
+        features = self.visual(image)
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def encode_text(self, text, normalize: bool = False):
+        cast_dtype = self.transformer.get_cast_dtype()
+
+        x = self.token_embedding(text).to(cast_dtype)  # [batch_size, n_ctx, d_model]
+
+        x = x + self.positional_embedding.to(cast_dtype)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x, attn_mask=self.attn_mask)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.ln_final(x)  # [batch_size, n_ctx, transformer.width]
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
+        return F.normalize(x, dim=-1) if normalize else x
+
+    def forward(
+            self,
+            image: Optional[torch.Tensor] = None,
+            text: Optional[torch.Tensor] = None,
+    ):
+        image_features = self.encode_image(image, normalize=True) if image is not None else None
+        text_features = self.encode_text(text, normalize=True) if text is not None else None
+        if self.output_dict:
+            return {
+                "image_features": image_features,
+                "text_features": text_features,
+                "logit_scale": self.logit_scale.exp()
+            }
+        return image_features, text_features, self.logit_scale.exp()
+
+
+class CustomTextCLIP(nn.Module):
+    output_dict: torch.jit.Final[bool]
+
+    def __init__(
+            self,
+            embed_dim: int,
+            vision_cfg: CLIPVisionCfg,
+            text_cfg: CLIPTextCfg,
+            quick_gelu: bool = False,
+            cast_dtype: Optional[torch.dtype] = None,
+            output_dict: bool = False,
+    ):
+        super().__init__()
+        self.output_dict = output_dict
+        self.visual = _build_vision_tower(embed_dim, vision_cfg, quick_gelu, cast_dtype)
+        self.text = _build_text_tower(embed_dim, text_cfg, quick_gelu, cast_dtype)
+        self.context_length = self.text.context_length
+        self.vocab_size = self.text.vocab_size
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+    def lock_image_tower(self, unlocked_groups=0, freeze_bn_stats=False):
+        # lock image tower as per LiT - https://arxiv.org/abs/2111.07991
+        self.visual.lock(unlocked_groups=unlocked_groups, freeze_bn_stats=freeze_bn_stats)
+
+    def lock_text_tower(self, unlocked_layers: int = 0, freeze_layer_norm: bool = True):
+        self.text.lock(unlocked_layers, freeze_layer_norm)
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.visual.set_grad_checkpointing(enable)
+        self.text.set_grad_checkpointing(enable)
+
+    def encode_image(self, image, normalize: bool = False):
+        features = self.visual(image)
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def encode_text(self, text, normalize: bool = False):
+        features = self.text(text)
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def forward(
+            self,
+            image: Optional[torch.Tensor] = None,
+            text: Optional[torch.Tensor] = None,
+    ):
+        image_features = self.encode_image(image, normalize=True) if image is not None else None
+        text_features = self.encode_text(text, normalize=True) if text is not None else None
+        if self.output_dict:
+            return {
+                "image_features": image_features,
+                "text_features": text_features,
+                "logit_scale": self.logit_scale.exp()
+            }
+        return image_features, text_features, self.logit_scale.exp()
+
+
+def convert_weights_to_lp(model: nn.Module, dtype=torch.float16):
+    """Convert applicable model parameters to low-precision (bf16 or fp16)"""
+
+    def _convert_weights(l):
+        if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
+            l.weight.data = l.weight.data.to(dtype)
+            if l.bias is not None:
+                l.bias.data = l.bias.data.to(dtype)
+
+        if isinstance(l, (nn.MultiheadAttention, Attention)):
+            for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]:
+                tensor = getattr(l, attr)
+                if tensor is not None:
+                    tensor.data = tensor.data.to(dtype)
+
+        if isinstance(l, (CLIP, TextTransformer)):
+            # convert text nn.Parameter projections
+            attr = getattr(l, "text_projection", None)
+            if attr is not None:
+                attr.data = attr.data.to(dtype)
+
+        if isinstance(l, VisionTransformer):
+            # convert vision nn.Parameter projections
+            attr = getattr(l, "proj", None)
+            if attr is not None:
+                attr.data = attr.data.to(dtype)
+
+    model.apply(_convert_weights)
+
+
+convert_weights_to_fp16 = convert_weights_to_lp  # backwards compat
+
+
+# used to maintain checkpoint compatibility
+def convert_to_custom_text_state_dict(state_dict: dict):
+    if 'text_projection' in state_dict:
+        # old format state_dict, move text tower -> .text
+        new_state_dict = {}
+        for k, v in state_dict.items():
+            if any(k.startswith(p) for p in (
+                'text_projection',
+                'positional_embedding',
+                'token_embedding',
+                'transformer',
+                'ln_final',
+            )):
+                k = 'text.' + k
+            new_state_dict[k] = v
+        return new_state_dict
+    return state_dict
+
+
+def build_model_from_openai_state_dict(
+        state_dict: dict,
+        quick_gelu=True,
+        cast_dtype=torch.float16,
+):
+    vit = "visual.proj" in state_dict
+
+    if vit:
+        vision_width = state_dict["visual.conv1.weight"].shape[0]
+        vision_layers = len(
+            [k for k in state_dict.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")])
+        vision_patch_size = state_dict["visual.conv1.weight"].shape[-1]
+        grid_size = round((state_dict["visual.positional_embedding"].shape[0] - 1) ** 0.5)
+        image_size = vision_patch_size * grid_size
+    else:
+        counts: list = [
+            len(set(k.split(".")[2] for k in state_dict if k.startswith(f"visual.layer{b}"))) for b in [1, 2, 3, 4]]
+        vision_layers = tuple(counts)
+        vision_width = state_dict["visual.layer1.0.conv1.weight"].shape[0]
+        output_width = round((state_dict["visual.attnpool.positional_embedding"].shape[0] - 1) ** 0.5)
+        vision_patch_size = None
+        assert output_width ** 2 + 1 == state_dict["visual.attnpool.positional_embedding"].shape[0]
+        image_size = output_width * 32
+
+    embed_dim = state_dict["text_projection"].shape[1]
+    context_length = state_dict["positional_embedding"].shape[0]
+    vocab_size = state_dict["token_embedding.weight"].shape[0]
+    transformer_width = state_dict["ln_final.weight"].shape[0]
+    transformer_heads = transformer_width // 64
+    transformer_layers = len(set(k.split(".")[2] for k in state_dict if k.startswith(f"transformer.resblocks")))
+
+    vision_cfg = CLIPVisionCfg(
+        layers=vision_layers,
+        width=vision_width,
+        patch_size=vision_patch_size,
+        image_size=image_size,
+    )
+    text_cfg = CLIPTextCfg(
+        context_length=context_length,
+        vocab_size=vocab_size,
+        width=transformer_width,
+        heads=transformer_heads,
+        layers=transformer_layers,
+    )
+    model = CLIP(
+        embed_dim,
+        vision_cfg=vision_cfg,
+        text_cfg=text_cfg,
+        quick_gelu=quick_gelu,  # OpenAI models were trained with QuickGELU
+        cast_dtype=cast_dtype,
+    )
+
+    for key in ["input_resolution", "context_length", "vocab_size"]:
+        state_dict.pop(key, None)
+
+    convert_weights_to_fp16(model)  # OpenAI state dicts are partially converted to float16
+    model.load_state_dict(state_dict)
+    return model.eval()
+
+
+def trace_model(model, batch_size=256, device=torch.device('cpu')):
+    model.eval()
+    image_size = model.visual.image_size
+    example_images = torch.ones((batch_size, 3, image_size, image_size), device=device)
+    example_text = torch.zeros((batch_size, model.context_length), dtype=torch.int, device=device)
+    model = torch.jit.trace_module(
+        model,
+        inputs=dict(
+            forward=(example_images, example_text),
+            encode_text=(example_text,),
+            encode_image=(example_images,)
+        ))
+    model.visual.image_size = image_size
+    return model
+
+
+def resize_pos_embed(state_dict, model, interpolation: str = 'bicubic', antialias: bool = True):
+    # Rescale the grid of position embeddings when loading from state_dict
+    old_pos_embed = state_dict.get('visual.positional_embedding', None)
+    if old_pos_embed is None or not hasattr(model.visual, 'grid_size'):
+        return
+    grid_size = to_2tuple(model.visual.grid_size)
+    extra_tokens = 1  # FIXME detect different token configs (ie no class token, or more)
+    new_seq_len = grid_size[0] * grid_size[1] + extra_tokens
+    if new_seq_len == old_pos_embed.shape[0]:
+        return
+
+    if extra_tokens:
+        pos_emb_tok, pos_emb_img = old_pos_embed[:extra_tokens], old_pos_embed[extra_tokens:]
+    else:
+        pos_emb_tok, pos_emb_img = None, old_pos_embed
+    old_grid_size = to_2tuple(int(math.sqrt(len(pos_emb_img))))
+
+    logging.info('Resizing position embedding grid-size from %s to %s', old_grid_size, grid_size)
+    pos_emb_img = pos_emb_img.reshape(1, old_grid_size[0], old_grid_size[1], -1).permute(0, 3, 1, 2)
+    pos_emb_img = F.interpolate(
+        pos_emb_img,
+        size=grid_size,
+        mode=interpolation,
+        antialias=antialias,
+        align_corners=False,
+    )
+    pos_emb_img = pos_emb_img.permute(0, 2, 3, 1).reshape(1, grid_size[0] * grid_size[1], -1)[0]
+    if pos_emb_tok is not None:
+        new_pos_embed = torch.cat([pos_emb_tok, pos_emb_img], dim=0)
+    else:
+        new_pos_embed = pos_emb_img
+    state_dict['visual.positional_embedding'] = new_pos_embed

ext/open_clip/model_configs/EVA01-g-14-plus.json

@@ -0,0 +1,18 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 224,
+        "timm_model_name": "eva_giant_patch14_224",
+        "timm_model_pretrained": false,
+        "timm_pool": "token",
+        "timm_proj": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 1024,
+        "heads": 16,
+        "layers": 24
+    },
+    "custom_text": true
+}

ext/open_clip/model_configs/EVA01-g-14.json

@@ -0,0 +1,18 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 224,
+        "timm_model_name": "eva_giant_patch14_224",
+        "timm_model_pretrained": false,
+        "timm_pool": "token",
+        "timm_proj": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 768,
+        "heads": 12,
+        "layers": 12
+    },
+    "custom_text": true
+}

ext/open_clip/model_configs/EVA02-B-16.json

@@ -0,0 +1,18 @@
+{
+    "embed_dim": 512,
+    "vision_cfg": {
+        "image_size": 224,
+        "timm_model_name": "eva02_base_patch16_clip_224",
+        "timm_model_pretrained": false,
+        "timm_pool": "token",
+        "timm_proj": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    },
+    "custom_text": true
+}

ext/open_clip/model_configs/EVA02-E-14-plus.json

@@ -0,0 +1,18 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 224,
+        "timm_model_name": "eva02_enormous_patch14_clip_224",
+        "timm_model_pretrained": false,
+        "timm_pool": "token",
+        "timm_proj": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 1280,
+        "heads": 20,
+        "layers": 32
+    },
+    "custom_text": true
+}

ext/open_clip/model_configs/EVA02-E-14.json

@@ -0,0 +1,18 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 224,
+        "timm_model_name": "eva02_enormous_patch14_clip_224",
+        "timm_model_pretrained": false,
+        "timm_pool": "token",
+        "timm_proj": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 1024,
+        "heads": 16,
+        "layers": 24
+    },
+    "custom_text": true
+}

ext/open_clip/model_configs/EVA02-L-14-336.json

@@ -0,0 +1,18 @@
+{
+    "embed_dim": 768,
+    "vision_cfg": {
+        "image_size": 336,
+        "timm_model_name": "eva02_large_patch14_clip_336",
+        "timm_model_pretrained": false,
+        "timm_pool": "token",
+        "timm_proj": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 768,
+        "heads": 12,
+        "layers": 12
+    },
+    "custom_text": true
+}

ext/open_clip/model_configs/EVA02-L-14.json

@@ -0,0 +1,18 @@
+{
+    "embed_dim": 768,
+    "vision_cfg": {
+        "image_size": 224,
+        "timm_model_name": "eva02_large_patch14_clip_224",
+        "timm_model_pretrained": false,
+        "timm_pool": "token",
+        "timm_proj": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 768,
+        "heads": 12,
+        "layers": 12
+    },
+    "custom_text": true
+}

ext/open_clip/model_configs/RN101-quickgelu.json

@@ -0,0 +1,22 @@
+{
+    "embed_dim": 512,
+    "quick_gelu": true,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": [
+            3,
+            4,
+            23,
+            3
+        ],
+        "width": 64,
+        "patch_size": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}

ext/open_clip/model_configs/RN101.json

@@ -0,0 +1,21 @@
+{
+    "embed_dim": 512,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": [
+            3,
+            4,
+            23,
+            3
+        ],
+        "width": 64,
+        "patch_size": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}

ext/open_clip/model_configs/RN50-quickgelu.json

@@ -0,0 +1,22 @@
+{
+    "embed_dim": 1024,
+    "quick_gelu": true,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": [
+            3,
+            4,
+            6,
+            3
+        ],
+        "width": 64,
+        "patch_size": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}

ext/open_clip/model_configs/RN50.json

@@ -0,0 +1,21 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": [
+            3,
+            4,
+            6,
+            3
+        ],
+        "width": 64,
+        "patch_size": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}

ext/open_clip/model_configs/RN50x16.json

@@ -0,0 +1,21 @@
+{
+    "embed_dim": 768,
+    "vision_cfg": {
+        "image_size": 384,
+        "layers": [
+            6,
+            8,
+            18,
+            8
+        ],
+        "width": 96,
+        "patch_size": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 768,
+        "heads": 12,
+        "layers": 12
+    }
+}

ext/open_clip/model_configs/RN50x4.json

@@ -0,0 +1,21 @@
+{
+    "embed_dim": 640,
+    "vision_cfg": {
+        "image_size": 288,
+        "layers": [
+            4,
+            6,
+            10,
+            6
+        ],
+        "width": 80,
+        "patch_size": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 640,
+        "heads": 10,
+        "layers": 12
+    }
+}