Upload folder using huggingface_hub

README.md

@@ -0,0 +1,39 @@
+---
+license: cc-by-nc-sa-4.0
+language:
+- en
+---
+# What's New
+
+[2024/10] A new [VideoCLIP-XL-v2](https://huggingface.co/alibaba-pai/VideoCLIP-XL-v2) model has been released.
+
+[2024/10] Initial commit for the [VideoCLIP-XL](https://huggingface.co/alibaba-pai/VideoCLIP-XL) model, the [VILD](https://huggingface.co/alibaba-pai/VILD) dataset, and the [LVDR](https://huggingface.co/alibaba-pai/LVDR) benchmark.
+
+# VideoCLIP-XL (eXtra Length)
+
+This model is proposed from [VideoCLIP-XL paper](https://arxiv.org/abs/2410.00741). 
+It aims to advance long description understanding for video CLIP Models.
+
+# Install
+~~~
+# 1. Create your environment
+# 2. Install torch
+# 3. Then:
+pip install -r requirements.txt
+~~~
+
+# Usage
+Please refer to ```demo.py```.
+
+# Source
+~~~
+@misc{wang2024videoclipxladvancinglongdescription,
+      title={VideoCLIP-XL: Advancing Long Description Understanding for Video CLIP Models}, 
+      author={Jiapeng Wang and Chengyu Wang and Kunzhe Huang and Jun Huang and Lianwen Jin},
+      year={2024},
+      eprint={2410.00741},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL},
+      url={https://arxiv.org/abs/2410.00741}, 
+}
+~~~

VideoCLIP-XL.bin

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

data_dataloaders.py

@@ -0,0 +1,57 @@
+import torch
+from torch.utils.data import DataLoader
+from dataloaders.dataloader_msvd_retrieval import MSVD_DataLoader
+
+
+def dataloader_msvd_train(args, tokenizer):
+    msvd_dataset = MSVD_DataLoader(
+        subset="train",
+        data_path=args.data_path,
+        features_path=args.features_path,
+        csv_path=args.train_csv,
+        max_words=args.max_words,
+        feature_framerate=args.feature_framerate,
+        tokenizer=tokenizer,
+        max_frames=args.max_frames,
+        frame_order=args.train_frame_order,
+        slice_framepos=args.slice_framepos,
+    )
+
+    train_sampler = torch.utils.data.distributed.DistributedSampler(msvd_dataset)
+    dataloader = DataLoader(
+        msvd_dataset,
+        batch_size=args.batch_size // args.n_gpu,
+        num_workers=args.num_thread_reader,
+        pin_memory=False,
+        shuffle=(train_sampler is None),
+        sampler=train_sampler,
+        drop_last=True,
+    )
+
+    return dataloader, len(msvd_dataset), train_sampler
+
+def dataloader_msvd_test(args, tokenizer, subset="test"):
+    msvd_testset = MSVD_DataLoader(
+        subset=subset,
+        data_path=args.data_path,
+        features_path=args.features_path,
+        csv_path=args.val_csv,
+        max_words=args.max_words,
+        feature_framerate=args.feature_framerate,
+        tokenizer=tokenizer,
+        max_frames=args.max_frames,
+        frame_order=args.eval_frame_order,
+        slice_framepos=args.slice_framepos,
+    )
+    dataloader_msrvtt = DataLoader(
+        msvd_testset,
+        batch_size=args.batch_size_val,
+        num_workers=args.num_thread_reader,
+        shuffle=False,
+        drop_last=False,
+    )
+    return dataloader_msrvtt, len(msvd_testset)
+
+
+DATALOADER_DICT = {}
+DATALOADER_DICT["msvd"] = {"train":dataloader_msvd_train, "val":dataloader_msvd_test, "test":dataloader_msvd_test}

dataloader_msvd_retrieval.py

@@ -0,0 +1,183 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import unicode_literals
+from __future__ import print_function
+
+import os
+from torch.utils.data import Dataset
+import numpy as np
+import pickle
+from dataloaders.tubeletvideo_util import TubeletVideoExtractor
+
+class MSVD_DataLoader(Dataset):
+    """MSVD dataset loader."""
+    def __init__(
+            self,
+            subset,
+            data_path,
+            features_path,
+            csv_path,
+            tokenizer,
+            max_words=30,
+            feature_framerate=1.0,
+            max_frames=100,
+            image_resolution=224,
+            frame_order=0,
+            slice_framepos=0,
+    ):
+        self.data_path = data_path
+        self.features_path = features_path
+        self.feature_framerate = feature_framerate
+        self.max_words = max_words
+        self.max_frames = max_frames
+        self.tokenizer = tokenizer
+        # 0: ordinary order; 1: reverse order; 2: random order.
+        self.frame_order = frame_order
+        assert self.frame_order in [0, 1, 2]
+        # 0: cut from head frames; 1: cut from tail frames; 2: extract frames uniformly.
+        self.slice_framepos = slice_framepos
+        assert self.slice_framepos in [0, 1, 2]
+
+        self.subset = subset
+        assert self.subset in ["train", "val", "test"]
+        video_id_path_dict = {}
+        video_id_path_dict["train"] = os.path.join(self.data_path, "train_list.txt")
+        video_id_path_dict["val"] = os.path.join(self.data_path, "val_list.txt")
+        video_id_path_dict["test"] = os.path.join(self.data_path, "test_list.txt")
+        caption_file = os.path.join(self.data_path, "raw-captions.pkl")
+
+        with open(video_id_path_dict[self.subset], 'r') as fp:
+            video_ids = [itm.strip() for itm in fp.readlines()]
+
+        with open(caption_file, 'rb') as f:
+            captions = pickle.load(f)
+
+        video_dict = {}
+        for root, dub_dir, video_files in os.walk(self.features_path):
+            for video_file in video_files:
+                video_id_ = ".".join(video_file.split(".")[:-1])
+                if video_id_ not in video_ids:
+                    continue
+                file_path_ = os.path.join(root, video_file)
+                video_dict[video_id_] = file_path_
+        self.video_dict = video_dict
+
+        self.sample_len = 0
+        self.sentences_dict = {}
+        self.cut_off_points = []
+        for video_id in video_ids:
+            assert video_id in captions
+            for cap in captions[video_id]:
+                cap_txt = " ".join(cap)
+                self.sentences_dict[len(self.sentences_dict)] = (video_id, cap_txt)
+            self.cut_off_points.append(len(self.sentences_dict))
+
+        ## below variables are used to multi-sentences retrieval
+        # self.cut_off_points: used to tag the label when calculate the metric
+        # self.sentence_num: used to cut the sentence representation
+        # self.video_num: used to cut the video representation
+        self.multi_sentence_per_video = True    # !!! important tag for eval
+        if self.subset == "val" or self.subset == "test":
+            self.sentence_num = len(self.sentences_dict)
+            self.video_num = len(video_ids)
+            assert len(self.cut_off_points) == self.video_num
+            print("For {}, sentence number: {}".format(self.subset, self.sentence_num))
+            print("For {}, video number: {}".format(self.subset, self.video_num))
+
+        print("Video number: {}".format(len(self.video_dict)))
+        print("Total Paire: {}".format(len(self.sentences_dict)))
+
+        self.sample_len = len(self.sentences_dict)
+        self.rawVideoExtractor = TubeletVideoExtractor(
+            csv_path=csv_path, 
+            framerate=feature_framerate, size=image_resolution)
+        self.SPECIAL_TOKEN = {"CLS_TOKEN": "<|startoftext|>", "SEP_TOKEN": "<|endoftext|>",
+                              "MASK_TOKEN": "[MASK]", "UNK_TOKEN": "[UNK]", "PAD_TOKEN": "[PAD]"}
+
+    def __len__(self):
+        return self.sample_len
+
+    def _get_text(self, video_id, caption):
+        k = 1
+        choice_video_ids = [video_id]
+        pairs_text = np.zeros((k, self.max_words), dtype=np.long)
+        pairs_mask = np.zeros((k, self.max_words), dtype=np.long)
+        pairs_segment = np.zeros((k, self.max_words), dtype=np.long)
+
+        for i, video_id in enumerate(choice_video_ids):
+            words = self.tokenizer.tokenize(caption)
+
+            words = [self.SPECIAL_TOKEN["CLS_TOKEN"]] + words
+            total_length_with_CLS = self.max_words - 1
+            if len(words) > total_length_with_CLS:
+                words = words[:total_length_with_CLS]
+            words = words + [self.SPECIAL_TOKEN["SEP_TOKEN"]]
+
+            input_ids = self.tokenizer.convert_tokens_to_ids(words)
+            input_mask = [1] * len(input_ids)
+            segment_ids = [0] * len(input_ids)
+            while len(input_ids) < self.max_words:
+                input_ids.append(0)
+                input_mask.append(0)
+                segment_ids.append(0)
+            assert len(input_ids) == self.max_words
+            assert len(input_mask) == self.max_words
+            assert len(segment_ids) == self.max_words
+
+            pairs_text[i] = np.array(input_ids)
+            pairs_mask[i] = np.array(input_mask)
+            pairs_segment[i] = np.array(segment_ids)
+
+        return pairs_text, pairs_mask, pairs_segment, choice_video_ids
+
+    def _get_rawvideo(self, choice_video_ids):
+        video_mask = np.zeros((len(choice_video_ids), self.max_frames), dtype=np.long)
+        max_video_length = [0] * len(choice_video_ids)
+
+        # Pair x L x T x 3 x H x W
+        video = np.zeros((len(choice_video_ids), self.max_frames, 1, 3,
+                          self.rawVideoExtractor.size, self.rawVideoExtractor.size), dtype=np.float)
+
+        for i, video_id in enumerate(choice_video_ids):
+            video_path = self.video_dict[video_id]
+
+            raw_video_data = self.rawVideoExtractor.get_video_data(video_path)
+            raw_video_data = raw_video_data['video']
+
+            if len(raw_video_data.shape) > 3:
+                raw_video_data_clip = raw_video_data
+                # L x T x 3 x H x W
+                raw_video_slice = self.rawVideoExtractor.process_raw_data(raw_video_data_clip)
+                if self.max_frames < raw_video_slice.shape[0]:
+                    if self.slice_framepos == 0:
+                        video_slice = raw_video_slice[:self.max_frames, ...]
+                    elif self.slice_framepos == 1:
+                        video_slice = raw_video_slice[-self.max_frames:, ...]
+                    else:
+                        sample_indx = np.linspace(0, raw_video_slice.shape[0] - 1, num=self.max_frames, dtype=int)
+                        video_slice = raw_video_slice[sample_indx, ...]
+                else:
+                    video_slice = raw_video_slice
+
+                video_slice = self.rawVideoExtractor.process_frame_order(video_slice, frame_order=self.frame_order)
+
+                slice_len = video_slice.shape[0]
+                max_video_length[i] = max_video_length[i] if max_video_length[i] > slice_len else slice_len
+                if slice_len < 1:
+                    pass
+                else:
+                    video[i][:slice_len, ...] = video_slice
+            else:
+                print("video path: {} error. video id: {}".format(video_path, video_id))
+
+        for i, v_length in enumerate(max_video_length):
+            video_mask[i][:v_length] = [1] * v_length
+
+        return video, video_mask
+
+    def __getitem__(self, idx):
+        video_id, caption = self.sentences_dict[idx]
+
+        pairs_text, pairs_mask, pairs_segment, choice_video_ids = self._get_text(video_id, caption)
+        video, video_mask = self._get_rawvideo(choice_video_ids)
+        return pairs_text, pairs_mask, pairs_segment, video, video_mask

demo.py

@@ -0,0 +1,435 @@
+import logging
+import os
+from typing import List
+import cv2
+import numpy as np
+import torch
+from torch import nn
+from PIL import Image
+import torch.nn.functional as F
+from tqdm import tqdm
+import threading
+from torch._utils import ExceptionWrapper
+
+from utils.text_encoder.simple_tokenizer import SimpleTokenizer as ClipTokenizer
+from modeling import VideoCLIP_XL
+from utils.text_encoder import text_encoder
+import argparse
+from data_dataloaders import DATALOADER_DICT
+
+args_parser = argparse.ArgumentParser()
+args_parser.add_argument("--datatype", type=str, default="msvd", help="dataset name")
+args_parser.add_argument("--local-rank", default=0, type=int, help="distribted training")
+
+args = args_parser.parse_args()
+
+logger = logging.getLogger(__name__)
+
+
+def get_a_var(obj):
+    if isinstance(obj, torch.Tensor):
+        return obj
+
+    if isinstance(obj, list) or isinstance(obj, tuple):
+        for result in map(get_a_var, obj):
+            if isinstance(result, torch.Tensor):
+                return result
+    if isinstance(obj, dict):
+        for result in map(get_a_var, obj.items()):
+            if isinstance(result, torch.Tensor):
+                return result
+    return None
+
+
+def parallel_apply(fct, model, inputs, device_ids):
+    modules = nn.parallel.replicate(model, device_ids)
+    assert len(modules) == len(inputs)
+    lock = threading.Lock()
+    results = {}
+    grad_enabled = torch.is_grad_enabled()
+
+    def _worker(i, module, input):
+        torch.set_grad_enabled(grad_enabled)
+        device = get_a_var(input).get_device()
+        try:
+            with torch.cuda.device(device):
+                # this also avoids accidental slicing of `input` if it is a Tensor
+                if not isinstance(input, (list, tuple)):
+                    input = (input,)
+                output = fct(module, *input)
+            with lock:
+                results[i] = output
+        except Exception:
+            with lock:
+                results[i] = ExceptionWrapper(where="in replica {} on device {}".format(i, device))
+
+    if len(modules) > 1:
+        threads = [threading.Thread(target=_worker, args=(i, module, input))
+                   for i, (module, input) in enumerate(zip(modules, inputs))]
+
+        for thread in threads:
+            thread.start()
+        for thread in threads:
+            thread.join()
+    else:
+        _worker(0, modules[0], inputs[0])
+
+    outputs = []
+    for i in range(len(inputs)):
+        output = results[i]
+        if isinstance(output, ExceptionWrapper):
+            output.reraise()
+        outputs.append(output)
+    return outputs
+
+
+def parallel_apply(fct, model, inputs, device_ids):
+    modules = nn.parallel.replicate(model, device_ids)
+    assert len(modules) == len(inputs)
+    lock = threading.Lock()
+    results = {}
+    grad_enabled = torch.is_grad_enabled()
+
+    def _worker(i, module, input):
+        torch.set_grad_enabled(grad_enabled)
+        device = get_a_var(input).get_device()
+        try:
+            with torch.cuda.device(device):
+                # this also avoids accidental slicing of `input` if it is a Tensor
+                if not isinstance(input, (list, tuple)):
+                    input = (input,)
+                output = fct(module, *input)
+            with lock:
+                results[i] = output
+        except Exception:
+            with lock:
+                results[i] = ExceptionWrapper(where="in replica {} on device {}".format(i, device))
+
+    if len(modules) > 1:
+        threads = [threading.Thread(target=_worker, args=(i, module, input))
+                   for i, (module, input) in enumerate(zip(modules, inputs))]
+
+        for thread in threads:
+            thread.start()
+        for thread in threads:
+            thread.join()
+    else:
+        _worker(0, modules[0], inputs[0])
+
+    outputs = []
+    for i in range(len(inputs)):
+        output = results[i]
+        if isinstance(output, ExceptionWrapper):
+            output.reraise()
+        outputs.append(output)
+    return outputs
+
+
+def tensor_video_to_text_sim(sim_tensor):
+    if not torch.is_tensor(sim_tensor):
+      sim_tensor = torch.tensor(sim_tensor)
+    # Code to avoid nans
+    sim_tensor[sim_tensor != sim_tensor] = float('-inf')
+    # Forms a similarity matrix for use with rank at k
+    values, _ = torch.max(sim_tensor, dim=1, keepdim=True)
+    return torch.squeeze(values).T
+
+
+def tensor_text_to_video_metrics(sim_tensor, top_k = [1,5,10]):
+    if not torch.is_tensor(sim_tensor):
+      sim_tensor = torch.tensor(sim_tensor)
+
+    # Permute sim_tensor so it represents a sequence of text-video similarity matrices.
+    # Then obtain the double argsort to position the rank on the diagonal
+    stacked_sim_matrices = sim_tensor.permute(1, 0, 2)
+    first_argsort = torch.argsort(stacked_sim_matrices, dim = -1, descending= True)
+    second_argsort = torch.argsort(first_argsort, dim = -1, descending= False)
+
+    # Extracts ranks i.e diagonals
+    ranks = torch.flatten(torch.diagonal(second_argsort, dim1 = 1, dim2 = 2))
+
+    # Now we need to extract valid ranks, as some belong to inf padding values
+    permuted_original_data = torch.flatten(torch.diagonal(sim_tensor, dim1 = 0, dim2 = 2))
+    mask = ~ torch.logical_or(torch.isinf(permuted_original_data), torch.isnan(permuted_original_data))
+    valid_ranks = ranks[mask]
+    # A quick dimension check validates our results, there may be other correctness tests pending
+    # Such as dot product localization, but that is for other time.
+    #assert int(valid_ranks.shape[0]) ==  sum([len(text_dict[k]) for k in text_dict])
+    if not torch.is_tensor(valid_ranks):
+      valid_ranks = torch.tensor(valid_ranks)
+    results = {f"R{k}": float(torch.sum(valid_ranks < k) * 100 / len(valid_ranks)) for k in top_k}
+    results["MedianR"] = float(torch.median(valid_ranks + 1))
+    results["MeanR"] = float(np.mean(valid_ranks.numpy() + 1))
+    results["Std_Rank"] = float(np.std(valid_ranks.numpy() + 1))
+    results['MR'] = results["MedianR"]
+    return results
+
+
+def _run_on_single_gpu(model, batch_list_t, batch_list_v, batch_sequence_output_list, batch_visual_output_list):
+    sim_matrix = []
+    for idx1, b1 in enumerate(batch_list_t):
+        input_mask, segment_ids, *_tmp = b1
+        sequence_output = batch_sequence_output_list[idx1]
+        each_row = []
+        for idx2, b2 in enumerate(batch_list_v):
+            video_mask, *_tmp = b2
+            visual_output = batch_visual_output_list[idx2]
+            b1b2_logits, *_tmp = model.get_similarity_logits(sequence_output, visual_output, input_mask, video_mask,
+                                                                     loose_type=model.loose_type)
+            b1b2_logits = b1b2_logits.cpu().detach().numpy()
+            each_row.append(b1b2_logits)
+        each_row = np.concatenate(tuple(each_row), axis=-1)
+        sim_matrix.append(each_row)
+    return sim_matrix
+
+
+def compute_metrics(x):
+    sx = np.sort(-x, axis=1)
+    d = np.diag(-x)
+    d = d[:, np.newaxis]
+    ind = sx - d
+    ind = np.where(ind == 0)
+    ind = ind[1]
+    metrics = {}
+    metrics['R1'] = float(np.sum(ind == 0)) * 100 / len(ind)
+    metrics['R5'] = float(np.sum(ind < 5)) * 100 / len(ind)
+    metrics['R10'] = float(np.sum(ind < 10)) * 100 / len(ind)
+    metrics['MR'] = np.median(ind) + 1
+    metrics["MedianR"] = metrics['MR']
+    metrics["MeanR"] = np.mean(ind) + 1
+    metrics["cols"] = [int(i) for i in list(ind)]
+    return metrics
+
+
+def eval_epoch(args, model, test_dataloader, device, n_gpu):
+
+    if hasattr(model, 'module'):
+        model = model.module.to(device)
+    else:
+        model = model.to(device)
+
+    # #################################################################
+    ## below variables are used to multi-sentences retrieval
+    # multi_sentence_: important tag for eval
+    # cut_off_points: used to tag the label when calculate the metric
+    # sentence_num: used to cut the sentence representation
+    # video_num: used to cut the video representation
+    # #################################################################
+    multi_sentence_ = False
+    cut_off_points_, sentence_num_, video_num_ = [], -1, -1
+    if hasattr(test_dataloader.dataset, 'multi_sentence_per_video') \
+            and test_dataloader.dataset.multi_sentence_per_video:
+        multi_sentence_ = True
+        cut_off_points_ = test_dataloader.dataset.cut_off_points
+        sentence_num_ = test_dataloader.dataset.sentence_num
+        video_num_ = test_dataloader.dataset.video_num
+        cut_off_points_ = [itm - 1 for itm in cut_off_points_]
+
+    if multi_sentence_:
+        logger.warning("Eval under the multi-sentence per video clip setting.")
+        logger.warning("sentence num: {}, video num: {}".format(sentence_num_, video_num_))
+
+    model.eval()
+    with torch.no_grad():
+        batch_list_t = []
+        batch_list_v = []
+        batch_sequence_output_list, batch_visual_output_list = [], []
+        total_video_num = 0
+
+        # ----------------------------
+        # 1. cache the features
+        # ----------------------------
+        for bid, batch in enumerate(tqdm(test_dataloader)):
+            batch = tuple(t.to(device) for t in batch)
+            input_ids, input_mask, segment_ids, video, video_mask = batch
+
+            if multi_sentence_:
+                # multi-sentences retrieval means: one clip has two or more descriptions.
+                b, *_t = video.shape
+                sequence_output = model.get_sequence_output(input_ids, segment_ids, input_mask)
+                batch_sequence_output_list.append(sequence_output)
+                batch_list_t.append((input_mask, segment_ids,))
+
+                s_, e_ = total_video_num, total_video_num + b
+                filter_inds = [itm - s_ for itm in cut_off_points_ if itm >= s_ and itm < e_]
+
+                if len(filter_inds) > 0:
+                    video, video_mask = video[filter_inds, ...], video_mask[filter_inds, ...]
+                    visual_output = model.get_visual_output(video, video_mask)
+                    batch_visual_output_list.append(visual_output)
+                    batch_list_v.append((video_mask,))
+                total_video_num += b
+            else:
+                sequence_output, visual_output = model.get_sequence_visual_output(input_ids, segment_ids, input_mask, video, video_mask)
+
+                batch_sequence_output_list.append(sequence_output)
+                batch_list_t.append((input_mask, segment_ids,))
+
+                batch_visual_output_list.append(visual_output)
+                batch_list_v.append((video_mask,))
+
+            print("{}/{}\r".format(bid, len(test_dataloader)), end="")
+
+        # ----------------------------------
+        # 2. calculate the similarity
+        # ----------------------------------
+        if n_gpu > 1:
+            device_ids = list(range(n_gpu))
+            batch_list_t_splits = []
+            batch_list_v_splits = []
+            batch_t_output_splits = []
+            batch_v_output_splits = []
+            bacth_len = len(batch_list_t)
+            split_len = (bacth_len + n_gpu - 1) // n_gpu
+            for dev_id in device_ids:
+                s_, e_ = dev_id * split_len, (dev_id + 1) * split_len
+                if dev_id == 0:
+                    batch_list_t_splits.append(batch_list_t[s_:e_])
+                    batch_list_v_splits.append(batch_list_v)
+
+                    batch_t_output_splits.append(batch_sequence_output_list[s_:e_])
+                    batch_v_output_splits.append(batch_visual_output_list)
+                else:
+                    devc = torch.device('cuda:{}'.format(str(dev_id)))
+                    devc_batch_list = [tuple(t.to(devc) for t in b) for b in batch_list_t[s_:e_]]
+                    batch_list_t_splits.append(devc_batch_list)
+                    devc_batch_list = [tuple(t.to(devc) for t in b) for b in batch_list_v]
+                    batch_list_v_splits.append(devc_batch_list)
+
+                    devc_batch_list = [b.to(devc) for b in batch_sequence_output_list[s_:e_]]
+                    batch_t_output_splits.append(devc_batch_list)
+                    devc_batch_list = [b.to(devc) for b in batch_visual_output_list]
+                    batch_v_output_splits.append(devc_batch_list)
+
+            parameters_tuple_list = [(batch_list_t_splits[dev_id], batch_list_v_splits[dev_id],
+                                      batch_t_output_splits[dev_id], batch_v_output_splits[dev_id]) for dev_id in device_ids]
+            parallel_outputs = parallel_apply(_run_on_single_gpu, model, parameters_tuple_list, device_ids)
+            sim_matrix = []
+            for idx in range(len(parallel_outputs)):
+                sim_matrix += parallel_outputs[idx]
+            sim_matrix = np.concatenate(tuple(sim_matrix), axis=0)
+        else:
+            sim_matrix = _run_on_single_gpu(model, batch_list_t, batch_list_v, batch_sequence_output_list, batch_visual_output_list)
+            sim_matrix = np.concatenate(tuple(sim_matrix), axis=0)
+
+    if multi_sentence_:
+        logger.info("before reshape, sim matrix size: {} x {}".format(sim_matrix.shape[0], sim_matrix.shape[1]))
+        cut_off_points2len_ = [itm + 1 for itm in cut_off_points_]
+        max_length = max([e_-s_ for s_, e_ in zip([0]+cut_off_points2len_[:-1], cut_off_points2len_)])
+        sim_matrix_new = []
+        for s_, e_ in zip([0] + cut_off_points2len_[:-1], cut_off_points2len_):
+            sim_matrix_new.append(np.concatenate((sim_matrix[s_:e_],
+                                                  np.full((max_length-e_+s_, sim_matrix.shape[1]), -np.inf)), axis=0))
+        sim_matrix = np.stack(tuple(sim_matrix_new), axis=0)
+        logger.info("after reshape, sim matrix size: {} x {} x {}".
+                    format(sim_matrix.shape[0], sim_matrix.shape[1], sim_matrix.shape[2]))
+
+        tv_metrics = tensor_text_to_video_metrics(sim_matrix)
+        vt_metrics = compute_metrics(tensor_video_to_text_sim(sim_matrix))
+    else:
+        logger.info("sim matrix size: {}, {}".format(sim_matrix.shape[0], sim_matrix.shape[1]))
+        tv_metrics = compute_metrics(sim_matrix)
+        vt_metrics = compute_metrics(sim_matrix.T)
+        logger.info('\t Length-T: {}, Length-V:{}'.format(len(sim_matrix), len(sim_matrix[0])))
+
+    logger.info("Text-to-Video:")
+    logger.info('\t>>>  R@1: {:.1f} - R@5: {:.1f} - R@10: {:.1f} - Median R: {:.1f} - Mean R: {:.1f}'.
+                format(tv_metrics['R1'], tv_metrics['R5'], tv_metrics['R10'], tv_metrics['MR'], tv_metrics['MeanR']))
+    logger.info("Video-to-Text:")
+    logger.info('\t>>>  V2T$R@1: {:.1f} - V2T$R@5: {:.1f} - V2T$R@10: {:.1f} - V2T$Median R: {:.1f} - V2T$Mean R: {:.1f}'.
+                format(vt_metrics['R1'], vt_metrics['R5'], vt_metrics['R10'], vt_metrics['MR'], vt_metrics['MeanR']))
+
+    R1 = tv_metrics['R1']
+    return R1
+
+
+def _frame_from_video(video):
+    while video.isOpened():
+        success, frame = video.read()
+        if success:
+            yield frame
+        else:
+            break
+        
+        
+v_mean = np.array([0.485, 0.456, 0.406]).reshape(1,1,3)
+v_std = np.array([0.229, 0.224, 0.225]).reshape(1,1,3)
+def normalize(data):
+    return (data / 255.0 - v_mean) / v_std
+        
+    
+def video_preprocessing(video_path, fnum=8):
+    video = cv2.VideoCapture(video_path)
+    frames = [x for x in _frame_from_video(video)]
+    step = len(frames) // fnum
+    frames = frames[::step][:fnum]
+
+    vid_tube = []
+    for fr in frames:
+        fr = fr[:,:,::-1]
+        fr = cv2.resize(fr, (224, 224))
+        fr = np.expand_dims(normalize(fr), axis=(0, 1))
+        vid_tube.append(fr) 
+    vid_tube = np.concatenate(vid_tube, axis=1)
+    vid_tube = np.transpose(vid_tube, (0, 1, 4, 2, 3))
+    vid_tube = torch.from_numpy(vid_tube)
+    
+    return vid_tube
+
+
+videoclip_xl = VideoCLIP_XL()
+state_dict = torch.load("./VideoCLIP-XL.bin", map_location="cpu")
+videoclip_xl.load_state_dict(state_dict)
+videoclip_xl.cuda().eval()
+
+       
+videos = [
+    "/path/to/video-1.mp4",
+    "/path/to/video-2.mp4",
+]
+
+texts = [
+    "text-1",
+    "text-2",
+    "text-3",
+]
+
+# with torch.no_grad():
+#     video_inputs = torch.cat([video_preprocessing(video) for video in videos], 0).float().cuda()
+#     video_features = videoclip_xl.vision_model.get_vid_features(video_inputs).float()
+#     video_features = video_features / video_features.norm(dim=-1, keepdim=True)
+
+#     text_inputs = text_encoder.tokenize(texts, truncate=True).cuda()
+#     text_features = videoclip_xl.text_model.encode_text(text_inputs).float()
+#     text_features = text_features / text_features.norm(dim=-1, keepdim=True)
+
+# Tmp = 100.
+
+# sim_matrix = (text_features @ video_features.T) * Tmp
+
+# print(f"{type(sim_matrix)=}")
+
+# tv_metrics = compute_metrics(sim_matrix)
+
+# print("Text-to-Video:")
+# print(f'\t>>>  R@1: {tv_metrics['R1']:.1f} - R@5: {tv_metrics['R5']:.1f} - R@10: {tv_metrics['R10']:.1f} - Median R: {tv_metrics['MR']:.1f} - Mean R: {tv_metrics['MeanR']:.1f}')
+
+
+tokenizer = ClipTokenizer()
+
+test_dataloader, test_length = None, 0
+if DATALOADER_DICT[args.datatype]["test"] is not None:
+    test_dataloader, test_length = DATALOADER_DICT[args.datatype]["test"](args, tokenizer)
+
+if DATALOADER_DICT[args.datatype]["val"] is not None:
+    val_dataloader, val_length = DATALOADER_DICT[args.datatype]["val"](args, tokenizer, subset="val")
+else:
+    val_dataloader, val_length = test_dataloader, test_length
+
+
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu", args.local_rank)
+n_gpu = torch.cuda.device_count()
+
+if args.local_rank == 0:
+    eval_epoch(args, videoclip_xl, test_dataloader, device, n_gpu)

modeling.py

@@ -0,0 +1,586 @@
+import os
+from typing import List
+
+import cv2
+import numpy as np
+import torch
+import torch.nn as nn
+import logging
+from PIL import Image
+
+from utils.text_encoder import text_encoder
+from utils.vision_encoder import get_vision_encoder
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from modules.until_module import PreTrainedModel, AllGather, CrossEn
+from modules.module_cross import CrossModel, CrossConfig, Transformer as TransformerClip
+
+from modules.module_clip import CLIP, convert_weights
+from torch.nn.utils.rnn import pad_packed_sequence, pack_padded_sequence
+
+logger = logging.getLogger(__name__)
+allgather = AllGather.apply
+
+class CLIP4ClipPreTrainedModel(PreTrainedModel, nn.Module):
+    """ An abstract class to handle weights initialization and
+        a simple interface for dowloading and loading pretrained models.
+    """
+    def __init__(self, cross_config, *inputs, **kwargs):
+        super(CLIP4ClipPreTrainedModel, self).__init__(cross_config)
+        self.cross_config = cross_config
+        self.clip = None
+        self.cross = None
+
+    @classmethod
+    def from_pretrained(cls, cross_model_name, state_dict=None, cache_dir=None, type_vocab_size=2, *inputs, **kwargs):
+
+        task_config = None
+        if "task_config" in kwargs.keys():
+            task_config = kwargs["task_config"]
+            if not hasattr(task_config, "local_rank"):
+                task_config.__dict__["local_rank"] = 0
+            elif task_config.local_rank == -1:
+                task_config.local_rank = 0
+
+        if state_dict is None: state_dict = {}
+        pretrained_clip_name = "ViT-B/32"
+        if hasattr(task_config, 'pretrained_clip_name'):
+            pretrained_clip_name = task_config.pretrained_clip_name
+        clip_state_dict = CLIP.get_config(pretrained_clip_name=pretrained_clip_name)
+        for key, val in clip_state_dict.items():
+            new_key = "clip." + key
+            if new_key not in state_dict:
+                state_dict[new_key] = val.clone()
+
+        cross_config, _ = CrossConfig.get_config(cross_model_name, cache_dir, type_vocab_size, state_dict=None, task_config=task_config)
+
+        model = cls(cross_config, clip_state_dict, *inputs, **kwargs)
+
+        ## ===> Initialization trick [HARD CODE]
+        if model.linear_patch == "3d":
+            contain_conv2 = False
+            for key in state_dict.keys():
+                if key.find("visual.conv2.weight") > -1:
+                    contain_conv2 = True
+                    break
+            if contain_conv2 is False and hasattr(model.clip.visual, "conv2"):
+                cp_weight = state_dict["clip.visual.conv1.weight"].clone()
+                kernel_size = model.clip.visual.conv2.weight.size(2)
+                conv2_size = model.clip.visual.conv2.weight.size()
+                conv2_size = list(conv2_size)
+
+                left_conv2_size = conv2_size.copy()
+                right_conv2_size = conv2_size.copy()
+                left_conv2_size[2] = (kernel_size - 1) // 2
+                right_conv2_size[2] = kernel_size - 1 - left_conv2_size[2]
+
+                left_zeros, right_zeros = None, None
+                if left_conv2_size[2] > 0:
+                    left_zeros = torch.zeros(*tuple(left_conv2_size), dtype=cp_weight.dtype, device=cp_weight.device)
+                if right_conv2_size[2] > 0:
+                    right_zeros = torch.zeros(*tuple(right_conv2_size), dtype=cp_weight.dtype, device=cp_weight.device)
+
+                cat_list = []
+                if left_zeros != None: cat_list.append(left_zeros)
+                cat_list.append(cp_weight.unsqueeze(2))
+                if right_zeros != None: cat_list.append(right_zeros)
+                cp_weight = torch.cat(cat_list, dim=2)
+
+                state_dict["clip.visual.conv2.weight"] = cp_weight
+
+        if model.sim_header == 'tightTransf':
+            contain_cross = False
+            for key in state_dict.keys():
+                if key.find("cross.transformer") > -1:
+                    contain_cross = True
+                    break
+            if contain_cross is False:
+                for key, val in clip_state_dict.items():
+                    if key == "positional_embedding":
+                        state_dict["cross.embeddings.position_embeddings.weight"] = val.clone()
+                        continue
+                    if key.find("transformer.resblocks") == 0:
+                        num_layer = int(key.split(".")[2])
+
+                        # cut from beginning
+                        if num_layer < task_config.cross_num_hidden_layers:
+                            state_dict["cross."+key] = val.clone()
+                            continue
+
+        if model.sim_header == "seqLSTM" or model.sim_header == "seqTransf":
+            contain_frame_position = False
+            for key in state_dict.keys():
+                if key.find("frame_position_embeddings") > -1:
+                    contain_frame_position = True
+                    break
+            if contain_frame_position is False:
+                for key, val in clip_state_dict.items():
+                    if key == "positional_embedding":
+                        state_dict["frame_position_embeddings.weight"] = val.clone()
+                        continue
+                    if model.sim_header == "seqTransf" and key.find("transformer.resblocks") == 0:
+                        num_layer = int(key.split(".")[2])
+                        # cut from beginning
+                        if num_layer < task_config.cross_num_hidden_layers:
+                            state_dict[key.replace("transformer.", "transformerClip.")] = val.clone()
+                            continue
+        ## <=== End of initialization trick
+
+        if state_dict is not None:
+            model = cls.init_preweight(model, state_dict, task_config=task_config)
+
+        return model
+
+def show_log(task_config, info):
+    if task_config is None or task_config.local_rank == 0:
+        logger.warning(info)
+
+def update_attr(target_name, target_config, target_attr_name, source_config, source_attr_name, default_value=None):
+    if hasattr(source_config, source_attr_name):
+        if default_value is None or getattr(source_config, source_attr_name) != default_value:
+            setattr(target_config, target_attr_name, getattr(source_config, source_attr_name))
+            show_log(source_config, "Set {}.{}: {}.".format(target_name,
+                                                            target_attr_name, getattr(target_config, target_attr_name)))
+    return target_config
+
+def check_attr(target_name, task_config):
+    return hasattr(task_config, target_name) and task_config.__dict__[target_name]
+
+class CLIP4Clip(CLIP4ClipPreTrainedModel):
+    def __init__(self, cross_config, clip_state_dict, task_config):
+        super(CLIP4Clip, self).__init__(cross_config)
+        self.task_config = task_config
+        self.ignore_video_index = -1
+
+        assert self.task_config.max_words + self.task_config.max_frames <= cross_config.max_position_embeddings
+
+        self._stage_one = True
+        self._stage_two = False
+
+        show_log(task_config, "Stage-One:{}, Stage-Two:{}".format(self._stage_one, self._stage_two))
+
+        self.loose_type = False
+        if self._stage_one and check_attr('loose_type', self.task_config):
+            self.loose_type = True
+            show_log(task_config, "Test retrieval by loose type.")
+
+        # CLIP Encoders: From OpenAI: CLIP [https://github.com/openai/CLIP] ===>
+        vit = "visual.proj" in clip_state_dict
+        assert vit
+        if vit:
+            vision_width = clip_state_dict["visual.conv1.weight"].shape[0]
+            vision_layers = len(
+                [k for k in clip_state_dict.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")])
+            vision_patch_size = clip_state_dict["visual.conv1.weight"].shape[-1]
+            grid_size = round((clip_state_dict["visual.positional_embedding"].shape[0] - 1) ** 0.5)
+            image_resolution = vision_patch_size * grid_size
+        else:
+            counts: list = [len(set(k.split(".")[2] for k in clip_state_dict if k.startswith(f"visual.layer{b}"))) for b in
+                            [1, 2, 3, 4]]
+            vision_layers = tuple(counts)
+            vision_width = clip_state_dict["visual.layer1.0.conv1.weight"].shape[0]
+            output_width = round((clip_state_dict["visual.attnpool.positional_embedding"].shape[0] - 1) ** 0.5)
+            vision_patch_size = None
+            assert output_width ** 2 + 1 == clip_state_dict["visual.attnpool.positional_embedding"].shape[0]
+            image_resolution = output_width * 32
+
+        embed_dim = clip_state_dict["text_projection"].shape[1]
+        context_length = clip_state_dict["positional_embedding"].shape[0]
+        vocab_size = clip_state_dict["token_embedding.weight"].shape[0]
+        transformer_width = clip_state_dict["ln_final.weight"].shape[0]
+        transformer_heads = transformer_width // 64
+        transformer_layers = len(set(k.split(".")[2] for k in clip_state_dict if k.startswith(f"transformer.resblocks")))
+
+        show_log(task_config, "\t embed_dim: {}".format(embed_dim))
+        show_log(task_config, "\t image_resolution: {}".format(image_resolution))
+        show_log(task_config, "\t vision_layers: {}".format(vision_layers))
+        show_log(task_config, "\t vision_width: {}".format(vision_width))
+        show_log(task_config, "\t vision_patch_size: {}".format(vision_patch_size))
+        show_log(task_config, "\t context_length: {}".format(context_length))
+        show_log(task_config, "\t vocab_size: {}".format(vocab_size))
+        show_log(task_config, "\t transformer_width: {}".format(transformer_width))
+        show_log(task_config, "\t transformer_heads: {}".format(transformer_heads))
+        show_log(task_config, "\t transformer_layers: {}".format(transformer_layers))
+
+        self.linear_patch = '2d'
+        if hasattr(task_config, "linear_patch"):
+            self.linear_patch = task_config.linear_patch
+            show_log(task_config, "\t\t linear_patch: {}".format(self.linear_patch))
+
+        # use .float() to avoid overflow/underflow from fp16 weight. https://github.com/openai/CLIP/issues/40
+        cut_top_layer = 0
+        show_log(task_config, "\t cut_top_layer: {}".format(cut_top_layer))
+        self.clip = CLIP(
+            embed_dim,
+            image_resolution, vision_layers-cut_top_layer, vision_width, vision_patch_size,
+            context_length, vocab_size, transformer_width, transformer_heads, transformer_layers-cut_top_layer,
+            linear_patch=self.linear_patch
+        ).float()
+
+        for key in ["input_resolution", "context_length", "vocab_size"]:
+            if key in clip_state_dict:
+                del clip_state_dict[key]
+
+        convert_weights(self.clip)
+        # <=== End of CLIP Encoders
+
+        self.sim_header = 'meanP'
+        if hasattr(task_config, "sim_header"):
+            self.sim_header = task_config.sim_header
+            show_log(task_config, "\t sim_header: {}".format(self.sim_header))
+        if self.sim_header == "tightTransf": assert self.loose_type is False
+
+        cross_config.max_position_embeddings = context_length
+        if self.loose_type is False:
+            # Cross Encoder ===>
+            cross_config = update_attr("cross_config", cross_config, "num_hidden_layers", self.task_config, "cross_num_hidden_layers")
+            self.cross = CrossModel(cross_config)
+            # <=== End of Cross Encoder
+            self.similarity_dense = nn.Linear(cross_config.hidden_size, 1)
+
+        if self.sim_header == "seqLSTM" or self.sim_header == "seqTransf":
+            self.frame_position_embeddings = nn.Embedding(cross_config.max_position_embeddings, cross_config.hidden_size)
+        if self.sim_header == "seqTransf":
+            self.transformerClip = TransformerClip(width=transformer_width, layers=self.task_config.cross_num_hidden_layers,
+                                                   heads=transformer_heads, )
+        if self.sim_header == "seqLSTM":
+            self.lstm_visual = nn.LSTM(input_size=cross_config.hidden_size, hidden_size=cross_config.hidden_size,
+                                       batch_first=True, bidirectional=False, num_layers=1)
+
+        self.loss_fct = CrossEn()
+
+        self.apply(self.init_weights)
+
+    def forward(self, input_ids, token_type_ids, attention_mask, video, video_mask=None):
+        input_ids = input_ids.view(-1, input_ids.shape[-1])
+        token_type_ids = token_type_ids.view(-1, token_type_ids.shape[-1])
+        attention_mask = attention_mask.view(-1, attention_mask.shape[-1])
+        video_mask = video_mask.view(-1, video_mask.shape[-1])
+
+        # T x 3 x H x W
+        video = torch.as_tensor(video).float()
+        b, pair, bs, ts, channel, h, w = video.shape
+        video = video.view(b * pair * bs * ts, channel, h, w)
+        video_frame = bs * ts
+
+        sequence_output, visual_output = self.get_sequence_visual_output(input_ids, token_type_ids, attention_mask,
+                                                                         video, video_mask, shaped=True, video_frame=video_frame)
+
+        if self.training:
+            loss = 0.
+            sim_matrix, *_tmp = self.get_similarity_logits(sequence_output, visual_output, attention_mask, video_mask,
+                                                    shaped=True, loose_type=self.loose_type)
+            sim_loss1 = self.loss_fct(sim_matrix)
+            sim_loss2 = self.loss_fct(sim_matrix.T)
+            sim_loss = (sim_loss1 + sim_loss2) / 2
+            loss += sim_loss
+
+            return loss
+        else:
+            return None
+
+    def get_sequence_output(self, input_ids, token_type_ids, attention_mask, shaped=False):
+        if shaped is False:
+            input_ids = input_ids.view(-1, input_ids.shape[-1])
+            token_type_ids = token_type_ids.view(-1, token_type_ids.shape[-1])
+            attention_mask = attention_mask.view(-1, attention_mask.shape[-1])
+
+        bs_pair = input_ids.size(0)
+        sequence_hidden = self.clip.encode_text(input_ids).float()
+        sequence_hidden = sequence_hidden.view(bs_pair, -1, sequence_hidden.size(-1))
+
+        return sequence_hidden
+
+    def get_visual_output(self, video, video_mask, shaped=False, video_frame=-1):
+        if shaped is False:
+            video_mask = video_mask.view(-1, video_mask.shape[-1])
+            video = torch.as_tensor(video).float()
+            b, pair, bs, ts, channel, h, w = video.shape
+            video = video.view(b * pair * bs * ts, channel, h, w)
+            video_frame = bs * ts
+
+        bs_pair = video_mask.size(0)
+        visual_hidden = self.clip.encode_image(video, video_frame=video_frame).float()
+        visual_hidden = visual_hidden.view(bs_pair, -1, visual_hidden.size(-1))
+
+        return visual_hidden
+
+    def get_sequence_visual_output(self, input_ids, token_type_ids, attention_mask, video, video_mask, shaped=False, video_frame=-1):
+        if shaped is False:
+            input_ids = input_ids.view(-1, input_ids.shape[-1])
+            token_type_ids = token_type_ids.view(-1, token_type_ids.shape[-1])
+            attention_mask = attention_mask.view(-1, attention_mask.shape[-1])
+            video_mask = video_mask.view(-1, video_mask.shape[-1])
+
+            video = torch.as_tensor(video).float()
+            b, pair, bs, ts, channel, h, w = video.shape
+            video = video.view(b * pair * bs * ts, channel, h, w)
+            video_frame = bs * ts
+
+        sequence_output = self.get_sequence_output(input_ids, token_type_ids, attention_mask, shaped=True)
+        visual_output = self.get_visual_output(video, video_mask, shaped=True, video_frame=video_frame)
+
+        return sequence_output, visual_output
+
+    def _get_cross_output(self, sequence_output, visual_output, attention_mask, video_mask):
+
+        concat_features = torch.cat((sequence_output, visual_output), dim=1)  # concatnate tokens and frames
+        concat_mask = torch.cat((attention_mask, video_mask), dim=1)
+        text_type_ = torch.zeros_like(attention_mask)
+        video_type_ = torch.ones_like(video_mask)
+        concat_type = torch.cat((text_type_, video_type_), dim=1)
+
+        cross_layers, pooled_output = self.cross(concat_features, concat_type, concat_mask, output_all_encoded_layers=True)
+        cross_output = cross_layers[-1]
+
+        return cross_output, pooled_output, concat_mask
+
+    def _mean_pooling_for_similarity_sequence(self, sequence_output, attention_mask):
+        attention_mask_un = attention_mask.to(dtype=torch.float).unsqueeze(-1)
+        attention_mask_un[:, 0, :] = 0.
+        sequence_output = sequence_output * attention_mask_un
+        text_out = torch.sum(sequence_output, dim=1) / torch.sum(attention_mask_un, dim=1, dtype=torch.float)
+        return text_out
+
+    def _mean_pooling_for_similarity_visual(self, visual_output, video_mask,):
+        video_mask_un = video_mask.to(dtype=torch.float).unsqueeze(-1)
+        visual_output = visual_output * video_mask_un
+        video_mask_un_sum = torch.sum(video_mask_un, dim=1, dtype=torch.float)
+        video_mask_un_sum[video_mask_un_sum == 0.] = 1.
+        video_out = torch.sum(visual_output, dim=1) / video_mask_un_sum
+        return video_out
+
+    def _mean_pooling_for_similarity(self, sequence_output, visual_output, attention_mask, video_mask,):
+        text_out = self._mean_pooling_for_similarity_sequence(sequence_output, attention_mask)
+        video_out = self._mean_pooling_for_similarity_visual(visual_output, video_mask)
+
+        return text_out, video_out
+
+    def _loose_similarity(self, sequence_output, visual_output, attention_mask, video_mask, sim_header="meanP"):
+        sequence_output, visual_output = sequence_output.contiguous(), visual_output.contiguous()
+
+        if sim_header == "meanP":
+            # Default: Parameter-free type
+            pass
+        elif sim_header == "seqLSTM":
+            # Sequential type: LSTM
+            visual_output_original = visual_output
+            visual_output = pack_padded_sequence(visual_output, torch.sum(video_mask, dim=-1).cpu(),
+                                                 batch_first=True, enforce_sorted=False)
+            visual_output, _ = self.lstm_visual(visual_output)
+            if self.training: self.lstm_visual.flatten_parameters()
+            visual_output, _ = pad_packed_sequence(visual_output, batch_first=True)
+            visual_output = torch.cat((visual_output, visual_output_original[:, visual_output.size(1):, ...].contiguous()), dim=1)
+            visual_output = visual_output + visual_output_original
+        elif sim_header == "seqTransf":
+            # Sequential type: Transformer Encoder
+            visual_output_original = visual_output
+            seq_length = visual_output.size(1)
+            position_ids = torch.arange(seq_length, dtype=torch.long, device=visual_output.device)
+            position_ids = position_ids.unsqueeze(0).expand(visual_output.size(0), -1)
+            frame_position_embeddings = self.frame_position_embeddings(position_ids)
+            visual_output = visual_output + frame_position_embeddings
+
+            extended_video_mask = (1.0 - video_mask.unsqueeze(1)) * -1000000.0
+            extended_video_mask = extended_video_mask.expand(-1, video_mask.size(1), -1)
+            visual_output = visual_output.permute(1, 0, 2)  # NLD -> LND
+            visual_output = self.transformerClip(visual_output, extended_video_mask)
+            visual_output = visual_output.permute(1, 0, 2)  # LND -> NLD
+            visual_output = visual_output + visual_output_original
+
+        if self.training:
+            visual_output = allgather(visual_output, self.task_config)
+            video_mask = allgather(video_mask, self.task_config)
+            sequence_output = allgather(sequence_output, self.task_config)
+            torch.distributed.barrier()
+
+        visual_output = visual_output / visual_output.norm(dim=-1, keepdim=True)
+        visual_output = self._mean_pooling_for_similarity_visual(visual_output, video_mask)
+        visual_output = visual_output / visual_output.norm(dim=-1, keepdim=True)
+
+        sequence_output = sequence_output.squeeze(1)
+        sequence_output = sequence_output / sequence_output.norm(dim=-1, keepdim=True)
+
+        logit_scale = self.clip.logit_scale.exp()
+        retrieve_logits = logit_scale * torch.matmul(sequence_output, visual_output.t())
+        return retrieve_logits
+
+    def _cross_similarity(self, sequence_output, visual_output, attention_mask, video_mask):
+        sequence_output, visual_output = sequence_output.contiguous(), visual_output.contiguous()
+
+        b_text, s_text, h_text = sequence_output.size()
+        b_visual, s_visual, h_visual = visual_output.size()
+
+        retrieve_logits_list = []
+
+        step_size = b_text      # set smaller to reduce memory cost
+        split_size = [step_size] * (b_text // step_size)
+        release_size = b_text - sum(split_size)
+        if release_size > 0:
+            split_size += [release_size]
+
+        # due to clip text branch retrun the last hidden
+        attention_mask = torch.ones(sequence_output.size(0), 1)\
+            .to(device=attention_mask.device, dtype=attention_mask.dtype)
+
+        sequence_output_splits = torch.split(sequence_output, split_size, dim=0)
+        attention_mask_splits = torch.split(attention_mask, split_size, dim=0)
+        for i in range(len(split_size)):
+            sequence_output_row = sequence_output_splits[i]
+            attention_mask_row = attention_mask_splits[i]
+            sequence_output_l = sequence_output_row.unsqueeze(1).repeat(1, b_visual, 1, 1)
+            sequence_output_l = sequence_output_l.view(-1, s_text, h_text)
+            attention_mask_l = attention_mask_row.unsqueeze(1).repeat(1, b_visual, 1)
+            attention_mask_l = attention_mask_l.view(-1, s_text)
+
+            step_truth = sequence_output_row.size(0)
+            visual_output_r = visual_output.unsqueeze(0).repeat(step_truth, 1, 1, 1)
+            visual_output_r = visual_output_r.view(-1, s_visual, h_visual)
+            video_mask_r = video_mask.unsqueeze(0).repeat(step_truth, 1, 1)
+            video_mask_r = video_mask_r.view(-1, s_visual)
+
+            cross_output, pooled_output, concat_mask = \
+                self._get_cross_output(sequence_output_l, visual_output_r, attention_mask_l, video_mask_r)
+            retrieve_logits_row = self.similarity_dense(pooled_output).squeeze(-1).view(step_truth, b_visual)
+
+            retrieve_logits_list.append(retrieve_logits_row)
+
+        retrieve_logits = torch.cat(retrieve_logits_list, dim=0)
+        return retrieve_logits
+
+    def get_similarity_logits(self, sequence_output, visual_output, attention_mask, video_mask, shaped=False, loose_type=False):
+        if shaped is False:
+            attention_mask = attention_mask.view(-1, attention_mask.shape[-1])
+            video_mask = video_mask.view(-1, video_mask.shape[-1])
+
+        contrastive_direction = ()
+        if loose_type:
+            assert self.sim_header in ["meanP", "seqLSTM", "seqTransf"]
+            retrieve_logits = self._loose_similarity(sequence_output, visual_output, attention_mask, video_mask, sim_header=self.sim_header)
+        else:
+            assert self.sim_header in ["tightTransf"]
+            retrieve_logits = self._cross_similarity(sequence_output, visual_output, attention_mask, video_mask, )
+
+        return retrieve_logits, contrastive_direction
+
+
+class VideoCLIP_XL(nn.Module):
+    def __init__(self):
+        super(VideoCLIP_XL, self).__init__()
+        self.text_model = text_encoder.load().float()
+        self.vision_model = get_vision_encoder().float()
+
+    def get_sequence_output(self, input_ids, token_type_ids, attention_mask, shaped=False):
+        if not shaped:
+            input_ids = input_ids.view(-1, input_ids.shape[-1])
+
+        # (B, D)
+        text_feat = self.text_model.encode_text(input_ids)
+
+        # match CLIP4Clip shape: (B, 1, D)
+        text_feat = text_feat.unsqueeze(1)
+
+        return text_feat
+
+    def get_visual_output(self, video, video_mask, shaped=False, video_frame=-1):
+        if not shaped:
+            video_mask = video_mask.view(-1, video_mask.shape[-1])
+
+            video = torch.as_tensor(video).float()
+            b, pair, bs, ts, c, h, w = video.shape
+
+            # reshape to (B, C, T, H, W)
+            video = video.view(b * pair, bs * ts, c, h, w)
+            video = video.permute(0, 2, 1, 3, 4).contiguous()
+
+        bs_pair = video_mask.size(0)
+
+        # forward through vision encoder
+        visual_feat = self.vision_model.get_vid_features(video)
+
+        # ---- normalize shape ----
+        if visual_feat.dim() == 2:
+            # (B, D) → (B, 1, D)
+            visual_feat = visual_feat.unsqueeze(1)
+
+        elif visual_feat.dim() == 3:
+            # already (B, T or tokens, D)
+            pass
+
+        else:
+            raise ValueError(f"Unexpected vision output shape: {visual_feat.shape}")
+
+        return visual_feat
+
+    def get_sequence_visual_output(
+        self,
+        input_ids,
+        token_type_ids,
+        attention_mask,
+        video,
+        video_mask,
+        shaped=False,
+        video_frame=-1
+    ):
+        if not shaped:
+            input_ids = input_ids.view(-1, input_ids.shape[-1])
+            attention_mask = attention_mask.view(-1, attention_mask.shape[-1])
+            video_mask = video_mask.view(-1, video_mask.shape[-1])
+
+            video = torch.as_tensor(video).float()
+            b, pair, bs, ts, c, h, w = video.shape
+            video = video.view(b * pair * bs * ts, c, h, w)
+            video_frame = bs * ts
+
+        sequence_output = self.get_sequence_output(
+            input_ids, token_type_ids, attention_mask, shaped=True
+        )
+
+        visual_output = self.get_visual_output(
+            video, video_mask, shaped=True, video_frame=video_frame
+        )
+
+        return sequence_output, visual_output
+
+    def get_similarity_logits(
+        self,
+        sequence_output,
+        visual_output,
+        attention_mask,
+        video_mask,
+        shaped=False,
+        loose_type=False
+    ):
+        if not shaped:
+            attention_mask = attention_mask.view(-1, attention_mask.shape[-1])
+            video_mask = video_mask.view(-1, video_mask.shape[-1])
+
+        # --- normalize frame features ---
+        visual_output = visual_output / visual_output.norm(dim=-1, keepdim=True)
+
+        # mean pooling over frames
+        video_mask_un = video_mask.unsqueeze(-1).float()
+        visual_output = visual_output * video_mask_un
+
+        denom = video_mask_un.sum(dim=1)
+        denom[denom == 0] = 1.0
+
+        video_feat = visual_output.sum(dim=1) / denom
+        video_feat = video_feat / video_feat.norm(dim=-1, keepdim=True)
+
+        # --- text ---
+        if sequence_output.dim() == 3:
+            text_feat = sequence_output[:, 0]  # CLS or pooled
+        else:
+            text_feat = sequence_output
+
+        text_feat = text_feat / text_feat.norm(dim=-1, keepdim=True)
+
+        # --- similarity ---
+        logit_scale = 1.0  # or learnable like CLIP
+        logits = logit_scale * torch.matmul(text_feat, video_feat.t())
+
+        return logits, ()

requirements.txt

@@ -0,0 +1,6 @@
+opencv-python
+ftfy
+regex
+timm
+decord
+einops

utils/text_encoder/__init__.py

@@ -0,0 +1 @@
+from .text_encoder import *

utils/text_encoder/bpe_simple_vocab_16e6.txt.gz

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

utils/text_encoder/model_text_encoder.py

@@ -0,0 +1,395 @@
+from collections import OrderedDict
+from typing import Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1):
+        super().__init__()
+
+        # all conv layers have stride 1. an avgpool is performed after the second convolution when stride > 1
+        self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu1 = nn.ReLU(inplace=True)
+
+        self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.relu2 = nn.ReLU(inplace=True)
+
+        self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity()
+
+        self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.relu3 = nn.ReLU(inplace=True)
+
+        self.downsample = None
+        self.stride = stride
+
+        if stride > 1 or inplanes != planes * Bottleneck.expansion:
+            # downsampling layer is prepended with an avgpool, and the subsequent convolution has stride 1
+            self.downsample = nn.Sequential(OrderedDict([
+                ("-1", nn.AvgPool2d(stride)),
+                ("0", nn.Conv2d(inplanes, planes * self.expansion, 1, stride=1, bias=False)),
+                ("1", nn.BatchNorm2d(planes * self.expansion))
+            ]))
+
+    def forward(self, x: torch.Tensor):
+        identity = x
+
+        out = self.relu1(self.bn1(self.conv1(x)))
+        out = self.relu2(self.bn2(self.conv2(out)))
+        out = self.avgpool(out)
+        out = self.bn3(self.conv3(out))
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu3(out)
+        return out
+
+
+class AttentionPool2d(nn.Module):
+    def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None):
+        super().__init__()
+        self.positional_embedding = nn.Parameter(torch.randn(spacial_dim ** 2 + 1, embed_dim) / embed_dim ** 0.5)
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)
+        self.num_heads = num_heads
+
+    def forward(self, x):
+        x = x.flatten(start_dim=2).permute(2, 0, 1)  # NCHW -> (HW)NC
+        x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0)  # (HW+1)NC
+        x = x + self.positional_embedding[:, None, :].to(x.dtype)  # (HW+1)NC
+        x, _ = F.multi_head_attention_forward(
+            query=x[:1], key=x, value=x,
+            embed_dim_to_check=x.shape[-1],
+            num_heads=self.num_heads,
+            q_proj_weight=self.q_proj.weight,
+            k_proj_weight=self.k_proj.weight,
+            v_proj_weight=self.v_proj.weight,
+            in_proj_weight=None,
+            in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]),
+            bias_k=None,
+            bias_v=None,
+            add_zero_attn=False,
+            dropout_p=0,
+            out_proj_weight=self.c_proj.weight,
+            out_proj_bias=self.c_proj.bias,
+            use_separate_proj_weight=True,
+            training=self.training,
+            need_weights=False
+        )
+        return x.squeeze(0)
+
+
+class ModifiedResNet(nn.Module):
+    """
+    A ResNet class that is similar to torchvision's but contains the following changes:
+    - There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max pool.
+    - Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions with stride > 1
+    - The final pooling layer is a QKV attention instead of an average pool
+    """
+
+    def __init__(self, layers, output_dim, heads, input_resolution=224, width=64):
+        super().__init__()
+        self.output_dim = output_dim
+        self.input_resolution = input_resolution
+
+        # the 3-layer stem
+        self.conv1 = nn.Conv2d(3, width // 2, kernel_size=3, stride=2, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(width // 2)
+        self.relu1 = nn.ReLU(inplace=True)
+        self.conv2 = nn.Conv2d(width // 2, width // 2, kernel_size=3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(width // 2)
+        self.relu2 = nn.ReLU(inplace=True)
+        self.conv3 = nn.Conv2d(width // 2, width, kernel_size=3, padding=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(width)
+        self.relu3 = nn.ReLU(inplace=True)
+        self.avgpool = nn.AvgPool2d(2)
+
+        # residual layers
+        self._inplanes = width  # this is a *mutable* variable used during construction
+        self.layer1 = self._make_layer(width, layers[0])
+        self.layer2 = self._make_layer(width * 2, layers[1], stride=2)
+        self.layer3 = self._make_layer(width * 4, layers[2], stride=2)
+        self.layer4 = self._make_layer(width * 8, layers[3], stride=2)
+
+        embed_dim = width * 32  # the ResNet feature dimension
+        self.attnpool = AttentionPool2d(input_resolution // 32, embed_dim, heads, output_dim)
+
+    def _make_layer(self, planes, blocks, stride=1):
+        layers = [Bottleneck(self._inplanes, planes, stride)]
+
+        self._inplanes = planes * Bottleneck.expansion
+        for _ in range(1, blocks):
+            layers.append(Bottleneck(self._inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        def stem(x):
+            x = self.relu1(self.bn1(self.conv1(x)))
+            x = self.relu2(self.bn2(self.conv2(x)))
+            x = self.relu3(self.bn3(self.conv3(x)))
+            x = self.avgpool(x)
+            return x
+
+        x = x.type(self.conv1.weight.dtype)
+        x = stem(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = self.attnpool(x)
+
+        return x
+
+
+class LayerNorm(nn.LayerNorm):
+    """Subclass torch's LayerNorm to handle fp16."""
+
+    def forward(self, x: torch.Tensor):
+        orig_type = x.dtype
+        ret = super().forward(x.type(torch.float32))
+        return ret.type(orig_type)
+
+
+class QuickGELU(nn.Module):
+    def forward(self, x: torch.Tensor):
+        return x * torch.sigmoid(1.702 * x)
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None):
+        super().__init__()
+
+        self.attn = nn.MultiheadAttention(d_model, n_head)
+        self.ln_1 = LayerNorm(d_model)
+        self.mlp = nn.Sequential(OrderedDict([
+            ("c_fc", nn.Linear(d_model, d_model * 4)),
+            ("gelu", QuickGELU()),
+            ("c_proj", nn.Linear(d_model * 4, d_model))
+        ]))
+        self.ln_2 = LayerNorm(d_model)
+        self.attn_mask = attn_mask
+
+    def attention(self, x: torch.Tensor):
+        self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None
+        return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0]
+
+    def forward(self, x: torch.Tensor):
+        x = x + self.attention(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(self, width: int, layers: int, heads: int, attn_mask: torch.Tensor = None):
+        super().__init__()
+        self.width = width
+        self.layers = layers
+        self.resblocks = nn.Sequential(*[ResidualAttentionBlock(width, heads, attn_mask) for _ in range(layers)])
+
+    def forward(self, x: torch.Tensor):
+        return self.resblocks(x)
+
+
+class VisionTransformer(nn.Module):
+    def __init__(self, input_resolution: int, patch_size: int, width: int, layers: int, heads: int, output_dim: int):
+        super().__init__()
+        self.input_resolution = input_resolution
+        self.output_dim = output_dim
+        self.conv1 = nn.Conv2d(in_channels=3, out_channels=width, kernel_size=patch_size, stride=patch_size, bias=False)
+
+        scale = width ** -0.5
+        self.class_embedding = nn.Parameter(scale * torch.randn(width))
+        self.positional_embedding = nn.Parameter(scale * torch.randn((input_resolution // patch_size) ** 2 + 1, width))
+        self.ln_pre = LayerNorm(width)
+
+        self.transformer = Transformer(width, layers, heads)
+
+        self.ln_post = LayerNorm(width)
+        self.proj = nn.Parameter(scale * torch.randn(width, output_dim))
+
+    def forward(self, x: torch.Tensor):
+        x = self.conv1(x)  # shape = [*, width, grid, grid]
+        x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
+        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
+        x = torch.cat([self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x], dim=1)  # shape = [*, grid ** 2 + 1, width]
+        x = x + self.positional_embedding.to(x.dtype)
+        x = self.ln_pre(x)
+
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+
+        x = self.ln_post(x[:, 0, :])
+
+        if self.proj is not None:
+            x = x @ self.proj
+
+        return x
+
+
+class CLIP(nn.Module):
+    def __init__(self,
+                 embed_dim: int,
+                 # vision
+                 image_resolution: int,
+                 vision_layers: Union[Tuple[int, int, int, int], int],
+                 vision_width: int,
+                 vision_patch_size: int,
+                 # text
+                 context_length: int,
+                 vocab_size: int,
+                 transformer_width: int,
+                 transformer_heads: int,
+                 transformer_layers: int, 
+                 load_from_clip: bool
+                 ):
+        super().__init__()
+
+        self.context_length = 248
+
+        self.transformer = Transformer(
+            width=transformer_width,
+            layers=transformer_layers,
+            heads=transformer_heads,
+            attn_mask=self.build_attention_mask()
+        )
+
+        self.vocab_size = vocab_size
+        self.token_embedding = nn.Embedding(vocab_size, transformer_width)
+
+        if load_from_clip == False:
+            self.positional_embedding = nn.Parameter(torch.empty(248, transformer_width))
+            self.positional_embedding_res = nn.Parameter(torch.empty(248, transformer_width))
+
+        else:
+            self.positional_embedding = nn.Parameter(torch.empty(77, transformer_width))
+
+        self.ln_final = LayerNorm(transformer_width)
+
+        self.text_projection = nn.Parameter(torch.empty(transformer_width, embed_dim))
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+        self.initialize_parameters()
+        self.mask1 = torch.zeros([248, 1])
+        self.mask1[:20, :] = 1
+        self.mask2 = torch.zeros([248, 1])
+        self.mask2[20:, :] = 1
+
+    
+    def initialize_parameters(self):
+        nn.init.normal_(self.token_embedding.weight, std=0.02)
+        nn.init.normal_(self.positional_embedding, std=0.01)
+
+        proj_std = (self.transformer.width ** -0.5) * ((2 * self.transformer.layers) ** -0.5)
+        attn_std = self.transformer.width ** -0.5
+        fc_std = (2 * self.transformer.width) ** -0.5
+        for block in self.transformer.resblocks:
+            nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
+            nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
+            nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
+            nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)
+
+        if self.text_projection is not None:
+            nn.init.normal_(self.text_projection, std=self.transformer.width ** -0.5)
+
+    def build_attention_mask(self):
+        # lazily create causal attention mask, with full attention between the vision tokens
+        # pytorch uses additive attention mask; fill with -inf
+        mask = torch.empty(self.context_length, self.context_length)
+        mask.fill_(float("-inf"))
+        mask.triu_(1)  # zero out the lower diagonal
+        return mask
+
+    @property
+    def dtype(self):
+        return self.token_embedding.weight.dtype
+
+    def encode_text(self, text): 
+        x = self.token_embedding(text).type(self.dtype)  # [batch_size, n_ctx, d_model]
+        
+        x = x + (self.positional_embedding.to(x.device) * self.mask1.to(x.device)).type(self.dtype).to(x.device) + (self.positional_embedding_res.to(x.device) * self.mask2.to(x.device)).type(self.dtype).to(x.device) 
+        
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.ln_final(x).type(self.dtype)
+
+        # x.shape = [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 x
+
+    def encode_text_full(self, text): 
+        x = self.token_embedding(text).type(self.dtype)  # [batch_size, n_ctx, d_model]
+        
+        x = x + (self.positional_embedding.to(x.device) * self.mask1.to(x.device)).type(self.dtype).to(x.device) + (self.positional_embedding_res.to(x.device) * self.mask2.to(x.device)).type(self.dtype).to(x.device) 
+        
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.ln_final(x).type(self.dtype)
+
+        return x
+    
+
+def convert_weights(model: nn.Module):
+    """Convert applicable model parameters to fp16"""
+
+    def _convert_weights_to_fp16(l):
+        if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
+            l.weight.data = l.weight.data.half()
+            if l.bias is not None:
+                l.bias.data = l.bias.data.half()
+
+        if isinstance(l, nn.MultiheadAttention):
+            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.half()
+
+        for name in ["text_projection", "proj"]:
+            if hasattr(l, name):
+                attr = getattr(l, name)
+                if attr is not None:
+                    attr.data = attr.data.half()
+
+    model.apply(_convert_weights_to_fp16)
+
+
+def build_model(load_from_clip: bool):
+    
+    vision_width = 1024
+    vision_layers = 24
+    vision_patch_size = 14
+    grid_size = 16
+    image_resolution = 224
+    
+    embed_dim = 768
+    context_length = 248
+    vocab_size = 49408
+    transformer_width = 768
+    transformer_heads = 12
+    transformer_layers = 12
+
+    model = CLIP(
+        embed_dim,
+        image_resolution, vision_layers, vision_width, vision_patch_size,
+        context_length, vocab_size, transformer_width, transformer_heads, transformer_layers, load_from_clip
+    )
+
+    convert_weights(model)
+    return model.eval()

utils/text_encoder/simple_tokenizer.py

@@ -0,0 +1,132 @@
+import gzip
+import html
+import os
+from functools import lru_cache
+
+import ftfy
+import regex as re
+
+
+@lru_cache()
+def default_bpe():
+    return os.path.join(os.path.dirname(os.path.abspath(__file__)), "bpe_simple_vocab_16e6.txt.gz")
+
+
+@lru_cache()
+def bytes_to_unicode():
+    """
+    Returns list of utf-8 byte and a corresponding list of unicode strings.
+    The reversible bpe codes work on unicode strings.
+    This means you need a large # of unicode characters in your vocab if you want to avoid UNKs.
+    When you're at something like a 10B token dataset you end up needing around 5K for decent coverage.
+    This is a signficant percentage of your normal, say, 32K bpe vocab.
+    To avoid that, we want lookup tables between utf-8 bytes and unicode strings.
+    And avoids mapping to whitespace/control characters the bpe code barfs on.
+    """
+    bs = list(range(ord("!"), ord("~")+1))+list(range(ord("¡"), ord("¬")+1))+list(range(ord("®"), ord("ÿ")+1))
+    cs = bs[:]
+    n = 0
+    for b in range(2**8):
+        if b not in bs:
+            bs.append(b)
+            cs.append(2**8+n)
+            n += 1
+    cs = [chr(n) for n in cs]
+    return dict(zip(bs, cs))
+
+
+def get_pairs(word):
+    """Return set of symbol pairs in a word.
+    Word is represented as tuple of symbols (symbols being variable-length strings).
+    """
+    pairs = set()
+    prev_char = word[0]
+    for char in word[1:]:
+        pairs.add((prev_char, char))
+        prev_char = char
+    return pairs
+
+
+def basic_clean(text):
+    text = ftfy.fix_text(text)
+    text = html.unescape(html.unescape(text))
+    return text.strip()
+
+
+def whitespace_clean(text):
+    text = re.sub(r'\s+', ' ', text)
+    text = text.strip()
+    return text
+
+
+class SimpleTokenizer(object):
+    def __init__(self, bpe_path: str = default_bpe()):
+        self.byte_encoder = bytes_to_unicode()
+        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
+        merges = gzip.open(bpe_path).read().decode("utf-8").split('\n')
+        merges = merges[1:49152-256-2+1]
+        merges = [tuple(merge.split()) for merge in merges]
+        vocab = list(bytes_to_unicode().values())
+        vocab = vocab + [v+'</w>' for v in vocab]
+        for merge in merges:
+            vocab.append(''.join(merge))
+        vocab.extend(['<|startoftext|>', '<|endoftext|>'])
+        self.encoder = dict(zip(vocab, range(len(vocab))))
+        self.decoder = {v: k for k, v in self.encoder.items()}
+        self.bpe_ranks = dict(zip(merges, range(len(merges))))
+        self.cache = {'<|startoftext|>': '<|startoftext|>', '<|endoftext|>': '<|endoftext|>'}
+        self.pat = re.compile(r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""", re.IGNORECASE)
+
+    def bpe(self, token):
+        if token in self.cache:
+            return self.cache[token]
+        word = tuple(token[:-1]) + ( token[-1] + '</w>',)
+        pairs = get_pairs(word)
+
+        if not pairs:
+            return token+'</w>'
+
+        while True:
+            bigram = min(pairs, key = lambda pair: self.bpe_ranks.get(pair, float('inf')))
+            if bigram not in self.bpe_ranks:
+                break
+            first, second = bigram
+            new_word = []
+            i = 0
+            while i < len(word):
+                try:
+                    j = word.index(first, i)
+                    new_word.extend(word[i:j])
+                    i = j
+                except:
+                    new_word.extend(word[i:])
+                    break
+
+                if word[i] == first and i < len(word)-1 and word[i+1] == second:
+                    new_word.append(first+second)
+                    i += 2
+                else:
+                    new_word.append(word[i])
+                    i += 1
+            new_word = tuple(new_word)
+            word = new_word
+            if len(word) == 1:
+                break
+            else:
+                pairs = get_pairs(word)
+        word = ' '.join(word)
+        self.cache[token] = word
+        return word
+
+    def encode(self, text):
+        bpe_tokens = []
+        text = whitespace_clean(basic_clean(text)).lower()
+        for token in re.findall(self.pat, text):
+            token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8'))
+            bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' '))
+        return bpe_tokens
+
+    def decode(self, tokens):
+        text = ''.join([self.decoder[token] for token in tokens])
+        text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors="replace").replace('</w>', ' ')
+        return text

utils/text_encoder/text_encoder.py

@@ -0,0 +1,75 @@
+import hashlib
+import os
+import urllib
+import warnings
+from typing import Any, Union, List
+from pkg_resources import packaging
+from torch import nn
+import torch
+from PIL import Image
+from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
+
+from .model_text_encoder import build_model
+from .simple_tokenizer import SimpleTokenizer as _Tokenizer
+
+try:
+    from torchvision.transforms import InterpolationMode
+    BICUBIC = InterpolationMode.BICUBIC
+except ImportError:
+    BICUBIC = Image.BICUBIC
+
+    
+_tokenizer = _Tokenizer()
+
+
+def _convert_image_to_rgb(image):
+    return image.convert("RGB")
+
+
+def load():
+    model = build_model(load_from_clip = False)
+
+    return model
+        
+
+def tokenize(texts: Union[str, List[str]], context_length: int = 77*4-60, truncate: bool = False) -> Union[torch.IntTensor, torch.LongTensor]:
+    """
+    Returns the tokenized representation of given input string(s)
+
+    Parameters
+    ----------
+    texts : Union[str, List[str]]
+        An input string or a list of input strings to tokenize
+
+    context_length : int
+        The context length to use; all CLIP models use 77 as the context length
+
+    truncate: bool
+        Whether to truncate the text in case its encoding is longer than the context length
+
+    Returns
+    -------
+    A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length].
+    We return LongTensor when torch version is <1.8.0, since older index_select requires indices to be long.
+    """
+    if isinstance(texts, str):
+        texts = [texts]
+
+    sot_token = _tokenizer.encoder["<|startoftext|>"]
+    eot_token = _tokenizer.encoder["<|endoftext|>"]
+    all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts]
+    if packaging.version.parse(torch.__version__) < packaging.version.parse("1.8.0"):
+        result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
+    else:
+        result = torch.zeros(len(all_tokens), context_length, dtype=torch.int)
+
+    for i, tokens in enumerate(all_tokens):
+        if len(tokens) > context_length:
+            if truncate:
+                tokens = tokens[:context_length]
+                tokens[-1] = eot_token
+            else:
+                raise RuntimeError(f"Input {texts[i]} is too long for context length {context_length}")
+        result[i, :len(tokens)] = torch.tensor(tokens)
+
+    return result

utils/vision_encoder/__init__.py

@@ -0,0 +1,11 @@
+import torch
+import numpy as np
+import cv2
+import os
+
+from .model_vision_encoder import VisionEncoder
+
+def get_vision_encoder():
+    vision_encoder = VisionEncoder()
+    
+    return vision_encoder

utils/vision_encoder/clip_vision.py

@@ -0,0 +1,327 @@
+#!/usr/bin/env python
+import os
+import logging
+from collections import OrderedDict
+
+import torch
+from torch import nn
+from einops import rearrange
+from timm.models.layers import DropPath
+from timm.models.registry import register_model
+
+import torch.utils.checkpoint as checkpoint
+
+logger = logging.getLogger(__name__)
+
+def load_temp_embed_with_mismatch(temp_embed_old, temp_embed_new, add_zero=True):
+    """
+    Add/Remove extra temporal_embeddings as needed.
+    https://arxiv.org/abs/2104.00650 shows adding zero paddings works.
+
+    temp_embed_old: (1, num_frames_old, 1, d)
+    temp_embed_new: (1, num_frames_new, 1, d)
+    add_zero: bool, if True, add zero, else, interpolate trained embeddings.
+    """
+    # TODO zero pad
+    num_frms_new = temp_embed_new.shape[1]
+    num_frms_old = temp_embed_old.shape[1]
+    logger.info(f"Load temporal_embeddings, lengths: {num_frms_old}-->{num_frms_new}")
+    if num_frms_new > num_frms_old:
+        if add_zero:
+            temp_embed_new[
+                :, :num_frms_old
+            ] = temp_embed_old  # untrained embeddings are zeros.
+        else:
+            temp_embed_new = interpolate_temporal_pos_embed(temp_embed_old, num_frms_new)
+    elif num_frms_new < num_frms_old:
+        temp_embed_new = temp_embed_old[:, :num_frms_new]
+    else:  # =
+        temp_embed_new = temp_embed_old
+    return temp_embed_new
+
+
+class QuickGELU(nn.Module):
+    def forward(self, x):
+        return x * torch.sigmoid(1.702 * x)
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(self, d_model, n_head, drop_path=0., attn_mask=None, dropout=0.):
+        super().__init__()
+
+        self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.drop_path2 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        # logger.info(f'Droppath: {drop_path}')
+        self.attn = nn.MultiheadAttention(d_model, n_head, dropout=dropout)
+        self.ln_1 = nn.LayerNorm(d_model)
+        self.mlp = nn.Sequential(OrderedDict([
+            ("c_fc", nn.Linear(d_model, d_model * 4)),
+            ("gelu", QuickGELU()),
+            ("drop1", nn.Dropout(dropout)),
+            ("c_proj", nn.Linear(d_model * 4, d_model)),
+            ("drop2", nn.Dropout(dropout)),
+        ]))
+        self.ln_2 = nn.LayerNorm(d_model)
+        self.attn_mask = attn_mask
+
+    def attention(self, x):
+        self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None
+        return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0]
+
+    def forward(self, x):
+        x = x + self.drop_path1(self.attention(self.ln_1(x)))
+        x = x + self.drop_path2(self.mlp(self.ln_2(x)))
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(self, width, layers, heads, drop_path=0., checkpoint_num=0, dropout=0.):
+        super().__init__()
+        dpr = [x.item() for x in torch.linspace(0, drop_path, layers)]
+        self.resblocks = nn.ModuleList()
+        for idx in range(layers):
+            self.resblocks.append(ResidualAttentionBlock(width, heads, drop_path=dpr[idx], dropout=dropout))
+        self.checkpoint_num = checkpoint_num
+
+    def forward(self, x):
+        for idx, blk in enumerate(self.resblocks):
+            if idx < self.checkpoint_num:
+                x = checkpoint.checkpoint(blk, x)
+            else:
+                x = blk(x)
+        return x
+
+
+class VisionTransformer(nn.Module):
+    def __init__(
+        self, input_resolution, patch_size, width, layers, heads, output_dim=None, 
+        kernel_size=1, num_frames=8, drop_path=0, checkpoint_num=0, dropout=0.,
+        temp_embed=True,
+    ):
+        super().__init__()
+        self.output_dim = output_dim
+        self.conv1 = nn.Conv3d(
+            3, width, 
+            (kernel_size, patch_size, patch_size), 
+            (kernel_size, patch_size, patch_size), 
+            (0, 0, 0), bias=False
+        )
+
+        scale = width ** -0.5
+        self.class_embedding = nn.Parameter(scale * torch.randn(width))
+        self.positional_embedding = nn.Parameter(scale * torch.randn((input_resolution // patch_size) ** 2 + 1, width))
+        self.ln_pre = nn.LayerNorm(width)
+        if temp_embed:
+            self.temporal_positional_embedding = nn.Parameter(torch.zeros(1, num_frames, width))
+        
+        self.transformer = Transformer(
+            width, layers, heads, drop_path=drop_path, checkpoint_num=checkpoint_num,
+            dropout=dropout)
+
+        self.ln_post = nn.LayerNorm(width)
+        if output_dim is not None:
+            self.proj = nn.Parameter(torch.empty(width, output_dim))
+        else:
+            self.proj = None
+        
+        self.dropout = nn.Dropout(dropout)
+
+    def get_num_layers(self):
+        return len(self.transformer.resblocks)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'positional_embedding', 'class_embedding', 'temporal_positional_embedding'}
+    
+    def mask_tokens(self, inputs, masking_prob=0.0):
+        B, L, _ = inputs.shape
+
+        # This is different from text as we are masking a fix number of tokens
+        Lm = int(masking_prob * L)
+        masked_indices = torch.zeros(B, L)
+        indices = torch.argsort(torch.rand_like(masked_indices), dim=-1)[:, :Lm]
+        batch_indices = (
+            torch.arange(masked_indices.shape[0]).unsqueeze(-1).expand_as(indices)
+        )
+        masked_indices[batch_indices, indices] = 1
+
+        masked_indices = masked_indices.bool()
+
+        return inputs[~masked_indices].reshape(B, -1, inputs.shape[-1])
+
+    def forward(self, x, masking_prob=0.0):
+        x = self.conv1(x)  # shape = [*, width, grid, grid]
+        B, C, T, H, W = x.shape
+        x = x.permute(0, 2, 3, 4, 1).reshape(B * T, H * W, C)
+
+        x = torch.cat([self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x], dim=1)  # shape = [*, grid ** 2 + 1, width]
+        x = x + self.positional_embedding.to(x.dtype)
+
+        # temporal pos
+        cls_tokens = x[:B, :1, :]
+        x = x[:, 1:]
+        x = rearrange(x, '(b t) n m -> (b n) t m', b=B, t=T)
+        if hasattr(self, 'temporal_positional_embedding'):
+            if x.size(1) == 1:
+                # This is a workaround for unused parameter issue
+                x = x + self.temporal_positional_embedding.mean(1)
+            else:
+                x = x + self.temporal_positional_embedding
+        x = rearrange(x, '(b n) t m -> b (n t) m', b=B, t=T)
+
+        if masking_prob > 0.0:
+            x = self.mask_tokens(x, masking_prob)
+
+        x = torch.cat((cls_tokens, x), dim=1)
+
+        x = self.ln_pre(x)
+
+        x = x.permute(1, 0, 2)  #BND -> NBD
+        x = self.transformer(x)
+
+        x = self.ln_post(x)
+
+        if self.proj is not None:
+            x = self.dropout(x[0]) @ self.proj
+        else:
+            x = x.permute(1, 0, 2)  #NBD -> BND
+
+        return x
+
+
+def inflate_weight(weight_2d, time_dim, center=True):
+    logger.info(f'Init center: {center}')
+    if center:
+        weight_3d = torch.zeros(*weight_2d.shape)
+        weight_3d = weight_3d.unsqueeze(2).repeat(1, 1, time_dim, 1, 1)
+        middle_idx = time_dim // 2
+        weight_3d[:, :, middle_idx, :, :] = weight_2d
+    else:
+        weight_3d = weight_2d.unsqueeze(2).repeat(1, 1, time_dim, 1, 1)
+        weight_3d = weight_3d / time_dim
+    return weight_3d
+
+
+def load_state_dict(model, state_dict, input_resolution=224, patch_size=16, center=True):
+    state_dict_3d = model.state_dict()
+    for k in state_dict.keys():
+        if k in state_dict_3d.keys() and state_dict[k].shape != state_dict_3d[k].shape:
+            if len(state_dict_3d[k].shape) <= 2:
+                logger.info(f'Ignore: {k}')
+                continue
+            logger.info(f'Inflate: {k}, {state_dict[k].shape} => {state_dict_3d[k].shape}')
+            time_dim = state_dict_3d[k].shape[2]
+            state_dict[k] = inflate_weight(state_dict[k], time_dim, center=center)
+
+    pos_embed_checkpoint = state_dict['positional_embedding']
+    embedding_size = pos_embed_checkpoint.shape[-1]
+    num_patches = (input_resolution // patch_size) ** 2
+    orig_size = int((pos_embed_checkpoint.shape[-2] - 1) ** 0.5)
+    new_size = int(num_patches ** 0.5)
+    if orig_size != new_size:
+        logger.info(f'Pos_emb from {orig_size} to {new_size}')
+        extra_tokens = pos_embed_checkpoint[:1]
+        pos_tokens = pos_embed_checkpoint[1:]
+        pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
+        pos_tokens = torch.nn.functional.interpolate(
+            pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
+        pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(0, 2)
+        new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=0)
+        state_dict['positional_embedding'] = new_pos_embed
+    
+    message = model.load_state_dict(state_dict, strict=False)
+    logger.info(f"Load pretrained weights: {message}")
+
+
+@register_model
+def clip_joint_b16(
+    pretrained=False, input_resolution=224, kernel_size=1,
+    center=True, num_frames=8, drop_path=0., checkpoint_num=0,
+    dropout=0.,
+):
+    model = VisionTransformer(
+        input_resolution=input_resolution, patch_size=16, 
+        width=768, layers=12, heads=12, output_dim=512,
+        kernel_size=kernel_size, num_frames=num_frames, 
+        drop_path=drop_path, checkpoint_num=checkpoint_num,
+        dropout=dropout,
+    )
+    if pretrained:
+        if isinstance(pretrained, str):
+            model_name = pretrained
+        else:
+            model_name = "ViT-B/16"
+        
+        logger.info('load pretrained weights')
+        state_dict = torch.load(_MODELS[model_name], map_location='cpu')
+        load_state_dict(model, state_dict, input_resolution=input_resolution, patch_size=16, center=center)
+    return model.eval()
+
+
+@register_model
+def clip_joint_l14(
+    pretrained=False, input_resolution=224, kernel_size=1,
+    center=True, num_frames=8, drop_path=0., checkpoint_num=0,
+    dropout=0.,
+):
+    model = VisionTransformer(
+        input_resolution=input_resolution, patch_size=14,
+        width=1024, layers=24, heads=16, output_dim=768,
+        kernel_size=kernel_size, num_frames=num_frames, 
+        drop_path=drop_path, checkpoint_num=checkpoint_num,
+        dropout=dropout,
+    )
+    
+    if pretrained:
+        if isinstance(pretrained, str):
+            model_name = pretrained
+        else:
+            model_name = "ViT-L/14"
+        logger.info('load pretrained weights')
+        state_dict = torch.load(_MODELS[model_name], map_location='cpu')
+        load_state_dict(model, state_dict, input_resolution=input_resolution, patch_size=14, center=center)
+    return model.eval()
+
+
+@register_model
+def clip_joint_l14_336(
+    pretrained=True, input_resolution=336, kernel_size=1,
+    center=True, num_frames=8, drop_path=0.
+):
+    raise NotImplementedError
+    model = VisionTransformer(
+        input_resolution=input_resolution, patch_size=14, 
+        width=1024, layers=24, heads=16, output_dim=768,
+        kernel_size=kernel_size, num_frames=num_frames,
+        drop_path=drop_path,
+    )
+    if pretrained:
+        logger.info('load pretrained weights')
+        state_dict = torch.load(_MODELS["ViT-L/14_336"], map_location='cpu')
+        load_state_dict(model, state_dict, input_resolution=input_resolution, patch_size=14, center=center)
+    return model.eval()
+
+
+def interpolate_pos_embed_vit(state_dict, new_model):
+    key = "vision_encoder.temporal_positional_embedding"
+    if key in state_dict:
+        vision_temp_embed_new = new_model.state_dict()[key]
+        vision_temp_embed_new = vision_temp_embed_new.unsqueeze(2)  # [1, n, d] -> [1, n, 1, d]
+        vision_temp_embed_old = state_dict[key]
+        vision_temp_embed_old = vision_temp_embed_old.unsqueeze(2)
+
+        state_dict[key] = load_temp_embed_with_mismatch(
+            vision_temp_embed_old, vision_temp_embed_new, add_zero=False
+        ).squeeze(2)
+
+    key = "text_encoder.positional_embedding"
+    if key in state_dict:
+        text_temp_embed_new = new_model.state_dict()[key]
+        text_temp_embed_new = text_temp_embed_new.unsqueeze(0).unsqueeze(2)  # [n, d] -> [1, n, 1, d]
+        text_temp_embed_old = state_dict[key]
+        text_temp_embed_old = text_temp_embed_old.unsqueeze(0).unsqueeze(2)
+
+        state_dict[key] = load_temp_embed_with_mismatch(
+            text_temp_embed_old, text_temp_embed_new, add_zero=False
+        ).squeeze(2).squeeze(0)
+    return state_dict

utils/vision_encoder/model_vision_encoder.py

@@ -0,0 +1,84 @@
+import os
+import logging
+import torch
+from torch import nn
+import math
+
+from .clip_vision import clip_joint_l14, clip_joint_b16
+
+logger = logging.getLogger(__name__)
+
+
+class VisionEncoder(nn.Module):
+
+    def __init__(self):
+        super(VisionEncoder, self).__init__()
+
+        self.vision_encoder_name = 'vit_l14'
+        self.vision_encoder_pretrained = False
+        self.inputs_image_res = 224
+        self.vision_encoder_kernel_size = 1
+        self.vision_encoder_center = True
+        self.video_input_num_frames = 8
+        self.vision_encoder_drop_path_rate = 0.1
+        self.vision_encoder_checkpoint_num = 24
+
+        self.vision_width = 1024
+        self.embed_dim = 768 
+        self.masking_prob = 0.9
+
+        self.vision_encoder = self.build_vision_encoder()
+
+        self.temp = nn.parameter.Parameter(torch.ones([]) * 1 / 100.0)
+        self.temp_min = 1 / 100.0
+
+    def no_weight_decay(self):
+        ret = {"temp"}
+        ret.update(
+            {"vision_encoder." + k for k in self.vision_encoder.no_weight_decay()}
+        )
+
+        return ret
+
+
+    def encode_vision(self, image, test=False):
+        if image.ndim == 5:
+            image = image.permute(0, 2, 1, 3, 4).contiguous()
+        else:
+            image = image.unsqueeze(2)
+
+        if not test and self.masking_prob > 0.0:
+            return self.vision_encoder(
+                image, masking_prob=self.masking_prob
+            )
+
+        return self.vision_encoder(image)
+
+
+    @torch.no_grad()
+    def clip_contrastive_temperature(self, min_val=0.001, max_val=0.5):
+        """Seems only used during pre-training"""
+        self.temp.clamp_(min=self.temp_min)
+
+    def build_vision_encoder(self):
+        """build vision encoder
+        Returns: (vision_encoder, vision_layernorm). Each is a `nn.Module`.
+
+        """
+        vision_encoder = clip_joint_l14(
+            pretrained=self.vision_encoder_pretrained,
+            input_resolution=self.inputs_image_res,
+            kernel_size=self.vision_encoder_kernel_size,
+            center=self.vision_encoder_center,
+            num_frames=self.video_input_num_frames,
+            drop_path=self.vision_encoder_drop_path_rate,
+            checkpoint_num=self.vision_encoder_checkpoint_num,
+        )
+
+        return vision_encoder
+
+
+    def get_vid_features(self, input_frames):
+        clip_feat = self.encode_vision(input_frames, test=True).float()
+        
+        return clip_feat