Now the patchioner code is installed via pip, the model weights are loaded through repo id from huggingface

.gitignore

@@ -1,3 +1,6 @@
 *.pyc
 *.pyo
 *.pyd
+venv/
+.gradio/
+.venv/

README.md

@@ -13,4 +13,8 @@ short_description: 'Repo for the Paper "One Patch to Caption Them All: ...'
 
 ArXiv: arxiv.org/abs/2510.02898
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Demo of the Patch-ioner framework, from the paper "One Patch to Caption Them All: A Unified Zero-shot Captioning Framework".
+
+The project page is at [paciosoft.com/Patch-ioner](https://paciosoft.com/Patch-ioner).
+
+

app.py

@@ -24,8 +24,7 @@ from PIL import Image
 import numpy as np
 from typing import List, Dict
 
-# Import the Patchioner model from the src directory
-from src.model import Patchioner
+from patchioner import Patchioner
 
 # Global variable to store the loaded model
 loaded_model = None
@@ -33,7 +32,7 @@ model_config_path = None
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Default model configuration
-DEFAULT_MODEL_CONFIG = "mlp.viecap.k.yaml"
+DEFAULT_MODEL_CONFIG = "https://huggingface.co/Ruggero1912/Patch-ioner_talk2dino_decap_COCO_Captions"
 
 # Example images directory
 current_dir = os.path.dirname(__file__)
@@ -50,13 +49,16 @@ def initialize_default_model() -> str:
         default_config_path = CONFIGS_DIR / DEFAULT_MODEL_CONFIG
         
         if not default_config_path.exists():
-            return f"❌ Default config file not found: {default_config_path}"
+            print( f"❌ Default config file not found: {default_config_path}" )
+            config = DEFAULT_MODEL_CONFIG  # Assume it's a URL or model identifier
+            print( f"Attempting to load model from identifier: {config}" )
+        
+        else:
+            config = default_config_path
         
         print(f"Loading default model: {DEFAULT_MODEL_CONFIG}")
         
-        # Load and parse the config
-        with open(default_config_path, 'r') as f:
-            config = yaml.safe_load(f)
+        
         
         # Load the model using the from_config class method
         model = Patchioner.from_config(config, device=device)
@@ -553,7 +555,7 @@ def generate_bbox_caption(image_data, image) -> str:
         return error_msg
 
 
-def create_gradio_interface():
+def create_gradio_interface(model_config_name : str):
     """Create and configure the Gradio interface."""
     
     # Get example files
@@ -593,7 +595,7 @@ def create_gradio_interface():
     ) as demo:
         #gr.HTML(custom_js)  # inject custom JS
         
-        gr.Markdown("""
+        gr.Markdown(f"""
         # 🎯 Patchioner Trace Captioning Demo
         
         This demo allows you to:
@@ -608,7 +610,7 @@ def create_gradio_interface():
         3. Use the appropriate tool to mark areas of interest in the image
         4. Click "Generate Caption" to get AI-generated descriptions
         
-        **Model:** Using `mlp.karpathy.yaml` configuration (automatically loaded)
+        **Model:** Using `{model_config_name}` configuration (automatically loaded)
         """)
         
         # Initialize model status
@@ -730,7 +732,7 @@ def create_gradio_interface():
                 outputs=[image_editor, image_annotator]
             )
         
-        gr.Markdown("""
+        gr.Markdown(f"""
         ### 💡 Tips:
         - **Mode Selection**: Switch between trace and bounding box modes based on your needs
         - **Trace Mode**: Draw continuous lines over areas you want to describe
@@ -741,7 +743,7 @@ def create_gradio_interface():
         ### 🔧 Technical Details:
         - **Trace Mode**: Converts drawings to normalized (x, y) coordinates with timestamps
         - **BBox Mode**: Uses bounding box coordinates for region-specific captioning
-        - **Model Architecture**: Uses `mlp.karpathy.yaml` configuration with CLIP and ViT components
+        - **Model Architecture**: Uses `{model_config_name}` configuration with CLIP and ViT components
         - **Processing**: Each trace/bbox is processed separately to generate corresponding captions
         """)
     
@@ -762,7 +764,7 @@ if __name__ == "__main__":
     print(f"Example images directory: {EXAMPLE_IMAGES_DIR}")
     print(f"Configs directory: {CONFIGS_DIR}")
     
-    demo = create_gradio_interface()
+    demo = create_gradio_interface(DEFAULT_MODEL_CONFIG)
     if not args.local:
         demo.launch()
     else:

configs/mlp.k.yaml

@@ -1,8 +0,0 @@
-decap_weights: 'weights/decap-talk2dino-coco_karpathy-009.pt'
-prefix_size: 768 
-linear_talk2dino: False 
-support_memory_size: 591753
-dino_model: 'dinov2_vitb14_reg'
-normalize: True
-kkv_attention: False
-projection_type: '/raid/datasets/im2txtmemories/coco_train_karpathy.json'

configs/mlp.viecap.k.yaml

@@ -1,31 +0,0 @@
-decap_weights: null
-prefix_size: 768 
-linear_talk2dino: False 
-support_memory_size: 0
-dino_model: 'dinov2_vitb14_reg'
-normalize: False
-kkv_attention: False
-use_talk2dino_project: False
-clip_model_name: "ViT-B/16"
-
-
-# nested config
-viecap:
-  clip_hidden_size: 768
-  suffix: ViT-B16_t2d_
-  project_length: 10
-  temperature: 0.01
-  top_k: 3
-  threshold: 0.4
-  language_model: 'gpt2'
-  name_of_entities_text: coco_entities #vinvl_vgoi_entities
-  files_path: 'weights/viecap_files/'
-  prompt_ensemble: True
-  weight_path: 'weights/viecap-talk2dino-coco_karpathy-0014.pt'
-  using_hard_prompt: True
-  soft_prompt_first: True
-  only_hard_prompt: False
-  using_greedy_search: True #if false, use beam search
-  beam_width: 5
-  text_prompt: None
-

requirements.txt

@@ -1,36 +1,3 @@
-# Core dependencies - absolutely required
-torch
-transformers==4.46.3
-gradio>=4.0.0
-gradio_image_annotation
-
-# Image processing - required for the demo
-pillow
-torchvision
-
-# Data handling - required 
-numpy
-tqdm
-
-# CLIP - essential for the model
-git+https://github.com/openai/CLIP.git
-
-# Model dependencies - needed for core functionality
-timm
-
-# Hugging Face model hosting
-huggingface_hub
-
-h5py
-
-# Optional: Only include if specifically needed
-# h5py  # Only needed for some data formats - can be installed conditionally
-# scikit-learn  # Only for evaluation - not needed for inference
-# plotly  # Only for plotting - not needed for basic demo  
-# pandas  # Only for data analysis - not needed for basic demo
-# matplotlib  # Only for plotting - not needed for basic demo
-# pycocotools  # Only for COCO evaluation - not needed for basic demo
-# nbformat  # Only for notebooks - not needed for basic demo
-# speaksee  # Only for evaluation - not needed for basic demo
-# munkres  # Only for specific evaluation metrics - not needed for basic demo
-# open_clip_torch  # Only if using open_clip models - not needed for basic demo
+git+https://github.com/Ruggero1912/Patch-ioner
+gradio==5.48.0
+gradio_image_annotation

src/INViTE/clipfolder/__init__.py

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

src/INViTE/clipfolder/bpe_simple_vocab_16e6.txt.gz

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

src/INViTE/clipfolder/clip.py

@@ -1,238 +0,0 @@
-import hashlib
-import os
-import urllib
-import warnings
-from typing import Any, Union, List
-from pkg_resources import packaging
-
-import torch
-from PIL import Image
-from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
-from tqdm import tqdm
-
-from .model import build_model
-from .simple_tokenizer import SimpleTokenizer as _Tokenizer
-
-try:
-    from torchvision.transforms import InterpolationMode
-    BICUBIC = InterpolationMode.BICUBIC
-except ImportError:
-    BICUBIC = Image.BICUBIC
-
-
-if packaging.version.parse(torch.__version__) < packaging.version.parse("1.7.1"):
-    warnings.warn("PyTorch version 1.7.1 or higher is recommended")
-
-
-__all__ = ["available_models", "load", "tokenize"]
-_tokenizer = _Tokenizer()
-
-_MODELS = {
-    "RN50": "https://openaipublic.azureedge.net/clip/models/afeb0e10f9e5a86da6080e35cf09123aca3b358a0c3e3b6c78a7b63bc04b6762/RN50.pt",
-    "RN101": "https://openaipublic.azureedge.net/clip/models/8fa8567bab74a42d41c5915025a8e4538c3bdbe8804a470a72f30b0d94fab599/RN101.pt",
-    "RN50x4": "https://openaipublic.azureedge.net/clip/models/7e526bd135e493cef0776de27d5f42653e6b4c8bf9e0f653bb11773263205fdd/RN50x4.pt",
-    "RN50x16": "https://openaipublic.azureedge.net/clip/models/52378b407f34354e150460fe41077663dd5b39c54cd0bfd2b27167a4a06ec9aa/RN50x16.pt",
-    "RN50x64": "https://openaipublic.azureedge.net/clip/models/be1cfb55d75a9666199fb2206c106743da0f6468c9d327f3e0d0a543a9919d9c/RN50x64.pt",
-    "ViT-B/32": "https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt",
-    "ViT-B/16": "https://openaipublic.azureedge.net/clip/models/5806e77cd80f8b59890b7e101eabd078d9fb84e6937f9e85e4ecb61988df416f/ViT-B-16.pt",
-    "ViT-L/14": "https://openaipublic.azureedge.net/clip/models/b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836/ViT-L-14.pt",
-    "ViT-L/14@336px": "https://openaipublic.azureedge.net/clip/models/3035c92b350959924f9f00213499208652fc7ea050643e8b385c2dac08641f02/ViT-L-14-336px.pt",
-}
-
-
-def _download(url: str, root: str):
-    os.makedirs(root, exist_ok=True)
-    filename = os.path.basename(url)
-
-    expected_sha256 = url.split("/")[-2]
-    download_target = os.path.join(root, filename)
-
-    if os.path.exists(download_target) and not os.path.isfile(download_target):
-        raise RuntimeError(f"{download_target} exists and is not a regular file")
-
-    if os.path.isfile(download_target):
-        if hashlib.sha256(open(download_target, "rb").read()).hexdigest() == expected_sha256:
-            return download_target
-        else:
-            warnings.warn(f"{download_target} exists, but the SHA256 checksum does not match; re-downloading the file")
-
-    with urllib.request.urlopen(url) as source, open(download_target, "wb") as output:
-        with tqdm(total=int(source.info().get("Content-Length")), ncols=80, unit='iB', unit_scale=True, unit_divisor=1024) as loop:
-            while True:
-                buffer = source.read(8192)
-                if not buffer:
-                    break
-
-                output.write(buffer)
-                loop.update(len(buffer))
-
-    if hashlib.sha256(open(download_target, "rb").read()).hexdigest() != expected_sha256:
-        raise RuntimeError("Model has been downloaded but the SHA256 checksum does not not match")
-
-    return download_target
-
-
-def _convert_image_to_rgb(image):
-    return image.convert("RGB")
-
-
-def _transform(n_px):
-    return Compose([
-        Resize(n_px, interpolation=BICUBIC),
-        CenterCrop(n_px),
-        _convert_image_to_rgb,
-        ToTensor(),
-        Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
-    ])
-
-
-def available_models() -> List[str]:
-    """Returns the names of available CLIP models"""
-    return list(_MODELS.keys())
-
-
-def load(name: str, device: Union[str, torch.device] = "cuda" if torch.cuda.is_available() else "cpu", jit: bool = False,
-         download_root: str = None, extract_last_k_th_token: int = -1, viz: bool = False, image_resolution: int = None):
-    """Load a CLIP model
-
-    Parameters
-    ----------
-    name : str
-        A model name listed by `clip.available_models()`, or the path to a model checkpoint containing the state_dict
-
-    device : Union[str, torch.device]
-        The device to put the loaded model
-
-    jit : bool
-        Whether to load the optimized JIT model or more hackable non-JIT model (default).
-
-    download_root: str
-        path to download the model files; by default, it uses "~/.cache/clip"
-
-    Returns
-    -------
-    model : torch.nn.Module
-        The CLIP model
-
-    preprocess : Callable[[PIL.Image], torch.Tensor]
-        A torchvision transform that converts a PIL image into a tensor that the returned model can take as its input
-    """
-    if name in _MODELS:
-        model_path = _download(_MODELS[name], download_root or os.path.expanduser("~/.cache/clip"))
-    elif os.path.isfile(name):
-        model_path = name
-    else:
-        raise RuntimeError(f"Model {name} not found; available models = {available_models()}")
-
-    with open(model_path, 'rb') as opened_file:
-        try:
-            # loading JIT archive
-            model = torch.jit.load(opened_file, map_location=device if jit else "cpu").eval()
-            state_dict = None
-        except RuntimeError:
-            # loading saved state dict
-            if jit:
-                warnings.warn(f"File {model_path} is not a JIT archive. Loading as a state dict instead")
-                jit = False
-            state_dict = torch.load(opened_file, map_location="cpu")
-
-    if not jit:
-        model = build_model(state_dict or model.state_dict(), extract_last_k_th_token, viz, image_resolution=image_resolution).to(device)
-        if str(device) == "cpu":
-            model.float()
-        return model, _transform(model.visual.input_resolution)
-
-    # patch the device names
-    device_holder = torch.jit.trace(lambda: torch.ones([]).to(torch.device(device)), example_inputs=[])
-    device_node = [n for n in device_holder.graph.findAllNodes("prim::Constant") if "Device" in repr(n)][-1]
-
-    def patch_device(module):
-        try:
-            graphs = [module.graph] if hasattr(module, "graph") else []
-        except RuntimeError:
-            graphs = []
-
-        if hasattr(module, "forward1"):
-            graphs.append(module.forward1.graph)
-
-        for graph in graphs:
-            for node in graph.findAllNodes("prim::Constant"):
-                if "value" in node.attributeNames() and str(node["value"]).startswith("cuda"):
-                    node.copyAttributes(device_node)
-
-    model.apply(patch_device)
-    patch_device(model.encode_image)
-    patch_device(model.encode_text)
-
-    # patch dtype to float32 on CPU
-    if str(device) == "cpu":
-        float_holder = torch.jit.trace(lambda: torch.ones([]).float(), example_inputs=[])
-        float_input = list(float_holder.graph.findNode("aten::to").inputs())[1]
-        float_node = float_input.node()
-
-        def patch_float(module):
-            try:
-                graphs = [module.graph] if hasattr(module, "graph") else []
-            except RuntimeError:
-                graphs = []
-
-            if hasattr(module, "forward1"):
-                graphs.append(module.forward1.graph)
-
-            for graph in graphs:
-                for node in graph.findAllNodes("aten::to"):
-                    inputs = list(node.inputs())
-                    for i in [1, 2]:  # dtype can be the second or third argument to aten::to()
-                        if inputs[i].node()["value"] == 5:
-                            inputs[i].node().copyAttributes(float_node)
-
-        model.apply(patch_float)
-        patch_float(model.encode_image)
-        patch_float(model.encode_text)
-
-        model.float()
-
-    return model, _transform(model.input_resolution.item())
-
-
-def tokenize(texts: Union[str, List[str]], context_length: int = 77, truncate: bool = False) -> Union[torch.IntTensor, torch.LongTensor]:
-    """
-    Returns the tokenized representation of given input string(s)
-
-    Parameters
-    ----------
-    texts : Union[str, List[str]]
-        An input string or a list of input strings to tokenize
-
-    context_length : int
-        The context length to use; all CLIP models use 77 as the context length
-
-    truncate: bool
-        Whether to truncate the text in case its encoding is longer than the context length
-
-    Returns
-    -------
-    A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length].
-    We return LongTensor when torch version is <1.8.0, since older index_select requires indices to be long.
-    """
-    if isinstance(texts, str):
-        texts = [texts]
-
-    sot_token = _tokenizer.encoder["<|startoftext|>"]
-    eot_token = _tokenizer.encoder["<|endoftext|>"]
-    all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts]
-    if packaging.version.parse(torch.__version__) < packaging.version.parse("1.8.0"):
-        result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
-    else:
-        result = torch.zeros(len(all_tokens), context_length, dtype=torch.int)
-
-    for i, tokens in enumerate(all_tokens):
-        if len(tokens) > context_length:
-            if truncate:
-                tokens = tokens[:context_length]
-                tokens[-1] = eot_token
-            else:
-                raise RuntimeError(f"Input {texts[i]} is too long for context length {context_length}")
-        result[i, :len(tokens)] = torch.tensor(tokens)
-
-    return result

src/INViTE/clipfolder/model.py

@@ -1,515 +0,0 @@
-from collections import OrderedDict
-from typing import Tuple, Union
-
-import numpy as np
-import torch
-import torch.nn.functional as F
-from torch import nn
-
-
-class Bottleneck(nn.Module):
-    expansion = 4
-
-    def __init__(self, inplanes, planes, stride=1):
-        super().__init__()
-
-        # all conv layers have stride 1. an avgpool is performed after the second convolution when stride > 1
-        self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
-        self.bn1 = nn.BatchNorm2d(planes)
-        self.relu1 = nn.ReLU(inplace=True)
-
-        self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
-        self.bn2 = nn.BatchNorm2d(planes)
-        self.relu2 = nn.ReLU(inplace=True)
-
-        self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity()
-
-        self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False)
-        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
-        self.relu3 = nn.ReLU(inplace=True)
-
-        self.downsample = None
-        self.stride = stride
-
-        if stride > 1 or inplanes != planes * Bottleneck.expansion:
-            # downsampling layer is prepended with an avgpool, and the subsequent convolution has stride 1
-            self.downsample = nn.Sequential(OrderedDict([
-                ("-1", nn.AvgPool2d(stride)),
-                ("0", nn.Conv2d(inplanes, planes * self.expansion, 1, stride=1, bias=False)),
-                ("1", nn.BatchNorm2d(planes * self.expansion))
-            ]))
-
-    def forward(self, x: torch.Tensor):
-        identity = x
-
-        out = self.relu1(self.bn1(self.conv1(x)))
-        out = self.relu2(self.bn2(self.conv2(out)))
-        out = self.avgpool(out)
-        out = self.bn3(self.conv3(out))
-
-        if self.downsample is not None:
-            identity = self.downsample(x)
-
-        out += identity
-        out = self.relu3(out)
-        return out
-
-
-class AttentionPool2d(nn.Module):
-    def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None):
-        super().__init__()
-        self.positional_embedding = nn.Parameter(torch.randn(spacial_dim ** 2 + 1, embed_dim) / embed_dim ** 0.5)
-        self.k_proj = nn.Linear(embed_dim, embed_dim)
-        self.q_proj = nn.Linear(embed_dim, embed_dim)
-        self.v_proj = nn.Linear(embed_dim, embed_dim)
-        self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)
-        self.num_heads = num_heads
-
-    def forward(self, x):
-        x = x.flatten(start_dim=2).permute(2, 0, 1)  # NCHW -> (HW)NC
-        x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0)  # (HW+1)NC
-        x = x + self.positional_embedding[:, None, :].to(x.dtype)  # (HW+1)NC
-        x, _ = F.multi_head_attention_forward(
-            query=x[:1], key=x, value=x,
-            embed_dim_to_check=x.shape[-1],
-            num_heads=self.num_heads,
-            q_proj_weight=self.q_proj.weight,
-            k_proj_weight=self.k_proj.weight,
-            v_proj_weight=self.v_proj.weight,
-            in_proj_weight=None,
-            in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]),
-            bias_k=None,
-            bias_v=None,
-            add_zero_attn=False,
-            dropout_p=0,
-            out_proj_weight=self.c_proj.weight,
-            out_proj_bias=self.c_proj.bias,
-            use_separate_proj_weight=True,
-            training=self.training,
-            need_weights=False
-        )
-        return x.squeeze(0)
-
-
-class ModifiedResNet(nn.Module):
-    """
-    A ResNet class that is similar to torchvision's but contains the following changes:
-    - There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max pool.
-    - Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions with stride > 1
-    - The final pooling layer is a QKV attention instead of an average pool
-    """
-
-    def __init__(self, layers, output_dim, heads, input_resolution=224, width=64):
-        super().__init__()
-        self.output_dim = output_dim
-        self.input_resolution = input_resolution
-
-        # the 3-layer stem
-        self.conv1 = nn.Conv2d(3, width // 2, kernel_size=3, stride=2, padding=1, bias=False)
-        self.bn1 = nn.BatchNorm2d(width // 2)
-        self.relu1 = nn.ReLU(inplace=True)
-        self.conv2 = nn.Conv2d(width // 2, width // 2, kernel_size=3, padding=1, bias=False)
-        self.bn2 = nn.BatchNorm2d(width // 2)
-        self.relu2 = nn.ReLU(inplace=True)
-        self.conv3 = nn.Conv2d(width // 2, width, kernel_size=3, padding=1, bias=False)
-        self.bn3 = nn.BatchNorm2d(width)
-        self.relu3 = nn.ReLU(inplace=True)
-        self.avgpool = nn.AvgPool2d(2)
-
-        # residual layers
-        self._inplanes = width  # this is a *mutable* variable used during construction
-        self.layer1 = self._make_layer(width, layers[0])
-        self.layer2 = self._make_layer(width * 2, layers[1], stride=2)
-        self.layer3 = self._make_layer(width * 4, layers[2], stride=2)
-        self.layer4 = self._make_layer(width * 8, layers[3], stride=2)
-
-        embed_dim = width * 32  # the ResNet feature dimension
-        self.attnpool = AttentionPool2d(input_resolution // 32, embed_dim, heads, output_dim)
-
-    def _make_layer(self, planes, blocks, stride=1):
-        layers = [Bottleneck(self._inplanes, planes, stride)]
-
-        self._inplanes = planes * Bottleneck.expansion
-        for _ in range(1, blocks):
-            layers.append(Bottleneck(self._inplanes, planes))
-
-        return nn.Sequential(*layers)
-
-    def forward(self, x):
-        def stem(x):
-            x = self.relu1(self.bn1(self.conv1(x)))
-            x = self.relu2(self.bn2(self.conv2(x)))
-            x = self.relu3(self.bn3(self.conv3(x)))
-            x = self.avgpool(x)
-            return x
-
-        x = x.type(self.conv1.weight.dtype)
-        x = stem(x)
-        x = self.layer1(x)
-        x = self.layer2(x)
-        x = self.layer3(x)
-        x = self.layer4(x)
-        x = self.attnpool(x)
-
-        return x
-
-
-class LayerNorm(nn.LayerNorm):
-    """Subclass torch's LayerNorm to handle fp16."""
-
-    def forward(self, x: torch.Tensor):
-        orig_type = x.dtype
-        ret = super().forward(x.type(torch.float32))
-        return ret.type(orig_type)
-
-
-class QuickGELU(nn.Module):
-    def forward(self, x: torch.Tensor):
-        return x * torch.sigmoid(1.702 * x)
-
-
-class ResidualAttentionBlock(nn.Module):
-    def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None, viz: bool = False):
-        super().__init__()
-
-        if viz:
-            self.attn = nn.MultiheadAttentionViz(d_model, n_head)
-        else:
-            self.attn = nn.MultiheadAttention(d_model, n_head)
-
-        self.ln_1 = LayerNorm(d_model)
-        self.mlp = nn.Sequential(OrderedDict([
-            ("c_fc", nn.Linear(d_model, d_model * 4)),
-            ("gelu", QuickGELU()),
-            ("c_proj", nn.Linear(d_model * 4, d_model))
-        ]))
-        self.ln_2 = LayerNorm(d_model)
-        self.attn_mask = attn_mask
-
-        """attn_mask – If specified, a 2D or 3D mask preventing attention to certain positions.
-         Must be of shape (L,S)(L, S)(L,S) or (N⋅num_heads,L,S)(N\cdot\text{num\_heads}, L, S)(N⋅num_heads,L,S), 
-         where NNN is the batch size, LLL is the target sequence length, and SSS is the source sequence length. 
-         A 2D mask will be broadcasted across the batch while a 3D mask allows for a different mask for each entry 
-         in the batch. Binary, byte, and float masks are supported. For a binary mask, 
-         a True value indicates that the corresponding position is not allowed to attend. 
-         For a byte mask, a non-zero value indicates that the corresponding position is not allowed to attend. 
-         For a float mask, the mask values will be added to the attention weight."""
-
-    def attention(self, x: torch.Tensor):
-        self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None
-        return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0]
-
-    def forward(self, x: torch.Tensor):
-        x = x + self.attention(self.ln_1(x))
-        x = x + self.mlp(self.ln_2(x))
-        return x
-
-
-
-class Transformer(nn.Module):
-    def __init__(self, width: int, layers: int, heads: int, attn_mask: torch.Tensor = None,
-                 extract_last_k_th_token: int=-1, viz: bool = False, num_tokens: int = 50):
-        super().__init__()
-        self.width = width
-        self.layers = layers
-        print('\n\n\n\n\ntransformer total layers', layers)
-        if extract_last_k_th_token>0:
-            start_mask_layer = layers - extract_last_k_th_token
-
-            ans = []
-            for cnt in range(layers):
-                if cnt < start_mask_layer:
-                    ans.append(ResidualAttentionBlock(width, heads, attn_mask, viz))
-                else:
-                    print(' mask for layer {}'.format(cnt))
-                    mask = torch.empty(num_tokens, num_tokens)
-                    mask.fill_(float("-inf"))
-                    mask.fill_diagonal_(0)
-                    ans.append(ResidualAttentionBlock(width, heads, mask.cuda(), viz))
-                    # TODO: here is hard coded 50 sequence length
-                    # only attend to themselves
-
-            self.resblocks = nn.Sequential(*ans)
-        else:
-            self.resblocks = nn.Sequential(*[ResidualAttentionBlock(width, heads, attn_mask, viz) for _ in range(layers)])
-
-    def forward(self, x: torch.Tensor):
-        return self.resblocks(x)
-
-
-class VisionTransformer(nn.Module):
-    def __init__(self, input_resolution: int, patch_size: int, width: int, layers: int, heads: int, output_dim: int,
-                 extract_last_k_th_token: int, viz: bool):
-        super().__init__()
-        self.input_resolution = input_resolution
-        self.output_dim = output_dim
-        self.conv1 = nn.Conv2d(in_channels=3, out_channels=width, kernel_size=patch_size, stride=patch_size, bias=False)
-
-        scale = width ** -0.5
-        self.class_embedding = nn.Parameter(scale * torch.randn(width))
-        self.positional_embedding = nn.Parameter(scale * torch.randn((input_resolution // patch_size) ** 2 + 1, width))
-        self.ln_pre = LayerNorm(width)
-
-        self.transformer = Transformer(width, layers, heads, extract_last_k_th_token=extract_last_k_th_token, viz=viz, num_tokens=(input_resolution // patch_size) ** 2 + 1)
-
-        self.ln_post = LayerNorm(width)
-        self.proj = nn.Parameter(scale * torch.randn(width, output_dim))
-
-    def forward(self, x: torch.Tensor, get_all_last: bool):
-        # convert x to conv1 dtype
-        x = x.type(self.conv1.weight.dtype)
-        x = self.conv1(x)  # shape = [*, width, grid, grid]
-        x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
-        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
-        x = torch.cat([self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x], dim=1)  # shape = [*, grid ** 2 + 1, width]
-        x = x + self.positional_embedding.to(x.dtype)
-        x = self.ln_pre(x)
-
-        x = x.permute(1, 0, 2)  # NLD -> LND
-        x = self.transformer(x)
-        x = x.permute(1, 0, 2)  # LND -> NLD
-
-        if get_all_last:
-            # take all tokens, x is of shape [*, grid ** 2 + 1, width]
-            # and we apply layer norm to each token separately
-            # x is of shape [*, grid ** 2 + 1, width]
-            x = torch.cat([self.ln_post(x[:, idx, :]).unsqueeze(1) for idx in range(x.size(1))], dim=1)
-        else:
-            # take the first token (CLS token), x is of shape [*, grid ** 2 + 1, width]
-            x = self.ln_post(x[:, 0, :])
-
-        if self.proj is not None:
-            x = x @ self.proj
-        # the returned x is of shape [*, output_dim] where * is the batch size or if get_all_last is True, [*, grid ** 2 + 1, output_dim]
-        return x
-
-
-class CLIP(nn.Module):
-    def __init__(self,
-                 embed_dim: int,
-                 # vision
-                 image_resolution: int,
-                 vision_layers: Union[Tuple[int, int, int, int], int],
-                 vision_width: int,
-                 vision_patch_size: int,
-                 # text
-                 context_length: int,
-                 vocab_size: int,
-                 transformer_width: int,
-                 transformer_heads: int,
-                 transformer_layers: int,
-                 extract_last_k_th_token: int = -1,
-                 viz: bool = False
-                 ):
-        super().__init__()
-
-        self.context_length = context_length
-
-        if isinstance(vision_layers, (tuple, list)):
-            vision_heads = vision_width * 32 // 64
-            self.visual = ModifiedResNet(
-                layers=vision_layers,
-                output_dim=embed_dim,
-                heads=vision_heads,
-                input_resolution=image_resolution,
-                width=vision_width
-            )
-        else:
-            vision_heads = vision_width // 64
-            self.visual = VisionTransformer(
-                input_resolution=image_resolution,
-                patch_size=vision_patch_size,
-                width=vision_width,
-                layers=vision_layers,
-                heads=vision_heads,
-                output_dim=embed_dim,
-                extract_last_k_th_token=extract_last_k_th_token,
-                viz=viz
-            )
-
-        self.transformer = Transformer(
-            width=transformer_width,
-            layers=transformer_layers,
-            heads=transformer_heads,
-            attn_mask=self.build_attention_mask()
-        )
-
-        self.vocab_size = vocab_size
-        self.token_embedding = nn.Embedding(vocab_size, transformer_width)
-        self.positional_embedding = nn.Parameter(torch.empty(self.context_length, transformer_width))
-        self.ln_final = LayerNorm(transformer_width)
-
-        self.text_projection = nn.Parameter(torch.empty(transformer_width, embed_dim))
-        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
-
-        self.initialize_parameters()
-
-    def initialize_parameters(self):
-        nn.init.normal_(self.token_embedding.weight, std=0.02)
-        nn.init.normal_(self.positional_embedding, std=0.01)
-
-        if isinstance(self.visual, ModifiedResNet):
-            if self.visual.attnpool is not None:
-                std = self.visual.attnpool.c_proj.in_features ** -0.5
-                nn.init.normal_(self.visual.attnpool.q_proj.weight, std=std)
-                nn.init.normal_(self.visual.attnpool.k_proj.weight, std=std)
-                nn.init.normal_(self.visual.attnpool.v_proj.weight, std=std)
-                nn.init.normal_(self.visual.attnpool.c_proj.weight, std=std)
-
-            for resnet_block in [self.visual.layer1, self.visual.layer2, self.visual.layer3, self.visual.layer4]:
-                for name, param in resnet_block.named_parameters():
-                    if name.endswith("bn3.weight"):
-                        nn.init.zeros_(param)
-
-        proj_std = (self.transformer.width ** -0.5) * ((2 * self.transformer.layers) ** -0.5)
-        attn_std = self.transformer.width ** -0.5
-        fc_std = (2 * self.transformer.width) ** -0.5
-        for block in self.transformer.resblocks:
-            nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
-            nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
-            nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
-            nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)
-
-        if self.text_projection is not None:
-            nn.init.normal_(self.text_projection, std=self.transformer.width ** -0.5)
-
-    def build_attention_mask(self):
-        # lazily create causal attention mask, with full attention between the vision tokens
-        # pytorch uses additive attention mask; fill with -inf
-        mask = torch.empty(self.context_length, self.context_length)
-        mask.fill_(float("-inf"))
-        mask.triu_(1)  # zero out the lower diagonal
-        return mask
-
-    @property
-    def dtype(self):
-        return self.visual.conv1.weight.dtype
-
-    def encode_image(self, image, get_all_last):
-        return self.visual(image.type(self.dtype), get_all_last)
-
-    def encode_text(self, text):
-        x = self.token_embedding(text).type(self.dtype)  # [batch_size, n_ctx, d_model]
-
-        x = x + self.positional_embedding.type(self.dtype)
-        x = x.permute(1, 0, 2)  # NLD -> LND
-        x = self.transformer(x)
-        x = x.permute(1, 0, 2)  # LND -> NLD
-        x = self.ln_final(x).type(self.dtype)
-
-        # x.shape = [batch_size, n_ctx, transformer.width]
-        # take features from the eot embedding (eot_token is the highest number in each sequence)
-        x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
-
-        return x
-
-    def forward(self, image, text, get_all_last=False):
-        image_features = self.encode_image(image, get_all_last)
-        text_features = self.encode_text(text)
-
-        # normalized features
-        image_features = image_features / image_features.norm(dim=1, keepdim=True)
-        text_features = text_features / text_features.norm(dim=1, keepdim=True)
-
-        # cosine similarity as logits
-        logit_scale = self.logit_scale.exp()
-
-        if get_all_last:
-            return logit_scale * image_features, text_features
-
-
-        logits_per_image = logit_scale * image_features @ text_features.t()
-        logits_per_text = logits_per_image.t()
-
-        # shape = [global_batch_size, global_batch_size]
-        return logits_per_image, logits_per_text
-
-
-def convert_weights(model: nn.Module):
-    """Convert applicable model parameters to fp16"""
-
-    def _convert_weights_to_fp16(l):
-        if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
-            l.weight.data = l.weight.data.half()
-            if l.bias is not None:
-                l.bias.data = l.bias.data.half()
-
-        if isinstance(l, nn.MultiheadAttention):
-            for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]:
-                tensor = getattr(l, attr)
-                if tensor is not None:
-                    tensor.data = tensor.data.half()
-
-        for name in ["text_projection", "proj"]:
-            if hasattr(l, name):
-                attr = getattr(l, name)
-                if attr is not None:
-                    attr.data = attr.data.half()
-
-    model.apply(_convert_weights_to_fp16)
-
-import torch
-import torch.nn.functional as F
-
-def resize_pos_embed(old_pe: torch.Tensor, new_shape: int) -> torch.Tensor:
-    # old_pe: [old_num_patches + 1, C]
-    # new_shape: new_num_patches + 1
-    cls_token = old_pe[:1]
-    patch_pe = old_pe[1:]
-    old_num = int(patch_pe.shape[0] ** 0.5)
-    new_num = int((new_shape - 1) ** 0.5)
-
-    patch_pe = patch_pe.reshape(1, old_num, old_num, -1).permute(0, 3, 1, 2)  # (1, C, H, W)
-    patch_pe = F.interpolate(patch_pe, size=(new_num, new_num), mode='bicubic', align_corners=False)
-    patch_pe = patch_pe.permute(0, 2, 3, 1).reshape(1, new_num * new_num, -1)
-
-    return torch.cat([cls_token.unsqueeze(0), patch_pe], dim=1).squeeze(0)
-
-def build_model(state_dict: dict, extract_last_k_th_token, viz, image_resolution: int = None) -> CLIP:
-    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)
-        if image_resolution is None:
-            image_resolution = 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]
-        if image_resolution is None:
-            image_resolution = 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("transformer.resblocks")))
-
-    model = CLIP(
-        embed_dim,
-        image_resolution, vision_layers, vision_width, vision_patch_size,
-        context_length, vocab_size, transformer_width, transformer_heads, transformer_layers, extract_last_k_th_token, viz
-    )
-
-    for key in ["input_resolution", "context_length", "vocab_size"]:
-        if key in state_dict:
-            del state_dict[key]
-
-    convert_weights(model)
-
-    pretrained_pe = state_dict['visual.positional_embedding']
-    model_pe = model.visual.positional_embedding
-
-    if vit and (pretrained_pe.shape != model_pe.shape):
-        print(f"Interpolating positional embedding from {pretrained_pe.shape} to {model_pe.shape}")
-        state_dict['visual.positional_embedding'] = resize_pos_embed(pretrained_pe, model_pe.shape[0])
-
-    model.load_state_dict(state_dict)
-    return model.eval()

src/INViTE/clipfolder/simple_tokenizer.py

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

src/INViTE/loader.py

@@ -1,72 +0,0 @@
-import torch
-from typing import Union
-from .clipfolder.clip import load as invite_clip_load, tokenize as invite_clip_tokenize
-
-
-def load_invite_clip(config: dict, device: Union[str, torch.device] = "cuda" if torch.cuda.is_available() else "cpu"):
-    """
-    Load an INViTE CLIP model based on the provided configuration.
-    
-    This method loads an INViTE CLIP model similar to how RegionCLIP is loaded in the Patchioner class.
-    
-    Args:
-        config (dict): Configuration dictionary containing the following keys:
-            - name (str): Model name listed by `clip.available_models()`, or path to a model checkpoint
-            - jit (bool, optional): Whether to load the optimized JIT model. Defaults to False
-            - download_root (str, optional): Path to download model files. Defaults to '/raid/datasets/models_weights/INViTE'
-            - extract_last_k_th_token (int, optional): Extract last k-th token. Defaults to -1
-            - viz (bool, optional): Visualization flag. Defaults to False
-        device (Union[str, torch.device], optional): Device to load the model on. 
-            Defaults to "cuda" if available, else "cpu"
-    
-    Returns:
-        tuple: (model, preprocess_transform, tokenize_fn)
-            - model: The loaded INViTE CLIP model
-            - preprocess_transform: Torchvision transform for preprocessing images
-            - tokenize_fn: Tokenization function for text processing
-    
-    Raises:
-        KeyError: If required 'name' key is missing from config
-        RuntimeError: If model loading fails
-    
-    Example:
-        config = {
-            'name': 'ViT-B/32',
-            'jit': False,
-            'download_root': '/raid/datasets/models_weights/INViTE',  # optional, this is the default
-            'extract_last_k_th_token': -1,
-            'viz': False
-        }
-        model, preprocess, tokenize = load_invite_clip(config, device='cuda')
-    """
-    
-    # Validate required parameters
-    if 'name' not in config:
-        raise KeyError("'name' key is required in config dictionary")
-    
-    # Extract parameters with defaults
-    name = config['name']
-    jit = config.get('jit', False)
-    download_root = config.get('download_root', '/raid/datasets/models_weights/INViTE')
-    extract_last_k_th_token = config.get('extract_last_k_th_token', -1)
-    viz = config.get('viz', False)
-
-    image_resolution = config.get('resolution', None)  # Default resolution if not specified
-    
-    # Load the INViTE CLIP model using the clip.load function
-    try:
-        model, preprocess_transform = invite_clip_load(
-            name=name,
-            device=device,
-            jit=jit,
-            download_root=download_root,
-            extract_last_k_th_token=extract_last_k_th_token,
-            viz=viz,
-            image_resolution=image_resolution
-        )
-        
-        # Return model, preprocess transform, and tokenize function
-        return model, preprocess_transform, invite_clip_tokenize
-        
-    except Exception as e:
-        raise RuntimeError(f"Failed to load INViTE CLIP model '{name}': {str(e)}")

src/alphaclip/INSTALL.md

@@ -1,113 +0,0 @@
-# AlphaCLIP Standalone - Installation Guide
-
-## Quick Installation
-
-### Prerequisites
-- Python 3.7 or higher
-- pip package manager
-
-### Step 1: Install Dependencies
-
-```bash
-cd alphaclip-standalone
-pip install -r requirements.txt
-```
-
-### Step 2: Install the Package
-
-```bash
-# Install in development mode (recommended for testing)
-pip install -e .
-
-# OR install normally
-pip install .
-```
-
-### Step 3: Test Installation
-
-```bash
-python test_installation.py
-```
-
-### Step 4: Run Example
-
-```bash
-python example.py
-```
-
-## Manual Dependency Installation
-
-If you encounter issues with the requirements.txt, install dependencies manually:
-
-```bash
-# Core PyTorch (choose appropriate version for your system)
-pip install torch torchvision torchaudio
-
-# Text processing
-pip install ftfy regex tqdm
-
-# LoRA support
-pip install loralib
-
-# Image processing
-pip install Pillow
-
-# Utilities
-pip install numpy packaging
-```
-
-## GPU Support
-
-For CUDA support, make sure you install PyTorch with CUDA:
-
-```bash
-# For CUDA 11.8
-pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
-
-# For CUDA 12.1
-pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
-
-# Check your CUDA version with: nvidia-smi
-```
-
-## Verification
-
-After installation, verify everything works:
-
-```python
-from alphaclip_loader import AlphaCLIPLoader
-
-# This should work without errors
-loader = AlphaCLIPLoader()
-models = loader.available_models()
-print("Available models:", models)
-```
-
-## Troubleshooting
-
-### Common Issues
-
-1. **ImportError: No module named 'loralib'**
-   ```bash
-   pip install loralib
-   ```
-
-2. **CUDA out of memory**
-   - Use CPU: `AlphaCLIPLoader(default_device="cpu")`
-   - Or use a smaller model like "ViT-B/32"
-
-3. **Model download fails**
-   - Check internet connection
-   - Ensure you have enough disk space (~1GB per model)
-   - Models are cached in `~/.cache/clip/`
-
-4. **Permission errors**
-   - Use `--user` flag: `pip install --user -e .`
-
-### Getting Help
-
-If you encounter issues:
-1. Check that all dependencies are properly installed
-2. Run the test script: `python test_installation.py`
-3. Check CUDA compatibility if using GPU
-4. Ensure Python version is 3.7+

src/alphaclip/LICENSE

@@ -1,201 +0,0 @@
-                                 Apache License
-                           Version 2.0, January 2004
-                        http://www.apache.org/licenses/
-
-   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
-
-   1. Definitions.
-
-      "License" shall mean the terms and conditions for use, reproduction,
-      and distribution as defined by Sections 1 through 9 of this document.
-
-      "Licensor" shall mean the copyright owner or entity authorized by
-      the copyright owner that is granting the License.
-
-      "Legal Entity" shall mean the union of the acting entity and all
-      other entities that control, are controlled by, or are under common
-      control with that entity. For the purposes of this definition,
-      "control" means (i) the power, direct or indirect, to cause the
-      direction or management of such entity, whether by contract or
-      otherwise, or (ii) ownership of fifty percent (50%) or more of the
-      outstanding shares, or (iii) beneficial ownership of such entity.
-
-      "You" (or "Your") shall mean an individual or Legal Entity
-      exercising permissions granted by this License.
-
-      "Source" form shall mean the preferred form for making modifications,
-      including but not limited to software source code, documentation
-      source, and configuration files.
-
-      "Object" form shall mean any form resulting from mechanical
-      transformation or translation of a Source form, including but
-      not limited to compiled object code, generated documentation,
-      and conversions to other media types.
-
-      "Work" shall mean the work of authorship, whether in Source or
-      Object form, made available under the License, as indicated by a
-      copyright notice that is included in or attached to the work
-      (an example is provided in the Appendix below).
-
-      "Derivative Works" shall mean any work, whether in Source or Object
-      form, that is based on (or derived from) the Work and for which the
-      editorial revisions, annotations, elaborations, or other modifications
-      represent, as a whole, an original work of authorship. For the purposes
-      of this License, Derivative Works shall not include works that remain
-      separable from, or merely link (or bind by name) to the interfaces of,
-      the Work and Derivative Works thereof.
-
-      "Contribution" shall mean any work of authorship, including
-      the original version of the Work and any modifications or additions
-      to that Work or Derivative Works thereof, that is intentionally
-      submitted to Licensor for inclusion in the Work by the copyright owner
-      or by an individual or Legal Entity authorized to submit on behalf of
-      the copyright owner. For the purposes of this definition, "submitted"
-      means any form of electronic, verbal, or written communication sent
-      to the Licensor or its representatives, including but not limited to
-      communication on electronic mailing lists, source code control systems,
-      and issue tracking systems that are managed by, or on behalf of, the
-      Licensor for the purpose of discussing and improving the Work, but
-      excluding communication that is conspicuously marked or otherwise
-      designated in writing by the copyright owner as "Not a Contribution."
-
-      "Contributor" shall mean Licensor and any individual or Legal Entity
-      on behalf of whom a Contribution has been received by Licensor and
-      subsequently incorporated within the Work.
-
-   2. Grant of Copyright License. Subject to the terms and conditions of
-      this License, each Contributor hereby grants to You a perpetual,
-      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
-      copyright license to reproduce, prepare Derivative Works of,
-      publicly display, publicly perform, sublicense, and distribute the
-      Work and such Derivative Works in Source or Object form.
-
-   3. Grant of Patent License. Subject to the terms and conditions of
-      this License, each Contributor hereby grants to You a perpetual,
-      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
-      (except as stated in this section) patent license to make, have made,
-      use, offer to sell, sell, import, and otherwise transfer the Work,
-      where such license applies only to those patent claims licensable
-      by such Contributor that are necessarily infringed by their
-      Contribution(s) alone or by combination of their Contribution(s)
-      with the Work to which such Contribution(s) was submitted. If You
-      institute patent litigation against any entity (including a
-      cross-claim or counterclaim in a lawsuit) alleging that the Work
-      or a Contribution incorporated within the Work constitutes direct
-      or contributory patent infringement, then any patent licenses
-      granted to You under this License for that Work shall terminate
-      as of the date such litigation is filed.
-
-   4. Redistribution. You may reproduce and distribute copies of the
-      Work or Derivative Works thereof in any medium, with or without
-      modifications, and in Source or Object form, provided that You
-      meet the following conditions:
-
-      (a) You must give any other recipients of the Work or
-          Derivative Works a copy of this License; and
-
-      (b) You must cause any modified files to carry prominent notices
-          stating that You changed the files; and
-
-      (c) You must retain, in the Source form of any Derivative Works
-          that You distribute, all copyright, patent, trademark, and
-          attribution notices from the Source form of the Work,
-          excluding those notices that do not pertain to any part of
-          the Derivative Works; and
-
-      (d) If the Work includes a "NOTICE" text file as part of its
-          distribution, then any Derivative Works that You distribute must
-          include a readable copy of the attribution notices contained
-          within such NOTICE file, excluding those notices that do not
-          pertain to any part of the Derivative Works, in at least one
-          of the following places: within a NOTICE text file distributed
-          as part of the Derivative Works; within the Source form or
-          documentation, if provided along with the Derivative Works; or,
-          within a display generated by the Derivative Works, if and
-          wherever such third-party notices normally appear. The contents
-          of the NOTICE file are for informational purposes only and
-          do not modify the License. You may add Your own attribution
-          notices within Derivative Works that You distribute, alongside
-          or as an addendum to the NOTICE text from the Work, provided
-          that such additional attribution notices cannot be construed
-          as modifying the License.
-
-      You may add Your own copyright statement to Your modifications and
-      may provide additional or different license terms and conditions
-      for use, reproduction, or distribution of Your modifications, or
-      for any such Derivative Works as a whole, provided Your use,
-      reproduction, and distribution of the Work otherwise complies with
-      the conditions stated in this License.
-
-   5. Submission of Contributions. Unless You explicitly state otherwise,
-      any Contribution intentionally submitted for inclusion in the Work
-      by You to the Licensor shall be under the terms and conditions of
-      this License, without any additional terms or conditions.
-      Notwithstanding the above, nothing herein shall supersede or modify
-      the terms of any separate license agreement you may have executed
-      with Licensor regarding such Contributions.
-
-   6. Trademarks. This License does not grant permission to use the trade
-      names, trademarks, service marks, or product names of the Licensor,
-      except as required for reasonable and customary use in describing the
-      origin of the Work and reproducing the content of the NOTICE file.
-
-   7. Disclaimer of Warranty. Unless required by applicable law or
-      agreed to in writing, Licensor provides the Work (and each
-      Contributor provides its Contributions) on an "AS IS" BASIS,
-      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
-      implied, including, without limitation, any warranties or conditions
-      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
-      PARTICULAR PURPOSE. You are solely responsible for determining the
-      appropriateness of using or redistributing the Work and assume any
-      risks associated with Your exercise of permissions under this License.
-
-   8. Limitation of Liability. In no event and under no legal theory,
-      whether in tort (including negligence), contract, or otherwise,
-      unless required by applicable law (such as deliberate and grossly
-      negligent acts) or agreed to in writing, shall any Contributor be
-      liable to You for damages, including any direct, indirect, special,
-      incidental, or consequential damages of any character arising as a
-      result of this License or out of the use or inability to use the
-      Work (including but not limited to damages for loss of goodwill,
-      work stoppage, computer failure or malfunction, or any and all
-      other commercial damages or losses), even if such Contributor
-      has been advised of the possibility of such damages.
-
-   9. Accepting Warranty or Additional Liability. While redistributing
-      the Work or Derivative Works thereof, You may choose to offer,
-      and charge a fee for, acceptance of support, warranty, indemnity,
-      or other liability obligations and/or rights consistent with this
-      License. However, in accepting such obligations, You may act only
-      on Your own behalf and on Your sole responsibility, not on behalf
-      of any other Contributor, and only if You agree to indemnify,
-      defend, and hold each Contributor harmless for any liability
-      incurred by, or claims asserted against, such Contributor by reason
-      of your accepting any such warranty or additional liability.
-
-   END OF TERMS AND CONDITIONS
-
-   APPENDIX: How to apply the Apache License to your work.
-
-      To apply the Apache License to your work, attach the following
-      boilerplate notice, with the fields enclosed by brackets "[]"
-      replaced with your own identifying information. (Don't include
-      the brackets!)  The text should be enclosed in the appropriate
-      comment syntax for the file format. We also recommend that a
-      file or class name and description of purpose be included on the
-      same "printed page" as the copyright notice for easier
-      identification within third-party archives.
-
-   Copyright [Zeyi Sun] [name of copyright owner]
-
-   Licensed under the Apache License, Version 2.0 (the "License");
-   you may not use this file except in compliance with the License.
-   You may obtain a copy of the License at
-
-       http://www.apache.org/licenses/LICENSE-2.0
-
-   Unless required by applicable law or agreed to in writing, software
-   distributed under the License is distributed on an "AS IS" BASIS,
-   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-   See the License for the specific language governing permissions and
-   limitations under the License.

src/alphaclip/MANIFEST.in

@@ -1,7 +0,0 @@
-include README.md
-include requirements.txt
-include LICENSE
-recursive-include alpha_clip *.py
-recursive-include alpha_clip *.gz
-include example.py
-include test_installation.py

src/alphaclip/README.md

@@ -1,266 +0,0 @@
-# AlphaCLIP Standalone
-
-A standalone, easy-to-use version of AlphaCLIP that can be integrated into any project without complex dependencies or setup.
-
-## Overview
-
-AlphaCLIP is an enhanced version of OpenAI's CLIP model that provides improved vision-language understanding capabilities. This standalone package makes it easy to use AlphaCLIP in your projects with minimal setup.
-
-## Features
-
-- **Easy Installation**: Simple pip install with minimal dependencies
-- **Clean API**: Intuitive interface for loading models and processing data
-- **Device Flexibility**: Automatic CUDA/CPU detection with manual override options
-- **Model Variety**: Support for multiple AlphaCLIP model variants
-- **Preprocessing Included**: Built-in image preprocessing and text tokenization
-
-## Installation
-
-### Requirements
-
-- Python 3.7 or higher
-- PyTorch 1.7.1 or higher
-- CUDA (optional, for GPU acceleration)
-
-### Install from source
-
-```bash
-# Clone or download this standalone package
-cd alphaclip-standalone
-
-# Install dependencies
-pip install -r requirements.txt
-
-# Install the package
-pip install -e .
-```
-
-### Core Dependencies
-
-The package requires the following core dependencies:
-
-```
-torch>=1.7.1
-torchvision
-ftfy
-regex
-tqdm
-loralib
-Pillow
-numpy
-packaging
-```
-
-## Quick Start
-
-### Basic Usage
-
-```python
-from alphaclip_loader import AlphaCLIPLoader
-
-# Initialize the loader
-loader = AlphaCLIPLoader()
-
-# Load a model (this will download the model if not cached)
-model, preprocess = loader.load_model("ViT-B/16")
-
-# Tokenize text
-text_tokens = loader.tokenize("A photo of a cat")
-
-# Get text embeddings
-text_features = loader.encode_text(model, "A photo of a cat")
-
-print(f"Text features shape: {text_features.shape}")
-```
-
-### Advanced Usage
-
-```python
-import torch
-from PIL import Image
-from alphaclip_loader import AlphaCLIPLoader
-
-# Initialize with specific device
-loader = AlphaCLIPLoader(default_device="cuda")
-
-# Load model with custom options
-model, preprocess = loader.load_model(
-    "ViT-B/16",
-    device="cuda",
-    lora_adapt=False,
-    rank=16
-)
-
-# Process an image
-image = Image.open("your_image.jpg")
-image_tensor = preprocess(image).unsqueeze(0)  # Add batch dimension
-
-# Get embeddings
-with torch.no_grad():
-    image_features = loader.encode_image(model, image_tensor)
-    text_features = loader.encode_text(model, ["A photo of a cat", "A dog playing"])
-
-# Compute similarities
-similarities = loader.get_similarity(text_features, image_features)
-print(f"Similarities: {similarities}")
-```
-
-### One-line Model Loading
-
-```python
-from alphaclip_loader import load_alphaclip
-
-# Quick loading function
-loader, model, preprocess = load_alphaclip("ViT-B/16", device="cuda")
-```
-
-## Available Models
-
-You can check available models using:
-
-```python
-from alphaclip_loader import AlphaCLIPLoader
-
-loader = AlphaCLIPLoader()
-models = loader.available_models()
-print("Available models:", models)
-```
-
-Typically includes:
-- `ViT-B/32`
-- `ViT-B/16`
-- `ViT-L/14`
-- `ViT-L/14@336px`
-- `RN50`, `RN101`, `RN50x4`, `RN50x16`, `RN50x64`
-
-## API Reference
-
-### AlphaCLIPLoader Class
-
-#### Methods
-
-- **`__init__(default_device=None)`**: Initialize loader with optional default device
-- **`available_models()`**: Get list of available model names
-- **`load_model(name, **kwargs)`**: Load a model with preprocessing function
-- **`tokenize(texts, context_length=77, truncate=True)`**: Tokenize text input
-- **`encode_text(model, texts)`**: Encode text to embeddings
-- **`encode_image(model, images)`**: Encode images to embeddings
-- **`get_similarity(text_features, image_features)`**: Compute cosine similarity
-
-#### load_model Parameters
-
-- `name`: Model name or checkpoint path
-- `alpha_vision_ckpt_pth`: Additional vision checkpoint path (default: "None")
-- `device`: Device to load on (default: auto-detect)
-- `jit`: Use JIT compilation (default: False)
-- `download_root`: Model download directory (default: ~/.cache/clip)
-- `lora_adapt`: Use LoRA adaptation (default: False)
-- `rank`: LoRA rank if enabled (default: 16)
-
-## Example Use Cases
-
-### Image-Text Similarity
-
-```python
-from alphaclip_loader import load_alphaclip
-from PIL import Image
-import torch
-
-loader, model, preprocess = load_alphaclip()
-
-# Load and preprocess image
-image = Image.open("cat.jpg")
-image_input = preprocess(image).unsqueeze(0)
-
-# Define candidate texts
-texts = ["a cat", "a dog", "a bird", "a car"]
-
-# Get features
-image_features = loader.encode_image(model, image_input)
-text_features = loader.encode_text(model, texts)
-
-# Calculate similarities
-similarities = loader.get_similarity(text_features, image_features)
-
-# Find best match
-best_match_idx = similarities.argmax()
-print(f"Best match: {texts[best_match_idx]} (score: {similarities[best_match_idx]:.3f})")
-```
-
-### Batch Processing
-
-```python
-from alphaclip_loader import AlphaCLIPLoader
-import torch
-
-loader = AlphaCLIPLoader()
-model, preprocess = loader.load_model("ViT-B/16")
-
-# Process multiple texts at once
-texts = [
-    "A red apple on a table",
-    "A dog running in the park", 
-    "A beautiful sunset"
-]
-
-# Batch tokenization and encoding
-text_features = loader.encode_text(model, texts)
-print(f"Batch text features shape: {text_features.shape}")  # [3, 512]
-```
-
-## Performance Tips
-
-1. **GPU Usage**: Use CUDA for better performance with larger models
-2. **Batch Processing**: Process multiple texts/images together when possible
-3. **Model Caching**: Models are automatically cached after first download
-4. **Memory Management**: Use `torch.no_grad()` during inference to save memory
-
-## Troubleshooting
-
-### Common Issues
-
-1. **CUDA Out of Memory**: Reduce batch size or use CPU
-2. **Model Download Fails**: Check internet connection and disk space
-3. **Import Errors**: Ensure all dependencies are installed
-
-### Dependencies Issues
-
-If you encounter import errors, try:
-
-```bash
-pip install --upgrade torch torchvision
-pip install ftfy regex tqdm loralib
-```
-
-## File Structure
-
-```
-alphaclip-standalone/
-├── __init__.py              # Package initialization
-├── alphaclip_loader.py      # Main loader class
-├── requirements.txt         # Dependencies
-├── setup.py                # Package setup
-├── README.md               # This file
-└── alpha_clip/             # Core AlphaCLIP modules
-    ├── __init__.py
-    ├── alpha_clip.py       # Main AlphaCLIP functions
-    ├── model.py           # Model architectures
-    ├── simple_tokenizer.py # Text tokenization
-    └── bpe_simple_vocab_16e6.txt.gz  # Tokenizer vocabulary
-```
-
-## License
-
-This standalone package maintains the same license as the original AlphaCLIP project.
-
-## Contributing
-
-This is a standalone distribution. For contributions to the core AlphaCLIP model, please refer to the main AlphaCLIP repository.
-
-## Changelog
-
-### Version 1.0.0
-- Initial standalone release
-- Clean API with AlphaCLIPLoader class
-- Comprehensive documentation and examples
-- Easy installation and setup

src/alphaclip/__init__.py

@@ -1,14 +0,0 @@
-"""
-AlphaCLIP Standalone Package
-
-A standalone version of AlphaCLIP that can be used independently.
-"""
-
-from .alphaclip_loader import AlphaCLIPLoader, load_alphaclip
-
-# Version info
-__version__ = "1.0.0"
-__author__ = "AlphaCLIP Team"
-
-# Make main classes available at package level
-__all__ = ['AlphaCLIPLoader', 'load_alphaclip']

src/alphaclip/alpha_clip/__init__.py

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

src/alphaclip/alpha_clip/alpha_clip.py

@@ -1,254 +0,0 @@
-import hashlib
-import os
-import urllib
-import warnings
-from typing import Any, Union, List
-from pkg_resources import packaging
-
-import torch
-from PIL import Image
-from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
-from tqdm import tqdm
-
-from .model import build_model
-from .simple_tokenizer import SimpleTokenizer as _Tokenizer
-
-try:
-    from torchvision.transforms import InterpolationMode
-    BICUBIC = InterpolationMode.BICUBIC
-except ImportError:
-    BICUBIC = Image.BICUBIC
-
-
-if packaging.version.parse(torch.__version__) < packaging.version.parse("1.7.1"):
-    warnings.warn("PyTorch version 1.7.1 or higher is recommended")
-
-
-__all__ = ["available_models", "load", "tokenize"]
-_tokenizer = _Tokenizer()
-
-_MODELS = {
-    "RN50": "https://openaipublic.azureedge.net/clip/models/afeb0e10f9e5a86da6080e35cf09123aca3b358a0c3e3b6c78a7b63bc04b6762/RN50.pt",
-    "RN101": "https://openaipublic.azureedge.net/clip/models/8fa8567bab74a42d41c5915025a8e4538c3bdbe8804a470a72f30b0d94fab599/RN101.pt",
-    "RN50x4": "https://openaipublic.azureedge.net/clip/models/7e526bd135e493cef0776de27d5f42653e6b4c8bf9e0f653bb11773263205fdd/RN50x4.pt",
-    "RN50x16": "https://openaipublic.azureedge.net/clip/models/52378b407f34354e150460fe41077663dd5b39c54cd0bfd2b27167a4a06ec9aa/RN50x16.pt",
-    "RN50x64": "https://openaipublic.azureedge.net/clip/models/be1cfb55d75a9666199fb2206c106743da0f6468c9d327f3e0d0a543a9919d9c/RN50x64.pt",
-    "ViT-B/32": "https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt",
-    "ViT-B/16": "https://openaipublic.azureedge.net/clip/models/5806e77cd80f8b59890b7e101eabd078d9fb84e6937f9e85e4ecb61988df416f/ViT-B-16.pt",
-    "ViT-L/14": "https://openaipublic.azureedge.net/clip/models/b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836/ViT-L-14.pt",
-    "ViT-L/14@336px": "https://openaipublic.azureedge.net/clip/models/3035c92b350959924f9f00213499208652fc7ea050643e8b385c2dac08641f02/ViT-L-14-336px.pt",
-}
-
-
-def _download(url: str, root: str):
-    os.makedirs(root, exist_ok=True)
-    filename = os.path.basename(url)
-
-    expected_sha256 = url.split("/")[-2]
-    download_target = os.path.join(root, filename)
-
-    if os.path.exists(download_target) and not os.path.isfile(download_target):
-        raise RuntimeError(f"{download_target} exists and is not a regular file")
-
-    if os.path.isfile(download_target):
-        if hashlib.sha256(open(download_target, "rb").read()).hexdigest() == expected_sha256:
-            return download_target
-        else:
-            warnings.warn(f"{download_target} exists, but the SHA256 checksum does not match; re-downloading the file")
-
-    with urllib.request.urlopen(url) as source, open(download_target, "wb") as output:
-        with tqdm(total=int(source.info().get("Content-Length")), ncols=80, unit='iB', unit_scale=True, unit_divisor=1024) as loop:
-            while True:
-                buffer = source.read(8192)
-                if not buffer:
-                    break
-
-                output.write(buffer)
-                loop.update(len(buffer))
-
-    if hashlib.sha256(open(download_target, "rb").read()).hexdigest() != expected_sha256:
-        raise RuntimeError("Model has been downloaded but the SHA256 checksum does not not match")
-
-    return download_target
-
-
-def _convert_image_to_rgb(image):
-    return image.convert("RGB")
-
-
-def _transform(n_px):
-    return Compose([
-        Resize(n_px, interpolation=BICUBIC),
-        CenterCrop(n_px),
-        _convert_image_to_rgb,
-        ToTensor(),
-        Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
-    ])
-
-
-def available_models() -> List[str]:
-    """Returns the names of available CLIP models"""
-    return list(_MODELS.keys())
-
-
-def load(name: str, alpha_vision_ckpt_pth="None", device: Union[str, torch.device] = "cuda" if torch.cuda.is_available() else "cpu", jit: bool = False, download_root: str = None, lora_adapt=False, rank=16):
-    """Load a CLIP model
-
-    Parameters
-    ----------
-    name : str
-        A model name listed by `clip.available_models()`, or the path to a model checkpoint containing the state_dict
-
-    alpha_vision_ckpt_pth: str
-        only changed when inferencing model instead of training
-
-    device : Union[str, torch.device]
-        The device to put the loaded model
-
-    jit : bool
-        Whether to load the optimized JIT model or more hackable non-JIT model (default).
-
-    download_root: str
-        path to download the model files; by default, it uses "~/.cache/clip"
-
-    Returns
-    -------
-    model : torch.nn.Module
-        The CLIP model
-
-    preprocess : Callable[[PIL.Image], torch.Tensor]
-        A torchvision transform that converts a PIL image into a tensor that the returned model can take as its input
-    """
-    if name in _MODELS:
-        model_path = _download(_MODELS[name], download_root or os.path.expanduser("~/.cache/clip"))
-    elif os.path.isfile(name):
-        model_path = name
-    else:
-        raise RuntimeError(f"Model {name} not found; available models = {available_models()}")
-
-    with open(model_path, 'rb') as opened_file:
-        try:
-            # loading JIT archive
-            model = torch.jit.load(opened_file, map_location=device if jit else "cpu").eval()
-            state_dict = None
-        except RuntimeError:
-            # loading saved state dict
-            if jit:
-                warnings.warn(f"File {model_path} is not a JIT archive. Loading as a state dict instead")
-                jit = False
-            state_dict = torch.load(opened_file, map_location="cpu")
-
-    if not jit:
-        model = build_model(state_dict or model.state_dict(), lora_adapt=lora_adapt, rank=rank).to(device)
-        if str(device) == "cpu":
-            model.float()
-        # If a separate checkpoint is provided for the visual encoder (e.g., CLIP), load it
-        if alpha_vision_ckpt_pth != "None":
-            # Load the visual encoder weights from the given checkpoint path
-            model.visual.load_state_dict(torch.load(alpha_vision_ckpt_pth))
-            # Set the model to evaluation mode
-            # Note: If LoRA is used, it may merge LoRA weights into the base model here for inference
-            model.eval() # merge lora params if exists (for inference only)
-        return model, _transform(model.visual.input_resolution)
-
-    # patch the device names
-    device_holder = torch.jit.trace(lambda: torch.ones([]).to(torch.device(device)), example_inputs=[])
-    device_node = [n for n in device_holder.graph.findAllNodes("prim::Constant") if "Device" in repr(n)][-1]
-
-    def _node_get(node: torch._C.Node, key: str):
-        """Gets attributes of a node which is polymorphic over return type.
-        
-        From https://github.com/pytorch/pytorch/pull/82628
-        """
-        sel = node.kindOf(key)
-        return getattr(node, sel)(key)
-
-    def patch_device(module):
-        try:
-            graphs = [module.graph] if hasattr(module, "graph") else []
-        except RuntimeError:
-            graphs = []
-
-        if hasattr(module, "forward1"):
-            graphs.append(module.forward1.graph)
-
-        for graph in graphs:
-            for node in graph.findAllNodes("prim::Constant"):
-                if "value" in node.attributeNames() and str(_node_get(node, "value")).startswith("cuda"):
-                    node.copyAttributes(device_node)
-
-    model.apply(patch_device)
-    patch_device(model.encode_image)
-    patch_device(model.encode_text)
-
-    # patch dtype to float32 on CPU
-    if str(device) == "cpu":
-        float_holder = torch.jit.trace(lambda: torch.ones([]).float(), example_inputs=[])
-        float_input = list(float_holder.graph.findNode("aten::to").inputs())[1]
-        float_node = float_input.node()
-
-        def patch_float(module):
-            try:
-                graphs = [module.graph] if hasattr(module, "graph") else []
-            except RuntimeError:
-                graphs = []
-
-            if hasattr(module, "forward1"):
-                graphs.append(module.forward1.graph)
-
-            for graph in graphs:
-                for node in graph.findAllNodes("aten::to"):
-                    inputs = list(node.inputs())
-                    for i in [1, 2]:  # dtype can be the second or third argument to aten::to()
-                        if _node_get(inputs[i].node(), "value") == 5:
-                            inputs[i].node().copyAttributes(float_node)
-
-        model.apply(patch_float)
-        patch_float(model.encode_image)
-        patch_float(model.encode_text)
-
-        model.float()
-    return model, _transform(model.input_resolution.item())
-
-
-def tokenize(texts: Union[str, List[str]], context_length: int = 77, truncate: bool = True) -> Union[torch.IntTensor, torch.LongTensor]:
-    """
-    Returns the tokenized representation of given input string(s)
-
-    Parameters
-    ----------
-    texts : Union[str, List[str]]
-        An input string or a list of input strings to tokenize
-
-    context_length : int
-        The context length to use; all CLIP models use 77 as the context length
-
-    truncate: bool
-        Whether to truncate the text in case its encoding is longer than the context length
-
-    Returns
-    -------
-    A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length].
-    We return LongTensor when torch version is <1.8.0, since older index_select requires indices to be long.
-    """
-    if isinstance(texts, str):
-        texts = [texts]
-
-    sot_token = _tokenizer.encoder["<|startoftext|>"]
-    eot_token = _tokenizer.encoder["<|endoftext|>"]
-    all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts]
-    if packaging.version.parse(torch.__version__) < packaging.version.parse("1.8.0"):
-        result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
-    else:
-        result = torch.zeros(len(all_tokens), context_length, dtype=torch.int)
-
-    for i, tokens in enumerate(all_tokens):
-        if len(tokens) > context_length:
-            if truncate:
-                tokens = tokens[:context_length]
-                tokens[-1] = eot_token
-            else:
-                raise RuntimeError(f"Input {texts[i]} is too long for context length {context_length}")
-        result[i, :len(tokens)] = torch.tensor(tokens)
-
-    return result

src/alphaclip/alpha_clip/bpe_simple_vocab_16e6.txt.gz

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

src/alphaclip/alpha_clip/model.py

@@ -1,609 +0,0 @@
-from collections import OrderedDict
-from typing import Tuple, Union
-
-import numpy as np
-import torch
-import torch.nn.functional as F
-from torch import nn
-import loralib as lora
-import math
-import collections
-
-class Bottleneck(nn.Module):
-    expansion = 4
-
-    def __init__(self, inplanes, planes, stride=1):
-        super().__init__()
-
-        # all conv layers have stride 1. an avgpool is performed after the second convolution when stride > 1
-        self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
-        self.bn1 = nn.BatchNorm2d(planes)
-        self.relu1 = nn.ReLU(inplace=True)
-
-        self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
-        self.bn2 = nn.BatchNorm2d(planes)
-        self.relu2 = nn.ReLU(inplace=True)
-
-        self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity()
-
-        self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False)
-        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
-        self.relu3 = nn.ReLU(inplace=True)
-
-        self.downsample = None
-        self.stride = stride
-
-        if stride > 1 or inplanes != planes * Bottleneck.expansion:
-            # downsampling layer is prepended with an avgpool, and the subsequent convolution has stride 1
-            self.downsample = nn.Sequential(OrderedDict([
-                ("-1", nn.AvgPool2d(stride)),
-                ("0", nn.Conv2d(inplanes, planes * self.expansion, 1, stride=1, bias=False)),
-                ("1", nn.BatchNorm2d(planes * self.expansion))
-            ]))
-
-    def forward(self, x: torch.Tensor):
-        identity = x
-
-        out = self.relu1(self.bn1(self.conv1(x)))
-        out = self.relu2(self.bn2(self.conv2(out)))
-        out = self.avgpool(out)
-        out = self.bn3(self.conv3(out))
-
-        if self.downsample is not None:
-            identity = self.downsample(x)
-
-        out += identity
-        out = self.relu3(out)
-        return out
-
-
-class AttentionPool2d(nn.Module):
-    def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None):
-        super().__init__()
-        self.positional_embedding = nn.Parameter(torch.randn(spacial_dim ** 2 + 1, embed_dim) / embed_dim ** 0.5)
-        self.k_proj = nn.Linear(embed_dim, embed_dim)
-        self.q_proj = nn.Linear(embed_dim, embed_dim)
-        self.v_proj = nn.Linear(embed_dim, embed_dim)
-        self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)
-        self.num_heads = num_heads
-
-    def forward(self, x):
-        x = x.flatten(start_dim=2).permute(2, 0, 1)  # NCHW -> (HW)NC
-        x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0)  # (HW+1)NC
-        x = x + self.positional_embedding[:, None, :].to(x.dtype)  # (HW+1)NC
-        x, _ = F.multi_head_attention_forward(
-            query=x[:1], key=x, value=x,
-            embed_dim_to_check=x.shape[-1],
-            num_heads=self.num_heads,
-            q_proj_weight=self.q_proj.weight,
-            k_proj_weight=self.k_proj.weight,
-            v_proj_weight=self.v_proj.weight,
-            in_proj_weight=None,
-            in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]),
-            bias_k=None,
-            bias_v=None,
-            add_zero_attn=False,
-            dropout_p=0,
-            out_proj_weight=self.c_proj.weight,
-            out_proj_bias=self.c_proj.bias,
-            use_separate_proj_weight=True,
-            training=self.training,
-            need_weights=False
-        )
-        return x.squeeze(0)
-
-
-class ModifiedResNet(nn.Module):
-    """
-    A ResNet class that is similar to torchvision's but contains the following changes:
-    - There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max pool.
-    - Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions with stride > 1
-    - The final pooling layer is a QKV attention instead of an average pool
-    """
-
-    def __init__(self, layers, output_dim, heads, input_resolution=224, width=64):
-        super().__init__()
-        self.output_dim = output_dim
-        self.input_resolution = input_resolution
-
-        # the 3-layer stem
-        self.conv1 = nn.Conv2d(3, width // 2, kernel_size=3, stride=2, padding=1, bias=False)
-        self.conv1_alpha = nn.Conv2d(in_channels=1, out_channels=width // 2, kernel_size=3, stride=2, padding=1, bias=False)
-        self.bn1 = nn.BatchNorm2d(width // 2)
-        self.relu1 = nn.ReLU(inplace=True)
-        self.conv2 = nn.Conv2d(width // 2, width // 2, kernel_size=3, padding=1, bias=False)
-        self.bn2 = nn.BatchNorm2d(width // 2)
-        self.relu2 = nn.ReLU(inplace=True)
-        self.conv3 = nn.Conv2d(width // 2, width, kernel_size=3, padding=1, bias=False)
-        self.bn3 = nn.BatchNorm2d(width)
-        self.relu3 = nn.ReLU(inplace=True)
-        self.avgpool = nn.AvgPool2d(2)
-
-        # residual layers
-        self._inplanes = width  # this is a *mutable* variable used during construction
-        self.layer1 = self._make_layer(width, layers[0])
-        self.layer2 = self._make_layer(width * 2, layers[1], stride=2)
-        self.layer3 = self._make_layer(width * 4, layers[2], stride=2)
-        self.layer4 = self._make_layer(width * 8, layers[3], stride=2)
-
-        embed_dim = width * 32  # the ResNet feature dimension
-        self.attnpool = AttentionPool2d(input_resolution // 32, embed_dim, heads, output_dim)
-
-    def _make_layer(self, planes, blocks, stride=1):
-        layers = [Bottleneck(self._inplanes, planes, stride)]
-
-        self._inplanes = planes * Bottleneck.expansion
-        for _ in range(1, blocks):
-            layers.append(Bottleneck(self._inplanes, planes))
-
-        return nn.Sequential(*layers)
-
-    def forward(self, x, alpha=None):
-        def stem(x):
-            x = self.relu1(self.bn1(self.conv1(x) + self.conv1_alpha(alpha)))
-            x = self.relu2(self.bn2(self.conv2(x)))
-            x = self.relu3(self.bn3(self.conv3(x)))
-            x = self.avgpool(x)
-            return x
-
-        x = x.type(self.conv1.weight.dtype)
-        x = stem(x)
-        x = self.layer1(x)
-        x = self.layer2(x)
-        x = self.layer3(x)
-        x = self.layer4(x)
-        x = self.attnpool(x)
-
-        return x
-
-
-class LayerNorm(nn.LayerNorm):
-    """Subclass torch's LayerNorm to handle fp16."""
-
-    def forward(self, x: torch.Tensor):
-        orig_type = x.dtype
-        ret = super().forward(x.type(torch.float32))
-        return ret.type(orig_type)
-
-
-class QuickGELU(nn.Module):
-    def forward(self, x: torch.Tensor):
-        return x * torch.sigmoid(1.702 * x)
-
-class Attention(nn.Module):
-    def __init__(
-            self,
-            dim,
-            num_heads=8,
-            qkv_bias=True,
-            scaled_cosine=False,
-            scale_heads=False,
-            logit_scale_max=math.log(1. / 0.01),
-            attn_drop=0.,
-            proj_drop=0.,
-            lora_adapt=False, 
-            rank=16
-    ):
-        super().__init__()
-        self.scaled_cosine = scaled_cosine
-        self.scale_heads = scale_heads
-        assert dim % num_heads == 0, 'dim should be divisible by num_heads'
-        self.num_heads = num_heads
-        self.head_dim = dim // num_heads
-        self.scale = self.head_dim ** -0.5
-        self.logit_scale_max = logit_scale_max
-
-        # keeping in_proj in this form (instead of nn.Linear) to match weight scheme of original
-        if lora_adapt:
-            print("!!!!!!!!!!using lora for qkv projection!!!!!!!!!!")
-            self.in_proj = lora.MergedLinear(dim, 3*dim, r=rank, enable_lora=[True, False, True])
-        else:
-            self.in_proj = nn.Linear(dim, dim * 3)
-        # self.in_proj_weight = nn.Parameter(torch.randn((dim * 3, dim)) * self.scale)
-        # if qkv_bias:
-        #     self.in_proj_bias = nn.Parameter(torch.zeros(dim * 3))
-        # else:
-        #     self.in_proj_bias = None
-
-        if self.scaled_cosine:
-            self.logit_scale = nn.Parameter(torch.log(10 * torch.ones((num_heads, 1, 1))))
-        else:
-            self.logit_scale = None
-        self.attn_drop = nn.Dropout(attn_drop)
-        if self.scale_heads:
-            self.head_scale = nn.Parameter(torch.ones((num_heads, 1, 1)))
-        else:
-            self.head_scale = None
-        self.out_proj = nn.Linear(dim, dim) if not lora_adapt else lora.Linear(dim, dim, r=rank)
-        self.out_drop = nn.Dropout(proj_drop)
-
-    def forward(self, x, attn_mask = None):
-        L, N, C = x.shape
-        q, k, v = self.in_proj(x).chunk(3, dim=-1)
-        q = q.contiguous().view(L, N * self.num_heads, -1).transpose(0, 1)
-        k = k.contiguous().view(L, N * self.num_heads, -1).transpose(0, 1)
-        v = v.contiguous().view(L, N * self.num_heads, -1).transpose(0, 1)
-
-        if self.logit_scale is not None:
-            attn = torch.bmm(F.normalize(q, dim=-1), F.normalize(k, dim=-1).transpose(-1, -2))
-            logit_scale = torch.clamp(self.logit_scale, max=self.logit_scale_max).exp()
-            attn = attn.view(N, self.num_heads, L, L) * logit_scale
-            attn = attn.view(-1, L, L)
-        else:
-            q = q * self.scale
-            attn = torch.bmm(q, k.transpose(-2, -1))
-
-        if attn_mask is not None:
-            if attn_mask.dtype == torch.bool:
-                new_attn_mask = torch.zeros_like(attn_mask, dtype=q.dtype)
-                new_attn_mask.masked_fill_(attn_mask, float("-inf"))
-                attn_mask = new_attn_mask
-            attn += attn_mask
-
-        attn = attn.softmax(dim=-1)
-        attn = self.attn_drop(attn)
-
-        x = torch.bmm(attn, v)
-        if self.head_scale is not None:
-            x = x.view(N, self.num_heads, L, C) * self.head_scale
-            x = x.view(-1, L, C)
-        x = x.transpose(0, 1).reshape(L, N, C)
-        x = self.out_proj(x)
-        x = self.out_drop(x)
-        return x, attn
-
-
-class CustomResidualAttentionBlock(nn.Module):
-    def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None, lora_adapt=False, rank=16):
-        super().__init__()
-        
-        self.attn = Attention(d_model, n_head, lora_adapt=lora_adapt, rank=rank)
-        self.ln_1 = LayerNorm(d_model)
-        self.mlp = nn.Sequential(OrderedDict([
-            ("c_fc", nn.Linear(d_model, d_model * 4) if not lora_adapt else lora.Linear(d_model, d_model*4, r=rank)),
-            ("gelu", QuickGELU()),
-            ("c_proj", nn.Linear(d_model * 4, d_model) if not lora_adapt else lora.Linear(d_model*4, d_model, r=rank))
-        ]))
-        self.ln_2 = LayerNorm(d_model)
-        self.attn_mask = attn_mask
-
-    def attention(self, x: torch.Tensor):
-        self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None
-        return self.attn(x, attn_mask=self.attn_mask)
-
-    def forward(self, x: torch.Tensor, return_attn=False):
-        attn_out, attn = self.attention(self.ln_1(x))
-        x = x + attn_out
-        x = x + self.mlp(self.ln_2(x))
-        if return_attn:
-            return x, attn
-        else:
-            return x
-
-class ResidualAttentionBlock(nn.Module):
-    def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None):
-        super().__init__()
-        
-        self.attn = nn.MultiheadAttention(d_model, n_head)
-        self.ln_1 = LayerNorm(d_model)
-        self.mlp = nn.Sequential(OrderedDict([
-            ("c_fc", nn.Linear(d_model, d_model * 4)),
-            ("gelu", QuickGELU()),
-            ("c_proj", nn.Linear(d_model * 4, d_model))
-        ]))
-        self.ln_2 = LayerNorm(d_model)
-        self.attn_mask = attn_mask
-
-    def attention(self, x: torch.Tensor):
-        self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None
-        return self.attn(x, x, x, attn_mask=self.attn_mask)[0]
-
-    def forward(self, x: torch.Tensor):
-        x = x + self.attention(self.ln_1(x))
-        x = x + self.mlp(self.ln_2(x))
-        return x
-
-class Transformer(nn.Module):
-    def __init__(self, width: int, layers: int, heads: int, attn_mask: torch.Tensor = None):
-        super().__init__()
-        self.width = width
-        self.layers = layers
-        self.resblocks = nn.Sequential(*[ResidualAttentionBlock(width, heads, attn_mask) for _ in range(layers)])
-
-    def forward(self, x: torch.Tensor):
-        return self.resblocks(x)
-
-class CustomTransformer(nn.Module):
-    def __init__(self, width: int, layers: int, heads: int, attn_mask: torch.Tensor = None, lora_adapt=False, rank=16):
-        super().__init__()
-        self.width = width
-        self.layers = layers
-        self.resblocks = nn.Sequential(*[CustomResidualAttentionBlock(width, heads, attn_mask, lora_adapt=lora_adapt, rank=rank) for _ in range(layers)])
-
-    def forward(self, x: torch.Tensor, return_attn=False):
-        if return_attn:
-            for i, block in enumerate(self.resblocks):
-                if i == len(self.resblocks) - 1:
-                    return block(x, return_attn=True)
-                else:
-                    x = block(x)
-            assert False
-        return self.resblocks(x)
-
-class VisionTransformer(nn.Module):
-    def __init__(self, input_resolution: int, patch_size: int, width: int, layers: int, heads: int, output_dim: int, lora_adapt=False, rank=16):
-        super().__init__()
-        self.input_resolution = input_resolution
-        self.output_dim = output_dim
-        self.conv1 = nn.Conv2d(in_channels=3, out_channels=width, kernel_size=patch_size, stride=patch_size, bias=False)
-        self.conv1_alpha = nn.Conv2d(in_channels=1, out_channels=width, kernel_size=patch_size, stride=patch_size, bias=False)
-
-        scale = width ** -0.5
-        self.class_embedding = nn.Parameter(scale * torch.randn(width))
-        self.positional_embedding = nn.Parameter(scale * torch.randn((input_resolution // patch_size) ** 2 + 1, width))
-        self.ln_pre = LayerNorm(width)
-
-        self.transformer = CustomTransformer(width, layers, heads, lora_adapt=lora_adapt, rank=rank)
-
-        self.ln_post = LayerNorm(width)
-        self.proj = nn.Parameter(scale * torch.randn(width, output_dim))
-
-    def forward(self, x: torch.Tensor, alpha=None, return_attn=False, return_patches=False):
-        # if x dtype is different from conv1, convert it
-        if x.dtype != self.conv1.weight.dtype:
-            x = x.type(self.conv1.weight.dtype)
-
-        if alpha.dtype != self.conv1_alpha.weight.dtype:
-            alpha = alpha.type(self.conv1_alpha.weight.dtype)
-        
-        x = self.conv1(x)  # shape = [*, width, grid, grid]
-        # ASSUME alpha is always not None!
-        x = x + self.conv1_alpha(alpha)
-        x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
-        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
-        x = torch.cat([self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x], dim=1)  # shape = [*, grid ** 2 + 1, width]
-        x = x + self.positional_embedding.to(x.dtype)
-        x = self.ln_pre(x)
-
-        x = x.permute(1, 0, 2)  # NLD -> LND
-        if return_attn:
-            x, attn_last = self.transformer(x, return_attn=True)
-        else:
-            x = self.transformer(x, return_attn=False)
-        x = x.permute(1, 0, 2)  # LND -> NLD
-
-        if not return_patches:
-            x = self.ln_post(x[:, 0, :])
-        else:
-            x = self.ln_post(x)
-
-        if self.proj is not None:
-            x = x @ self.proj
-        if return_attn:
-            return x, attn_last
-        else:
-            return x
-
-
-class CLIP(nn.Module):
-    def __init__(self,
-                 embed_dim: int,
-                 # vision
-                 image_resolution: int,
-                 vision_layers: Union[Tuple[int, int, int, int], int],
-                 vision_width: int,
-                 vision_patch_size: int,
-                 # text
-                 context_length: int,
-                 vocab_size: int,
-                 transformer_width: int,
-                 transformer_heads: int,
-                 transformer_layers: int,
-                 lora_adapt = False,
-                 rank = 16,
-                 ):
-        super().__init__()
-
-        self.context_length = context_length
-
-        if isinstance(vision_layers, (tuple, list)):
-            vision_heads = vision_width * 32 // 64
-            self.visual = ModifiedResNet(
-                layers=vision_layers,
-                output_dim=embed_dim,
-                heads=vision_heads,
-                input_resolution=image_resolution,
-                width=vision_width
-            )
-        else:
-            vision_heads = vision_width // 64
-            self.visual = VisionTransformer(
-                input_resolution=image_resolution,
-                patch_size=vision_patch_size,
-                width=vision_width,
-                layers=vision_layers,
-                heads=vision_heads,
-                output_dim=embed_dim,
-                lora_adapt=lora_adapt,
-                rank=rank
-            )
-
-        self.transformer = Transformer(
-            width=transformer_width,
-            layers=transformer_layers,
-            heads=transformer_heads,
-            attn_mask=self.build_attention_mask()
-        )
-
-        self.vocab_size = vocab_size
-        self.token_embedding = nn.Embedding(vocab_size, transformer_width)
-        self.positional_embedding = nn.Parameter(torch.empty(self.context_length, transformer_width))
-        self.ln_final = LayerNorm(transformer_width)
-
-        self.text_projection = nn.Parameter(torch.empty(transformer_width, embed_dim))
-        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
-
-        self.initialize_parameters()
-
-    def initialize_parameters(self):
-        nn.init.normal_(self.token_embedding.weight, std=0.02)
-        nn.init.normal_(self.positional_embedding, std=0.01)
-
-        if isinstance(self.visual, ModifiedResNet):
-            if self.visual.attnpool is not None:
-                std = self.visual.attnpool.c_proj.in_features ** -0.5
-                nn.init.normal_(self.visual.attnpool.q_proj.weight, std=std)
-                nn.init.normal_(self.visual.attnpool.k_proj.weight, std=std)
-                nn.init.normal_(self.visual.attnpool.v_proj.weight, std=std)
-                nn.init.normal_(self.visual.attnpool.c_proj.weight, std=std)
-
-            for resnet_block in [self.visual.layer1, self.visual.layer2, self.visual.layer3, self.visual.layer4]:
-                for name, param in resnet_block.named_parameters():
-                    if name.endswith("bn3.weight"):
-                        nn.init.zeros_(param)
-
-        proj_std = (self.transformer.width ** -0.5) * ((2 * self.transformer.layers) ** -0.5)
-        attn_std = self.transformer.width ** -0.5
-        fc_std = (2 * self.transformer.width) ** -0.5
-        for block in self.transformer.resblocks:
-            nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
-            nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
-            nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
-            nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)
-
-        if self.text_projection is not None:
-            nn.init.normal_(self.text_projection, std=self.transformer.width ** -0.5)
-
-    def build_attention_mask(self):
-        # lazily create causal attention mask, with full attention between the vision tokens
-        # pytorch uses additive attention mask; fill with -inf
-        mask = torch.empty(self.context_length, self.context_length)
-        mask.fill_(float("-inf"))
-        mask.triu_(1)  # zero out the lower diagonal
-        return mask
-
-    @property
-    def dtype(self):
-        if not hasattr(self.visual, "conv1"):
-            return self.visual.module.conv1.weight.dtype
-        return self.visual.conv1.weight.dtype
-
-    def encode_image(self, image, alpha):
-        assert alpha is not None
-        return self.visual(image.type(self.dtype), alpha.type(self.dtype))
-
-    def encode_text(self, text):
-        x = self.token_embedding(text).type(self.dtype)  # [batch_size, n_ctx, d_model]
-
-        x = x + self.positional_embedding.type(self.dtype)
-        x = x.permute(1, 0, 2)  # NLD -> LND
-        x = self.transformer(x)
-        x = x.permute(1, 0, 2)  # LND -> NLD
-        x = self.ln_final(x).type(self.dtype)
-
-        # x.shape = [batch_size, n_ctx, transformer.width]
-        # take features from the eot embedding (eot_token is the highest number in each sequence)
-        x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
-
-        return x
-
-    def forward(self, image, text, alpha):
-
-        image_features = self.encode_image(image, alpha)
-        text_features = self.encode_text(text)
-
-        # normalized features
-        image_features = image_features / image_features.norm(dim=1, keepdim=True)
-        text_features = text_features / text_features.norm(dim=1, keepdim=True)
-
-        # cosine similarity as logits
-        logit_scale = self.logit_scale.exp()
-        logits_per_image = logit_scale * image_features @ text_features.t()
-        logits_per_text = logits_per_image.t()
-
-        # shape = [global_batch_size, global_batch_size]
-        return logits_per_image, logits_per_text
-
-
-def convert_weights(model: nn.Module):
-    """Convert applicable model parameters to fp16"""
-
-    def _convert_weights_to_fp16(l):
-        if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
-            l.weight.data = l.weight.data.half()
-            if l.bias is not None:
-                l.bias.data = l.bias.data.half()
-
-        if isinstance(l, nn.MultiheadAttention):
-            for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]:
-                tensor = getattr(l, attr)
-                if tensor is not None:
-                    tensor.data = tensor.data.half()
-
-        for name in ["text_projection", "proj"]:
-            if hasattr(l, name):
-                attr = getattr(l, name)
-                if attr is not None:
-                    attr.data = attr.data.half()
-
-    model.apply(_convert_weights_to_fp16)
-
-
-def build_model(state_dict: dict, lora_adapt=False, rank=16):
-    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_resolution = 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_resolution = 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("transformer.resblocks")))
-
-    # always load lora version
-    model = CLIP(
-        embed_dim,
-        image_resolution, vision_layers, vision_width, vision_patch_size,
-        context_length, vocab_size, transformer_width, transformer_heads, transformer_layers,
-        lora_adapt=lora_adapt, rank=rank,
-    )
-
-    for key in ["input_resolution", "context_length", "vocab_size"]:
-        if key in state_dict:
-            del state_dict[key]
-    # para_wb to linear
-    new_state_dict = collections.OrderedDict()
-    for k, v in state_dict.items():
-        if 'visual' in k:
-            if 'in_proj_weight' in k:
-                new_state_dict[k.replace('in_proj_weight', 'in_proj.weight')] = v
-            elif 'in_proj_bias' in k:
-                new_state_dict[k.replace('in_proj_bias', 'in_proj.bias')] = v
-            else:
-                new_state_dict[k] = v
-        else:
-            new_state_dict[k] = v
-                
-    state_dict = new_state_dict
-    # add rgba_conv_weight
-    if 'visual.conv1_alpha.weight' not in state_dict.keys(): # zero initialization on alpha channel
-        rgb_weight = state_dict['visual.conv1.weight'].clone().detach()
-        rgba_weigth = torch.zeros_like(rgb_weight)[:, 0:1, :, :]
-        state_dict['visual.conv1_alpha.weight'] = rgba_weigth
-    convert_weights(model)
-    model.load_state_dict(state_dict, strict=False)
-    return model.eval()

src/alphaclip/alpha_clip/simple_tokenizer.py

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

src/alphaclip/alpha_mask_utils.py

@@ -1,111 +0,0 @@
-"""
-Utility functions for converting bboxes and traces to alpha masks for AlphaClip.
-"""
-
-import torch
-import math
-
-
-def bbox_to_alpha_mask(bbox, grid_size, patch_size, crop_dim):
-    """
-    Convert a single bounding box to an alpha mask for AlphaClip.
-    
-    Args:
-        bbox: [x_min, y_min, w, h] format in original coordinates
-        grid_size: Number of patches per side (e.g., 37 for 518/14)
-        patch_size: Size of each patch in pixels
-        crop_dim: Size of the cropped image
-    
-    Returns:
-        alpha_mask: Binary mask of shape (grid_size, grid_size)
-    """
-    alpha_mask = torch.zeros((grid_size, grid_size))
-    
-    # Convert bbox to patch coordinates
-    x_min, y_min, w, h = bbox
-    x_max = x_min + w
-    y_max = y_min + h
-    
-    # Scale to patch grid coordinates
-    x1_patch = int(x_min // patch_size)
-    y1_patch = int(y_min // patch_size)
-    x2_patch = int(x_max // patch_size)
-    y2_patch = int(y_max // patch_size)
-    
-    # Clamp to grid bounds
-    x1_patch = max(0, min(x1_patch, grid_size - 1))
-    y1_patch = max(0, min(y1_patch, grid_size - 1))
-    x2_patch = max(0, min(x2_patch, grid_size))  # Allow up to grid_size for exclusive end
-    y2_patch = max(0, min(y2_patch, grid_size))
-    
-    # Set the region to 1 (using slice notation for proper indexing)
-    if x2_patch > x1_patch and y2_patch > y1_patch:
-        alpha_mask[y1_patch:y2_patch, x1_patch:x2_patch] = 1.0
-    
-    return alpha_mask
-
-
-def bboxes_to_alpha_mask(bboxes, grid_size, patch_size, crop_dim):
-    """
-    Convert multiple bboxes to a single OR-ed alpha mask.
-    
-    Args:
-        bboxes: Tensor of bboxes in [x_min, y_min, w, h] format, shape [n_boxes, 4]
-        grid_size: Number of patches per side
-        patch_size: Size of each patch in pixels 
-        crop_dim: Size of the cropped image
-        
-    Returns:
-        alpha_mask: Binary mask of shape (grid_size, grid_size)
-    """
-    alpha_mask = torch.zeros((grid_size, grid_size))
-    
-    for bbox in bboxes:
-        # Skip dummy boxes (negative values)
-        if bbox.sum().item() < 0:
-            continue
-            
-        bbox_mask = bbox_to_alpha_mask(bbox, grid_size, patch_size, crop_dim)
-        alpha_mask = torch.logical_or(alpha_mask, bbox_mask).float()
-    
-    return alpha_mask
-
-
-def trace_to_alpha_mask(trace, grid_size):
-    """
-    Convert a trace to an alpha mask using the existing map_traces_to_grid function.
-    
-    Args:
-        trace: List of trace points with 'x' and 'y' coordinates (normalized 0-1)
-        grid_size: Number of patches per side
-        
-    Returns:
-        alpha_mask: Binary mask of shape (grid_size, grid_size) 
-    """
-    from src.bbox_utils import map_traces_to_grid
-    
-    alpha_mask = map_traces_to_grid(trace, grid_size)
-    # Convert to binary (any value > 0 becomes 1)
-    alpha_mask = (alpha_mask > 0).float()
-    
-    return alpha_mask
-
-
-def traces_to_alpha_mask(traces, grid_size):
-    """
-    Convert multiple traces to a single OR-ed alpha mask.
-    
-    Args:
-        traces: List of traces
-        grid_size: Number of patches per side
-        
-    Returns:
-        alpha_mask: Binary mask of shape (grid_size, grid_size)
-    """
-    alpha_mask = torch.zeros((grid_size, grid_size))
-    
-    for trace in traces:
-        trace_mask = trace_to_alpha_mask(trace, grid_size)
-        alpha_mask = torch.logical_or(alpha_mask, trace_mask).float()
-    
-    return alpha_mask

src/alphaclip/alphaclip_loader.py

@@ -1,233 +0,0 @@
-"""
-AlphaCLIP Standalone Loader
-
-This module provides a simple interface to load and use AlphaCLIP models.
-It exposes the core functionality of AlphaCLIP in a standalone package.
-
-Usage:
-    from alphaclip_loader import AlphaCLIPLoader
-    
-    # Initialize the loader
-    loader = AlphaCLIPLoader()
-    
-    # Load a model
-    model, preprocess = loader.load_model("ViT-B/16")
-    
-    # Tokenize text
-    tokens = loader.tokenize("A photo of a cat")
-    
-    # Get available models
-    models = loader.available_models()
-"""
-
-import os
-import sys
-from typing import Union, List, Tuple, Optional
-
-# Check for critical dependencies
-missing_deps = []
-try:
-    import torch
-except ImportError:
-    missing_deps.append("torch")
-
-try:
-    from PIL import Image
-except ImportError:
-    missing_deps.append("Pillow")
-
-if missing_deps:
-    raise ImportError(f"Missing required dependencies: {', '.join(missing_deps)}. "
-                     f"Please install them with: pip install {' '.join(missing_deps)}")
-
-# Add the alpha_clip directory to the path
-_current_dir = os.path.dirname(os.path.abspath(__file__))
-_alpha_clip_dir = os.path.join(_current_dir, 'alpha_clip')
-if _alpha_clip_dir not in sys.path:
-    sys.path.insert(0, _alpha_clip_dir)
-
-# Import the alpha_clip modules
-try:
-    #import .alpha_clip
-    from .alpha_clip import available_models, load, tokenize
-except ImportError as e:
-    raise ImportError(f"Failed to import alpha_clip modules: {e}. Please ensure all dependencies are installed.")
-
-
-class AlphaCLIPLoader:
-    """
-    A convenience wrapper for AlphaCLIP functionality.
-    
-    This class provides a clean interface to load AlphaCLIP models and 
-    perform text tokenization.
-    """
-    
-    def __init__(self, default_device: Optional[str] = None):
-        """
-        Initialize the AlphaCLIP loader.
-        
-        Args:
-            default_device: Default device to load models on. If None, will use
-                          CUDA if available, otherwise CPU.
-        """
-        if default_device is None:
-            self.default_device = "cuda" if torch.cuda.is_available() else "cpu"
-        else:
-            self.default_device = default_device
-    
-    def available_models(self) -> List[str]:
-        """
-        Get list of available AlphaCLIP model names.
-        
-        Returns:
-            List of model names that can be used with load_model()
-        """
-        return available_models()
-    
-    def load_model(
-        self, 
-        name: str,
-        alpha_vision_ckpt_pth: str = "None",
-        device: Optional[Union[str, torch.device]] = None,
-        jit: bool = False,
-        download_root: Optional[str] = None,
-        lora_adapt: bool = False,
-        rank: int = 16
-    ) -> Tuple[torch.nn.Module, callable]:
-        """
-        Load an AlphaCLIP model.
-        
-        Args:
-            name: Model name (e.g., "ViT-B/16") or path to checkpoint
-            alpha_vision_ckpt_pth: Path to additional vision checkpoint
-            device: Device to load model on (defaults to self.default_device)
-            jit: Whether to load JIT optimized model
-            download_root: Directory to download models to
-            lora_adapt: Whether to use LoRA adaptation
-            rank: LoRA rank if lora_adapt is True
-            
-        Returns:
-            Tuple of (model, preprocess_function)
-        """
-        if device is None:
-            device = self.default_device
-            
-        return load(
-            name=name,
-            alpha_vision_ckpt_pth=alpha_vision_ckpt_pth,
-            device=device,
-            jit=jit,
-            download_root=download_root,
-            lora_adapt=lora_adapt,
-            rank=rank
-        )
-    
-    def tokenize(
-        self, 
-        texts: Union[str, List[str]], 
-        context_length: int = 77, 
-        truncate: bool = True
-    ) -> torch.Tensor:
-        """
-        Tokenize text for use with AlphaCLIP models.
-        
-        Args:
-            texts: String or list of strings to tokenize
-            context_length: Maximum token length (default 77)
-            truncate: Whether to truncate long texts
-            
-        Returns:
-            Tensor of tokenized text
-        """
-        return tokenize(texts, context_length, truncate)
-    
-    def encode_text(self, model: torch.nn.Module, texts: Union[str, List[str]]) -> torch.Tensor:
-        """
-        Convenience method to tokenize and encode text.
-        
-        Args:
-            model: Loaded AlphaCLIP model
-            texts: Text(s) to encode
-            
-        Returns:
-            Text embeddings tensor
-        """
-        tokens = self.tokenize(texts)
-        if hasattr(model, 'token_embedding'):
-            # Move tokens to same device as model
-            device = next(model.parameters()).device
-            tokens = tokens.to(device)
-        
-        with torch.no_grad():
-            text_features = model.encode_text(tokens)
-            
-        return text_features
-    
-    def encode_image(self, model: torch.nn.Module, images: torch.Tensor) -> torch.Tensor:
-        """
-        Convenience method to encode images.
-        
-        Args:
-            model: Loaded AlphaCLIP model
-            images: Preprocessed image tensor
-            
-        Returns:
-            Image embeddings tensor
-        """
-        with torch.no_grad():
-            image_features = model.encode_image(images)
-            
-        return image_features
-    
-    def get_similarity(self, text_features: torch.Tensor, image_features: torch.Tensor) -> torch.Tensor:
-        """
-        Compute cosine similarity between text and image features.
-        
-        Args:
-            text_features: Text embedding tensor
-            image_features: Image embedding tensor
-            
-        Returns:
-            Similarity scores tensor
-        """
-        # Normalize features
-        text_features = text_features / text_features.norm(dim=-1, keepdim=True)
-        image_features = image_features / image_features.norm(dim=-1, keepdim=True)
-        
-        # Compute similarity
-        similarity = (text_features @ image_features.T)
-        return similarity
-
-
-# Convenience function for quick model loading
-def load_alphaclip(
-    model_name: str = "ViT-B/16",
-    device: Optional[str] = None,
-    alpha_vision_ckpt_pth: str = "None",
-    download_root = '/raid/datasets/models_weights/alphaclip',
-    **kwargs
-) -> Tuple[AlphaCLIPLoader, torch.nn.Module, callable]:
-    """
-    Quick function to load AlphaCLIP with a loader instance.
-    
-    Args:
-        model_name: Name of the model to load
-        device: Device to use
-        **kwargs: Additional arguments for model loading
-        
-    Returns:
-        Tuple of (loader, model, preprocess_function)
-    """
-    loader = AlphaCLIPLoader(default_device=device)
-    model, preprocess = loader.load_model(model_name, **kwargs)
-    return loader, model, preprocess
-
-
-# Make key functions available at module level
-__all__ = [
-    'AlphaCLIPLoader',
-    'load_alphaclip',
-    'available_models',
-    'load',
-    'tokenize'
-]

src/alphaclip/example.py

@@ -1,76 +0,0 @@
-#!/usr/bin/env python3
-"""
-Example usage of AlphaCLIP Standalone
-
-This script demonstrates basic usage of the AlphaCLIP standalone package.
-"""
-
-import torch
-import numpy as np
-from alphaclip_loader import AlphaCLIPLoader, load_alphaclip
-
-def main():
-    print("AlphaCLIP Standalone Example")
-    print("=" * 40)
-    
-    # Check if CUDA is available
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    print(f"Using device: {device}")
-    
-    # Method 1: Using the loader class
-    print("\n1. Using AlphaCLIPLoader class:")
-    loader = AlphaCLIPLoader(default_device=device)
-    
-    # Show available models
-    models = loader.available_models()
-    print(f"Available models: {models}")
-    
-    # Load a model
-    print("\nLoading ViT-B/16 model...")
-    model, preprocess = loader.load_model("ViT-B/16")
-    print(f"Model loaded successfully!")
-    
-    # Test text encoding
-    test_texts = [
-        "a photo of a cat",
-        "a dog running in the park",
-        "a beautiful sunset over the ocean"
-    ]
-    
-    print(f"\nEncoding {len(test_texts)} texts...")
-    text_features = loader.encode_text(model, test_texts)
-    print(f"Text features shape: {text_features.shape}")
-    
-    # Compute similarities between texts
-    print("\nComputing text-to-text similarities:")
-    similarities = loader.get_similarity(text_features, text_features)
-    
-    for i, text1 in enumerate(test_texts):
-        for j, text2 in enumerate(test_texts):
-            if i <= j:  # Only show upper triangle
-                sim = similarities[i, j].item()
-                print(f"  '{text1}' <-> '{text2}': {sim:.3f}")
-    
-    # Method 2: Using the quick loader function
-    print("\n\n2. Using quick loader function:")
-    loader2, model2, preprocess2 = load_alphaclip("ViT-B/16", device=device)
-    
-    # Test single text
-    single_text = "a red apple on a wooden table"
-    single_features = loader2.encode_text(model2, single_text)
-    print(f"Single text '{single_text}' encoded to shape: {single_features.shape}")
-    
-    # Test tokenization
-    print("\n3. Tokenization example:")
-    tokens = loader.tokenize(test_texts)
-    print(f"Tokenized {len(test_texts)} texts to shape: {tokens.shape}")
-    
-    # Show some token examples
-    print("First few tokens for each text:")
-    for i, text in enumerate(test_texts):
-        print(f"  '{text}': {tokens[i][:10].tolist()}...")
-    
-    print("\nExample completed successfully!")
-
-if __name__ == "__main__":
-    main()

src/alphaclip/requirements.txt

@@ -1,10 +0,0 @@
-# Core dependencies for AlphaCLIP standalone
-torch>=1.7.1
-torchvision
-ftfy
-regex
-tqdm
-loralib
-Pillow
-numpy
-packaging

src/alphaclip/setup.py

@@ -1,47 +0,0 @@
-"""
-Setup script for AlphaCLIP Standalone
-"""
-
-from setuptools import setup, find_packages
-import os
-
-# Read requirements
-with open('requirements.txt', 'r') as f:
-    requirements = [line.strip() for line in f if line.strip() and not line.startswith('#')]
-
-# Read README if it exists
-readme_content = ""
-if os.path.exists('README.md'):
-    with open('README.md', 'r', encoding='utf-8') as f:
-        readme_content = f.read()
-
-setup(
-    name="alphaclip-standalone",
-    version="1.0.0",
-    author="AlphaCLIP Team",
-    description="Standalone version of AlphaCLIP for easy integration",
-    long_description=readme_content,
-    long_description_content_type="text/markdown",
-    packages=find_packages(),
-    package_data={
-        'alpha_clip': ['*.gz'],  # Include the tokenizer vocabulary file
-    },
-    include_package_data=True,
-    install_requires=requirements,
-    python_requires=">=3.7",
-    classifiers=[
-        "Development Status :: 4 - Beta",
-        "Intended Audience :: Developers",
-        "Intended Audience :: Science/Research",
-        "License :: OSI Approved :: MIT License",
-        "Operating System :: OS Independent",
-        "Programming Language :: Python :: 3",
-        "Programming Language :: Python :: 3.7",
-        "Programming Language :: Python :: 3.8",
-        "Programming Language :: Python :: 3.9",
-        "Programming Language :: Python :: 3.10",
-        "Topic :: Scientific/Engineering :: Artificial Intelligence",
-        "Topic :: Software Development :: Libraries :: Python Modules",
-    ],
-    keywords="clip, vision, language, deep learning, pytorch",
-)

src/alphaclip/test_installation.py

@@ -1,149 +0,0 @@
-#!/usr/bin/env python3
-"""
-Test script for AlphaCLIP Standalone
-
-This script tests the basic functionality of the standalone package
-to ensure everything is working correctly.
-"""
-
-import sys
-import os
-
-def test_imports():
-    """Test that all required modules can be imported."""
-    print("Testing imports...")
-    
-    try:
-        import torch
-        print(f"✓ PyTorch {torch.__version__} imported successfully")
-    except ImportError as e:
-        print(f"✗ Failed to import PyTorch: {e}")
-        return False
-    
-    try:
-        import torchvision
-        print(f"✓ Torchvision imported successfully")
-    except ImportError as e:
-        print(f"✗ Failed to import torchvision: {e}")
-        return False
-    
-    try:
-        from alphaclip_loader import AlphaCLIPLoader
-        print("✓ AlphaCLIPLoader imported successfully")
-    except ImportError as e:
-        print(f"✗ Failed to import AlphaCLIPLoader: {e}")
-        return False
-    
-    try:
-        import loralib
-        print("✓ LoraLib imported successfully")
-    except ImportError as e:
-        print(f"✗ Failed to import loralib: {e}")
-        return False
-    
-    return True
-
-def test_model_loading():
-    """Test loading a model."""
-    print("\nTesting model loading...")
-    
-    try:
-        from alphaclip_loader import AlphaCLIPLoader
-        
-        loader = AlphaCLIPLoader(default_device="cpu")  # Use CPU for testing
-        models = loader.available_models()
-        print(f"✓ Available models: {models}")
-        
-        # Try to load the smallest model for testing
-        print("Loading ViT-B/32 model (this may take a while for first download)...")
-        model, preprocess = loader.load_model("ViT-B/32", device="cpu")
-        print("✓ Model loaded successfully")
-        
-        return True
-        
-    except Exception as e:
-        print(f"✗ Failed to load model: {e}")
-        return False
-
-def test_tokenization():
-    """Test text tokenization."""
-    print("\nTesting tokenization...")
-    
-    try:
-        from alphaclip_loader import AlphaCLIPLoader
-        
-        loader = AlphaCLIPLoader()
-        test_text = "a photo of a cat"
-        tokens = loader.tokenize(test_text)
-        print(f"✓ Tokenized '{test_text}' to shape {tokens.shape}")
-        
-        # Test batch tokenization
-        test_texts = ["a cat", "a dog", "a bird"]
-        batch_tokens = loader.tokenize(test_texts)
-        print(f"✓ Batch tokenized {len(test_texts)} texts to shape {batch_tokens.shape}")
-        
-        return True
-        
-    except Exception as e:
-        print(f"✗ Failed tokenization test: {e}")
-        return False
-
-def test_text_encoding():
-    """Test text encoding with a loaded model."""
-    print("\nTesting text encoding...")
-    
-    try:
-        from alphaclip_loader import AlphaCLIPLoader
-        
-        loader = AlphaCLIPLoader(default_device="cpu")
-        model, preprocess = loader.load_model("ViT-B/32", device="cpu")
-        
-        test_text = "a photo of a cat"
-        features = loader.encode_text(model, test_text)
-        print(f"✓ Encoded text to features with shape {features.shape}")
-        
-        # Test batch encoding
-        test_texts = ["a cat", "a dog"]
-        batch_features = loader.encode_text(model, test_texts)
-        print(f"✓ Batch encoded {len(test_texts)} texts to shape {batch_features.shape}")
-        
-        return True
-        
-    except Exception as e:
-        print(f"✗ Failed text encoding test: {e}")
-        return False
-
-def main():
-    """Run all tests."""
-    print("AlphaCLIP Standalone Test Suite")
-    print("=" * 40)
-    
-    tests = [
-        test_imports,
-        test_tokenization,
-        test_model_loading,
-        test_text_encoding,
-    ]
-    
-    passed = 0
-    total = len(tests)
-    
-    for test in tests:
-        try:
-            if test():
-                passed += 1
-        except Exception as e:
-            print(f"✗ Test {test.__name__} failed with exception: {e}")
-    
-    print(f"\n{'='*40}")
-    print(f"Test Results: {passed}/{total} tests passed")
-    
-    if passed == total:
-        print("🎉 All tests passed! AlphaCLIP Standalone is working correctly.")
-        return 0
-    else:
-        print("❌ Some tests failed. Please check the error messages above.")
-        return 1
-
-if __name__ == "__main__":
-    sys.exit(main())

src/bbox_utils.py

@@ -1,421 +0,0 @@
-import torch
-from copy import deepcopy
-from PIL import ImageDraw
-import itertools
-import random
-
-
-def extract_bboxes_feats(patch_embeddings, bboxes, gaussian_avg=False, 
-                         gaussian_bbox_variance=0.5, get_single_embedding_per_image=False,
-                         patch_size=14, attention_map=None):
-    """
-    if get_single_embedding_per_image is True, the weights of all the bounding boxes patches on an image will be summed and the function will return the patch weights depending on this map
-    """
-    N = patch_embeddings.shape[0]
-    N_boxes = bboxes.shape[1]
-    grid_size = int(patch_embeddings.shape[1]**0.5)
-    device = patch_embeddings.device
-
-    bboxes //= patch_size
-    bboxes = bboxes.int()
-
-    # Reshape patches to grid
-    patch_embeddings = patch_embeddings.view(N, grid_size, grid_size, -1)  # Shape (N, grid_size, grid_size, embed_dim)
-    if attention_map is not None:
-        attention_map = attention_map.view(N, grid_size, grid_size)  # Shape (N, grid_size, grid_size)
-    # Grid of the sum of the gaussian weights
-    total_patch_weights = torch.zeros(N, grid_size, grid_size)
-
-    # Extract boxes
-    x1, y1, w, h = bboxes.unbind(-1)  # Separate box dimensions (N, N_boxes)
-
-    # Create mesh grid for slicing
-    x2 = x1 + w  # Exclusive end x
-    y2 = y1 + h  # Exclusive end y
-
-    means = []
-    for i in range(N):
-        image_means = []
-        for j in range(N_boxes):
-            if bboxes[i, j].sum().item() < 0 and get_single_embedding_per_image:
-                # this is the case where we receive a dummy box
-                continue
-            # Extract the region for each box
-            region_patches = patch_embeddings[i, y1[i, j]:y2[i, j] + 1, x1[i, j]:x2[i, j] + 1, :]  # (h, w, embed_dim)
-            
-            if attention_map is not None:
-                patch_weights = attention_map[i, y1[i, j]:y2[i, j] + 1, x1[i, j]:x2[i, j] + 1]
-                patch_weights /= patch_weights.sum()
-                total_patch_weights[i, y1[i, j]:y2[i, j] + 1, x1[i, j]:x2[i, j] + 1] += patch_weights
-                
-                weighted_patches = region_patches * patch_weights.to(device).unsqueeze(-1)  # (h, w, embed_dim)
-                region_mean = weighted_patches.sum(dim=(0, 1))  # Weighted mean
-                
-            elif gaussian_avg:
-                # Create Gaussian weights
-                h_span, w_span = region_patches.shape[:2]
-                y_coords, x_coords = torch.meshgrid(
-                    torch.linspace(-1, 1, h_span),
-                    torch.linspace(-1, 1, w_span),
-                    indexing="ij"
-                )
-                if gaussian_bbox_variance == 0:
-                    patch_weights = torch.zeros((h_span, w_span))
-                    # Determine central indices
-                    center_y = [h_span // 2] if h_span % 2 == 1 else [h_span // 2 - 1, h_span // 2]
-                    center_x = [w_span // 2] if w_span % 2 == 1 else [w_span // 2 - 1, w_span // 2]
-                    # Randomly select one of the central elements in even case
-                    center_y = random.choice(center_y)
-                    center_x = random.choice(center_x)
-                    # Set the selected central element to 1
-                    patch_weights[center_y, center_x] = 1.0
-                else:
-                    distances = x_coords**2 + y_coords**2
-                    patch_weights = torch.exp(-distances / gaussian_bbox_variance)
-                    patch_weights = patch_weights / patch_weights.sum()  # Normalize to sum to 1
-
-                # Apply Gaussian weights to region patches
-                weighted_patches = region_patches * patch_weights.to(device).unsqueeze(-1)  # (h, w, embed_dim)
-                region_mean = weighted_patches.sum(dim=(0, 1))  # Weighted mean
-                
-                # Recording the bbox weight inside the image patch weight map
-                total_patch_weights[i, y1[i, j]:y2[i, j] + 1, x1[i, j]:x2[i, j] + 1] += patch_weights
-            else:
-                # Mean pooling case: create uniform weights
-                h_span, w_span = region_patches.shape[:2]
-                uniform_weights = torch.ones(h_span, w_span) / (h_span * w_span)
-                
-                # Update total_patch_weights for mean pooling
-                total_patch_weights[i, y1[i,j]:y2[i,j]+1, x1[i,j]:x2[i,j]+1] += uniform_weights
-                
-                # Compute mean of the region
-                region_mean = region_patches.mean(dim=(0, 1))
-
-            # Store the mean
-            image_means.append(region_mean)
-        if not get_single_embedding_per_image:
-            means.append(torch.stack(image_means))
-
-    # Normalizing the weight map so the sum is equal to 1
-    total_patch_weights /= total_patch_weights.sum(dim=(1,2), keepdim=True)
-    if not get_single_embedding_per_image:
-        return torch.stack(means)  # Shape (N, N_boxes, embed_dim)
-    else:
-        # Expand dimensions to match embeddings
-        total_patch_weights = total_patch_weights.unsqueeze(-1).to(device)
-
-        # Compute weighted sum
-        weighted_patch_mean = (total_patch_weights * patch_embeddings).sum(dim=(1, 2))  
-        return  weighted_patch_mean 
-# Shape (N, embed_dim)
-
-#def adjust_bbox_for_transform(image, bbox, resize_dim, crop_dim):
-#    """
-#    Adjusts the bounding box for a resized and center-cropped image.
-#
-#    Args:
-#        image (PIL.Image): The input image.
-#        bbox (list): The bounding box in [x1, y1, w, h] format.
-#        resize_dim (int): The dimension of the shortest side after resizing.
-#        crop_dim (int): The size of the square crop.
-#
-#    Returns:
-#        list: The adjusted bounding box in [x1, y1, w, h] format.
-#    """
-#    x1, y1, w, h = bbox
-#    orig_width, orig_height = image.size
-#
-#    # Calculate resize scale for the shortest side
-#    if orig_width < orig_height:
-#        scale = resize_dim / orig_width
-#        resized_width, resized_height = resize_dim, int(orig_height * scale)
-#    else:
-#        scale = resize_dim / orig_height
-#        resized_width, resized_height = int(orig_width * scale), resize_dim
-#
-#    # Scale the bounding box
-#    x1 *= scale
-#    y1 *= scale
-#    w *= scale
-#    h *= scale
-#
-#    # Calculate cropping offsets
-#    crop_x = (resized_width - crop_dim) // 2
-#    crop_y = (resized_height - crop_dim) // 2
-#
-#    # Adjust bounding box for cropping
-#    x1 -= crop_x
-#    y1 -= crop_y
-#
-#    # Clamp the bounding box to the cropped area
-#    x1 = max(0, x1)
-#    y1 = max(0, y1)
-#    w = min(w, crop_dim - x1)
-#    h = min(h, crop_dim - y1)
-#
-#    return [x1, y1, w, h]
-
-def map_traces_to_grid(traces, n_patch):
-    grid = torch.zeros((n_patch, n_patch))
-    patch_size = 1.0 / n_patch
-    
-    for trace in traces:
-        x, y = trace['x'], trace['y']
-        if 0 <= x <= 1 and 0 <= y <= 1:
-            grid_x, grid_y = int(x / patch_size), int(y / patch_size)
-            grid[min(grid_y, n_patch - 1), min(grid_x, n_patch - 1)] += 1
-    
-    return grid
-
-def adjust_bbox_for_transform(image, bbox, resize_dim, crop_dim):
-    """
-    Adjusts the bounding box for a resized and center-cropped image.
-
-    Args:
-        image (PIL.Image): The input image.
-        bbox (list): The bounding box in [x1, y1, w, h] format.
-        resize_dim (int): The dimension of the shortest side after resizing.
-        crop_dim (int): The size of the square crop.
-
-    Returns:
-        list: The adjusted bounding box in [x1, y1, w, h] format.
-    """
-    x1, y1, w, h = bbox
-    orig_width, orig_height = image.size
-
-    # Scale factors for resizing
-    if orig_width < orig_height:
-        scale_w = resize_dim / orig_width
-        scale_h = (resize_dim * orig_height) / orig_width / orig_height
-    else:
-        scale_h = resize_dim / orig_height
-        scale_w = (resize_dim * orig_width) / orig_height / orig_width
-
-    # New dimensions after resize
-    new_width = int(orig_width * scale_w)
-    new_height = int(orig_height * scale_h)
-
-    # Update bounding box for resizing
-    x1 = x1 * scale_w
-    y1 = y1 * scale_h
-    w = w * scale_w
-    h = h * scale_h
-
-    # Compute cropping offsets
-    crop_x_offset = max(0, (new_width - crop_dim) // 2)
-    crop_y_offset = max(0, (new_height - crop_dim) // 2)
-
-    # Adjust bounding box for cropping
-    x1 -= crop_x_offset
-    y1 -= crop_y_offset
-
-    # Clip bounding box to crop dimensions
-    x1 = max(0, min(x1, crop_dim - 1))
-    y1 = max(0, min(y1, crop_dim - 1))
-    w = max(0, min(w, crop_dim - x1))
-    h = max(0, min(h, crop_dim - y1))
-
-    return [x1, y1, w, h]
-
-
-
-def adjust_bbox_for_transform_no_scale(image, bbox, target_width, target_height):
-    """
-    - Does not preserve the image scale.
-    Adjusts the bounding box for an image resized to a fixed width and height.
-
-    Args:
-        image (PIL.Image): The original image.
-        bbox (list): The bounding box in [x1, y1, w, h] format.
-        target_width (int): The width of the resized image.
-        target_height (int): The height of the resized image.
-
-    Returns:
-        list: The adjusted bounding box in [x1, y1, w, h] format.
-    """
-    x1, y1, w, h = bbox
-    orig_width, orig_height = image.size
-
-    # Calculate scale factors for width and height
-    scale_w = target_width / orig_width
-    scale_h = target_height / orig_height
-
-    # Adjust the bounding box
-    x1 = x1 * scale_w
-    y1 = y1 * scale_h
-    w = w * scale_w
-    h = h * scale_h
-
-    # Return the adjusted bounding box
-    return [x1, y1, w, h]
-
-
-def draw_bounding_boxes(input_image, bounding_boxes, captions=[""], color="red", width=2, text_background=True, boxes_to_show = None):
-    """
-    Draws bounding boxes on an image.
-
-    Args:
-        image (PIL.Image): The image to draw on.
-        bounding_boxes (list): A list of bounding boxes, each as [x1, y1, x2, y2].
-        color (str): The color of the bounding boxes (default is red).
-        width (int): The width of the bounding box lines (default is 2).
-
-    Returns:
-        PIL.Image: The image with bounding boxes drawn.
-    """
-    # Create a drawing context
-    image = deepcopy(input_image)
-    draw = ImageDraw.Draw( image )
-
-    #scale = 720.0 / max(image.size)
-    if boxes_to_show is not None:
-        if isinstance(boxes_to_show, int):
-            indexes_to_show = random.sample(range(len(bounding_boxes)), boxes_to_show)
-        else:
-            indexes_to_show = boxes_to_show
-    
-    for i, (bbox, cap ) in enumerate(itertools.zip_longest(bounding_boxes, captions, fillvalue="")):
-
-        if boxes_to_show is not None:
-            if i not in indexes_to_show: continue
-        #bbox = [ i / scale for i in bbox ]
-        #x1, y1, w, h = bbox
-        x1, y1, x2, y2 = bbox
-        
-        #x2, y2 = x1 + w, y1 + h  # Convert width/height to bottom-right corner
-        try:
-            draw.rectangle([x1, y1, x2, y2], outline=color, width=width)
-            if cap != "":
-                if text_background:
-                    left,top,right,bottom = draw.multiline_textbbox((x1,y1), cap) #textbbox
-                    draw.rectangle((left-5, top-5, right+5, bottom+5), fill="white")
-                draw.multiline_text((x1,y1), cap, fill=color)   #text
-            
-        except Exception as e:
-            print("exception, i: ", i, f"{x1 = } {y1 = } {x2 = }, {y2 = }")
-            print(e)
-    
-    return image
-
-def extract_bboxes_feats_double_dino(dino_model, patch_embeddings, bboxes, cls_token, registers_tokens, patch_size, return_type="cls", gaussian_bbox_variance=0.5):
-    """
-    Perform a forward pass of the last DINO layer with selected features, batched.
-
-    Args:
-        dino_model: The DINO model.
-        patch_embeddings: Patch embeddings before the last layer.
-        bboxes: Bounding boxes for each image in the batch (BS x N_BOX_MAX x 4).
-        cls_token: CLS token embedding.
-        return_type: Type of feature to return ('cls', 'avg', 'gaussian_avg').
-        gaussian_bbox_variance: Variance for Gaussian averaging.
-
-    Returns:
-        bbox_features: Features for each bounding box based on return_type.
-    """
-    N = patch_embeddings.shape[0]  # Batch size
-    N_boxes = bboxes.shape[1]      # Number of bounding boxes
-    grid_size = int(patch_embeddings.shape[1] ** 0.5)  # Assuming square grid
-    embed_dim = patch_embeddings.shape[-1]
-
-    bboxes_patch_indexes = bboxes.clone()
-    bboxes_patch_indexes //= patch_size  # Scale down bbox coordinates to match patch grid
-    bboxes_patch_indexes = bboxes_patch_indexes.int()
-
-    # Reshape patches to grid
-    patch_embeddings = patch_embeddings.view(N, grid_size, grid_size, embed_dim)  # (N, grid_size, grid_size, embed_dim)
-
-    if cls_token is not None:
-        cls_tokens = cls_token.view(N, embed_dim)
-        if registers_tokens is not None:
-            patches_offset = 5
-        else:
-            patches_offset = 1
-    else:
-        assert return_type != "cls"
-        patches_offset = 0
-    batch_outputs = []
-
-    #batch_inputs = []
-
-    means = []
-    for i in range(N):  # Iterate over batch
-        image_means = []
-        
-        if cls_token is not None:
-            cls_cur_img = cls_tokens[i].reshape(1, 1, embed_dim)
-            if registers_tokens is not None:
-                cur_img_register_tokens = registers_tokens[i].reshape(1, 4, embed_dim)
-            
-        for j in range(N_boxes):  # Iterate over bounding boxes
-            # Extract the region for the bounding box
-            region_patches_xy = patch_embeddings[i, bboxes_patch_indexes[i, j, 1]:bboxes_patch_indexes[i, j, 3] + 1, bboxes_patch_indexes[i, j, 0]:bboxes_patch_indexes[i, j, 2] + 1, :]
-            #region_patches = region_patches.reshape(-1, embed_dim)  # Flatten to (num_patches, embed_dim)
-
-            #region_patches = region_patches.view(-1, embed_dim)  # Flatten to (num_patches, embed_dim)
-            #cls_cur_img = cls_tokens[i].unsqueeze(0)  # Add batch dimension (1, embed_dim)
-            #region_patches = region_patches.unsqueeze(0)  # Add batch dimension (1, num_patches, embed_dim)
-            region_patches = region_patches_xy.reshape(1,-1, embed_dim)
-            if cls_token is not None:
-                inputs = torch.cat([cls_cur_img, region_patches], dim=1)  # Concatenate along the token dimension (1, num_patches + 1, embed_dim)
-                if registers_tokens is not None:
-                    inputs = torch.cat([cls_cur_img, cur_img_register_tokens, region_patches], dim=1)  # Concatenate along the token dimension (1, num_patches + 5, embed_dim)
-            else:
-                inputs = torch.cat([region_patches], dim=1)  # Concatenate along the token dimension (1, num_patches + 1, embed_dim)
-                
-            outputs = dino_model.blocks[-1](inputs)  # Forward pass
-            # shape (1, 1 + len(region_patches), 768)
-            #cls_cur_img = cls_tokens[i]
-            #cls_cur_img = cls_cur_img.reshape(1, embed_dim)
-            #inputs = torch.cat([cls_cur_img, region_patches], dim=0)  # Add CLS token to inputs
-            #outputs = dino_model.blocks[-1](inputs)  # Forward pass
-            
-            batch_outputs.append(outputs)
-
-            region_patches = outputs[0, patches_offset: ,] #(1,45,768) -> (1,1,768)
-            
-            if return_type == "gaussian_avg":
-                #region_patches = outputs[5: ,]
-                h_span, w_span = region_patches_xy.shape[:2]
-                y_coords, x_coords = torch.meshgrid(
-                    torch.linspace(-1, 1, h_span),
-                    torch.linspace(-1, 1, w_span),
-                    indexing="ij"
-                )
-                distances = x_coords**2 + y_coords**2
-                gaussian_weights = torch.exp(-distances / gaussian_bbox_variance)  # Adjust 0.1 for variance control
-                gaussian_weights = gaussian_weights / gaussian_weights.sum()  # Normalize to sum to 1
-
-                # Apply Gaussian weights to region patches
-                weighted_patches = region_patches_xy * gaussian_weights.to(next(dino_model.parameters()).device).unsqueeze(-1)  # (h, w, embed_dim)
-                region_mean = weighted_patches.sum(dim=(0,1))  # Weighted mean
-                #image_means.append(region_mean)
-            elif return_type == "avg":
-                # Compute mean of the region
-                region_mean = region_patches.mean(dim=(0))  # Mean over h, w
-            elif return_type == "cls":
-                region_mean = outputs[0, 0, ]
-            image_means.append(region_mean)
-
-        means.append(torch.stack(image_means))
-
-    stacked_means = torch.stack(means)
-    #stacked_means = stacked_means.reshape(-1, embed_dim)
-    return stacked_means
-
-
-def process_bboxes(imgs, bboxes, transform):
-    transformed_bboxes = []
-    bboxes = bboxes.tolist()
-    for img, img_bboxes in zip(imgs, bboxes):
-        for bbox in img_bboxes:
-            # Crop the region defined by bbox
-            x_min, y_min, w, h = bbox
-            x_max = x_min + w
-            y_max = y_min + h
-            cropped_region = img.crop((x_min, y_min, x_max, y_max))
-            
-            # Apply the transform to the cropped region
-            transformed_region = transform(cropped_region)
-            transformed_bboxes.append(transformed_region)
-    
-    return torch.stack(transformed_bboxes)

src/clipcap/CLIPCAP_INTEGRATION.md

@@ -1,206 +0,0 @@
-# ClipCap Integration with Patchioner Class
-
-This document describes how ClipCap models have been integrated into the Patchioner class for DINO feature-based image captioning.
-
-## Overview
-
-ClipCap support has been added to the Patchioner class following the same pattern as other captioning models (VieCap, MeaCap, etc.). This integration allows you to use trained ClipCap models with DINO features for image captioning tasks.
-
-## Architecture
-
-### Files Added/Modified
-
-1. **`src/clipcap/entrypoint.py`** - Main ClipCap integration module
-   - `ClipCapModel` class for DINO feature-based captioning
-   - Model classes: `ClipCaptionModel`, `ClipCaptionPrefix`, `MLP`, `TransformerMapper`
-   - Text generation utilities
-
-2. **`src/model.py`** - Modified Patchioner class
-   - Added `clipcap_config` parameter to constructor
-   - Added ClipCap initialization logic
-   - Added ClipCap support to `caption_tokens` method
-
-3. **Configuration Files**
-   - `configs/clipcap_dino_vitb14.k.yaml` - DINOv2-B/14 configuration
-   - `configs/clipcap_dino_vitl14.k.yaml` - DINOv2-L/14 configuration
-
-## Configuration
-
-### YAML Configuration Format
-
-```yaml
-decap_weights: '/path/to/decap/weights.pt'
-prefix_size: 768  # DINO feature dimension
-support_memory_size: 0
-dino_model: 'dinov2_vitb14'
-normalize: True
-resize_dim: 518
-crop_dim: 518
-use_talk2dino_project: False
-
-# ClipCap configuration
-clipcap:
-  language_model: 'gpt2'
-  prefix_length: 10            # Sequence length for prefix
-  clip_length: 10              # CLIP sequence length (for transformer mapping)
-  num_layers: 8                # Number of transformer layers (for transformer mapping)
-  mapping_type: 'mlp'          # 'mlp' or 'transformer'
-  only_prefix: True            # Train only prefix mapping vs full model
-  temperature: 1.0             # Sampling temperature
-  top_p: 0.8                   # Nucleus sampling parameter
-  entry_length: 67             # Maximum caption length
-  stop_token: '.'              # Stop token for generation
-  weight_path: '/path/to/trained/clipcap/model.pt'
-```
-
-### Supported DINO Models
-
-The integration automatically detects DINO feature dimensions:
-
-- **DINOv2-S/14**: 384 dimensions (`dinov2_vits14`)
-- **DINOv2-B/14**: 768 dimensions (`dinov2_vitb14`)
-- **DINOv2-L/14**: 1024 dimensions (`dinov2_vitl14`)
-- **DINOv2-G/14**: 1536 dimensions (`dinov2_vitg14`)
-
-## Usage
-
-### 1. Training ClipCap Models
-
-First, train your ClipCap model with DINO features:
-
-```bash
-# Extract DINO features
-python clipcap_dino_parse_coco.py --dino_model_type dinov2_vitb14
-
-# Train ClipCap model
-python clipcapTraining.py \
-    --use_dino \
-    --dino_model_type dinov2_vitb14 \
-    --prefix_length 10 \
-    --mapping_type mlp \
-    --only_prefix \
-    --epochs 10
-```
-
-### 2. Using ClipCap with Patchioner
-
-```python
-import torch
-from src.model import Patchioner
-
-# Load model with ClipCap configuration
-device = torch.device('cuda')
-model = Patchioner.from_config('configs/clipcap_dino_vitb14.k.yaml', device)
-
-# Generate captions from images
-imgs = torch.randn(2, 3, 518, 518).to(device)  # Example batch
-results = model.forward(imgs, get_cls_capt=True)
-captions = results['cls_capt']
-
-print("Generated captions:")
-for i, caption in enumerate(captions):
-    print(f"Image {i+1}: {caption}")
-```
-
-### 3. Using ClipCap Directly
-
-```python
-from src.clipcap.entrypoint import ClipCapModel
-import torch
-
-# Configuration
-config = {
-    'language_model': 'gpt2',
-    'prefix_length': 10,
-    'mapping_type': 'mlp',
-    'only_prefix': True,
-    'weight_path': '/path/to/trained/model.pt'
-}
-
-# Initialize model
-device = torch.device('cuda')
-clipcap = ClipCapModel(config, device, dino_feature_dim=768)
-
-# Generate captions from DINO features
-dino_features = torch.randn(2, 768).to(device)
-captions = clipcap.forward(dino_features)
-
-print(captions)
-```
-
-## Performance Improvements
-
-### Batched Text Generation
-
-The ClipCap integration includes an efficient batched text generation implementation:
-
-- **`generate_batched()`**: Processes entire batches simultaneously
-- **Significant speedup**: 2-8x faster than sequential processing
-- **Memory efficient**: Optimized for GPU memory usage
-- **Configurable**: Can fallback to sequential mode if needed
-
-### Configuration Options
-
-```yaml
-clipcap:
-  use_batched_generation: True  # Enable batched generation (recommended)
-  temperature: 1.0              # Sampling temperature
-  top_p: 0.8                   # Nucleus sampling parameter
-  entry_length: 67             # Maximum sequence length
-```
-
-## Model Architecture Details
-
-### ClipCap Model Structure
-
-1. **Input**: DINO features (384/768/1024/1536 dimensions)
-2. **Mapping Layer**: 
-   - **MLP**: `DINO_dim → GPT2_dim * prefix_length`
-   - **Transformer**: Multi-layer transformer mapping
-3. **GPT-2 Decoder**: Pretrained GPT-2 for text generation
-4. **Output**: Natural language captions
-
-### Key Components
-
-- **`ClipCapModel`**: Main class for DINO-to-text captioning
-- **`MLP`/`TransformerMapper`**: Feature mapping from DINO to GPT-2 space
-- **Text Generation**: Nucleus sampling with configurable parameters
-
-## Integration with Existing Pipeline
-
-The ClipCap integration follows the established pattern:
-
-1. **Configuration**: YAML-based configuration like other models
-2. **Initialization**: Automatic DINO dimension detection
-3. **Forward Pass**: Seamless integration with existing forward methods
-4. **Scoring**: Optional confidence scoring support
-
-## Testing
-
-Run the integration test:
-
-```bash
-python test_clipcap_integration.py
-```
-
-This test verifies:
-- Configuration loading from YAML
-- Model instantiation with ClipCap
-- Caption generation with dummy DINO features
-- Score computation functionality
-
-
-## Troubleshooting
-
-### Common Issues
-
-1. **Dimension Mismatch**: Ensure `prefix_size` matches DINO model dimension
-2. **Missing Weights**: Verify `weight_path` points to trained ClipCap model
-3. **Memory Issues**: Use `only_prefix=True` for lower memory usage
-4. **Generation Quality**: Tune `temperature`, `top_p`, and `entry_length`
-
-## References
-
-- [ClipCap Paper](https://arxiv.org/abs/2111.09734)
-- [DINO Paper](https://arxiv.org/abs/2104.14294)
-- [DINOv2 Paper](https://arxiv.org/abs/2304.07193)

src/clipcap/clipcapTrainREADME.md

@@ -1,301 +0,0 @@
-# ClipCap Training with DINO Features - README
-
-This guide provides instructions for training ClipCap with DINO visual features instead of CLIP features.
-
-## Prerequisites
-
-1. Ensure you have the required dependencies installed:
-   - PyTorch
-   - torchvision
-   - transformers
-   - tqdm
-   - Pillow
-   - scikit-image
-
-2. Prepare your COCO dataset with the following structure:
-   ```
-   ./data/coco/
-   ├── annotations/
-   │   └── train_caption.json
-   ├── train2014/
-   │   └── COCO_train2014_*.jpg
-   └── val2014/
-       └── COCO_val2014_*.jpg
-   ```
-
-## Required Files for DINO Feature Extraction
-
-To start the DINO feature extraction for the COCO dataset, you need:
-
-### 1. **COCO Dataset Structure**:
-```
-/raid/datasets/coco/                     # Main COCO directory (default)
-├── train2014/                           # REQUIRED: Training images
-│   └── COCO_train2014_*.jpg            # Image files
-├── val2014/                             # REQUIRED: Validation images
-│   └── COCO_val2014_*.jpg               # Image files
-└── train_split_karpathy.json            # REQUIRED: Karpathy format annotations (default)
-```
-
-### 2. **Required Files**:
-- **`train_split_karpathy.json`**: COCO caption annotations in Karpathy format (default)
-- **Training images**: COCO 2014 training set (COCO_train2014_*.jpg)
-- **Validation images**: COCO 2014 validation set (COCO_val2014_*.jpg)
-
-### 3. **Annotation Format Support**:
-
-The script supports two annotation formats:
-
-#### **A. Karpathy Format** (default, recommended):
-```json
-{
-    "images": [
-        {"id": 522418, "file_name": "COCO_val2014_000000522418.jpg"}
-    ],
-    "annotations": [
-        {"image_id": 522418, "id": 0, "caption": "A woman wearing a net..."}
-    ]
-}
-```
-
-#### **B. ClipCap Format** (legacy):
-```json
-[
-    {"image_id": 522418, "caption": "A woman wearing a net..."}
-]
-```
-
-### 3. **Specifying Custom Input/Output Paths**:
-
-You can customize the paths using command-line arguments:
-
-```bash
-python clipcap_dino_parse_coco.py \
-    --dino_model_type dinov2_vitb14 \
-    --coco_images_dir "/path/to/your/coco/dataset" \
-    --captions_file "/path/to/your/train_caption.json" \
-    --output_file "/path/to/output/dino_features.pkl"
-```
-
-**Available path arguments**:
-- `--coco_images_dir`: Path to COCO images directory (should contain `train2014/` and `val2014/` subdirs) - **Default: `/raid/datasets/coco`**
-- `--captions_file`: Path to COCO captions JSON file (supports both Karpathy and ClipCap formats) - **Default: `/raid/datasets/coco/train_split_karpathy.json`**
-- `--output_file`: Custom output file path (optional, auto-generated if not specified)
-
-### 4. **Default Behavior** (if no paths specified):
-```bash
-# This will use default paths for your setup:
-python clipcap_dino_parse_coco.py --dino_model_type dinov2_vitb14
-
-# Equivalent to:
-python clipcap_dino_parse_coco.py \
-    --dino_model_type dinov2_vitb14 \
-    --coco_images_dir "/raid/datasets/coco" \
-    --captions_file "/raid/datasets/coco/train_split_karpathy.json" \
-    --output_file "/raid/datasets/coco/coco_karpathy_split_dinov2_vitb14_train.pkl"
-```
-
-## Step 1: Extract DINO Features
-
-First, extract DINO features from the COCO images using the modified feature extraction script:
-
-### For DINOv2-B/14 (768-dim features):
-```bash
-# Default paths (uses /raid/datasets/coco and Karpathy annotations)
-python clipcap_dino_parse_coco.py --dino_model_type dinov2_vitb14 --resize_dim 518 --crop_dim 518
-
-# Custom paths
-python clipcap_dino_parse_coco.py \
-    --dino_model_type dinov2_vitb14 \
-    --coco_images_dir "/your/coco/path" \
-    --captions_file "/your/coco/train_split_karpathy.json" \
-    --output_file "/your/output/dino_vitb14_features.pkl"
-```
-
-### For DINOv2-L/14 (1024-dim features):
-```bash
-# Default paths
-python clipcap_dino_parse_coco.py --dino_model_type dinov2_vitl14 --resize_dim 518 --crop_dim 518
-
-# Custom paths
-python clipcap_dino_parse_coco.py \
-    --dino_model_type dinov2_vitl14 \
-    --coco_images_dir "/your/coco/path" \
-    --output_file "/your/output/dino_vitl14_features.pkl"
-```
-
-### For DINOv2-S/14 (384-dim features):
-```bash
-python clipcap_dino_parse_coco.py --dino_model_type dinov2_vits14 --resize_dim 518 --crop_dim 518
-```
-
-### For DINOv2-G/14 (1536-dim features):
-```bash
-python clipcap_dino_parse_coco.py --dino_model_type dinov2_vitg14 --resize_dim 518 --crop_dim 518
-```
-
-**Output**: This will create a file like `/raid/datasets/models_weights/clipcap/training-features/coco_karpathy_split_dinov2_vitb14_train.pkl` (or your custom path) containing the DINO features and captions.
-
-### Check Available Arguments:
-```bash
-python clipcap_dino_parse_coco.py --help
-```
-
-## Step 2: Train ClipCap with DINO Features
-
-### Basic Training Command (MLP with sequence length 10):
-
-For **DINOv2-B/14** with **MLP mapping** and **prefix length 10**:
-```bash
-python clipcapTraining.py \
-    --data /raid/datasets/models_weights/clipcap/training-features/coco_karpathy_split_dinov2_vitb14_train.pkl \
-    --out_dir ./checkpoints_dino_vitb14_mlp_len10 \
-    --prefix dino_vitb14_mlp_len10 \
-    --epochs 10 \
-    --save_every 2 \
-    --prefix_length 10 \
-    --bs 32 \
-    --mapping_type mlp \
-    --use_dino \
-    --dino_model_type dinov2_vitb14 \
-    --only_prefix
-```
-
-### Training Options for Different DINO Models:
-
-#### DINOv2-L/14 (1024-dim):
-```bash
-python clipcapTraining.py \
-    --data ./data/coco/coco_karpathy_split_dinov2_vitl14_train.pkl \
-    --out_dir ./checkpoints_dino_vitl14_mlp_len10 \
-    --prefix dino_vitl14_mlp_len10 \
-    --epochs 10 \
-    --save_every 2 \
-    --prefix_length 10 \
-    --bs 32 \
-    --mapping_type mlp \
-    --use_dino \
-    --dino_model_type dinov2_vitl14 \
-    --only_prefix
-```
-
-#### DINOv2-S/14 (384-dim):
-```bash
-python clipcapTraining.py \
-    --data ./data/coco/coco_karpathy_split_dinov2_vits14_train.pkl \
-    --out_dir ./checkpoints_dino_vits14_mlp_len10 \
-    --prefix dino_vits14_mlp_len10 \
-    --epochs 10 \
-    --save_every 2 \
-    --prefix_length 10 \
-    --bs 32 \
-    --mapping_type mlp \
-    --use_dino \
-    --dino_model_type dinov2_vits14 \
-    --only_prefix
-```
-
-### Advanced Training Options:
-
-#### Train both prefix and GPT (full model):
-```bash
-python clipcapTraining.py \
-    --data /raid/datasets/models_weights/clipcap/training-features/coco_karpathy_split_dinov2_vitb14_train.pkl \
-    --out_dir ./checkpoints_dino_vitb14_mlp_len10_full \
-    --prefix dino_vitb14_mlp_len10_full \
-    --epochs 10 \
-    --save_every 2 \
-    --prefix_length 10 \
-    --bs 16 \
-    --mapping_type mlp \
-    --use_dino \
-    --dino_model_type dinov2_vitb14
-```
-
-#### Use Transformer mapping instead of MLP:
-```bash
-python clipcapTraining.py \
-    --data /raid/datasets/models_weights/clipcap/training-features/coco_karpathy_split_dinov2_vitb14_train.pkl \
-    --out_dir ./checkpoints_dino_vitb14_transformer_len10 \
-    --prefix dino_vitb14_transformer_len10 \
-    --epochs 10 \
-    --save_every 2 \
-    --prefix_length 10 \
-    --bs 32 \
-    --mapping_type transformer \
-    --num_layers 8 \
-    --use_dino \
-    --dino_model_type dinov2_vitb14 \
-    --only_prefix
-```
-
-#### Custom feature dimension (if needed):
-```bash
-python clipcapTraining.py \
-    --data ./data/coco/coco_karpathy_split_dinov2_vitb14_train.pkl \
-    --out_dir ./checkpoints_dino_custom \
-    --prefix dino_custom \
-    --epochs 10 \
-    --prefix_length 10 \
-    --bs 32 \
-    --mapping_type mlp \
-    --use_dino \
-    --dino_model_type dinov2_vitb14 \
-    --dino_feature_dim 768 \
-    --only_prefix
-```
-
-## Key Parameters Explanation:
-
-- `--use_dino`: Enable DINO mode (required for DINO training)
-- `--dino_model_type`: Specify which DINO model was used for feature extraction
-- `--dino_feature_dim`: Override automatic feature dimension detection
-- `--prefix_length`: Number of prefix tokens (set to 10 as requested)
-- `--mapping_type`: Choose between 'mlp' or 'transformer' mapping
-- `--only_prefix`: Train only the mapping layer, freeze GPT-2
-- `--bs`: Batch size (adjust based on GPU memory)
-- `--epochs`: Number of training epochs
-- `--save_every`: Save checkpoint every N epochs
-
-## Expected Feature Dimensions:
-
-- **DINOv2-S/14**: 384 dimensions
-- **DINOv2-B/14**: 768 dimensions  
-- **DINOv2-L/14**: 1024 dimensions
-- **DINOv2-G/14**: 1536 dimensions
-
-## Training Tips:
-
-1. **Memory Usage**: DINO features are typically larger than CLIP features, so you might need to reduce batch size
-2. **Convergence**: DINO-based models may require different learning rates or longer training
-3. **Prefix Length**: Experiment with different prefix lengths (5, 10, 20) for optimal performance
-4. **Mapping Type**: MLP is faster, Transformer might give better results but requires more memory
-
-## Output:
-
-The training will save checkpoints in the specified output directory:
-- `{prefix}-{epoch:03d}.pt`: Model checkpoint for each epoch
-- `{prefix}_latest.pt`: Latest model checkpoint (updated every 10k iterations)
-- `{prefix}.json`: Training configuration
-
-## Example Full Workflow:
-
-```bash
-# 1. Extract DINO features
-python clipcap_dino_parse_coco.py --dino_model_type dinov2_vitb14
-
-# 2. Train ClipCap with DINO features (MLP, length 10, prefix-only)
-python clipcapTraining.py \
-    --data /raid/datasets/models_weights/clipcap/training-features/coco_karpathy_split_dinov2_vitb14_train.pkl \
-    --out_dir ./checkpoints_dino_vitb14_mlp_len10 \
-    --prefix dino_vitb14_mlp_len10 \
-    --epochs 10 \
-    --prefix_length 10 \
-    --bs 32 \
-    --mapping_type mlp \
-    --use_dino \
-    --dino_model_type dinov2_vitb14 \
-    --only_prefix
-```
-
-This will train a ClipCap model using DINO features with MLP mapping and sequence length 10 as requested.

src/clipcap/clipcapTraining.py

@@ -1,405 +0,0 @@
-import torch
-import torch.nn as nn
-from torch.nn import functional as nnf
-from torch.utils.data import Dataset, DataLoader
-from enum import Enum
-from transformers import GPT2Tokenizer, GPT2LMHeadModel, AdamW, get_linear_schedule_with_warmup
-from tqdm import tqdm
-import os
-import pickle
-import sys
-import argparse
-import json
-from typing import Tuple, Optional, Union
-
-
-class MappingType(Enum):
-    MLP = 'mlp'
-    Transformer = 'transformer'
-
-
-class ClipCocoDataset(Dataset):
-
-    def __len__(self) -> int:
-        return len(self.captions_tokens)
-
-    def pad_tokens(self, item: int):
-        tokens = self.captions_tokens[item]
-        padding = self.max_seq_len - tokens.shape[0]
-        if padding > 0:
-            tokens = torch.cat((tokens, torch.zeros(padding, dtype=torch.int64) - 1))
-            self.captions_tokens[item] = tokens
-        elif padding < 0:
-            tokens = tokens[:self.max_seq_len]
-            self.captions_tokens[item] = tokens
-        mask = tokens.ge(0)  # mask is zero where we out of sequence
-        tokens[~mask] = 0
-        mask = mask.float()
-        mask = torch.cat((torch.ones(self.prefix_length), mask), dim=0)  # adding prefix mask
-        return tokens, mask
-
-    def __getitem__(self, item: int) -> Tuple[torch.Tensor, ...]:
-        tokens, mask = self.pad_tokens(item)
-        prefix = self.prefixes[self.caption2embedding[item]]
-        if self.normalize_prefix:
-            prefix = prefix.float()
-            prefix = prefix / prefix.norm(2, -1)
-        return tokens, mask, prefix
-
-    def __init__(self, data_path: str,  prefix_length: int, gpt2_type: str = "gpt2",
-                 normalize_prefix=False):
-        self.tokenizer = GPT2Tokenizer.from_pretrained(gpt2_type)
-        self.prefix_length = prefix_length
-        self.normalize_prefix = normalize_prefix
-        with open(data_path, 'rb') as f:
-            all_data = pickle.load(f)
-        print("Data size is %0d" % len(all_data["clip_embedding"]))
-        sys.stdout.flush()
-        self.prefixes = all_data["clip_embedding"]
-        captions_raw = all_data["captions"]
-        self.image_ids = [caption["image_id"] for caption in captions_raw]
-        self.captions = [caption['caption'] for caption in captions_raw]
-        if os.path.isfile(f"{data_path[:-4]}_tokens.pkl"):
-            with open(f"{data_path[:-4]}_tokens.pkl", 'rb') as f:
-                self.captions_tokens, self.caption2embedding, self.max_seq_len = pickle.load(f)
-        else:
-            self.captions_tokens = []
-            self.caption2embedding = []
-            max_seq_len = 0
-            for caption in captions_raw:
-                self.captions_tokens.append(torch.tensor(self.tokenizer.encode(caption['caption']), dtype=torch.int64))
-                self.caption2embedding.append(caption["clip_embedding"])
-                max_seq_len = max(max_seq_len, self.captions_tokens[-1].shape[0])
-            # self.max_seq_len = max_seq_len
-            with open(f"{data_path[:-4]}_tokens.pkl", 'wb') as f:
-                pickle.dump([self.captions_tokens, self.caption2embedding, max_seq_len], f)
-        all_len = torch.tensor([len(self.captions_tokens[i]) for i in range(len(self))]).float()
-        self.max_seq_len = min(int(all_len.mean() + all_len.std() * 10), int(all_len.max()))
-
-
-class MLP(nn.Module):
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.model(x)
-
-    def __init__(self, sizes: Tuple[int, ...], bias=True, act=nn.Tanh):
-        super(MLP, self).__init__()
-        layers = []
-        for i in range(len(sizes) - 1):
-            layers.append(nn.Linear(sizes[i], sizes[i + 1], bias=bias))
-            if i < len(sizes) - 2:
-                layers.append(act())
-        self.model = nn.Sequential(*layers)
-
-
-class MlpTransformer(nn.Module):
-    def __init__(self, in_dim, h_dim, out_d: Optional[int] = None, act=nnf.relu, dropout=0.):
-        super().__init__()
-        out_d = out_d if out_d is not None else in_dim
-        self.fc1 = nn.Linear(in_dim, h_dim)
-        self.act = act
-        self.fc2 = nn.Linear(h_dim, out_d)
-        self.dropout = nn.Dropout(dropout)
-
-    def forward(self, x):
-        x = self.fc1(x)
-        x = self.act(x)
-        x = self.dropout(x)
-        x = self.fc2(x)
-        x = self.dropout(x)
-        return x
-
-class MultiHeadAttention(nn.Module):
-
-    def __init__(self, dim_self, dim_ref, num_heads, bias=True, dropout=0.):
-        super().__init__()
-        self.num_heads = num_heads
-        head_dim = dim_self // num_heads
-        self.scale = head_dim ** -0.5
-        self.to_queries = nn.Linear(dim_self, dim_self, bias=bias)
-        self.to_keys_values = nn.Linear(dim_ref, dim_self * 2, bias=bias)
-        self.project = nn.Linear(dim_self, dim_self)
-        self.dropout = nn.Dropout(dropout)
-
-    def forward(self, x, y=None, mask=None):
-        y = y if y is not None else x
-        b, n, c = x.shape
-        _, m, d = y.shape
-        # b n h dh
-        queries = self.to_queries(x).reshape(b, n, self.num_heads, c // self.num_heads)
-        # b m 2 h dh
-        keys_values = self.to_keys_values(y).reshape(b, m, 2, self.num_heads, c // self.num_heads)
-        keys, values = keys_values[:, :, 0], keys_values[:, :, 1]
-        attention = torch.einsum('bnhd,bmhd->bnmh', queries, keys) * self.scale
-        if mask is not None:
-            if mask.dim() == 2:
-                mask = mask.unsqueeze(1)
-            attention = attention.masked_fill(mask.unsqueeze(3), float("-inf"))
-        attention = attention.softmax(dim=2)
-        out = torch.einsum('bnmh,bmhd->bnhd', attention, values).reshape(b, n, c)
-        out = self.project(out)
-        return out, attention
-
-
-class TransformerLayer(nn.Module):
-
-    def forward_with_attention(self, x, y=None, mask=None):
-        x_, attention = self.attn(self.norm1(x), y, mask)
-        x = x + x_
-        x = x + self.mlp(self.norm2(x))
-        return x, attention
-
-    def forward(self, x, y=None, mask=None):
-        x = x + self.attn(self.norm1(x), y, mask)[0]
-        x = x + self.mlp(self.norm2(x))
-        return x
-
-    def __init__(self, dim_self, dim_ref, num_heads, mlp_ratio=4., bias=False, dropout=0., act=nnf.relu,
-                 norm_layer: nn.Module = nn.LayerNorm):
-        super().__init__()
-        self.norm1 = norm_layer(dim_self)
-        self.attn = MultiHeadAttention(dim_self, dim_ref, num_heads, bias=bias, dropout=dropout)
-        self.norm2 = norm_layer(dim_self)
-        self.mlp = MlpTransformer(dim_self, int(dim_self * mlp_ratio), act=act, dropout=dropout)
-
-
-class Transformer(nn.Module):
-
-    def forward_with_attention(self, x, y=None, mask=None):
-        attentions = []
-        for layer in self.layers:
-            x, att = layer.forward_with_attention(x, y, mask)
-            attentions.append(att)
-        return x, attentions
-
-    def forward(self, x, y=None, mask=None):
-        for i, layer in enumerate(self.layers):
-            if i % 2 == 0 and self.enc_dec: # cross
-                x = layer(x, y)
-            elif self.enc_dec:  # self
-                x = layer(x, x, mask)
-            else:  # self or cross
-                x = layer(x, y, mask)
-        return x
-
-    def __init__(self, dim_self: int, num_heads: int, num_layers: int, dim_ref: Optional[int] = None,
-                 mlp_ratio: float = 2., act=nnf.relu, norm_layer: nn.Module = nn.LayerNorm, enc_dec: bool = False):
-        super(Transformer, self).__init__()
-        dim_ref = dim_ref if dim_ref is not None else dim_self
-        self.enc_dec = enc_dec
-        if enc_dec:
-            num_layers = num_layers * 2
-        layers = []
-        for i in range(num_layers):
-            if i % 2 == 0 and enc_dec:  # cross
-                layers.append(TransformerLayer(dim_self, dim_ref, num_heads, mlp_ratio, act=act, norm_layer=norm_layer))
-            elif enc_dec:  # self
-                layers.append(TransformerLayer(dim_self, dim_self, num_heads, mlp_ratio, act=act, norm_layer=norm_layer))
-            else:  # self or cross
-                layers.append(TransformerLayer(dim_self, dim_ref, num_heads, mlp_ratio, act=act, norm_layer=norm_layer))
-        self.layers = nn.ModuleList(layers)
-
-
-class TransformerMapper(nn.Module):
-
-    def forward(self, x):
-        x = self.linear(x).view(x.shape[0], self.clip_length, -1)
-        prefix = self.prefix_const.unsqueeze(0).expand(x.shape[0], *self.prefix_const.shape)
-        prefix = torch.cat((x, prefix), dim=1)
-        out = self.transformer(prefix)[:, self.clip_length:]
-        return out
-
-    def __init__(self, dim_clip: int, dim_embedding: int, prefix_length: int, clip_length: int, num_layers: int = 8):
-        super(TransformerMapper, self).__init__()
-        self.clip_length = clip_length
-        self.transformer = Transformer(dim_embedding, 8, num_layers)
-        self.linear = nn.Linear(dim_clip, clip_length * dim_embedding)
-        self.prefix_const = nn.Parameter(torch.randn(prefix_length, dim_embedding), requires_grad=True)
-
-
-class ClipCaptionModel(nn.Module):
-
-    def get_dummy_token(self, batch_size: int, device: torch.device) -> torch.Tensor:
-        return torch.zeros(batch_size, self.prefix_length, dtype=torch.int64, device=device)
-
-    def forward(self, tokens: torch.Tensor, prefix: torch.Tensor, mask: Optional[torch.Tensor] = None,
-                labels: Optional[torch.Tensor] = None):
-        embedding_text = self.gpt.transformer.wte(tokens)
-        prefix_projections = self.clip_project(prefix).view(-1, self.prefix_length, self.gpt_embedding_size)
-        embedding_cat = torch.cat((prefix_projections, embedding_text), dim=1)
-        if labels is not None:
-            dummy_token = self.get_dummy_token(tokens.shape[0], tokens.device)
-            labels = torch.cat((dummy_token, tokens), dim=1)
-        out = self.gpt(inputs_embeds=embedding_cat, labels=labels, attention_mask=mask)
-        return out
-
-    def __init__(self, prefix_length: int, clip_length: Optional[int] = None, prefix_size: int = 512,
-                 num_layers: int = 8, mapping_type: MappingType = MappingType.MLP):
-        super(ClipCaptionModel, self).__init__()
-        self.prefix_length = prefix_length
-        self.gpt = GPT2LMHeadModel.from_pretrained('gpt2')
-        self.gpt_embedding_size = self.gpt.transformer.wte.weight.shape[1]
-        if mapping_type == MappingType.MLP:
-            self.clip_project = MLP((prefix_size, (self.gpt_embedding_size * prefix_length) // 2,
-                                     self.gpt_embedding_size * prefix_length))
-        else:
-            self.clip_project = TransformerMapper(prefix_size, self.gpt_embedding_size, prefix_length,
-                                                                     clip_length, num_layers)
-
-
-class ClipCaptionPrefix(ClipCaptionModel):
-
-    def parameters(self, recurse: bool = True):
-        return self.clip_project.parameters()
-
-    def train(self, mode: bool = True):
-        super(ClipCaptionPrefix, self).train(mode)
-        self.gpt.eval()
-        return self
-
-
-def save_config(args: argparse.Namespace):
-    config = {}
-    for key, item in args._get_kwargs():
-        config[key] = item
-    out_path = os.path.join(args.out_dir, f"{args.prefix}.json")
-    with open(out_path, 'w') as outfile:
-        json.dump(config, outfile)
-
-
-def load_model(config_path: str, epoch_or_latest: Union[str, int] = '_latest'):
-    with open(config_path) as f:
-        config = json.load(f)
-    parser = argparse.ArgumentParser()
-    parser.set_defaults(**config)
-    args = parser.parse_args()
-    if type(epoch_or_latest) is int:
-        epoch_or_latest = f"-{epoch_or_latest:03d}"
-    model_path = os.path.join(args.out_dir, f"{args.prefix}{epoch_or_latest}.pt")
-    if args.only_prefix:
-        model = ClipCaptionPrefix(args.prefix_length)
-    else:
-        model = ClipCaptionModel(args.prefix_length)
-    if os.path.isfile(model_path):
-        print(f"loading model from {model_path}")
-        model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
-    else:
-        print(f"{model_path} is not exist")
-    return model, parser
-
-
-def train(dataset: ClipCocoDataset, model: ClipCaptionModel, args,
-          lr: float = 2e-5, warmup_steps: int = 5000, output_dir: str = ".", output_prefix: str = "", device = torch.device('cuda:0')):
-
-    batch_size = args.bs
-    epochs = args.epochs
-    if not os.path.exists(output_dir):
-        os.makedirs(output_dir)
-    model = model.to(device)
-    model.train()
-    optimizer = AdamW(model.parameters(), lr=lr)
-    train_dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, drop_last=True)
-    scheduler = get_linear_schedule_with_warmup(
-        optimizer, num_warmup_steps=warmup_steps, num_training_steps=epochs * len(train_dataloader)
-    )
-    # save_config(args)
-    for epoch in range(epochs):
-        print(f">>> Training epoch {epoch}")
-        sys.stdout.flush()
-        progress = tqdm(total=len(train_dataloader), desc=output_prefix)
-        for idx, (tokens, mask, prefix) in enumerate(train_dataloader):
-            model.zero_grad()
-            tokens, mask, prefix = tokens.to(device), mask.to(device), prefix.to(device, dtype=torch.float32)
-            outputs = model(tokens, prefix, mask)
-            logits = outputs.logits[:, dataset.prefix_length - 1: -1]
-            loss = nnf.cross_entropy(logits.reshape(-1, logits.shape[-1]), tokens.flatten(), ignore_index=0)
-            loss.backward()
-            optimizer.step()
-            scheduler.step()
-            optimizer.zero_grad()
-            progress.set_postfix({"loss": loss.item()})
-            progress.update()
-            if (idx + 1) % 10000 == 0:
-                torch.save(
-                    model.state_dict(),
-                    os.path.join(output_dir, f"{output_prefix}_latest.pt"),
-                )
-        progress.close()
-        if epoch % args.save_every == 0 or epoch == epochs - 1:
-            torch.save(
-                model.state_dict(),
-                os.path.join(output_dir, f"{output_prefix}-{epoch:03d}.pt"),
-            )
-    return model
-
-
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--data', default='/raid/datasets/models_weights/clipcap/training-features/coco_karpathy_split_train.pkl')
-    parser.add_argument('--out_dir', default='/raid/datasets/models_weights/clipcap/checkpoints/dinov2b14/')
-    parser.add_argument('--prefix', default='coco_prefix', help='prefix for saved filenames')
-    parser.add_argument('--epochs', type=int, default=10)
-    parser.add_argument('--save_every', type=int, default=1)
-    parser.add_argument('--prefix_length', type=int, default=10)
-    parser.add_argument('--prefix_length_clip', type=int, default=10)
-    parser.add_argument('--bs', type=int, default=40)
-    parser.add_argument('--only_prefix', dest='only_prefix', action='store_true')
-    parser.add_argument('--mapping_type', type=str, default='mlp', help='mlp/transformer')
-    parser.add_argument('--num_layers', type=int, default=8)
-    parser.add_argument('--is_rn', dest='is_rn', action='store_true')
-    parser.add_argument('--normalize_prefix', dest='normalize_prefix', action='store_true')
-    # DINO-specific arguments
-    parser.add_argument('--use_dino', action='store_true', default=False, help='Use DINO features instead of CLIP')
-    parser.add_argument('--dino_model_type', type=str, default='dinov2_vitb14', 
-                       choices=['dinov2_vits14', 'dinov2_vitb14', 'dinov2_vitl14', 'dinov2_vitg14'], 
-                       help='DINO model type')
-    parser.add_argument('--dino_feature_dim', type=int, default=None, 
-                       help='DINO feature dimension (auto-detected if None)')
-    parser.add_argument('--device', type=str, default='cuda:0', help='Device to use for training')
-    args = parser.parse_args()
-
-    if isinstance(args.device, str):
-        if not args.device.startswith('cuda') and not args.device.startswith('cpu'):
-            # if it is an integer index, convert to f'cuda:{args.device}'
-            if args.device.isdigit():
-                args.device = f'cuda:{args.device}'
-            else:
-                raise ValueError(f"Invalid device string: {args.device}")
-        args.device = torch.device(args.device)
-
-    prefix_length = args.prefix_length
-    dataset = ClipCocoDataset(args.data, prefix_length, normalize_prefix=args.normalize_prefix)
-    
-    # Determine prefix dimension based on model type
-    if args.use_dino:
-        if args.dino_feature_dim is not None:
-            prefix_dim = args.dino_feature_dim
-        else:
-            # Auto-detect DINO feature dimensions
-            dino_dims = {
-                'dinov2_vits14': 384,
-                'dinov2_vitb14': 768,
-                'dinov2_vitl14': 1024,
-                'dinov2_vitg14': 1536
-            }
-            prefix_dim = dino_dims.get(args.dino_model_type, 768)
-        print(f"Using DINO features with dimension: {prefix_dim}")
-    else:
-        prefix_dim = 640 if args.is_rn else 512
-        print(f"Using CLIP features with dimension: {prefix_dim}")
-    
-    args.mapping_type = {'mlp': MappingType.MLP, 'transformer': MappingType.Transformer}[args.mapping_type]
-    if args.only_prefix:
-        model = ClipCaptionPrefix(prefix_length, clip_length=args.prefix_length_clip, prefix_size=prefix_dim,
-                                  num_layers=args.num_layers, mapping_type=args.mapping_type)
-        print("Train only prefix")
-    else:
-        model = ClipCaptionModel(prefix_length, clip_length=args.prefix_length_clip, prefix_size=prefix_dim,
-                                  num_layers=args.num_layers, mapping_type=args.mapping_type)
-        print("Train both prefix and GPT")
-        sys.stdout.flush()
-    train(dataset, model, args, output_dir=args.out_dir, output_prefix=args.prefix, device=args.device)
-
-
-if __name__ == '__main__':
-    main()

src/clipcap/clipcap_dino_parse_coco.py

@@ -1,613 +0,0 @@
-import torch
-import torch.nn.functional as F
-import skimage.io as io
-from PIL import Image
-import pickle
-import json
-import os
-from tqdm import tqdm
-import argparse
-import torchvision.transforms as T
-import numpy as np
-import yaml
-import clip
-import sys
-
-# Add the src directory to the path so we can import ProjectionLayer
-sys.path.append(os.path.join(os.path.dirname(__file__), '../..', 'src'))
-
-
-# Container to store intermediate outputs for feature extraction
-feats = {}
-
-def get_self_attention(module, input, output):
-    """Hook to capture self-attention weights"""
-    global qkv_attention_out
-    qkv_attention_out = output
-
-def get_layer_n_output(module, input, output):
-    """Hook to capture intermediate layer output"""
-    feats['intermediate_output'] = output
-
-def transform_to_standard_dino_out(x, model, num_global_tokens=1):
-    """Transform raw DINO output to standardized format"""
-    x_norm = model.norm(x)
-    if num_global_tokens == 1:
-        # Standard model without registers
-        return {
-            "x_norm_clstoken": x_norm[:, 0],
-            "x_norm_regtokens": None,
-            "x_norm_patchtokens": x_norm[:, 1:],
-            "x_prenorm": x,
-        }
-    else:
-        # Model with registers (num_global_tokens = 5)
-        return {
-            "x_norm_clstoken": x_norm[:, 0],
-            "x_norm_regtokens": x_norm[:, 1:num_global_tokens],
-            "x_norm_patchtokens": x_norm[:, num_global_tokens:],
-            "x_prenorm": x,
-        }
-
-def process_self_attention(output, batch_size, num_tokens, num_attn_heads, embed_dim, scale, num_global_tokens, ret_self_attn_maps=False):
-    """Process self-attention output to compute attention weights"""
-    qkv = output.reshape(batch_size, num_tokens, 3, num_attn_heads, embed_dim // num_attn_heads).permute(2, 0, 3, 1, 4)
-    q, k, v = qkv[0] * scale, qkv[1], qkv[2]
-    attn = q @ k.transpose(-2, -1)
-    self_attn_maps = attn[:, :, 0, num_global_tokens:]  # CLS token attention to patches
-    self_attn = self_attn_maps.mean(dim=1)  # Average over attention heads
-    self_attn = self_attn.softmax(dim=-1)
-    if ret_self_attn_maps:
-        return self_attn, self_attn_maps
-    else:
-        return self_attn
-
-
-# Global variables to store hook outputs
-dino_layer_n_output = None
-qkv_attention_out = None
-
-def get_layer_n_output(module, input, output):
-    """Hook to capture intermediate layer output"""
-    global dino_layer_n_output
-    dino_layer_n_output = output
-
-
-def select_most_significant_patch(dino_outs, self_attn, criteria, cls_token=None, caption_embedding=None):
-    """
-    Select the most significant patch token based on different criteria.
-    
-    Args:
-        dino_outs: Dictionary containing normalized DINO outputs
-        self_attn: Self-attention weights from CLS to patches [batch_size, num_patches]
-        criteria: Selection criteria ('max_attention', 'most_similar_to_cls', etc.)
-        cls_token: CLS token embeddings [batch_size, embed_dim]
-        caption_embedding: Text caption embeddings [batch_size, embed_dim]
-        
-    Returns:
-        selected_patches: [batch_size, embed_dim] - Selected patch embeddings
-    """
-    patch_tokens = dino_outs['x_norm_patchtokens']  # [batch_size, num_patches, embed_dim]
-    batch_size, num_patches, embed_dim = patch_tokens.shape
-    
-    if criteria == "max_attention":
-        # Select patch with highest attention weight from CLS token
-        if self_attn is None:
-            raise ValueError("self_attn required for max_attention criteria")
-        max_attn_indices = self_attn.argmax(dim=1)  # [batch_size]
-        selected_patches = patch_tokens[torch.arange(batch_size), max_attn_indices]
-        
-    elif criteria == "most_similar_to_cls":
-        # Select patch most similar to CLS token using cosine similarity
-        if cls_token is None:
-            raise ValueError("cls_token required for most_similar_to_cls criteria")
-        # Compute cosine similarity between CLS and all patches
-        cls_normalized = F.normalize(cls_token, p=2, dim=1)  # [batch_size, embed_dim]
-        patches_normalized = F.normalize(patch_tokens, p=2, dim=2)  # [batch_size, num_patches, embed_dim]
-        similarities = torch.bmm(patches_normalized, cls_normalized.unsqueeze(2)).squeeze(2)  # [batch_size, num_patches]
-        max_sim_indices = similarities.argmax(dim=1)  # [batch_size]
-        selected_patches = patch_tokens[torch.arange(batch_size), max_sim_indices]
-        
-    elif criteria == "most_similar_to_caption":
-        # Select patch most similar to caption embedding
-        if caption_embedding is None:
-            raise ValueError("caption_embedding required for most_similar_to_caption criteria")
-        caption_normalized = F.normalize(caption_embedding, p=2, dim=1)  # [batch_size, embed_dim]
-        patches_normalized = F.normalize(patch_tokens, p=2, dim=2)  # [batch_size, num_patches, embed_dim]
-        similarities = torch.bmm(patches_normalized, caption_normalized.unsqueeze(2)).squeeze(2)  # [batch_size, num_patches]
-        max_sim_indices = similarities.argmax(dim=1)  # [batch_size]
-        selected_patches = patch_tokens[torch.arange(batch_size), max_sim_indices]
-        
-    elif criteria == "max_norm":
-        # Select patch with highest L2 norm
-        patch_norms = torch.norm(patch_tokens, p=2, dim=2)  # [batch_size, num_patches]
-        max_norm_indices = patch_norms.argmax(dim=1)  # [batch_size]
-        selected_patches = patch_tokens[torch.arange(batch_size), max_norm_indices]
-        
-    elif criteria == "centroid_distance":
-        # Select patch farthest from the centroid of all patches
-        centroid = patch_tokens.mean(dim=1, keepdim=True)  # [batch_size, 1, embed_dim]
-        distances = torch.norm(patch_tokens - centroid, p=2, dim=2)  # [batch_size, num_patches]
-        max_dist_indices = distances.argmax(dim=1)  # [batch_size]
-        selected_patches = patch_tokens[torch.arange(batch_size), max_dist_indices]
-        
-    else:
-        raise ValueError(f"Unknown patch selection criteria: {criteria}")
-    
-    return selected_patches
-
-
-def load_text_encoder(text_encoder_path, device, config_path=None):
-    """
-    Load a text encoder model for caption similarity.
-    Supports Talk2Dino, CLIP, and DINO.txt-based text encoders.
-    """
-    if text_encoder_path is None:
-        return None
-    
-    print(f"Loading text encoder from: {text_encoder_path}")
-    
-    # Check for DINO.txt model
-    if text_encoder_path.lower() == 'dinotxt' or text_encoder_path.lower() == 'dino.txt':
-        # Load DINO.txt model
-        try:
-            from src.dinotxt_utils import get_tokenizer
-            
-            print("Loading DINO.txt model...")
-            dinotxt_model = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitl14_reg4_dinotxt_tet1280d20h24l')
-            dinotxt_model.eval()
-            dinotxt_model.to(device)
-            
-            tokenizer = get_tokenizer()
-            
-            return {
-                'type': 'dinotxt',
-                'model': dinotxt_model,
-                'tokenizer': tokenizer
-            }
-            
-        except ImportError:
-            raise ImportError("Could not import dinotxt_utils. Make sure src/dinotxt_utils.py is accessible.")
-        except Exception as e:
-            raise RuntimeError(f"Failed to load DINO.txt model: {e}")
-    
-    # Check if it's a Talk2Dino model (expect config and weights)
-    elif text_encoder_path.endswith('.pth') or text_encoder_path.endswith('.pt'):
-        # Use provided config or auto-find
-        if config_path is None:
-            # Look for corresponding config file
-            base_path = text_encoder_path.rsplit('.', 1)[0]
-            config_path = base_path + '.yaml'
-            
-            # Alternative config path patterns
-            if not os.path.exists(config_path):
-                # Try configs_talk2dino directory
-                config_name = os.path.basename(base_path) + '.yaml'
-                config_path = os.path.join(os.path.dirname(__file__), 'configs_talk2dino', config_name)
-                
-        if not os.path.exists(config_path):
-            raise FileNotFoundError(f"Could not find config file for {text_encoder_path}. "
-                                   f"Expected at {config_path} or specify --text_encoder_config.")
-        
-        # Load Talk2Dino model
-        try:
-            from src.model import ProjectionLayer
-            
-            print(f"Using config: {config_path}")
-            
-            # Load the projection layer
-            talk2dino = ProjectionLayer.from_config(config_path)
-            talk2dino.load_state_dict(torch.load(text_encoder_path, map_location=device))
-            talk2dino.to(device)
-            talk2dino.eval()
-            
-            # Load CLIP model for text encoding
-            clip_model, _ = clip.load("ViT-B/32", device=device)
-            clip_model.eval()
-            
-            return {
-                'type': 'talk2dino',
-                'talk2dino': talk2dino,
-                'clip_model': clip_model,
-                'config_path': config_path
-            }
-            
-        except ImportError:
-            raise ImportError("Could not import ProjectionLayer. Make sure src/model.py is accessible.")
-    
-    else:
-        # Assume it's a direct model path (CLIP or other)
-        try:
-            # Try loading as a CLIP model
-            clip_model, _ = clip.load(text_encoder_path, device=device)
-            clip_model.eval()
-            
-            return {
-                'type': 'clip',
-                'clip_model': clip_model
-            }
-        except:
-            raise ValueError(f"Could not load text encoder from {text_encoder_path}. "
-                           f"Supported formats: 1) 'dinotxt' or 'dino.txt' for DINO.txt model, "
-                           f"2) Talk2Dino (.pth/.pt), 3) CLIP model names.")
-
-
-def encode_caption(caption, text_encoder, device):
-    """
-    Encode a text caption using the loaded text encoder.
-    """
-    if text_encoder is None:
-        return None
-    
-    if text_encoder['type'] == 'dinotxt':
-        # Use DINO.txt pipeline: tokenize + encode + extract patch-aligned features
-        with torch.no_grad():
-            # Tokenize with DINO.txt tokenizer
-            text_tokens = text_encoder['tokenizer'].tokenize([caption]).to(device)
-            
-            # Encode with DINO.txt model
-            dinotxt_features = text_encoder['model'].encode_text(text_tokens)
-            
-            # Extract patch-aligned text embeddings (dimensions 1024:)
-            # DINO.txt concatenates standard text features [0:1024] and patch-aligned features [1024:]
-            patch_aligned_features = dinotxt_features[:, 1024:]
-            
-            # Normalize the features to match DINO feature space
-            patch_aligned_features = F.normalize(patch_aligned_features, p=2, dim=-1)
-            return patch_aligned_features
-            
-    elif text_encoder['type'] == 'talk2dino':
-        # Use Talk2Dino pipeline: CLIP text encoding + Talk2Dino projection
-        with torch.no_grad():
-            # Tokenize and encode with CLIP
-            text_tokens = clip.tokenize([caption]).to(device)
-            clip_text_features = text_encoder['clip_model'].encode_text(text_tokens)
-            
-            # Project through Talk2Dino to DINO space
-            dino_text_features = text_encoder['talk2dino'].project_clip_txt(clip_text_features)
-            
-            # Normalize the encoded text to match DINO feature space
-            dino_text_features = F.normalize(dino_text_features, p=2, dim=-1)
-            return dino_text_features
-            
-    elif text_encoder['type'] == 'clip':
-        # Use CLIP directly
-        with torch.no_grad():
-            text_tokens = clip.tokenize([caption]).to(device)
-            clip_text_features = text_encoder['clip_model'].encode_text(text_tokens)
-            
-            # Normalize the features
-            clip_text_features = F.normalize(clip_text_features, p=2, dim=-1)
-            return clip_text_features
-    
-    else:
-        raise ValueError(f"Unknown text encoder type: {text_encoder['type']}")
-
-
-def main(dino_model_type: str, resize_dim: int = 518, crop_dim: int = 518, 
-         coco_images_dir: str = "/raid/datasets/coco/", captions_file: str = "/raid/datasets/coco/train_split_karpathy.json",
-         output_file: str = None, feature_type: str = "cls", extract_attention: bool = False,
-         patch_selection_criteria: str = "max_attention", text_encoder_path: str = None, text_encoder_config: str = None):
-    """
-    Extract DINO features from COCO images for ClipCap training.
-    
-    Args:
-        feature_type: Type of features to extract
-            - "cls": CLS token features (default)
-            - "avg_patch": Mean pooled patch token features
-            - "avg_self_attn": Self-attention weighted patch token features
-            - "most_significant_patch": Single most important patch token
-        extract_attention: Whether to extract self-attention weights (required for avg_self_attn)
-        patch_selection_criteria: Criteria for selecting most significant patch
-        text_encoder_path: Path to text encoder for caption similarity
-    """
-    device = torch.device('cuda:0')
-    dino_model_name = dino_model_type.replace('/', '_')
-    
-    # Determine model properties
-    num_global_tokens = 1 if "reg" not in dino_model_type else 5
-    patch_size = 14  # DINOv2 uses 14x14 patches
-    num_patch_tokens = (crop_dim // patch_size) * (crop_dim // patch_size)
-    num_tokens = num_global_tokens + num_patch_tokens
-    
-    # Get embedding dimension based on model type
-    if 'vitl' in dino_model_type:
-        embed_dim = 1024
-        num_attn_heads = 16
-    elif 'vitb' in dino_model_type:
-        embed_dim = 768
-        num_attn_heads = 12
-    elif 'vits' in dino_model_type:
-        embed_dim = 384
-        num_attn_heads = 6
-    elif 'vitg' in dino_model_type:
-        embed_dim = 1536
-        num_attn_heads = 24
-    else:
-        raise ValueError(f"Unknown model type: {dino_model_type}")
-    
-    scale = (embed_dim // num_attn_heads) ** -0.5
-    
-    # Set default output path if not specified
-    if output_file is None:
-        if feature_type == "cls":
-            feature_suffix = ""
-        elif feature_type == "most_significant_patch":
-            if patch_selection_criteria == "most_similar_to_caption" and text_encoder_path is not None:
-                # Determine text encoder type from path to create unique filename
-                if text_encoder_path.lower() in ['dinotxt', 'dino.txt']:
-                    text_encoder_suffix = "_dinotxt"
-                elif text_encoder_path.endswith('.pth') or text_encoder_path.endswith('.pt'):
-                    text_encoder_suffix = "_t2d"  # Talk2Dino
-                else:
-                    # CLIP or other models
-                    text_encoder_suffix = "_clip"
-                feature_suffix = f"_{feature_type}_{patch_selection_criteria}{text_encoder_suffix}"
-            else:
-                feature_suffix = f"_{feature_type}_{patch_selection_criteria}"
-        else:
-            feature_suffix = f"_{feature_type}"
-        output_file = f"/raid/datasets/models_weights/clipcap/training-features/coco_karpathy_split_{dino_model_name}{feature_suffix}_train.pkl"
-    
-    # Create output directory if it doesn't exist
-    os.makedirs(os.path.dirname(output_file), exist_ok=True)
-    
-    # Load DINO model
-    print(f"Loading DINO model: {dino_model_type}")
-    print(f"Feature type: {feature_type}")
-    if feature_type == "most_significant_patch":
-        print(f"Patch selection criteria: {patch_selection_criteria}")
-    print(f"Model properties: embed_dim={embed_dim}, num_heads={num_attn_heads}, num_global_tokens={num_global_tokens}")
-    
-    if 'dinov2' in dino_model_type:
-        model_family = 'facebookresearch/dinov2'
-        dino_model = torch.hub.load(model_family, dino_model_type)
-    else:
-        raise ValueError(f"Unsupported DINO model type: {dino_model_type}")
-    
-    # Setup transforms for DINO
-    image_transforms = T.Compose([
-        T.Resize(resize_dim, interpolation=T.InterpolationMode.BICUBIC),
-        T.CenterCrop(crop_dim),
-        T.ToTensor(),
-        T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
-    ])
-    
-    dino_model.eval()
-    dino_model.to(device)
-    
-    # Register hooks if we need attention or intermediate outputs
-    if feature_type == "avg_self_attn" or extract_attention or \
-       (feature_type == "most_significant_patch" and patch_selection_criteria in ["max_attention", "most_similar_to_caption"]):
-        print("Registering hooks for attention extraction...")
-        dino_model.blocks[-1].attn.qkv.register_forward_hook(get_self_attention)
-    
-    if feature_type in ["avg_patch", "avg_self_attn", "most_significant_patch"]:
-        print("Registering hooks for intermediate output extraction...")
-        dino_model.blocks[-1].register_forward_hook(get_layer_n_output)
-    
-    # Load caption data
-    print(f"Loading captions from: {captions_file}")
-    with open(captions_file, 'r') as f:
-        data = json.load(f)
-    
-    # Handle different annotation formats
-    if isinstance(data, list):
-        # Original ClipCap format: list of dicts with 'image_id' and 'caption'
-        annotations = data
-        print(f"{len(annotations)} captions loaded from json (ClipCap format)")
-    elif isinstance(data, dict) and 'annotations' in data:
-        # Karpathy format: dict with 'annotations' key
-        annotations = data['annotations']
-        print(f"{len(annotations)} captions loaded from json (Karpathy format)")
-        
-        # Create image ID to filename mapping for faster lookup
-        if 'images' in data:
-            image_id_to_filename = {img['id']: img['file_name'] for img in data['images']}
-        else:
-            image_id_to_filename = {}
-    else:
-        raise ValueError("Unsupported annotation format")
-    
-    # Load text encoder if needed for caption similarity
-    text_encoder = None
-    if feature_type == "most_significant_patch" and patch_selection_criteria == "most_similar_to_caption":
-        if text_encoder_path is None:
-            raise ValueError("text_encoder_path required for most_similar_to_caption criteria")
-        text_encoder = load_text_encoder(text_encoder_path, device, text_encoder_config)
-        print(f"Loaded text encoder from: {text_encoder_path}")
-    
-    all_embeddings = []
-    all_captions = []
-    
-    print(f"Processing images from: {coco_images_dir}")
-    print(f"Output will be saved to: {output_file}")
-    
-    for i, annotation in enumerate(tqdm(annotations)):
-        img_id = annotation["image_id"]
-        
-        # Determine filename based on format
-        if isinstance(data, list):
-            # Original format: construct filename from image_id
-            filename = os.path.join(coco_images_dir, "train2014", f"COCO_train2014_{int(img_id):012d}.jpg")
-            if not os.path.isfile(filename):
-                filename = os.path.join(coco_images_dir, "val2014", f"COCO_val2014_{int(img_id):012d}.jpg")
-        else:
-            # Karpathy format: use filename from images mapping or construct it
-            if img_id in image_id_to_filename:
-                if 'train' in image_id_to_filename[img_id]:
-                    fold = "train2014"
-                else: 
-                    fold = "val2014"
-                filename = os.path.join(coco_images_dir, fold, image_id_to_filename[img_id])
-            else:
-                # Fallback: try to construct filename
-                filename = os.path.join(coco_images_dir, "train2014", f"COCO_train2014_{int(img_id):012d}.jpg")
-                if not os.path.isfile(filename):
-                    filename = os.path.join(coco_images_dir, "val2014", f"COCO_val2014_{int(img_id):012d}.jpg")
-        
-        if not os.path.isfile(filename):
-            print(f"Warning: Image not found: {filename}")
-            continue
-        
-        # Load and process image
-        try:
-            image = io.imread(filename)
-            if len(image.shape) == 2:  # grayscale
-                image = Image.fromarray(image).convert('RGB')
-            else:
-                image = Image.fromarray(image)
-        except Exception as e:
-            print(f"Warning: Failed to load image {filename}: {e}")
-            continue
-        
-        # Apply DINO transforms
-        image_tensor = image_transforms(image).unsqueeze(0).to(device)
-        
-        with torch.no_grad():
-            # Clear any previous stored data
-            global dino_layer_n_output, qkv_attention_out
-            dino_layer_n_output = None
-            qkv_attention_out = None
-            
-            # Extract DINO features
-            if feature_type == "cls":
-                # Standard CLS token extraction
-                features = dino_model(image_tensor)
-                # For DINOv2, the output is the CLS token by default
-                if len(features.shape) == 3:  # If we get [batch, seq_len, dim]
-                    features = features[:, 0, :]  # Take CLS token
-                prefix = features.cpu()
-            else:
-                # For patch-based features, we need intermediate outputs
-                _ = dino_model(image_tensor)  # Forward pass to trigger hooks
-                
-                if dino_layer_n_output is None:
-                    raise RuntimeError("No intermediate output captured. Check hook registration.")
-                
-                # Transform to standard format
-                dino_outs = transform_to_standard_dino_out(dino_layer_n_output, dino_model, num_global_tokens)
-                
-                if feature_type == "avg_patch":
-                    # Average of patch tokens (excluding global tokens)
-                    prefix = dino_outs['x_norm_patchtokens'].mean(dim=1)  # [B, D]
-                elif feature_type == "avg_self_attn":
-                    # Self-attention weighted average of patch tokens
-                    if qkv_attention_out is None:
-                        raise RuntimeError("No attention output captured. Check hook registration.")
-                    
-                    # Process self-attention to get attention weights
-                    batch_size = qkv_attention_out.shape[0]
-                    self_attn = process_self_attention(
-                        qkv_attention_out, 
-                        batch_size,
-                        num_tokens, 
-                        num_attn_heads, 
-                        embed_dim,
-                        scale, 
-                        num_global_tokens
-                    )
-                    
-                    # Compute attention-weighted average
-                    prefix = (self_attn.unsqueeze(-1) * dino_outs['x_norm_patchtokens']).mean(dim=1)
-                elif feature_type == "most_significant_patch":
-                    # Select single most significant patch based on criteria
-                    self_attn = None
-                    cls_token = None
-                    caption_embedding = None
-                    
-                    # Prepare required inputs based on criteria
-                    if patch_selection_criteria in ["max_attention", "most_similar_to_caption"]:
-                        if qkv_attention_out is None:
-                            raise RuntimeError("No attention output captured. Check hook registration.")
-                        batch_size = qkv_attention_out.shape[0]
-                        self_attn = process_self_attention(
-                            qkv_attention_out, 
-                            batch_size,
-                            num_tokens, 
-                            num_attn_heads, 
-                            embed_dim,
-                            scale, 
-                            num_global_tokens
-                        )
-                    
-                    if patch_selection_criteria == "most_similar_to_cls":
-                        cls_token = dino_outs['x_norm_clstoken']
-                    
-                    if patch_selection_criteria == "most_similar_to_caption":
-                        if text_encoder is not None:
-                            caption_embedding = encode_caption(annotation["caption"], text_encoder, device)
-                    
-                    # Select the most significant patch
-                    prefix = select_most_significant_patch(
-                        dino_outs, 
-                        self_attn, 
-                        patch_selection_criteria,
-                        cls_token=cls_token,
-                        caption_embedding=caption_embedding
-                    )
-                else:
-                    raise ValueError(f"Unknown feature type: {feature_type}")
-                
-                prefix = prefix.cpu()
-        
-        # Create annotation in ClipCap format for compatibility
-        caption_entry = {
-            "image_id": img_id,
-            "caption": annotation["caption"],
-            "clip_embedding": i  # Index for the embedding
-        }
-        
-        all_embeddings.append(prefix)
-        all_captions.append(caption_entry)
-        
-        if (i + 1) % 10000 == 0:
-            # Create output directory if it doesn't exist
-            os.makedirs(os.path.dirname(output_file), exist_ok=True)
-            with open(output_file, 'wb') as f:
-                pickle.dump({"clip_embedding": torch.cat(all_embeddings, dim=0), "captions": all_captions}, f)
-
-    # Create output directory if it doesn't exist
-    os.makedirs(os.path.dirname(output_file), exist_ok=True)
-    with open(output_file, 'wb') as f:
-        pickle.dump({"clip_embedding": torch.cat(all_embeddings, dim=0), "captions": all_captions}, f)
-
-    print('Done')
-    print("%0d embeddings saved " % len(all_embeddings))
-    print(f"Feature dimension: {all_embeddings[0].shape[-1]}")
-    return 0
-
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description='Extract DINO features from COCO images for ClipCap training')
-    parser.add_argument('--dino_model_type', default="dinov2_vitb14", 
-                       choices=('dinov2_vits14', 'dinov2_vitb14', 'dinov2_vitl14', 'dinov2_vitg14',
-                               'dinov2_vits14_reg', 'dinov2_vitb14_reg', 'dinov2_vitl14_reg', 'dinov2_vitg14_reg'),
-                       help='DINO model type to use for feature extraction')
-    parser.add_argument('--feature_type', default="cls", 
-                       choices=('cls', 'avg_patch', 'avg_self_attn', 'most_significant_patch'),
-                       help='Type of features to extract: cls (CLS token), avg_patch (mean pooled patches), avg_self_attn (attention-weighted patches), most_significant_patch (single most important patch)')
-    parser.add_argument('--patch_selection_criteria', default="max_attention",
-                       choices=('max_attention', 'most_similar_to_cls', 'most_similar_to_caption', 'max_norm', 'centroid_distance'),
-                       help='Criteria for selecting the most significant patch (only used with most_significant_patch feature_type)')
-    parser.add_argument('--text_encoder_path', type=str, default=None,
-                       help='Path to text encoder for caption similarity. Supports: 1) "dinotxt" or "dino.txt" for DINO.txt model, 2) Talk2Dino weights (.pth/.pt) - will auto-find config, 3) CLIP model names (e.g., "ViT-B/32")')
-    parser.add_argument('--text_encoder_config', type=str, default=None,
-                       help='Optional: explicit config path for Talk2Dino models (if not auto-found)')
-    parser.add_argument('--resize_dim', type=int, default=518, help='Resize dimension for images')
-    parser.add_argument('--crop_dim', type=int, default=518, help='Crop dimension for images')
-    parser.add_argument('--coco_images_dir', type=str, default="/raid/datasets/coco", 
-                       help='Path to COCO images directory (should contain train2014/ and val2014/ subdirs)')
-    parser.add_argument('--captions_file', type=str, default="/raid/datasets/coco/train_split_karpathy.json",
-                       help='Path to COCO captions JSON file (supports both Karpathy and ClipCap formats)')
-    parser.add_argument('--output_file', type=str, default=None,
-                       help='Output pickle file path (default: auto-generated based on model and feature type)')
-    parser.add_argument('--extract_attention', action='store_true',
-                       help='Extract attention weights (automatically enabled for avg_self_attn feature type)')
-    
-    args = parser.parse_args()
-    
-    main(args.dino_model_type, args.resize_dim, args.crop_dim, 
-         args.coco_images_dir, args.captions_file, args.output_file, 
-         args.feature_type, args.extract_attention,
-         args.patch_selection_criteria, args.text_encoder_path, args.text_encoder_config)

src/clipcap/clipcap_parse_coco.py

@@ -1,51 +0,0 @@
-import torch
-import skimage.io as io
-import clip
-from PIL import Image
-import pickle
-import json
-import os
-from tqdm import tqdm
-import argparse
-
-
-def main(clip_model_type: str):
-    device = torch.device('cuda:0')
-    clip_model_name = clip_model_type.replace('/', '_')
-    out_path = f"./data/coco/oscar_split_{clip_model_name}_train.pkl"
-    clip_model, preprocess = clip.load(clip_model_type, device=device, jit=False)
-    with open('./data/coco/annotations/train_caption.json', 'r') as f:
-        data = json.load(f)
-    print("%0d captions loaded from json " % len(data))
-    all_embeddings = []
-    all_captions = []
-    for i in tqdm(range(len(data))):
-        d = data[i]
-        img_id = d["image_id"]
-        filename = f"./data/coco/train2014/COCO_train2014_{int(img_id):012d}.jpg"
-        if not os.path.isfile(filename):
-            filename = f"./data/coco/val2014/COCO_val2014_{int(img_id):012d}.jpg"
-        image = io.imread(filename)
-        image = preprocess(Image.fromarray(image)).unsqueeze(0).to(device)
-        with torch.no_grad():
-            prefix = clip_model.encode_image(image).cpu()
-        d["clip_embedding"] = i
-        all_embeddings.append(prefix)
-        all_captions.append(d)
-        if (i + 1) % 10000 == 0:
-            with open(out_path, 'wb') as f:
-                pickle.dump({"clip_embedding": torch.cat(all_embeddings, dim=0), "captions": all_captions}, f)
-
-    with open(out_path, 'wb') as f:
-        pickle.dump({"clip_embedding": torch.cat(all_embeddings, dim=0), "captions": all_captions}, f)
-
-    print('Done')
-    print("%0d embeddings saved " % len(all_embeddings))
-    return 0
-
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--clip_model_type', default="ViT-B/32", choices=('RN50', 'RN101', 'RN50x4', 'ViT-B/32'))
-    args = parser.parse_args()
-    exit(main(args.clip_model_type))

src/clipcap/entrypoint.py

@@ -1,564 +0,0 @@
-import torch
-from torch import nn
-import json
-import os
-from transformers import GPT2Tokenizer, GPT2LMHeadModel
-from typing import List, Optional, Tuple, Union
-from argparse import Namespace
-from enum import Enum
-
-import torch.nn.functional as nnf
-
-class MappingType(Enum):
-    MLP = 'mlp'
-    Transformer = 'transformer'
-
-
-class MLP(nn.Module):
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.model(x)
-
-    def __init__(self, sizes: Tuple[int, ...], bias=True, act=nn.Tanh):
-        super(MLP, self).__init__()
-        layers = []
-        for i in range(len(sizes) - 1):
-            layers.append(nn.Linear(sizes[i], sizes[i + 1], bias=bias))
-            if i < len(sizes) - 2:
-                layers.append(act())
-        self.model = nn.Sequential(*layers)
-
-
-class MlpTransformer(nn.Module):
-    def __init__(self, in_dim, h_dim, out_d: Optional[int] = None, act=nnf.relu, dropout=0.):
-        super().__init__()
-        out_d = out_d if out_d is not None else in_dim
-        self.fc1 = nn.Linear(in_dim, h_dim)
-        self.act = act
-        self.fc2 = nn.Linear(h_dim, out_d)
-        self.dropout = nn.Dropout(dropout)
-
-    def forward(self, x):
-        x = self.fc1(x)
-        x = self.act(x)
-        x = self.dropout(x)
-        x = self.fc2(x)
-        x = self.dropout(x)
-        return x
-
-class MultiHeadAttention(nn.Module):
-
-    def __init__(self, dim_self, dim_ref, num_heads, bias=True, dropout=0.):
-        super().__init__()
-        self.num_heads = num_heads
-        head_dim = dim_self // num_heads
-        self.scale = head_dim ** -0.5
-        self.to_queries = nn.Linear(dim_self, dim_self, bias=bias)
-        self.to_keys_values = nn.Linear(dim_ref, dim_self * 2, bias=bias)
-        self.project = nn.Linear(dim_self, dim_self)
-        self.dropout = nn.Dropout(dropout)
-
-    def forward(self, x, y=None, mask=None):
-        y = y if y is not None else x
-        b, n, c = x.shape
-        _, m, d = y.shape
-        # b n h dh
-        queries = self.to_queries(x).reshape(b, n, self.num_heads, c // self.num_heads)
-        # b m 2 h dh
-        keys_values = self.to_keys_values(y).reshape(b, m, 2, self.num_heads, c // self.num_heads)
-        keys, values = keys_values[:, :, 0], keys_values[:, :, 1]
-        attention = torch.einsum('bnhd,bmhd->bnmh', queries, keys) * self.scale
-        if mask is not None:
-            if mask.dim() == 2:
-                mask = mask.unsqueeze(1)
-            attention = attention.masked_fill(mask.unsqueeze(3), float("-inf"))
-        attention = attention.softmax(dim=2)
-        out = torch.einsum('bnmh,bmhd->bnhd', attention, values).reshape(b, n, c)
-        out = self.project(out)
-        return out, attention
-
-
-class TransformerLayer(nn.Module):
-
-    def forward_with_attention(self, x, y=None, mask=None):
-        x_, attention = self.attn(self.norm1(x), y, mask)
-        x = x + x_
-        x = x + self.mlp(self.norm2(x))
-        return x, attention
-
-    def forward(self, x, y=None, mask=None):
-        x = x + self.attn(self.norm1(x), y, mask)[0]
-        x = x + self.mlp(self.norm2(x))
-        return x
-
-    def __init__(self, dim_self, dim_ref, num_heads, mlp_ratio=4., bias=False, dropout=0., act=nnf.relu,
-                 norm_layer: nn.Module = nn.LayerNorm):
-        super().__init__()
-        self.norm1 = norm_layer(dim_self)
-        self.attn = MultiHeadAttention(dim_self, dim_ref, num_heads, bias=bias, dropout=dropout)
-        self.norm2 = norm_layer(dim_self)
-        self.mlp = MlpTransformer(dim_self, int(dim_self * mlp_ratio), act=act, dropout=dropout)
-
-
-class Transformer(nn.Module):
-
-    def forward_with_attention(self, x, y=None, mask=None):
-        attentions = []
-        for layer in self.layers:
-            x, att = layer.forward_with_attention(x, y, mask)
-            attentions.append(att)
-        return x, attentions
-
-    def forward(self, x, y=None, mask=None):
-        for i, layer in enumerate(self.layers):
-            if i % 2 == 0 and self.enc_dec: # cross
-                x = layer(x, y)
-            elif self.enc_dec:  # self
-                x = layer(x, x, mask)
-            else:  # self or cross
-                x = layer(x, y, mask)
-        return x
-
-    def __init__(self, dim_self: int, num_heads: int, num_layers: int, dim_ref: Optional[int] = None,
-                 mlp_ratio: float = 2., act=nnf.relu, norm_layer: nn.Module = nn.LayerNorm, enc_dec: bool = False):
-        super(Transformer, self).__init__()
-        dim_ref = dim_ref if dim_ref is not None else dim_self
-        self.enc_dec = enc_dec
-        if enc_dec:
-            num_layers = num_layers * 2
-        layers = []
-        for i in range(num_layers):
-            if i % 2 == 0 and enc_dec:  # cross
-                layers.append(TransformerLayer(dim_self, dim_ref, num_heads, mlp_ratio, act=act, norm_layer=norm_layer))
-            elif enc_dec:  # self
-                layers.append(TransformerLayer(dim_self, dim_self, num_heads, mlp_ratio, act=act, norm_layer=norm_layer))
-            else:  # self or cross
-                layers.append(TransformerLayer(dim_self, dim_ref, num_heads, mlp_ratio, act=act, norm_layer=norm_layer))
-        self.layers = nn.ModuleList(layers)
-
-class TransformerMapper(nn.Module):
-    def forward(self, x):
-        x = self.linear(x).view(x.shape[0], self.clip_length, -1)
-        prefix = self.prefix_const.unsqueeze(0).expand(x.shape[0], *self.prefix_const.shape)
-        prefix = torch.cat((x, prefix), dim=1)
-        out = self.transformer(prefix)[:, self.clip_length:]
-        return out
-
-    def __init__(self, dim_clip: int, dim_embedding: int, prefix_length: int, clip_length: int, num_layers: int = 8):
-        super(TransformerMapper, self).__init__()
-        self.clip_length = clip_length
-        self.transformer = Transformer(dim_embedding, 8, num_layers) #nn.Transformer(d_model=dim_embedding, nhead=8, num_encoder_layers=num_layers)
-        self.linear = nn.Linear(dim_clip, clip_length * dim_embedding)
-        self.prefix_const = nn.Parameter(torch.randn(prefix_length, dim_embedding), requires_grad=True)
-
-
-class ClipCaptionModel(nn.Module):
-
-    def get_dummy_token(self, batch_size: int, device: torch.device) -> torch.Tensor:
-        return torch.zeros(batch_size, self.prefix_length, dtype=torch.int64, device=device)
-
-    def forward(self, tokens: torch.Tensor, prefix: torch.Tensor, mask: Optional[torch.Tensor] = None,
-                labels: Optional[torch.Tensor] = None):
-        embedding_text = self.gpt.transformer.wte(tokens)
-        prefix_projections = self.clip_project(prefix).view(-1, self.prefix_length, self.gpt_embedding_size)
-        embedding_cat = torch.cat((prefix_projections, embedding_text), dim=1)
-        if labels is not None:
-            dummy_token = self.get_dummy_token(tokens.shape[0], tokens.device)
-            labels = torch.cat((dummy_token, tokens), dim=1)
-        out = self.gpt(inputs_embeds=embedding_cat, labels=labels, attention_mask=mask)
-        return out
-
-    def __init__(self, prefix_length: int, clip_length: Optional[int] = None, prefix_size: int = 512,
-                 num_layers: int = 8, mapping_type: MappingType = MappingType.MLP):
-        super(ClipCaptionModel, self).__init__()
-        self.prefix_length = prefix_length
-        self.gpt = GPT2LMHeadModel.from_pretrained('gpt2')
-        self.gpt_embedding_size = self.gpt.transformer.wte.weight.shape[1]
-        if mapping_type == MappingType.MLP:
-            self.clip_project = MLP((prefix_size, (self.gpt_embedding_size * prefix_length) // 2,
-                                     self.gpt_embedding_size * prefix_length))
-        else:
-            self.clip_project = TransformerMapper(prefix_size, self.gpt_embedding_size, prefix_length,
-                                                  clip_length, num_layers)
-
-
-class ClipCaptionPrefix(ClipCaptionModel):
-
-    def parameters(self, recurse: bool = True):
-        return self.clip_project.parameters()
-
-    def train(self, mode: bool = True):
-        super(ClipCaptionPrefix, self).train(mode)
-        self.gpt.eval()
-        return self
-
-
-def generate_batched(
-        model,
-        tokenizer,
-        prefix_embeds,
-        entry_length=67,
-        top_p=0.8,
-        temperature=1.0,
-        stop_token: str = '.',
-):
-    """
-    Batched text generation for ClipCap models.
-    
-    Args:
-        model: ClipCap model
-        tokenizer: GPT2 tokenizer
-        prefix_embeds: (batch_size, prefix_length, embedding_dim) - prefix embeddings
-        entry_length: Maximum sequence length to generate
-        top_p: Nucleus sampling parameter
-        temperature: Sampling temperature
-        stop_token: Token to stop generation
-    
-    Returns:
-        List[str]: Generated captions for each item in batch
-    """
-    model.eval()
-    device = next(model.parameters()).device
-    batch_size = prefix_embeds.shape[0]
-    
-    # Initialize
-    stop_token_index = tokenizer.encode(stop_token)[0]
-    filter_value = -float("Inf")
-    
-    # Track which sequences are still generating
-    active_sequences = torch.ones(batch_size, dtype=torch.bool, device=device)
-    
-    # Initialize token sequences - start with None
-    tokens = None
-    generated_embeds = prefix_embeds  # Start with prefix embeddings
-    
-    with torch.no_grad():
-        for step in range(entry_length):
-            # Forward pass for all active sequences
-            outputs = model.gpt(inputs_embeds=generated_embeds)
-            logits = outputs.logits[:, -1, :]  # Get logits for last token: (batch_size, vocab_size)
-            
-            # Apply temperature
-            logits = logits / (temperature if temperature > 0 else 1.0)
-            
-            # Apply nucleus sampling for each sequence in batch
-            for i in range(batch_size):
-                if not active_sequences[i]:
-                    continue
-                    
-                # Sort logits for this sequence
-                sorted_logits, sorted_indices = torch.sort(logits[i], descending=True)
-                cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
-                
-                # Find indices to remove (above top_p threshold)
-                sorted_indices_to_remove = cumulative_probs > top_p
-                sorted_indices_to_remove[1:] = sorted_indices_to_remove[:-1].clone()
-                sorted_indices_to_remove[0] = 0
-                
-                # Set logits to -inf for tokens to remove
-                indices_to_remove = sorted_indices[sorted_indices_to_remove]
-                logits[i, indices_to_remove] = filter_value
-            
-            # Clamp logits to avoid extreme values
-            logits = torch.clamp(logits, min=-1e9, max=1e9)  # keep values bounded
-            # Sample next tokens for all sequences
-            probs = torch.softmax(logits, dim=-1)
-
-            # if some sequences probs tensor contains NaNs (e.g. all logits were -inf), set stop_token_index prob to 1
-            for i in range(batch_size):
-                if torch.isnan(probs[i]).all(): #if not torch.isfinite(probs[i]).any() or probs[i].sum() == 0:
-                    probs[i] = torch.zeros_like(probs[i])
-                    probs[i, stop_token_index] = 1.0
-
-            next_tokens = torch.multinomial(probs, num_samples=1)  # (batch_size, 1)
-            
-            # Get embeddings for next tokens
-            next_token_embeds = model.gpt.transformer.wte(next_tokens)  # (batch_size, 1, embed_dim)
-            
-            # Update token sequences
-            if tokens is None:
-                tokens = next_tokens
-            else:
-                tokens = torch.cat((tokens, next_tokens), dim=1)
-            
-            # Update generated embeddings
-            generated_embeds = torch.cat((generated_embeds, next_token_embeds), dim=1)
-            
-            # Check for stop tokens and update active sequences
-            for i in range(batch_size):
-                if active_sequences[i] and next_tokens[i].item() == stop_token_index:
-                    active_sequences[i] = False
-            
-            # If all sequences have stopped, break early
-            if not active_sequences.any():
-                break
-    
-    # Decode all sequences
-    captions = []
-    for i in range(batch_size):
-        if tokens is not None:
-            token_list = tokens[i].cpu().numpy().tolist()
-            # Remove padding and decode
-            caption = tokenizer.decode(token_list)
-            # Clean up the caption
-            caption = caption.split(stop_token)[0] + stop_token
-            captions.append(caption)
-        else:
-            captions.append("")
-    
-    return captions
-
-
-def generate2(
-        model,
-        tokenizer,
-        tokens=None,
-        prompt=None,
-        embed=None,
-        entry_count=1,
-        entry_length=67,  # maximum number of words
-        top_p=0.8,
-        temperature=1.,
-        stop_token: str = '.',
-):
-    """
-    Legacy single-sequence generation function.
-    For new code, use generate_batched instead.
-    """
-    model.eval()
-    generated_num = 0
-    generated_list = []
-    stop_token_index = tokenizer.encode(stop_token)[0]
-    filter_value = -float("Inf")
-    device = next(model.parameters()).device
-
-    with torch.no_grad():
-
-        for entry_idx in range(entry_count):
-            if embed is not None:
-                generated = embed
-            else:
-                if tokens is None:
-                    tokens = torch.tensor(tokenizer.encode(prompt))
-                    tokens = tokens.unsqueeze(0).to(device)
-
-                generated = model.gpt.transformer.wte(tokens)
-
-            for i in range(entry_length):
-
-                outputs = model.gpt(inputs_embeds=generated)
-                logits = outputs.logits
-                logits = logits[:, -1, :] / (temperature if temperature > 0 else 1.0)
-                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
-                cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
-                sorted_indices_to_remove = cumulative_probs > top_p
-                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[
-                                                    ..., :-1
-                                                    ].clone()
-                sorted_indices_to_remove[..., 0] = 0
-
-                indices_to_remove = sorted_indices[sorted_indices_to_remove]
-                logits[:, indices_to_remove] = filter_value
-                next_token = torch.multinomial(torch.softmax(logits, dim=-1), num_samples=1)
-                next_token_embed = model.gpt.transformer.wte(next_token)
-                if tokens is None:
-                    tokens = next_token
-                else:
-                    tokens = torch.cat((tokens, next_token), dim=1)
-                generated = torch.cat((generated, next_token_embed), dim=1)
-                if stop_token_index == next_token.item():
-                    break
-
-            output_list = list(tokens.squeeze().cpu().numpy())
-            output_text = tokenizer.decode(output_list)
-            generated_list.append(output_text)
-
-    return generated_list[0]
-
-
-class ClipCapModel(torch.nn.Module):
-    """
-    ClipCap integration for the Patchioner class.
-    """
-
-    def __init__(self, args, device, dino_feature_dim=768):
-        super(ClipCapModel, self).__init__()
-        args_dict = args.copy()
-        self.args = args = self.load_config(args)
-        self.device = device
-        self.dino_feature_dim = dino_feature_dim
-        
-        # Initialize tokenizer
-        self.tokenizer = GPT2Tokenizer.from_pretrained(args.language_model)
-        if self.tokenizer.pad_token_id is None:
-            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
-        
-        # Determine mapping type
-        mapping_type = MappingType.MLP if args.mapping_type.lower() == 'mlp' else MappingType.Transformer
-        
-        # Initialize model with DINO feature dimensions
-        if args.only_prefix:
-            self.model = ClipCaptionPrefix(
-                prefix_length=args.prefix_length,
-                clip_length=args.clip_length,
-                prefix_size=dino_feature_dim,
-                num_layers=args.num_layers,
-                mapping_type=mapping_type
-            )
-        else:
-            self.model = ClipCaptionModel(
-                prefix_length=args.prefix_length,
-                clip_length=args.clip_length,
-                prefix_size=dino_feature_dim,
-                num_layers=args.num_layers,
-                mapping_type=mapping_type
-            )
-        
-        # Load trained weights
-        print(f"Loading ClipCap weights from: {args.weight_path}")
-        checkpoint = torch.load(args.weight_path, map_location=device)
-        self.model.load_state_dict(checkpoint, strict=False)
-        self.model.to(device)
-        self.model.eval()
-
-    defaults = {
-        "language_model": "gpt2",
-        "prefix_length": 10,
-        "clip_length": 10,
-        "num_layers": 8,
-        "mapping_type": "mlp",
-        "only_prefix": True,
-        "temperature": 1.0,
-        "top_p": 0.8,
-        "entry_length": 67,
-        "stop_token": ".",
-        "use_batched_generation": True,  # Use batched generation by default
-        "normalize_prefix": False,  # Whether to L2 normalize the input features
-        "weight_path": "/raid/datasets/models_weights/clipcap/training-features/clipcap_dino_vitb14_len10_mlp.pt"
-    }
-
-    def load_config(self, args_dict: dict) -> Namespace:
-        def dict_to_namespace(d):
-            if isinstance(d, dict):
-                return Namespace(**{k: dict_to_namespace(v) for k, v in d.items()})
-            return d
-
-        # Apply defaults
-        for key, value in self.defaults.items():
-            if isinstance(value, dict):
-                for sub_key, sub_value in value.items():
-                    args_dict.setdefault(key, {}).setdefault(sub_key, sub_value)
-            else:
-                args_dict.setdefault(key, value)
-
-        args = dict_to_namespace(args_dict)
-        return args
-
-    def forward(self, dino_features, compute_scores: bool = False) -> List[str]:
-        """
-        DINO Features: (batch_size, dino_feature_dim)
-        - returns: List[str] of generated captions
-        """
-        if self.args.use_batched_generation:
-            return self.forward_batched(dino_features, compute_scores)
-        else:
-            return self.forward_sequential(dino_features, compute_scores)
-    
-    def forward_batched(self, dino_features, compute_scores: bool = False) -> List[str]:
-        """
-        Efficient batched generation for multiple sequences.
-        """
-        batch_size = dino_features.shape[0]
-        
-        # Apply normalization if specified (to match training)
-        if self.args.normalize_prefix:
-            dino_features = dino_features / dino_features.norm(dim=-1, keepdim=True)
-        
-        # Generate prefix embeddings for entire batch
-        with torch.no_grad():
-            prefix_embeds = self.model.clip_project(dino_features).view(
-                batch_size, self.args.prefix_length, -1
-            )
-            
-            # Generate captions for entire batch
-            captions = generate_batched(
-                model=self.model,
-                tokenizer=self.tokenizer,
-                prefix_embeds=prefix_embeds,
-                entry_length=self.args.entry_length,
-                temperature=self.args.temperature,
-                top_p=self.args.top_p,
-                stop_token=self.args.stop_token
-            )
-            
-            if compute_scores:
-                # Compute perplexity scores for generated captions
-                scores = self.compute_perplexity_scores(captions)
-                return captions, scores
-            else:
-                return captions
-    
-    def forward_sequential(self, dino_features, compute_scores: bool = False) -> List[str]:
-        """
-        Sequential generation for backward compatibility or debugging.
-        """
-        batch_size = dino_features.shape[0]
-        captions = []
-        scores = []
-
-        # Process each feature in the batch sequentially
-        for i in range(batch_size):
-            feature = dino_features[i:i+1]  # Keep batch dimension
-            
-            # Apply normalization if enabled
-            if self.args.normalize_prefix:
-                feature = feature / feature.norm(dim=-1, keepdim=True)
-            
-            # Generate prefix embeddings
-            with torch.no_grad():
-                prefix_embed = self.model.clip_project(feature).view(1, self.args.prefix_length, -1)
-                
-                # Generate caption using legacy function
-                caption = generate2(
-                    model=self.model,
-                    tokenizer=self.tokenizer,
-                    embed=prefix_embed,
-                    entry_length=self.args.entry_length,
-                    temperature=self.args.temperature,
-                    top_p=self.args.top_p,
-                    stop_token=self.args.stop_token
-                )
-                
-                captions.append(caption)
-                if compute_scores:
-                    # Compute perplexity for this caption
-                    score = self.compute_perplexity_scores([caption])[0]
-                    scores.append(score)
-
-        return captions if not compute_scores else (captions, scores)
-    
-    def compute_perplexity_scores(self, captions: List[str]) -> List[float]:
-        """
-        Compute perplexity scores for generated captions.
-        """
-        scores = []
-        self.model.eval()
-        
-        with torch.no_grad():
-            for caption in captions:
-                try:
-                    # Tokenize caption
-                    tokens = self.tokenizer.encode(caption, return_tensors='pt').to(self.device)
-                    
-                    # Compute loss (negative log-likelihood)
-                    outputs = self.model.gpt(input_ids=tokens, labels=tokens)
-                    loss = outputs.loss
-                    
-                    # Convert to perplexity (lower is better, but we'll use 1/perplexity as score)
-                    perplexity = torch.exp(loss).item()
-                    score = 1.0 / perplexity if perplexity > 0 else 1.0
-                    scores.append(score)
-                except:
-                    # Fallback score if computation fails
-                    scores.append(1.0)
-        
-        return scores

src/clipcap/predict.py

@@ -1,302 +0,0 @@
-# Prediction interface for Cog ⚙️
-# Reference: https://github.com/replicate/cog/blob/main/docs/python.md
-
-import clip
-import os
-from torch import nn
-import numpy as np
-import torch
-import torch.nn.functional as nnf
-import sys
-from typing import Tuple, List, Union, Optional
-from transformers import (
-    GPT2Tokenizer,
-    GPT2LMHeadModel,
-    AdamW,
-    get_linear_schedule_with_warmup,
-)
-import skimage.io as io
-import PIL.Image
-
-import cog
-
-# import torch
-
-N = type(None)
-V = np.array
-ARRAY = np.ndarray
-ARRAYS = Union[Tuple[ARRAY, ...], List[ARRAY]]
-VS = Union[Tuple[V, ...], List[V]]
-VN = Union[V, N]
-VNS = Union[VS, N]
-T = torch.Tensor
-TS = Union[Tuple[T, ...], List[T]]
-TN = Optional[T]
-TNS = Union[Tuple[TN, ...], List[TN]]
-TSN = Optional[TS]
-TA = Union[T, ARRAY]
-
-WEIGHTS_PATHS = {
-    "coco": "coco_weights.pt",
-    "conceptual-captions": "conceptual_weights.pt",
-}
-
-D = torch.device
-CPU = torch.device("cpu")
-
-
-class Predictor(cog.Predictor):
-    def setup(self):
-        """Load the model into memory to make running multiple predictions efficient"""
-        self.device = torch.device("cuda")
-        self.clip_model, self.preprocess = clip.load(
-            "ViT-B/32", device=self.device, jit=False
-        )
-        self.tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
-
-        self.models = {}
-        self.prefix_length = 10
-        for key, weights_path in WEIGHTS_PATHS.items():
-            model = ClipCaptionModel(self.prefix_length)
-            model.load_state_dict(torch.load(weights_path, map_location=CPU))
-            model = model.eval()
-            model = model.to(self.device)
-            self.models[key] = model
-
-    @cog.input("image", type=cog.Path, help="Input image")
-    @cog.input(
-        "model",
-        type=str,
-        options=WEIGHTS_PATHS.keys(),
-        default="coco",
-        help="Model to use",
-    )
-    @cog.input(
-        "use_beam_search",
-        type=bool,
-        default=False,
-        help="Whether to apply beam search to generate the output text",
-    )
-    def predict(self, image, model, use_beam_search):
-        """Run a single prediction on the model"""
-        image = io.imread(image)
-        model = self.models[model]
-        pil_image = PIL.Image.fromarray(image)
-        image = self.preprocess(pil_image).unsqueeze(0).to(self.device)
-        with torch.no_grad():
-            prefix = self.clip_model.encode_image(image).to(
-                self.device, dtype=torch.float32
-            )
-            prefix_embed = model.clip_project(prefix).reshape(1, self.prefix_length, -1)
-        if use_beam_search:
-            return generate_beam(model, self.tokenizer, embed=prefix_embed)[0]
-        else:
-            return generate2(model, self.tokenizer, embed=prefix_embed)
-
-
-class MLP(nn.Module):
-    def forward(self, x: T) -> T:
-        return self.model(x)
-
-    def __init__(self, sizes: Tuple[int, ...], bias=True, act=nn.Tanh):
-        super(MLP, self).__init__()
-        layers = []
-        for i in range(len(sizes) - 1):
-            layers.append(nn.Linear(sizes[i], sizes[i + 1], bias=bias))
-            if i < len(sizes) - 2:
-                layers.append(act())
-        self.model = nn.Sequential(*layers)
-
-
-class ClipCaptionModel(nn.Module):
-
-    # @functools.lru_cache #FIXME
-    def get_dummy_token(self, batch_size: int, device: D) -> T:
-        return torch.zeros(
-            batch_size, self.prefix_length, dtype=torch.int64, device=device
-        )
-
-    def forward(
-        self, tokens: T, prefix: T, mask: Optional[T] = None, labels: Optional[T] = None
-    ):
-        embedding_text = self.gpt.transformer.wte(tokens)
-        prefix_projections = self.clip_project(prefix).view(
-            -1, self.prefix_length, self.gpt_embedding_size
-        )
-        # print(embedding_text.size()) #torch.Size([5, 67, 768])
-        # print(prefix_projections.size()) #torch.Size([5, 1, 768])
-        embedding_cat = torch.cat((prefix_projections, embedding_text), dim=1)
-        if labels is not None:
-            dummy_token = self.get_dummy_token(tokens.shape[0], tokens.device)
-            labels = torch.cat((dummy_token, tokens), dim=1)
-        out = self.gpt(inputs_embeds=embedding_cat, labels=labels, attention_mask=mask)
-        return out
-
-    def __init__(self, prefix_length: int, prefix_size: int = 512):
-        super(ClipCaptionModel, self).__init__()
-        self.prefix_length = prefix_length
-        self.gpt = GPT2LMHeadModel.from_pretrained("gpt2")
-        self.gpt_embedding_size = self.gpt.transformer.wte.weight.shape[1]
-        if prefix_length > 10:  # not enough memory
-            self.clip_project = nn.Linear(
-                prefix_size, self.gpt_embedding_size * prefix_length
-            )
-        else:
-            self.clip_project = MLP(
-                (
-                    prefix_size,
-                    (self.gpt_embedding_size * prefix_length) // 2,
-                    self.gpt_embedding_size * prefix_length,
-                )
-            )
-
-
-class ClipCaptionPrefix(ClipCaptionModel):
-    def parameters(self, recurse: bool = True):
-        return self.clip_project.parameters()
-
-    def train(self, mode: bool = True):
-        super(ClipCaptionPrefix, self).train(mode)
-        self.gpt.eval()
-        return self
-
-
-def generate_beam(
-    model,
-    tokenizer,
-    beam_size: int = 5,
-    prompt=None,
-    embed=None,
-    entry_length=67,
-    temperature=1.0,
-    stop_token: str = ".",
-):
-
-    model.eval()
-    stop_token_index = tokenizer.encode(stop_token)[0]
-    tokens = None
-    scores = None
-    device = next(model.parameters()).device
-    seq_lengths = torch.ones(beam_size, device=device)
-    is_stopped = torch.zeros(beam_size, device=device, dtype=torch.bool)
-    with torch.no_grad():
-        if embed is not None:
-            generated = embed
-        else:
-            if tokens is None:
-                tokens = torch.tensor(tokenizer.encode(prompt))
-                tokens = tokens.unsqueeze(0).to(device)
-                generated = model.gpt.transformer.wte(tokens)
-        for i in range(entry_length):
-            outputs = model.gpt(inputs_embeds=generated)
-            logits = outputs.logits
-            logits = logits[:, -1, :] / (temperature if temperature > 0 else 1.0)
-            logits = logits.softmax(-1).log()
-            if scores is None:
-                scores, next_tokens = logits.topk(beam_size, -1)
-                generated = generated.expand(beam_size, *generated.shape[1:])
-                next_tokens, scores = next_tokens.permute(1, 0), scores.squeeze(0)
-                if tokens is None:
-                    tokens = next_tokens
-                else:
-                    tokens = tokens.expand(beam_size, *tokens.shape[1:])
-                    tokens = torch.cat((tokens, next_tokens), dim=1)
-            else:
-                logits[is_stopped] = -float(np.inf)
-                logits[is_stopped, 0] = 0
-                scores_sum = scores[:, None] + logits
-                seq_lengths[~is_stopped] += 1
-                scores_sum_average = scores_sum / seq_lengths[:, None]
-                scores_sum_average, next_tokens = scores_sum_average.view(-1).topk(
-                    beam_size, -1
-                )
-                next_tokens_source = next_tokens // scores_sum.shape[1]
-                seq_lengths = seq_lengths[next_tokens_source]
-                next_tokens = next_tokens % scores_sum.shape[1]
-                next_tokens = next_tokens.unsqueeze(1)
-                tokens = tokens[next_tokens_source]
-                tokens = torch.cat((tokens, next_tokens), dim=1)
-                generated = generated[next_tokens_source]
-                scores = scores_sum_average * seq_lengths
-                is_stopped = is_stopped[next_tokens_source]
-            next_token_embed = model.gpt.transformer.wte(next_tokens.squeeze()).view(
-                generated.shape[0], 1, -1
-            )
-            generated = torch.cat((generated, next_token_embed), dim=1)
-            is_stopped = is_stopped + next_tokens.eq(stop_token_index).squeeze()
-            if is_stopped.all():
-                break
-    scores = scores / seq_lengths
-    output_list = tokens.cpu().numpy()
-    output_texts = [
-        tokenizer.decode(output[: int(length)])
-        for output, length in zip(output_list, seq_lengths)
-    ]
-    order = scores.argsort(descending=True)
-    output_texts = [output_texts[i] for i in order]
-    return output_texts
-
-
-def generate2(
-    model,
-    tokenizer,
-    tokens=None,
-    prompt=None,
-    embed=None,
-    entry_count=1,
-    entry_length=67,  # maximum number of words
-    top_p=0.8,
-    temperature=1.0,
-    stop_token: str = ".",
-):
-    model.eval()
-    generated_num = 0
-    generated_list = []
-    stop_token_index = tokenizer.encode(stop_token)[0]
-    filter_value = -float("Inf")
-    device = next(model.parameters()).device
-
-    with torch.no_grad():
-
-        for entry_idx in range(entry_count):
-            if embed is not None:
-                generated = embed
-            else:
-                if tokens is None:
-                    tokens = torch.tensor(tokenizer.encode(prompt))
-                    tokens = tokens.unsqueeze(0).to(device)
-
-                generated = model.gpt.transformer.wte(tokens)
-
-            for i in range(entry_length):
-
-                outputs = model.gpt(inputs_embeds=generated)
-                logits = outputs.logits
-                logits = logits[:, -1, :] / (temperature if temperature > 0 else 1.0)
-                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
-                cumulative_probs = torch.cumsum(
-                    nnf.softmax(sorted_logits, dim=-1), dim=-1
-                )
-                sorted_indices_to_remove = cumulative_probs > top_p
-                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[
-                    ..., :-1
-                ].clone()
-                sorted_indices_to_remove[..., 0] = 0
-
-                indices_to_remove = sorted_indices[sorted_indices_to_remove]
-                logits[:, indices_to_remove] = filter_value
-                next_token = torch.argmax(logits, -1).unsqueeze(0)
-                next_token_embed = model.gpt.transformer.wte(next_token)
-                if tokens is None:
-                    tokens = next_token
-                else:
-                    tokens = torch.cat((tokens, next_token), dim=1)
-                generated = torch.cat((generated, next_token_embed), dim=1)
-                if stop_token_index == next_token.item():
-                    break
-
-            output_list = list(tokens.squeeze().cpu().numpy())
-            output_text = tokenizer.decode(output_list)
-            generated_list.append(output_text)
-
-    return generated_list[0]

src/dataset.py

@@ -1,94 +0,0 @@
-import torch
-from torch.utils.data import Dataset
-
-from tqdm import tqdm
-import json
-from typing import Tuple
-import clip
-import random
-import json
-import random
-from tqdm import tqdm
-
-class ClipCocoDataset(Dataset):
-
-    def __len__(self) -> int:
-        return len(self.captions_tokens)
-    
-    def pad_tokens(self, item: int):
-        tokens = self.captions_tokens[item]
-        padding = self.max_seq_len - tokens.shape[0]
-        if padding > 0:
-            tokens = torch.cat((tokens, torch.zeros(padding, dtype=torch.int64)))
-        elif padding < 0:
-            tokens = tokens[:self.max_seq_len]
-        return tokens
-
-
-    def __getitem__(self, item: int) -> Tuple[torch.Tensor, ...]:
-        # tokens = self.captions_tokens[item]
-        
-        clip_tokens = self.pad_tokens(item)
-        if self.feats is None:
-            clip_tokens_77 = self.captions_tokens[item]
-            return clip_tokens, clip_tokens_77
-        else:
-            return clip_tokens, self.feats[item]
-
-    def __init__(self, data_path: str,  clip_model=None, talk2dino=None, use_dino_feats=False, tokenizer=None):
-        if tokenizer is not None:
-            self.clip_tokenizer = tokenizer
-        else:
-            print(f"Using default tokenizer")
-            self.clip_tokenizer = clip.tokenize
-        self.prefix_length = 10
-        self.max_seq_len = 20
-        self.feats = None
-        
-        if clip_model is not None:
-            device = next(clip_model.parameters()).device
-            print("Pre-extracting features...")
-
-        if not use_dino_feats:
-            with open(data_path, 'r') as f:
-                self.captions = [ann['caption'] for ann in json.load(f)['annotations']]
-        else:
-            data = torch.load(data_path)
-            self.captions = [ann['caption'] for ann in data['annotations']]
-            self.feats = [ann['features'] for ann in data['annotations']]
-
-                    
-        random.shuffle(self.captions)
-        self.captions_tokens = []
-        
-        batch_size = 64
-        batched_captions = [self.captions[i:i + batch_size] for i in range(0, len(self.captions), batch_size)]
-
-        for batch in tqdm(batched_captions):
-            try:
-                # Tokenize the batch of captions
-                batch_tokens = [torch.tensor(self.clip_tokenizer(caption)[0], dtype=torch.int64) for caption in batch]
-                
-                # Pad tokens to the same length for batching
-                batch_tokens_padded = torch.nn.utils.rnn.pad_sequence(batch_tokens, batch_first=True)
-                self.captions_tokens.extend(batch_tokens)
-
-                if clip_model is not None:
-                    with torch.no_grad():
-                        # Encode the text batch
-                        feats = clip_model.encode_text(batch_tokens_padded.to(device))
-                        
-                        if talk2dino is not None:
-                            # Project to desired feature space
-                            feats = talk2dino.project_clip_txt(feats).to('cpu')
-
-                        # Concatenate features
-                        if self.feats is None:
-                            self.feats = feats
-                        else:
-                            self.feats = torch.cat((self.feats, feats))
-            except Exception as e:
-                print(f"Error processing batch: {e}")
-        print(len(self.captions_tokens))
-
-    

src/datasetMix.py

@@ -1,153 +0,0 @@
-import torch
-from torch.utils.data import Dataset
-
-from tqdm import tqdm
-import json
-from typing import Tuple
-import clip
-import random
-import json
-import random
-from tqdm import tqdm
-
-from pycocotools.coco import COCO
-
-class ClipCocoDatasetMix(Dataset):
-
-    def __len__(self) -> int:
-        return len(self.image_index_list)
-
-    def _pad_tokens(self, tokens: torch.Tensor) -> torch.Tensor:
-        padding = self.max_seq_len - tokens.shape[0]
-        if padding > 0:
-            tokens = torch.cat((tokens, torch.zeros(padding, dtype=torch.int64)))
-        elif padding < 0:
-            tokens = tokens[:self.max_seq_len]
-        return tokens
-
-
-    def __getitem__(self, item: int) -> Tuple[torch.Tensor, ...]:
-        
-        # get the image index for the item
-        img_idx = self.image_index_list[item]
-        # get the caption index for that image
-        first_caption_idx = self.image_index_list.index(img_idx)
-
-        # the caption index is the item - the first caption index
-        caption_idx = item - first_caption_idx
-
-        # how many captions are there for that image?
-        num_captions = len(self.captions_list_of_lists[img_idx])
-        try:
-            tokens = self.captions_tokens_list_of_lists[img_idx][caption_idx] #self.captions_list_of_lists[img_idx][caption_idx]
-        except IndexError:
-            print(f"{len(self.captions_tokens_list_of_lists)= } - {len(self.captions_tokens_list_of_lists[img_idx])= }")
-            print(f"IndexError: {img_idx}, {caption_idx}, {num_captions}")
-            raise
-        padded_tokens = self._pad_tokens(tokens)
-
-        feats_same_img = self.feats[img_idx][random.choice(range(num_captions))]
-
-        if self.feats is None or len(self.feats) == 0:
-            raise Exception("Precomputed features required")
-        else:
-            return padded_tokens, feats_same_img
-
-    def __init__(self, data_path: str,  clip_model=None, talk2dino=None, use_precomputed_feats=False, tokenizer=None):
-        
-        batch_size = 64
-        self.max_seq_len = 20
-
-        if use_precomputed_feats:
-            raise Exception("Precomputed features not supported")
-
-        if tokenizer is not None:
-            self.clip_tokenizer = tokenizer
-        else:
-            print(f"Using default tokenizer")
-            self.clip_tokenizer = clip.tokenize
-
-        coco_data = COCO(data_path)
-        # I want to load the captions from the json file in a list of lists, 
-        # where each list contains the captions for a single image
-
-        self.captions_list_of_lists = []
-        
-        self.image_index_list = []
-
-        max_seq_len = 20
-
-        for img_idx, (img_id, image) in enumerate(list(coco_data.imgs.items())):
-            # get the captions for that image
-            captions = coco_data.imgToAnns[img_id]
-            # get the texts of the captions
-            captions = [cap['caption'] for cap in captions] #[coco_data.anns[cap]['caption'] for cap in captions]
-            self.captions_list_of_lists.append(captions)
-            self.image_index_list.append([img_idx] * len(captions))
-
-            #max_seq_len = max(max_seq_len, max([len(caption) for caption in captions]))
-        
-        self.max_seq_len = max_seq_len
-        print(f"Computed Max seq len: {max_seq_len}")
-
-        if clip_model is not None:
-            device = next(clip_model.parameters()).device
-            print("Pre-extracting features...")
-                    
-        #random.shuffle(self.captions_list_of_lists)
-        # should shuffle in the same way self.image_index_list and self.captions_list_of_lists
-        # Combine captions and image indices into a list of pairs
-        combined = list(zip(self.captions_list_of_lists, self.image_index_list, range(len(self.captions_list_of_lists))))
-
-        # Shuffle them together
-        random.shuffle(combined)
-
-        # Unzip the shuffled pairs back into two separate lists
-        self.captions_list_of_lists, self.image_index_list, img_idxes_shuffled = zip(*combined)
-        # Convert back to lists (zip returns tuples)
-        self.captions_list_of_lists = list(self.captions_list_of_lists)
-        self.image_index_list = list(self.image_index_list)
-        img_idxes_shuffled = list(img_idxes_shuffled)
-
-        # self.image_index_list is a list of lists, where each list contains the image index for each caption,
-        # so we need to flatten it
-        self.image_index_list = [img_idxes_shuffled.index(item) for sublist in self.image_index_list for item in sublist]
-
-
-        self.captions_tokens_list_of_lists = []
-        self.feats = [] # feats will be a list of tensors, each tensor will be (num_captions, embedding_dimension)
-        #ignore. # feats shape will be (num_images, num_captions, embedding_dimension)
-        
-        #batched_captions = [self.captions[i:i + batch_size] for i in range(0, len(self.captions), batch_size)]
-
-        for captions_list in tqdm(self.captions_list_of_lists, dynamic_ncols=True):
-            try:
-                # Tokenize the batch of captions
-                batch_tokens = [torch.tensor(self.clip_tokenizer(caption)[0], dtype=torch.int64) for caption in captions_list]
-                
-                # Pad tokens to the same length for batching
-                batch_tokens_padded = torch.nn.utils.rnn.pad_sequence(batch_tokens, batch_first=True)
-                self.captions_tokens_list_of_lists.append(batch_tokens)
-
-                # alternative:
-                # tokens = self.clip_tokenizer(captions_list, truncate=True).to(device)  # shape: (num_captions, context_length)
-
-
-                if clip_model is not None:
-                    with torch.no_grad():
-                        # Encode the text batch
-                        feats = clip_model.encode_text(batch_tokens_padded.to(device))
-                        
-                        if talk2dino is not None:
-                            # Project to desired feature space
-                            feats = talk2dino.project_clip_txt(feats).to('cpu')
-
-                        self.feats.append(feats.cpu())  # store (num_captions, embed_dim) for each image
-
-            except Exception as e:
-                print(f"Error processing batch: {e}")
-        
-        print(f"Dataset loaded with {len(self.captions_list_of_lists)} images")
-        print(f"Max seq len: {max_seq_len}")
-        print(f"Number of captions: {len(self.image_index_list)}")
-    

src/decap/decap.py

@@ -1,193 +0,0 @@
-import os
-from torch import nn
-import numpy as np
-import torch
-import torch.nn.functional as nnf
-import sys
-from typing import Tuple, List, Union, Optional
-from tqdm import tqdm, trange
-import pickle
-import PIL.Image as Image
-import json
-import random
-import sys
-import clip
-import PIL
-import random
-
-from torch.utils.data import Dataset, DataLoader
-from enum import Enum
-from transformers import GPT2Tokenizer, GPT2LMHeadModel, AdamW, get_linear_schedule_with_warmup
-from tqdm import tqdm
-import os
-import pickle
-import sys
-import argparse
-import json
-from typing import Tuple, Optional, Union
-
-import os
-from dotenv import load_dotenv
-
-load_dotenv()
-
-
-DECAP_DECODER_CONFIG_PATH = os.path.join(os.path.dirname(__file__), "decoder_config.pkl")
-DECAP_COCO_WEIGHTS_PATH = None#'../../thesis-data/decap/coco_model/coco_prefix-009.pt'
-        
-class MappingType(Enum):
-    MLP = 'mlp'
-    Transformer = 'transformer'
-
-    
-class MLP(nn.Module):
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.model(x)
-
-    def __init__(self, sizes: Tuple[int, ...], bias=True, act=nn.Tanh):
-        super(MLP, self).__init__()
-        layers = []
-        for i in range(len(sizes) - 1):
-            layers.append(nn.Linear(sizes[i], sizes[i + 1], bias=bias))
-            if i < len(sizes) - 2:
-                layers.append(act())
-        self.model = nn.Sequential(*layers)
-        
-
-class DeCap(nn.Module):
-
-    def __init__(self,prefix_size: int = 512):
-        super(DeCap, self).__init__()
-        # decoder: 4 layers transformer with 4 attention heads
-        # the decoder is not pretrained
-        with open(DECAP_DECODER_CONFIG_PATH,'rb') as f:
-            config = pickle.load(f)
-        self.decoder = GPT2LMHeadModel(config)
-        self.embedding_size = self.decoder.transformer.wte.weight.shape[1]
-        self.clip_project = MLP((prefix_size,self.embedding_size))
-        
-    def forward(self, clip_features,tokens):
-        embedding_text = self.decoder.transformer.wte(tokens)
-        embedding_clip = self.clip_project(clip_features)
-        embedding_clip = embedding_clip.reshape(-1,1,self.embedding_size)
-        embedding_cat = torch.cat([embedding_clip,embedding_text],dim=1)
-        out = self.decoder(inputs_embeds=embedding_cat)
-        return out
-
-from clip.simple_tokenizer import SimpleTokenizer as _Tokenizer
-_Tokenizer = _Tokenizer()
-
-def Decoding(model,clip_features):
-    model.eval()
-    embedding_cat = model.clip_project(clip_features).reshape(1,1,-1)
-    entry_length = 30
-    temperature = 1
-    tokens = None
-    for i in range(entry_length):
-        # print(location_token.shape)
-        outputs = model.decoder(inputs_embeds=embedding_cat)
-
-        logits = outputs.logits
-        logits = logits[:, -1, :] / (temperature if temperature > 0 else 1.0)
-        logits_max = logits.max()
-        logits = torch.nn.functional.softmax(logits, -1)
-        next_token = torch.argmax(logits, -1).unsqueeze(0)
-        next_token_embed = model.decoder.transformer.wte(next_token)
-
-        if tokens is None:
-            tokens = next_token
-
-        else:
-            tokens = torch.cat((tokens, next_token), dim=1)
-        if next_token.item()==49407:
-            break
-        embedding_cat = torch.cat((embedding_cat, next_token_embed), dim=1)
-    try:
-        output_list = list(tokens.squeeze().cpu().numpy())
-        output = _Tokenizer.decode(output_list)
-    except:
-        output = 'None'
-    return output
-
-def decoding_batched(model, clip_features, compute_scores : bool = False, decoding_method : callable = None, return_start_end_tokens : bool = False):
-    """
-    Returns the generated sequences for a batch of clip features.
-    - if compute_scores is True, also returns the scores of the generated sequences.
-    - returns a list of strings if compute_scores is False, otherwise a tuple of a list of strings and a list of floats.
-    """
-
-    model.eval()
-    embedding_cat = model.clip_project(clip_features).view(clip_features.shape[0], 1, -1)
-    entry_length = 30
-    temperature = 1
-    tokens = None
-    sequence_log_probs = None
-
-    for i in range(entry_length):
-        outputs = model.decoder(inputs_embeds=embedding_cat)
-
-        logits = outputs.logits[:, -1, :]
-        logits = logits / (temperature if temperature > 0 else 1.0)
-
-        probs = torch.nn.functional.softmax(logits, -1)
-
-        if compute_scores:
-            log_probs = torch.log(probs)  # Convert to log-probabilities
-
-        next_token = torch.argmax(probs, -1).unsqueeze(1)
-        next_token_embed = model.decoder.transformer.wte(next_token)
-
-        if tokens is None:
-            tokens = next_token
-            if compute_scores:
-                sequence_log_probs = log_probs.gather(1, next_token)  # Store log-prob of first token
-        else:
-            tokens = torch.cat((tokens, next_token), dim=1)
-            if compute_scores:
-                token_log_probs = log_probs.gather(1, next_token)  # Get log-prob of chosen token
-                sequence_log_probs = torch.cat((sequence_log_probs, token_log_probs), dim=1)  # Append
-        
-        # Append new token embedding to input
-        embedding_cat = torch.cat((embedding_cat, next_token_embed), dim=1)
-
-    if compute_scores:
-        # Compute total sequence scores
-        sequence_scores = sequence_log_probs.sum(dim=-1)  # Sum log-probs over sequence
-        final_scores = torch.exp(sequence_scores)  # Convert log-sum-prob to probability-like score
-    
-    try:
-        outputs = []
-        for tokens_elem in tokens:
-            output_list = list(tokens_elem.squeeze().cpu().numpy())
-            if decoding_method is not None:
-                output = decoding_method(output_list)
-            else:
-                output = _Tokenizer.decode(output_list)
-            
-
-
-            output = output.split('<|endoftext|>')[0]
-            if not return_start_end_tokens:
-                output = output.replace('<|startoftext|>', '')
-            else:
-                output += '<|endoftext|>'
-
-            outputs.append(output)
-    except:
-        outputs = None
-    
-    return (outputs, final_scores.cpu().numpy().tolist()) if compute_scores else outputs
-
-
-decap_model = None
-
-def get_decap_model(device, weights_path = DECAP_COCO_WEIGHTS_PATH, prefix_size=512):
-    #global decap_model
-    #if decap_model is not None:
-    #    return decap_model
-    decap_model = DeCap(prefix_size)
-    decap_model.load_state_dict(torch.load(weights_path,map_location= torch.device('cpu')), strict=False)
-    decap_model = decap_model.to(device)
-    decap_model = decap_model.eval()
-    return decap_model

src/decap/decoderTraining.py

@@ -1,464 +0,0 @@
-import torch
-from torch.utils.data import DataLoader
-from torch.utils.data.distributed import DistributedSampler
-from torch.nn.parallel import DistributedDataParallel as DDP
-import torch.distributed as dist
-
-from im2txtprojection.im2txtprojection import Im2TxtProjector, ProjectionType
-from transformers import get_linear_schedule_with_warmup
-from torch.optim import AdamW
-from tqdm import tqdm
-from decap import get_decap_model
-import os
-import sys
-import argparse
-import json
-from typing import Union
-import sys
-import clip
-import json
-
-import csv
-
-
-from src.dataset import ClipCocoDataset
-from src.datasetMix import ClipCocoDatasetMix
-from src.model import DeCap, ProjectionLayer
-
-DECAP_DECODER_CONFIG_PATH = os.path.join("./decoder_config.pkl")
-
-def save_config(args: argparse.Namespace):
-    config = {}
-    for key, item in args._get_kwargs():
-        config[key] = item
-    out_path = os.path.join(args.out_dir, f"{args.prefix}.json")
-    with open(out_path, 'w') as outfile:
-        json.dump(config, outfile)
-
-
-def load_model(config_path: str, epoch_or_latest: Union[str, int] = '_latest'):
-    with open(config_path) as f:
-        config = json.load(f)
-    parser = argparse.ArgumentParser()
-    parser.set_defaults(**config)
-    args = parser.parse_args()
-    if type(epoch_or_latest) is int:
-        epoch_or_latest = f"-{epoch_or_latest:03d}"
-    model_path = os.path.join(args.out_dir, f"{args.prefix}{epoch_or_latest}.pt")
-    if args.only_prefix:
-        model = ClipCaptionPrefix(args.prefix_length)
-    else:
-        model = ClipCaptionModel(args.prefix_length)
-    if os.path.isfile(model_path):
-        print(f"loading model from {model_path}")
-        model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
-    else:
-        print(f"{model_path} is not exist")
-    return model, parser
-
-
-    
-
-def train_decoder(args,
-          lr: float = 1e-5, warmup_steps: int = 1000, output_dir: str = ".", output_prefix: str = ""):
-
-    # device = torch.device('cuda:1')
-    batch_size = args.bs
-    epochs = args.epochs
-    if not os.path.exists(output_dir):
-        os.makedirs(output_dir)
-    args.is_master = ( args.local_rank == 0 or args.not_distributed != False)
-
-    # set the device
-    #torch.cuda.set_device(args.local_rank)
-    #device = torch.device('cuda:'+str(args.local_rank))
-    if args.not_distributed == False:
-        torch.cuda.set_device(args.local_rank)
-        device = torch.device('cuda:'+str(args.local_rank))
-        dist.init_process_group(backend='nccl', init_method='env://')
-    else:
-        device = torch.device('cuda:'+str(args.local_rank))
-        print(f"NOT DISTRIBUTED")
-    print(f"Using device {device}")
-    SEED=42
-    torch.cuda.manual_seed_all(SEED)
-    
-    if args.use_regionclip:
-        # RegionCLIP typically uses 1024 dimensions for ResNet-50 or 512 for ViT
-        # We'll determine this from the loaded model
-        prefix_size = 1024  # Default for RegionCLIP ResNet-50, but will be adjusted if needed
-    elif args.denseclip_config is not None:
-        # DenseClip typically uses 512 dimensions (similar to CLIP ViT-B)
-        from src.denseclip.loader import load_denseclip_config
-        denseclip_config_dict = load_denseclip_config(args.denseclip_config)
-        prefix_size = denseclip_config_dict.get('model', {}).get('text', {}).get('embed_dim', None)
-        if prefix_size is None:
-            print(f"Warning: Could not determine prefix_size from DenseClip config {args.denseclip_config}. Defaulting to 512.")
-            prefix_size = 512  # Fallback to a common size)
-
-    elif 'H' in args.clip_model or args.use_dinotxt:
-        prefix_size = 1024
-    elif args.talk2dino_weights is not None or args.use_dino_feats:
-        prefix_size = 768
-    else:
-        prefix_size = 512
-    
-    if args.im_proj:
-        memory_bank_path = os.path.abspath(args.dataset)
-        print(f"Using Im2TxtProjector with {memory_bank_path = }")
-        im_proj = Im2TxtProjector(
-            type=memory_bank_path,
-            use_talk2dino=True,
-            linear_talk2dino=False,
-            memory_bank_name='coco_karpathy',
-            device_str=device)
-
-    if args.use_regionclip:
-        from src.regionclip.loader import load_regionclip_from_checkpoint
-        from src.regionclip.datasets.clip_prompt_utils import tokenize as regionclip_tokenize
-        
-        print("Using RegionCLIP for text encoding.")
-        if args.regionclip_checkpoint is None:
-            raise ValueError("RegionCLIP checkpoint path must be provided when using --use-regionclip")
-        
-        clip_model = load_regionclip_from_checkpoint(
-            args.regionclip_checkpoint, 
-            device=device, 
-            config=args.regionclip_config
-        )
-        tokenizer = regionclip_tokenize
-        preprocess = None  # RegionCLIP doesn't need preprocessing for text-only training
-        
-        # Determine the actual embedding dimension from the loaded model
-        if hasattr(clip_model, 'text_projection'):
-            actual_prefix_size = clip_model.text_projection.shape[1]
-            print(f"RegionCLIP text embedding dimension: {actual_prefix_size}")
-            if actual_prefix_size != prefix_size:
-                print(f"Updating prefix_size from {prefix_size} to {actual_prefix_size}")
-                prefix_size = actual_prefix_size
-        
-        # Test RegionCLIP text encoding to ensure it works
-        try:
-            test_text = ["A test sentence"]
-            test_tokens = tokenizer(test_text)
-            test_features = clip_model.encode_text(test_tokens.to(device))
-            print(f"RegionCLIP test encoding successful. Output shape: {test_features.shape}")
-        except Exception as e:
-            print(f"Warning: RegionCLIP test encoding failed: {e}")
-            print("This might cause issues during training.")
-        
-    elif args.denseclip_config is not None:
-        from src.denseclip.loader import load_denseclip
-        
-        print(f"Using DenseClip for text encoding with config: {args.denseclip_config}")
-        
-        try:
-            clip_model = load_denseclip(
-                config_name=args.denseclip_config,
-                device=device
-            )
-            
-            # Try to use DenseClip's tokenizer first
-            try:
-                from src.denseclip.loader import DenseCLIP_tokenize
-                tokenizer = DenseCLIP_tokenize
-                print("Using DenseClip tokenizer")
-            except ImportError:
-                # Fallback to CLIP tokenizer if DenseClip tokenizer is not available
-                import clip
-                tokenizer = clip.tokenize
-                print("Warning: DenseClip tokenizer not available, using CLIP tokenizer")
-            
-            preprocess = None  # DenseClip doesn't need preprocessing for text-only training
-            
-            # Determine the actual embedding dimension from the loaded model
-            if hasattr(clip_model, 'text_encoder') and hasattr(clip_model.text_encoder, 'embed_dim'):
-                actual_prefix_size = clip_model.text_encoder.embed_dim
-                print(f"DenseClip text embedding dimension: {actual_prefix_size}")
-                if actual_prefix_size != prefix_size:
-                    print(f"Updating prefix_size from {prefix_size} to {actual_prefix_size}")
-                    prefix_size = actual_prefix_size
-            
-            # Test DenseClip text encoding to ensure it works
-            test_text = ["A test sentence"]
-            test_tokens = tokenizer(test_text)
-            if hasattr(test_tokens, 'to'):
-                test_tokens = test_tokens.to(device)
-            test_features = clip_model.encode_text(test_tokens)
-            print(f"DenseClip test encoding successful. Output shape: {test_features.shape}")
-            
-        except Exception as e:
-            print(f"Error loading DenseClip model: {e}")
-            raise e
-        
-    elif args.use_open_clip:
-        from open_clip import create_model_and_transforms, tokenize
-        print("Using open_clip for model loading.")
-        clip_model, preprocess_train, preprocess_val = create_model_and_transforms(model_name=args.clip_model, pretrained="laion2b_s32b_b79k", device=device)
-        preprocess = preprocess_train
-        tokenizer = tokenize
-
-    elif args.use_dinotxt:
-        from src.dinotxt_utils import get_tokenizer
-        clip_model = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitl14_reg4_dinotxt_tet1280d20h24l')
-        tokenizer = get_tokenizer().tokenize
-    else:
-        clip_model, preprocess = clip.load(args.clip_model, device=device, jit=False)
-        tokenizer = clip.tokenize
-    clip_model.eval()
-    clip_model.to(device)
-    
-    # Create model after determining the correct prefix_size
-    if args.decap_weights is None:
-        model = DeCap(prefix_size)
-    else:
-        model = get_decap_model(device, args.decap_weights, prefix_size)
-    
-    if args.talk2dino_weights is not None:
-        # loading Talk2DINO
-        print(f"Loading Talk2DINO weights from {args.talk2dino_weights}")
-        talk2dino = ProjectionLayer.from_config(args.talk2dino_config)
-        talk2dino.load_state_dict(torch.load(args.talk2dino_weights, device))
-        talk2dino.to(device)
-        talk2dino.eval()
-
-    else:
-        talk2dino = None
-        
-        
-    loss_ce = torch.nn.CrossEntropyLoss(ignore_index=0,label_smoothing=0.1)
-    model.to(device)
-
-    if args.not_distributed == False:
-        model = DDP(
-            model,
-            device_ids=[args.local_rank],
-            output_device=args.local_rank,
-            find_unused_parameters=True
-        )
-    
-    if not args.pre_extract_features:
-        print("Features pre-extraction de-activated")
-        if args.mix_captions:
-            print("Using mix captions")
-            dataset = ClipCocoDatasetMix(args.dataset, use_precomputed_feats=args.use_dino_feats, tokenizer=tokenizer)
-        else:
-            dataset = ClipCocoDataset(args.dataset, use_dino_feats=args.use_dino_feats, tokenizer=tokenizer)
-    else:
-        if args.mix_captions:
-            print("Using mix captions")
-            dataset = ClipCocoDatasetMix(args.dataset, clip_model=clip_model, talk2dino=talk2dino, tokenizer=tokenizer)
-        else:
-            dataset = ClipCocoDataset(args.dataset, clip_model=clip_model, talk2dino=talk2dino, tokenizer=tokenizer)
-        
-    
-    optimizer = AdamW(model.parameters(),lr=lr)
-    
-    print(f"Going to construct DataLoader with {len(dataset)} samples")
-    if args.not_distributed == False:
-        sampler = DistributedSampler(dataset)
-        train_dataloader = DataLoader(dataset, sampler=sampler, batch_size=batch_size, drop_last=True)
-    else:
-        train_dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, drop_last=True)
-
-    print("DataLoader constructed")
-    scheduler = get_linear_schedule_with_warmup(
-        optimizer, num_warmup_steps=warmup_steps, num_training_steps=epochs * len(train_dataloader)
-    )
-    
-    
-    for epoch in range(epochs):
-        
-        epoch_loss = 0.0
-        epoch_acc = 0.0
-        num_batches = 0
-        
-        loss_token_save,ac_save= 0,0
-        sys.stdout.flush()
-        if args.is_master:
-            print(f">>> Training epoch {epoch}")
-            progress = tqdm(total=int(len(train_dataloader)/10), desc=output_prefix, dynamic_ncols=True)
-        
-        if args.not_distributed == False:
-            dist.barrier()
-        
-        for idx,(clip_tokens, pipeline_input) in enumerate(train_dataloader):
-            
-
-            clip_tokens, pipeline_input = clip_tokens.to(device), pipeline_input.to(device)
-            
-            with torch.no_grad():
-                if not args.pre_extract_features and not args.use_dino_feats:
-                    if args.use_regionclip:
-                        # RegionCLIP text encoding
-                        feature_text = clip_model.encode_text(pipeline_input)
-                    elif args.denseclip_config is not None:
-                        # DenseClip text encoding
-                        feature_text = clip_model.encode_text(pipeline_input)
-                    else:
-                        # Standard CLIP or OpenCLIP text encoding
-                        feature_text = clip_model.encode_text(pipeline_input)
-                    
-                    if args.use_dinotxt:
-                        feature_text = feature_text[:, 1024:] # patch-aligned text embedding
-
-                    if args.talk2dino_weights is not None:
-                        feature_text = talk2dino.project_clip_txt(feature_text)
-                else:
-                    feature_text = pipeline_input
-                    if args.im_proj:
-                        feature_text = im_proj.project(feature_text, normalize=True)
-                
-                feature_text /= feature_text.norm(dim=-1, keepdim=True)
-                
-                if args.gaussian_noise != 0:
-                    feature_text += args.gaussian_noise * torch.randn(feature_text.shape).to(device)
-                    feature_text /= feature_text.norm(dim=-1, keepdim=True)
-                    
-
-            outputs = model(feature_text.float(),clip_tokens)
-            logits = outputs
-            
-            logits = logits.logits
-
-            logits = logits[:,: -1]
-            clip_tokens = clip_tokens.flatten()
-            logits = logits.reshape(-1, logits.shape[-1])
-            
-            loss_token = loss_ce(logits, clip_tokens)
-            ac=((logits.argmax(1)==clip_tokens)*(clip_tokens>0)).sum()/(clip_tokens>0).sum()
-            optimizer.zero_grad()
-            loss_all = loss_token
-            loss_all.backward()
-            optimizer.step()
-            scheduler.step()
-            
-            epoch_loss += loss_token.item()
-            epoch_acc += ac.item()
-            num_batches += 1
-            
-            if args.is_master:
-                
-                if(idx+1) %10 == 0:
-                    progress.set_postfix({"loss_token": loss_token_save/10.0,"acc_token":ac_save/10.0})
-                    progress.update()
-                    loss_token_save,ac_save= 0,0
-                else:
-                    loss_token_save += loss_token.item()
-                    ac_save += ac.item()
-
-        if args.is_master:
-            log_dir = os.path.join('./log', f"{args.dataset}.txt")#'./log/'+args.dataset+'.txt'
-            with open(log_dir,'w') as f:
-                f.writelines('epoch ' +str(epoch) +': '+ progress.postfix+'\r\n')
-            progress.close()
-            if epoch % args.save_every == 0 or epoch == epochs - 1:
-                torch.save(
-                    model.state_dict(),
-                    os.path.join(output_dir, f"{output_prefix}-{epoch:03d}.pt"),
-                )
-        
-        # after the epoch, we need to synchronize the loss and accuracy across all processes
-        loss_tensor = torch.tensor(epoch_loss, device=device)
-        acc_tensor = torch.tensor(epoch_acc, device=device)
-        count_tensor = torch.tensor(num_batches, device=device)
-
-        if args.not_distributed == False:
-            # sum on all processes
-            torch.distributed.all_reduce(loss_tensor, op=torch.distributed.ReduceOp.SUM)
-            torch.distributed.all_reduce(acc_tensor, op=torch.distributed.ReduceOp.SUM)
-            torch.distributed.all_reduce(count_tensor, op=torch.distributed.ReduceOp.SUM)
-
-        # compute global mean
-        avg_loss = loss_tensor.item() / count_tensor.item()
-        avg_acc = acc_tensor.item() / count_tensor.item()
-
-        if args.is_master:
-            epoch_loss_current = {'epoch': epoch, 'loss': avg_loss, 'accuracy': avg_acc}
-            #epoch_losses.append(epoch_loss_current)
-            print(f"Epoch {epoch} loss: {avg_loss}, accuracy: {avg_acc}")
-
-            loss_csv_path = os.path.join(output_dir, f"{output_prefix}_epoch_losses.csv")
-            with open(loss_csv_path, 'a', newline='') as csvfile:
-                writer = csv.DictWriter(csvfile, fieldnames=['epoch', 'loss', 'accuracy'])
-                # Write the header only if the file is empty
-                if os.stat(loss_csv_path).st_size == 0:
-                    writer.writeheader()
-                writer.writerow(epoch_loss_current)
-    return model
-
-# DeCap CLIP B16 karpathy train split:
-#python decapTraining.py --out_dir weights_clip_b16_karpathy --not-distributed 1 --local-rank 0 --dataset coco_train_karpathy.json --prefix coco_karpathy
-# DECAP with proj -> ma in realtà non serve.
-#python decapTraining.py --out_dir weights_clip_b16_proj_karpathy --not-distributed 1 --local-rank 0 --dataset coco_train_karpathy.json --prefix coco_karpathy --im_proj
-
-# Patchioner DINOv2 karpathy train split with proj:
-#python decapTraining.py --out_dir weights_dino_b14_proj_karpathy --not-distributed 1 --local-rank 1 --dataset coco_train_karpathy.json --prefix coco_karpathy --talk2dino_weights weights_talk2dino/vitb_mlp_infonce.pth --talk2dino_config configs_talk2dino/vitb_mlp_infonce.yaml --pre_extract_features --im_proj
-# Patchioner DINOv2 karpathy train split
-#python decapTraining.py --out_dir weights_dino_b14_karpathy --not-distributed 1 --local-rank 1 --dataset coco_train_karpathy.json --prefix coco_karpathy --talk2dino_weights weights_talk2dino/vitb_mlp_infonce.pth --talk2dino_config configs_talk2dino/vitb_mlp_infonce.yaml
-#python decapTraining.py --out_dir weights_dino_b14_karpathy --not-distributed 1 --local-rank 1 --dataset coco_train_karpathy.json --prefix coco_karpathy --talk2dino_weights weights_talk2dino/vitb_mlp_infonce.pth --talk2dino_config configs_talk2dino/vitb_mlp_infonce.yaml --use_dino_feats --pre_extract_features
-
-# DeCap CLIP B32 karpathy train split:
-#python decapTraining.py --out_dir weights_clip_b32_karpathy --not-distributed 1 --local-rank 0 --dataset coco_train_karpathy.json --prefix coco_karpathy --clip_model ViT-B/32
-
-# DeCap with RegionCLIP text encoder:
-#python decoderTraining.py --out_dir weights_regionclip_karpathy --not-distributed 1 --local-rank 0 --dataset coco_train_karpathy.json --prefix coco_karpathy --use-regionclip
-
-# DeCap with DenseClip text encoder:
-#python decoderTraining.py --out_dir weights_denseclip_segmentation_vitb16_karpathy --not-distributed 1 --local-rank 0 --dataset coco_train_karpathy.json --prefix coco_karpathy --denseclip-config denseclip_segmentation_vitb16
-
-
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--decap_weights', type=str, default=None, help="If setted the Decap initialization is not random")
-    parser.add_argument('--clip_model', type=str, default='ViT-B/16', help="CLIP configuration")
-    parser.add_argument('--use_dinotxt', default=None, action='store_true', help="CLIP configuration")
-    parser.add_argument('--gaussian_noise', type=float, default=0, help="Standard deviation of the Gaussian noise to apply to the text input")
-    parser.add_argument('--out_dir', default='./coco_model')
-    parser.add_argument('--prefix', default='./coco_prefix', help='prefix for saved filenames')
-    parser.add_argument('--dataset', default='coco', help='coco or cc3m or bookcorpus')
-    parser.add_argument('--epochs', type=int, default=10)
-    parser.add_argument('--save_every', type=int, default=1)
-    parser.add_argument('--prefix_length', type=int, default=1)
-    parser.add_argument('--prefix_length_clip', type=int, default=1)
-    parser.add_argument('--bs', type=int, default=64)
-    parser.add_argument('--talk2dino_weights', type=str, default=None, help="Talk2DINO weights. If None, the training will be performed without Talk2DINO.")
-    parser.add_argument('--talk2dino_config', type=str, default=None, help="Talk2DINO configs. Valid only if the weights are setted.")
-    parser.add_argument('--use_dino_feats', action="store_true", default=False, help="If setted, we use the pre-extracted features of DINOv2")
-    parser.add_argument('--im_proj', action="store_true", default=False, help="If setted, we use the projection on the input features")
-    parser.add_argument('--pre_extract_features', action="store_true", default=False, help="If setted, the features will be extracted during the dataloading")
-    parser.add_argument('--only_prefix', dest='only_prefix', action='store_true')
-    parser.add_argument('--mapping_type', type=str, default='mlp', help='mlp/transformer')
-    parser.add_argument('--num_layers', type=int, default=8)
-    parser.add_argument('--is_rn', dest='is_rn', action='store_true')
-    parser.add_argument('--normalize_prefix', dest='normalize_prefix', action='store_true')
-    parser.add_argument('--local-rank', type=int, default=-1, metavar='N', help='Local process rank.') 
-    parser.add_argument('--not-distributed', type=int, default=False, metavar='N', help='Not Distributed toggle.')
-    parser.add_argument('--use-open-clip', action='store_true', default=False, help='Use OpenCLIP instead of CLIP')
-    parser.add_argument('--mix-captions', action='store_true', default=False, help='Mix captions from the same image')
-    parser.add_argument('--use-regionclip', action='store_true', default=False, help='Use RegionCLIP for text encoding')
-    parser.add_argument('--regionclip-checkpoint', type=str, default='/raid/datasets/models_weights/regionclip/regionclip_pretrained-cc_rn50x4.pth', help='Path to RegionCLIP checkpoint file')
-    parser.add_argument('--regionclip-config', type=str, default='pretrain/RegionCLIP_RN50x4.yaml', help='Path to RegionCLIP config file or config name')
-    parser.add_argument('--denseclip-config', type=str, default=None, help='Path to DenseClip config file or config name')
-    args = parser.parse_args()
-    
-    # Validate RegionCLIP arguments
-    if args.use_regionclip and args.regionclip_checkpoint is None:
-        parser.error("--regionclip-checkpoint is required when using --use-regionclip")
-    
-    if args.use_regionclip and args.use_open_clip:
-        parser.error("Cannot use both --use-regionclip and --use-open-clip at the same time")
-    
-    # Validate DenseClip arguments
-    if args.denseclip_config is not None and args.use_regionclip:
-        parser.error("Cannot use both --denseclip-config and --use-regionclip at the same time")
-    
-    if args.denseclip_config is not None and args.use_open_clip:
-        parser.error("Cannot use both --denseclip-config and --use-open-clip at the same time")
-    
-
-    train_decoder(args, output_dir=args.out_dir, output_prefix=args.prefix)
-
-
-if __name__ == '__main__':
-    main()

src/decap/decoder_config.pkl

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

src/decap/im2txtprojection/im2txtprojection.py

@@ -1,500 +0,0 @@
-from enum import Enum
-import numpy as np
-import math
-import json
-import random
-import torch
-from tqdm import tqdm
-import os
-import h5py
-from typing import Tuple
-from dotenv import load_dotenv
-from src.dinotxt_utils import get_tokenizer
-
-load_dotenv()
-
-class ProjectionType(Enum):
-    COCO_CAPTIONS       = 'coco_captions'
-    MS_MARCO_QUERIES_A  = 'ms_marco_queries_a'
-    CC3M_BLIP           = 'cc3m_blip_captions'
-    VISUAL_GENOME       = 'vg_captions'
-    VISUAL_GENOME_TEST  = "vg_dense_captions_test"
-    ONLINE_TEXTS        = "online_texts"
-
-class Im2TxtProjector:
-    """
-    Im2TxtProjector creates and manages text embedding memory banks for different models:
-    - Standard CLIP models
-    - OpenCLIP models  
-    - RegionCLIP models
-    - DenseClip models
-    - Talk2DINO projected embeddings
-    
-    For RegionCLIP usage, pass regionclip_config as a dict with:
-    {
-        'checkpoint': '/path/to/regionclip_checkpoint.pth',
-        'config_name': 'RegionCLIP_RN50.yaml'  # optional
-    }
-    
-    For DenseClip usage, pass denseclip_config as a string with the config file name:
-    'denseclip_vitb16'  # or other valid DenseClip config name
-    """
-
-    SUPPORT_MEMORY_SIZE = 500000
-
-    __IM2TXT_MEMORY_PATH = os.getenv("IM2TXT_MEMORY_PATH")
-    
-    if __IM2TXT_MEMORY_PATH is None:
-        default_path = "weights/im2txtmemories" #os.path.join(os.path.dirname(__file__), "../../../im2txtmemories")
-        print(f"[!] Warning: IM2TXT_MEMORY_PATH not set in environment variables, using '{default_path}' [!]")
-        __IM2TXT_MEMORY_PATH = default_path
-    
-    __DECAP_FOLDER = os.path.join(os.path.dirname(__file__), "../")
-    __TALK2DINO_CONFIG_WEIGHTS_PATH = __DECAP_FOLDER
-
-    captions_dataType = 'train2017'
-    ANNOTATIONS_CAPTION_FILE_PATH = os.path.join(__DECAP_FOLDER, 'captions_{}.json'.format(captions_dataType))
-    VG_ANNOTATIONS_DENSE_CAPTIONS_FILE_PATH = '/raid/datasets/densecaptioning-annotations/data/vg/controlcap/vg1.2/train.json'
-    VG_ANNOTATIONS_DENSE_CAPTIONS_TEST_FILE_PATH = '/raid/datasets/densecaptioning-annotations/data/vg/controlcap/vg1.2/test.json'
-
-    CC3M_BLIP_FILE_PATH = os.path.join(__DECAP_FOLDER, "blipv2_captions.txt")
-    MS_MARCO_QUERIES_FILE_PATH = '/raid/datasets/MSMarco/queries/queries.train.tsv'
-
-    @staticmethod
-    def create_regionclip_config(checkpoint_path: str, config_name: str = None):
-        """
-        Helper method to create RegionCLIP configuration dictionary.
-        
-        Args:
-            checkpoint_path (str): Path to RegionCLIP checkpoint file
-            config_name (str, optional): RegionCLIP config name (e.g., 'RegionCLIP_RN50.yaml')
-        
-        Returns:
-            dict: Configuration dictionary for RegionCLIP
-        """
-        return {
-            'checkpoint': checkpoint_path,
-            'config_name': config_name
-        }
-
-    def __init__(self, type = ProjectionType.COCO_CAPTIONS, verbose : bool = True, device_str = "cpu", use_talk2dino : bool = True, 
-                 support_memory_size : int = SUPPORT_MEMORY_SIZE, batch_size=1000, 
-                 clip_modelname = None, linear_talk2dino : bool = False, 
-                 normalize_memory_embs : bool = False, talk2dino_attn_type='qkv', online_texts=None,
-                 memory_bank_name = None, use_open_clip = False, regionclip_config=None, invite_config=None, denseclip_config=None) -> None:
-        """
-        - normalize_memory_embs -> normalizes the embeddings memory (required for projection in CLIP space)
-        - type : ProjectionType -> the type of the support memory to be built . Can either be the path to the file containing the captions or the type of the support memory to be built
-
-        """
-        # check if hdf5 already exists, otherwhise builds the support memory for that kind
-        
-        #if type not in ProjectionType.mro()
-
-        self.type = type
-        self.device_str = device_str
-        self.device = torch.device(self.device_str)
-        self.use_talk2dino = use_talk2dino
-        self.linear_talk2dino = linear_talk2dino
-        self.talk2dino_attn_type = talk2dino_attn_type
-        self.online_texts = online_texts
-        self.use_open_clip = use_open_clip
-        self.regionclip_config = regionclip_config
-        self.invite_config = invite_config
-        self.denseclip_config = denseclip_config
-        
-        if use_open_clip:
-            assert use_talk2dino is False, "use_open_clip and use_talk2dino cannot be used together"
-        
-        if regionclip_config is not None:
-            assert use_talk2dino is False, "regionclip_config and use_talk2dino cannot be used together"
-            assert use_open_clip is False, "regionclip_config and use_open_clip cannot be used together"
-
-        if invite_config is not None:
-            # overwrite clip_modelname with invite_config['name'] if provided
-            clip_modelname = invite_config.get('name', clip_modelname)
-            assert use_talk2dino is False, "invite_config and use_talk2dino cannot be used together"
-            
-        if denseclip_config is not None:
-            assert use_talk2dino is False, "denseclip_config and use_talk2dino cannot be used together"
-            assert use_open_clip is False, "denseclip_config and use_open_clip cannot be used together"
-            assert regionclip_config is None, "denseclip_config and regionclip_config cannot be used together"
-            
-
-        if clip_modelname is None:
-            if self.use_talk2dino:
-                clip_modelname = "ViT-B/16"
-            elif regionclip_config is not None:
-                # For RegionCLIP, we'll use a generic identifier since the model type is in the config
-                clip_modelname = "RegionCLIP"
-            elif denseclip_config is not None:
-                # For DenseClip, we'll use a generic identifier since the model type is in the config
-                clip_modelname = "DenseClip"
-            else:
-                clip_modelname = "ViT-B/32"
-        self.clip_modelname = clip_modelname
-
-        self.SUPPORT_MEMORY_SIZE = support_memory_size
-        if use_talk2dino:
-            prefix = ""
-            postfix = '-B16' if use_talk2dino is True else use_talk2dino
-            if linear_talk2dino:
-                postfix += "-linear"
-        elif regionclip_config is not None:
-            prefix = "regionclip-"
-            postfix = ""
-        elif denseclip_config is not None:
-            prefix = "denseclip-"
-            postfix = ""
-        else:
-            prefix = "clip-"
-            postfix = ""
-        if talk2dino_attn_type != 'qkv':
-            self.talk2dino_attn_type_str = f"_{talk2dino_attn_type}"
-        else:
-            self.talk2dino_attn_type_str = ''
-        if isinstance(type, ProjectionType):
-            dataset_name = type.value 
-        elif memory_bank_name is not None:
-            dataset_name = memory_bank_name
-        else:
-            dataset_name = 'coco'
-
-        if use_open_clip:
-            postfix += "-open_clip"
-        elif regionclip_config is not None:
-            postfix += "-regionclip"
-            # Add checkpoint identifier to make filename unique
-            checkpoint_path = regionclip_config.get('checkpoint', '')
-            checkpoint_name = os.path.basename(checkpoint_path).replace('.pth', '').replace('.pt', '')
-            if checkpoint_name:
-                postfix += f"-{checkpoint_name}"
-        elif denseclip_config is not None:
-            postfix += "-denseclip"
-            # Add config identifier to make filename unique
-            config_name = os.path.basename(denseclip_config).replace('.yaml', '').replace('.yml', '')
-            if config_name:
-                postfix += f"-{config_name}"
-
-        self.H5PY_FILE_PATH = os.path.join( self.__IM2TXT_MEMORY_PATH, prefix + f'{dataset_name}{self.talk2dino_attn_type_str}_text_embeddings{postfix}-{clip_modelname.replace("/", ".")}-{self.SUPPORT_MEMORY_SIZE}.h5' )
-        self.H5PY_EMBEDDINGS_DATASET_NAME = '{}-embeddings'.format(dataset_name)
-        self.H5PY_TEXT_DATASET_NAME = '{}-text'.format(dataset_name)
-
-        embs_dataset, text_dataset = self._load_support_memory()
-
-        if text_dataset is None:
-            if verbose: 
-                model_type = "RegionCLIP" if regionclip_config is not None else ("DenseClip" if denseclip_config is not None else ("OpenCLIP" if use_open_clip else "CLIP"))
-                print(f"[+] Going to build support memory for the given data type: {type} using {model_type} [+]")
-            embs_dataset, text_dataset = self._build_support_memory(batch_size)
-            if verbose: print(f"[+] Done [+]")
-        
-        if self.type != ProjectionType.ONLINE_TEXTS:
-            embs_dataset, text_dataset = self._load_support_memory()
-        
-        print(f"[-] loaded memory from {os.path.abspath( self.H5PY_FILE_PATH )} [-]")
-        if regionclip_config is not None:
-            print(f"[-] Using RegionCLIP text embeddings from checkpoint: {regionclip_config.get('checkpoint', 'Unknown')} [-]")
-        elif denseclip_config is not None:
-            print(f"[-] Using DenseClip text embeddings from config: {denseclip_config} [-]")
-
-        self.text_dataset = text_dataset
-        self.embs_dataset = torch.tensor(embs_dataset[:]).to(self.device)
-        self.embs_dataset = self.embs_dataset[self.embs_dataset.norm(dim=-1) != 0]
-
-
-        if normalize_memory_embs:
-            self.embs_dataset /= self.embs_dataset.norm(dim=-1,keepdim=True).float()
-        
-
-
-    def project(self, image_embedding, temperature : float = 0.01, normalize : bool = False, return_argmax_text : bool = False, return_n_best_sims=None) -> torch.TensorType:
-        if not isinstance(image_embedding, torch.Tensor):
-            print(f"the type of image_embedding is '{type(image_embedding)}' converting it to torch tensor")
-            image_embedding = torch.tensor(image_embedding, dtype=torch.float).to(self.device)
-        
-        orig_device = image_embedding.device
-        
-        if image_embedding.device != self.device:
-            image_embedding = image_embedding.to(self.device)
-        
-        if image_embedding.dtype != float:
-            #print(f"[-] image_embedding.dtype is {image_embedding.dtype}, converting it to float [-]")
-            image_embedding = image_embedding.float()
-        
-        embs_dataset = self.embs_dataset / self.embs_dataset.norm(dim=-1, keepdim=True)
-        image_embedding /= image_embedding.norm(dim=-1,keepdim=True)
-        
-        sim = image_embedding@embs_dataset.T.float()
-        if return_argmax_text:
-            argmax_texts = [self.text_dataset[idx].decode() for idx in sim.argmax(dim=-1)]
-            if return_n_best_sims:
-                return argmax_texts, sim.sort(dim=-1, descending=True).values[:, :return_n_best_sims].tolist()
-            return argmax_texts
-        softmax_sim = (sim / temperature).softmax(dim=-1)
-        prefix_embedding = [email protected]_dataset.float()
-    
-        if normalize:
-            prefix_embedding /= prefix_embedding.norm(dim=-1,keepdim=True)
-        
-        if return_n_best_sims:
-            return prefix_embedding.to(orig_device), sim.sort(dim=-1, descending=True).values[:, :return_n_best_sims].tolist()
-
-        return prefix_embedding.to(orig_device)
-
-    def _load_support_memory(self) -> Tuple[np.ndarray, np.ndarray]:
-        if self.type == ProjectionType.ONLINE_TEXTS:
-            print(f"[-] _load_support_memory: support memory for provided texts will be constructed [-]")
-            return None, None          
-        if not os.path.exists(self.H5PY_FILE_PATH):
-            print(f"[-] _load_support_memory: the path '{self.H5PY_FILE_PATH}' does not exist [-]")
-            return None, None
-
-        with h5py.File(self.H5PY_FILE_PATH, 'r') as hf:
-
-            if self.H5PY_EMBEDDINGS_DATASET_NAME in hf:
-                embeddings_dataset = hf[self.H5PY_EMBEDDINGS_DATASET_NAME][:]
-                text_dataset = hf[self.H5PY_TEXT_DATASET_NAME][:]
-            else:
-                embeddings_dataset = None
-                text_dataset = None
-        if 'DINO.txt' in self.clip_modelname:
-            embeddings_dataset = embeddings_dataset[:, 1024:] # Get patch-aligned text embeddings
-        return embeddings_dataset, text_dataset
-
-
-        
-    def _build_support_memory(self, batch_size = 1000) -> Tuple[np.ndarray, np.ndarray]:
-        ## construct the support memory
-
-        self._load_models()
-
-        if self.type == ProjectionType.COCO_CAPTIONS:
-            from pycocotools.coco import COCO
-            coco_obj = COCO(Im2TxtProjector.ANNOTATIONS_CAPTION_FILE_PATH)
-            data = random.sample(list(coco_obj.anns.values()), k=self.SUPPORT_MEMORY_SIZE)
-            data = [ d['caption'] for d in data ]
-        elif self.type == ProjectionType.VISUAL_GENOME:
-            from pycocotools.coco import COCO
-            coco_obj = COCO(Im2TxtProjector.VG_ANNOTATIONS_DENSE_CAPTIONS_FILE_PATH)
-            # data = random.sample(list(coco_obj.anns.values()), k=self.SUPPORT_MEMORY_SIZE)
-            data = list(coco_obj.anns.values())[:self.SUPPORT_MEMORY_SIZE]
-            data = [ d['caption'] for d in data ]
-        elif self.type == ProjectionType.VISUAL_GENOME_TEST:
-            from pycocotools.coco import COCO
-            coco_obj = COCO(Im2TxtProjector.VG_ANNOTATIONS_DENSE_CAPTIONS_TEST_FILE_PATH)
-            # data = random.sample(list(coco_obj.anns.values()), k=self.SUPPORT_MEMORY_SIZE)
-            data = list(coco_obj.anns.values())[:self.SUPPORT_MEMORY_SIZE]
-            data = [ d['caption'] for d in data ]
-        elif self.type == ProjectionType.MS_MARCO_QUERIES_A:
-            print(f"Loading MSMarco queries from file ", Im2TxtProjector.MS_MARCO_QUERIES_FILE_PATH)
-            with open(Im2TxtProjector.MS_MARCO_QUERIES_FILE_PATH, "r") as input_file:
-                lines = input_file.readlines()
-            data = random.sample(lines, k=self.SUPPORT_MEMORY_SIZE)
-            data = [ d.split("\t")[1].replace("\n", "") for d in data ]
-            print(f"Loaded from file '{self.SUPPORT_MEMORY_SIZE}' lines, example of line: '{data[0]}'")
-        elif self.type == ProjectionType.CC3M_BLIP:
-            print(f"Loading cc3m captions txt file ", Im2TxtProjector.CC3M_BLIP_FILE_PATH)
-            with open(Im2TxtProjector.CC3M_BLIP_FILE_PATH, "r") as input_file:
-                lines = input_file.readlines()
-            data = random.sample(lines, k=self.SUPPORT_MEMORY_SIZE)
-            data = [ d.replace("\n", "") for d in data ]
-            print(f"Loaded from file '{len(data)}' lines, example of line: '{data[0]}'")
-        elif self.type == ProjectionType.CC3M_BLIP:
-            print(f"Loading cc3m captions txt file ", Im2TxtProjector.CC3M_BLIP_FILE_PATH)
-            with open(Im2TxtProjector.CC3M_BLIP_FILE_PATH, "r") as input_file:
-                lines = input_file.readlines()
-            data = random.sample(lines, k=self.SUPPORT_MEMORY_SIZE)
-            data = [ d.replace("\n", "") for d in data ]
-            print(f"Loaded from file '{len(data)}' lines, example of line: '{data[0]}'") 
-        elif self.type == ProjectionType.ONLINE_TEXTS:
-            data = self.online_texts
-            print(f"Loaded online_texts '{len(data)}' lines, example of line: '{data[0]}'") 
-        elif type(self.type) == str:
-            if os.path.exists(self.type):
-                path = self.type
-                from pycocotools.coco import COCO
-                coco_obj = COCO(path)
-                data = random.sample(list(coco_obj.anns.values()), k=min(self.SUPPORT_MEMORY_SIZE, len(coco_obj.anns)))
-                data = [ d['caption'] for d in data ]
-        else:
-            #data = random.sample(data,500000)
-            print(f"[!] Unimplemented data type '{self.type}'[!]")
-            return None, None
-        
-        text_features = []
-        captions = []
-        
-        self.clip_model.eval()
-        
-        n_txts = len(data)
-        n_batch = math.ceil(n_txts / batch_size)
-        for i in tqdm(range(n_batch)):
-            start = i * batch_size
-            end = start + batch_size if i < n_batch - 1 else n_txts
-            
-            texts = data[start:end]
-            with torch.no_grad():
-                texts_token = self.tokenizer(texts).to(self.device)
-                text_feature = self.clip_model.encode_text(texts_token)
-                if self.use_talk2dino:
-                    text_feature = self.talk2dino.project_clip_txt(text_feature)
-                text_features.append(text_feature)
-                captions.extend(texts)
-
-        text_features = torch.cat(text_features,dim=0)
-
-        
-        #text_features /= text_features.norm(dim=-1,keepdim=True).float()
-        
-        # store captions and text features in hdf5 dataset
-
-        text_features_ndarray = text_features.cpu().numpy()
-
-        assert len(text_features_ndarray) == len(captions), f"len(text_features_ndarray) = {len(text_features_ndarray)} != len(captions) = {len(captions)}"
-
-        #if not os.path.exists(self.H5PY_FILE_PATH):
-        #    print(f"os.path '{self.H5PY_FILE_PATH}' does not exists")
-
-        EMBEDDINGS_DIMENSION = text_features_ndarray.shape[1]
-            
-        if self.type != ProjectionType.ONLINE_TEXTS:
-            with h5py.File(self.H5PY_FILE_PATH, 'w') as hf:
-
-                if self.H5PY_EMBEDDINGS_DATASET_NAME in hf:
-                    embeddings_dataset = hf[self.H5PY_EMBEDDINGS_DATASET_NAME]
-                    text_dataset = hf[self.H5PY_TEXT_DATASET_NAME]
-                    print(f"[!] Dataset '{self.H5PY_EMBEDDINGS_DATASET_NAME}' already exists! Going to overwrite [!]")
-                else:
-                    embeddings_dataset = hf.create_dataset(self.H5PY_EMBEDDINGS_DATASET_NAME, shape=(self.SUPPORT_MEMORY_SIZE, EMBEDDINGS_DIMENSION), dtype='float32')
-                    text_dataset = hf.create_dataset(self.H5PY_TEXT_DATASET_NAME, shape=(self.SUPPORT_MEMORY_SIZE, ), dtype=h5py.string_dtype(encoding='utf-8'))    #, dtype='str'
-            
-                for num_row in range(len(text_features_ndarray)):
-                    embeddings_dataset[num_row] = text_features_ndarray[num_row]
-                    text_dataset[num_row] = captions[num_row]
-        else:
-            embeddings_dataset = text_features_ndarray
-            text_dataset = [x.encode() for x in captions]
-
-        return embeddings_dataset, text_dataset
-
-    clip_model = None
-    def _load_models(self):
-
-        if self.clip_model is not None:
-            # case already done
-            return
-        
-        if self.use_open_clip:
-            print("[-] loading open_clip model [-]")
-            assert self.clip_modelname is not None, "clip_modelname must be provided when using open_clip"
-            from open_clip import create_model_and_transforms, tokenize
-            self.clip_model, preprocess_train, preprocess_val = create_model_and_transforms(self.clip_modelname, pretrained="laion2b_s32b_b79k", device=self.device)
-            self.preprocess = preprocess_train
-            self.tokenizer = tokenize
-            return
-        
-        if self.regionclip_config is not None:
-            print("[-] loading RegionCLIP model [-]")
-            from src.regionclip.loader import load_regionclip_from_checkpoint
-            from src.regionclip.datasets.clip_prompt_utils import tokenize as regionclip_tokenize
-            
-            regionclip_checkpoint = self.regionclip_config.get('checkpoint', None)
-            if regionclip_checkpoint is None:
-                raise ValueError("RegionCLIP checkpoint not specified in the configuration")
-            regionclip_config_name = self.regionclip_config.get('config_name', None)
-            
-            print(f"[-] Loading RegionCLIP from checkpoint: {regionclip_checkpoint} [-]")
-            if regionclip_config_name:
-                print(f"[-] Using RegionCLIP config: {regionclip_config_name} [-]")
-            
-            self.clip_model = load_regionclip_from_checkpoint(
-                regionclip_checkpoint, 
-                device=self.device, 
-                config=regionclip_config_name
-            )
-            self.tokenizer = regionclip_tokenize
-            self.preprocess = None  # RegionCLIP doesn't need preprocessing for text encoding
-            
-            # Test RegionCLIP text encoding to ensure it works
-            try:
-                test_text = ["A test sentence for RegionCLIP"]
-                test_tokens = self.tokenizer(test_text)
-                test_features = self.clip_model.encode_text(test_tokens.to(self.device))
-                print(f"[-] RegionCLIP text encoding test successful. Output shape: {test_features.shape} [-]")
-            except Exception as e:
-                print(f"[!] Warning: RegionCLIP text encoding test failed: {e} [!]")
-                raise e
-            
-            return
-
-        if self.denseclip_config is not None:
-            print("[-] loading DenseClip model [-]")
-            from src.denseclip.loader import load_denseclip, DenseCLIP_tokenize
-            
-            print(f"[-] Loading DenseClip from config: {self.denseclip_config} [-]")
-            
-            # Load DenseClip model
-            self.clip_model = load_denseclip(
-                config_name=self.denseclip_config,
-                device=self.device                
-            )
-            
-            # DenseClip should have encode_text method and a tokenizer
-            # We need to check if DenseClip has a tokenizer method
-            if DenseCLIP_tokenize is not None:
-                self.tokenizer = DenseCLIP_tokenize
-            else:
-                # Fallback to CLIP tokenizer if DenseClip doesn't provide one
-                import clip
-                self.tokenizer = clip.tokenize
-                print("[!] Warning: DenseClip model doesn't have tokenizer, using CLIP tokenizer [!]")
-            
-            self.preprocess = None  # DenseClip doesn't need preprocessing for text encoding
-            
-            # Test DenseClip text encoding to ensure it works
-            try:
-                test_text = ["A test sentence for DenseClip"]
-                test_tokens = self.tokenizer(test_text)
-                if hasattr(test_tokens, 'to'):
-                    test_tokens = test_tokens.to(self.device)
-                test_features = self.clip_model.encode_text(test_tokens)
-                print(f"[-] DenseClip text encoding test successful. Output shape: {test_features.shape} [-]")
-            except Exception as e:
-                print(f"[!] Warning: DenseClip text encoding test failed: {e} [!]")
-                raise e
-            
-            return
-
-        import clip
-        if self.clip_modelname is None:
-            clip_model_name = "ViT-B/16" if self.use_talk2dino else "ViT-B/32"
-        else:
-            clip_model_name = self.clip_modelname
-        if 'DINO.txt' not in clip_model_name:
-            self.clip_model, self.preprocess = clip.load(clip_model_name, device=self.device, jit=False)
-            self.tokenizer = clip.tokenize
-            if self.use_talk2dino:
-                # loading Talk2DINO
-                if type(self.use_talk2dino) == str:
-                    proj_name = self.use_talk2dino
-                elif self.linear_talk2dino is False:
-                    proj_name = 'vitb_mlp_infonce'
-                else:
-                    proj_name = 'vitb_linear_infonce'
-                
-                
-                config = os.path.join(self.__TALK2DINO_CONFIG_WEIGHTS_PATH, "configs_talk2dino", proj_name + '.yaml')
-                weights = os.path.join(self.__TALK2DINO_CONFIG_WEIGHTS_PATH, "weights_talk2dino", proj_name + self.talk2dino_attn_type_str + '.pth')
-                #import sys
-                #import os
-                #add_path = os.path.abspath( os.path.dirname("../"))
-                ##print(add_path)
-                #sys.path.insert(1, add_path )
-                from src.model import ProjectionLayer
-                self.talk2dino = ProjectionLayer.from_config(config)
-                self.talk2dino.load_state_dict(torch.load((weights), self.device))
-                self.talk2dino.to(self.device)
-        else:
-            self.clip_model = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitl14_reg4_dinotxt_tet1280d20h24l').to(self.device)
-            self.tokenizer = get_tokenizer().tokenize

src/denseclip/clip_loader/README.md

@@ -1,233 +0,0 @@
-# DenseCLIP to CLIP Loader
-
-A simple interface for loading DenseCLIP checkpoints as CLIP-like models for text and image encoding.
-
-## Overview
-
-This module provides a clean API to load DenseCLIP models and use them like standard CLIP models for encoding text and images. It abstracts away the complexity of DenseCLIP's detection/segmentation components and exposes only the core vision-language encoding functionality.
-
-## Features
-
-- ✅ **Simple API**: Load DenseCLIP models with a single function call
-- ✅ **CLIP-like Interface**: Familiar `encode_text()` and `encode_image()` methods
-- ✅ **Flexible Configuration**: YAML-based configuration system
-- ✅ **Multiple Input Types**: Support for PIL Images, image tensors, strings, and text lists
-- ✅ **Automatic Preprocessing**: Built-in image preprocessing pipeline
-- ✅ **Device Management**: Automatic GPU/CPU detection and placement
-
-## Quick Start
-
-```python
-from clip_loader import load_clip
-
-# Load DenseCLIP model with default configuration
-model = load_clip('denseclip_segmentation_vitb16')
-
-# Encode text
-texts = ["a photo of a cat", "a photo of a dog"]
-text_features = model.encode_text(texts)
-
-# Encode images (PIL Images)
-from PIL import Image
-images = [Image.open("cat.jpg"), Image.open("dog.jpg")]
-image_features = model.encode_image(images)
-
-# Compute similarities
-similarities = model.compute_similarity(image_features, text_features)
-print(f"Image-text similarities: {similarities}")
-```
-
-## Configuration
-
-Models are configured using YAML files in the `configs/` directory. The main configuration for DenseCLIP ViT-B/16 is in `configs/denseclip_vitb16.yaml`.
-
-### Configuration Structure
-
-```yaml
-model:
-  name: "denseclip_vitb16"
-  type: "vit"
-  
-  vision:
-    image_resolution: 224
-    vision_layers: 12
-    vision_width: 768
-    vision_patch_size: 16
-    embed_dim: 512
-  
-  text:
-    context_length: 13  # DenseCLIP uses shorter context
-    vocab_size: 49408
-    transformer_width: 512
-    transformer_heads: 8
-    transformer_layers: 12
-    embed_dim: 512
-
-checkpoint:
-  path: "/path/to/denseclip/checkpoint.pth"
-  format: "denseclip"
-  
-preprocessing:
-  image_mean: [0.48145466, 0.4578275, 0.40821073]
-  image_std: [0.26862954, 0.26130258, 0.27577711]
-  normalize: true
-```
-
-## API Reference
-
-### Core Functions
-
-#### `load_clip(config_name, checkpoint_path=None, device='auto')`
-
-Load a DenseCLIP model with the specified configuration.
-
-**Parameters:**
-- `config_name` (str): Name of config file (without .yaml extension)
-- `checkpoint_path` (str, optional): Path to checkpoint file (overrides config)
-- `device` (str): Device to load on ('auto', 'cpu', 'cuda')
-
-**Returns:**
-- `DenseCLIPModel`: Loaded model ready for inference
-
-#### `load_denseclip_model(config_path, checkpoint_path=None, device='auto')`
-
-Load a DenseCLIP model from configuration file path.
-
-### DenseCLIPModel Methods
-
-#### `encode_text(texts)`
-
-Encode text into feature vectors.
-
-**Parameters:**
-- `texts` (str or List[str]): Text string(s) to encode
-
-**Returns:**
-- `torch.Tensor`: Normalized text features [batch_size, embed_dim]
-
-#### `encode_image(images)`
-
-Encode images into feature vectors.
-
-**Parameters:**
-- `images`: PIL Image, List[PIL.Image], or preprocessed tensor
-
-**Returns:**
-- `torch.Tensor`: Normalized image features [batch_size, embed_dim]
-
-#### `compute_similarity(image_features, text_features, temperature=1.0)`
-
-Compute similarity between image and text features.
-
-**Parameters:**
-- `image_features` (torch.Tensor): Image features [N, embed_dim]
-- `text_features` (torch.Tensor): Text features [M, embed_dim]
-- `temperature` (float): Temperature scaling factor
-
-**Returns:**
-- `torch.Tensor`: Similarity matrix [N, M]
-
-## Examples
-
-### Basic Text-Image Retrieval
-
-```python
-from clip_loader import load_clip
-from PIL import Image
-
-# Load model
-model = load_clip('denseclip_vitb16')
-
-# Load and encode images
-images = [
-    Image.open("cat.jpg"),
-    Image.open("dog.jpg"),
-    Image.open("car.jpg")
-]
-image_features = model.encode_image(images)
-
-# Encode text queries
-queries = [
-    "a cute cat",
-    "a happy dog", 
-    "a red car"
-]
-text_features = model.encode_text(queries)
-
-# Find best matches
-similarities = model.compute_similarity(image_features, text_features)
-best_matches = similarities.argmax(dim=1)
-
-for i, query in enumerate(queries):
-    best_image_idx = best_matches[i]
-    score = similarities[best_image_idx, i].item()
-    print(f"Query '{query}' -> Image {best_image_idx} (score: {score:.3f})")
-```
-
-### Zero-Shot Classification
-
-```python
-from clip_loader import load_clip
-from PIL import Image
-
-model = load_clip('denseclip_vitb16')
-
-# Load test image
-image = Image.open("test_image.jpg")
-image_features = model.encode_image(image)
-
-# Define class labels
-class_labels = [
-    "a photo of a cat",
-    "a photo of a dog", 
-    "a photo of a bird",
-    "a photo of a car",
-    "a photo of a house"
-]
-
-# Encode labels
-text_features = model.encode_text(class_labels)
-
-# Classify
-similarities = model.compute_similarity(image_features, text_features)
-probabilities = similarities.softmax(dim=-1)
-
-# Show results
-for i, label in enumerate(class_labels):
-    prob = probabilities[0, i].item()
-    print(f"{label}: {prob:.3f}")
-```
-
-### Custom Configuration
-
-```python
-from clip_loader import load_denseclip_model
-
-# Load with custom config
-model = load_denseclip_model(
-    config_path='configs/custom_config.yaml',
-    checkpoint_path='/path/to/custom/checkpoint.pth',
-    device='cuda:1'
-)
-```
-
-## Requirements
-
-- PyTorch >= 1.9.0
-- torchvision >= 0.10.0
-- Pillow >= 8.0.0
-- PyYAML >= 5.4.0
-
-## Notes
-
-- DenseCLIP uses a shorter text context length (13) compared to standard CLIP (77)
-- The model preserves ~98% similarity with original CLIP text representations
-- Image preprocessing follows CLIP's standard normalization
-- All features are L2-normalized for cosine similarity computation
-
-## Supported Models
-
-Currently supported:
-- `denseclip_vitb16`: DenseCLIP with ViT-B/16 backbone
-
-To add support for other DenseCLIP variants, create new configuration files in the `configs/` directory.

src/denseclip/clip_loader/SUMMARY.md

@@ -1,78 +0,0 @@
-# DenseCLIP to CLIP Loader - Quick Start
-
-## ✅ Successfully Created!
-
-The `clip_loader` module provides a simple interface to load DenseCLIP checkpoints as CLIP-like models for text and image encoding.
-
-## 📁 Structure
-
-```
-/raid/homes/giacomo.pacini/DenseCLIP/clip_loader/
-├── __init__.py                    # Module initialization
-├── denseclip_loader.py           # Main loader implementation
-├── example_usage.py              # Example script
-├── requirements.txt              # Dependencies
-├── README.md                     # Full documentation
-└── configs/
-    └── denseclip_vitb16.yaml     # Configuration for ViT-B/16 model
-```
-
-## 🚀 Quick Usage
-
-```python
-from clip_loader import load_clip
-
-# Load DenseCLIP model with default configuration
-model = load_clip('denseclip_vitb16')
-
-# Encode text
-texts = ["a photo of a cat", "a photo of a dog"]
-text_features = model.encode_text(texts)  # Shape: [2, 512]
-
-# Encode images (if you have PIL Images)
-# image_features = model.encode_image(images)
-
-# Compute similarities
-similarities = model.compute_similarity(text_features, text_features)
-print(f"Cat-Dog similarity: {similarities[0, 1]:.3f}")
-```
-
-## ✅ Test Results
-
-- **✅ Model loads successfully** from DenseCLIP checkpoint
-- **✅ Text encoding works** (shape: [batch_size, 512])  
-- **✅ Features are normalized** (L2 norm = 1.0)
-- **✅ Similarities make sense** (Cat-Dog: 0.872, Car-Person: lower)
-- **✅ Zero-shot classification** shows logical patterns
-- **✅ Model has 157M parameters** (94M vision + 63M text)
-
-## 🔧 Key Features
-
-- **Simple API**: Just call `load_clip()` and start encoding
-- **Handles DenseCLIP specifics**: Automatically extracts weights from segmentation checkpoint
-- **CLIP-compatible**: Same interface as OpenAI CLIP
-- **Flexible configuration**: YAML-based configuration system
-- **GPU ready**: Automatic device detection and placement
-- **Context length**: Uses DenseCLIP's shorter context (13 vs 77)
-
-## 🎯 Use Cases
-
-1. **Text-Image Retrieval**: Encode both and compute similarities
-2. **Zero-Shot Classification**: Encode class descriptions and compare
-3. **Text Similarity**: Compare text representations  
-4. **Feature Extraction**: Get dense vector representations
-
-## 📝 Configuration
-
-The model uses `/raid/datasets/models_weights/denseclip/segmentation/semanticFPN/ViT-B-DenseCLIP.pth` by default. You can override this by modifying `configs/denseclip_vitb16.yaml` or passing a custom checkpoint path.
-
-## 🔍 What's Different from Standard CLIP
-
-- **Shorter context length**: 13 tokens vs 77
-- **Higher image resolution**: 640px vs 224px  
-- **Fine-tuned weights**: Adapted for dense prediction tasks
-- **High text similarity**: ~98% similarity with original CLIP representations
-
-## 🎉 Ready to Use!
-
-The loader is fully functional and ready for use in your projects. See `README.md` for detailed documentation and more examples.

src/denseclip/clip_loader/__init__.py

@@ -1,21 +0,0 @@
-"""
-DenseCLIP to CLIP Loader
-
-A simple interface for loading DenseCLIP checkpoints as CLIP-like models
-for text and image encoding.
-"""
-
-from .denseclip_loader import (
-    DenseCLIPModel,
-    load_denseclip_model,
-    load_clip,
-    load_config
-)
-
-__version__ = "1.0.0"
-__all__ = [
-    "DenseCLIPModel",
-    "load_denseclip_model", 
-    "load_clip",
-    "load_config"
-]

src/denseclip/clip_loader/bpe_simple_vocab_16e6.txt.gz

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

src/denseclip/clip_loader/configs/denseclip_segmentation_vitb16.yaml

@@ -1,41 +0,0 @@
-# DenseCLIP ViT-B/16 Configuration
-# Configuration for loading DenseCLIP checkpoint as a CLIP-like model
-
-model:
-  name: "denseclip_vitb16"
-  type: "vit"  # vision transformer
-  
-  # Vision encoder configuration
-  vision:
-    image_resolution: 640
-    vision_layers: 12
-    vision_width: 768
-    vision_patch_size: 16
-    embed_dim: 512
-  
-  # Text encoder configuration  
-  text:
-    context_length: 13  # DenseCLIP uses shorter context
-    vocab_size: 49408
-    transformer_width: 512
-    transformer_heads: 8
-    transformer_layers: 12
-    embed_dim: 512
-
-# Checkpoint information
-checkpoint:
-  path: "/raid/datasets/models_weights/denseclip/segmentation/semanticFPN/ViT-B-DenseCLIP.pth"
-  format: "denseclip"  # vs "openai_clip"
-  
-# Processing configuration
-preprocessing:
-  image_mean: [0.48145466, 0.4578275, 0.40821073]
-  image_std: [0.26862954, 0.26130258, 0.27577711]
-  normalize: true
-  
-# Optional overrides
-overrides:
-  # Set to true to use OpenAI CLIP tokenizer instead of DenseCLIP's
-  use_openai_tokenizer: false
-  # Set custom context length (will resize positional embeddings if needed)
-  custom_context_length: null

src/denseclip/clip_loader/configs/denseclip_segmentation_vitb16_long_ctx.yaml

@@ -1,41 +0,0 @@
-# DenseCLIP ViT-B/16 Configuration
-# Configuration for loading DenseCLIP checkpoint as a CLIP-like model
-
-model:
-  name: "denseclip_vitb16"
-  type: "vit"  # vision transformer
-  
-  # Vision encoder configuration
-  vision:
-    image_resolution: 640
-    vision_layers: 12
-    vision_width: 768
-    vision_patch_size: 16
-    embed_dim: 512
-  
-  # Text encoder configuration  
-  text:
-    context_length: 77  # DenseCLIP uses shorter context
-    vocab_size: 49408
-    transformer_width: 512
-    transformer_heads: 8
-    transformer_layers: 12
-    embed_dim: 512
-
-# Checkpoint information
-checkpoint:
-  path: "/raid/datasets/models_weights/denseclip/segmentation/semanticFPN/ViT-B-DenseCLIP_long_ctx.pth"
-  format: "denseclip"  # vs "openai_clip"
-  
-# Processing configuration
-preprocessing:
-  image_mean: [0.48145466, 0.4578275, 0.40821073]
-  image_std: [0.26862954, 0.26130258, 0.27577711]
-  normalize: true
-  
-# Optional overrides
-overrides:
-  # Set to true to use OpenAI CLIP tokenizer instead of DenseCLIP's
-  use_openai_tokenizer: false
-  # Set custom context length (will resize positional embeddings if needed)
-  custom_context_length: null

src/denseclip/clip_loader/denseclip_loader.py

@@ -1,316 +0,0 @@
-#!/usr/bin/env python3
-"""
-DenseCLIP to CLIP Loader
-
-A simple interface for loading DenseCLIP checkpoints as CLIP-like models
-for text and image encoding.
-"""
-
-import os
-import sys
-import yaml
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from typing import Union, List, Tuple, Optional, Dict, Any
-from PIL import Image
-import torchvision.transforms as transforms
-
-# Import local model components
-try:
-    from .models import CLIPVisionTransformer, CLIPTextEncoder, ResidualAttentionBlock, LayerNorm, QuickGELU
-    from .tokenizer import tokenize
-except ImportError:
-    # Fallback for direct execution
-    from models import CLIPVisionTransformer, CLIPTextEncoder, ResidualAttentionBlock, LayerNorm, QuickGELU
-    from tokenizer import tokenize
-
-
-class DenseCLIPModel(nn.Module):
-    """
-    A CLIP-like model loaded from DenseCLIP checkpoints.
-    Provides simple text and image encoding functionality.
-    """
-    
-    def __init__(self, config: Dict[str, Any]):
-        super().__init__()
-        
-        self.config = config
-        
-        # Initialize vision encoder
-        vision_config = config['model']['vision']
-        self.visual = CLIPVisionTransformer(
-            input_resolution=vision_config['image_resolution'],
-            patch_size=vision_config['vision_patch_size'],
-            width=vision_config['vision_width'],
-            layers=vision_config['vision_layers'],
-            heads=vision_config['vision_width'] // 64,
-            output_dim=vision_config['embed_dim']
-        )
-        
-        # Initialize text encoder
-        text_config = config['model']['text']
-        self.text_encoder = CLIPTextEncoder(
-            context_length=text_config['context_length'],
-            vocab_size=text_config['vocab_size'],
-            transformer_width=text_config['transformer_width'],
-            transformer_heads=text_config['transformer_heads'],
-            transformer_layers=text_config['transformer_layers'],
-            embed_dim=text_config['embed_dim']
-        )
-        
-        # Store configuration for preprocessing
-        self.context_length = text_config['context_length']
-        self.image_resolution = vision_config['image_resolution']
-        
-        # Initialize preprocessing
-        self._setup_preprocessing()
-        
-    def _setup_preprocessing(self):
-        """Setup image preprocessing pipeline"""
-        preprocess_config = self.config['preprocessing']
-        
-        self.preprocess = transforms.Compose([
-            transforms.Resize(self.image_resolution, interpolation=transforms.InterpolationMode.BICUBIC),
-            transforms.CenterCrop(self.image_resolution),
-            transforms.ToTensor(),
-            transforms.Normalize(
-                mean=preprocess_config['image_mean'],
-                std=preprocess_config['image_std']
-            )
-        ])
-    
-    def encode_image(self, images: Union[torch.Tensor, List[Image.Image], Image.Image]) -> torch.Tensor:
-        """
-        Encode images into feature vectors
-        
-        Args:
-            images: PIL Images, list of PIL Images, or preprocessed tensor
-            
-        Returns:
-            Normalized image features [batch_size, embed_dim]
-        """
-        if isinstance(images, (list, tuple)):
-            # List of PIL Images
-            image_tensors = torch.stack([self.preprocess(img) for img in images])
-        elif isinstance(images, Image.Image):
-            # Single PIL Image
-            image_tensors = self.preprocess(images).unsqueeze(0)
-        elif isinstance(images, torch.Tensor):
-            # Already preprocessed tensor
-            image_tensors = images
-        else:
-            raise ValueError(f"Unsupported image type: {type(images)}")
-        
-        # Move to same device as model
-        device = next(self.parameters()).device
-        image_tensors = image_tensors.to(device)
-        
-        # Encode
-        with torch.no_grad():
-            image_features = self.visual(image_tensors)
-            image_features = F.normalize(image_features, dim=-1)
-            
-        return image_features
-    
-    def encode_text(self, texts: Union[str, List[str]]) -> torch.Tensor:
-        """
-        Encode texts into feature vectors
-        
-        Args:
-            texts: Single text string or list of text strings
-            
-        Returns:
-            Normalized text features [batch_size, embed_dim]
-        """
-        if isinstance(texts, str):
-            texts = [texts]
-        
-        # Tokenize if necessary
-        if isinstance(texts, list):
-            tokens = tokenize(texts, context_length=self.context_length)
-        elif isinstance(texts, torch.Tensor):
-            if texts.dim() == 1:
-                # Single tokenized text
-                tokens = texts.unsqueeze(0)
-            else:
-                tokens = texts
-        else:
-            raise ValueError(f"Unsupported text type: {type(texts)}")    
-        # Move to same device as model
-        device = next(self.parameters()).device
-        tokens = tokens.to(device)
-        
-        # Encode
-        with torch.no_grad():
-            text_features = self.text_encoder(tokens)
-            text_features = F.normalize(text_features, dim=-1)
-            
-        return text_features
-    
-    def compute_similarity(self, 
-                         image_features: torch.Tensor, 
-                         text_features: torch.Tensor,
-                         temperature: float = 1.0) -> torch.Tensor:
-        """
-        Compute similarity between image and text features
-        
-        Args:
-            image_features: Normalized image features [N, embed_dim]
-            text_features: Normalized text features [M, embed_dim]
-            temperature: Temperature for scaling similarities
-            
-        Returns:
-            Similarity matrix [N, M]
-        """
-        return (image_features @ text_features.t()) / temperature
-    
-    def forward(self, images: torch.Tensor, texts: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
-        """
-        Forward pass for both image and text encoding
-        
-        Args:
-            images: Preprocessed image tensor [batch_size, 3, H, W]
-            texts: Tokenized text tensor [batch_size, context_length]
-            
-        Returns:
-            Tuple of (image_features, text_features)
-        """
-        image_features = self.visual(images)
-        text_features = self.text_encoder(texts)
-        
-        # Normalize features
-        image_features = F.normalize(image_features, dim=-1)
-        text_features = F.normalize(text_features, dim=-1)
-        
-        return image_features, text_features
-
-
-def load_config(config_path: str) -> Dict[str, Any]:
-    """Load configuration from YAML file"""
-    with open(config_path, 'r') as f:
-        config = yaml.safe_load(f)
-    return config
-
-
-def load_denseclip_weights(checkpoint_path: str) -> Dict[str, torch.Tensor]:
-    """Load DenseCLIP checkpoint and extract relevant weights"""
-    print(f"Loading DenseCLIP checkpoint from: {checkpoint_path}")
-    checkpoint = torch.load(checkpoint_path, map_location='cpu')
-    
-    if 'state_dict' not in checkpoint:
-        raise ValueError("Checkpoint doesn't contain 'state_dict'")
-    
-    state_dict = checkpoint['state_dict']
-    
-    # Extract vision and text encoder weights
-    vision_weights = {}
-    text_weights = {}
-    
-    for key, value in state_dict.items():
-        if key.startswith('backbone.'):
-            # Remove 'backbone.' prefix for vision encoder
-            new_key = key[len('backbone.'):]
-            vision_weights[new_key] = value
-        elif key.startswith('text_encoder.'):
-            # Remove 'text_encoder.' prefix  
-            new_key = key[len('text_encoder.'):]
-            text_weights[new_key] = value
-    
-    print(f"Extracted {len(vision_weights)} vision parameters")
-    print(f"Extracted {len(text_weights)} text parameters")
-    
-    return {
-        'vision': vision_weights,
-        'text': text_weights,
-        'full_state_dict': state_dict
-    }
-
-
-def load_denseclip_model(config_path: str, 
-                        checkpoint_path: Optional[str] = None,
-                        device: str = 'auto') -> DenseCLIPModel:
-    """
-    Load a DenseCLIP model from configuration and checkpoint
-    
-    Args:
-        config_path: Path to YAML configuration file
-        checkpoint_path: Optional path to checkpoint (overrides config)
-        device: Device to load model on ('auto', 'cpu', 'cuda')
-        
-    Returns:
-        Loaded DenseCLIPModel ready for inference
-    """
-    # Load configuration
-    config = load_config(config_path)
-    
-    # Override checkpoint path if provided
-    if checkpoint_path is not None:
-        config['checkpoint']['path'] = checkpoint_path
-    
-    # Create model
-    model = DenseCLIPModel(config)
-    
-    # Load weights
-    checkpoint_path = config['checkpoint']['path']
-    if os.path.exists(checkpoint_path):
-        weights = load_denseclip_weights(checkpoint_path)
-        
-        # Load vision encoder weights
-        if weights['vision']:
-            missing_v, unexpected_v = model.visual.load_state_dict(weights['vision'], strict=False)
-            if missing_v:
-                print(f"Missing vision keys: {len(missing_v)} (expected for FPN/post-norm components)")
-            if unexpected_v:
-                # Filter out expected mismatches
-                important_unexpected = [k for k in unexpected_v if not any(x in k for x in ['fpn', 'ln_post', 'proj'])]
-                if important_unexpected:
-                    print(f"Unexpected vision keys: {important_unexpected}")
-                else:
-                    print(f"✓ Vision weights loaded (ignoring {len(unexpected_v)} FPN/post-norm parameters)")
-        
-        # Load text encoder weights  
-        if weights['text']:
-            missing_t, unexpected_t = model.text_encoder.load_state_dict(weights['text'], strict=False)
-            if missing_t:
-                print(f"Missing text keys: {len(missing_t)}")
-            if unexpected_t:
-                print(f"Unexpected text keys: {unexpected_t}")
-        
-        print("✓ Model weights loaded successfully")
-    else:
-        print(f"⚠ Checkpoint not found at {checkpoint_path}, using random weights")
-    
-    # Setup device
-    if device == 'auto':
-        device = 'cuda' if torch.cuda.is_available() else 'cpu'
-    
-    model = model.to(device)
-    model.eval()
-    
-    print(f"✓ Model loaded on {device}")
-    return model
-
-
-# Convenience function
-def load_clip(config_name: str = 'denseclip_vitb16', 
-              checkpoint_path: Optional[str] = None,
-              device: str = 'auto') -> DenseCLIPModel:
-    """
-    Convenience function to load a DenseCLIP model
-    
-    Args:
-        config_name: Name of config file (without .yaml extension)
-        checkpoint_path: Optional path to checkpoint
-        device: Device to load on
-        
-    Returns:
-        Loaded DenseCLIPModel
-    """
-    current_dir = os.path.dirname(os.path.abspath(__file__))
-    config_path = os.path.join(current_dir, 'configs', f'{config_name}.yaml')
-    
-    if not os.path.exists(config_path):
-        raise FileNotFoundError(f"Config file not found: {config_path}")
-    
-    return load_denseclip_model(config_path, checkpoint_path, device)

src/denseclip/clip_loader/example_usage.py

@@ -1,108 +0,0 @@
-#!/usr/bin/env python3
-"""
-Example usage of the DenseCLIP to CLIP loader
-"""
-
-import sys
-import os
-
-# Add the clip_loader to path
-current_dir = os.path.dirname(os.path.abspath(__file__))
-sys.path.append(current_dir)
-
-from denseclip_loader import load_clip
-import torch
-
-
-def main():
-    print("🚀 DenseCLIP to CLIP Loader Example")
-    print("=" * 50)
-    
-    # Load model
-    print("Loading DenseCLIP model...")
-    try:
-        model = load_clip('denseclip_segmentation_vitb16')
-        print("✅ Model loaded successfully!")
-    except Exception as e:
-        print(f"❌ Error loading model: {e}")
-        return
-    
-    print(f"Model device: {next(model.parameters()).device}")
-    print(f"Text context length: {model.context_length}")
-    print(f"Image resolution: {model.image_resolution}")
-    
-    # Test text encoding
-    print("\n📝 Testing text encoding...")
-    texts = [
-        "a photo of a cat",
-        "a photo of a dog", 
-        "a photo of a car",
-        "a person walking",
-        "a beautiful sunset"
-    ]
-    
-    text_features = model.encode_text(texts)
-    print(f"Text features shape: {text_features.shape}")
-    print(f"Text features norm: {text_features.norm(dim=-1)}")  # Should be ~1.0 (normalized)
-    
-    # Test text-text similarities
-    print("\n🔍 Text-to-text similarities:")
-    text_similarities = model.compute_similarity(text_features, text_features)
-    
-    print(f"{'Text':<20} {'Self-sim':<10} {'vs Cat':<10} {'vs Dog':<10}")
-    print("-" * 50)
-    for i, text in enumerate(texts):
-        self_sim = text_similarities[i, i].item()
-        cat_sim = text_similarities[i, 0].item()
-        dog_sim = text_similarities[i, 1].item()
-        print(f"{text:<20} {self_sim:<10.3f} {cat_sim:<10.3f} {dog_sim:<10.3f}")
-    
-    # Test zero-shot classification concepts
-    print("\n🎯 Zero-shot classification example:")
-    test_queries = [
-        "an animal",
-        "a vehicle", 
-        "a person",
-        "nature scene"
-    ]
-    
-    query_features = model.encode_text(test_queries)
-    classification_similarities = model.compute_similarity(text_features, query_features)
-    
-    print(f"{'Original Text':<20} {'Animal':<8} {'Vehicle':<8} {'Person':<8} {'Nature':<8}")
-    print("-" * 60)
-    for i, text in enumerate(texts):
-        sims = classification_similarities[i]
-        print(f"{text:<20} {sims[0]:<8.3f} {sims[1]:<8.3f} {sims[2]:<8.3f} {sims[3]:<8.3f}")
-    
-    # Test feature statistics
-    print("\n📊 Feature statistics:")
-    print(f"Text feature mean: {text_features.mean():.6f}")
-    print(f"Text feature std: {text_features.std():.6f}")
-    print(f"Text feature min: {text_features.min():.6f}")
-    print(f"Text feature max: {text_features.max():.6f}")
-    
-    # Test model components
-    print("\n🔧 Model architecture:")
-    print(f"Vision encoder: {type(model.visual).__name__}")
-    print(f"Text encoder: {type(model.text_encoder).__name__}")
-    
-    # Count parameters
-    vision_params = sum(p.numel() for p in model.visual.parameters())
-    text_params = sum(p.numel() for p in model.text_encoder.parameters())
-    total_params = vision_params + text_params
-    
-    print(f"\n📈 Parameter count:")
-    print(f"Vision encoder: {vision_params:,}")
-    print(f"Text encoder: {text_params:,}")
-    print(f"Total: {total_params:,}")
-    
-    print("\n✅ All tests completed successfully!")
-    print("\n💡 Usage tip:")
-    print("   from clip_loader import load_clip")
-    print("   model = load_clip('denseclip_vitb16')")
-    print("   features = model.encode_text(['your text here'])")
-
-
-if __name__ == "__main__":
-    main()