chore: add project files

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+figures/*.jpg filter=lfs diff=lfs merge=lfs -text
+figures/*.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,169 @@
+# Project Specific
+data/
+data
+datalink
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+**private/*
+**ego_base_link/*
+**.vscode/*
+
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/

LLM_stage.py

@@ -0,0 +1,168 @@
+"""
+Evaluate a model on the egoschema dataset using LVNet (captions pre-generated)
+
+Sample Run:
+
+python3 LLM_stage.py \
+  --output-dir ego_base_link \
+  --captions data/ES_captions_gpt4o.jsonl \
+  --per-vid-captions 12 \
+  --gptmodel "gpt-4o" \
+  --temperature 0.0
+"""
+
+import argparse
+import json
+import os
+
+from tqdm import tqdm
+
+from src.run_gpt import run_gpt
+
+# You may add multiple keys to run parallel calls
+dict_api = {
+    "api_key": "ADD",
+}
+
+
+_PROMPT_TEMPLATE = (
+    "Here are descriptions of the video frames at specific times, noted in seconds."
+    "\n\n{Putdesc}.\n\nThe descriptions of the frames conclude. Think step-by-step"
+    " and I request your selection of the most appropriate response to the following"
+    " question\n\nQuestion:\n{Putquestion}\n\nOptions:\n{AllOptions}"
+)
+
+
+def eval_model(args):
+    # change split to split
+    captions_path, data_path, split, gptmodel, temp, base_dir, job_name = (
+        args.captions,
+        args.data,
+        args.per_vid_captions,
+        args.gptmodel,
+        args.temperature,
+        args.output_dir,
+        args.job_name,
+    )
+
+    prompt = _PROMPT_TEMPLATE
+
+    os.makedirs(base_dir, exist_ok=True)
+    output_dir = f"{base_dir}/egoschema/{job_name}"
+    output_dir = os.path.expanduser(output_dir)
+    os.makedirs(output_dir, exist_ok=True)
+
+    save_name = captions_path.rsplit("/", 2)[-1].replace(".jsonl", "")
+    output_summary_path = f"{output_dir}/{save_name}.jsonl"
+    print(f"Saving outputs to:{output_summary_path}")
+    output_summary = open(output_summary_path, "w")
+
+    input_summary = [
+        json.loads(q) for q in open(os.path.expanduser(captions_path), "r")
+    ]
+    dataset = json.load(open(os.path.expanduser(data_path), "r"))
+    input_len = len(input_summary)
+    assert (
+        input_len % split == 0
+    ), f"input_len%split:{input_len%split}, input_len:{input_len}, split:{split}"
+    groups = input_len // split
+
+    final_prompts = []
+    final_info = []
+    for i in tqdm(range(groups)):
+        sidx = i * split
+        eidx = (i + 1) * split
+
+        desc = ""
+        timeline = []
+        for idx, e in enumerate(input_summary[sidx:eidx]):
+            cur_data = dataset[e["q_uid"]]
+            desc += e["answer"] + " "
+            timeline.append(e["timeline"])
+
+            if idx == split - 1:  # the last of split
+                action_0 = cur_data["option 0"]
+                action_1 = cur_data["option 1"]
+                action_2 = cur_data["option 2"]
+                action_3 = cur_data["option 3"]
+                action_4 = cur_data["option 4"]
+
+                option_list = ""
+                option_number_candidate = ["one", "two", "three", "four", "five"]
+                option_number = option_number_candidate[4]
+                AllOptNumber = "option 0, option 1, option 2, option 3, option 4"
+                FocusOptions = ""
+
+                alloptions = f"option 0: {action_0}\noption 1: {action_1}\noption 2: {action_2}\noption 3: {action_3}\noption 4: {action_4}"
+                option_list = f"option 0: {action_0}\noption 1: {action_1}\noption 2: {action_2}\noption 3: {action_3}\noption 4: {action_4}"
+
+                FocusOptions += "option 0, option 1, option 2, option 3, option 4"
+
+                question = cur_data["question"]
+
+                curr_prompt = (
+                    prompt.replace("{Putdesc}", desc)
+                    .replace("{Putquestion}", question)
+                    .replace("{Putoptions}", option_list)
+                    .replace("{PutOptNumber}", option_number)
+                    .replace("{FocusOptions}", FocusOptions)
+                    .replace("{AllOptions}", alloptions)
+                    .replace("{PutAllOptNumber}", AllOptNumber)
+                )
+
+                final_prompts.append(curr_prompt)
+
+                CA_option = {}
+                if "CA" in cur_data:
+                    CA_option = {"CA": cur_data["CA"]}
+
+                info = {
+                    "q_uid": e["q_uid"],
+                    "prompt": curr_prompt,
+                    "timeline": timeline,
+                    "question": question,
+                    "option 0": action_0,
+                    "option 1": action_1,
+                    "option 2": action_2,
+                    "option 3": action_3,
+                    "option 4": action_4,
+                } | CA_option
+
+                final_info.append(info)
+
+    output_VLM = run_gpt(
+        texts=final_prompts,
+        api_keys=list(dict_api.values()),
+        max_tokens=2000,
+        model=gptmodel,
+        temperature=temp,
+        num_threads=20,  # Tune this
+        backoff_time=1 * 60,
+        silent=False,
+        dataset="egoschema",
+    )
+
+    output_VLM = list(output_VLM)
+
+    for q_idx, info in enumerate(tqdm(final_info)):  # prompt_list = # Q&C
+        info["answer"] = output_VLM[q_idx]
+        output_summary.write(json.dumps(info) + "\n")
+
+    # finish the summarization for the current question
+    output_summary.close()
+    print(f"output_summary_path:{output_summary_path}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--output-dir", type=str)
+    parser.add_argument("--job-name", type=str, default="run001")
+    parser.add_argument("--captions", type=str, default="data/ES_captions_gpt4o.jsonl")
+    parser.add_argument("--data", type=str, default="data/ES_qa_data.json")
+    parser.add_argument("--per-vid-captions", type=int, default=12)
+    parser.add_argument("--gptmodel", type=str, default="gpt-3.5-turbo-1106")
+    parser.add_argument("--temperature", type=float, default=None)
+
+    args = parser.parse_args()
+
+    eval_model(args)

README.md

@@ -1,3 +1,122 @@
 ---
 license: apache-2.0
+tags:
+- computer-vision
+- video-question-answering
+- zero-shot
+- 9 pages workshop at neurips2024
 ---
+
+# LVNet
+
+[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/too-many-frames-not-all-useful-efficient/zero-shot-video-question-answer-on-egoschema-1)](https://paperswithcode.com/sota/zero-shot-video-question-answer-on-egoschema-1?p=too-many-frames-not-all-useful-efficient)
+[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/too-many-frames-not-all-useful-efficient/zero-shot-video-question-answer-on-intentqa)](https://paperswithcode.com/sota/zero-shot-video-question-answer-on-intentqa?p=too-many-frames-not-all-useful-efficient)
+[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/too-many-frames-not-all-useful-efficient/zero-shot-video-question-answer-on-next-qa)](https://paperswithcode.com/sota/zero-shot-video-question-answer-on-next-qa?p=too-many-frames-not-all-useful-efficient)
+[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/too-many-frames-not-all-useful-efficient/zero-shot-video-question-answer-on-egoschema)](https://paperswithcode.com/sota/zero-shot-video-question-answer-on-egoschema?p=too-many-frames-not-all-useful-efficient)
+
+Official Code for **_Too Many Frames, Not All Useful_: Efficient Strategies for Long-Form Video QA**  
+Accepted to the 9 pages Workshop on Video-Language Models at *NeurIPS 2024*
+
+[**Paper Link**](https://arxiv.org/abs/2406.09396)
+
+---
+
+## Abstract
+
+Long-form videos that span wide temporal intervals are highly information-redundant and contain multiple distinct events or entities that are often loosely-related. Therefore, when performing long-form video question answering (LVQA), all information necessary to generate a correct response can often be contained within a small subset of frames. Recent literature explores the use of large language models (LLMs) in LVQA benchmarks, achieving exceptional performance, while relying on vision language models (VLMs) to convert all visual content within videos into natural language. Such VLMs often independently caption a large number of frames uniformly sampled from long videos, which is not efficient and can mostly be redundant.
+
+Questioning these decision choices, we explore optimal strategies for key-frame selection and sequence-aware captioning that can significantly reduce these redundancies. We propose two novel approaches that improve each aspect, namely **Hierarchical Keyframe Selector** and **Sequential Visual LLM**. Our resulting framework, called **LVNet**, achieves state-of-the-art performance across three benchmark LVQA datasets.
+
+---
+
+## Accuracy vs. Captions on the EgoSchema Subset
+
+- LVNet shows a SOTA **68.2% accuracy**, using only **12** captions.  
+- This highlights the quality of keyframes from the Hierarchical Keyframe Selector.
+
+<img src="figures/hkf_graph.png" alt="acc_captions" width="600"/>
+
+---
+
+## Hierarchical Keyframe Selector: Structural Overview
+
+- **Overall strategy**: Generate captions via a Hierarchical Keyframe Selector, then feed them to a separate LLM to answer the question.
+- **Temporal Scene Clustering (TSC)**: Divides the long video into multiple scenes, enabling per-scene subsampling.
+- **Coarse Keyframe Detector (CKD)**: Selects frames best-aligned with keywords relevant to the query.
+- **Fine Keyframe Detector (FKD)**: Refines the keyword alignments within a smaller set of frames via templated visual prompting.
+
+<img src="figures/architecture.png" alt="architecture" width="600"/>
+
+---
+
+## Operational Visualization of HKS
+
+- **Temporal Scene Clustering (TSC)**: 900 frames get clustered into scenes, then uniformly subsampled to produce about 280 frames.  
+- **Coarse Keyframe Detector (CKD)**: Among those, 32 frames are selected, based on alignment with query keywords.  
+- **Visual Templating**: The coarsely refined keyframes are ordered by confidence and temporal order, grouped into 4 chunks of 8 frames.  
+- **Fine Keyframe Detector (FKD)**: Selects the final 12 frames via further keyword alignment checks.
+
+<img src="figures/architecture_qualitative.png" alt="hks_visualization" width="800"/>
+
+---
+
+## Experiments: EgoSchema, NExT-QA, and IntentQA
+
+- LVNet achieves **61.1%**, **72.9%**, and **71.7%** on the three benchmarks, respectively, using **just 12** frames—on par with or exceeding models that use **many** more frames.  
+- Models with specialized video-caption pretraining or significantly more captions are shown in gray/light green for fairness comparison.
+
+<img src="tables/table_combined.png" alt="egoschema_table" width="900"/>
+
+---
+
+## Comparison with Other Keyframe Selection Methods
+
+Below is a side-by-side comparison of **LVNet** and **VideoAgent**.  
+- **LVNet** starts with uniform sampling but then refines keyframes via TSC, CKD, and FKD. This yields 12 frames, 8 of which show “phone usage,” the correct activity.  
+- **VideoAgent** continues uniform sampling due to insufficient initial frames, resulting in 0 relevant frames out of 9 and an incorrect final answer.
+
+<img src="figures/KFSelectionFlowComparison.jpg" alt="kf_selection_flow" width="900"/>
+
+---
+
+## Evaluation
+
+### Generating Answers Using LLM
+
+You can quickly run the LLM to produce answers once you have the keyframe-based captions:
+
+1. **Download the Captions for Dataset**
+
+* EgoSchema: `bash scripts/get_ES_captions.sh `
+
+2. **Run LLM** `bash scripts/eval_ES.sh`
+
+### Generate captions using our provided modules
+#### Hierarchical Keyframe Selector (HKS)
+- Temporal Scene Clustering (TSC): temporalSceneClustering.py </br>
+- Coarse Keyframe Detector (CKD): coarseKeyframeDetector.py </br>
+- Fine Keyframe detector (FKD): fineKeyframeDetector.py </br>
+
+1. **EgoSchema keyframe selection from images**: `bash config/run.sh `
+
+2. **Generate captions based on the keyframes**: `bash scripts/create_caption.sh`
+
+## Data
+### Hierarchical Keyframe Selector hyper-parameters & paths
+- [[LINK]](config/config.py)
+
+### coarseKeyframeDetector.py CLIP model checkpoint
+- ICCV 2023 [Perceptual Grouping in Contrastive Vision-Language Models](https://arxiv.org/abs/2210.09996)
+- Checkpoint: [Download](https://github.com/kahnchana/clippy/releases/download/v1.0/clippy_5k.pt)
+
+
+# Citation
+```
+@inproceedings{Park2024TooMF,
+  title={Too Many Frames, not all Useful: Efficient Strategies for Long-Form Video QA},
+  author={Jongwoo Park and Kanchana Ranasinghe and Kumara Kahatapitiya and Wonjeong Ryoo and Donghyun Kim and Michael S. Ryoo},
+  year={2024}
+}
+```
+
+For more details and updates, please see our [GitHub Repository](https://github.com/jongwoopark7978/LVNet).

VLM_stage.py

@@ -0,0 +1,154 @@
+import os
+import json
+import base64
+import random
+import argparse
+
+import natsort
+
+from PIL import Image
+from tqdm import tqdm
+
+import torch
+from torch.utils.data import Dataset, DataLoader
+
+from src.run_gpt import run_gpt
+
+random.seed(10)
+dict_api = {
+    "api_key":"ADD",
+}
+
+
+class CustomDatasetGPT(Dataset):
+    def __init__(self, questions, num_kf):
+        self.questions = questions
+        self.num_kf = num_kf
+
+    def __getitem__(self, index):
+        line = self.questions[index]
+        group = 4
+        newnum_per_group = self.num_kf // group
+        oldnum_per_group = len(line["VLM_path"]) // group
+        assert oldnum_per_group >= newnum_per_group, f"oldnum_per_group:{oldnum_per_group} is smaller than newnum_per_group:{newnum_per_group}"
+
+        new_kf_paths = []
+        new_kf_timelines = []
+        for i in range(group):
+            start_index = i * oldnum_per_group
+            end_index = start_index + oldnum_per_group
+
+            sub_kf_paths = line["VLM_path"][start_index:min(end_index, len(line["VLM_path"]))]
+            sub_kf_timelines = line["VLM_timeline"][start_index:min(end_index, len(line["VLM_timeline"]))]
+            new_kf_paths.extend(sub_kf_paths[:newnum_per_group])
+            new_kf_timelines.extend(sub_kf_timelines[:newnum_per_group])
+
+        kf_paths = natsort.natsorted(new_kf_paths)
+        kf_timelines = natsort.natsorted(new_kf_timelines)
+
+        images = []
+        images_base64 = []
+
+        for e in kf_paths:
+            images.append(Image.open(e).convert('RGB'))
+            images_base64.append(encode_image(e))
+
+        return images_base64, kf_paths, kf_timelines
+
+    def __len__(self):
+        return len(self.questions)
+
+
+def encode_image(image_path):
+    with open(image_path, "rb") as image_file:
+        return base64.b64encode(image_file.read()).decode('utf-8')
+
+def create_data_loader_gpt(questions, num_kf, batch_size=1, num_workers=4):
+    assert batch_size == 1, "batch_size must be 1"
+
+    dataset = CustomDatasetGPT(questions, num_kf)
+    data_loader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers, shuffle=False)
+
+    return data_loader, dataset
+
+def eval_model(args):
+    base_dir, question_path, vlm, num_kf, temp = (
+        args.output_dir,
+        args.question_path,
+        args.gptmodel,
+        args.num_kf,
+        args.temp,
+    )
+
+    questions = [json.loads(q) for q in open(os.path.expanduser(question_path), "r")]
+
+    fname = question_path.split('/')[-1]
+    answer_path = f"{base_dir}/egoschema/{num_kf}/{fname}"
+    os.makedirs(os.path.dirname(answer_path), exist_ok=True)
+    print(f"question_path:{question_path}\nanswer_path:{answer_path}")
+
+    ans_file = open(answer_path, "w")
+    data_loader, dataset = create_data_loader_gpt(questions, num_kf)
+
+    for (base64_image, kf_paths, kf_timelines), line in tqdm(zip(data_loader, questions), total=len(questions)):
+        idx = line["q_uid"]
+        CA = line["CA"] if "CA" in line else None
+        option0 = line['option 0']
+        option1 = line['option 1']
+        option2 = line['option 2']
+        option3 = line['option 3']
+        option4 = line['option 4']
+        question = line['question']
+
+        lenwords = "50"
+        prompt = f"'C' stands for the cameraman. Describe the activity depicted in this first-person perspective image in less than {lenwords} words. In your answer, don't mention that the image is in first-person perspective, as we already know this."
+        prompts = [prompt] * num_kf
+        
+        image_paths = [e[0] for e in kf_paths]
+        image_timelines = [e[0] for e in kf_timelines]
+
+        output_VLM = run_gpt(
+                images=image_paths,
+                texts=prompts,
+                api_keys=list(dict_api.values()),
+                max_tokens=2000,
+                model=vlm,
+                temperature=temp,
+                num_threads=20,  # Tune this
+                backoff_time=1 * 60,
+                silent=False,
+                dataset="egoschema",
+                verbose=False,
+        )
+
+        output_VLM = list(output_VLM)
+
+        for j, e in enumerate(image_timelines):
+            line_frame = line.copy()
+            line_frame["answer"] = f"At {str(e)} seconds, {output_VLM[j]}"
+            line_frame["AR-VLM_model_id"] = vlm
+            line_frame["AR-VLM_prompt"] = prompts[j]
+            line_frame["timeline"] = float(e)
+            line_frame["frame_idx"] = j
+            line_frame["image_paths"] = image_paths
+
+            if "imgidx_kw_dict" in line_frame.keys(): line_frame.pop("imgidx_kw_dict")
+            if "google_drive_id" in line_frame.keys(): line_frame.pop("google_drive_id")
+
+            ans_file.write(json.dumps(line_frame)+"\n")
+
+    print(f"question.\nquestion_path:{question_path}\nanswer_path:{answer_path}")
+
+    ans_file.close()
+    return "job is done"
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--output-dir", type=str)
+    parser.add_argument("--question-path", type=str, default="")
+    parser.add_argument("--num-kf", type=int) 
+    parser.add_argument("--gptmodel", type=str, default="gpt-4o")
+    parser.add_argument("--temp", type=float, default=None)
+    args = parser.parse_args()
+    eval_model(args)

coarseKeyframeDetector.py

@@ -0,0 +1,254 @@
+import os
+import json
+import shutil
+
+from tqdm import tqdm
+from PIL import Image
+
+import natsort
+
+import numpy as np
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+
+from config import config
+from src.open_clip import create_model_and_transforms
+
+
+class loading_img(Dataset):
+    def __init__(self, img_list):
+        self.img_list = img_list
+
+    def __len__(self):
+        return len(self.img_list)
+
+    def __getitem__(self, idx):
+        return self.img_list[idx].squeeze(0)
+
+class CustomDataset(Dataset):
+    def __init__(self, questions, clippy, preprocess_val, clip_size, base_dir):
+        self.questions      = questions
+        self.clippy         = clippy
+        self.clip_size      = clip_size
+        self.preprocess_val = preprocess_val
+        self.device         = next(clippy.parameters()).device
+        self.base_dir       = base_dir
+
+    def __getitem__(self, index):
+        line = self.questions[index]
+        images_dir = f"{line['q_uid']}"
+
+        if line["Activity"] == "" or ("Activity" not in line): ref1 = []
+
+        else:
+            if isinstance(line["Activity"], list): ref1 = line["Activity"]
+            else: ref1 = line["Activity"].split(', ')
+        
+        keywords = ref1
+        clip_size = self.clip_size
+        clippy = self.clippy
+        preprocess_val = self.preprocess_val
+        
+        images = []
+        timelines = []
+        timelines_int = []
+        img_names = []
+        image_list = []
+
+        nframes_paths = line["filepath"]
+        total_len = len(nframes_paths)
+        nframes_paths = natsort.natsorted(nframes_paths)
+
+        img_paths = []
+        for img_path in nframes_paths:
+            img_path = self.base_dir + "/" + "/".join(img_path.split("/")[-4:])
+            img_paths.append(img_path)
+
+            img_names.append(img_path.split('/')[-1].split('.')[0])
+            cur_img = Image.open(img_path).resize(clip_size)
+            image_list.append(preprocess_val(cur_img))
+
+            timeline = f"{img_names[-1].split('_')[-2]}.{img_names[-1].split('_')[-1]} seconds"
+            timeline_int = float(f"{img_names[-1].split('_')[-2]}.{img_names[-1].split('_')[-1]}")
+            timelines.append(timeline)
+            timelines_int.append(timeline_int)
+
+        return image_list, img_paths, timelines, timelines_int, keywords, img_names
+
+    def __len__(self):
+        return len(self.questions)
+
+
+def disable_torch_init():
+    setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
+    setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
+
+def SortSimilarity(q_uid, simmat, keywords, nimgtokens, nframes_paths, maximgslen):
+    sort_simmat, sort_idx = torch.sort(simmat, dim=-1, descending=True)
+    sort_idx = torch.floor(sort_idx/nimgtokens).to(int)
+
+    curimgslen = 0
+
+    imgidx_kw_dict = dict()
+    numrow, numcol = sort_simmat.shape
+    
+    row_col_list = [0 for _ in range(numrow)]
+    token = True
+
+    while token:
+        j = 0
+        while j < numrow:
+            k = 0
+            i = row_col_list[j]
+
+            while k < numcol-i:
+                col_idx = i+k
+                k += 1
+
+                simvalue = sort_simmat[j, col_idx].item()
+                img_idx = sort_idx[j, col_idx].item()
+
+                curr_keyword = keywords[j]
+                curr_kfpath = nframes_paths[img_idx]
+
+                if img_idx in imgidx_kw_dict: continue
+
+                else:
+                    imgidx_kw_dict[img_idx] = {"kw": curr_keyword, "simvalue": simvalue, "kf_path": curr_kfpath, "kw_others": []}
+                    curimgslen += 1
+
+                    row_col_list[j] = col_idx + 1
+                    if curimgslen == maximgslen: return imgidx_kw_dict
+                    else: break
+
+            j += 1
+
+        if sum(row_col_list) >= numrow*(numcol-1): token = False
+
+def create_data_loader(questions, clippy, preprocess_val, clip_size, base_dir, batch_size=1, num_workers=16):
+    assert batch_size == 1, "batch_size must be 1"
+    dataset = CustomDataset(questions, clippy, preprocess_val, clip_size, base_dir)
+    data_loader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers, shuffle=False)
+    return data_loader
+
+def eval_model():
+    disable_torch_init()
+    question_path, maximgslen, base_dir, concatname, modelpath, answerpath, concatdir = config.question_path, config.maximgslen, config.base_dir, config.concatname, config.modelpath, config.answerpath, config.concatdir
+
+    pretrained_ckpt = f"{modelpath}"
+    clippy, preprocess_train, preprocess_val = create_model_and_transforms(
+            "clippy-B-16",
+            device="cuda",
+            pretrained=pretrained_ckpt
+    )
+    clip_size = (224,224)
+    device = next(clippy.parameters()).device
+
+    questions = [json.loads(q) for q in open(os.path.expanduser(question_path), "r")] 
+
+    answer_path = f"{answerpath}"
+    print(f"\nquestion_path:{question_path}\nanswer_path:{answer_path}")
+    os.makedirs(os.path.dirname(answer_path), exist_ok=True)
+
+    with open(answer_path, "w") as ans_file:
+        data_loader = create_data_loader(questions, clippy, preprocess_val, clip_size, base_dir)
+        concatimg_dir_base = f"{concatdir}"
+
+        with torch.no_grad():
+            for (image_list, nframes_paths, timelines, timelines_int, keywords, img_names), line in tqdm(zip(data_loader, questions), total=len(questions)):
+                q_uid = line["q_uid"]
+                CA = line["CA"] if "CA" in line else None
+                option0 = line['option 0']
+                option1 = line['option 1']
+                option2 = line['option 2']
+                option3 = line['option 3']
+                option4 = line['option 4']
+                question = line['question']
+
+                pastobj = None
+                past_VLM_path = None
+                past_VLM_timeline = None
+
+                img_embed = []
+                nframes_paths = [e[0] for e in nframes_paths]
+
+                image_set = loading_img(image_list)
+                image_loader = DataLoader(image_set, batch_size=64, shuffle=False, num_workers=16)
+                for e in image_loader: img_embed.append(clippy.encode_image(e.to(device), pool=False)[:, 1:])
+                img_embed = torch.concat(img_embed, dim=0)
+
+                limit_keywords = config.limit_keywords
+                keywords = [e[0] for e in keywords][:limit_keywords]
+                keyword_embed = clippy.text.encode(keywords, convert_to_tensor=True)
+
+                nframe, nimgtokens, channels = img_embed.shape
+                keyword_embed = keyword_embed.unsqueeze(1)
+                img_embed = img_embed.flatten(0, 1).unsqueeze(0) 
+
+                simmat = F.cosine_similarity(keyword_embed, img_embed, dim=-1).to(torch.float)
+                imgidx_kw_dict = SortSimilarity(q_uid, simmat, keywords, nimgtokens, nframes_paths, maximgslen=maximgslen)
+
+                # order of simvalue
+                simvalue = np.array([e["simvalue"] for e in imgidx_kw_dict.values()])
+                ordered_idx = np.argsort(simvalue)
+                simvalue = simvalue[ordered_idx]
+                kf_paths = np.array([e["kf_path"] for e in imgidx_kw_dict.values()])[ordered_idx]
+                matchingkw = np.array([e["kw"] for e in imgidx_kw_dict.values()])[ordered_idx]
+
+                #order by timeline
+                time_kf_paths = np.array(kf_paths[:16])
+                timelines_int = np.array([float(f"{e.replace('.jpg', '').split('/')[-1].split('_')[1]}" + "."+ f"{e.replace('.jpg', '').split('/')[-1].split('_')[2]}") for e in time_kf_paths])
+                time_ordered_idx = np.argsort(timelines_int)
+
+                timelines_int = timelines_int[time_ordered_idx]
+                time_simvalue = np.array(simvalue[:16])[time_ordered_idx]
+                time_kf_paths = np.array(time_kf_paths)[time_ordered_idx]
+                time_matchingkw = np.array(matchingkw[:16])[time_ordered_idx]
+
+                simvalue[:16] = time_simvalue
+                kf_paths[:16] = time_kf_paths
+                matchingkw[:16] = time_matchingkw
+
+                segment_timeline = f"{timelines[0][0].split(' seconds')[0]}-{timelines[-1][0].split(' seconds')[0]}"
+
+                imgw, imgh = Image.open(kf_paths[0]).size
+                redwidth = 20
+                newimgw, newimgh = (imgw+redwidth) * 4 + redwidth, (imgh+redwidth) * 2 + redwidth
+                concatimg = np.zeros((newimgh, newimgw, 3), dtype=np.uint8)
+                concatimg[:, :, 0] = 255
+                concatimglist = []
+                concatimg_dir = f"{concatimg_dir_base}/{q_uid}"
+
+                for i, cpath in enumerate(kf_paths):
+                    cur_img = np.array(Image.open(cpath))
+                    whole_frame = 8
+                    remainder =  i % whole_frame
+                    rowremainder = i % (whole_frame//2)
+                    startwidth = redwidth + (imgw + redwidth)*rowremainder
+                    endwidth = startwidth + imgw
+
+                    if remainder / whole_frame < 0.5: concatimg[redwidth:redwidth+imgh, startwidth:endwidth, :] = cur_img
+                    else: concatimg[redwidth+imgh+redwidth:newimgh-redwidth, startwidth:endwidth, :] = cur_img
+
+                    if remainder == whole_frame - 1: concatimglist.append(Image.fromarray(concatimg))
+
+                if os.path.exists(concatimg_dir): shutil.rmtree(concatimg_dir)
+                os.makedirs(f"{concatimg_dir}", exist_ok=True)
+                for i, img in enumerate(concatimglist): img.save(f"{concatimg_dir}/concat_{i}.jpg")
+
+                line["kf_paths"] = kf_paths.tolist()
+                line["keywords"] = matchingkw.tolist()
+                line["simvalue"] = simvalue.tolist()
+                line["imgidx_kw_dict"] = imgidx_kw_dict
+                line["segment_timeline"] = segment_timeline
+                line["concatimg_dir"] = concatimg_dir
+
+                ans_file.write(json.dumps(line) + "\n")
+
+    print(f"question_path:{question_path}\nanswer_path:{answer_path}")
+
+
+if __name__ == "__main__":
+    eval_model()

config/config.py

@@ -0,0 +1,38 @@
+# base
+dict_api = {
+    "api_key":"ADD",
+}
+base_dir = "your_path" # your data path ex) img_folder_path/egoschema
+
+
+# scene clustering
+divlam     = 12
+f_path     = "your_path" # files from keywords dir
+q_path     = "your_path" # files from questions dir
+a_path     = "your_path"
+img_folder = "your_path" # your img folder path ex) img_folder_path/egoschema/frames_900_4531/q_uid/image_sec_millisec.jpg
+
+
+# coarse key frame detector
+maximgslen     = 32
+limit_keywords = 25
+concatname     = "LVnet"
+modelpath      = "your_path" # model path
+question_path  = "your_path" # recommend using the same path with scene clustering answer path
+answerpath     = f"{base_dir}/kwkfmatching/kf_{concatname}.jsonl"  # kwkfmatching is not necessary.
+concatdir      = f"{base_dir}/kwkfmatching/concatimg_{concatname}" # kwkfmatching is not necessary.
+
+
+# fine key frame detector
+kf_vlm            = "gpt-4o"
+kf_temp           = None
+kf_num_select     = 3
+kf_num_input_imgs = 32
+kf_question_path  = "your_path" # recommend using the same path with coarse key frame detector answer path
+kf_answer_path    = f"{base_dir}/kf_VLM/kf_VLM{kf_num_input_imgs}sel{kf_num_select}_{kf_question_path.split('/')[-1].split('.')[0]}.jsonl" # kf_VLM is not necessary.
+
+
+# fine key frame detector refine
+refine_num_group   = 4
+refine_kflen       = 12
+refine_output_path = f"{base_dir}/kf_VLM/refine/" + kf_answer_path.split('/')[-1] # kf_VLM is not necessary.

config/run.sh

@@ -0,0 +1,4 @@
+devnum=${1:-0}
+CUDA_VISIBLE_DEVICES=$devnum python3 temporalSceneClustering.py
+CUDA_VISIBLE_DEVICES=$devnum python3 coarseKeyframeDetector.py
+CUDA_VISIBLE_DEVICES=$devnum python3 fineKeyframeDetector.py

extractKeyword.py

@@ -0,0 +1,132 @@
+import argparse
+import json
+import os
+
+from tqdm import tqdm
+from src.run_gpt import run_gpt
+
+"""
+Extract keywords from the given question and options
+
+Sample Run
+python3 extractKeyword.py --output-dir ego_base_link --question questions/500questions.jsonl --gptmodel "gpt-4-1106-preview"
+
+"""
+
+
+# You may add multiple keys if you want parallel calls
+dict_api = {
+    "api_key": "ADD",
+}
+
+PROMPT = (
+    "Think step-by-step and for each option, identify all the specified activities. "
+    "Each description of activity should use active voice with plain verbs, contain fewer than six words, "
+    "and retains as many original terms from the options as possible.\n"
+    "Here are the options:\n\n"
+    "option 0: {Putop0}\n"
+    "option 1: {Putop1}\n"
+    "option 2: {Putop2}\n"
+    "option 3: {Putop3}\n"
+    "option 4: {Putop4}\n"
+    "option 5: {Putquestion}.\n"
+    "All the options were introduced. 'C' represents the camera operator in the options.  "
+    "Your answer should follow the JSON format shown below and should only include the JSON result. "
+    "Do not output any warnings or notes under any circumstances. Instead, adhere strictly to the provided JSON format example.\n"
+    "This is one example output format.\n"
+    "{\"option 0\": [\"plays soccer\", \"go to school\"], \"option 1\": [\"go to the gym\", \"go to school\"], "
+    "\"option 2\": [\"go to school\", \"dry hair\"], \"option 3\": [\"plays basketball\", \"look at the tree\"], "
+    "\"option 4\": [\"plays soccer\", \"drop the ball\"], \"option 5\": [\"turn the table\", \"go to school\"]}"
+)
+
+
+def main(args):
+    # 1. Create output directories
+    os.makedirs(args.output_dir, exist_ok=True)
+    job_dir = os.path.join(args.output_dir, "extractedKeywords")
+    os.makedirs(job_dir, exist_ok=True)
+
+
+    # 2. Build the output file name (based on --question)
+    question_file_name = os.path.basename(args.question).replace(".jsonl", "")
+    output_summary_path = os.path.join(job_dir, f"{question_file_name}.jsonl")
+    print(f"Saving outputs to: {output_summary_path}")
+
+    # 3. Read the question file
+    with open(os.path.expanduser(args.question), "r") as f:
+        question_data = [json.loads(line) for line in f]
+
+    # 4. Construct final prompts
+    final_prompts = []
+    final_info = []
+    for entry in tqdm(question_data, desc="Building prompts"):
+        q_uid = entry["q_uid"]
+        # Insert each option + question into the embedded prompt
+        cur_prompt = (
+            PROMPT
+            .replace("{Putop0}", entry["option 0"])
+            .replace("{Putop1}", entry["option 1"])
+            .replace("{Putop2}", entry["option 2"])
+            .replace("{Putop3}", entry["option 3"])
+            .replace("{Putop4}", entry["option 4"])
+            .replace("{Putquestion}", entry["question"])
+        )
+
+        final_prompts.append(cur_prompt)
+
+        # Track data for JSON output
+        info = {
+            "q_uid": q_uid,
+            "prompt": cur_prompt,
+            "option 0": entry["option 0"],
+            "option 1": entry["option 1"],
+            "option 2": entry["option 2"],
+            "option 3": entry["option 3"],
+            "option 4": entry["option 4"],
+            "question": entry["question"],
+        }
+
+        # Include ground-truth label if present
+        if "CA" in entry:
+            info["CA"] = entry["CA"]
+
+        final_info.append(info)
+
+    # 5. Call GPT
+    print("Sending prompts to GPT. This may take a while...")
+    output_results = run_gpt(
+        texts=final_prompts,
+        api_keys=list(dict_api.values()),
+        max_tokens=2000,
+        model=args.gptmodel,
+        temperature=args.temperature,
+        num_threads=5,    # Adjust as needed
+        backoff_time=60,   # Adjust as needed
+        silent=False,
+        dataset="extractKeyword",
+    )
+
+    output_results = list(output_results)
+
+    # 6. Save results
+    with open(output_summary_path, "w") as outfile:
+        for i, info in enumerate(final_info):
+            info["answer"] = output_results[i]
+            outfile.write(json.dumps(info) + "\n")
+
+    print(f"Done! Results written to {output_summary_path}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--output-dir", type=str, required=True,
+                        help="Directory to store the resulting JSONL file.")
+    parser.add_argument("--question", type=str, required=True,
+                        help="Path to the JSONL file with question data (e.g., 500questions.jsonl).")
+    parser.add_argument("--gptmodel", type=str, default="gpt-4-1106-preview",
+                        help="The GPT model to call.")
+    parser.add_argument("--temperature", type=float, default=None,
+                        help="Temperature parameter for GPT.")
+
+    args = parser.parse_args()
+    main(args)

fineKeyframeDetector.py

@@ -0,0 +1,189 @@
+import os
+import json
+import base64
+import natsort
+
+import numpy as np
+
+from PIL import Image
+from tqdm import tqdm
+
+import torch
+from torch.utils.data import Dataset, DataLoader
+
+from config import config
+from src.refine import refine_answer
+from src.run_gpt import run_gpt
+
+class CustomDatasetGPT(Dataset):
+    def __init__(self, questions, num_input_imgs, num_select):
+        self.questions      = questions
+        self.num_input_imgs = num_input_imgs
+        self.num_select     = num_select
+
+    def __getitem__(self, index):
+        line           = self.questions[index]
+        num_select     = self.num_select
+        num_input_imgs = self.num_input_imgs
+
+        giter      = 0
+        imgs_group = 8
+        num_groups = num_input_imgs//imgs_group
+
+        kf_paths      = line["kf_paths"]
+        keywords      = line["keywords"]
+        simvalue      = line["simvalue"]
+        concatimg_dir = line['concatimg_dir']
+
+        concatimg_paths = natsort.natsorted([f"{concatimg_dir}/{im}" for im in os.listdir(concatimg_dir) if "ipynb" not in im])
+
+        concatimages        = []
+        concatimages_base64 = []
+        qs_org              = []
+        kw_perconcat        = []
+        kf_paths_perconcat  = []
+        simvalue_perconcat  = []
+        segment_timeline    = []
+
+        for concatidx, img_path in enumerate(concatimg_paths):
+            concatimages.append(Image.open(img_path).convert('RGB'))
+            concatimages_base64.append(img_path)
+
+            kw_sidx = imgs_group*(concatidx)
+            kw_eidx = imgs_group*(concatidx+1)
+
+            concat_kw = keywords[kw_sidx:kw_eidx]
+            qs_org_ = create_question(concat_kw, num_select)
+
+            kw_perconcat.append(concat_kw)
+            qs_org.append(qs_org_)
+            kf_paths_perconcat.append(kf_paths[kw_sidx:kw_eidx])
+            simvalue_perconcat.append(simvalue[kw_sidx:kw_eidx])
+            segment_timeline.append(line["segment_timeline"])
+
+        concatimg_paths     = concatimg_paths[-num_groups:]
+        concatimages_base64 = concatimages_base64[-num_groups:] 
+        qs_org              = qs_org[-num_groups:]
+        kw_perconcat        = kw_perconcat[-num_groups:]
+        kf_paths_perconcat  = kf_paths_perconcat[-num_groups:]
+        simvalue_perconcat  = simvalue_perconcat[-num_groups:]
+        segment_timeline    = segment_timeline[-num_groups:]
+
+        return concatimages_base64, concatimages[0].size, kw_perconcat, kf_paths_perconcat, qs_org, segment_timeline, concatimg_paths, simvalue_perconcat
+
+    def __len__(self):
+        return len(self.questions)
+
+def create_question(concat_kw, num_select):
+    imgkw_sen = ""
+
+    for i, e in enumerate(concat_kw):
+        if i < len(concat_kw) - 1: imgkw_sen = imgkw_sen + f"Image_{i}: '{e}', "
+        else: imgkw_sen = imgkw_sen + f"Image_{i}: '{e}'."
+
+    if num_select == 3:
+        prompt = f"Eight images, having egocentric perspectives, are juxtaposed, separated by a red vertical line and red horizontal line. In the first row, the images from left to right are named as image_0, image_1, image_2, image_3. In the second row, the images from left to right are named as image_4, image_5, image_6, image_7. Here are images and their associated guess words: {imgkw_sen}. Think step-by-step and list only the names of the {num_select} images most closely related to the guessed words. Do not select blurry images in your answer. If none of the images correspond to the provided guess words, choose any two images at random. Your answer should follow the JSON format shown below and should only include the JSON result. Do not output any warnings or notes under any circumstances. Instead, adhere strictly to the provided JSON format example.\n" 
+        prompt += "{\"image name\": write reason for your selection in 10 words}."
+        prompt += " This is one example output format. {\n  \"image_0\": \"Person washing a plate; linked to dish cleaning.\",\n  \"image_2\": \"Person washing a bowl; linked to dish cleaning.\",\n  \"image_6\": \"Person running water on a sponge; related to dish cleaning.\"\n}"
+
+    elif num_select == 4:
+        prompt = f"Eight images, having egocentric perspectives, are juxtaposed, separated by a red vertical line and red horizontal line. In the first row, the images from left to right are named as image_0, image_1, image_2, image_3. In the second row, the images from left to right are named as image_4, image_5, image_6, image_7. Here are images and their associated guess words: {imgkw_sen}. Think step-by-step and list only the names of the {num_select} images most closely related to the guessed words. Do not select blurry images in your answer. If none of the images correspond to the provided guess words, choose any two images at random. Your answer should follow the JSON format shown below and should only include the JSON result. Do not output any warnings or notes under any circumstances. Instead, adhere strictly to the provided JSON format example.\n" 
+        prompt += "{\"image name\": write reason for your selection in 10 words}."
+        prompt += " This is one example output format. {\n  \"image_0\": \"Person washing a plate; linked to dish cleaning.\",\n  \"image_2\": \"Person washing a bowl; linked to dish cleaning.\",\n  \"image_6\": \"Person running water on a sponge; related to dish cleaning.\",\n  \"image_7\": \"Person moves mouse; linked to working.\"\n}"
+
+    else: assert False, f"num_select:{num_select} is not defined yet"
+
+    return prompt
+
+def encode_image(image_path):
+    with open(image_path, "rb") as image_file:
+        return base64.b64encode(image_file.read()).decode('utf-8')
+
+def create_data_loader_gpt(questions, num_input_imgs, num_select, batch_size=1, num_workers=4):
+    assert batch_size == 1, "batch_size must be 1"
+    dataset = CustomDatasetGPT(questions, num_input_imgs, num_select)
+    data_loader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers, shuffle=False)
+    return data_loader, dataset
+
+def eval_model():
+    question_path, vlm, num_input_imgs, num_select, temp = config.kf_question_path, config.kf_vlm, config.kf_num_input_imgs, config.kf_num_select, config.kf_temp
+    questions = [json.loads(q) for q in open(os.path.expanduser(question_path), "r")]
+    num_questions = len(questions)
+    giter = 0
+
+    answer_path = config.kf_answer_path
+    os.makedirs(os.path.dirname(answer_path), exist_ok=True)
+
+    print(f"question_path:{question_path}\nanswer_path:{answer_path}")
+    ans_file = open(answer_path, "w")
+    data_loader, dataset = create_data_loader_gpt(questions, num_input_imgs, num_select)
+
+    outputs = ""
+    for (image_paths, image_sizes, kw_perconcat, kf_paths_perconcat, cur_prompts, segment_timeline, concatimg_paths, simvalue_perconcat), line in tqdm(zip(data_loader, questions), total=len(questions)):
+        idx, q_uid = line["q_uid"], line["q_uid"]
+        CA         = line["CA"] if "CA" in line else None
+        option0    = line['option 0']
+        option1    = line['option 1']
+        option2    = line['option 2']
+        option3    = line['option 3']
+        option4    = line['option 4']
+        question   = line['question']
+
+        pastobj           = None
+        past_VLM_path     = None
+        past_VLM_timeline = None
+
+        kw_VLM       = []
+        kf_paths_VLM = []
+        kf_timeline  = []
+
+        kw_VLM_ordered       = []
+        kf_paths_VLM_ordered = []
+        kf_timeline_ordered  = []
+
+        prompts     = [x[0] for x in cur_prompts]
+        image_paths = [x[0] for x in image_paths]
+
+        output_VLM = run_gpt(
+            images=image_paths,
+            texts=prompts,
+            api_keys = list(config.dict_api.values()),
+            max_tokens=2000,
+            model=vlm,
+            temperature=temp,
+            num_threads=20,
+            backoff_time=1*60,
+            silent=False,
+            dataset="egoschema",
+            verbose=False,
+        )
+        output_VLM = list(output_VLM)
+
+        for j, _ in enumerate(cur_prompts):
+            kf_paths_perconcat_ = kf_paths_perconcat[j]
+            kf_timeline.append([f"{e[0].split('_')[-2]}.{e[0].split('_')[-1].split('.')[0]}" for e in kf_paths_perconcat_])
+
+        line_frame                      = line.copy()
+
+        line_frame["output_VLM"]        = output_VLM
+        line_frame["concatimg_paths"]   = concatimg_paths
+        line_frame["kf_paths_VLM"]      = kf_paths_perconcat
+        line_frame["kf_timeline"]       = kf_timeline
+        line_frame["kw_perconcat_clip"] = kw_perconcat
+        line_frame["iter"]              = giter
+
+        line_frame.pop("filepath")
+        line_frame.pop("kf_paths")
+        line_frame.pop("google_drive_id")
+
+        try: ans_file.write(json.dumps(line_frame) + "\n")
+        except: assert False, f"line_frame:{line_frame}"
+
+    ans_file.close()
+    print(f"question_path:{question_path}\nanswer_path:{answer_path}")
+    print("job is done")
+
+
+if __name__ == "__main__":
+    eval_model()
+    refine_answer()

requirements.txt

@@ -0,0 +1,24 @@
+torch==2.1.2
+torchvision==0.16.2
+transformers==4.37.2
+tokenizers==0.15.1
+sentencepiece==0.1.99
+shortuuid
+
+peft
+
+numpy
+
+requests
+
+uvicorn
+fastapi
+
+rp
+tqdm
+
+shutil
+
+PIL
+natsort
+clip

scripts/create_caption.sh

@@ -0,0 +1,11 @@
+# run this script from the root of the repo
+# fix paths to data and update GPT keys in code
+
+export PYTHONPATH=$PYTHONPATH:$PWD
+
+python3 VLM_stage.py \
+  --output-dir ego_base_link_bigtensor \
+  --question-path your_question_path.jsonl \
+  --gptmodel "gpt-4o" \
+  --num-kf 12 \
+  --temp 0

scripts/eval_ES.sh

@@ -0,0 +1,23 @@
+# run this script from the root of the repo
+# fix paths to data and update GPT keys in code
+
+export PYTHONPATH=$PYTHONPATH:$PWD
+
+# Evaluate on ES subset (uncomment it if you want to run it)
+python3 LLM_stage.py \
+  --output-dir ego_base_link \
+  --captions data/ESsub_captions_gpt4o.jsonl \
+  --data data/ESsub_qa_data.json \
+  --per-vid-captions 12 \
+  --gptmodel "gpt-4o" \
+  --temperature 0.0
+
+
+# Evaluate on ES full dataset (uncomment it if you want to run it)
+# python3 LLM_stage.py \
+#   --output-dir ego_base_link \
+#   --captions data/ES_captions_gpt4o.jsonl \
+#   --data data/ES_qa_data.json \
+#   --per-vid-captions 12 \
+#   --gptmodel "gpt-4o" \
+#   --temperature 0.0

scripts/get_ES_captions.sh

@@ -0,0 +1,8 @@
+# downloaded our pre-generated EgoSchema captions
+
+wget https://github.com/jongwoopark7978/LVNet_dev/releases/download/v1.0/ESsub_captions_gpt4o.jsonl
+wget TBA
+wget TBA
+wget TBA
+
+# TODO: update the paths when the repo is public.

src/open_clip/__init__.py

@@ -0,0 +1,11 @@
+from llava.open_clip.constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
+from llava.open_clip.factory import create_model, create_model_and_transforms, create_model_from_pretrained, get_tokenizer
+from llava.open_clip.factory import list_models, add_model_config, get_model_config, load_checkpoint
+from llava.open_clip.loss import ClipLoss
+from llava.open_clip.model import CLIP, CustomTextCLIP, CLIPTextCfg, CLIPVisionCfg,\
+    convert_weights_to_lp, convert_weights_to_fp16, trace_model, get_cast_dtype
+from llava.open_clip.openai import load_openai_model, list_openai_models
+from llava.open_clip.pretrained import list_pretrained, list_pretrained_models_by_tag, list_pretrained_tags_by_model,\
+    get_pretrained_url, download_pretrained_from_url, is_pretrained_cfg, get_pretrained_cfg, download_pretrained
+from llava.open_clip.tokenizer import SimpleTokenizer, tokenize
+from llava.open_clip.transform import image_transform

src/open_clip/bpe_simple_vocab_16e6.txt.gz

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

src/open_clip/constants.py

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

src/open_clip/factory.py

@@ -0,0 +1,287 @@
+import json
+import logging
+import os
+import pathlib
+import re
+from copy import deepcopy
+from pathlib import Path
+from typing import Optional, Tuple, Union
+
+import torch
+
+from llava.open_clip.constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
+from llava.open_clip.model import CLIP, CustomTextCLIP, convert_weights_to_lp, convert_to_custom_text_state_dict,\
+    resize_pos_embed, get_cast_dtype
+from llava.open_clip.openai import load_openai_model
+from llava.open_clip.pretrained import is_pretrained_cfg, get_pretrained_cfg, download_pretrained, list_pretrained_tags_by_model
+from llava.open_clip.transform import image_transform
+from llava.open_clip.tokenizer import HFTokenizer, tokenize
+
+
+_MODEL_CONFIG_PATHS = [Path(__file__).parent / f"model_configs/"]
+# _MODEL_CONFIG_PATHS =  ["/home/mryoo/llava_16/llava/open_clip/model_config/"]
+_MODEL_CONFIGS = {}  # directory (model_name: config) of model architecture configs
+
+
+def _natural_key(string_):
+    return [int(s) if s.isdigit() else s for s in re.split(r'(\d+)', string_.lower())]
+
+
+def _rescan_model_configs():
+    global _MODEL_CONFIGS
+
+    config_ext = ('.json',)
+    config_files = []
+
+    # print(f"_MODEL_CONFIG_PATHS:{_MODEL_CONFIG_PATHS}")
+    for config_path in _MODEL_CONFIG_PATHS:
+        if config_path.is_file() and config_path.suffix in config_ext:
+            config_files.append(config_path)
+        elif config_path.is_dir():
+            for ext in config_ext:
+                config_files.extend(config_path.glob(f'*{ext}'))
+        # for ext in config_ext:
+        #     config_files.extend(config_path.glob(f'*{ext}'))
+
+    for cf in config_files:
+        with open(cf, 'r') as f:
+            model_cfg = json.load(f)
+            if all(a in model_cfg for a in ('embed_dim', 'vision_cfg', 'text_cfg')):
+                _MODEL_CONFIGS[cf.stem] = model_cfg
+
+    _MODEL_CONFIGS = {k: v for k, v in sorted(_MODEL_CONFIGS.items(), key=lambda x: _natural_key(x[0]))}
+    # print(f"_MODEL_CONFIGS:{_MODEL_CONFIGS}")
+
+
+_rescan_model_configs()  # initial populate of model config registry
+
+
+def list_models():
+    """ enumerate available model architectures based on config files """
+    return list(_MODEL_CONFIGS.keys())
+
+
+def add_model_config(path):
+    """ add model config path or file and update registry """
+    if not isinstance(path, Path):
+        path = Path(path)
+    _MODEL_CONFIG_PATHS.append(path)
+    _rescan_model_configs()
+
+
+def get_model_config(model_name):
+    # print(f"_MODEL_CONFIGS:{_MODEL_CONFIGS}")
+    if model_name in _MODEL_CONFIGS:
+        return deepcopy(_MODEL_CONFIGS[model_name])
+    else:
+        return None
+
+
+def get_tokenizer(model_name):
+    config = get_model_config(model_name)
+    tokenizer = HFTokenizer(config['text_cfg']['hf_tokenizer_name']) if 'hf_tokenizer_name' in config['text_cfg'] else tokenize
+    return tokenizer
+
+
+def load_state_dict(checkpoint_path: str, map_location='cpu'):
+    checkpoint = torch.load(checkpoint_path, map_location=map_location)
+    if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
+        state_dict = checkpoint['state_dict']
+    else:
+        state_dict = checkpoint
+    if next(iter(state_dict.items()))[0].startswith('module'):
+        state_dict = {k[7:]: v for k, v in state_dict.items()}
+    return state_dict
+
+
+def load_checkpoint(model, checkpoint_path, strict=True):
+    state_dict = load_state_dict(checkpoint_path)
+    # detect old format and make compatible with new format
+    if 'positional_embedding' in state_dict and not hasattr(model, 'positional_embedding'):
+        state_dict = convert_to_custom_text_state_dict(state_dict)
+    resize_pos_embed(state_dict, model)
+    incompatible_keys = model.load_state_dict(state_dict, strict=strict)
+    return incompatible_keys
+
+
+def create_model(
+        model_name: str,
+        pretrained: Optional[str] = None,
+        precision: str = 'fp32',
+        device: Union[str, torch.device] = 'cpu',
+        jit: bool = False,
+        force_quick_gelu: bool = False,
+        force_custom_text: bool = False,
+        force_patch_dropout: Optional[float] = None,
+        pretrained_image: bool = False,
+        pretrained_hf: bool = True,
+        cache_dir: Optional[str] = None,
+):
+    model_name = model_name.replace('/', '-')  # for callers using old naming with / in ViT names
+    if isinstance(device, str):
+        device = torch.device(device)
+
+    if pretrained and pretrained.lower() == 'openai':
+        logging.info(f'Loading pretrained {model_name} from OpenAI.')
+        model = load_openai_model(
+            model_name,
+            precision=precision,
+            device=device,
+            jit=jit,
+            cache_dir=cache_dir,
+        )
+    else:
+        model_cfg = get_model_config(model_name)
+        if model_cfg is not None:
+            logging.info(f'Loaded {model_name} model config.')
+        else:
+            logging.error(f'Model config for {model_name} not found; available models {list_models()}.')
+            raise RuntimeError(f'Model config for {model_name} not found.')
+
+        if force_quick_gelu:
+            # override for use of QuickGELU on non-OpenAI transformer models
+            model_cfg["quick_gelu"] = True
+
+        if force_patch_dropout is not None:
+            # override the default patch dropout value
+            model_cfg["vision_cfg"]["patch_dropout"] = force_patch_dropout
+
+        if pretrained_image:
+            if 'timm_model_name' in model_cfg.get('vision_cfg', {}):
+                # pretrained weight loading for timm models set via vision_cfg
+                model_cfg['vision_cfg']['timm_model_pretrained'] = True
+            else:
+                assert False, 'pretrained image towers currently only supported for timm models'
+
+        cast_dtype = get_cast_dtype(precision)
+        custom_text = model_cfg.pop('custom_text', False) or force_custom_text or ('hf_model_name' in model_cfg.get('text_cfg', {}))
+
+        if custom_text:
+            if 'hf_model_name' in model_cfg.get('text_cfg', {}):
+                model_cfg['text_cfg']['hf_model_pretrained'] = pretrained_hf
+            model = CustomTextCLIP(**model_cfg, cast_dtype=cast_dtype)
+        else:
+            model = CLIP(**model_cfg, cast_dtype=cast_dtype)
+
+        pretrained_cfg = {}
+        if pretrained:
+            checkpoint_path = ''
+            pretrained_cfg = get_pretrained_cfg(model_name, pretrained)
+            if pretrained_cfg:
+                checkpoint_path = download_pretrained(pretrained_cfg, cache_dir=cache_dir)
+            elif os.path.exists(pretrained):
+                checkpoint_path = pretrained
+
+            if checkpoint_path:
+                logging.info(f'Loading pretrained {model_name} weights ({pretrained}).')
+                load_checkpoint(model, checkpoint_path)
+            else:
+                error_str = (
+                    f'Pretrained weights ({pretrained}) not found for model {model_name}.'
+                    f'Available pretrained tags ({list_pretrained_tags_by_model(model_name)}.')
+                logging.warning(error_str)
+                raise RuntimeError(error_str)
+
+        model.to(device=device)
+        if precision in ("fp16", "bf16"):
+            convert_weights_to_lp(model, dtype=torch.bfloat16 if precision == 'bf16' else torch.float16)
+
+        # set image / mean metadata from pretrained_cfg if available, or use default
+        model.visual.image_mean = pretrained_cfg.get('mean', None) or OPENAI_DATASET_MEAN
+        model.visual.image_std = pretrained_cfg.get('std', None) or OPENAI_DATASET_STD
+
+        if jit:
+            model = torch.jit.script(model)
+
+    return model
+
+
+def create_model_and_transforms(
+        model_name: str,
+        pretrained: Optional[str] = None,
+        precision: str = 'fp32',
+        device: Union[str, torch.device] = 'cpu',
+        jit: bool = False,
+        force_quick_gelu: bool = False,
+        force_custom_text: bool = False,
+        force_patch_dropout: Optional[float] = None,
+        pretrained_image: bool = False,
+        pretrained_hf: bool = True,
+        image_mean: Optional[Tuple[float, ...]] = None,
+        image_std: Optional[Tuple[float, ...]] = None,
+        cache_dir: Optional[str] = None,
+):
+    model = create_model(
+        model_name,
+        pretrained,
+        precision=precision,
+        device=device,
+        jit=jit,
+        force_quick_gelu=force_quick_gelu,
+        force_custom_text=force_custom_text,
+        force_patch_dropout=force_patch_dropout,
+        pretrained_image=pretrained_image,
+        pretrained_hf=pretrained_hf,
+        cache_dir=cache_dir,
+    )
+
+    image_mean = image_mean or getattr(model.visual, 'image_mean', None)
+    image_std = image_std or getattr(model.visual, 'image_std', None)
+    preprocess_train = image_transform(
+        model.visual.image_size,
+        is_train=True,
+        mean=image_mean,
+        std=image_std
+    )
+    preprocess_val = image_transform(
+        model.visual.image_size,
+        is_train=False,
+        mean=image_mean,
+        std=image_std
+    )
+
+    return model, preprocess_train, preprocess_val
+
+
+def create_model_from_pretrained(
+        model_name: str,
+        pretrained: str,
+        precision: str = 'fp32',
+        device: Union[str, torch.device] = 'cpu',
+        jit: bool = False,
+        force_quick_gelu: bool = False,
+        force_custom_text: bool = False,
+        return_transform: bool = True,
+        image_mean: Optional[Tuple[float, ...]] = None,
+        image_std: Optional[Tuple[float, ...]] = None,
+        cache_dir: Optional[str] = None,
+):
+    if not is_pretrained_cfg(model_name, pretrained) and not os.path.exists(pretrained):
+        raise RuntimeError(
+            f'{pretrained} is not a valid pretrained cfg or checkpoint for {model_name}.'
+            f' Use open_clip.list_pretrained() to find one.')
+
+    model = create_model(
+        model_name,
+        pretrained,
+        precision=precision,
+        device=device,
+        jit=jit,
+        force_quick_gelu=force_quick_gelu,
+        force_custom_text=force_custom_text,
+        cache_dir=cache_dir,
+    )
+
+    if not return_transform:
+        return model
+
+    image_mean = image_mean or getattr(model.visual, 'image_mean', None)
+    image_std = image_std or getattr(model.visual, 'image_std', None)
+    preprocess = image_transform(
+        model.visual.image_size,
+        is_train=False,
+        mean=image_mean,
+        std=image_std
+    )
+
+    return model, preprocess

src/open_clip/hf_configs.py

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

src/open_clip/hf_model.py

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

src/open_clip/loss.py

@@ -0,0 +1,121 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+try:
+    import torch.distributed.nn
+    from torch import distributed as dist
+    has_distributed = True
+except ImportError:
+    has_distributed = False
+
+try:
+    import horovod.torch as hvd
+except ImportError:
+    hvd = None
+
+
+def gather_features(
+        image_features,
+        text_features,
+        local_loss=False,
+        gather_with_grad=False,
+        rank=0,
+        world_size=1,
+        use_horovod=False
+):
+    assert has_distributed, 'torch.distributed did not import correctly, please use a PyTorch version with support.'
+    if use_horovod:
+        assert hvd is not None, 'Please install horovod'
+        if gather_with_grad:
+            all_image_features = hvd.allgather(image_features)
+            all_text_features = hvd.allgather(text_features)
+        else:
+            with torch.no_grad():
+                all_image_features = hvd.allgather(image_features)
+                all_text_features = hvd.allgather(text_features)
+            if not local_loss:
+                # ensure grads for local rank when all_* features don't have a gradient
+                gathered_image_features = list(all_image_features.chunk(world_size, dim=0))
+                gathered_text_features = list(all_text_features.chunk(world_size, dim=0))
+                gathered_image_features[rank] = image_features
+                gathered_text_features[rank] = text_features
+                all_image_features = torch.cat(gathered_image_features, dim=0)
+                all_text_features = torch.cat(gathered_text_features, dim=0)
+    else:
+        # We gather tensors from all gpus
+        if gather_with_grad:
+            all_image_features = torch.cat(torch.distributed.nn.all_gather(image_features), dim=0)
+            all_text_features = torch.cat(torch.distributed.nn.all_gather(text_features), dim=0)
+        else:
+            gathered_image_features = [torch.zeros_like(image_features) for _ in range(world_size)]
+            gathered_text_features = [torch.zeros_like(text_features) for _ in range(world_size)]
+            dist.all_gather(gathered_image_features, image_features)
+            dist.all_gather(gathered_text_features, text_features)
+            if not local_loss:
+                # ensure grads for local rank when all_* features don't have a gradient
+                gathered_image_features[rank] = image_features
+                gathered_text_features[rank] = text_features
+            all_image_features = torch.cat(gathered_image_features, dim=0)
+            all_text_features = torch.cat(gathered_text_features, dim=0)
+
+    return all_image_features, all_text_features
+
+
+class ClipLoss(nn.Module):
+
+    def __init__(
+            self,
+            local_loss=False,
+            gather_with_grad=False,
+            cache_labels=False,
+            rank=0,
+            world_size=1,
+            use_horovod=False,
+    ):
+        super().__init__()
+        self.local_loss = local_loss
+        self.gather_with_grad = gather_with_grad
+        self.cache_labels = cache_labels
+        self.rank = rank
+        self.world_size = world_size
+        self.use_horovod = use_horovod
+
+        # cache state
+        self.prev_num_logits = 0
+        self.labels = {}
+
+    def forward(self, image_features, text_features, logit_scale):
+        device = image_features.device
+        if self.world_size > 1:
+            all_image_features, all_text_features = gather_features(
+                image_features, text_features,
+                self.local_loss, self.gather_with_grad, self.rank, self.world_size, self.use_horovod)
+
+            if self.local_loss:
+                logits_per_image = logit_scale * image_features @ all_text_features.T
+                logits_per_text = logit_scale * text_features @ all_image_features.T
+            else:
+                logits_per_image = logit_scale * all_image_features @ all_text_features.T
+                logits_per_text = logits_per_image.T
+        else:
+            logits_per_image = logit_scale * image_features @ text_features.T
+            logits_per_text = logit_scale * text_features @ image_features.T
+
+        # calculated ground-truth and cache if enabled
+        num_logits = logits_per_image.shape[0]
+        if self.prev_num_logits != num_logits or device not in self.labels:
+            labels = torch.arange(num_logits, device=device, dtype=torch.long)
+            if self.world_size > 1 and self.local_loss:
+                labels = labels + num_logits * self.rank
+            if self.cache_labels:
+                self.labels[device] = labels
+                self.prev_num_logits = num_logits
+        else:
+            labels = self.labels[device]
+
+        total_loss = (
+            F.cross_entropy(logits_per_image, labels) +
+            F.cross_entropy(logits_per_text, labels)
+            ) / 2
+        return total_loss

src/open_clip/model.py

@@ -0,0 +1,413 @@
+""" CLIP Model
+
+Adapted from https://github.com/openai/CLIP. Originally MIT License, Copyright (c) 2021 OpenAI.
+"""
+from dataclasses import dataclass
+import logging
+import math
+from typing import Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.utils.checkpoint import checkpoint
+from sentence_transformers import SentenceTransformer
+
+from llava.open_clip.hf_model import HFTextEncoder
+from llava.open_clip.modified_resnet import ModifiedResNet
+from llava.open_clip.timm_model import TimmModel
+from llava.open_clip.transformer import LayerNormFp32, LayerNorm, QuickGELU, Attention, VisionTransformer, TextTransformer
+from llava.open_clip.utils import to_2tuple
+
+
+@dataclass
+class CLIPVisionCfg:
+    layers: Union[Tuple[int, int, int, int], int] = 12
+    width: int = 768
+    head_width: int = 64
+    mlp_ratio: float = 4.0
+    patch_size: int = 16
+    image_size: Union[Tuple[int, int], int] = 224
+    ls_init_value: Optional[float] = None  # layer scale initial value
+    patch_dropout: float = 0.  # what fraction of patches to dropout during training (0 would mean disabled and no patches dropped) - 0.5 to 0.75 recommended in the paper for optimal results
+    global_average_pool: bool = False  # whether to global average pool the last embedding layer, instead of using CLS token (https://arxiv.org/abs/2205.01580)
+    global_max_pool: bool = False  # whether to max pool, from CLIPpy paper
+    timm_model_name: str = None  # a valid model name overrides layers, width, patch_size
+    timm_model_pretrained: bool = False  # use (imagenet) pretrained weights for named model
+    timm_pool: str = 'avg'  # feature pooling for timm model ('abs_attn', 'rot_attn', 'avg', '')
+    timm_proj: str = 'linear'  # linear projection for timm model output ('linear', 'mlp', '')
+    timm_proj_bias: bool = False  # enable bias final projection
+
+
+@dataclass
+class CLIPTextCfg:
+    context_length: int = 77
+    vocab_size: int = 49408
+    width: int = 512
+    heads: int = 8
+    layers: int = 12
+    ls_init_value: Optional[float] = None  # layer scale initial value
+    hf_model_name: str = None
+    hf_tokenizer_name: str = None
+    hf_model_pretrained: bool = True
+    proj: str = 'mlp'
+    pooler_type: str = 'mean_pooler'
+
+
+def get_cast_dtype(precision: str):
+    cast_dtype = None
+    if precision == 'bf16':
+        cast_dtype = torch.bfloat16
+    elif precision == 'fp16':
+        cast_dtype = torch.float16
+    return cast_dtype
+
+
+def _build_vision_tower(
+        embed_dim: int,
+        vision_cfg: CLIPVisionCfg,
+        quick_gelu: bool = False,
+        cast_dtype: Optional[torch.dtype] = None
+):
+    if isinstance(vision_cfg, dict):
+        vision_cfg = CLIPVisionCfg(**vision_cfg)
+
+    # OpenAI models are pretrained w/ QuickGELU but native nn.GELU is both faster and more
+    # memory efficient in recent PyTorch releases (>= 1.10).
+    # NOTE: timm models always use native GELU regardless of quick_gelu flag.
+    act_layer = QuickGELU if quick_gelu else nn.GELU
+
+    if vision_cfg.timm_model_name:
+        visual = TimmModel(
+            vision_cfg.timm_model_name,
+            pretrained=vision_cfg.timm_model_pretrained,
+            pool=vision_cfg.timm_pool,
+            proj=vision_cfg.timm_proj,
+            proj_bias=vision_cfg.timm_proj_bias,
+            embed_dim=embed_dim,
+            image_size=vision_cfg.image_size
+        )
+        act_layer = nn.GELU  # so that text transformer doesn't use QuickGELU w/ timm models
+    elif isinstance(vision_cfg.layers, (tuple, list)):
+        vision_heads = vision_cfg.width * 32 // vision_cfg.head_width
+        visual = ModifiedResNet(
+            layers=vision_cfg.layers,
+            output_dim=embed_dim,
+            heads=vision_heads,
+            image_size=vision_cfg.image_size,
+            width=vision_cfg.width
+        )
+    else:
+        vision_heads = vision_cfg.width // vision_cfg.head_width
+        norm_layer = LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm
+        visual = VisionTransformer(
+            image_size=vision_cfg.image_size,
+            patch_size=vision_cfg.patch_size,
+            width=vision_cfg.width,
+            layers=vision_cfg.layers,
+            heads=vision_heads,
+            mlp_ratio=vision_cfg.mlp_ratio,
+            ls_init_value=vision_cfg.ls_init_value,
+            patch_dropout=vision_cfg.patch_dropout,
+            global_average_pool=vision_cfg.global_average_pool,
+            global_max_pool=vision_cfg.global_max_pool,
+            output_dim=embed_dim,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+        )
+
+    return visual
+
+
+def _build_text_tower(
+        embed_dim: int,
+        text_cfg: CLIPTextCfg,
+        quick_gelu: bool = False,
+        cast_dtype: Optional[torch.dtype] = None,
+):
+    if isinstance(text_cfg, dict):
+        text_cfg = CLIPTextCfg(**text_cfg)
+
+    if text_cfg.hf_model_name:
+        if text_cfg.hf_model_name == "sentence-transformers/sentence-t5-base":
+            text = SentenceTransformer("sentence-transformers/sentence-t5-base")
+        else:
+            text = HFTextEncoder(
+                text_cfg.hf_model_name,
+                output_dim=embed_dim,
+                proj=text_cfg.proj,
+                pooler_type=text_cfg.pooler_type,
+                pretrained=text_cfg.hf_model_pretrained
+            )
+    else:
+        act_layer = QuickGELU if quick_gelu else nn.GELU
+        norm_layer = LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm
+
+        text = TextTransformer(
+            context_length=text_cfg.context_length,
+            vocab_size=text_cfg.vocab_size,
+            width=text_cfg.width,
+            heads=text_cfg.heads,
+            layers=text_cfg.layers,
+            ls_init_value=text_cfg.ls_init_value,
+            output_dim=embed_dim,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+        )
+    return text
+
+
+class CLIP(nn.Module):
+    def __init__(
+            self,
+            embed_dim: int,
+            vision_cfg: CLIPVisionCfg,
+            text_cfg: CLIPTextCfg,
+            quick_gelu: bool = False,
+            cast_dtype: Optional[torch.dtype] = None,
+    ):
+        super().__init__()
+        self.visual = _build_vision_tower(embed_dim, vision_cfg, quick_gelu, cast_dtype)
+
+        text = _build_text_tower(embed_dim, text_cfg, quick_gelu, cast_dtype)
+        self.transformer = text.transformer
+        self.vocab_size = text.vocab_size
+        self.token_embedding = text.token_embedding
+        self.positional_embedding = text.positional_embedding
+        self.ln_final = text.ln_final
+        self.text_projection = text.text_projection
+        self.register_buffer('attn_mask', text.attn_mask, persistent=False)
+
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+    def lock_image_tower(self, unlocked_groups=0, freeze_bn_stats=False):
+        # lock image tower as per LiT - https://arxiv.org/abs/2111.07991
+        self.visual.lock(unlocked_groups=unlocked_groups, freeze_bn_stats=freeze_bn_stats)
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.visual.set_grad_checkpointing(enable)
+        self.transformer.grad_checkpointing = enable
+
+    def encode_image(self, image, normalize: bool = False):
+        features = self.visual(image)
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def encode_text(self, text, normalize: bool = False):
+        cast_dtype = self.transformer.get_cast_dtype()
+
+        x = self.token_embedding(text).to(cast_dtype)  # [batch_size, n_ctx, d_model]
+
+        x = x + self.positional_embedding.to(cast_dtype)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x, attn_mask=self.attn_mask)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.ln_final(x)  # [batch_size, n_ctx, transformer.width]
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
+        return F.normalize(x, dim=-1) if normalize else x
+
+    def forward(self, image, text):
+        image_features = self.encode_image(image, normalize=True)
+        text_features = self.encode_text(text, normalize=True)
+        return image_features, text_features, self.logit_scale.exp()
+
+
+class CustomTextCLIP(nn.Module):
+    def __init__(
+            self,
+            embed_dim: int,
+            vision_cfg: CLIPVisionCfg,
+            text_cfg: CLIPTextCfg,
+            quick_gelu: bool = False,
+            cast_dtype: Optional[torch.dtype] = None,
+    ):
+        super().__init__()
+        self.visual = _build_vision_tower(embed_dim, vision_cfg, quick_gelu, cast_dtype)
+        self.text = _build_text_tower(embed_dim, text_cfg, quick_gelu, cast_dtype)
+        if isinstance(text_cfg, dict):
+            text_cfg = CLIPTextCfg(**text_cfg)
+        if text_cfg.hf_model_name:
+            self.use_st = text_cfg.hf_model_name == "sentence-transformers/sentence-t5-base"
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+    def lock_image_tower(self, unlocked_groups=0, freeze_bn_stats=False):
+        # lock image tower as per LiT - https://arxiv.org/abs/2111.07991
+        self.visual.lock(unlocked_groups=unlocked_groups, freeze_bn_stats=freeze_bn_stats)
+
+    def lock_text_tower(self, unlocked_layers: int = 0, freeze_layer_norm: bool = True):
+        self.text.lock(unlocked_layers, freeze_layer_norm)
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.visual.set_grad_checkpointing(enable)
+        self.text.set_grad_checkpointing(enable)
+
+    def encode_image(self, image, normalize: bool = False, pool: bool = True):
+        features = self.visual(image, pool=pool)
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def encode_text(self, text, normalize: bool = False):
+        features = self.text(text)
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def forward(self, image, text):
+        image_features = self.encode_image(image, normalize=True)
+        text_features = self.encode_text(text, normalize=True)
+        return image_features, text_features, self.logit_scale.exp()
+
+
+def convert_weights_to_lp(model: nn.Module, dtype=torch.float16):
+    """Convert applicable model parameters to low-precision (bf16 or fp16)"""
+
+    def _convert_weights(l):
+        if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
+            l.weight.data = l.weight.data.to(dtype)
+            if l.bias is not None:
+                l.bias.data = l.bias.data.to(dtype)
+
+        if isinstance(l, (nn.MultiheadAttention, Attention)):
+            for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]:
+                tensor = getattr(l, attr)
+                if tensor is not None:
+                    tensor.data = tensor.data.to(dtype)
+
+        for name in ["text_projection", "proj"]:
+            if hasattr(l, name):
+                attr = getattr(l, name)
+                if attr is not None:
+                    attr.data = attr.data.to(dtype)
+
+    model.apply(_convert_weights)
+
+
+convert_weights_to_fp16 = convert_weights_to_lp  # backwards compat
+
+
+# used to maintain checkpoint compatibility
+def convert_to_custom_text_state_dict(state_dict: dict):
+    if 'text_projection' in state_dict:
+        # old format state_dict, move text tower -> .text
+        new_state_dict = {}
+        for k, v in state_dict.items():
+            if any(k.startswith(p) for p in (
+                'text_projection',
+                'positional_embedding',
+                'token_embedding',
+                'transformer',
+                'ln_final',
+            )):
+                k = 'text.' + k
+            new_state_dict[k] = v
+        return new_state_dict
+    return state_dict
+
+
+def build_model_from_openai_state_dict(
+        state_dict: dict,
+        quick_gelu=True,
+        cast_dtype=torch.float16,
+):
+    vit = "visual.proj" in state_dict
+
+    if vit:
+        vision_width = state_dict["visual.conv1.weight"].shape[0]
+        vision_layers = len(
+            [k for k in state_dict.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")])
+        vision_patch_size = state_dict["visual.conv1.weight"].shape[-1]
+        grid_size = round((state_dict["visual.positional_embedding"].shape[0] - 1) ** 0.5)
+        image_size = vision_patch_size * grid_size
+    else:
+        counts: list = [
+            len(set(k.split(".")[2] for k in state_dict if k.startswith(f"visual.layer{b}"))) for b in [1, 2, 3, 4]]
+        vision_layers = tuple(counts)
+        vision_width = state_dict["visual.layer1.0.conv1.weight"].shape[0]
+        output_width = round((state_dict["visual.attnpool.positional_embedding"].shape[0] - 1) ** 0.5)
+        vision_patch_size = None
+        assert output_width ** 2 + 1 == state_dict["visual.attnpool.positional_embedding"].shape[0]
+        image_size = output_width * 32
+
+    embed_dim = state_dict["text_projection"].shape[1]
+    context_length = state_dict["positional_embedding"].shape[0]
+    vocab_size = state_dict["token_embedding.weight"].shape[0]
+    transformer_width = state_dict["ln_final.weight"].shape[0]
+    transformer_heads = transformer_width // 64
+    transformer_layers = len(set(k.split(".")[2] for k in state_dict if k.startswith(f"transformer.resblocks")))
+
+    vision_cfg = CLIPVisionCfg(
+        layers=vision_layers,
+        width=vision_width,
+        patch_size=vision_patch_size,
+        image_size=image_size,
+    )
+    text_cfg = CLIPTextCfg(
+        context_length=context_length,
+        vocab_size=vocab_size,
+        width=transformer_width,
+        heads=transformer_heads,
+        layers=transformer_layers
+    )
+    model = CLIP(
+        embed_dim,
+        vision_cfg=vision_cfg,
+        text_cfg=text_cfg,
+        quick_gelu=quick_gelu,  # OpenAI models were trained with QuickGELU
+        cast_dtype=cast_dtype,
+    )
+
+    for key in ["input_resolution", "context_length", "vocab_size"]:
+        state_dict.pop(key, None)
+
+    convert_weights_to_fp16(model)  # OpenAI state dicts are partially converted to float16
+    model.load_state_dict(state_dict)
+    return model.eval()
+
+
+def trace_model(model, batch_size=256, device=torch.device('cpu')):
+    model.eval()
+    image_size = model.visual.image_size
+    example_images = torch.ones((batch_size, 3, image_size, image_size), device=device)
+    example_text = torch.zeros((batch_size, model.context_length), dtype=torch.int, device=device)
+    model = torch.jit.trace_module(
+        model,
+        inputs=dict(
+            forward=(example_images, example_text),
+            encode_text=(example_text,),
+            encode_image=(example_images,)
+        ))
+    model.visual.image_size = image_size
+    return model
+
+
+def resize_pos_embed(state_dict, model, interpolation: str = 'bicubic', seq_dim=1):
+    # Rescale the grid of position embeddings when loading from state_dict
+    old_pos_embed = state_dict.get('visual.positional_embedding', None)
+    if old_pos_embed is None or not hasattr(model.visual, 'grid_size'):
+        return
+    grid_size = to_2tuple(model.visual.grid_size)
+    extra_tokens = 1  # FIXME detect different token configs (ie no class token, or more)
+    new_seq_len = grid_size[0] * grid_size[1] + extra_tokens
+    if new_seq_len == old_pos_embed.shape[0]:
+        return
+
+    if extra_tokens:
+        pos_emb_tok, pos_emb_img = old_pos_embed[:extra_tokens], old_pos_embed[extra_tokens:]
+    else:
+        pos_emb_tok, pos_emb_img = None, old_pos_embed
+    old_grid_size = to_2tuple(int(math.sqrt(len(pos_emb_img))))
+
+    logging.info('Resizing position embedding grid-size from %s to %s', old_grid_size, grid_size)
+    pos_emb_img = pos_emb_img.reshape(1, old_grid_size[0], old_grid_size[1], -1).permute(0, 3, 1, 2)
+    pos_emb_img = F.interpolate(
+        pos_emb_img,
+        size=grid_size,
+        mode=interpolation,
+        align_corners=True,
+    )
+    pos_emb_img = pos_emb_img.permute(0, 2, 3, 1).reshape(1, grid_size[0] * grid_size[1], -1)[0]
+    if pos_emb_tok is not None:
+        new_pos_embed = torch.cat([pos_emb_tok, pos_emb_img], dim=0)
+    else:
+        new_pos_embed = pos_emb_img
+    state_dict['visual.positional_embedding'] = new_pos_embed

src/open_clip/model_configs/RN101-quickgelu.json

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

src/open_clip/model_configs/RN101.json

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

src/open_clip/model_configs/RN50-quickgelu.json

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

src/open_clip/model_configs/RN50.json

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

src/open_clip/model_configs/RN50x16.json

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

src/open_clip/model_configs/RN50x4.json

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

src/open_clip/model_configs/RN50x64.json

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

src/open_clip/model_configs/ViT-L-14-280.json

@@ -0,0 +1,16 @@
+{
+    "embed_dim": 768,
+    "vision_cfg": {
+        "image_size": 280,
+        "layers": 24,
+        "width": 1024,
+        "patch_size": 14
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 768,
+        "heads": 12,
+        "layers": 12
+    }
+}

src/open_clip/model_configs/ViT-L-14-336.json

@@ -0,0 +1,16 @@
+{
+    "embed_dim": 768,
+    "vision_cfg": {
+        "image_size": 336,
+        "layers": 24,
+        "width": 1024,
+        "patch_size": 14
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 768,
+        "heads": 12,
+        "layers": 12
+    }
+}

src/open_clip/model_configs/ViT-L-14.json

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

src/open_clip/model_configs/ViT-L-16-320.json

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

src/open_clip/model_configs/ViT-L-16.json

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