add rap_sam

.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

@@ -4,10 +4,8 @@ emoji: 🌍
 colorFrom: purple
 colorTo: purple
 sdk: gradio
-sdk_version: 4.13.0
-app_file: app.py
+sdk_version: 4.7.1
+app_file: main.py
 pinned: false
-license: mit
+python_version: 3.10
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app/configs/rap_sam_r50_12e_adaptor.py

@@ -0,0 +1,88 @@
+from mmdet.models import ResNet, MaskFormerFusionHead, CrossEntropyLoss, DiceLoss
+
+from app.models.detectors import YOSOVideoSam
+from app.models.heads import RapSAMVideoHead
+from app.models.necks import YOSONeck
+
+num_things_classes = 80
+num_stuff_classes = 53
+ov_model_name = 'convnext_large_d_320'
+ov_datasets_name = 'CocoPanopticOVDataset'
+num_classes = num_things_classes + num_stuff_classes
+model = dict(
+    type=YOSOVideoSam,
+    data_preprocessor=None,
+    backbone=dict(
+        type=ResNet,
+        depth=50,
+        num_stages=4,
+        out_indices=(0, 1, 2, 3),
+        frozen_stages=-1,
+        norm_cfg=dict(type='BN', requires_grad=True),
+        norm_eval=True,
+        init_cfg=None,
+    ),
+    neck=dict(
+        type=YOSONeck,
+        agg_dim=128,
+        hidden_dim=256,
+        backbone_shape=[256, 512, 1024, 2048],
+    ),
+    panoptic_head=dict(
+        type=RapSAMVideoHead,
+        prompt_with_kernel_updator=False,
+        panoptic_with_kernel_updator=True,
+        use_adaptor=True,
+        use_kernel_updator=True,
+        sphere_cls=True,
+        ov_classifier_name=f'{ov_model_name}_{ov_datasets_name}',
+        num_stages=3,
+        feat_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=100,
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * num_classes + [0.1]),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0)
+    ),
+    panoptic_fusion_head=dict(
+        type=MaskFormerFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None
+    ),
+    train_cfg=None,
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=dict(
+        type='Pretrained',
+        checkpoint='models/rapsam_r50_12e.pth'
+    )
+)

app/models/detectors/__init__.py

@@ -0,0 +1 @@
+from .rapsam import YOSOVideoSam

app/models/detectors/mask2former_vid.py

@@ -0,0 +1,281 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from typing import Dict, List, Tuple
+
+import torch
+from mmengine.structures import InstanceData
+from torch import Tensor
+import torch.nn.functional as F
+
+from mmdet.registry import MODELS
+from mmdet.structures import SampleList, OptSampleList, TrackDataSample
+from mmdet.utils import ConfigType, OptConfigType, OptMultiConfig
+from mmdet.models.detectors.single_stage import SingleStageDetector
+
+from app.models.utils import mask_pool
+
+
+@MODELS.register_module()
+class Mask2formerVideo(SingleStageDetector):
+    r"""Implementation of `Per-Pixel Classification is
+    NOT All You Need for Semantic Segmentation
+    <https://arxiv.org/pdf/2107.06278>`_."""
+    OVERLAPPING = None
+
+    def __init__(self,
+                 backbone: ConfigType,
+                 neck: OptConfigType = None,
+                 panoptic_head: OptConfigType = None,
+                 panoptic_fusion_head: OptConfigType = None,
+                 train_cfg: OptConfigType = None,
+                 test_cfg: OptConfigType = None,
+                 data_preprocessor: OptConfigType = None,
+                 inference_sam: bool = False,
+                 init_cfg: OptMultiConfig = None
+                 ):
+        super(SingleStageDetector, self).__init__(
+            data_preprocessor=data_preprocessor, init_cfg=init_cfg)
+        self.backbone = MODELS.build(backbone)
+        if neck is not None:
+            self.neck = MODELS.build(neck)
+
+        panoptic_head_ = panoptic_head.deepcopy()
+        panoptic_head_.update(train_cfg=train_cfg)
+        panoptic_head_.update(test_cfg=test_cfg)
+        self.panoptic_head = MODELS.build(panoptic_head_)
+
+        panoptic_fusion_head_ = panoptic_fusion_head.deepcopy()
+        panoptic_fusion_head_.update(test_cfg=test_cfg)
+        self.panoptic_fusion_head = MODELS.build(panoptic_fusion_head_)
+
+        self.num_things_classes = self.panoptic_head.num_things_classes
+        self.num_stuff_classes = self.panoptic_head.num_stuff_classes
+        self.num_classes = self.panoptic_head.num_classes
+
+        self.train_cfg = train_cfg
+        self.test_cfg = test_cfg
+
+        self.alpha = 0.4
+        self.beta = 0.8
+
+        self.inference_sam = inference_sam
+
+    def predict(self,
+                batch_inputs: Tensor,
+                batch_data_samples: SampleList,
+                rescale: bool = True) -> SampleList:
+        """Predict results from a batch of inputs and data samples with post-
+        processing.
+
+        Args:
+            batch_inputs (Tensor): Inputs with shape (N, C, H, W).
+            batch_data_samples (List[:obj:`DetDataSample`]): The Data
+                Samples. It usually includes information such as
+                `gt_instance`, `gt_panoptic_seg` and `gt_sem_seg`.
+            rescale (bool): Whether to rescale the results.
+                Defaults to True.
+
+        Returns:
+            list[:obj:`DetDataSample`]: Detection results of the
+            input images. Each DetDataSample usually contain
+            'pred_instances' and `pred_panoptic_seg`. And the
+            ``pred_instances`` usually contains following keys.
+
+                - scores (Tensor): Classification scores, has a shape
+                    (num_instance, )
+                - labels (Tensor): Labels of bboxes, has a shape
+                    (num_instances, ).
+                - bboxes (Tensor): Has a shape (num_instances, 4),
+                    the last dimension 4 arrange as (x1, y1, x2, y2).
+                - masks (Tensor): Has a shape (num_instances, H, W).
+
+            And the ``pred_panoptic_seg`` contains the following key
+
+                - sem_seg (Tensor): panoptic segmentation mask, has a
+                    shape (1, h, w).
+        """
+        if isinstance(batch_data_samples[0], TrackDataSample):
+            bs, num_frames, three, h, w = batch_inputs.shape
+            assert three == 3, "Only supporting images with 3 channels."
+            x = batch_inputs.reshape((bs * num_frames, three, h, w))
+            feats = self.extract_feat(x)
+        else:
+            num_frames = 0
+            bs = batch_inputs.shape[0]
+            feats = self.extract_feat(batch_inputs)
+
+        mask_cls_results, mask_pred_results, iou_results = self.panoptic_head.predict(feats, batch_data_samples)
+
+        if self.inference_sam:
+            for i, data_sample in enumerate(batch_data_samples):
+                meta = data_sample.metainfo
+                img_height, img_width = meta['img_shape'][:2]
+                mask_pred_result = mask_pred_results[i][:, :img_height, :img_width]
+                mask_pred_result = mask_pred_result.view(-1, img_height, img_width) > 0
+                all_pred_instances = InstanceData(masks=mask_pred_result)
+                batch_data_samples[i].pred_instances = all_pred_instances
+
+            return batch_data_samples
+
+        if self.OVERLAPPING is not None:
+            assert len(self.OVERLAPPING) == self.num_classes
+            mask_cls_results = self.open_voc_inference(feats, mask_cls_results, mask_pred_results)
+
+        if num_frames > 0:
+            for frame_id in range(num_frames):
+                results_list_img = self.panoptic_fusion_head.predict(
+                    mask_cls_results,
+                    mask_pred_results[:, :, frame_id],
+                    [batch_data_samples[idx][frame_id] for idx in range(bs)],
+                    rescale=rescale
+                )
+                _ = self.add_track_pred_to_datasample(
+                    [batch_data_samples[idx][frame_id] for idx in range(bs)], results_list_img
+                )
+            results = batch_data_samples
+        else:
+            results_list = self.panoptic_fusion_head.predict(
+                mask_cls_results,
+                mask_pred_results,
+                batch_data_samples,
+                iou_results=iou_results,
+                rescale=rescale
+            )
+            results = self.add_pred_to_datasample(batch_data_samples, results_list)
+
+        return results
+
+    def add_pred_to_datasample(self, data_samples: SampleList,
+                               results_list: List[dict]) -> SampleList:
+        """Add predictions to `DetDataSample`.
+
+        Args:
+            data_samples (list[:obj:`DetDataSample`], optional): A batch of
+                data samples that contain annotations and predictions.
+            results_list (List[dict]): Instance segmentation, segmantic
+                segmentation and panoptic segmentation results.
+
+        Returns:
+            list[:obj:`DetDataSample`]: Detection results of the
+            input images. Each DetDataSample usually contain
+            'pred_instances' and `pred_panoptic_seg`. And the
+            ``pred_instances`` usually contains following keys.
+
+                - scores (Tensor): Classification scores, has a shape
+                    (num_instance, )
+                - labels (Tensor): Labels of bboxes, has a shape
+                    (num_instances, ).
+                - bboxes (Tensor): Has a shape (num_instances, 4),
+                    the last dimension 4 arrange as (x1, y1, x2, y2).
+                - masks (Tensor): Has a shape (num_instances, H, W).
+
+            And the ``pred_panoptic_seg`` contains the following key
+
+                - sem_seg (Tensor): panoptic segmentation mask, has a
+                    shape (1, h, w).
+        """
+        for data_sample, pred_results in zip(data_samples, results_list):
+            if 'pan_results' in pred_results:
+                data_sample.pred_panoptic_seg = pred_results['pan_results']
+
+            if 'ins_results' in pred_results:
+                data_sample.pred_instances = pred_results['ins_results']
+
+            assert 'sem_results' not in pred_results
+
+        return data_samples
+
+    def add_track_pred_to_datasample(self, data_samples: SampleList, results_list: List[dict]) -> SampleList:
+        for data_sample, pred_results in zip(data_samples, results_list):
+            if 'pan_results' in pred_results:
+                assert self.num_stuff_classes > 0
+                data_sample.pred_track_panoptic_seg = pred_results['pan_results']
+
+            if 'ins_results' in pred_results:
+                bboxes = pred_results['ins_results']['bboxes']
+                labels = pred_results['ins_results']['labels']
+                track_ids = torch.arange(len(bboxes), dtype=labels.dtype, device=bboxes.device) + 1
+                pred_results['ins_results']['instances_id'] = track_ids
+                data_sample.pred_track_instances = pred_results['ins_results']
+
+            if 'pro_results' in pred_results:
+                data_sample.pred_track_proposal = pred_results['pro_results']
+
+            assert 'sem_results' not in pred_results
+
+        return data_samples
+
+    def _forward(
+            self,
+            batch_inputs: Tensor,
+            batch_data_samples: OptSampleList = None) -> Tuple[List[Tensor]]:
+        """Network forward process. Usually includes backbone, neck and head
+        forward without any post-processing.
+
+         Args:
+            batch_inputs (Tensor): Inputs with shape (N, C, H, W).
+            batch_data_samples (list[:obj:`DetDataSample`]): The batch
+                data samples. It usually includes information such
+                as `gt_instance` or `gt_panoptic_seg` or `gt_sem_seg`.
+
+        Returns:
+            tuple[List[Tensor]]: A tuple of features from ``panoptic_head``
+            forward.
+        """
+        if isinstance(batch_data_samples[0], TrackDataSample):
+            bs, num_frames, three, h, w = batch_inputs.shape
+            assert three == 3, "Only supporting images with 3 channels."
+
+            x = batch_inputs.reshape((bs * num_frames, three, h, w))
+            feats = self.extract_feat(x)
+        else:
+            feats = self.extract_feat(batch_inputs)
+        results = self.panoptic_head.forward(feats, batch_data_samples)
+        return results
+
+    def open_voc_inference(self, feats, mask_cls_results, mask_pred_results):
+        if len(mask_pred_results.shape) == 5:
+            batch_size = mask_cls_results.shape[0]
+            num_frames = mask_pred_results.shape[2]
+            mask_pred_results = mask_pred_results.permute(0, 2, 1, 3, 4).flatten(0, 1)
+        else:
+            batch_size = mask_cls_results.shape[0]
+            num_frames = 0
+        clip_feat = self.backbone.get_clip_feature(feats[-1])
+        clip_feat_mask = F.interpolate(
+            mask_pred_results,
+            size=clip_feat.shape[-2:],
+            mode='bilinear',
+            align_corners=False
+        )
+        if num_frames > 0:
+            clip_feat_mask = clip_feat_mask.unflatten(0, (batch_size, num_frames)).permute(0, 2, 1, 3, 4).flatten(2, 3)
+            clip_feat = clip_feat.unflatten(0, (batch_size, num_frames)).permute(0, 2, 1, 3, 4).flatten(2, 3)
+        instance_feat = mask_pool(clip_feat, clip_feat_mask)
+        instance_feat = self.backbone.forward_feat(instance_feat)
+        clip_logit = self.panoptic_head.forward_logit(instance_feat)
+        clip_logit = clip_logit[..., :-1]
+        query_logit = mask_cls_results[..., :-1]
+
+        clip_logit = clip_logit.softmax(-1)
+        query_logit = query_logit.softmax(-1)
+        overlapping_mask = torch.tensor(self.OVERLAPPING, dtype=torch.float32, device=clip_logit.device)
+
+        valid_masking = ((clip_feat_mask > 0).to(dtype=torch.float32).flatten(-2).sum(-1) > 0).to(
+            torch.float32)[..., None]
+        alpha = torch.ones_like(clip_logit) * self.alpha * valid_masking
+        beta = torch.ones_like(clip_logit) * self.beta * valid_masking
+
+        cls_logits_seen = (
+            (query_logit ** (1 - alpha) * clip_logit ** alpha).log()
+            * overlapping_mask
+        )
+        cls_logits_unseen = (
+            (query_logit ** (1 - beta) * clip_logit ** beta).log()
+            * (1 - overlapping_mask)
+        )
+        cls_results = cls_logits_seen + cls_logits_unseen
+        is_void_prob = F.softmax(mask_cls_results, dim=-1)[..., -1:]
+        mask_cls_results = torch.cat([
+            cls_results.softmax(-1) * (1.0 - is_void_prob), is_void_prob], dim=-1)
+        mask_cls_results = torch.log(mask_cls_results + 1e-8)
+        return mask_cls_results

app/models/detectors/rapsam.py

@@ -0,0 +1,66 @@
+from mmdet.registry import MODELS
+from mmdet.utils import ConfigType, OptConfigType, OptMultiConfig
+
+from mmdet.models.detectors import SingleStageDetector
+
+from .mask2former_vid import Mask2formerVideo
+
+
+@MODELS.register_module()
+class YOSOVideoSam(Mask2formerVideo):
+    OVERLAPPING = None
+
+    def __init__(self,
+                 backbone: ConfigType,
+                 neck: OptConfigType = None,
+                 panoptic_head: OptConfigType = None,
+                 panoptic_fusion_head: OptConfigType = None,
+                 train_cfg: OptConfigType = None,
+                 test_cfg: OptConfigType = None,
+                 data_preprocessor: OptConfigType = None,
+                 inference_sam: bool = False,
+                 init_cfg: OptMultiConfig = None
+                 ):
+        super(SingleStageDetector, self).__init__(
+            data_preprocessor=data_preprocessor, init_cfg=init_cfg)
+        self.backbone = MODELS.build(backbone)
+        if neck is not None:
+            self.neck = MODELS.build(neck)
+
+        panoptic_head_ = panoptic_head.deepcopy()
+        panoptic_head_.update(train_cfg=train_cfg)
+        panoptic_head_.update(test_cfg=test_cfg)
+        self.panoptic_head = MODELS.build(panoptic_head_)
+
+        panoptic_fusion_head_ = panoptic_fusion_head.deepcopy()
+        panoptic_fusion_head_.update(test_cfg=test_cfg)
+        self.panoptic_fusion_head = MODELS.build(panoptic_fusion_head_)
+
+        self.num_things_classes = self.panoptic_head.num_things_classes
+        self.num_stuff_classes = self.panoptic_head.num_stuff_classes
+        self.num_classes = self.panoptic_head.num_classes
+
+        self.train_cfg = train_cfg
+        self.test_cfg = test_cfg
+
+        self.alpha = 0.4
+        self.beta = 0.8
+
+        self.inference_sam = inference_sam
+
+    def predict_with_point(self, x, batch_data_samples):
+        feats = self.extract_feat(x)
+        mask_cls_results, mask_pred_results, iou_results = self.panoptic_head.predict(feats, batch_data_samples)
+
+        if 'gt_instances_collected' not in batch_data_samples[0]:
+            results_list = self.panoptic_fusion_head.predict(
+                mask_cls_results,
+                mask_pred_results,
+                batch_data_samples,
+                iou_results=iou_results,
+                rescale=False
+            )
+            mask_pred_results = results_list[0]['pan_results'].sem_seg[None]
+            mask_cls_results = mask_cls_results
+
+        return mask_pred_results.cpu().numpy(), mask_cls_results.cpu().numpy()

app/models/heads/__init__.py

@@ -0,0 +1 @@
+from .rapsam_head import RapSAMVideoHead

app/models/heads/mask2former_vid.py

@@ -0,0 +1,616 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import copy
+import os
+from typing import List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.distributed as dist
+from mmcv.cnn import Conv2d
+from mmdet.models import Mask2FormerTransformerDecoder
+from mmengine.dist import get_dist_info
+from mmengine.model import caffe2_xavier_init, ModuleList
+from torch import Tensor
+from mmdet.models.layers import MLP, inverse_sigmoid
+from mmdet.models.layers import coordinate_to_encoding
+from mmdet.structures.bbox import bbox_xyxy_to_cxcywh
+
+from mmdet.registry import MODELS, TASK_UTILS
+from mmdet.structures import SampleList, TrackDataSample
+from mmdet.utils import (ConfigType, OptConfigType, OptMultiConfig)
+from mmdet.models.layers import SinePositionalEncoding3D
+from mmdet.models.dense_heads.anchor_free_head import AnchorFreeHead
+from mmcv.cnn.bricks.transformer import MultiheadAttention
+from app.models.utils import mask_pool
+
+
+@MODELS.register_module()
+class Mask2FormerVideoHead(AnchorFreeHead):
+    """Implements the Mask2Former head.
+
+    See `Masked-attention Mask Transformer for Universal Image
+    Segmentation <https://arxiv.org/pdf/2112.01527>`_ for details.
+
+    Args:
+        in_channels (list[int]): Number of channels in the input feature map.
+        feat_channels (int): Number of channels for features.
+        out_channels (int): Number of channels for output.
+        num_things_classes (int): Number of things.
+        num_stuff_classes (int): Number of stuff.
+        num_queries (int): Number of query in Transformer decoder.
+        pixel_decoder (:obj:`ConfigDict` or dict): Config for pixel
+            decoder. Defaults to None.
+        enforce_decoder_input_project (bool, optional): Whether to add
+            a layer to change the embed_dim of tranformer encoder in
+            pixel decoder to the embed_dim of transformer decoder.
+            Defaults to False.
+        transformer_decoder (:obj:`ConfigDict` or dict): Config for
+            transformer decoder. Defaults to None.
+        positional_encoding (:obj:`ConfigDict` or dict): Config for
+            transformer decoder position encoding. Defaults to
+            dict(num_feats=128, normalize=True).
+        loss_cls (:obj:`ConfigDict` or dict): Config of the classification
+            loss. Defaults to None.
+        loss_mask (:obj:`ConfigDict` or dict): Config of the mask loss.
+            Defaults to None.
+        loss_dice (:obj:`ConfigDict` or dict): Config of the dice loss.
+            Defaults to None.
+        train_cfg (:obj:`ConfigDict` or dict, optional): Training config of
+            Mask2Former head.
+        test_cfg (:obj:`ConfigDict` or dict, optional): Testing config of
+            Mask2Former head.
+        init_cfg (:obj:`ConfigDict` or dict or list[:obj:`ConfigDict` or \
+            dict], optional): Initialization config dict. Defaults to None.
+    """
+
+    def __init__(self,
+                 in_channels: List[int],
+                 feat_channels: int,
+                 out_channels: int,
+                 num_mask_tokens: int = 1,
+                 num_things_classes: int = 80,
+                 num_stuff_classes: int = 53,
+                 num_queries: int = 100,
+                 num_transformer_feat_level: int = 3,
+                 pixel_decoder: ConfigType = ...,
+                 enforce_decoder_input_project: bool = False,
+                 transformer_decoder: ConfigType = ...,
+                 positional_encoding: ConfigType = None,
+                 loss_cls: ConfigType = None,
+                 loss_mask: ConfigType = None,
+                 loss_dice: ConfigType = None,
+                 train_cfg: OptConfigType = None,
+                 test_cfg: OptConfigType = None,
+                 init_cfg: OptMultiConfig = None,
+                 # ov configs
+                 sphere_cls: bool = False,
+                 ov_classifier_name: Optional[str] = None,
+                 logit: Optional[int] = None,
+                 use_adaptor = False,
+                 **kwargs) -> None:
+        super(AnchorFreeHead, self).__init__(init_cfg=init_cfg)
+        self.use_adaptor = use_adaptor
+
+        self.num_mask_tokens = num_mask_tokens
+        self.mask_tokens = nn.Embedding(num_mask_tokens, feat_channels)
+        self.pb_embedding = nn.Embedding(2, feat_channels)
+        self.pos_linear = nn.Linear(2 * feat_channels, feat_channels)
+
+        self.num_things_classes = num_things_classes
+        self.num_stuff_classes = num_stuff_classes
+        self.num_classes = self.num_things_classes + self.num_stuff_classes
+        self.num_queries = num_queries
+        self.num_transformer_feat_level = num_transformer_feat_level
+        self.num_heads = transformer_decoder.layer_cfg.cross_attn_cfg.num_heads
+        self.num_transformer_decoder_layers = transformer_decoder.num_layers
+        # assert pixel_decoder.encoder.layer_cfg. \
+        #            self_attn_cfg.num_levels == num_transformer_feat_level
+        pixel_decoder_ = copy.deepcopy(pixel_decoder)
+        pixel_decoder_.update(
+            in_channels=in_channels,
+            feat_channels=feat_channels,
+            out_channels=out_channels)
+        self.pixel_decoder = MODELS.build(pixel_decoder_)
+        self.transformer_decoder = Mask2FormerTransformerDecoder(
+            **transformer_decoder)
+        self.decoder_embed_dims = self.transformer_decoder.embed_dims
+
+        self.decoder_input_projs = ModuleList()
+        # from low resolution to high resolution
+        for _ in range(num_transformer_feat_level):
+            if (self.decoder_embed_dims != feat_channels
+                    or enforce_decoder_input_project):
+                self.decoder_input_projs.append(
+                    Conv2d(
+                        feat_channels, self.decoder_embed_dims, kernel_size=1))
+            else:
+                self.decoder_input_projs.append(nn.Identity())
+        self.decoder_positional_encoding = SinePositionalEncoding3D(
+            **positional_encoding)
+        self.query_embed = nn.Embedding(self.num_queries, feat_channels)
+        self.query_feat = nn.Embedding(self.num_queries, feat_channels)
+        # from low resolution to high resolution
+        self.level_embed = nn.Embedding(self.num_transformer_feat_level,
+                                        feat_channels)
+
+        if not sphere_cls:
+            self.cls_embed = nn.Linear(feat_channels, self.num_classes + 1)
+        self.mask_embed = nn.Sequential(
+            nn.Linear(feat_channels, feat_channels), nn.ReLU(inplace=True),
+            nn.Linear(feat_channels, feat_channels), nn.ReLU(inplace=True),
+            nn.Linear(feat_channels, out_channels))
+
+        self.iou_embed = nn.Sequential(
+            nn.Linear(feat_channels, feat_channels), nn.ReLU(inplace=True),
+            nn.Linear(feat_channels, feat_channels), nn.ReLU(inplace=True),
+            nn.Linear(feat_channels, 1))
+
+        self.test_cfg = test_cfg
+        self.train_cfg = train_cfg
+        if train_cfg:
+            self.assigner = TASK_UTILS.build(self.train_cfg['assigner'])
+            self.sampler = TASK_UTILS.build(
+                self.train_cfg['sampler'], default_args=dict(context=self))
+            self.num_points = self.train_cfg.get('num_points', 12544)
+            self.oversample_ratio = self.train_cfg.get('oversample_ratio', 3.0)
+            self.importance_sample_ratio = self.train_cfg.get(
+                'importance_sample_ratio', 0.75)
+
+        self.class_weight = loss_cls.class_weight
+        self.loss_cls = MODELS.build(loss_cls)
+        self.loss_mask = MODELS.build(loss_mask)
+        self.loss_dice = MODELS.build(loss_dice)
+
+        # prepare OV things
+        # OV cls embed
+        if sphere_cls:
+            rank, world_size = get_dist_info()
+            if ov_classifier_name is None:
+                _dim = 1024  # temporally hard code
+                cls_embed = torch.empty(self.num_classes, _dim)
+                torch.nn.init.orthogonal_(cls_embed)
+                cls_embed = cls_embed[:, None]
+            else:
+                # ov_path = os.path.join(os.path.expanduser('~/.cache/embd'), f"{ov_classifier_name}.pth")
+                ov_path = os.path.join('./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))
+            if self.loss_cls and self.loss_cls.use_sigmoid:
+                pass
+            else:
+                _dim = cls_embed.size(2)
+                _prototypes = cls_embed.size(1)
+
+                if rank == 0:
+                    back_token = torch.zeros(1, _dim, dtype=torch.float32, device='cuda')
+                    # back_token = back_token / back_token.norm(p=2, dim=-1, keepdim=True)
+                else:
+                    back_token = torch.empty(1, _dim, dtype=torch.float32, device='cuda')
+                if world_size > 1:
+                    dist.broadcast(back_token, src=0)
+                back_token = back_token.to(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)
+
+            # cls embd proj
+            cls_embed_dim = self.cls_embed.size(0)
+            self.cls_proj = nn.Sequential(
+                nn.Linear(feat_channels, feat_channels), nn.ReLU(inplace=True),
+                nn.Linear(feat_channels, feat_channels), nn.ReLU(inplace=True),
+                nn.Linear(feat_channels, cls_embed_dim)
+            )
+
+            # Haobo Yuan:
+            # For the logit_scale, I refer to this issue.
+            # https://github.com/openai/CLIP/issues/46#issuecomment-945062212
+            # https://github.com/openai/CLIP/issues/46#issuecomment-782558799
+            # Based on my understanding, it is a mistake of CLIP.
+            # Because they mention that they refer to InstDisc (Wu, 2018) paper.
+            # InstDisc set a non-learnable temperature to np.log(1 / 0.07).
+            # 4.6052 is np.log(1 / 0.01)
+            # np.log(1 / 0.07) will be fast converged to np.log(1 / 0.01)
+            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)
+
+            # Mask Pooling
+            self.mask_pooling = mask_pool
+            self.mask_pooling_proj = nn.Sequential(
+                nn.LayerNorm(feat_channels),
+                nn.Linear(feat_channels, feat_channels)
+            )
+
+        if use_adaptor:
+            cross_attn_cfg = dict(embed_dims=256, batch_first=True, num_heads=8)
+            self.panoptic_attn = MultiheadAttention(**cross_attn_cfg)
+            self.panoptic_norm = nn.LayerNorm(256)
+            if sphere_cls:
+                cls_embed_dim = self.cls_embed.size(0)
+                self.panoptic_cls = nn.Sequential(
+                    nn.Linear(feat_channels, cls_embed_dim)
+                )
+            else:
+                raise NotImplementedError
+            self.prompt_attn = MultiheadAttention(**cross_attn_cfg)
+            self.prompt_norm = nn.LayerNorm(256)
+            self.prompt_iou = nn.Linear(256, 1)
+
+    def init_weights(self) -> None:
+        for m in self.decoder_input_projs:
+            if isinstance(m, Conv2d):
+                caffe2_xavier_init(m, bias=0)
+
+        self.pixel_decoder.init_weights()
+
+        for p in self.transformer_decoder.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_normal_(p)
+
+    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 _forward_head(self, decoder_out: Tensor, mask_feature: Tensor,
+                      attn_mask_target_size: Tuple[int, int],
+                      num_frames: int = 0) -> Tuple[Tensor]:
+        """Forward for head part which is called after every decoder layer.
+
+        Args:
+            decoder_out (Tensor): in shape (batch_size, num_queries, c).
+            mask_feature (Tensor): in shape (batch_size, c, h, w).
+            attn_mask_target_size (tuple[int, int]): target attention
+                mask size.
+
+        Returns:
+            tuple: A tuple contain three elements.
+
+                - cls_pred (Tensor): Classification scores in shape \
+                    (batch_size, num_queries, cls_out_channels). \
+                    Note `cls_out_channels` should includes background.
+                - mask_pred (Tensor): Mask scores in shape \
+                    (batch_size, num_queries,h, w).
+                - attn_mask (Tensor): Attention mask in shape \
+                    (batch_size * num_heads, num_queries, h, w).
+                - num_frames: How many frames are there in video.
+        """
+        decoder_out = self.transformer_decoder.post_norm(decoder_out)
+        # shape (num_queries, batch_size, c)
+        if isinstance(self.cls_embed, nn.Module):
+            cls_pred = self.cls_embed(decoder_out)
+        # shape (num_queries, batch_size, c)
+        mask_embed = self.mask_embed(decoder_out)
+        # shape (num_queries, batch_size, h, w)
+        mask_pred = torch.einsum('bqc,bchw->bqhw', mask_embed, mask_feature)
+
+        if not isinstance(self.cls_embed, nn.Module):
+            maskpool_embd = self.mask_pooling(x=mask_feature, mask=mask_pred.detach())
+            maskpool_embd = self.mask_pooling_proj(maskpool_embd)
+            cls_embd = self.cls_proj(maskpool_embd + decoder_out)
+            cls_pred = self.forward_logit(cls_embd)
+
+        iou_pred = self.iou_embed(decoder_out)
+
+        if num_frames > 0:
+            assert len(mask_pred.shape) == 4
+            assert mask_pred.shape[2] % num_frames == 0
+            frame_h = mask_pred.shape[2] // num_frames
+            num_q = mask_pred.shape[1]
+            _mask_pred = mask_pred.unflatten(-2, (num_frames, frame_h)).flatten(1, 2)
+            attn_mask = F.interpolate(
+                _mask_pred,
+                attn_mask_target_size,
+                mode='bilinear',
+                align_corners=False)
+            attn_mask = attn_mask.unflatten(1, (num_q, num_frames)).flatten(2, 3)
+        else:
+            attn_mask = F.interpolate(
+                mask_pred,
+                attn_mask_target_size,
+                mode='bilinear',
+                align_corners=False)
+        # shape (num_queries, batch_size, h, w) ->
+        #   (batch_size * num_head, num_queries, h, w)
+        attn_mask = attn_mask.flatten(2).unsqueeze(1).repeat(
+            (1, self.num_heads, 1, 1)).flatten(0, 1)
+        attn_mask = attn_mask.sigmoid() < 0.5
+        attn_mask = attn_mask.detach()
+
+        return cls_pred, mask_pred, iou_pred, attn_mask
+
+    def forward(self, x: List[Tensor], batch_data_samples: SampleList) -> Tuple[List[Tensor]]:
+        """Forward function.
+
+        Args:
+            x (list[Tensor]): Multi scale Features from the
+                upstream network, each is a 4D-tensor.
+            batch_data_samples (List[:obj:`DetDataSample`]): The Data
+                Samples. It usually includes information such as
+                `gt_instance`, `gt_panoptic_seg` and `gt_sem_seg`.
+
+        Returns:
+            tuple[list[Tensor]]: A tuple contains two elements.
+
+                - cls_pred_list (list[Tensor)]: Classification logits \
+                    for each decoder layer. Each is a 3D-tensor with shape \
+                    (batch_size, num_queries, cls_out_channels). \
+                    Note `cls_out_channels` should includes background.
+                - mask_pred_list (list[Tensor]): Mask logits for each \
+                    decoder layer. Each with shape (batch_size, num_queries, \
+                    h, w).
+        """
+        batch_img_metas = []
+        if isinstance(batch_data_samples[0], TrackDataSample):
+            for track_sample in batch_data_samples:
+                cur_list = []
+                for det_sample in track_sample:
+                    cur_list.append(det_sample.metainfo)
+                batch_img_metas.append(cur_list)
+            num_frames = len(batch_img_metas[0])
+        else:
+            for data_sample in batch_data_samples:
+                batch_img_metas.append(data_sample.metainfo)
+            num_frames = 0
+        batch_size = len(batch_img_metas)
+        #(bs_nf, c, h,w)
+        mask_features, multi_scale_memorys = self.pixel_decoder(x)
+        if num_frames > 0:
+            mask_features = mask_features.unflatten(0, (batch_size, num_frames))
+            mask_features = mask_features.transpose(1, 2).flatten(2, 3) #(bs, c, nf*h,w)
+        decoder_inputs = []
+        decoder_positional_encodings = []
+        for i in range(self.num_transformer_feat_level):
+            decoder_input = self.decoder_input_projs[i](multi_scale_memorys[i]) #(bs_nf, c, h,w)
+            decoder_input = decoder_input.flatten(2).permute(0, 2, 1) #(bs_nf,h*w, c)
+            if num_frames > 0:
+                decoder_input = decoder_input.unflatten(0, (batch_size, num_frames))
+                decoder_input = decoder_input.flatten(1, 2) #(bs, nf*h*w, c)
+            level_embed = self.level_embed.weight[i].view(1, 1, -1)
+            decoder_input = decoder_input + level_embed
+
+            # shape (batch_size, c, h, w) -> (batch_size, h*w, c)
+            num_frames_real = 1 if num_frames == 0 else num_frames
+            mask = decoder_input.new_zeros(
+                (batch_size, num_frames_real) + multi_scale_memorys[i].shape[-2:],
+                dtype=torch.bool)
+            decoder_positional_encoding = self.decoder_positional_encoding(
+                mask)
+            decoder_positional_encoding = decoder_positional_encoding.transpose(
+                1, 2).flatten(2).permute(0, 2, 1)
+            decoder_inputs.append(decoder_input) #(bs, nf*h*w, c)
+            decoder_positional_encodings.append(decoder_positional_encoding) #(bs, nf*h*w, c)
+
+        if self.prompt_training:
+            query_feat, input_query_bbox, self_attn_mask, _ = self.prepare_for_dn_mo(
+                batch_data_samples)
+            query_embed = coordinate_to_encoding(input_query_bbox.sigmoid())
+            query_embed = self.pos_linear(query_embed)
+        else:
+            query_feat = self.query_feat.weight.unsqueeze(0).repeat((batch_size, 1, 1))
+            query_embed = self.query_embed.weight.unsqueeze(0).repeat((batch_size, 1, 1))
+            self_attn_mask = None
+
+        cls_pred_list = []
+        mask_pred_list = []
+        iou_pred_list = []
+        cls_pred, mask_pred, iou_pred, attn_mask = self._forward_head(
+            query_feat, mask_features, multi_scale_memorys[0].shape[-2:],
+            num_frames=num_frames
+        )
+        cls_pred_list.append(cls_pred)
+        iou_pred_list.append(iou_pred)
+        if num_frames > 0: #(bs, 100, nf*h, w)-->(bs, 100, nf, h, w)
+            mask_pred = mask_pred.unflatten(2, (num_frames, -1))
+        mask_pred_list.append(mask_pred)
+
+        for i in range(self.num_transformer_decoder_layers):
+            level_idx = i % self.num_transformer_feat_level
+            # if a mask is all True(all background), then set it all False.
+            attn_mask[torch.where(attn_mask.sum(-1) == attn_mask.shape[-1])] = False
+
+            # cross_attn + self_attn
+            layer = self.transformer_decoder.layers[i]
+            query_feat = layer(
+                query=query_feat, #(bs, 100, c)
+                key=decoder_inputs[level_idx], #(bs, nf*h*w, c)
+                value=decoder_inputs[level_idx],
+                query_pos=query_embed,
+                key_pos=decoder_positional_encodings[level_idx],
+                cross_attn_mask=attn_mask,
+                self_attn_mask=self_attn_mask,
+                query_key_padding_mask=None,
+                # here we do not apply masking on padded region
+                key_padding_mask=None)
+            cls_pred, mask_pred, iou_pred, attn_mask = self._forward_head(
+                query_feat, mask_features, multi_scale_memorys[(i + 1) % self.num_transformer_feat_level].shape[-2:],
+                num_frames=num_frames
+            )
+
+            cls_pred_list.append(cls_pred)
+            iou_pred_list.append(iou_pred)
+            if num_frames > 0:
+                mask_pred = mask_pred.unflatten(2, (num_frames, -1))
+            mask_pred_list.append(mask_pred)
+        
+        if self.use_adaptor:
+            keys = mask_features.flatten(2).transpose(1, 2).contiguous()
+            h, w = mask_features.shape[-2] // num_frames_real, mask_features.shape[-1]
+            mask = decoder_input.new_zeros((batch_size, num_frames_real, h, w), dtype=torch.bool)
+            key_pos = self.decoder_positional_encoding(mask)
+            key_pos = key_pos.transpose(1, 2).flatten(2).permute(0, 2, 1)
+            if not self.prompt_training:
+                object_kernels = self.panoptic_attn(query_feat, keys, key_pos=key_pos, query_pos=query_embed)
+                object_kernels = self.panoptic_norm(object_kernels)
+                mask_preds = torch.einsum('bnc,bchw->bnhw', object_kernels, mask_features)
+                
+                cls_embd = self.panoptic_cls(object_kernels)
+                cls_scores = torch.einsum('bnc,ckp->bnkp', F.normalize(cls_embd, dim=-1), self.cls_embed)
+                cls_scores = cls_scores.max(-1).values
+                cls_scores = self.logit_scale.exp() * cls_scores
+                
+                if num_frames > 0: 
+                    mask_pred_list.append(mask_preds.unflatten(2, (num_frames, -1)))
+                else:
+                    mask_pred_list.append(mask_preds)
+                cls_pred_list.append(cls_scores)
+                iou_pred_list.append(iou_pred_list[-1])
+            else:
+                object_kernels = self.prompt_attn(query_feat, keys, key_pos=key_pos, query_pos=query_embed)
+                object_kernels = self.prompt_norm(object_kernels)
+                iou_preds = self.prompt_iou(object_kernels)
+                mask_preds = torch.einsum('bnc,bchw->bnhw', object_kernels, mask_features)
+                
+                if num_frames > 0: 
+                    mask_pred_list.append(mask_preds.unflatten(2, (num_frames, -1)))
+                else:
+                    mask_pred_list.append(mask_preds)
+                cls_pred_list.append(cls_pred_list[-1])
+                iou_pred_list.append(iou_preds)
+
+        return cls_pred_list, mask_pred_list, iou_pred_list, query_feat
+
+    def predict(self, x: Tuple[Tensor],
+                batch_data_samples: SampleList,
+                return_query=False,
+                ) -> Tuple[Tensor, ...]:
+        """Test without augmentaton.
+
+        Args:
+            return_query:
+            x (tuple[Tensor]): Multi-level features from the
+                upstream network, each is a 4D-tensor.
+            batch_data_samples (List[:obj:`DetDataSample`]): The Data
+                Samples. It usually includes information such as
+                `gt_instance`, `gt_panoptic_seg` and `gt_sem_seg`.
+
+        Returns:
+            tuple[Tensor]: A tuple contains two tensors.
+
+                - mask_cls_results (Tensor): Mask classification logits,\
+                    shape (batch_size, num_queries, cls_out_channels).
+                    Note `cls_out_channels` should includes background.
+                - mask_pred_results (Tensor): Mask logits, shape \
+                    (batch_size, num_queries, h, w).
+        """
+        self.prompt_training = False
+        data_sample = batch_data_samples[0]
+        if isinstance(data_sample, TrackDataSample):
+            img_shape = data_sample[0].metainfo['batch_input_shape']
+            num_frames = len(data_sample)
+        else:
+            if 'gt_instances_collected' in data_sample:
+                self.prompt_training = True
+            img_shape = data_sample.metainfo['batch_input_shape']
+            num_frames = 0
+        all_cls_scores, all_mask_preds, all_iou_preds, query_feat = self(x, batch_data_samples)
+        mask_cls_results = all_cls_scores[-1]
+        mask_pred_results = all_mask_preds[-1]
+        iou_results = all_iou_preds[-1]
+
+        if num_frames > 0:
+            mask_pred_results = mask_pred_results.flatten(1, 2)
+        mask_pred_results = F.interpolate(
+            mask_pred_results,
+            size=(img_shape[0], img_shape[1]),
+            mode='bilinear',
+            align_corners=False)
+        if num_frames > 0:
+            num_queries = mask_cls_results.shape[1]
+            mask_pred_results = mask_pred_results.unflatten(1, (num_queries, num_frames))
+
+        if return_query:
+            return mask_cls_results, mask_pred_results, query_feat, iou_results
+        else:
+            return mask_cls_results, mask_pred_results, iou_results
+
+    def prepare_for_dn_mo(self, batch_data_samples):
+        scalar, noise_scale = 100, 0.4
+        gt_instances = [t.gt_instances_collected for t in batch_data_samples]
+
+        point_coords = torch.stack([inst.point_coords for inst in gt_instances])
+        pb_labels = torch.stack([inst['pb_labels'] for inst in gt_instances])
+        labels = torch.zeros_like(pb_labels).long()
+
+        boxes = point_coords  # + boxes
+
+        factors = []
+        for i, data_sample in enumerate(batch_data_samples):
+            h, w, = data_sample.metainfo['img_shape']
+            factor = boxes[i].new_tensor([w, h, w, h]).unsqueeze(0).repeat(boxes[i].size(0), 1)
+            factors.append(factor)
+        factors = torch.stack(factors, 0)
+
+        boxes = bbox_xyxy_to_cxcywh(boxes / factors)  # xyxy / factor or xywh / factor ????
+        # box_start = [t['box_start'] for t in targets]
+        box_start = [len(point) for point in point_coords]
+
+        known_labels = labels
+        known_pb_labels = pb_labels
+        known_bboxs = boxes
+
+        known_labels_expaned = known_labels.clone()
+        known_pb_labels_expaned = known_pb_labels.clone()
+        known_bbox_expand = known_bboxs.clone()
+
+        if noise_scale > 0 and self.training:
+            diff = torch.zeros_like(known_bbox_expand)
+            diff[:, :, :2] = known_bbox_expand[:, :, 2:] / 2
+            diff[:, :, 2:] = known_bbox_expand[:, :, 2:]
+            # add very small noise to input points; no box
+            sc = 0.01
+            for i, st in enumerate(box_start):
+                diff[i, :st] = diff[i, :st] * sc
+            known_bbox_expand += torch.mul(
+                (torch.rand_like(known_bbox_expand) * 2 - 1.0),
+                diff) * noise_scale
+
+            known_bbox_expand = known_bbox_expand.clamp(min=0.0, max=1.0)
+
+        input_label_embed = self.pb_embedding(known_pb_labels_expaned)
+
+        input_bbox_embed = inverse_sigmoid(known_bbox_expand)
+
+        input_label_embed = input_label_embed.repeat_interleave(
+            self.num_mask_tokens,
+            1) + self.mask_tokens.weight.unsqueeze(0).repeat(
+                input_label_embed.shape[0], input_label_embed.shape[1], 1)
+        input_bbox_embed = input_bbox_embed.repeat_interleave(
+            self.num_mask_tokens, 1)
+
+        single_pad = self.num_mask_tokens
+
+        # NOTE scalar is modified to 100, each click cannot see each other
+        scalar = int(input_label_embed.shape[1] / self.num_mask_tokens)
+
+        pad_size = input_label_embed.shape[1]
+
+        if input_label_embed.shape[1] > 0:
+            input_query_label = input_label_embed
+            input_query_bbox = input_bbox_embed
+
+        tgt_size = pad_size
+        attn_mask = torch.ones(tgt_size, tgt_size).to('cuda') < 0
+        # match query cannot see the reconstruct
+        attn_mask[pad_size:, :pad_size] = True
+        # reconstruct cannot see each other
+        for i in range(scalar):
+            if i == 0:
+                attn_mask[single_pad * i:single_pad * (i + 1), single_pad * (i + 1):pad_size] = True
+            if i == scalar - 1:
+                attn_mask[single_pad * i:single_pad * (i + 1), :single_pad * i] = True
+            else:
+                attn_mask[single_pad * i:single_pad * (i + 1), single_pad * (i + 1):pad_size] = True
+                attn_mask[single_pad * i:single_pad * (i + 1), :single_pad * i] = True
+        mask_dict = {
+            'known_lbs_bboxes': (known_labels, known_bboxs),
+            'pad_size': pad_size,
+            'scalar': scalar,
+        }
+        return input_query_label, input_query_bbox, attn_mask, mask_dict

app/models/heads/rapsam_head.py

@@ -0,0 +1,227 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from typing import List, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from mmdet.models.layers import coordinate_to_encoding
+
+from mmdet.registry import MODELS, TASK_UTILS
+from mmdet.structures import SampleList, TrackDataSample
+from mmdet.utils import (ConfigType, OptConfigType, OptMultiConfig)
+from mmdet.models.dense_heads.anchor_free_head import AnchorFreeHead
+
+from mmcv.cnn.bricks.transformer import MultiheadAttention
+
+from .mask2former_vid import Mask2FormerVideoHead
+from .yoso_head import CrossAttenHead, KernelUpdator
+
+@MODELS.register_module()
+class RapSAMVideoHead(Mask2FormerVideoHead):
+
+    def __init__(self,
+                 frozen_head=False,
+                 frozen_pred=False,
+                 use_adaptor=False,
+                 prompt_with_kernel_updator=False,
+                 panoptic_with_kernel_updator=False,
+                 num_mask_tokens = 1,
+                 num_stages = 3,
+                 use_kernel_updator=False,
+                 sphere_cls = False,
+                 ov_classifier_name = None,
+                 temperature=0.1,
+                 feat_channels=256,
+                 num_things_classes: int = 80,
+                 num_stuff_classes: int = 53,
+                 num_queries: int = 100,
+                 loss_cls: ConfigType = None,
+                 loss_mask: ConfigType = None,
+                 loss_dice: ConfigType = None,
+                 train_cfg: OptConfigType = None,
+                 test_cfg: OptConfigType = None,
+                 init_cfg: OptMultiConfig = None,
+                 matching_whole_map: bool = False,
+                 enable_box_query: bool = False,
+                 **kwargs) -> None:
+        super(AnchorFreeHead, self).__init__(init_cfg=init_cfg)
+        self.prompt_with_kernel_updator = prompt_with_kernel_updator
+        self.panoptic_with_kernel_updator = panoptic_with_kernel_updator
+        self.use_adaptor = use_adaptor
+
+        self.num_mask_tokens = num_mask_tokens
+        self.mask_tokens = nn.Embedding(num_mask_tokens, feat_channels)
+        self.pb_embedding = nn.Embedding(2, feat_channels)
+        self.pos_linear = nn.Linear(2 * feat_channels, feat_channels)
+
+        self.matching_whole_map = matching_whole_map
+        self.enable_box_query = enable_box_query
+
+        self.num_things_classes = num_things_classes
+        self.num_stuff_classes = num_stuff_classes
+        self.num_classes = self.num_things_classes + self.num_stuff_classes
+        self.num_queries = num_queries
+        self.feat_channels = feat_channels
+        self.num_stages = num_stages
+        self.kernels = nn.Embedding(self.num_queries, feat_channels)
+        self.mask_heads = nn.ModuleList()
+        for _ in range(self.num_stages):
+            self.mask_heads.append(CrossAttenHead(
+                self.num_classes, self.feat_channels, self.num_queries,
+                use_kernel_updator=use_kernel_updator,
+                frozen_head=frozen_head, frozen_pred=frozen_pred,
+                sphere_cls=sphere_cls,
+                ov_classifier_name=ov_classifier_name, with_iou_pred=True))
+        self.temperature = temperature
+
+        if use_adaptor:
+            cross_attn_cfg = dict(embed_dims=256, batch_first=True, num_heads=8)
+            if self.panoptic_with_kernel_updator:
+                self.panoptic_attn = KernelUpdator(feat_channels=256)
+                self.panoptic_norm = nn.Identity()
+                if sphere_cls:
+                    cls_embed_dim = self.mask_heads[0].fc_cls.size(0)
+                    self.panoptic_cls = nn.Sequential(
+                        nn.Linear(feat_channels, cls_embed_dim)
+                    )
+                else:
+                    raise NotImplementedError
+                    self.panoptic_cls = nn.Linear(256, self.num_classes+1)
+            else:
+                self.panoptic_attn = MultiheadAttention(**cross_attn_cfg)
+                self.panoptic_norm = nn.LayerNorm(256)
+                if sphere_cls:
+                    cls_embed_dim = self.mask_heads[0].fc_cls.size(0)
+                    self.panoptic_cls = nn.Sequential(
+                        nn.Linear(feat_channels, cls_embed_dim)
+                    )
+                else:
+                    raise NotImplementedError
+                    self.panoptic_cls = nn.Linear(256, self.num_classes+1)
+            
+            if self.prompt_with_kernel_updator:
+                self.prompt_attn = KernelUpdator(feat_channels=256)
+                self.prompt_norm = nn.Identity()
+                self.prompt_iou = nn.Linear(256, 1)
+            else:
+                self.prompt_attn = MultiheadAttention(**cross_attn_cfg)
+                self.prompt_norm = nn.LayerNorm(256)
+                self.prompt_iou = nn.Linear(256, 1)
+
+        self.test_cfg = test_cfg
+        self.train_cfg = train_cfg
+        if train_cfg:
+            self.assigner = TASK_UTILS.build(self.train_cfg['assigner'])
+            self.sampler = TASK_UTILS.build(
+                self.train_cfg['sampler'], default_args=dict(context=self))
+            self.num_points = self.train_cfg.get('num_points', 12544)
+            self.oversample_ratio = self.train_cfg.get('oversample_ratio', 3.0)
+            self.importance_sample_ratio = self.train_cfg.get(
+                'importance_sample_ratio', 0.75)
+
+        self.class_weight = loss_cls.class_weight
+        self.loss_cls = MODELS.build(loss_cls)
+        self.loss_mask = MODELS.build(loss_mask)
+        self.loss_dice = MODELS.build(loss_dice)
+        
+
+    def init_weights(self) -> None:
+        pass
+    
+    def forward(self, x, batch_data_samples: SampleList) -> Tuple[List[Tensor]]:
+        batch_img_metas = []
+        if isinstance(batch_data_samples[0], TrackDataSample):
+            for track_sample in batch_data_samples:
+                cur_list = []
+                for det_sample in track_sample:
+                    cur_list.append(det_sample.metainfo)
+                batch_img_metas.append(cur_list)
+            num_frames = len(batch_img_metas[0])
+        else:
+            for data_sample in batch_data_samples:
+                batch_img_metas.append(data_sample.metainfo)
+            num_frames = 0
+        bs = len(batch_img_metas)
+        
+        all_cls_scores = []
+        all_masks_preds = []
+        all_iou_preds = []
+        if self.prompt_training:
+            input_query_label, input_query_bbox, self_attn_mask, mask_dict = self.prepare_for_dn_mo(
+                batch_data_samples)
+            pos_embed = coordinate_to_encoding(input_query_bbox.sigmoid())
+            pos_embed = self.pos_linear(pos_embed)
+            object_kernels = input_query_label + pos_embed
+        else:
+            object_kernels = self.kernels.weight[None].repeat(bs, 1, 1)
+            self_attn_mask = None
+        mask_features = x
+        if num_frames > 0: # (bs*num_frames, c, h, w) -> (bs, c, num_frames*h, w)
+            mask_features = mask_features.unflatten(0, (bs, num_frames))
+            mask_features = mask_features.transpose(1, 2).flatten(2, 3)
+        
+        mask_preds = torch.einsum('bnc,bchw->bnhw', object_kernels, mask_features)
+        for stage in range(self.num_stages):
+            mask_head = self.mask_heads[stage]
+            cls_scores, mask_preds, iou_preds, object_kernels = mask_head(
+                mask_features, object_kernels, mask_preds, self_attn_mask)
+            cls_scores = cls_scores / self.temperature
+            all_iou_preds.append(iou_preds)
+            all_cls_scores.append(cls_scores)
+            if num_frames > 0: 
+                #(bs,num_query, num_frames*h, w) --> (bs,num_query,num_frames,h,w)
+                all_masks_preds.append(mask_preds.unflatten(2, (num_frames, -1)))
+            else:
+                all_masks_preds.append(mask_preds)
+        
+        if self.use_adaptor:
+            keys = mask_features.flatten(2).transpose(1, 2).contiguous()
+            if not self.prompt_training:
+                if self.panoptic_with_kernel_updator:
+                    hard_sigmoid_masks = (mask_preds.sigmoid() > 0.5).float()
+                    f = torch.einsum('bnhw,bchw->bnc', hard_sigmoid_masks, mask_features)
+                    object_kernels = self.panoptic_attn(f, object_kernels)
+                    object_kernels = self.panoptic_norm(object_kernels)
+                    mask_preds = torch.einsum('bnc,bchw->bnhw', object_kernels, mask_features)
+                else:
+                    object_kernels = self.panoptic_attn(object_kernels, keys)
+                    object_kernels = self.panoptic_norm(object_kernels)
+                    mask_preds = torch.einsum('bnc,bchw->bnhw', object_kernels, mask_features)
+                cls_embd = self.panoptic_cls(object_kernels)
+                cls_scores = torch.einsum('bnc,ckp->bnkp', F.normalize(cls_embd, dim=-1), self.mask_heads[0].fc_cls)
+                cls_scores = cls_scores.max(-1).values
+                cls_scores = self.mask_heads[0].logit_scale.exp() * cls_scores
+                
+                if num_frames > 0: 
+                    all_masks_preds.append(mask_preds.unflatten(2, (num_frames, -1)))
+                else:
+                    all_masks_preds.append(mask_preds)
+                all_cls_scores.append(cls_scores)
+                all_iou_preds.append(all_iou_preds[-1])
+            else:
+                if self.prompt_with_kernel_updator:
+                    hard_sigmoid_masks = (mask_preds.sigmoid() > 0.5).float()
+                    f = torch.einsum('bnhw,bchw->bnc', hard_sigmoid_masks, mask_features)
+                    object_kernels = self.prompt_attn(f, object_kernels)
+                    object_kernels = self.prompt_norm(object_kernels)
+                    iou_preds = self.prompt_iou(object_kernels)
+                    mask_preds = torch.einsum('bnc,bchw->bnhw', object_kernels, mask_features)
+                else:
+                    object_kernels = self.prompt_attn(object_kernels, keys)
+                    object_kernels = self.prompt_norm(object_kernels)
+                    iou_preds = self.prompt_iou(object_kernels)
+                    mask_preds = torch.einsum('bnc,bchw->bnhw', object_kernels, mask_features)
+                if num_frames > 0: 
+                    all_masks_preds.append(mask_preds.unflatten(2, (num_frames, -1)))
+                else:
+                    all_masks_preds.append(mask_preds)
+                all_cls_scores.append(all_cls_scores[-1])
+                all_iou_preds.append(iou_preds)
+        return all_cls_scores, all_masks_preds, all_iou_preds, object_kernels
+
+    def get_targets(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def loss_by_feat(self, *args, **kwargs):
+        raise NotImplementedError

app/models/heads/yoso_head.py

@@ -0,0 +1,531 @@
+from typing import List, Tuple
+import os
+import torch.distributed as dist
+from torch import Tensor
+from mmdet.registry import MODELS, TASK_UTILS
+from mmdet.models.dense_heads import AnchorFreeHead
+from mmdet.structures import SampleList
+from mmdet.models.dense_heads import Mask2FormerHead
+import math
+from mmengine.model.weight_init import trunc_normal_
+import torch
+from torch import nn
+import torch.nn.functional as F
+from mmcv.cnn import build_activation_layer, build_norm_layer
+
+from mmengine.dist import get_dist_info
+
+
+@MODELS.register_module()
+class YOSOHead(Mask2FormerHead):
+    def __init__(self,
+                 num_cls_fcs=1,
+                 num_mask_fcs=1,
+                 sphere_cls=False,
+                 ov_classifier_name=None,
+                 use_kernel_updator=False,
+                 num_stages=3,
+                 feat_channels=256,
+                 out_channels=256,
+                 num_things_classes=80,
+                 num_stuff_classes=53,
+                 num_classes=133,
+                 num_queries=100,
+                 temperature=0.1,
+                 loss_cls=dict(
+                     type='CrossEntropyLoss',
+                     use_sigmoid=False,
+                     loss_weight=2.0,
+                     reduction='mean',
+                     class_weight=[1.0] * 133 + [0.1]),
+                 loss_mask=dict(
+                     type='CrossEntropyLoss',
+                     use_sigmoid=True,
+                     reduction='mean',
+                     loss_weight=5.0),
+                 loss_dice=dict(
+                     type='DiceLoss',
+                     use_sigmoid=True,
+                     activate=True,
+                     reduction='mean',
+                     naive_dice=True,
+                     eps=1.0,
+                     loss_weight=5.0),
+                 train_cfg=None,
+                 test_cfg=None,
+                 init_cfg=None):
+        super(AnchorFreeHead, self).__init__(init_cfg=init_cfg)
+        self.num_stages = num_stages
+        self.feat_channels = feat_channels
+        self.out_channels = out_channels
+        self.num_things_classes = num_things_classes
+        self.num_stuff_classes = num_stuff_classes
+        self.num_classes = num_classes
+        self.num_queries = num_queries
+        self.temperature = temperature
+
+        self.test_cfg = test_cfg
+        self.train_cfg = train_cfg
+        if train_cfg:
+            self.assigner = TASK_UTILS.build(self.train_cfg['assigner'])
+            self.sampler = TASK_UTILS.build(
+                self.train_cfg['sampler'], default_args=dict(context=self))
+            self.num_points = self.train_cfg.get('num_points', 12544)
+            self.oversample_ratio = self.train_cfg.get('oversample_ratio', 3.0)
+            self.importance_sample_ratio = self.train_cfg.get(
+                'importance_sample_ratio', 0.75)
+
+        self.class_weight = loss_cls.class_weight
+        self.loss_cls = MODELS.build(loss_cls)
+        self.loss_mask = MODELS.build(loss_mask)
+        self.loss_dice = MODELS.build(loss_dice)
+
+        self.kernels = nn.Embedding(self.num_queries, self.feat_channels)
+
+        self.mask_heads = nn.ModuleList()
+        for _ in range(self.num_stages):
+            self.mask_heads.append(CrossAttenHead(
+                self.num_classes, self.feat_channels, self.num_queries,
+                use_kernel_updator=use_kernel_updator,
+                sphere_cls=sphere_cls, ov_classifier_name=ov_classifier_name,
+                num_cls_fcs=num_cls_fcs, num_mask_fcs=num_mask_fcs
+            ))
+
+    def init_weights(self) -> None:
+        super(AnchorFreeHead, self).init_weights()
+
+    def forward(self, x: List[Tensor],
+                batch_data_samples: SampleList) -> Tuple[List[Tensor]]:
+        all_cls_scores = []
+        all_masks_preds = []
+        proposal_kernels = self.kernels.weight
+        object_kernels = proposal_kernels[None].repeat(x.shape[0], 1, 1)
+        mask_preds = torch.einsum('bnc,bchw->bnhw', object_kernels, x)
+
+        for stage in range(self.num_stages):
+            mask_head = self.mask_heads[stage]
+            cls_scores, mask_preds, iou_pred, object_kernels = mask_head(x, object_kernels, mask_preds)
+            cls_scores = cls_scores / self.temperature
+
+            all_cls_scores.append(cls_scores)
+            all_masks_preds.append(mask_preds)
+
+        return all_cls_scores, all_masks_preds
+
+    def predict(self, x: Tuple[Tensor], batch_data_samples: SampleList) -> Tuple[Tensor]:
+        batch_img_metas = [
+            data_sample.metainfo for data_sample in batch_data_samples
+        ]
+        all_cls_scores, all_mask_preds = self(x, batch_data_samples)
+        mask_cls_results = all_cls_scores[-1]
+        mask_pred_results = all_mask_preds[-1]
+
+        # upsample masks
+        img_shape = batch_img_metas[0]['batch_input_shape']
+        mask_pred_results = F.interpolate(
+            mask_pred_results,
+            size=(img_shape[0], img_shape[1]),
+            mode='bilinear',
+            align_corners=False)
+
+        return mask_cls_results, mask_pred_results
+
+
+class FFN(nn.Module):
+
+    def __init__(self,
+                 embed_dims=256,
+                 feedforward_channels=1024,
+                 num_fcs=2,
+                 add_identity=True):
+        super(FFN, self).__init__()
+        self.embed_dims = embed_dims
+        self.feedforward_channels = feedforward_channels
+        self.num_fcs = num_fcs
+
+        layers = []
+        in_channels = embed_dims
+        for _ in range(num_fcs - 1):
+            layers.append(nn.Sequential(
+                nn.Linear(in_channels, feedforward_channels),
+                nn.ReLU(True),
+                nn.Dropout(0.0)))
+            in_channels = feedforward_channels
+        layers.append(nn.Linear(feedforward_channels, embed_dims))
+        layers.append(nn.Dropout(0.0))
+        self.layers = nn.Sequential(*layers)
+        self.add_identity = add_identity
+        self.dropout_layer = nn.Dropout(0.0)
+
+    def forward(self, x, identity=None):
+        out = self.layers(x)
+        if not self.add_identity:
+            return self.dropout_layer(out)
+        if identity is None:
+            identity = x
+        return identity + self.dropout_layer(out)
+
+
+class DySepConvAtten(nn.Module):
+    def __init__(self, hidden_dim, num_proposals, conv_kernel_size_1d):
+        super(DySepConvAtten, self).__init__()
+        self.hidden_dim = hidden_dim
+        self.num_proposals = num_proposals
+        self.kernel_size = conv_kernel_size_1d
+
+        self.weight_linear = nn.Linear(self.hidden_dim, self.num_proposals + self.kernel_size)
+        self.norm = nn.LayerNorm(self.hidden_dim)
+
+    def forward(self, query, value):
+        assert query.shape == value.shape
+        B, N, C = query.shape
+
+        dy_conv_weight = self.weight_linear(query)
+        dy_depth_conv_weight = dy_conv_weight[:, :, :self.kernel_size].view(B, self.num_proposals, 1, self.kernel_size)
+        dy_point_conv_weight = dy_conv_weight[:, :, self.kernel_size:].view(B, self.num_proposals, self.num_proposals,
+                                                                            1)
+
+        res = []
+        value = value.unsqueeze(1)
+        for i in range(B):
+            out = F.relu(F.conv1d(input=value[i], weight=dy_depth_conv_weight[i], groups=N, padding='same'))
+            out = F.conv1d(input=out, weight=dy_point_conv_weight[i], padding='same')
+            res.append(out)
+
+        point_out = torch.cat(res, dim=0)
+        point_out = self.norm(point_out)
+        return point_out
+
+
+class KernelUpdator(nn.Module):
+
+    def __init__(self,
+                 in_channels=256,
+                 feat_channels=64,
+                 out_channels=None,
+                 input_feat_shape=3,
+                 gate_sigmoid=True,
+                 gate_norm_act=False,
+                 activate_out=False,
+                 act_cfg=dict(type='ReLU', inplace=True),
+                 norm_cfg=dict(type='LN')):
+        super(KernelUpdator, self).__init__()
+        self.in_channels = in_channels
+        self.feat_channels = feat_channels
+        self.out_channels_raw = out_channels
+        self.gate_sigmoid = gate_sigmoid
+        self.gate_norm_act = gate_norm_act
+        self.activate_out = activate_out
+        if isinstance(input_feat_shape, int):
+            input_feat_shape = [input_feat_shape] * 2
+        self.input_feat_shape = input_feat_shape
+        self.act_cfg = act_cfg
+        self.norm_cfg = norm_cfg
+        self.out_channels = out_channels if out_channels else in_channels
+
+        self.num_params_in = self.feat_channels
+        self.num_params_out = self.feat_channels
+        self.dynamic_layer = nn.Linear(
+            self.in_channels, self.num_params_in + self.num_params_out)
+        self.input_layer = nn.Linear(self.in_channels,
+                                     self.num_params_in + self.num_params_out,
+                                     1)
+        self.input_gate = nn.Linear(self.in_channels, self.feat_channels, 1)
+        self.update_gate = nn.Linear(self.in_channels, self.feat_channels, 1)
+        if self.gate_norm_act:
+            self.gate_norm = build_norm_layer(norm_cfg, self.feat_channels)[1]
+
+        self.norm_in = build_norm_layer(norm_cfg, self.feat_channels)[1]
+        self.norm_out = build_norm_layer(norm_cfg, self.feat_channels)[1]
+        self.input_norm_in = build_norm_layer(norm_cfg, self.feat_channels)[1]
+        self.input_norm_out = build_norm_layer(norm_cfg, self.feat_channels)[1]
+
+        self.activation = build_activation_layer(act_cfg)
+
+        self.fc_layer = nn.Linear(self.feat_channels, self.out_channels, 1)
+        self.fc_norm = build_norm_layer(norm_cfg, self.out_channels)[1]
+
+    def forward(self, update_feature, input_feature):
+        """
+        Args:
+            update_feature (torch.Tensor): [bs, num_proposals, in_channels]
+            input_feature (torch.Tensor): [bs, num_proposals, in_channels]
+        """
+        bs, num_proposals, _ = update_feature.shape
+
+        parameters = self.dynamic_layer(update_feature)
+        param_in = parameters[..., :self.num_params_in]
+        param_out = parameters[..., -self.num_params_out:]
+
+        input_feats = self.input_layer(input_feature)
+        input_in = input_feats[..., :self.num_params_in]
+        input_out = input_feats[..., -self.num_params_out:]
+
+        gate_feats = input_in * param_in
+        if self.gate_norm_act:
+            gate_feats = self.activation(self.gate_norm(gate_feats))
+
+        input_gate = self.input_norm_in(self.input_gate(gate_feats))
+        update_gate = self.norm_in(self.update_gate(gate_feats))
+        if self.gate_sigmoid:
+            input_gate = input_gate.sigmoid()
+            update_gate = update_gate.sigmoid()
+        param_out = self.norm_out(param_out)
+        input_out = self.input_norm_out(input_out)
+
+        if self.activate_out:
+            param_out = self.activation(param_out)
+            input_out = self.activation(input_out)
+
+        # param_out has shape (batch_size, feat_channels, out_channels)
+        features = update_gate * param_out + input_gate * input_out
+
+        features = self.fc_layer(features)
+        features = self.fc_norm(features)
+        features = self.activation(features)
+
+        return features
+
+
+class CrossAttenHead(nn.Module):
+
+    def __init__(self,
+                 num_classes,
+                 in_channels,
+                 num_proposals,
+                 frozen_head=False,
+                 frozen_pred=False,
+                 with_iou_pred=False,
+                 sphere_cls=False,
+                 ov_classifier_name=None,
+                 num_cls_fcs=1,
+                 num_mask_fcs=1,
+                 conv_kernel_size_1d=3,
+                 conv_kernel_size_2d=1,
+                 use_kernel_updator=False):
+        super(CrossAttenHead, self).__init__()
+        self.sphere_cls = sphere_cls
+        self.with_iou_pred = with_iou_pred
+        self.frozen_head = frozen_head
+        self.frozen_pred = frozen_pred
+        self.num_cls_fcs = num_cls_fcs
+        self.num_mask_fcs = num_mask_fcs
+        self.num_classes = num_classes
+        self.conv_kernel_size_2d = conv_kernel_size_2d
+
+        self.hidden_dim = in_channels
+        self.feat_channels = in_channels
+        self.num_proposals = num_proposals
+        self.hard_mask_thr = 0.5
+        self.use_kernel_updator = use_kernel_updator
+        # assert use_kernel_updator
+        if use_kernel_updator:
+            self.kernel_update = KernelUpdator(
+                in_channels=256,
+                feat_channels=256,
+                out_channels=256,
+                input_feat_shape=3,
+                act_cfg=dict(type='ReLU', inplace=True),
+                norm_cfg=dict(type='LN')
+            )
+        else:
+            self.f_atten = DySepConvAtten(self.feat_channels, self.num_proposals, conv_kernel_size_1d)
+            self.f_dropout = nn.Dropout(0.0)
+            self.f_atten_norm = nn.LayerNorm(self.hidden_dim * self.conv_kernel_size_2d ** 2)
+            self.k_atten = DySepConvAtten(self.feat_channels, self.num_proposals, conv_kernel_size_1d)
+            self.k_dropout = nn.Dropout(0.0)
+            self.k_atten_norm = nn.LayerNorm(self.hidden_dim * self.conv_kernel_size_2d ** 2)
+
+        self.s_atten = nn.MultiheadAttention(embed_dim=self.hidden_dim *
+                                                       self.conv_kernel_size_2d ** 2,
+                                             num_heads=8,
+                                             dropout=0.0)
+        self.s_dropout = nn.Dropout(0.0)
+        self.s_atten_norm = nn.LayerNorm(self.hidden_dim * self.conv_kernel_size_2d ** 2)
+
+        self.ffn = FFN(self.hidden_dim, feedforward_channels=2048, num_fcs=2)
+        self.ffn_norm = nn.LayerNorm(self.hidden_dim)
+
+        self.cls_fcs = nn.ModuleList()
+        for _ in range(self.num_cls_fcs):
+            self.cls_fcs.append(nn.Linear(self.hidden_dim, self.hidden_dim, bias=False))
+            self.cls_fcs.append(nn.LayerNorm(self.hidden_dim))
+            self.cls_fcs.append(nn.ReLU(True))
+
+        if sphere_cls:
+            rank, world_size = get_dist_info()
+            if ov_classifier_name is None:
+                _dim = 1024  # temporally hard code
+                cls_embed = torch.empty(self.num_classes, _dim)
+                torch.nn.init.orthogonal_(cls_embed)
+                cls_embed = cls_embed[:, None]
+            else:
+                ov_path = os.path.join(os.path.expanduser('~/.cache/embd'), 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))
+
+            # background class
+            _dim = cls_embed.size(2)
+            _prototypes = cls_embed.size(1)
+            if rank == 0:
+                back_token = torch.zeros(1, _dim, dtype=torch.float32, device='cuda')
+            else:
+                back_token = torch.empty(1, _dim, dtype=torch.float32, device='cuda')
+            if world_size > 1:
+                dist.broadcast(back_token, src=0)
+            back_token = back_token.to(device='cpu')
+            cls_embed = torch.cat([
+                cls_embed, back_token.repeat(_prototypes, 1)[None]
+            ], dim=0)
+            self.register_buffer('fc_cls', cls_embed.permute(2, 0, 1).contiguous(), persistent=False)
+
+            # cls embd proj
+            cls_embed_dim = self.fc_cls.size(0)
+            self.cls_proj = nn.Sequential(
+                nn.Linear(self.hidden_dim, self.hidden_dim), nn.ReLU(inplace=True),
+                nn.Linear(self.hidden_dim, self.hidden_dim), nn.ReLU(inplace=True),
+                nn.Linear(self.hidden_dim, cls_embed_dim)
+            )
+
+            logit_scale = torch.tensor(4.6052, dtype=torch.float32)
+            self.register_buffer('logit_scale', logit_scale, persistent=False)
+        else:
+            self.fc_cls = nn.Linear(self.hidden_dim, self.num_classes + 1)
+
+        self.mask_fcs = nn.ModuleList()
+        for _ in range(self.num_mask_fcs):
+            self.mask_fcs.append(nn.Linear(self.hidden_dim, self.hidden_dim, bias=False))
+            self.mask_fcs.append(nn.LayerNorm(self.hidden_dim))
+            self.mask_fcs.append(nn.ReLU(True))
+        self.fc_mask = nn.Linear(self.hidden_dim, self.hidden_dim)
+
+        if self.with_iou_pred:
+            self.iou_embed = nn.Sequential(
+                nn.Linear(self.hidden_dim, self.hidden_dim),
+                nn.ReLU(inplace=True),
+                nn.Linear(self.hidden_dim, self.hidden_dim),
+                nn.ReLU(inplace=True),
+                nn.Linear(self.hidden_dim, 1),
+            )
+        prior_prob = 0.01
+        self.bias_value = -math.log((1 - prior_prob) / prior_prob)
+
+        self.apply(self._init_weights)
+        if not sphere_cls:
+            nn.init.constant_(self.fc_cls.bias, self.bias_value)
+
+        if self.frozen_head:
+            self._frozen_head()
+        if self.frozen_pred:
+            self._frozen_pred()
+
+    def _init_weights(self, m):
+        # print("init weights")
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def _frozen_head(self):
+        for n, p in self.kernel_update.named_parameters():
+            p.requires_grad = False
+        for n, p in self.s_atten.named_parameters():
+            p.requires_grad = False
+        for n, p in self.s_dropout.named_parameters():
+            p.requires_grad = False
+        for n, p in self.s_atten_norm.named_parameters():
+            p.requires_grad = False
+        for n, p in self.ffn.named_parameters():
+            p.requires_grad = False
+        for n, p in self.ffn_norm.named_parameters():
+            p.requires_grad = False
+
+    def _frozen_pred(self):
+        # frozen cls_fcs, fc_cls, mask_fcs, fc_mask
+        for n, p in self.cls_fcs.named_parameters():
+            p.requires_grad = False
+        for n, p in self.fc_cls.named_parameters():
+            p.requires_grad = False
+        for n, p in self.mask_fcs.named_parameters():
+            p.requires_grad = False
+        for n, p in self.fc_mask.named_parameters():
+            p.requires_grad = False
+
+    def train(self, mode):
+        super().train(mode)
+        if self.frozen_head:
+            self.kernel_update.eval()
+            self.s_atten.eval()
+            self.s_dropout.eval()
+            self.s_atten_norm.eval()
+            self.ffn.eval()
+            self.ffn_norm.eval()
+        if self.frozen_pred:
+            self.cls_fcs.eval()
+            self.fc_cls.eval()
+            self.mask_fcs.eval()
+            self.fc_mask.eval()
+
+    def forward(self, features, proposal_kernels, mask_preds, self_attn_mask=None):
+        B, C, H, W = features.shape
+
+        soft_sigmoid_masks = mask_preds.sigmoid()
+        nonzero_inds = soft_sigmoid_masks > self.hard_mask_thr
+        hard_sigmoid_masks = nonzero_inds.float()
+
+        # [B, N, C]
+        f = torch.einsum('bnhw,bchw->bnc', hard_sigmoid_masks, features)
+        # [B, N, C, K, K] -> [B, N, C * K * K]
+        num_proposals = proposal_kernels.shape[1]
+        k = proposal_kernels.view(B, num_proposals, -1)
+
+        # ----
+        if self.use_kernel_updator:
+            k = self.kernel_update(f, k)
+        else:
+            f_tmp = self.f_atten(k, f)
+            f = f + self.f_dropout(f_tmp)
+            f = self.f_atten_norm(f)
+
+            f_tmp = self.k_atten(k, f)
+            f = f + self.k_dropout(f_tmp)
+            k = self.k_atten_norm(f)
+
+        # [N, B, C]
+        k = k.permute(1, 0, 2)
+
+        k_tmp = self.s_atten(query=k, key=k, value=k, attn_mask=self_attn_mask)[0]
+        k = k + self.s_dropout(k_tmp)
+        k = self.s_atten_norm(k.permute(1, 0, 2))
+
+        obj_feat = self.ffn_norm(self.ffn(k))
+
+        cls_feat = obj_feat
+        mask_feat = obj_feat
+
+        for cls_layer in self.cls_fcs:
+            cls_feat = cls_layer(cls_feat)
+
+        if self.sphere_cls:
+            cls_embd = self.cls_proj(cls_feat)  # FIXME Too much cls linear (cls_fcs + cls_proj)
+            cls_score = torch.einsum('bnc,ckp->bnkp', F.normalize(cls_embd, dim=-1), self.fc_cls)
+            cls_score = cls_score.max(-1).values
+            cls_score = self.logit_scale.exp() * cls_score
+        else:
+            cls_score = self.fc_cls(cls_feat)
+        for reg_layer in self.mask_fcs:
+            mask_feat = reg_layer(mask_feat)
+        # [B, N, K * K, C] -> [B, N, C]
+        mask_kernels = self.fc_mask(mask_feat)
+
+        new_mask_preds = torch.einsum("bqc,bchw->bqhw", mask_kernels, features)
+        if self.with_iou_pred:
+            iou_pred = self.iou_embed(mask_feat)
+            iou_pred = iou_pred
+        else:
+            iou_pred = None
+        return cls_score, new_mask_preds, iou_pred, obj_feat

app/models/necks/__init__.py

@@ -0,0 +1 @@
+from .ramsam_neck import YOSONeck

app/models/necks/ramsam_neck.py

@@ -0,0 +1,196 @@
+import math
+import torch
+from torch import nn
+import torch.nn.functional as F
+from mmengine.model import kaiming_init
+from mmdet.registry import MODELS
+from mmcv.ops import DeformConv2d, ModulatedDeformConv2d
+
+
+class DeformLayer(nn.Module):
+
+    def __init__(self,
+                 in_planes,
+                 out_planes,
+                 deconv_kernel=4,
+                 deconv_stride=2,
+                 deconv_pad=1,
+                 deconv_out_pad=0,
+                 modulate_deform=True,
+                 num_groups=1,
+                 deform_num_groups=1,
+                 dilation=1):
+        super(DeformLayer, self).__init__()
+        self.deform_modulated = modulate_deform
+        if modulate_deform:
+            deform_conv_op = ModulatedDeformConv2d
+            offset_channels = 27
+        else:
+            deform_conv_op = DeformConv2d
+            offset_channels = 18
+
+        self.dcn_offset = nn.Conv2d(in_planes, offset_channels * deform_num_groups, kernel_size=3, stride=1, padding=1 * dilation, dilation=dilation)
+        self.dcn = deform_conv_op(in_planes, out_planes, kernel_size=3, stride=1, padding=1 * dilation, bias=False, groups=num_groups, dilation=dilation, deformable_groups=deform_num_groups)
+        for layer in [self.dcn]:
+            kaiming_init(layer)
+
+        nn.init.constant_(self.dcn_offset.weight, 0)
+        nn.init.constant_(self.dcn_offset.bias, 0)
+
+        # nn.GroupNorm(64, out_planes) # nn.BatchNorm2d(out_planes) #
+        self.dcn_bn = nn.SyncBatchNorm(out_planes)  
+        self.up_sample = nn.ConvTranspose2d(in_channels=out_planes, out_channels=out_planes, kernel_size=deconv_kernel, stride=deconv_stride, padding=deconv_pad, output_padding=deconv_out_pad, bias=False)
+        self._deconv_init()
+        # nn.GroupNorm(64, out_planes) # nn.BatchNorm2d(out_planes) #
+        self.up_bn = nn.SyncBatchNorm(out_planes)  
+        self.relu = nn.ReLU()
+
+    def forward(self, x):
+        out = x
+        if self.deform_modulated:
+            offset_mask = self.dcn_offset(out)
+            offset_x, offset_y, mask = torch.chunk(offset_mask, 3, dim=1)
+            offset = torch.cat((offset_x, offset_y), dim=1)
+            mask = mask.sigmoid()
+            out = self.dcn(out, offset, mask)
+        else:
+            offset = self.dcn_offset(out)
+            out = self.dcn(out, offset)
+        x = out
+
+        x = self.dcn_bn(x)
+        x = self.relu(x)
+        x = self.up_sample(x)
+        x = self.up_bn(x)
+        x = self.relu(x)
+        return x
+
+    def _deconv_init(self):
+        w = self.up_sample.weight.data
+        f = math.ceil(w.size(2) / 2)
+        c = (2 * f - 1 - f % 2) / (2. * f)
+        for i in range(w.size(2)):
+            for j in range(w.size(3)):
+                w[0, 0, i, j] = \
+                    (1 - math.fabs(i / f - c)) * (1 - math.fabs(j / f - c))
+        for c in range(1, w.size(0)):
+            w[c, 0, :, :] = w[0, 0, :, :]
+
+class LiteDeformConv(nn.Module):
+    def __init__(self, agg_dim, backbone_shape):
+        super(LiteDeformConv, self).__init__()
+        in_channels = []
+        out_channels = [agg_dim]
+        for feat in backbone_shape:
+            in_channels.append(feat)
+            out_channels.append(feat//2)
+        
+        self.lateral_conv0 = nn.Conv2d(in_channels=in_channels[-1], out_channels=out_channels[-1], kernel_size=1, stride=1, padding=0)
+
+        self.deform_conv1 = DeformLayer(in_planes=out_channels[-1], out_planes=out_channels[-2])
+
+        self.lateral_conv1 = nn.Conv2d(in_channels=in_channels[-2], out_channels=out_channels[-2], kernel_size=1, stride=1, padding=0)
+
+        self.deform_conv2 = DeformLayer(in_planes=out_channels[-2], out_planes=out_channels[-3])
+
+        self.lateral_conv2 = nn.Conv2d(in_channels=in_channels[-3], out_channels=out_channels[-3], kernel_size=1, stride=1, padding=0)
+
+        self.deform_conv3 = DeformLayer(in_planes=out_channels[-3], out_planes=out_channels[-4])
+
+        self.lateral_conv3 = nn.Conv2d(in_channels=in_channels[-4], out_channels=out_channels[-4], kernel_size=1, stride=1, padding=0)
+
+        # self.fuse_conv = nn.Conv2d(in_channels=sum(out_channels[1:]), out_channels=out_channels[-5], kernel_size=3, stride=1, padding=1)
+        self.output_conv = nn.Conv2d(in_channels=out_channels[-5], out_channels=out_channels[-5], kernel_size=3, stride=1, padding=1)
+
+        self.bias = nn.Parameter(torch.FloatTensor(1,out_channels[-5],1,1), requires_grad=True)
+        self.bias.data.fill_(0.0)
+
+        self.conv_a5 = nn.Conv2d(in_channels=out_channels[-1], out_channels=out_channels[-5], kernel_size=1, stride=1, padding=0, bias=False)
+        self.conv_a4 = nn.Conv2d(in_channels=out_channels[-2], out_channels=out_channels[-5], kernel_size=1, stride=1, padding=0, bias=False)
+        self.conv_a3 = nn.Conv2d(in_channels=out_channels[-3], out_channels=out_channels[-5], kernel_size=1, stride=1, padding=0, bias=False)
+        self.conv_a2 = nn.Conv2d(in_channels=out_channels[-4], out_channels=out_channels[-5], kernel_size=1, stride=1, padding=0, bias=False)
+
+    def forward(self, features_list):
+        p5 = self.lateral_conv0(features_list[-1])
+        x5 = p5
+        x = self.deform_conv1(x5)
+
+        p4 = self.lateral_conv1(features_list[-2])
+        x4 = p4 + x
+        x = self.deform_conv2(x4)
+
+        p3 = self.lateral_conv2(features_list[-3])
+        x3 = p3 + x
+        x = self.deform_conv3(x3)
+
+        p2 = self.lateral_conv3(features_list[-4])
+        x2 = p2 + x
+
+        # CFA
+        x5 = self.conv_a5(x5)
+        x4 = self.conv_a4(x4)
+        x3 = self.conv_a3(x3)
+
+        _x5 = F.interpolate(x5, scale_factor=8, align_corners=False, mode='bilinear')
+        _x4 = F.interpolate(x4, scale_factor=4, align_corners=False, mode='bilinear')
+        _x3 = F.interpolate(x3, scale_factor=2, align_corners=False, mode='bilinear')
+        x2 = self.conv_a2(x2)
+        x = _x5 + _x4 + _x3 + x2 + self.bias
+
+        x = self.output_conv(x)
+
+        return x, (x5, x4, x3)
+
+
+
+@MODELS.register_module()
+class YOSONeck(nn.Module):
+
+    def __init__(self,
+                 agg_dim,
+                 hidden_dim,
+                 backbone_shape,
+                 return_multi_scale=False,
+                 return_single_scale=False,
+                 #Just for compatibility with Mask2Former, not actually used
+                 in_channels=None,
+                 feat_channels=None,
+                 out_channels=None
+    ):
+        super().__init__()
+        # in_channels == backbone_shape
+        # hidden_dim == feat_channels == out_channels == 256
+        self.return_single_scale = return_single_scale
+        self.return_multi_scale = return_multi_scale
+        self.deconv = LiteDeformConv(agg_dim=agg_dim, backbone_shape=backbone_shape)
+
+        self.loc_conv = nn.Conv2d(in_channels=agg_dim + 2, out_channels=hidden_dim, kernel_size=1, stride=1)
+        self.init_weights()
+
+    def init_weights(self) -> None:
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def generate_coord(self, input_feat):
+        x_range = torch.linspace(-1, 1, input_feat.shape[-1], device=input_feat.device)
+        y_range = torch.linspace(-1, 1, input_feat.shape[-2], device=input_feat.device)
+        y, x = torch.meshgrid(y_range, x_range)
+        y = y.expand([input_feat.shape[0], 1, -1, -1])
+        x = x.expand([input_feat.shape[0], 1, -1, -1])
+        coord_feat = torch.cat([x, y], 1)
+        return coord_feat
+
+    def forward(self,
+                features_list,
+                batch_img_metas = None,
+                num_frames = None):
+        features, multi_scale = self.deconv(features_list)
+        coord_feat = self.generate_coord(features)
+        features = torch.cat([features, coord_feat], 1)
+        features = self.loc_conv(features)
+        if self.return_single_scale: # maskformer
+            return features, multi_scale[0]
+        if self.return_multi_scale: # mask2former
+            return features, multi_scale
+        return features

app/models/utils/__init__.py

@@ -0,0 +1,3 @@
+from .video_gt_preprocess import preprocess_video_panoptic_gt
+from .mask_pool import mask_pool
+from .no_obj import NO_OBJ

app/models/utils/load_checkpoint.py

@@ -0,0 +1,38 @@
+from mmengine.runner.checkpoint import CheckpointLoader
+
+
+def load_checkpoint_with_prefix(filename, prefix=None, map_location='cpu', logger='current'):
+    """Load partial pretrained model with specific prefix.
+
+    Args:
+        prefix (str): The prefix of sub-module.
+        filename (str): Accept local filepath, URL, ``torchvision://xxx``,
+            ``open-mmlab://xxx``. Please refer to ``docs/model_zoo.md`` for
+            details.
+        map_location (str | None): Same as :func:`torch.load`.
+            Defaults to None.
+        logger: logger
+
+    Returns:
+        dict or OrderedDict: The loaded checkpoint.
+    """
+
+    checkpoint = CheckpointLoader.load_checkpoint(filename, map_location=map_location, logger=logger)
+
+    if 'state_dict' in checkpoint:
+        state_dict = checkpoint['state_dict']
+    else:
+        state_dict = checkpoint
+    if not prefix:
+        return state_dict
+    if not prefix.endswith('.'):
+        prefix += '.'
+    prefix_len = len(prefix)
+
+    state_dict = {
+        k[prefix_len:]: v
+        for k, v in state_dict.items() if k.startswith(prefix)
+    }
+
+    assert state_dict, f'{prefix} is not in the pretrained model'
+    return state_dict

app/models/utils/mask_pool.py

@@ -0,0 +1,27 @@
+import torch
+import torch.nn.functional as F
+
+
+# https://github.com/NVlabs/ODISE/blob/e97b06c424c575fec9fc5368dd4b3e050d91abc4/odise/modeling/meta_arch/odise.py#L923
+
+def mask_pool(x, mask):
+    """
+    Args:
+        x: [B, C, H, W]
+        mask: [B, Q, H, W]
+    """
+    if not x.shape[-2:] == mask.shape[-2:]:
+        # reshape mask to x
+        mask = F.interpolate(mask, size=x.shape[-2:], mode='bilinear', align_corners=False)
+    with torch.no_grad():
+        mask = mask.detach()
+        mask = (mask > 0).to(mask.dtype)
+        denorm = mask.sum(dim=(-1, -2), keepdim=True) + 1e-8
+
+    mask_pooled_x = torch.einsum(
+        "bchw,bqhw->bqc",
+        x,
+        mask / denorm,
+    )
+    return mask_pooled_x
+

app/models/utils/no_obj.py

@@ -0,0 +1 @@
+NO_OBJ = 65535

app/models/utils/video_gt_preprocess.py

@@ -0,0 +1,87 @@
+import torch
+
+
+def preprocess_video_panoptic_gt(
+        gt_labels,
+        gt_masks,
+        gt_semantic_seg,
+        gt_instance_ids,
+        num_things,
+        num_stuff,
+):
+    num_classes = num_things + num_stuff
+    num_frames = len(gt_masks)
+    mask_size = gt_masks[0].masks.shape[-2:]
+
+    thing_masks_list = []
+    for frame_id in range(num_frames):
+        thing_masks_list.append(gt_masks[frame_id].pad(
+            mask_size, pad_val=0).to_tensor(
+            dtype=torch.bool, device=gt_labels.device)
+        )
+    instances = torch.unique(gt_instance_ids[:, 1])
+    things_masks = []
+    labels = []
+    for instance in instances:
+        pos_ins = torch.nonzero(torch.eq(gt_instance_ids[:, 1], instance), as_tuple=True)[0]  # 0 is for redundant tuple
+        labels_instance = gt_labels[:, 1][pos_ins]
+        assert torch.allclose(labels_instance, labels_instance[0])
+        labels.append(labels_instance[0])
+        instance_frame_ids = gt_instance_ids[:, 0][pos_ins].to(dtype=torch.int32).tolist()
+        instance_masks = []
+        for frame_id in range(num_frames):
+            frame_instance_ids = gt_instance_ids[gt_instance_ids[:, 0] == frame_id, 1]
+            if frame_id not in instance_frame_ids:
+                empty_mask = torch.zeros(
+                    mask_size,
+                    dtype=thing_masks_list[frame_id].dtype, device=thing_masks_list[frame_id].device
+                )
+                instance_masks.append(empty_mask)
+            else:
+                pos_inner_frame = torch.nonzero(torch.eq(frame_instance_ids, instance), as_tuple=True)[0].item()
+                frame_mask = thing_masks_list[frame_id][pos_inner_frame]
+                instance_masks.append(frame_mask)
+        things_masks.append(torch.stack(instance_masks))
+
+    if len(instances) == 0:
+        things_masks = torch.stack(thing_masks_list, dim=1)
+        labels = torch.empty_like(instances)
+    else:
+        things_masks = torch.stack(things_masks)
+        labels = torch.stack(labels)
+    assert torch.all(torch.less(labels, num_things))
+
+    if gt_semantic_seg is not None:
+        things_labels = labels
+        gt_semantic_seg = gt_semantic_seg.squeeze(1)
+
+        semantic_labels = torch.unique(
+            gt_semantic_seg,
+            sorted=False,
+            return_inverse=False,
+            return_counts=False)
+        stuff_masks_list = []
+        stuff_labels_list = []
+        for label in semantic_labels:
+            if label < num_things or label >= num_classes:
+                continue
+            stuff_mask = gt_semantic_seg == label
+            stuff_masks_list.append(stuff_mask)
+            stuff_labels_list.append(label)
+
+        if len(stuff_masks_list) > 0:
+            stuff_masks = torch.stack(stuff_masks_list, dim=0)
+            stuff_labels = torch.stack(stuff_labels_list, dim=0)
+            assert torch.all(torch.ge(stuff_labels, num_things)) and torch.all(torch.less(stuff_labels, num_classes))
+            labels = torch.cat([things_labels, stuff_labels], dim=0)
+            masks = torch.cat([things_masks, stuff_masks], dim=0)
+        else:
+            labels = things_labels
+            masks = things_masks
+        assert len(labels) == len(masks)
+    else:
+        masks = things_masks
+
+    labels = labels.to(dtype=torch.long)
+    masks = masks.to(dtype=torch.long)
+    return labels, masks

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
+    }
+}

ext/open_clip/model_configs/RN50x64.json

@@ -0,0 +1,21 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 448,
+        "layers": [
+            3,
+            15,
+            36,
+            10
+        ],
+        "width": 128,
+        "patch_size": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 1024,
+        "heads": 16,
+        "layers": 12
+    }
+}

ext/open_clip/model_configs/ViT-B-16-plus-240.json

@@ -0,0 +1,16 @@
+{
+    "embed_dim": 640,
+    "vision_cfg": {
+        "image_size": 240,
+        "layers": 12,
+        "width": 896,
+        "patch_size": 16
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 640,
+        "heads": 10,
+        "layers": 12
+    }
+}

ext/open_clip/model_configs/ViT-B-16-plus.json

@@ -0,0 +1,16 @@
+{
+    "embed_dim": 640,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 12,
+        "width": 896,
+        "patch_size": 16
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 640,
+        "heads": 10,
+        "layers": 12
+    }
+}

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

@@ -0,0 +1,16 @@
+{
+    "embed_dim": 512,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 12,
+        "width": 768,
+        "patch_size": 16
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}

ext/open_clip/model_configs/ViT-B-32-plus-256.json

@@ -0,0 +1,16 @@
+{
+    "embed_dim": 640,
+    "vision_cfg": {
+        "image_size": 256,
+        "layers": 12,
+        "width": 896,
+        "patch_size": 32
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 640,
+        "heads": 10,
+        "layers": 12
+    }
+}

ext/open_clip/model_configs/ViT-B-32-quickgelu.json

@@ -0,0 +1,17 @@
+{
+    "embed_dim": 512,
+    "quick_gelu": true,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 12,
+        "width": 768,
+        "patch_size": 32
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}

ext/open_clip/model_configs/ViT-B-32.json

@@ -0,0 +1,16 @@
+{
+    "embed_dim": 512,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 12,
+        "width": 768,
+        "patch_size": 32
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}

ext/open_clip/model_configs/ViT-H-14.json

@@ -0,0 +1,17 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 32,
+        "width": 1280,
+        "head_width": 80,
+        "patch_size": 14
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 1024,
+        "heads": 16,
+        "layers": 24
+    }
+}

ext/open_clip/model_configs/ViT-H-16.json

@@ -0,0 +1,17 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 32,
+        "width": 1280,
+        "head_width": 80,
+        "patch_size": 16
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 1024,
+        "heads": 16,
+        "layers": 24
+    }
+}